euzu/tuliprox
1
0
Fork 0
mirror of https://github.com/euzu/tuliprox.git synced 2026-10-17 05:12:34 +02:00
Code Issues Packages Projects Releases 5 Wiki Activity

Merge pull request #401 from euzu/feature/connection_handling

Browse Source 
Connection handling fixed for user sessions
This commit is contained in:
euzu
2025-11-05 12:07:14 +01:00
committed by GitHub
parent 041d7ddb16 cb1692b9f6
commit cd2df9ad46
7 changed files with 126 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/src/api/endpoints/websocket_api.rs
+4 -4
View File
@@ -4,7 +4,7 @@ use crate::api::model::EventMessage;
use crate::auth::{verify_token_admin, verify_token_user};
use axum::extract::ws::CloseFrame;
use axum::{extract::ws::{Message, WebSocket, WebSocketUpgrade},response::IntoResponse};
use log::{error, info};
use log::{error, trace};
use shared::model::{ProtocolHandler, ProtocolHandlerMemory, ProtocolMessage, UserRole, WsCloseCode, PROTOCOL_VERSION};
use std::sync::Arc;
use shared::utils::{concat_path_leading_slash};
@@ -14,7 +14,7 @@ async fn websocket_handler(
axum::extract::State(app_state): axum::extract::State<Arc<AppState>>,
ws: WebSocketUpgrade,
) -> impl IntoResponse {
info!("Websocket connected");
trace!("Websocket connected");
ws.on_upgrade(move |socket| handle_socket(socket, app_state, false))
}
@@ -23,7 +23,7 @@ async fn websocket_handler_auth(
axum::extract::State(app_state): axum::extract::State<Arc<AppState>>,
ws: WebSocketUpgrade,
) -> impl IntoResponse {
info!("Websocket connected");
trace!("Websocket connected");
ws.on_upgrade(move |socket| handle_socket(socket, app_state, true))
}
@@ -134,7 +134,7 @@ async fn handle_protocol_message(
}
}
Ok(_) => {
error!("Unexpected protocol message after handshake");
trace!("Unexpected protocol message after handshake");
None
}
Err(e) => {
backend/src/api/model/active_user_manager.rs
+102 -86
View File
@@ -73,16 +73,20 @@ impl UserConnectionData {
}
}
#[derive(Debug, Default)]
struct UserConnections {
by_key: HashMap<String, UserConnectionData>,
key_by_addr: HashMap<SocketAddr, String>,
}
pub struct ActiveUserManager {
grace_period_millis: AtomicU64,
grace_period_timeout_secs: AtomicU64,
log_active_user: AtomicBool,
user: Arc<RwLock<HashMap<String, UserConnectionData>>>,
user_by_addr: Arc<RwLock<HashMap<SocketAddr, String>>>,
gc_ts: Option<AtomicU64>,
connections: RwLock<UserConnections>,
event_manager: Arc<EventManager>,
geo_ip: Arc<ArcSwapOption<GeoIp>>,
}
impl ActiveUserManager {
@@ -96,24 +100,20 @@ impl ActiveUserManager {
grace_period_millis: AtomicU64::new(grace_period_millis),
grace_period_timeout_secs: AtomicU64::new(grace_period_timeout_secs),
log_active_user: AtomicBool::new(log_active_user),
user: Arc::new(RwLock::new(HashMap::new())),
user_by_addr: Arc::new(RwLock::new(HashMap::new())),
connections: RwLock::new(UserConnections::default()),
gc_ts: Some(AtomicU64::new(current_time_secs())),
geo_ip: Arc::clone(geoip),
event_manager: Arc::clone(event_manager),
}
}
#[inline]
async fn get_active_connections(user: &Arc<RwLock<HashMap<String, UserConnectionData>>>) -> usize {
user.read().await.iter().filter(|(_, c)| c.connections > 0).map(|(_, c)| c.connections as usize).sum()
}
async fn log_active_user(&self) {
let user = Arc::clone(&self.user);
let is_log_user_enabled = self.is_log_user_enabled();
let user_connection_count = Self::get_active_connections(&user).await;
let user_count = user.read().await.iter().filter(|(_, c)| c.connections > 0).count();
let (user_count, user_connection_count) = {
let user_connections = self.connections.read().await;
(user_connections.by_key.iter().filter(|(_, c)| c.connections > 0).count(),
user_connections.key_by_addr.len())
};
self.event_manager.send_event(EventMessage::ActiveUser(ActiveUserConnectionChange::Connections(user_count, user_connection_count)));
if is_log_user_enabled {
info!("Active Users: {user_count}, Active User Connections: {user_connection_count}");
@@ -121,28 +121,33 @@ impl ActiveUserManager {
}
pub async fn release_connection(&self, addr: &SocketAddr) {
let username_opt = {
self.user_by_addr.write().await.remove(addr)
};
let log_active_user = {
let mut user_connections = self.connections.write().await;
let username_opt = {
user_connections.key_by_addr.remove(addr)
};
if let Some(username) = username_opt.as_ref() {
{
let mut user = self.user.write().await;
if let Some(connection_data) = user.get_mut(username) {
if connection_data.connections > 0 {
connection_data.connections -= 1;
}
if let Some(username) = username_opt.as_ref() {
{
if let Some(connection_data) = user_connections.by_key.get_mut(username) {
if connection_data.connections > 0 {
connection_data.connections -= 1;
}
if connection_data.connections < connection_data.max_connections {
connection_data.granted_grace = false;
connection_data.grace_ts = 0;
if connection_data.connections < connection_data.max_connections {
connection_data.granted_grace = false;
connection_data.grace_ts = 0;
}
connection_data.streams.retain(|c| c.addr != *addr);
}
connection_data.streams.retain(|c| c.addr != *addr);
}
true
} else {
false
}
}
};
if username_opt.is_some() {
if log_active_user {
self.log_active_user().await;
}
}
@@ -156,7 +161,7 @@ impl ActiveUserManager {
}
pub async fn user_connections(&self, username: &str) -> u32 {
if let Some(connection_data) = self.user.read().await.get(username) {
if let Some(connection_data) = self.connections.read().await.by_key.get(username) {
return connection_data.connections;
}
0
@@ -204,7 +209,7 @@ impl ActiveUserManager {
max_connections: u32,
) -> UserConnectionPermission {
if max_connections > 0 {
if let Some(connection_data) = self.user.write().await.get_mut(username) {
if let Some(connection_data) = self.connections.write().await.by_key.get_mut(username) {
return self.check_connection_permission(username, connection_data);
}
}
@@ -212,11 +217,11 @@ impl ActiveUserManager {
}
pub async fn active_users(&self) -> usize {
self.user.read().await.iter().filter(|(_, c)| c.connections > 0).count()
self.connections.read().await.key_by_addr.len()
}
pub async fn active_connections(&self) -> usize {
Self::get_active_connections(&self.user).await
self.connections.read().await.by_key.values().map(|c| c.connections).sum::<u32>() as usize
}
#[allow(clippy::too_many_arguments)]
@@ -241,21 +246,18 @@ impl ActiveUserManager {
country,
);
{
let mut user_map = self.user.write().await;
if let Some(connection_data) = user_map.get_mut(username) {
let mut user_connections = self.connections.write().await;
if let Some(connection_data) = user_connections.by_key.get_mut(username) {
connection_data.connections += 1;
connection_data.max_connections = max_connections;
connection_data.streams.push(stream_info.clone());
} else {
let mut connection_data = UserConnectionData::new(1, max_connections);
connection_data.streams.push(stream_info.clone());
user_map.insert(username.to_string(), connection_data);
user_connections.by_key.insert(username.to_string(), connection_data);
}
}
{
let mut user_by_addr = self.user_by_addr.write().await;
user_by_addr.insert(fingerprint.addr, username.to_string());
user_connections.key_by_addr.insert(fingerprint.addr, username.to_string());
}
self.log_active_user().await;
@@ -287,8 +289,8 @@ impl ActiveUserManager {
self.gc();
let username = user.username.clone();
let mut user_map = self.user.write().await;
let connection_data = user_map.entry(username.clone()).or_insert_with(|| {
let mut user_connections = self.connections.write().await;
let connection_data = user_connections.by_key.entry(username.clone()).or_insert_with(|| {
debug!("Creating session for user {username} with token {session_token} {}", sanitize_sensitive_info(stream_url));
let mut data = UserConnectionData::new(0, user.max_connections);
let session = Self::new_user_session(session_token, virtual_id, provider, stream_url, addr, connection_permission);
@@ -321,28 +323,19 @@ impl ActiveUserManager {
}
pub async fn update_session_addr(&self, username: &str, token: &str, addr: &SocketAddr) {
let mut old_addr = None;
{
let mut user_map = self.user.write().await;
if let Some(connection_data) = user_map.get_mut(username) {
if let Some(session) = connection_data.sessions.iter_mut().find(|session| session.token == token) {
let previous_addr = session.addr;
session.addr = *addr;
for stream in &mut connection_data.streams {
if stream.addr == previous_addr {
stream.addr = *addr;
}
let mut user_connections = self.connections.write().await;
if let Some(connection_data) = user_connections.by_key.get_mut(username) {
if let Some(session) = connection_data.sessions.iter_mut().find(|s| s.token == token) {
let previous_addr = session.addr;
session.addr = *addr;
for stream in &mut connection_data.streams {
if stream.addr == previous_addr {
stream.addr = *addr;
}
old_addr = Some(previous_addr);
}
}
}
if let Some(previous_addr) = old_addr {
let mut user_by_addr = self.user_by_addr.write().await;
user_by_addr.remove(&previous_addr);
user_by_addr.insert(*addr, username.to_string());
}
}
pub async fn get_and_update_user_session(&self, username: &str, token: &str) -> Option<UserSession> {
@@ -350,35 +343,31 @@ impl ActiveUserManager {
}
async fn update_user_session(&self, username: &str, token: &str) -> Option<UserSession> {
let mut users = self.user.write().await;
if let Some(connection_data) = users.get_mut(username) {
connection_data.ts = current_time_secs();
let mut user_connections = self.connections.write().await;
// Search for index of session
if let Some(index) = connection_data
.sessions
.iter()
.position(|s| s.token == token)
{
// Refresh session last access
connection_data.sessions[index].ts = current_time_secs();
// Only re-evaluate permission for limited users during grace
if connection_data.max_connections > 0
&& connection_data.sessions[index].permission == UserConnectionPermission::GracePeriod
{
let new_permission = self.check_connection_permission(username, connection_data);
connection_data.sessions[index].permission = new_permission;
}
return Some(connection_data.sessions[index].clone());
}
let connection_data = user_connections.by_key.get_mut(username)?;
let now = current_time_secs();
connection_data.ts = now;
let session_index = connection_data.sessions.iter().position(|s| s.token == token)?;
connection_data.sessions[session_index].ts = now;
if connection_data.max_connections > 0
&& connection_data.sessions[session_index].permission == UserConnectionPermission::GracePeriod
{
let new_permission = self.check_connection_permission(username, connection_data);
connection_data.sessions[session_index].permission = new_permission;
}
None
Some(connection_data.sessions[session_index].clone())
}
pub async fn active_streams(&self) -> Vec<StreamInfo> {
let user_map = self.user.read().await;
let user_connections = self.connections.read().await;
let mut streams = Vec::new();
for (_username, connection_data) in user_map.iter() {
for connection_data in user_connections.by_key.values() {
for stream in &connection_data.streams {
streams.push(stream.clone());
}
@@ -392,12 +381,39 @@ impl ActiveUserManager {
let now = current_time_secs();
if now - ts > USER_GC_TTL {
if let Ok(mut users) = self.user.try_write() {
users.retain(|_k, v| now - v.ts < USER_CON_TTL && v.connections > 0);
for connection_data in users.values_mut() {
if let Ok(mut user_connections) = self.connections.try_write() {
let mut addrs_to_remove = Vec::new();
for (username, connection_data) in &mut user_connections.by_key {
let old_sessions: Vec<SocketAddr> = connection_data
.sessions
.iter()
.filter(|s| now - s.ts >= USER_CON_TTL)
.map(|s| s.addr)
.collect();
addrs_to_remove.extend(
old_sessions.into_iter().map(|addr| (addr, username.clone()))
);
connection_data.sessions.retain(|s| now - s.ts < USER_CON_TTL);
}
user_connections.by_key.retain(|_username, v| {
now - v.ts < USER_CON_TTL && v.connections > 0
});
for (addr, username) in addrs_to_remove {
if let Some(current_user) = user_connections.key_by_addr.get(&addr) {
if current_user == &username {
user_connections.key_by_addr.remove(&addr);
}
}
}
let valid_users: std::collections::HashSet<_> = user_connections.by_key.keys().cloned().collect();
user_connections.key_by_addr.retain(|_addr, username| valid_users.contains(username));
gc_ts.store(now, Ordering::Release);
}
}
backend/src/api/model/connection_manager.rs
+3 -13
View File
@@ -2,7 +2,7 @@ use std::borrow::Cow;
use std::net::SocketAddr;
use crate::api::model::{ActiveProviderManager, ActiveUserManager, EventManager, EventMessage, SharedStreamManager};
use std::sync::Arc;
use log::{debug};
use log::debug;
use shared::model::{ActiveUserConnectionChange, StreamChannel};
use crate::auth::Fingerprint;
@@ -36,9 +36,7 @@ impl ConnectionManager {
self.close_socket_signal_tx.subscribe()
}
fn handle_release_connection(&self, addr: &SocketAddr) {
self.event_manager.send_event(EventMessage::ActiveUser(ActiveUserConnectionChange::Disconnected(*addr)));
pub fn kick_connection(&self, addr: &SocketAddr) {
if let Err(e) = self.close_socket_signal_tx.send(*addr) {
debug!("No active receivers for close signal ({addr}): {e:?}");
}
@@ -48,15 +46,7 @@ impl ConnectionManager {
self.user_manager.release_connection(addr).await;
self.provider_manager.release_connection(addr).await;
self.shared_stream_manager.release_connection(addr, true).await;
self.handle_release_connection(addr);
}
pub async fn trigger_release_connection_from_user_manager(&self, addr: &SocketAddr) {
self.provider_manager.release_connection(addr).await;
self.shared_stream_manager.release_connection(addr, true).await;
self.handle_release_connection(addr);
self.event_manager.send_event(EventMessage::ActiveUser(ActiveUserConnectionChange::Disconnected(*addr)));
}
#[allow(clippy::too_many_arguments)]
backend/src/api/serve.rs
+1 -3
View File
@@ -152,10 +152,8 @@ where
conn.as_mut().graceful_shutdown();
}
Ok(msg) = addr_close_rx.recv() => {
// this comes from user manager itself when a user connection is closed
// no need to call `user_manager.remove_connection()`
if msg == addr {
connection_manager_clone.trigger_release_connection_from_user_manager(&addr).await;
connection_manager_clone.release_connection(&addr).await;
debug!("Forced client disconnect {msg}");
conn.as_mut().graceful_shutdown();
break;
frontend/scss/app/components/source_editor/_source_editor.scss
+1
View File
@@ -18,6 +18,7 @@
&-actions {
display: flex;
flex-flow: row wrap;
justify-content: flex-end;
box-sizing: border-box;
gap: var(--gap-small);
border-bottom: 2px solid var(--border-color);
frontend/src/app/components/dashboard/streams_view.rs
+13 -2
View File
@@ -1,7 +1,7 @@
use std::rc::Rc;
use yew::prelude::*;
use yew_i18n::use_translation;
use crate::app::components::StreamsTable;
use crate::app::components::{Card, StatusCard, StreamsTable};
use crate::app::StatusContext;
#[function_component]
@@ -22,7 +22,18 @@ pub fn StreamsView() -> Html {
<h1>{ translate.t("LABEL.STREAMS")}</h1>
</div>
<div class="tp__streams__body">
<StreamsTable streams={ (*memo_streams).clone() } />
<div class="tp__stats__body-group">
<Card><StatusCard title={translate.t("LABEL.ACTIVE_USERS")} data={status_ctx.status.as_ref().map_or_else(|| "n/a".to_string(), |status| status.active_users.to_string())} /></Card>
<Card><StatusCard title={translate.t("LABEL.ACTIVE_USER_CONNECTIONS")} data={status_ctx.status.as_ref().map_or_else(|| "n/a".to_string(), |status| status.active_user_connections.to_string())} /></Card>
<Card>
<StatusCard
title={translate.t("LABEL.ACTIVE_PROVIDER_CONNECTIONS")}
data={status_ctx.status.as_ref().map_or_else(|| "n/a".to_string(), |status|
status.active_provider_connections.as_ref().map(|map| map.values().sum::<usize>()).unwrap_or(0).to_string())}
/>
</Card>
</div>
<StreamsTable streams={ (*memo_streams).clone() } />
</div>
</div>
}
shared/src/utils/default_utils.rs
+2 -2
View File
@@ -8,6 +8,6 @@ pub const fn default_resolve_delay_secs() -> u16 { 2 }
// Default grace values to accommodate rapid channel changes and seek requests,
// helping avoid triggering hard max_connection enforcement.
pub const fn default_grace_period_millis() -> u64 { 400 }
pub const fn default_grace_period_timeout_secs() -> u64 { 2 }
pub const fn default_grace_period_millis() -> u64 { 2000 }
pub const fn default_grace_period_timeout_secs() -> u64 { 4 }
pub const fn default_connect_timeout_secs() -> u32 { 6 }