tailnet/derpclient_test.go

package tailnet

import (
	"context"
	"os"
	"testing"
	"time"

	"tailscale.com/types/key"
)

// TestDERPConnect verifies we can connect to a DERP relay obtained from MapPoll.
func TestDERPConnect(t *testing.T) {
	authKey := os.Getenv("TAILNET_AUTH_KEY")
	if authKey == "" {
		t.Skip("set TAILNET_AUTH_KEY to run DERP tests")
	}

	dir := t.TempDir()
	c, err := NewControlClient(ControlOptions{
		StateDir: dir,
		AuthKey:  authKey,
		Hostname: "derp-test",
		Logf:     t.Logf,
	})
	if err != nil {
		t.Fatalf("NewControlClient: %v", err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	// Register + MapPoll to get DERPMap
	if _, err := c.Register(ctx); err != nil {
		t.Fatalf("Register: %v", err)
	}
	ch, err := c.MapPoll(ctx)
	if err != nil {
		t.Fatalf("MapPoll: %v", err)
	}

	var derpMap = c.DERPMap()
	// Wait for a MapResponse with DERPMap
	for derpMap == nil {
		mr, ok := <-ch
		if !ok {
			t.Fatal("map channel closed without DERPMap")
		}
		if mr.DERPMap != nil {
			derpMap = mr.DERPMap
		}
	}
	t.Logf("got DERP map with %d regions", len(derpMap.Regions))

	// Pick first region
	var regionID int
	for id := range derpMap.Regions {
		regionID = id
		break
	}

	dc := NewDERPClient(c.nodeKey, t.Logf, nil)
	defer dc.Close()

	if err := dc.ConnectRegion(ctx, derpMap, regionID); err != nil {
		t.Fatalf("ConnectRegion %d: %v", regionID, err)
	}

	regions := dc.ConnectedRegions()
	if len(regions) != 1 || regions[0] != regionID {
		t.Fatalf("expected connected to region %d, got %v", regionID, regions)
	}
	t.Logf("connected to DERP region %d", regionID)
}

// TestDERPSendRecv verifies two clients can exchange packets via DERP.
func TestDERPSendRecv(t *testing.T) {
	authKey := os.Getenv("TAILNET_AUTH_KEY")
	if authKey == "" {
		t.Skip("set TAILNET_AUTH_KEY to run DERP tests")
	}

	dir := t.TempDir()
	ctrl, err := NewControlClient(ControlOptions{
		StateDir: dir,
		AuthKey:  authKey,
		Hostname: "derp-test-sr",
		Logf:     t.Logf,
	})
	if err != nil {
		t.Fatalf("NewControlClient: %v", err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	if _, err := ctrl.Register(ctx); err != nil {
		t.Fatalf("Register: %v", err)
	}
	ch, err := ctrl.MapPoll(ctx)
	if err != nil {
		t.Fatalf("MapPoll: %v", err)
	}

	var derpMap = ctrl.DERPMap()
	for derpMap == nil {
		mr, ok := <-ch
		if !ok {
			t.Fatal("map channel closed without DERPMap")
		}
		if mr.DERPMap != nil {
			derpMap = mr.DERPMap
		}
	}

	// Pick a region
	var regionID int
	for id := range derpMap.Regions {
		regionID = id
		break
	}
	t.Logf("using DERP region %d (%s)", regionID, derpMap.Regions[regionID].RegionName)

	// Two independent DERP clients with fresh keys
	keyA := key.NewNode()
	keyB := key.NewNode()

	clientA := NewDERPClient(keyA, t.Logf, nil)
	defer clientA.Close()
	clientB := NewDERPClient(keyB, t.Logf, nil)
	defer clientB.Close()

	if err := clientA.ConnectRegion(ctx, derpMap, regionID); err != nil {
		t.Fatalf("clientA connect: %v", err)
	}
	if err := clientB.ConnectRegion(ctx, derpMap, regionID); err != nil {
		t.Fatalf("clientB connect: %v", err)
	}

	// Give DERP server a moment to register both clients
	time.Sleep(500 * time.Millisecond)

	// A sends to B
	payload := []byte("hello from A to B via DERP")
	if err := clientA.Send(regionID, keyB.Public(), payload); err != nil {
		t.Fatalf("Send A→B: %v", err)
	}
	t.Logf("A sent %d bytes to B", len(payload))

	// B receives
	select {
	case result := <-clientB.RecvChan():
		if result.Source != keyA.Public() {
			t.Fatalf("expected source %s, got %s", keyA.Public().ShortString(), result.Source.ShortString())
		}
		if string(result.Data) != string(payload) {
			t.Fatalf("expected %q, got %q", payload, result.Data)
		}
		t.Logf("B received %d bytes from A: %q", len(result.Data), result.Data)
	case <-time.After(10 * time.Second):
		t.Fatal("timeout waiting for packet from A")
	}

	// B sends back to A
	reply := []byte("reply from B to A")
	if err := clientB.Send(regionID, keyA.Public(), reply); err != nil {
		t.Fatalf("Send B→A: %v", err)
	}

	select {
	case result := <-clientA.RecvChan():
		if result.Source != keyB.Public() {
			t.Fatalf("expected source %s, got %s", keyB.Public().ShortString(), result.Source.ShortString())
		}
		if string(result.Data) != string(reply) {
			t.Fatalf("expected %q, got %q", reply, result.Data)
		}
		t.Logf("A received %d bytes from B: %q", len(result.Data), result.Data)
	case <-time.After(10 * time.Second):
		t.Fatal("timeout waiting for reply from B")
	}
}

// TestDERPClientCloseIdempotent verifies that calling Close multiple times
// does not panic and returns nil.
func TestDERPClientCloseIdempotent(t *testing.T) {
	dc := NewDERPClient(key.NewNode(), t.Logf, nil)

	// First close
	if err := dc.Close(); err != nil {
		t.Fatalf("first Close: %v", err)
	}

	// Second close should not panic
	if err := dc.Close(); err != nil {
		t.Fatalf("second Close: %v", err)
	}

	// Third close for good measure
	if err := dc.Close(); err != nil {
		t.Fatalf("third Close: %v", err)
	}
	t.Logf("triple Close without panic: OK")
}

// TestDERPClientClosedSendError verifies that Send returns an error after Close.
func TestDERPClientClosedSendError(t *testing.T) {
	dc := NewDERPClient(key.NewNode(), t.Logf, nil)
	dc.Close()

	dstKey := key.NewNode().Public()
	err := dc.Send(1, dstKey, []byte("hello"))
	if err == nil {
		t.Fatal("expected error from Send after Close, got nil")
	}
	t.Logf("Send after Close correctly returned: %v", err)
}

// TestDERPClientConnectedRegionsEmpty verifies that ConnectedRegions returns
// an empty slice when no regions have been connected.
func TestDERPClientConnectedRegionsEmpty(t *testing.T) {
	dc := NewDERPClient(key.NewNode(), t.Logf, nil)
	defer dc.Close()

	regions := dc.ConnectedRegions()
	if len(regions) != 0 {
		t.Fatalf("expected 0 connected regions, got %v", regions)
	}
	t.Logf("ConnectedRegions with no connects: [] (OK)")
}