From e0bc011d8615ef1cc77a77da1d5fdff4aeca9836 Mon Sep 17 00:00:00 2001
From: Spike Curtis <spike@coder.com>
Date: Mon, 30 Oct 2023 16:20:56 +0400
Subject: [PATCH] fix: ensure we probe STUN regions on incremental netchecks

Signed-off-by: Spike Curtis <spike@coder.com>
---
 net/netcheck/netcheck.go | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/net/netcheck/netcheck.go b/net/netcheck/netcheck.go
index 4d1b30c58300b..f7ef5100e7c90 100644
--- a/net/netcheck/netcheck.go
+++ b/net/netcheck/netcheck.go
@@ -416,8 +416,11 @@ func makeProbePlan(dm *tailcfg.DERPMap, ifState *interfaces.State, last *Report)
 	had4 := len(last.RegionV4Latency) > 0
 	had6 := len(last.RegionV6Latency) > 0
 	hadBoth := have6if && had4 && had6
+	// 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) {
-		if ri == numIncrementalRegions {
+		if numSTUN == numIncrementalRegions {
 			break
 		}
 		var p4, p6 []probe
@@ -459,6 +462,10 @@ func makeProbePlan(dm *tailcfg.DERPMap, ifState *interfaces.State, last *Report)
 				do6 = false
 			}
 			n := reg.Nodes[try%len(reg.Nodes)]
+			// Coder: The probe won't be valid if the node doesn't have a STUNPort.
+			if n.STUNPort < 0 {
+				continue
+			}
 			prevLatency := cmpx.Or(
 				last.RegionLatency[reg.RegionID]*120/100,
 				defaultActiveRetransmitTime)
@@ -479,6 +486,9 @@ func makeProbePlan(dm *tailcfg.DERPMap, ifState *interfaces.State, last *Report)
 		if len(p6) > 0 {
 			plan[fmt.Sprintf("region-%d-v6", reg.RegionID)] = p6
 		}
+		if len(p4) > 0 || len(p6) > 0 {
+			numSTUN++
+		}
 	}
 	return plan
 }