From 6fac2903e1146164b2a587fd36d53a083ac7e98c Mon Sep 17 00:00:00 2001
From: Nick Khyl <nickk@tailscale.com>
Date: Mon, 13 Jan 2025 18:20:09 -0600
Subject: [PATCH] ipn/ipnserver: fix race condition where LocalBackend is reset
 after a different user connects

In this commit, we add a failing test to verify that ipn/ipnserver.Server correctly
sets and unsets the current user when two different clients send requests concurrently
(A sends request, B sends request, A's request completes, B's request completes).

The expectation is that the user who wins the race becomes the current user
from the LocalBackend's perspective, remaining in this state until they disconnect,
after which a different user should be able to connect and use the LocalBackend.

We then fix the second of two bugs in (*Server).addActiveHTTPRequest, where a race
condition causes the LocalBackend's state to be reset after a new client connects,
instead of after the last active request of the previous client completes and the server
becomes idle.

Fixes tailscale/corp#25804

Signed-off-by: Nick Khyl <nickk@tailscale.com>
---
 ipn/ipnserver/server.go      | 15 ++++----
 ipn/ipnserver/server_test.go | 67 ++++++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 8 deletions(-)

diff --git a/ipn/ipnserver/server.go b/ipn/ipnserver/server.go
index c0e99f5d8..a69e43067 100644
--- a/ipn/ipnserver/server.go
+++ b/ipn/ipnserver/server.go
@@ -394,11 +394,14 @@ func (s *Server) addActiveHTTPRequest(req *http.Request, actor ipnauth.Actor) (o
 
 	onDone = func() {
 		s.mu.Lock()
+		defer s.mu.Unlock()
 		delete(s.activeReqs, req)
-		remain := len(s.activeReqs)
-		s.mu.Unlock()
+		if len(s.activeReqs) != 0 {
+			// The server is not idle yet.
+			return
+		}
 
-		if remain == 0 && s.resetOnZero {
+		if s.resetOnZero {
 			if lb.InServerMode() {
 				s.logf("client disconnected; staying alive in server mode")
 			} else {
@@ -408,11 +411,7 @@ func (s *Server) addActiveHTTPRequest(req *http.Request, actor ipnauth.Actor) (o
 		}
 
 		// Wake up callers waiting for the server to be idle:
-		if remain == 0 {
-			s.mu.Lock()
-			s.zeroReqWaiter.wakeAll()
-			s.mu.Unlock()
-		}
+		s.zeroReqWaiter.wakeAll()
 	}
 
 	return onDone, nil
diff --git a/ipn/ipnserver/server_test.go b/ipn/ipnserver/server_test.go
index 7f6131328..97a616db8 100644
--- a/ipn/ipnserver/server_test.go
+++ b/ipn/ipnserver/server_test.go
@@ -12,6 +12,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"runtime"
+	"strconv"
 	"sync"
 	"sync/atomic"
 	"testing"
@@ -187,6 +188,72 @@ func TestSequentialOSUserSwitchingOnWindows(t *testing.T) {
 	connectDisconnectAsUser("UserB")
 }
 
+func TestConcurrentOSUserSwitchingOnWindows(t *testing.T) {
+	enableLogging := false
+	setGOOSForTest(t, "windows")
+
+	ctx := context.Background()
+	server := startDefaultTestIPNServer(t, ctx, enableLogging)
+
+	connectDisconnectAsUser := func(name string) {
+		// User connects and starts watching the IPN bus.
+		client := server.getClientAs(name)
+		watcher, cancelWatcher := client.WatchIPNBus(ctx, ipn.NotifyInitialState)
+		defer cancelWatcher()
+
+		runtime.Gosched()
+
+		// Get the current user from the LocalBackend's perspective
+		// as soon as we're connected.
+		gotUID, gotActor := server.mustBackend().CurrentUserForTest()
+
+		// Wait for the first notification to arrive.
+		// It will either be the initial state we've requested via [ipn.NotifyInitialState],
+		// returned by an actual handler, or a "fake" notification sent by the server
+		// itself to indicate that it is being used by someone else.
+		n, err := watcher.Next()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		// If our user lost the race and the IPN is in use by another user,
+		// we should just return. For the sake of this test, we're not
+		// interested in waiting for the server to become idle.
+		if n.State != nil && *n.State == ipn.InUseOtherUser {
+			return
+		}
+
+		// Otherwise, our user should have been the current user since the time we connected.
+		if gotUID != client.User.UID {
+			t.Errorf("CurrentUser(Initial): got UID %q; want %q", gotUID, client.User.UID)
+			return
+		}
+		if gotActor, ok := gotActor.(*ipnauth.TestActor); !ok || *gotActor != *client.User {
+			t.Errorf("CurrentUser(Initial): got %v; want %v", gotActor, client.User)
+			return
+		}
+
+		// And should still be the current user (as they're still connected)...
+		server.checkCurrentUser(client.User)
+	}
+
+	numIterations := 10
+	for range numIterations {
+		numGoRoutines := 100
+		var wg sync.WaitGroup
+		wg.Add(numGoRoutines)
+		for i := range numGoRoutines {
+			// User logs in, uses Tailscale for a bit, then logs out
+			// in parallel with other users doing the same.
+			go func() {
+				defer wg.Done()
+				connectDisconnectAsUser("User-" + strconv.Itoa(i))
+			}()
+		}
+		wg.Wait()
+	}
+}
+
 func setGOOSForTest(tb testing.TB, goos string) {
 	tb.Helper()
 	envknob.Setenv("TS_DEBUG_FAKE_GOOS", goos)