ipn/{ipnserver,localapi}: fix InUseOtherUser handling with WatchIPNBus

Updates tailscale/corp#8222

Change-Id: I2d6fa6514c7b8d0f89fded35a2d44e7df27e6fb1
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>

ipn/ipnserver/server.go

@@ -52,7 +52,8 @@ type Server struct {
 	mu            sync.Mutex
 	lastUserID    ipn.WindowsUserID // tracks last userid; on change, Reset state for paranoia
 	activeReqs    map[*http.Request]*ipnauth.ConnIdentity
-	backendWaiter set.HandleSet[context.CancelFunc] // values are wake-up funcs of lb waiters
+	backendWaiter waiterSet // of LocalBackend waiters
+	zeroReqWaiter waiterSet // of blockUntilZeroConnections waiters
 }
 
 func (s *Server) mustBackend() *ipnlocal.LocalBackend {
@@ -63,22 +64,47 @@ func (s *Server) mustBackend() *ipnlocal.LocalBackend {
 	return lb
 }
 
+// waiterSet is a set of callers waiting on something. Each item (map value) in
+// the set is a func that wakes up that waiter's context. The waiter is responsible
+// for removing itself from the set when woken up. The (*waiterSet).add method
+// returns a cleanup method which does that removal. The caller than defers that
+// cleanup.
+//
+// TODO(bradfitz): this is a generally useful pattern. Move elsewhere?
+type waiterSet set.HandleSet[context.CancelFunc]
+
+// add registers a new waiter in the set.
+// It aquires mu to add the waiter, and does so again when cleanup is called to remove it.
+// ready is closed when the waiter is ready (or ctx is done).
+func (s *waiterSet) add(mu *sync.Mutex, ctx context.Context) (ready <-chan struct{}, cleanup func()) {
+	ctx, cancel := context.WithCancel(ctx)
+	hs := (*set.HandleSet[context.CancelFunc])(s) // change method set
+	mu.Lock()
+	h := hs.Add(cancel)
+	mu.Unlock()
+	return ctx.Done(), func() {
+		mu.Lock()
+		delete(*hs, h)
+		mu.Unlock()
+		cancel()
+	}
+}
+
+// wakeAll wakes up all waiters in the set.
+func (w waiterSet) wakeAll() {
+	for _, cancel := range w {
+		cancel() // they'll remove themselves
+	}
+}
+
 func (s *Server) awaitBackend(ctx context.Context) (_ *ipnlocal.LocalBackend, ok bool) {
 	lb := s.lb.Load()
 	if lb != nil {
 		return lb, true
 	}
-	ctx, cancel := context.WithCancel(ctx)
-	defer cancel()
 
-	s.mu.Lock()
-	h := s.backendWaiter.Add(cancel)
-	s.mu.Unlock()
-	defer func() {
-		s.mu.Lock()
-		delete(s.backendWaiter, h)
-		s.mu.Unlock()
-	}()
+	ready, cleanup := s.backendWaiter.add(&s.mu, ctx)
+	defer cleanup()
 
 	// Try again, now that we've registered, in case there was a
 	// race.
@@ -87,7 +113,7 @@ func (s *Server) awaitBackend(ctx context.Context) (_ *ipnlocal.LocalBackend, ok
 		return lb, true
 	}
 
-	<-ctx.Done()
+	<-ready
 	lb = s.lb.Load()
 	return lb, lb != nil
 }
@@ -160,6 +186,11 @@ func (s *Server) serveHTTP(w http.ResponseWriter, r *http.Request) {
 
 	onDone, err := s.addActiveHTTPRequest(r, ci)
 	if err != nil {
+		if ou, ok := err.(inUseOtherUserError); ok && localapi.InUseOtherUserIPNStream(w, r, ou.Unwrap()) {
+			w.(http.Flusher).Flush()
+			s.blockWhileIdentityInUse(ctx, ci)
+			return
+		}
 		http.Error(w, err.Error(), http.StatusUnauthorized)
 		return
 	}
@@ -219,6 +250,30 @@ func (s *Server) checkConnIdentityLocked(ci *ipnauth.ConnIdentity) error {
 	return nil
 }
 
+// blockWhileIdentityInUse blocks while ci can't connect to the server because
+// the server is in use by a different user.
+//
+// This is primarily used for the Windows GUI, to block until one user's done
+// controlling the tailscaled process.
+func (s *Server) blockWhileIdentityInUse(ctx context.Context, ci *ipnauth.ConnIdentity) error {
+	inUse := func() bool {
+		s.mu.Lock()
+		defer s.mu.Unlock()
+		_, ok := s.checkConnIdentityLocked(ci).(inUseOtherUserError)
+		return ok
+	}
+	for inUse() {
+		// Check whenever the connection count drops down to zero.
+		ready, cleanup := s.zeroReqWaiter.add(&s.mu, ctx)
+		<-ready
+		cleanup()
+		if err := ctx.Err(); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 // localAPIPermissions returns the permissions for the given identity accessing
 // the Tailscale local daemon API.
 //
@@ -340,6 +395,13 @@ func (s *Server) addActiveHTTPRequest(req *http.Request, ci *ipnauth.ConnIdentit
 				lb.ResetForClientDisconnect()
 			}
 		}
+
+		// Wake up callers waiting for the server to be idle:
+		if remain == 0 {
+			s.mu.Lock()
+			s.zeroReqWaiter.wakeAll()
+			s.mu.Unlock()
+		}
 	}
 
 	return onDone, nil
@@ -373,9 +435,7 @@ func (s *Server) SetLocalBackend(lb *ipnlocal.LocalBackend) {
 	s.startBackendIfNeeded()
 
 	s.mu.Lock()
-	for _, wake := range s.backendWaiter {
-		wake() // they'll remove themselves when woken
-	}
+	s.backendWaiter.wakeAll()
 	s.mu.Unlock()
 
 	// TODO(bradfitz): send status update to GUI long poller waiter. See