prober: add a DERP bandwidth probe

Updates tailscale/corp#17912

Signed-off-by: Anton Tolchanov <anton@tailscale.com>

prober/derp_test.go

@@ -5,12 +5,19 @@ package prober
 
 import (
 	"context"
+	"crypto/sha256"
+	"crypto/tls"
 	"encoding/json"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"time"
 
+	"tailscale.com/derp"
+	"tailscale.com/derp/derphttp"
 	"tailscale.com/tailcfg"
+	"tailscale.com/types/key"
 )
 
 func TestDerpProber(t *testing.T) {
@@ -50,18 +57,21 @@ func TestDerpProber(t *testing.T) {
 	clk := newFakeTime()
 	p := newForTest(clk.Now, clk.NewTicker)
 	dp := &derpProber{
-		p:           p,
-		derpMapURL:  srv.URL,
-		tlsProbeFn:  func(_ string) ProbeFunc { return func(context.Context) error { return nil } },
-		udpProbeFn:  func(_ string, _ int) ProbeFunc { return func(context.Context) error { return nil } },
-		meshProbeFn: func(_, _ string) ProbeFunc { return func(context.Context) error { return nil } },
-		nodes:       make(map[string]*tailcfg.DERPNode),
-		probes:      make(map[string]*Probe),
+		p:            p,
+		derpMapURL:   srv.URL,
+		tlsInterval:  time.Second,
+		tlsProbeFn:   func(_ string) ProbeFunc { return func(context.Context) error { return nil } },
+		udpInterval:  time.Second,
+		udpProbeFn:   func(_ string, _ int) ProbeFunc { return func(context.Context) error { return nil } },
+		meshInterval: time.Second,
+		meshProbeFn:  func(_, _ string) ProbeFunc { return func(context.Context) error { return nil } },
+		nodes:        make(map[string]*tailcfg.DERPNode),
+		probes:       make(map[string]*Probe),
 	}
 	if err := dp.ProbeMap(context.Background()); err != nil {
 		t.Errorf("unexpected ProbeMap() error: %s", err)
 	}
-	if len(dp.nodes) != 2 || dp.nodes["derpn1.tailscale.test"] == nil || dp.nodes["derpn2.tailscale.test"] == nil {
+	if len(dp.nodes) != 2 || dp.nodes["n1"] == nil || dp.nodes["n2"] == nil {
 		t.Errorf("unexpected nodes: %+v", dp.nodes)
 	}
 	// Probes expected for two nodes:
@@ -103,3 +113,99 @@ func TestDerpProber(t *testing.T) {
 		t.Errorf("unexpected probes: %+v", dp.probes)
 	}
 }
+
+func TestRunDerpProbeNodePair(t *testing.T) {
+	// os.Setenv("DERP_DEBUG_LOGS", "true")
+	serverPrivateKey := key.NewNode()
+	s := derp.NewServer(serverPrivateKey, t.Logf)
+	defer s.Close()
+
+	httpsrv := &http.Server{
+		TLSNextProto: make(map[string]func(*http.Server, *tls.Conn, http.Handler)),
+		Handler:      derphttp.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)
+		}
+	}()
+	newClient := func() *derphttp.Client {
+		c, err := derphttp.NewClient(key.NewNode(), serverURL, t.Logf)
+		if err != nil {
+			t.Fatalf("NewClient: %v", err)
+		}
+		m, err := c.Recv()
+		if err != nil {
+			t.Fatalf("Recv: %v", err)
+		}
+		switch m.(type) {
+		case derp.ServerInfoMessage:
+		default:
+			t.Fatalf("unexpected first message type %T", m)
+		}
+		return c
+	}
+
+	c1 := newClient()
+	defer c1.Close()
+	c2 := newClient()
+	defer c2.Close()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+	defer cancel()
+
+	err = runDerpProbeNodePair(ctx, &tailcfg.DERPNode{Name: "c1"}, &tailcfg.DERPNode{Name: "c2"}, c1, c2, 100_000_000)
+	if err != nil {
+		t.Error(err)
+	}
+}
+
+func Test_packetsForSize(t *testing.T) {
+	tests := []struct {
+		name        string
+		size        int
+		wantPackets int
+		wantUnique  bool
+	}{
+		{"small_unqiue", 8, 1, true},
+		{"8k_unique", 8192, 1, true},
+		{"full_size_packet", derp.MaxPacketSize, 1, true},
+		{"larger_than_one", derp.MaxPacketSize + 1, 2, false},
+		{"large", 500000, 8, false},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			hashes := make(map[string]int)
+			for i := 0; i < 5; i++ {
+				pkts := packetsForSize(int64(tt.size))
+				if len(pkts) != tt.wantPackets {
+					t.Errorf("packetsForSize(%d) got %d packets, want %d", tt.size, len(pkts), tt.wantPackets)
+				}
+				var total int
+				hash := sha256.New()
+				for _, p := range pkts {
+					hash.Write(p)
+					total += len(p)
+				}
+				hashes[string(hash.Sum(nil))]++
+				if total != tt.size {
+					t.Errorf("packetsForSize(%d) returned %d bytes total", tt.size, total)
+				}
+			}
+			unique := len(hashes) > 1
+			if unique != tt.wantUnique {
+				t.Errorf("packetsForSize(%d) is unique=%v (returned %d different answers); want unique=%v", tt.size, unique, len(hashes), unique)
+			}
+		})
+	}
+}