From 1355f622beca0db5794201ab8802804ab1299e2f Mon Sep 17 00:00:00 2001
From: Percy Wegmann <percy@tailscale.com>
Date: Thu, 14 Nov 2024 14:21:30 -0600
Subject: [PATCH] cmd/derpprobe,prober: add ability to restrict derpprobe to a
 single region

Updates #24522

Co-authored-by: Mario Minardi <mario@tailscale.com>
Signed-off-by: Percy Wegmann <percy@tailscale.com>
---
 cmd/derpprobe/derpprobe.go |  4 ++++
 prober/derp.go             | 23 +++++++++++++++++++++++
 prober/derp_test.go        | 31 +++++++++++++++++++++++++++++--
 3 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/cmd/derpprobe/derpprobe.go b/cmd/derpprobe/derpprobe.go
index 5b7b77091..8f04326b0 100644
--- a/cmd/derpprobe/derpprobe.go
+++ b/cmd/derpprobe/derpprobe.go
@@ -29,6 +29,7 @@ var (
 	tlsInterval  = flag.Duration("tls-interval", 15*time.Second, "TLS probe interval")
 	bwInterval   = flag.Duration("bw-interval", 0, "bandwidth probe interval (0 = no bandwidth probing)")
 	bwSize       = flag.Int64("bw-probe-size-bytes", 1_000_000, "bandwidth probe size")
+	regionCode   = flag.String("region-code", "", "probe only this region (e.g. 'lax'); if left blank, all regions will be probed")
 )
 
 func main() {
@@ -47,6 +48,9 @@ func main() {
 	if *bwInterval > 0 {
 		opts = append(opts, prober.WithBandwidthProbing(*bwInterval, *bwSize))
 	}
+	if *regionCode != "" {
+		opts = append(opts, prober.WithRegion(*regionCode))
+	}
 	dp, err := prober.DERP(p, *derpMapURL, opts...)
 	if err != nil {
 		log.Fatal(err)
diff --git a/prober/derp.go b/prober/derp.go
index 0dadbe8c2..b1ebc590d 100644
--- a/prober/derp.go
+++ b/prober/derp.go
@@ -45,6 +45,9 @@ type derpProber struct {
 	bwInterval  time.Duration
 	bwProbeSize int64
 
+	// Optionally restrict probes to a single regionCode.
+	regionCode string
+
 	// Probe class for fetching & updating the DERP map.
 	ProbeMap ProbeClass
 
@@ -97,6 +100,14 @@ func WithTLSProbing(interval time.Duration) DERPOpt {
 	}
 }
 
+// WithRegion restricts probing to the specified region identified by its code
+// (e.g. "lax"). This is case sensitive.
+func WithRegion(regionCode string) DERPOpt {
+	return func(d *derpProber) {
+		d.regionCode = regionCode
+	}
+}
+
 // DERP creates a new derpProber.
 //
 // If derpMapURL is "local", the DERPMap is fetched via
@@ -135,6 +146,10 @@ func (d *derpProber) probeMapFn(ctx context.Context) error {
 	defer d.Unlock()
 
 	for _, region := range d.lastDERPMap.Regions {
+		if d.skipRegion(region) {
+			continue
+		}
+
 		for _, server := range region.Nodes {
 			labels := Labels{
 				"region":    region.RegionCode,
@@ -316,6 +331,10 @@ func (d *derpProber) updateMap(ctx context.Context) error {
 	d.lastDERPMapAt = time.Now()
 	d.nodes = make(map[string]*tailcfg.DERPNode)
 	for _, reg := range d.lastDERPMap.Regions {
+		if d.skipRegion(reg) {
+			continue
+		}
+
 		for _, n := range reg.Nodes {
 			if existing, ok := d.nodes[n.Name]; ok {
 				return fmt.Errorf("derpmap has duplicate nodes: %+v and %+v", existing, n)
@@ -338,6 +357,10 @@ func (d *derpProber) ProbeUDP(ipaddr string, port int) ProbeClass {
 	}
 }
 
+func (d *derpProber) skipRegion(region *tailcfg.DERPRegion) bool {
+	return d.regionCode != "" && region.RegionCode != d.regionCode
+}
+
 func derpProbeUDP(ctx context.Context, ipStr string, port int) error {
 	pc, err := net.ListenPacket("udp", ":0")
 	if err != nil {
diff --git a/prober/derp_test.go b/prober/derp_test.go
index a34292a23..c084803e9 100644
--- a/prober/derp_test.go
+++ b/prober/derp_test.go
@@ -44,6 +44,19 @@ func TestDerpProber(t *testing.T) {
 					},
 				},
 			},
+			1: {
+				RegionID:   1,
+				RegionCode: "one",
+				Nodes: []*tailcfg.DERPNode{
+					{
+						Name:     "n3",
+						RegionID: 0,
+						HostName: "derpn3.tailscale.test",
+						IPv4:     "1.1.1.1",
+						IPv6:     "::1",
+					},
+				},
+			},
 		},
 	}
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -68,6 +81,7 @@ func TestDerpProber(t *testing.T) {
 		meshProbeFn:  func(_, _ string) ProbeClass { return FuncProbe(func(context.Context) error { return nil }) },
 		nodes:        make(map[string]*tailcfg.DERPNode),
 		probes:       make(map[string]*Probe),
+		regionCode:   "zero",
 	}
 	if err := dp.probeMapFn(context.Background()); err != nil {
 		t.Errorf("unexpected probeMapFn() error: %s", err)
@@ -84,9 +98,9 @@ func TestDerpProber(t *testing.T) {
 
 	// Add one more node and check that probes got created.
 	dm.Regions[0].Nodes = append(dm.Regions[0].Nodes, &tailcfg.DERPNode{
-		Name:     "n3",
+		Name:     "n4",
 		RegionID: 0,
-		HostName: "derpn3.tailscale.test",
+		HostName: "derpn4.tailscale.test",
 		IPv4:     "1.1.1.1",
 		IPv6:     "::1",
 	})
@@ -113,6 +127,19 @@ func TestDerpProber(t *testing.T) {
 	if len(dp.probes) != 4 {
 		t.Errorf("unexpected probes: %+v", dp.probes)
 	}
+
+	// Stop filtering regions.
+	dp.regionCode = ""
+	if err := dp.probeMapFn(context.Background()); err != nil {
+		t.Errorf("unexpected probeMapFn() error: %s", err)
+	}
+	if len(dp.nodes) != 2 {
+		t.Errorf("unexpected nodes: %+v", dp.nodes)
+	}
+	// 6 regular probes + 2 mesh probe
+	if len(dp.probes) != 8 {
+		t.Errorf("unexpected probes: %+v", dp.probes)
+	}
 }
 
 func TestRunDerpProbeNodePair(t *testing.T) {