diff --git a/src/api/api_utils.rs b/src/api/api_utils.rs index 4e903d147..ee9d8abca 100644 --- a/src/api/api_utils.rs +++ b/src/api/api_utils.rs @@ -601,8 +601,8 @@ pub async fn stream_response(app_state: &AppState, if let Some(provider) = provider_name { if matches!(item_type, PlaylistItemType::LiveHls | PlaylistItemType::LiveDash | PlaylistItemType::Video | PlaylistItemType::Series) { - if let Some(token) = app_state.active_users.create_user_session(&user.username, virtual_id, &provider, &stream_url, connection_permission).await { - response = response.header(axum::http::header::SET_COOKIE, &create_session_cookie(token)); + if let Some(token) = app_state.active_users.create_user_session(user, virtual_id, &provider, stream_url, connection_permission).await { + response = response.header(axum::http::header::SET_COOKIE, create_session_cookie(token)); } } } @@ -769,7 +769,7 @@ pub fn redirect(url: &str) -> impl IntoResponse { .unwrap() } -pub async fn is_seek_request<'a>( +pub async fn is_seek_request( cluster: XtreamCluster, req_headers: &HeaderMap, ) -> bool { diff --git a/src/api/endpoints/hls_api.rs b/src/api/endpoints/hls_api.rs index bbaf670a3..42740a85d 100644 --- a/src/api/endpoints/hls_api.rs +++ b/src/api/endpoints/hls_api.rs @@ -60,7 +60,7 @@ pub(in crate::api) async fn handle_hls_stream_request(app_state: &Arc, match app_state.active_provider.get_next_provider(&input.name).await { Some(provider_cfg) => { let stream_url = get_stream_alternative_url(&url, input, &provider_cfg); - let session_token= app_state.active_users.create_user_session(&user.username, virtual_id, &provider_cfg.name, &stream_url, connection_permission).await; + let session_token= app_state.active_users.create_user_session(user, virtual_id, &provider_cfg.name, &stream_url, connection_permission).await; (stream_url, session_token) }, None => (url, None), @@ -77,10 +77,10 @@ pub(in crate::api) async fn handle_hls_stream_request(app_state: &Arc, hls_url: response_url, virtual_id, input_id: input.id, - user_token: session_token.unwrap_or_default(), + user_token: session_token, }; let hls_content = rewrite_hls(user, &rewrite_hls_props); - hls_response(hls_content, session_token.map(|token| create_session_cookie(token))).into_response() + hls_response(hls_content, session_token.map(create_session_cookie)).into_response() } Err(err) => { error!("Failed to download m3u8 {}", sanitize_sensitive_info(err.to_string().as_str())); @@ -119,16 +119,22 @@ async fn hls_api_stream( return create_custom_video_stream_response(&app_state.config, &CustomVideoStreamType::ProviderConnectionsExhausted).into_response(); } - let hls_url = if let Some((session_token_opt, hls_url)) = get_hls_session_token_and_url_from_token(&app_state.config.t_encrypt_secret, ¶ms.token) { - if let Some(session_token) = session_token_opt { - if session.token != session_token { - return axum::http::StatusCode::BAD_REQUEST.into_response(); - } - } - hls_url - } else { - return axum::http::StatusCode::BAD_REQUEST.into_response(); + // let hls_url = if let Some((session_token_opt, hls_url)) = get_hls_session_token_and_url_from_token(&app_state.config.t_encrypt_secret, ¶ms.token) { + // if let Some(session_token) = session_token_opt { + // if session.token != session_token { + // return axum::http::StatusCode::BAD_REQUEST.into_response(); + // } + // } + // hls_url + // } else { + // return axum::http::StatusCode::BAD_REQUEST.into_response(); + // }; + + let hls_url = match get_hls_session_token_and_url_from_token(&app_state.config.t_encrypt_secret, ¶ms.token) { + Some((Some(session_token), hls_url)) if session_token == session.token => hls_url, + _ => return axum::http::StatusCode::BAD_REQUEST.into_response(), }; + session.stream_url = hls_url; if session.virtual_id == virtual_id { if is_seek_request(XtreamCluster::Live, &req_headers).await { @@ -145,7 +151,7 @@ async fn hls_api_stream( } if is_hls_url(&session.stream_url) { - return handle_hls_stream_request(&app_state, &user, Some(&session), &session.stream_url, virtual_id, input, connection_permission).await.into_response(); + return handle_hls_stream_request(&app_state, &user, Some(session), &session.stream_url, virtual_id, input, connection_permission).await.into_response(); } force_provider_stream_response(&app_state, session, PlaylistItemType::LiveHls, &req_headers, input, &user).await.into_response() diff --git a/src/api/endpoints/m3u_api.rs b/src/api/endpoints/m3u_api.rs index 1a0937b7d..8dcc2131e 100644 --- a/src/api/endpoints/m3u_api.rs +++ b/src/api/endpoints/m3u_api.rs @@ -102,7 +102,7 @@ async fn m3u_api_stream( // partial request means we are in reverse proxy mode, seek happened return force_provider_stream_response(&app_state, session, pli.item_type, &req_headers, input, &user).await.into_response() } - }; + } let connection_permission = user.connection_permission(&app_state).await; if connection_permission == UserConnectionPermission::Exhausted { diff --git a/src/api/endpoints/xtream_api.rs b/src/api/endpoints/xtream_api.rs index a36f3bb95..1142302bf 100644 --- a/src/api/endpoints/xtream_api.rs +++ b/src/api/endpoints/xtream_api.rs @@ -216,7 +216,7 @@ async fn xtream_player_api_stream( // partial request means we are in reverse proxy mode, seek happened return force_provider_stream_response(app_state, session, pli.item_type, req_headers, input, &user).await.into_response() } - }; + } let connection_permission = user.connection_permission(app_state).await; if connection_permission == UserConnectionPermission::Exhausted { diff --git a/src/api/model/active_user_manager.rs b/src/api/model/active_user_manager.rs index 3a3beaaa9..3e9f7f02c 100644 --- a/src/api/model/active_user_manager.rs +++ b/src/api/model/active_user_manager.rs @@ -1,4 +1,4 @@ -use crate::model::api_proxy::UserConnectionPermission; +use crate::model::api_proxy::{ProxyUserCredentials, UserConnectionPermission}; use crate::model::config::Config; use crate::utils::default_utils::{default_grace_period_millis, default_grace_period_timeout_secs}; use crate::utils::time_utils::current_time_secs; @@ -7,6 +7,7 @@ use log::{debug, info}; use rand::RngCore; use std::collections::HashMap; use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; use tokio::sync::RwLock; pub struct UserConnectionGuard { @@ -42,10 +43,10 @@ struct UserConnectionData { } impl UserConnectionData { - fn new(max_connections: u32) -> Self { + fn new(connections: u32, max_connections: u32) -> Self { Self { max_connections, - connections: 1, + connections, granted_grace: false, grace_ts: 0, sessions: Vec::new(), @@ -58,6 +59,7 @@ pub struct ActiveUserManager { grace_period_timeout_secs: u64, log_active_user: bool, user: Arc>>, + gc_ts: Option, } impl ActiveUserManager { @@ -72,6 +74,7 @@ impl ActiveUserManager { grace_period_timeout_secs, log_active_user, user: Arc::new(RwLock::new(HashMap::new())), + gc_ts: Some(AtomicU64::new(current_time_secs())), } } @@ -81,6 +84,7 @@ impl ActiveUserManager { grace_period_timeout_secs: self.grace_period_timeout_secs, log_active_user: self.log_active_user, user: Arc::clone(&self.user), + gc_ts: None, } } @@ -91,40 +95,40 @@ impl ActiveUserManager { 0 } - async fn check_connection_permission(&self, username: &str, connection_data: &mut UserConnectionData) -> UserConnectionPermission { - let current_connections = connection_data.connections; + fn check_connection_permission(&self, username: &str, connection_data: &mut UserConnectionData) -> UserConnectionPermission { + let current_connections = connection_data.connections; - if current_connections < connection_data.max_connections { - // Reset grace period because user is back under max_connections - connection_data.granted_grace = false; - connection_data.grace_ts = 0; - return UserConnectionPermission::Allowed; + if current_connections < connection_data.max_connections { + // Reset grace period because user is back under max_connections + connection_data.granted_grace = false; + connection_data.grace_ts = 0; + return UserConnectionPermission::Allowed; + } + + let now = get_current_timestamp(); + // Check if user already used grace period + if connection_data.granted_grace { + if now - connection_data.grace_ts <= self.grace_period_timeout_secs { + // Grace timeout still active, deny connection + debug!("User access denied, grace exhausted, too many connections: {username}"); + return UserConnectionPermission::Exhausted; } + // Grace timeout expired, reset grace counters + connection_data.granted_grace = false; + connection_data.grace_ts = 0; + } - let now = get_current_timestamp(); - // Check if user already used grace period - if connection_data.granted_grace { - if now - connection_data.grace_ts <= self.grace_period_timeout_secs { - // Grace timeout still active, deny connection - debug!("User access denied, grace exhausted, too many connections: {username}"); - return UserConnectionPermission::Exhausted; - } - // Grace timeout expired, reset grace counters - connection_data.granted_grace = false; - connection_data.grace_ts = 0; - } + if self.grace_period_millis > 0 && current_connections == connection_data.max_connections { + // Allow grace period once + connection_data.granted_grace = true; + connection_data.grace_ts = now; + debug!("Granted grace period for user access: {username}"); + return UserConnectionPermission::GracePeriod; + } - if self.grace_period_millis > 0 && current_connections == connection_data.max_connections { - // Allow grace period once - connection_data.granted_grace = true; - connection_data.grace_ts = now; - debug!("Granted grace period for user access: {username}"); - return UserConnectionPermission::GracePeriod; - } - - // Too many connections, no grace allowed - debug!("User access denied, too many connections: {username}"); - UserConnectionPermission::Exhausted + // Too many connections, no grace allowed + debug!("User access denied, too many connections: {username}"); + UserConnectionPermission::Exhausted } pub async fn connection_permission( @@ -134,7 +138,7 @@ impl ActiveUserManager { ) -> UserConnectionPermission { if max_connections > 0 { if let Some(connection_data) = self.user.write().await.get_mut(username) { - return self.check_connection_permission(username, connection_data).await; + return self.check_connection_permission(username, connection_data); } } UserConnectionPermission::Allowed @@ -155,7 +159,7 @@ impl ActiveUserManager { connection_data.connections += 1; connection_data.max_connections = max_connections; } else { - lock.insert(username.to_string(), UserConnectionData::new(max_connections)); + lock.insert(username.to_string(), UserConnectionData::new(1, max_connections)); } drop(lock); @@ -189,39 +193,42 @@ impl ActiveUserManager { } fn find_user_session(token: u32, sessions: &[UserSession]) -> Option<&UserSession> { - for session in sessions { - if session.token == token { - return Some(session); - } - } - None + sessions.iter().find(|&session| session.token == token) } - pub async fn create_user_session(&self, username: &str, virtual_id: u32, provider: &str, stream_url: &str, connection_permission: UserConnectionPermission) -> Option { - let mut lock = self.user.write().await; - if let Some(connection_data) = lock.get_mut(username) { - let session_token = rand::rng().next_u32(); - let session = UserSession { - token: session_token, - virtual_id, - provider: provider.to_string(), - stream_url: stream_url.to_string(), - ts: current_time_secs(), - permission: connection_permission, - }; - connection_data.sessions.push(session); - return Some(session_token); + fn new_user_session(virtual_id: u32, provider: &str, stream_url: &str, connection_permission: UserConnectionPermission) -> UserSession { + let session_token = rand::rng().next_u32(); + UserSession { + token: session_token, + virtual_id, + provider: provider.to_string(), + stream_url: stream_url.to_string(), + ts: current_time_secs(), + permission: connection_permission, + } + } + + pub async fn create_user_session(&self, user: &ProxyUserCredentials, virtual_id: u32, provider: &str, stream_url: &str, connection_permission: UserConnectionPermission) -> Option { + let mut lock = self.user.write().await; + if let Some(connection_data) = lock.get_mut(&user.username) { + let session = Self::new_user_session(virtual_id, provider, stream_url, connection_permission); + self.gc().await; + let token = session.token; + connection_data.sessions.push(session); + Some(token) + } else { + let mut connection_data = UserConnectionData::new(0, user.max_connections); + let session = Self::new_user_session(virtual_id, provider, stream_url, connection_permission); + self.gc().await; + let token = session.token; + connection_data.sessions.push(session); + lock.insert(user.username.to_string(), connection_data); + Some(token) } - drop(lock); - None } pub async fn get_user_session(&self, username: &str, token: u32) -> Option { self.update_user_session(username, token).await - // let mut lock = self.user.write().await; - // lock.get_mut(username) - // .and_then(|conn| Self::find_user_session(token, &conn.sessions)) - // .cloned() // owned copy } async fn update_user_session(&self, username: &str, token: u32) -> Option { @@ -243,7 +250,7 @@ impl ActiveUserManager { if let Some(index) = found_session_index { let session_permission = connection_data.sessions[index].permission.clone(); if session_permission == UserConnectionPermission::GracePeriod { - let new_permission = self.check_connection_permission(username, connection_data).await; + let new_permission = self.check_connection_permission(username, connection_data); connection_data.sessions[index].permission = new_permission; } return Some(connection_data.sessions[index].clone()); @@ -252,7 +259,6 @@ impl ActiveUserManager { None } - async fn log_active_user(&self) { if self.log_active_user { let user_count = self.active_users().await; @@ -260,6 +266,18 @@ impl ActiveUserManager { info!("Active Users: {user_count}, Active User Connections: {user_connection_count}"); } } + + async fn gc(&self) { + if let Some(gc_ts) = &self.gc_ts { + let ts = gc_ts.load(Ordering::SeqCst); + if current_time_secs() - ts > 10_800 { // 3 hours + let mut lock = self.user.write().await; + for (_, connection_data) in lock.iter_mut() { + connection_data.sessions.retain(|s| current_time_secs() - s.ts < 10_800); + } + } + } + } } // diff --git a/src/api/model/provider_config.rs b/src/api/model/provider_config.rs index a40997e8f..fce730e62 100644 --- a/src/api/model/provider_config.rs +++ b/src/api/model/provider_config.rs @@ -152,6 +152,7 @@ impl ProviderConfig { } // is intended to use with redirects, to cycle through provider + // do not increment and connection counter! async fn get_next(&self, grace: bool, grace_period_timeout_secs: u64) -> bool { if self.max_connections == 0 { return true; @@ -162,7 +163,6 @@ impl ProviderConfig { if connections < self.max_connections { guard.granted_grace = false; guard.grace_ts = 0; - guard.current_connections += 1; } let now = get_current_timestamp(); diff --git a/src/processing/parser/hls.rs b/src/processing/parser/hls.rs index 9a2c479e8..eb4cf8a9a 100644 --- a/src/processing/parser/hls.rs +++ b/src/processing/parser/hls.rs @@ -28,7 +28,7 @@ pub struct RewriteHlsProps<'a> { pub hls_url: String, pub virtual_id: u32, pub input_id: u16, - pub user_token: u32, + pub user_token: Option, } fn rewrite_hls_url(input: &str, replacement: &str) -> String { @@ -51,8 +51,10 @@ fn rewrite_uri_attrib(line: &str, props: &RewriteHlsProps) -> String { if let Some(caps) = CONSTANTS.re_hls_uri.captures(line) { let uri = &caps[1]; let target_url = &rewrite_hls_url(&props.hls_url, uri); - if let Some(token) = create_hls_session_token_and_url(props.secret, props.user_token, target_url) { - return CONSTANTS.re_hls_uri.replace(line, format!(r#"URI="{token}""#)).to_string(); + if let Some(user_token) = &props.user_token { + if let Some(token) = create_hls_session_token_and_url(props.secret, *user_token, target_url) { + return CONSTANTS.re_hls_uri.replace(line, format!(r#"URI="{token}""#)).to_string(); + } } } line.to_string() @@ -76,17 +78,19 @@ pub fn rewrite_hls(user: &ProxyUserCredentials, props: &RewriteHlsProps) -> Stri } else { rewrite_hls_url(&props.hls_url, line) }; - if let Some(token) = create_hls_session_token_and_url(props.secret, props.user_token, &target_url) { - let url = format!( - "{}/{HLS_PREFIX}/{}/{}/{}/{}/{}", - props.base_url, - username, - password, - props.input_id, - props.virtual_id, - token - ); - result.push(url); + if let Some(user_token) = &props.user_token { + if let Some(token) = create_hls_session_token_and_url(props.secret, *user_token, &target_url) { + let url = format!( + "{}/{HLS_PREFIX}/{}/{}/{}/{}/{}", + props.base_url, + username, + password, + props.input_id, + props.virtual_id, + token + ); + result.push(url); + } } } result.push("\r\n".to_string()); diff --git a/src/utils/hash_utils.rs b/src/utils/hash_utils.rs index 4b660022e..d30ea1dfb 100644 --- a/src/utils/hash_utils.rs +++ b/src/utils/hash_utils.rs @@ -28,11 +28,11 @@ pub fn generate_playlist_uuid(key: &str, provider_id: &str, item_type: PlaylistI pub fn u32_to_base64(value: u32) -> String { // big-endian is safer and more portable when you care about consistent ordering or cross-platform data let bytes = value.to_be_bytes(); - general_purpose::STANDARD.encode(&bytes) + general_purpose::URL_SAFE_NO_PAD.encode(bytes) } pub fn base64_to_u32(encoded: &str) -> Option { - let decoded = general_purpose::STANDARD.decode(encoded).ok()?; + let decoded = general_purpose::URL_SAFE_NO_PAD.decode(encoded).ok()?; if decoded.len() != 4 { return None;