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) {