From 8fd032f158becf11044d41c39cb5616fcc1f67ef Mon Sep 17 00:00:00 2001
From: Johanna Larsson <johanna.a.m.larsson@gmail.com>
Date: Sun, 18 Oct 2026 11:38:16 +0100
Subject: [PATCH] Rate limit subscribeRepos

This endpoint was lacking rate limits, even though it can be a very costly operation. I set it to a 5 per 60s per ip rate limit, which should support normal behaviors and mitigate any misbehaving relays.
---
 crates/tranquil-pds/src/rate_limit/extractor.rs | 5 +++++
 crates/tranquil-pds/src/rate_limit/mod.rs       | 2 ++
 crates/tranquil-pds/src/state.rs                | 7 +++++++
 crates/tranquil-sync/src/subscribe_repos.rs     | 2 ++
 4 files changed, 16 insertions(+)

diff --git a/crates/tranquil-pds/src/rate_limit/extractor.rs b/crates/tranquil-pds/src/rate_limit/extractor.rs
index 68a5db75..71a7ed27 100644
--- a/crates/tranquil-pds/src/rate_limit/extractor.rs
+++ b/crates/tranquil-pds/src/rate_limit/extractor.rs
@@ -115,6 +115,11 @@ impl RateLimitPolicy for HandleVerificationLimit {
     const KIND: RateLimitKind = RateLimitKind::HandleVerification;
 }
 
+pub struct SubscribeReposLimit;
+impl RateLimitPolicy for SubscribeReposLimit {
+    const KIND: RateLimitKind = RateLimitKind::SubscribeRepos;
+}
+
 pub trait RateLimitRejection: IntoResponse + Send + 'static {
     fn new() -> Self;
 }
diff --git a/crates/tranquil-pds/src/rate_limit/mod.rs b/crates/tranquil-pds/src/rate_limit/mod.rs
index f14d1913..8e9c9914 100644
--- a/crates/tranquil-pds/src/rate_limit/mod.rs
+++ b/crates/tranquil-pds/src/rate_limit/mod.rs
@@ -39,6 +39,7 @@ pub struct RateLimiters {
     pub sso_unlink: Arc<KeyedRateLimiter>,
     pub oauth_register_complete: Arc<KeyedRateLimiter>,
     pub handle_verification: Arc<KeyedRateLimiter>,
+    pub subscribe_repos: Arc<KeyedRateLimiter>,
 }
 
 impl Default for RateLimiters {
@@ -70,6 +71,7 @@ impl RateLimiters {
             sso_unlink: keyed_limiter(RateLimitKind::SsoUnlink),
             oauth_register_complete: keyed_limiter(RateLimitKind::OAuthRegisterComplete),
             handle_verification: keyed_limiter(RateLimitKind::HandleVerification),
+            subscribe_repos: keyed_limiter(RateLimitKind::SubscribeRepos),
         }
     }
 
diff --git a/crates/tranquil-pds/src/state.rs b/crates/tranquil-pds/src/state.rs
index c2dd150f..7405a0d8 100644
--- a/crates/tranquil-pds/src/state.rs
+++ b/crates/tranquil-pds/src/state.rs
@@ -99,6 +99,7 @@ pub enum RateLimitKind {
     SsoUnlink,
     OAuthRegisterComplete,
     HandleVerification,
+    SubscribeRepos,
 }
 
 impl RateLimitKind {
@@ -124,6 +125,7 @@ impl RateLimitKind {
             Self::SsoUnlink => "sso_unlink",
             Self::OAuthRegisterComplete => "oauth_register_complete",
             Self::HandleVerification => "handle_verification",
+            Self::SubscribeRepos => "subscribe_repos",
         }
     }
 
@@ -209,6 +211,10 @@ impl RateLimitKind {
                 limit: 10,
                 window_ms: 60_000,
             },
+            Self::SubscribeRepos => RateLimitParams {
+                limit: 5,
+                window_ms: 60_000,
+            },
         }
     }
 }
@@ -489,6 +495,7 @@ impl AppState {
             RateLimitKind::SsoUnlink => &self.rate_limiters.sso_unlink,
             RateLimitKind::OAuthRegisterComplete => &self.rate_limiters.oauth_register_complete,
             RateLimitKind::HandleVerification => &self.rate_limiters.handle_verification,
+            RateLimitKind::SubscribeRepos => &self.rate_limiters.subscribe_repos,
         };
 
         if limiter.check_key(&client_ip.to_string()).is_err() {
diff --git a/crates/tranquil-sync/src/subscribe_repos.rs b/crates/tranquil-sync/src/subscribe_repos.rs
index ac1f90c9..9fbcd018 100644
--- a/crates/tranquil-sync/src/subscribe_repos.rs
+++ b/crates/tranquil-sync/src/subscribe_repos.rs
@@ -8,6 +8,7 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 use tokio::sync::broadcast::error::RecvError;
 use tracing::{error, info, warn};
 use tranquil_db_traits::SequenceNumber;
+use tranquil_pds::rate_limit::{RateLimited, SubscribeReposLimit};
 use tranquil_pds::state::AppState;
 use tranquil_pds::sync::firehose::SequencedEvent;
 use tranquil_pds::sync::frame::{ErrorFrameName, InfoFrameName};
@@ -30,6 +31,7 @@ pub async fn subscribe_repos(
     ws: WebSocketUpgrade,
     State(state): State<AppState>,
     Query(params): Query<SubscribeReposParams>,
+    _rate_limit: RateLimited<SubscribeReposLimit>,
 ) -> Response {
     ws.on_upgrade(move |socket| handle_socket(socket, state, params))
 }