Check that a stream may still write before it takes a topic back

Claude-Session: https://claude.ai/code/session_01DeXuhSjq5Gta5qPJgSceFe

server/server.go

@@ -1691,9 +1691,7 @@ func (s *Server) handleSubscribeHTTP(w http.ResponseWriter, r *http.Request, v *
 			for _, t := range topics {
 				t.Keepalive()
 			}
-			for _, t := range rateTopics {
-				t.KeepRateVisitor(v)
-			}
+			s.keepRateVisitors(r, v, rateTopics)
 			if err := sub(v, model.NewKeepaliveMessage(topicsStr)); err != nil { // Send keepalive message
 				return err
 			}
@@ -1790,9 +1788,7 @@ func (s *Server) handleSubscribeWS(w http.ResponseWriter, r *http.Request, v *vi
 				for _, t := range topics {
 					t.Keepalive()
 				}
-				for _, t := range rateTopics {
-					t.KeepRateVisitor(v)
-				}
+				s.keepRateVisitors(r, v, rateTopics)
 				if err := ping(); err != nil {
 					return err
 				}
@@ -1898,23 +1894,13 @@ func (s *Server) maybeSetRateVisitors(r *http.Request, v *visitor, topics []*top
 		return eligibleRateTopics, nil
 	}
 
-	// If access controls are enabled, only set rate visitor if
-	// - topic is reserved, and v.user is the owner
-	// - topic is not reserved, and v.user has write access
+	// If access controls are enabled, only set rate visitor if v may write to the topic (see mayBeRateVisitor)
 	writableRateTopics := make([]*topic, 0)
-	for _, t := range topics {
-		if !util.Contains(eligibleRateTopics, t) {
-			continue
-		}
-		ownerUserID, err := s.userManager.ReservationOwner(t.ID)
+	for _, t := range eligibleRateTopics {
+		ok, err := s.mayBeRateVisitor(v, t)
 		if err != nil {
 			return nil, err
-		}
-		if ownerUserID == "" {
-			if err := s.userManager.Authorize(v.User(), t.ID, user.PermissionWrite); err == nil {
-				writableRateTopics = append(writableRateTopics, t)
-			}
-		} else if ownerUserID == v.MaybeUserID() {
+		} else if ok {
 			writableRateTopics = append(writableRateTopics, t)
 		}
 	}
@@ -1922,6 +1908,44 @@ func (s *Server) maybeSetRateVisitors(r *http.Request, v *visitor, topics []*top
 	return writableRateTopics, nil
 }
 
+// mayBeRateVisitor returns true if v may be the topic's rate visitor:
+// - access controls are off
+// - or the topic is reserved, and v.user is the owner
+// - or the topic is not reserved, and v.user has write access
+func (s *Server) mayBeRateVisitor(v *visitor, t *topic) (bool, error) {
+	if s.userManager == nil {
+		return true, nil
+	}
+	ownerUserID, err := s.userManager.ReservationOwner(t.ID)
+	if err != nil {
+		return false, err
+	}
+	if ownerUserID == "" {
+		return s.userManager.Authorize(v.User(), t.ID, user.PermissionWrite) == nil, nil
+	}
+	return ownerUserID == v.MaybeUserID(), nil
+}
+
+// keepRateVisitors keeps a stream's visitor alive as the rate visitor of its topics (see topic.KeepRateVisitor),
+// and takes back a stale topic only if the stream may still be its rate visitor; write access may have been
+// revoked since it subscribed
+func (s *Server) keepRateVisitors(r *http.Request, v *visitor, rateTopics []*topic) {
+	for _, t := range rateTopics {
+		if !t.KeepRateVisitor(v) {
+			continue
+		}
+		ok, err := s.mayBeRateVisitor(v, t)
+		if err != nil {
+			logvr(v, r).Tag(tagSubscribe).With(t).Err(err).Warn("Cannot check if visitor may be rate visitor for topic %s", t.ID)
+			continue
+		} else if !ok {
+			continue
+		}
+		logvr(v, r).Tag(tagSubscribe).With(t).Debug("Taking back topic %s as rate visitor", t.ID)
+		t.TakeRateVisitor(v)
+	}
+}
+
 func (s *Server) setRateVisitors(r *http.Request, v *visitor, rateTopics []*topic) {
 	for _, t := range rateTopics {
 		logvr(v, r).

server/server_test.go

@@ -3782,6 +3782,58 @@ func TestServer_SubscriberRateLimiting_StreamReclaimsTopicFromGonePoller(t *test
 	require.Equal(t, int64(1), s.visitor(netip.MustParseAddr("2001:db8::1"), nil).Stats().Messages)
 }
 
+func TestServer_SubscriberRateLimiting_StreamDoesNotReclaimTopicAfterWriteRevoked(t *testing.T) {
+	// A stream whose write access was revoked while it is open must not take a stale topic back:
+	// publishes would be charged to (and rejected by) the quota of a subscriber who may no longer write
+	c := newTestConfigWithAuthFile(t, "")
+	c.VisitorSubscriberRateLimiting = true
+	c.VisitorMessageDailyLimit = 1
+	c.KeepaliveInterval = 20 * time.Millisecond
+	c.AuthDefault = user.PermissionDenyAll
+	s := newTestServer(t, c)
+	require.Nil(t, s.userManager.AddUser("alice", "alice", user.RoleUser, false))
+	require.Nil(t, s.userManager.AddUser("bob", "bob", user.RoleUser, false))
+	require.Nil(t, s.userManager.AllowAccess("alice", "upAAAAAAAAAAAA", user.PermissionReadWrite))
+	require.Nil(t, s.userManager.AllowAccess("bob", "upAAAAAAAAAAAA", user.PermissionReadWrite))
+	aliceHeaders := map[string]string{
+		"Authorization": util.BasicAuth("alice", "alice"),
+	}
+	bobHeaders := map[string]string{
+		"Authorization": util.BasicAuth("bob", "bob"),
+	}
+
+	// Alice streams and uses up her quota, then may only read
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	go func() {
+		request(t, s, "GET", "/upAAAAAAAAAAAA/json", "", aliceHeaders, func(r *http.Request) {
+			r.RemoteAddr = "1.1.1.1:1234"
+			*r = *r.WithContext(ctx)
+		})
+	}()
+	waitFor(t, func() bool {
+		to := lockedTopic(s, "upAAAAAAAAAAAA")
+		return to != nil && to.SubscribersCount() == 1
+	})
+	to := lockedTopic(s, "upAAAAAAAAAAAA")
+	require.True(t, to.RateVisitor().MessageAllowed())
+	require.Nil(t, s.userManager.AllowAccess("alice", "upAAAAAAAAAAAA", user.PermissionRead))
+
+	// Bob takes the topic with a poll, and is gone a day later
+	response := request(t, s, "GET", "/upAAAAAAAAAAAA/json?poll=1", "", bobHeaders, func(r *http.Request) {
+		r.RemoteAddr = "2.2.2.2:1234"
+	})
+	require.Equal(t, 200, response.Code)
+	ageRateVisitor(to, 25*time.Hour)
+	time.Sleep(5 * c.KeepaliveInterval)
+
+	// Alice's stream did not take the topic back
+	response = request(t, s, "PUT", "/upAAAAAAAAAAAA", "some message", bobHeaders, func(r *http.Request) {
+		r.RemoteAddr = "2.2.2.2:1234"
+	})
+	require.Equal(t, 200, response.Code)
+}
+
 // ageRateVisitor makes the topic's rate visitor look as if it was last seen d ago
 func ageRateVisitor(t *topic, d time.Duration) {
 	t.mu.Lock()

server/topic.go

@@ -132,14 +132,24 @@ func (t *topic) SetRateVisitor(v *visitor) {
 }
 
 // KeepRateVisitor is called on the keepalives of a stream whose visitor may be the topic's rate visitor.
-// A stream with the rate visitor's key keeps it alive, whichever of its subscribers set it. A stream
-// with another key takes the topic back once the rate visitor is stale.
-func (t *topic) KeepRateVisitor(v *visitor) {
+// A stream with the rate visitor's key keeps it alive, whichever of its subscribers set it. It returns
+// true if the rate visitor has another key and is stale, so the stream may take the topic back (see
+// TakeRateVisitor), once it is checked that it still may.
+func (t *topic) KeepRateVisitor(v *visitor) (stale bool) {
 	t.mu.Lock()
 	defer t.mu.Unlock()
 	if t.rateVisitor != nil && t.rateVisitor.key == v.key {
 		t.rateSeen = time.Now()
-	} else if t.rateVisitorStaleNoLock() {
+		return false
+	}
+	return t.rateVisitorStaleNoLock()
+}
+
+// TakeRateVisitor makes v the topic's rate visitor if the current one is still stale
+func (t *topic) TakeRateVisitor(v *visitor) {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	if t.rateVisitorStaleNoLock() {
 		t.rateVisitor = v
 		t.rateSeen = time.Now()
 	}