From 639b9d68f67816c40fe38b58323c0e9271d06f16 Mon Sep 17 00:00:00 2001
From: Colin Adler <colin1adler@gmail.com>
Date: Mon, 3 Apr 2023 15:24:10 -0500
Subject: [PATCH] fix(healthcheck): ensure STUNOnly nodes aren't marked as
 unhealthy

---
 coderd/healthcheck/derp.go      |  9 ++++++-
 coderd/healthcheck/derp_test.go | 43 +++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/coderd/healthcheck/derp.go b/coderd/healthcheck/derp.go
index 6c1426f600ba1..e4cb206f3ead0 100644
--- a/coderd/healthcheck/derp.go
+++ b/coderd/healthcheck/derp.go
@@ -171,7 +171,14 @@ func (r *DERPNodeReport) Run(ctx context.Context) error {
 	r.doExchangeMessage(ctx)
 	r.doSTUNTest(ctx)
 
-	if !r.CanExchangeMessages || r.UsesWebsocket || r.STUN.Error != nil {
+	// We can't exchange messages with the node,
+	if (!r.CanExchangeMessages && !r.Node.STUNOnly) ||
+		// A node may use websockets because `Upgrade: DERP` may be blocked on
+		// the load balancer. This is unhealthy because websockets are slower
+		// than the regular DERP protocol.
+		r.UsesWebsocket ||
+		// The node was marked as STUN compatible but the STUN test failed.
+		r.STUN.Error != nil {
 		r.Healthy = false
 	}
 	return nil
diff --git a/coderd/healthcheck/derp_test.go b/coderd/healthcheck/derp_test.go
index b976df8a4d659..fdc313e72bd28 100644
--- a/coderd/healthcheck/derp_test.go
+++ b/coderd/healthcheck/derp_test.go
@@ -186,6 +186,49 @@ func TestDERP(t *testing.T) {
 			}
 		}
 	})
+
+	t.Run("OK/STUNOnly", func(t *testing.T) {
+		t.Parallel()
+
+		var (
+			ctx    = context.Background()
+			report = healthcheck.DERPReport{}
+			opts   = &healthcheck.DERPReportOptions{
+				DERPMap: &tailcfg.DERPMap{Regions: map[int]*tailcfg.DERPRegion{
+					1: {
+						EmbeddedRelay: true,
+						RegionID:      999,
+						Nodes: []*tailcfg.DERPNode{{
+							Name:             "999stun0",
+							RegionID:         999,
+							HostName:         "stun.l.google.com",
+							STUNPort:         19302,
+							STUNOnly:         true,
+							InsecureForTests: true,
+							ForceHTTP:        true,
+						}},
+					},
+				}},
+			}
+		)
+
+		err := report.Run(ctx, opts)
+		require.NoError(t, err)
+
+		assert.True(t, report.Healthy)
+		for _, region := range report.Regions {
+			assert.True(t, region.Healthy)
+			for _, node := range region.NodeReports {
+				assert.True(t, node.Healthy)
+				assert.False(t, node.CanExchangeMessages)
+				assert.Len(t, node.ClientLogs, 0)
+
+				assert.True(t, node.STUN.Enabled)
+				assert.True(t, node.STUN.CanSTUN)
+				assert.NoError(t, node.STUN.Error)
+			}
+		}
+	})
 }
 
 func tsDERPMap(ctx context.Context, t testing.TB) *tailcfg.DERPMap {