From ebb6b9ade2016c8a04bcb5228d75b83707cf8dc0 Mon Sep 17 00:00:00 2001
From: Spike Curtis <spike@coder.com>
Date: Mon, 6 Nov 2023 07:31:59 +0400
Subject: [PATCH] fix: plan v4/v6 STUN probing based on relative index

Signed-off-by: Spike Curtis <spike@coder.com>
---
 net/netcheck/netcheck.go      |   9 ++-
 net/netcheck/netcheck_test.go | 114 ++++++++++++++++++++++++++++++++++
 2 files changed, 120 insertions(+), 3 deletions(-)

diff --git a/net/netcheck/netcheck.go b/net/netcheck/netcheck.go
index f7ef5100e7c90..9ce6118482686 100644
--- a/net/netcheck/netcheck.go
+++ b/net/netcheck/netcheck.go
@@ -419,7 +419,7 @@ func makeProbePlan(dm *tailcfg.DERPMap, ifState *interfaces.State, last *Report)
 	// Coder: Some regions don't have STUN, so we need to make sure we have probed
 	// enough STUN regions
 	numSTUN := 0
-	for ri, reg := range sortRegions(dm, last) {
+	for _, reg := range sortRegions(dm, last) {
 		if numSTUN == numIncrementalRegions {
 			break
 		}
@@ -430,14 +430,14 @@ func makeProbePlan(dm *tailcfg.DERPMap, ifState *interfaces.State, last *Report)
 		// By default, each node only gets one STUN packet sent,
 		// except the fastest two from the previous round.
 		tries := 1
-		isFastestTwo := ri < 2
+		isFastestTwo := numSTUN < 2
 
 		if isFastestTwo {
 			tries = 2
 		} else if hadBoth {
 			// For dual stack machines, make the 3rd & slower nodes alternate
 			// between.
-			if ri%2 == 0 {
+			if numSTUN%2 == 0 {
 				do4, do6 = true, false
 			} else {
 				do4, do6 = false, true
@@ -979,6 +979,9 @@ func (c *Client) GetReport(ctx context.Context, dm *tailcfg.DERPMap) (_ *Report,
 	}
 
 	plan := makeProbePlan(dm, ifState, last)
+	if len(plan) == 0 {
+		c.logf("empty probe plan; do we have STUN regions?")
+	}
 
 	// If we're doing a full probe, also check for a captive portal. We
 	// delay by a bit to wait for UDP STUN to finish, to avoid the probe if
diff --git a/net/netcheck/netcheck_test.go b/net/netcheck/netcheck_test.go
index 53fded1897ca8..2bd995af1a7a4 100644
--- a/net/netcheck/netcheck_test.go
+++ b/net/netcheck/netcheck_test.go
@@ -1202,3 +1202,117 @@ func TestNeverPickSTUNOnlyRegionAsPreferredDERP(t *testing.T) {
 		t.Errorf("got PreferredDERP=%d; want 1", r.PreferredDERP)
 	}
 }
+
+func Test_makeProbePlan_incremental(t *testing.T) {
+	type testcase struct {
+		name      string
+		haveV4    bool
+		haveV6    bool
+		stunPorts []int
+		expected  []string
+	}
+	testcases := []testcase{
+		{
+			name:      "slowSTUN",
+			haveV4:    true,
+			stunPorts: []int{-1, -1, -1, 4, 5, 6},
+			expected:  []string{"region-4-v4", "region-5-v4", "region-6-v4"},
+		},
+		{
+			name:      "fastSTUN",
+			haveV4:    true,
+			stunPorts: []int{1, 2, 3, -1, -1, -1},
+			expected:  []string{"region-1-v4", "region-2-v4", "region-2-v4"},
+		},
+		{
+			name:      "dualStackFast",
+			haveV4:    true,
+			haveV6:    true,
+			stunPorts: []int{1, 2, 3, -1, -1, -1},
+			expected: []string{
+				"region-1-v4", "region-2-v4", "region-2-v4",
+				"region-1-v6", "region-2-v6", // only fastest 2 have v6
+			},
+		},
+		{
+			name:      "dualStackSlow",
+			haveV4:    true,
+			haveV6:    true,
+			stunPorts: []int{-1, -1, -1, 4, 5, 6},
+			expected: []string{
+				"region-4-v4", "region-5-v4", "region-6-v4",
+				"region-4-v6", "region-5-v6", // only fastest 2 have v6
+			},
+		},
+		{
+			// pathological case: STUN regions are not in top 3 and even
+			// indexes re: latency
+			name:      "dualStackInterspersed",
+			haveV4:    true,
+			haveV6:    true,
+			stunPorts: []int{-1, -1, -1, 4, -1, 6, -1, 8},
+			expected: []string{
+				"region-4-v4", "region-6-v4", "region-8-v4",
+				"region-4-v6", "region-6-v6", // only fastest 2 have v6
+			},
+		},
+	}
+	for _, tc := range testcases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			state := &interfaces.State{
+				HaveV4: tc.haveV4,
+				HaveV6: tc.haveV6,
+			}
+			plan := makeProbePlan(
+				derpMapBySTUNPort(tc.stunPorts),
+				state,
+				reportLatencyAscending(len(tc.stunPorts), tc.haveV4, tc.haveV6),
+			)
+			for _, e := range tc.expected {
+				if _, ok := plan[e]; !ok {
+					t.Errorf("expected %s in the plan", e)
+				}
+			}
+		})
+	}
+
+}
+
+func derpMapBySTUNPort(ports []int) *tailcfg.DERPMap {
+	derpMap := &tailcfg.DERPMap{Regions: make(map[int]*tailcfg.DERPRegion)}
+	for i, p := range ports {
+		j := i + 1
+		var name string
+		if p < 0 {
+			name = fmt.Sprintf("noSTUN%d", j)
+		} else {
+			name = fmt.Sprintf("STUN%d", j)
+		}
+		derpMap.Regions[j] = &tailcfg.DERPRegion{
+			RegionID:   j,
+			RegionCode: name,
+			Nodes: []*tailcfg.DERPNode{
+				{Name: name, STUNPort: p},
+			}}
+	}
+	return derpMap
+}
+
+func reportLatencyAscending(n int, haveV4, haveV6 bool) *Report {
+	r := &Report{
+		RegionLatency:   make(map[int]time.Duration),
+		RegionV4Latency: make(map[int]time.Duration),
+		RegionV6Latency: make(map[int]time.Duration),
+	}
+	for i := 1; i <= n; i++ {
+		r.RegionLatency[i] = time.Duration(i) * time.Millisecond
+		if haveV4 {
+			r.RegionV4Latency[i] = time.Duration(i) * time.Millisecond
+		}
+		if haveV6 {
+			r.RegionV6Latency[i] = time.Duration(i) * time.Millisecond
+		}
+	}
+	return r
+}