o_sfu/runtime/websocket_server/
controller.rs

1//! websocket controller for one upgraded socket
2//!
3//! this module bounds upgrade admission before handing the socket to
4//! [`super::session::run`]
5
6use std::sync::Arc;
7
8use axum::{
9    extract::{FromRef, State, ws::WebSocketUpgrade},
10    http::StatusCode,
11    response::{IntoResponse, Response},
12};
13use tokio_util::{sync::CancellationToken, task::TaskTracker};
14use tracing::warn;
15
16use super::{admission::PreAuthWebSocketAdmissionRejection, io::MAX_CLIENT_FRAME_BYTES, session};
17use crate::{
18    config::{AuthConfig, UserConfig},
19    core::prelude::SfuCore,
20    runtime::{
21        RuntimeMetrics, RuntimeState, request_origin::RequestOrigin, room::RoomManager,
22        telemetry::schema::event as telemetry_event,
23    },
24};
25
26pub(crate) struct WebSocketServices {
27    pub(super) auth: AuthConfig,
28    pub(super) user: UserConfig,
29    pub(super) room_manager: Arc<RoomManager>,
30    pub(super) sfu_core: SfuCore,
31    pub(super) metrics: Arc<RuntimeMetrics>,
32    pub(super) shutdown: CancellationToken,
33    sessions: TaskTracker,
34    pre_auth_websocket_admission: super::PreAuthWebSocketAdmission,
35}
36
37impl FromRef<RuntimeState> for WebSocketServices {
38    fn from_ref(state: &RuntimeState) -> Self {
39        Self {
40            auth: state.config.auth.clone(),
41            user: state.config.user,
42            room_manager: Arc::clone(&state.room_manager),
43            sfu_core: state.sfu_core.clone(),
44            metrics: Arc::clone(&state.metrics),
45            shutdown: state.session_shutdown.clone(),
46            sessions: state.session_tasks.clone(),
47            pre_auth_websocket_admission: state.pre_auth_websocket_admission.clone(),
48        }
49    }
50}
51
52pub(crate) async fn upgrade(
53    State(services): State<WebSocketServices>,
54    origin: RequestOrigin,
55    websocket: WebSocketUpgrade,
56) -> Response {
57    let remote_address = Arc::<str>::from(origin.remote_address);
58    let pre_auth_permit = match services
59        .pre_auth_websocket_admission
60        .try_acquire(Arc::clone(&remote_address))
61    {
62        Ok(permit) => permit,
63        Err(rejection) => {
64            reject_pre_auth_admission(&services, remote_address.as_ref(), rejection);
65            return StatusCode::SERVICE_UNAVAILABLE.into_response();
66        }
67    };
68    let session_task = services.sessions.token();
69    websocket
70        .max_message_size(MAX_CLIENT_FRAME_BYTES)
71        .max_frame_size(MAX_CLIENT_FRAME_BYTES)
72        .on_upgrade(move |socket| async move {
73            session::run(socket, services, remote_address, pre_auth_permit).await;
74            drop(session_task);
75        })
76}
77
78fn reject_pre_auth_admission(
79    services: &WebSocketServices,
80    remote_address: &str,
81    rejection: PreAuthWebSocketAdmissionRejection,
82) {
83    match rejection {
84        PreAuthWebSocketAdmissionRejection::Global => {
85            warn!(
86                event = telemetry_event::WS_HANDSHAKE_REJECTED,
87                remote_address,
88                max_pre_auth_websocket_sessions = services.auth.max_pre_auth_websocket_sessions,
89                "rejecting websocket upgrade because global pre-auth admission is full"
90            );
91        }
92        PreAuthWebSocketAdmissionRejection::Origin => {
93            warn!(
94                event = telemetry_event::WS_HANDSHAKE_REJECTED,
95                remote_address,
96                max_pre_auth_websocket_sessions_per_origin =
97                    services.auth.max_pre_auth_websocket_sessions_per_origin,
98                "rejecting websocket upgrade because origin pre-auth admission is full"
99            );
100        }
101    }
102}
o_sfu/runtime/websocket_server/controller.rs

o_sfu/runtime/websocket_server/
controller.rs
