prober: support filtering regions by region ID in addition to code

Updates tailscale/corp#25758

Signed-off-by: Percy Wegmann <percy@tailscale.com>

cmd/derpprobe/derpprobe.go

@@ -32,7 +32,7 @@ var (
 	bwTUNIPv4Address   = flag.String("bw-tun-ipv4-addr", "", "if specified, bandwidth probes will be performed over a TUN device at this address in order to exercise TCP-in-TCP in similar fashion to TCP over Tailscale via DERP; we will use a /30 subnet including this IP address")
 	qdPacketsPerSecond = flag.Int("qd-packets-per-second", 0, "if greater than 0, queuing delay will be measured continuously using 260 byte packets (approximate size of a CallMeMaybe packet) sent at this rate per second")
 	qdPacketTimeout    = flag.Duration("qd-packet-timeout", 5*time.Second, "queuing delay packets arriving after this period of time from being sent are treated like dropped packets and don't count toward queuing delay timings")
-	regionCode         = flag.String("region-code", "", "probe only this region (e.g. 'lax'); if left blank, all regions will be probed")
+	regionCodeOrID     = flag.String("region-code", "", "probe only this region (e.g. 'lax' or '17'); if left blank, all regions will be probed")
 )
 
 func main() {
@@ -52,8 +52,8 @@ func main() {
 	if *bwInterval > 0 {
 		opts = append(opts, prober.WithBandwidthProbing(*bwInterval, *bwSize, *bwTUNIPv4Address))
 	}
-	if *regionCode != "" {
-		opts = append(opts, prober.WithRegion(*regionCode))
+	if *regionCodeOrID != "" {
+		opts = append(opts, prober.WithRegionCodeOrID(*regionCodeOrID))
 	}
 	dp, err := prober.DERP(p, *derpMapURL, opts...)
 	if err != nil {

prober/derp.go

@@ -60,8 +60,8 @@ type derpProber struct {
 	qdPacketsPerSecond int // in packets per second
 	qdPacketTimeout    time.Duration
 
-	// Optionally restrict probes to a single regionCode.
-	regionCode string
+	// Optionally restrict probes to a single regionCodeOrID.
+	regionCodeOrID string
 
 	// Probe class for fetching & updating the DERP map.
 	ProbeMap ProbeClass
@@ -135,11 +135,11 @@ 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 {
+// WithRegionCodeOrID restricts probing to the specified region identified by its code
+// (e.g. "lax") or its id (e.g. "17"). This is case sensitive.
+func WithRegionCodeOrID(regionCode string) DERPOpt {
 	return func(d *derpProber) {
-		d.regionCode = regionCode
+		d.regionCodeOrID = regionCode
 	}
 }
 
@@ -598,7 +598,7 @@ func (d *derpProber) ProbeUDP(ipaddr string, port int) ProbeClass {
 }
 
 func (d *derpProber) skipRegion(region *tailcfg.DERPRegion) bool {
-	return d.regionCode != "" && region.RegionCode != d.regionCode
+	return d.regionCodeOrID != "" && region.RegionCode != d.regionCodeOrID && strconv.Itoa(region.RegionID) != d.regionCodeOrID
 }
 
 func derpProbeUDP(ctx context.Context, ipStr string, port int) error {

prober/derp_test.go

@@ -71,17 +71,17 @@ func TestDerpProber(t *testing.T) {
 	clk := newFakeTime()
 	p := newForTest(clk.Now, clk.NewTicker)
 	dp := &derpProber{
-		p:            p,
-		derpMapURL:   srv.URL,
-		tlsInterval:  time.Second,
-		tlsProbeFn:   func(_ string) ProbeClass { return FuncProbe(func(context.Context) error { return nil }) },
-		udpInterval:  time.Second,
-		udpProbeFn:   func(_ string, _ int) ProbeClass { return FuncProbe(func(context.Context) error { return nil }) },
-		meshInterval: time.Second,
-		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",
+		p:              p,
+		derpMapURL:     srv.URL,
+		tlsInterval:    time.Second,
+		tlsProbeFn:     func(_ string) ProbeClass { return FuncProbe(func(context.Context) error { return nil }) },
+		udpInterval:    time.Second,
+		udpProbeFn:     func(_ string, _ int) ProbeClass { return FuncProbe(func(context.Context) error { return nil }) },
+		meshInterval:   time.Second,
+		meshProbeFn:    func(_, _ string) ProbeClass { return FuncProbe(func(context.Context) error { return nil }) },
+		nodes:          make(map[string]*tailcfg.DERPNode),
+		probes:         make(map[string]*Probe),
+		regionCodeOrID: "zero",
 	}
 	if err := dp.probeMapFn(context.Background()); err != nil {
 		t.Errorf("unexpected probeMapFn() error: %s", err)
@@ -129,7 +129,7 @@ func TestDerpProber(t *testing.T) {
 	}
 
 	// Stop filtering regions.
-	dp.regionCode = ""
+	dp.regionCodeOrID = ""
 	if err := dp.probeMapFn(context.Background()); err != nil {
 		t.Errorf("unexpected probeMapFn() error: %s", err)
 	}