diff --git a/Cargo.lock b/Cargo.lock index fa0d1d4..5d93403 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -868,7 +868,7 @@ checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" [[package]] name = "librtmp2" version = "0.1.0" -source = "git+https://github.com/OpenRTMP/librtmp2.git?branch=main#e80119900eaaed16a2c272a552d8c47954319d34" +source = "git+https://github.com/OpenRTMP/librtmp2.git?branch=main#e889f8ca61dc4fd879096da4facae230aac1ef15" dependencies = [ "cc", "libc", @@ -950,9 +950,9 @@ checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" [[package]] name = "memchr" -version = "2.8.2" +version = "2.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" [[package]] name = "mime" diff --git a/src/rtmp_bridge.rs b/src/rtmp_bridge.rs index 1c667d5..54b674a 100644 --- a/src/rtmp_bridge.rs +++ b/src/rtmp_bridge.rs @@ -818,6 +818,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()); diff --git a/src/server.rs b/src/server.rs index 27a1eaf..9ddbc9e 100644 --- a/src/server.rs +++ b/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>> = StdMutex::new(None); +thread_local! { + static RTMP_POLL_SERVER: Cell> = 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,39 @@ 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) { + 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 +588,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; } diff --git a/src/test_support.rs b/src/test_support.rs index 066a3b7..acff2a5 100644 --- a/src/test_support.rs +++ b/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 = 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; }