tailscale/derp/derphttp/derphttp_test.go

157 lines
3.2 KiB
Go
Raw Permalink Normal View History

// Copyright (c) 2020 Tailscale Inc & AUTHORS All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package derphttp
import (
"context"
"crypto/tls"
"net"
"net/http"
"sync"
"testing"
"time"
"tailscale.com/derp"
"tailscale.com/types/key"
)
func TestSendRecv(t *testing.T) {
serverPrivateKey := key.NewNode()
const numClients = 3
var clientPrivateKeys []key.NodePrivate
var clientKeys []key.NodePublic
for i := 0; i < numClients; i++ {
priv := key.NewNode()
clientPrivateKeys = append(clientPrivateKeys, priv)
clientKeys = append(clientKeys, priv.Public())
}
s := derp.NewServer(serverPrivateKey, t.Logf)
defer s.Close()
httpsrv := &http.Server{
TLSNextProto: make(map[string]func(*http.Server, *tls.Conn, http.Handler)),
Handler: Handler(s),
}
ln, err := net.Listen("tcp4", "localhost:0")
if err != nil {
t.Fatal(err)
}
serverURL := "http://" + ln.Addr().String()
t.Logf("server URL: %s", serverURL)
go func() {
if err := httpsrv.Serve(ln); err != nil {
if err == http.ErrServerClosed {
return
}
panic(err)
}
}()
var clients []*Client
var recvChs []chan []byte
done := make(chan struct{})
var wg sync.WaitGroup
defer func() {
close(done)
for _, c := range clients {
c.Close()
}
wg.Wait()
}()
for i := 0; i < numClients; i++ {
key := clientPrivateKeys[i]
c, err := NewClient(key, serverURL, t.Logf)
if err != nil {
t.Fatalf("client %d: %v", i, err)
}
if err := c.Connect(context.Background()); err != nil {
t.Fatalf("client %d Connect: %v", i, err)
}
waitConnect(t, c)
clients = append(clients, c)
recvChs = append(recvChs, make(chan []byte))
wg.Add(1)
go func(i int) {
defer wg.Done()
for {
select {
case <-done:
return
default:
}
m, err := c.Recv()
if err != nil {
select {
case <-done:
return
default:
}
t.Logf("client%d: %v", i, err)
break
}
switch m := m.(type) {
default:
t.Errorf("unexpected message type %T", m)
continue
case derp.PeerGoneMessage:
// Ignore.
case derp.ReceivedPacket:
recvChs[i] <- append([]byte(nil), m.Data...)
}
}
}(i)
}
recv := func(i int, want string) {
t.Helper()
select {
case b := <-recvChs[i]:
if got := string(b); got != want {
t.Errorf("client1.Recv=%q, want %q", got, want)
}
case <-time.After(5 * time.Second):
t.Errorf("client%d.Recv, got nothing, want %q", i, want)
}
}
recvNothing := func(i int) {
t.Helper()
select {
case b := <-recvChs[0]:
t.Errorf("client%d.Recv=%q, want nothing", i, string(b))
default:
}
}
msg1 := []byte("hello 0->1\n")
if err := clients[0].Send(clientKeys[1], msg1); err != nil {
t.Fatal(err)
}
recv(1, string(msg1))
recvNothing(0)
recvNothing(2)
msg2 := []byte("hello 1->2\n")
if err := clients[1].Send(clientKeys[2], msg2); err != nil {
t.Fatal(err)
}
recv(2, string(msg2))
recvNothing(0)
recvNothing(1)
}
func waitConnect(t testing.TB, c *Client) {
t.Helper()
if m, err := c.Recv(); err != nil {
t.Fatalf("client first Recv: %v", err)
} else if v, ok := m.(derp.ServerInfoMessage); !ok {
t.Fatalf("client first Recv was unexpected type %T", v)
}
}