Expands bandwidth limit to also include excessive polling

server/server.go

@@ -1424,6 +1424,10 @@ func (s *Server) handleSubscribeHTTP(w http.ResponseWriter, r *http.Request, v *
 	}
 	var wlock sync.Mutex
 	var closed bool
+	// Only messages replayed from the cache are charged against the visitor's daily bandwidth
+	// budget, the same one attachment traffic uses. This is set in the poll branch below, which
+	// returns before any Subscribe, so sub() is never called concurrently while it is true.
+	meterPollBandwidth := false
 	defer func() {
 		// This blocks until any in-flight sub() call finishes writing/flushing the response writer,
 		// then marks the connection as closed so future sub() calls are no-ops. This prevents a panic
@@ -1442,6 +1446,11 @@ func (s *Server) handleSubscribeHTTP(w http.ResponseWriter, r *http.Request, v *
 		if err != nil {
 			return err
 		}
+		// Charge before writing, so an exhausted budget fails the first message and surfaces as a
+		// clean 429 with nothing written.
+		if meterPollBandwidth && !v.BandwidthAllowed(int64(len(m))) {
+			return errHTTPTooManyRequestsLimitAttachmentBandwidth
+		}
 		wlock.Lock()
 		defer wlock.Unlock()
 		if closed {
@@ -1464,6 +1473,7 @@ func (s *Server) handleSubscribeHTTP(w http.ResponseWriter, r *http.Request, v *
 		for _, t := range topics {
 			t.Keepalive()
 		}
+		meterPollBandwidth = true
 		return s.sendOldMessages(topics, since, scheduled, v, sub)
 	}
 	ctx, cancel := context.WithCancel(context.Background())

server/server.yml

@@ -417,7 +417,10 @@
 
 # Rate limiting: Attachment size and bandwidth limits per visitor:
 # - visitor-attachment-total-size-limit is the total storage limit used for attachments per visitor
-# - visitor-attachment-daily-bandwidth-limit is the total daily attachment download/upload traffic limit per visitor
+# - visitor-attachment-daily-bandwidth-limit is the total daily traffic limit per visitor. It covers
+#   attachment downloads/uploads AND messages replayed from the message cache by poll requests. A
+#   poll without a "since" cursor returns the topic's entire cache, so a busy topic can be re-read
+#   for many times its own size; charging it here caps what one visitor can pull per day.
 #
 # visitor-attachment-total-size-limit: "100M"
 # visitor-attachment-daily-bandwidth-limit: "500M"

server/server_test.go

@@ -2810,6 +2810,38 @@ func TestServer_PublishAttachmentBandwidthLimit(t *testing.T) {
 	})
 }
 
+func TestServer_PollBandwidthLimit(t *testing.T) {
+	forEachBackend(t, func(t *testing.T, databaseURL string) {
+		// A poll without "since" replays the entire cache, so a topic that is cheap to fill is
+		// expensive to read over and over. Replayed bytes are charged against the same daily
+		// budget as attachment traffic. One message per poll keeps the accounting coarse: any
+		// shortfall hits the very first message, so the request is rejected before anything is
+		// written rather than truncated mid-stream.
+		c := newTestConfig(t, databaseURL)
+		c.VisitorAttachmentDailyBandwidthLimit = 9000 // Enough for two replays of the ~4 KB topic below, not three
+		s := newTestServer(t, c)
+
+		require.Equal(t, 200, request(t, s, "PUT", "/mytopic", util.RandomString(4000), nil).Code)
+
+		// Two full replays fit in the budget
+		for i := 1; i <= 2; i++ {
+			response := request(t, s, "GET", "/mytopic/json?poll=1", "", nil)
+			require.Equal(t, 200, response.Code)
+			require.Equal(t, 1, len(toMessages(t, response.Body.String())))
+		}
+
+		// The third is rejected before a single byte is written
+		response := request(t, s, "GET", "/mytopic/json?poll=1", "", nil)
+		require.Equal(t, 429, response.Code)
+		require.Equal(t, 42905, toHTTPError(t, response.Body.String()).Code)
+
+		// A subscription that replays nothing is not charged against the budget
+		response = request(t, s, "GET", "/mytopic/json?poll=1&since=none", "", nil)
+		require.Equal(t, 200, response.Code)
+		require.Empty(t, strings.TrimSpace(response.Body.String()))
+	})
+}
+
 func TestServer_PublishAttachmentBandwidthLimitUploadOnly(t *testing.T) {
 	forEachBackend(t, func(t *testing.T, databaseURL string) {
 		content := util.RandomString(5000) // > 4096

server/visitor.go

@@ -65,7 +65,7 @@ type visitor struct {
 	callsLimiter         *util.FixedLimiter // Rate limiter for calls
 	subscriptionLimiter  *util.FixedLimiter // Fixed limiter for active subscriptions (ongoing connections)
 	topicCreationLimiter *rate.Limiter      // Rate limiter for inserting new topics into the in-memory topic map
-	bandwidthLimiter     *util.RateLimiter  // Limiter for attachment bandwidth downloads
+	bandwidthLimiter     *util.RateLimiter  // Limiter for attachment downloads and cached-message replay (polls)
 	accountLimiter       *rate.Limiter      // Rate limiter for account actions (signup, password-reset requests), may be nil
 	authLimiter          *rate.Limiter      // Limiter for incorrect login attempts, may be nil
 	firebase             time.Time          // Next allowed Firebase message