fix(ci): rustfmt test_support imports for PR #59

src/rtmp_bridge.rs

@@ -147,6 +147,17 @@ impl DbRtmpBridge {
         }
     }
 
+    /// True once `on_connect` has recorded a remote IP for `conn`. Lets
+    /// callers that may need to register a connection mid-poll (before the
+    /// normal `on_connect` pass runs) skip the redundant re-registration and
+    /// its log line on every subsequent publish/play attempt.
+    pub(crate) fn is_registered(&self, conn: ConnId) -> bool {
+        self.conns
+            .lock()
+            .get(&conn)
+            .is_some_and(|cs| !cs.remote_ip.is_empty())
+    }
+
     fn is_auth_rate_limited(&self, remote_ip: &str) -> bool {
         let mut guard = self.auth_failures.lock();
         let now = Instant::now();
@@ -818,6 +829,38 @@ mod tests {
             .unwrap();
     }
 
+    #[test]
+    fn publish_failures_count_toward_auth_rate_limit_when_remote_ip_known() {
+        let db = Arc::new(Db::open(":memory:").unwrap());
+        let bridge = test_bridge(db);
+        let ip = "203.0.113.7:1935";
+
+        for conn in 0..RTMP_AUTH_MAX_FAILURES as u64 {
+            bridge.on_connect(conn, ip);
+            assert!(bridge.authorize_publish(conn, "live", "bogus").is_err());
+        }
+
+        bridge.on_connect(RTMP_AUTH_MAX_FAILURES as u64, ip);
+        assert!(bridge.is_auth_rate_limited(&remote_ip_of(ip)));
+    }
+
+    #[test]
+    fn publish_before_on_connect_skips_auth_failure_tracking() {
+        let db = Arc::new(Db::open(":memory:").unwrap());
+        let bridge = test_bridge(db);
+        let ip = "203.0.113.7:1935";
+
+        for conn in 0..(RTMP_AUTH_MAX_FAILURES as u64 + 2) {
+            assert!(bridge.authorize_publish(conn, "live", "bogus").is_err());
+            bridge.on_connect(conn, ip);
+        }
+
+        assert!(
+            !bridge.is_auth_rate_limited(&remote_ip_of(ip)),
+            "failed publish attempts before on_connect must not consume the per-IP auth budget"
+        );
+    }
+
     #[test]
     fn auth_rate_limit_survives_reconnect_from_same_ip() {
         let db = Arc::new(Db::open(":memory:").unwrap());
@@ -965,6 +1008,16 @@ mod tests {
         assert_eq!(db.publisher_list(Some("s1")).len(), 1);
     }
 
+    #[test]
+    fn is_registered_reflects_on_connect_state() {
+        let db = Arc::new(Db::open(":memory:").unwrap());
+        let bridge = test_bridge(db);
+
+        assert!(!bridge.is_registered(7));
+        bridge.on_connect(7, "127.0.0.1:1000");
+        assert!(bridge.is_registered(7));
+    }
+
     #[test]
     fn authorize_publish_switching_streams_deactivates_prior_publisher() {
         let db = Arc::new(Db::open(":memory:").unwrap());

src/server.rs

@@ -2,6 +2,7 @@
 //! and the RTMP listener(s), then runs until a shutdown signal arrives.
 
 use parking_lot::Mutex;
+use std::cell::Cell;
 use std::collections::{HashMap, HashSet};
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Arc, Mutex as StdMutex};
@@ -17,6 +18,22 @@ use crate::rtmp_bridge::{DbRtmpBridge, FrameInfo, FrameKind, RtmpEventHandler};
 /// registered on the RTMP thread before the poll loop starts.
 pub(crate) static RTMP_BRIDGE: StdMutex<Option<Arc<DbRtmpBridge>>> = StdMutex::new(None);
 
+thread_local! {
+    static RTMP_POLL_SERVER: Cell<Option<*mut librtmp2::server::Server>> = const {
+        Cell::new(None)
+    };
+}
+
+/// Pin the active RTMP server for the duration of `poll()` so publish/play
+/// callbacks can resolve `remote_addr` before auth rate limiting.
+pub(crate) fn set_rtmp_poll_server(server: *mut librtmp2::server::Server) {
+    RTMP_POLL_SERVER.with(|cell| cell.set(Some(server)));
+}
+
+pub(crate) fn clear_rtmp_poll_server() {
+    RTMP_POLL_SERVER.with(|cell| cell.set(None));
+}
+
 /// How often the poll loop wakes up to service the RTMP/RTMPS listener(s).
 pub(crate) const POLL_INTERVAL_MS: u64 = 50;
 
@@ -64,11 +81,48 @@ where
     }
 }
 
+/// Register the client IP on the bridge before publish/play auth runs. During
+/// `server.poll()` the publish/play callbacks can fire before
+/// `process_server_connections` reaches `on_connect`, which would otherwise
+/// skip per-IP auth-failure tracking and rate limiting.
+fn ensure_conn_registered_for_auth(conn_id: u64) {
+    // Skip the pointer walk entirely once the normal `on_connect` pass (or an
+    // earlier call from this same function) has already registered the
+    // remote IP. Without this check every publish/play attempt on an
+    // already-registered connection re-ran `on_connect` and its "new
+    // connection" log line, which was both misleading and needless lock
+    // contention on the hot path.
+    if with_rtmp_bridge(|bridge| bridge.is_registered(conn_id)).unwrap_or(true) {
+        return;
+    }
+    RTMP_POLL_SERVER.with(|cell| {
+        let Some(server_ptr) = cell.get() else {
+            return;
+        };
+        if server_ptr.is_null() {
+            return;
+        }
+        // SAFETY: `RTMP_POLL_SERVER` is set only on the RTMP thread for the
+        // duration of `server.poll()`, which exclusively owns `server`.
+        let server = unsafe { &*server_ptr };
+        let Some(conn) = server
+            .connections
+            .iter()
+            .find(|c| c.conn_id == conn_id && c.client_fd >= 0)
+        else {
+            return;
+        };
+        with_rtmp_bridge(|bridge| bridge.on_connect(conn_id, &conn.remote_addr));
+    });
+}
+
 pub(crate) fn rtmp_publish_cb(conn_id: u64, app: &str, stream_key: &str) -> bool {
+    ensure_conn_registered_for_auth(conn_id);
     with_rtmp_bridge(|b| b.authorize_publish(conn_id, app, stream_key).is_ok()).unwrap_or(false)
 }
 
 pub(crate) fn rtmp_play_cb(conn_id: u64, app: &str, play_key: &str) -> bool {
+    ensure_conn_registered_for_auth(conn_id);
     with_rtmp_bridge(|b| b.authorize_play(conn_id, app, play_key).is_ok()).unwrap_or(false)
 }
 
@@ -543,7 +597,10 @@ impl ServerApp {
                     break;
                 }
 
-                if let Err(e) = server.poll(0) {
+                set_rtmp_poll_server(&mut server);
+                let poll_result = server.poll(0);
+                clear_rtmp_poll_server();
+                if let Err(e) = poll_result {
                     crate::log_warn!("RTMP polling stopped: {e}");
                     break;
                 }

src/test_support.rs

@@ -17,8 +17,8 @@ use crate::http::{self, AppState};
 use crate::logger;
 use crate::rtmp_bridge::{DbRtmpBridge, RtmpEventHandler};
 use crate::server::{
-    POLL_INTERVAL_MS, RTMP_BRIDGE, TrackedConn, process_server_connections, rtmp_media_cb,
-    rtmp_play_cb, rtmp_publish_cb,
+    POLL_INTERVAL_MS, RTMP_BRIDGE, TrackedConn, clear_rtmp_poll_server, process_server_connections,
+    rtmp_media_cb, rtmp_play_cb, rtmp_publish_cb, set_rtmp_poll_server,
 };
 
 static TEST_RUNTIME: OnceLock<Runtime> = OnceLock::new();
@@ -140,7 +140,10 @@ impl TestServer {
                     break;
                 }
 
-                if server.poll(0).is_err() {
+                set_rtmp_poll_server(&mut server);
+                let poll_result = server.poll(0);
+                clear_rtmp_poll_server();
+                if poll_result.is_err() {
                     break;
                 }