hscontrol: prefer completed auth over expired ctx in followup wait

waitForFollowup selected on ctx.Done() and the verdict channel with equal
priority; when both were ready, select picked at random and discarded a
successful registration as a spurious 401 timeout. Check for a completed
verdict first, race the deadline only if none is ready.

Fixes #3385

hscontrol/auth.go

@@ -315,31 +315,41 @@ func (h *Headscale) waitForFollowup(
 	}
 
 	if reg, ok := h.state.GetAuthCacheEntry(followupReg); ok {
+		var verdict types.AuthVerdict
 		select {
-		case <-ctx.Done():
-			return nil, NewHTTPError(http.StatusUnauthorized, "registration timed out", err)
-		case verdict := <-reg.WaitForAuth():
-			if verdict.Accept() {
-				if !verdict.Node.Valid() {
-					// registration is expired in the cache, instruct the client to try a new registration
-					return h.reqToNewRegisterResponse(req, machineKey)
-				}
-
-				// The followup poll is only authenticated by the auth ID in the
-				// URL, so fail closed unless the Noise session asking for the
-				// result was started with the same machine key that opened the
-				// registration. [State.HandleNodeFromAuthPath] resolves the node
-				// from the cached [types.RegistrationData.MachineKey], so the two
-				// match on the normal path. [Headscale.handleRegister] and
-				// [Headscale.handleLogout] apply the same check.
-				err := machineKeyMismatch(verdict.Node, machineKey)
-				if err != nil {
-					return nil, err
-				}
-
-				return nodeToRegisterResponse(verdict.Node), nil
+		// Prefer a completed registration even if the context has also
+		// expired. When both are ready, a plain select picks at random and
+		// would discard a successful registration as a spurious timeout
+		// (issue #3385).
+		case verdict = <-reg.WaitForAuth():
+		default:
+			select {
+			case <-ctx.Done():
+				return nil, NewHTTPError(http.StatusUnauthorized, "registration timed out", ctx.Err())
+			case verdict = <-reg.WaitForAuth():
 			}
 		}
+
+		if verdict.Accept() {
+			if !verdict.Node.Valid() {
+				// registration is expired in the cache, instruct the client to try a new registration
+				return h.reqToNewRegisterResponse(req, machineKey)
+			}
+
+			// The followup poll is only authenticated by the auth ID in the
+			// URL, so fail closed unless the Noise session asking for the
+			// result was started with the same machine key that opened the
+			// registration. [State.HandleNodeFromAuthPath] resolves the node
+			// from the cached [types.RegistrationData.MachineKey], so the two
+			// match on the normal path. [Headscale.handleRegister] and
+			// [Headscale.handleLogout] apply the same check.
+			err := machineKeyMismatch(verdict.Node, machineKey)
+			if err != nil {
+				return nil, err
+			}
+
+			return nodeToRegisterResponse(verdict.Node), nil
+		}
 	}
 
 	// if the follow-up registration isn't found anymore, instruct the client to try a new registration

hscontrol/auth_test.go

@@ -4183,3 +4183,74 @@ func TestWaitForFollowupMachineKeyMismatch(t *testing.T) {
 		assert.NotEmpty(t, resp.User.DisplayName, "the owner's identity is returned on the legitimate path")
 	})
 }
+
+// TestFollowupWaitPrefersCompletedAuthOverExpiredContext reproduces
+// https://github.com/juanfont/headscale/issues/3385.
+//
+// Root cause: [Headscale.waitForFollowup] selects on ctx.Done() and the auth
+// verdict channel with equal priority. When the registration has ALREADY
+// completed (verdict buffered) but the request context has ALSO expired, Go's
+// select picks a ready case at random, so roughly half the time it returns
+// "registration timed out" and discards a successful registration.
+//
+// The v0.28.0 hscontrol test suite hit this because the followup context
+// timeout was only 100ms while the setup goroutine (create user + node in
+// SQLite) frequently took longer on slower/constrained builders (ppc64le,
+// Alpine CI). Both channels ended up ready at once and the flake surfaced as
+// TestAuthenticationFlows/followup_registration_success failing with
+// "http error[401]: registration timed out".
+//
+// The fix must give the completed-auth case priority over context
+// cancellation. This test forces both cases ready on every iteration; it must
+// never report a timeout.
+func TestFollowupWaitPrefersCompletedAuthOverExpiredContext(t *testing.T) {
+	app := createTestApp(t)
+
+	machineKey := key.NewMachine().Public()
+	nodeKey := key.NewNode().Public()
+
+	const iterations = 300
+
+	timeouts, authorized := 0, 0
+
+	for i := range iterations {
+		regID, err := types.NewAuthID()
+		require.NoError(t, err)
+
+		authReq := types.NewRegisterAuthRequest(&types.RegistrationData{
+			Hostname: "followup-race-node",
+		})
+		app.state.SetAuthCacheEntry(regID, authReq)
+
+		// Registration completes BEFORE we wait: verdict is buffered.
+		user := app.state.CreateUserForTest(fmt.Sprintf("followup-race-user-%d", i))
+		node := app.state.CreateNodeForTest(user, "followup-race-node")
+		// waitForFollowup fails closed unless the node carries the polling
+		// session's machine key; production sets this via the cached
+		// RegistrationData. CreateNodeForTest picks a random one.
+		node.MachineKey = machineKey
+		authReq.FinishAuth(types.AuthVerdict{Node: node.View()})
+
+		// Context is expired BEFORE we wait: both select cases are ready.
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+
+		req := tailcfg.RegisterRequest{
+			Followup: fmt.Sprintf("http://localhost:8080/register/%s", regID),
+			NodeKey:  nodeKey,
+		}
+
+		resp, err := app.waitForFollowup(ctx, req, machineKey)
+		switch {
+		case err != nil:
+			timeouts++
+		case resp != nil && resp.MachineAuthorized:
+			authorized++
+		}
+	}
+
+	assert.Zero(t, timeouts,
+		"waitForFollowup must never report a timeout when auth has already completed; got %d/%d timeouts",
+		timeouts, iterations)
+	assert.Equal(t, iterations, authorized, "every completed registration must be returned as authorized")
+}