o_sfu/runtime/websocket_server/
io.rs

1use std::{future::Future, time::Duration};
2
3use axum::extract::ws::{CloseFrame, Message, WebSocket};
4use futures_util::{
5    SinkExt,
6    stream::{SplitSink, SplitStream},
7};
8use o_sfu_protocol::wire::{
9    ClientEnvelope, Envelope, EnvelopeBatch, EnvelopeBatchDecodeError, EnvelopeDecodeError,
10    MAX_ENVELOPE_BATCH_LEN, ServerEnvelope, WebSocketCloseCode, decode_envelope_batch,
11};
12use tokio::time::timeout;
13
14use crate::application::user_session::UserOutput;
15
16pub(crate) type WsWriter = SplitSink<WebSocket, Message>;
17pub(crate) type WsReader = SplitStream<WebSocket>;
18
19pub const MAX_CLIENT_FRAME_BYTES: usize = 256 * 1024;
20
21pub const MAX_CLIENT_BATCH_ENVELOPES: usize = MAX_ENVELOPE_BATCH_LEN;
22
23const OUTBOUND_WRITE_TIMEOUT: Duration = Duration::from_secs(5);
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub enum ClientBatchDecodeFailureKind {
27    InvalidInput,
28    UnsupportedFeature,
29}
30
31#[derive(Debug, Clone, PartialEq, Eq)]
32pub enum ClientBatchDecodeError {
33    FrameTooLarge { actual: usize, limit: usize },
34    BatchTooLarge { actual: usize, limit: usize },
35    InvalidJson,
36    InvalidRoutingMetadata,
37    InvalidEnvelope(EnvelopeDecodeError),
38}
39
40impl ClientBatchDecodeError {
41    #[must_use]
42    pub const fn kind(&self) -> ClientBatchDecodeFailureKind {
43        match self {
44            Self::InvalidEnvelope(EnvelopeDecodeError::UnknownTag(_)) => {
45                ClientBatchDecodeFailureKind::UnsupportedFeature
46            }
47            Self::FrameTooLarge { .. }
48            | Self::BatchTooLarge { .. }
49            | Self::InvalidJson
50            | Self::InvalidRoutingMetadata
51            | Self::InvalidEnvelope(
52                EnvelopeDecodeError::InvalidPayload(_) | EnvelopeDecodeError::UnexpectedPayload(_),
53            ) => ClientBatchDecodeFailureKind::InvalidInput,
54        }
55    }
56}
57
58/// # Errors
59///
60/// Returns an error when the frame exceeds the byte limit, the batch exceeds
61/// the envelope limit, the payload is not valid envelope JSON, the route
62/// metadata is mixed, or any decoded envelope violates the protocol signaling
63/// contract.
64pub fn decode_client_batch(payload: &str) -> Result<Vec<ClientEnvelope>, ClientBatchDecodeError> {
65    if payload.len() > MAX_CLIENT_FRAME_BYTES {
66        return Err(ClientBatchDecodeError::FrameTooLarge {
67            actual: payload.len(),
68            limit: MAX_CLIENT_FRAME_BYTES,
69        });
70    }
71    let batch =
72        decode_envelope_batch(payload, MAX_CLIENT_BATCH_ENVELOPES).map_err(
73            |error| match error {
74                EnvelopeBatchDecodeError::InvalidJson => ClientBatchDecodeError::InvalidJson,
75                EnvelopeBatchDecodeError::BatchTooLarge { actual, limit } => {
76                    ClientBatchDecodeError::BatchTooLarge { actual, limit }
77                }
78                EnvelopeBatchDecodeError::InvalidRoutingMetadata => {
79                    ClientBatchDecodeError::InvalidRoutingMetadata
80                }
81            },
82        )?;
83    batch
84        .into_iter()
85        .map(|envelope| {
86            ClientEnvelope::decode(envelope).map_err(ClientBatchDecodeError::InvalidEnvelope)
87        })
88        .collect()
89}
90
91pub(super) async fn send_user_output_bounded(
92    writer: &mut WsWriter,
93    output: UserOutput,
94) -> Result<usize, WebSocketCloseCode> {
95    with_outbound_write_timeout(send_user_signals(writer, output)).await
96}
97
98pub(super) async fn send_message_bounded(
99    writer: &mut WsWriter,
100    message: Message,
101) -> Result<(), WebSocketCloseCode> {
102    with_outbound_write_timeout(async {
103        writer
104            .send(message)
105            .await
106            .map_err(|_error| WebSocketCloseCode::Error)
107    })
108    .await
109}
110
111pub(super) async fn close_writer_bounded(writer: &mut WsWriter, code: WebSocketCloseCode) {
112    let _closed = with_outbound_write_timeout(async {
113        let _result = writer
114            .send(Message::Close(Some(CloseFrame {
115                code: u16::from(code),
116                reason: "".into(),
117            })))
118            .await;
119        Ok(())
120    })
121    .await;
122}
123
124/// plain messages are batched until a synchronous request or response has to
125/// cross the socket so control-flow envelopes stay in order
126pub(super) async fn send_user_signals(
127    writer: &mut WsWriter,
128    signals: UserOutput,
129) -> Result<usize, WebSocketCloseCode> {
130    if signals.is_empty() {
131        return Ok(0);
132    }
133    let mut batch_count = 0;
134    let mut pending_messages = Vec::with_capacity(signals.len().min(MAX_ENVELOPE_BATCH_LEN));
135    for signal in signals {
136        match signal {
137            ServerEnvelope::Message(_) => {
138                pending_messages.push(
139                    signal
140                        .into_envelope()
141                        .map_err(|_error| WebSocketCloseCode::Error)?,
142                );
143            }
144            ServerEnvelope::Request { .. } | ServerEnvelope::Response { .. } => {
145                batch_count += send_pending_messages(writer, &mut pending_messages).await?;
146                let envelope = signal
147                    .into_envelope()
148                    .map_err(|_error| WebSocketCloseCode::Error)?;
149                send_serialized_batch(writer, &[envelope]).await?;
150                batch_count += 1;
151            }
152        }
153    }
154    batch_count += send_pending_messages(writer, &mut pending_messages).await?;
155    Ok(batch_count)
156}
157
158async fn send_pending_messages(
159    writer: &mut WsWriter,
160    pending_messages: &mut EnvelopeBatch,
161) -> Result<usize, WebSocketCloseCode> {
162    if pending_messages.is_empty() {
163        return Ok(0);
164    }
165    let mut batch_count = 0;
166    for batch in pending_messages.chunks(MAX_ENVELOPE_BATCH_LEN) {
167        send_serialized_batch(writer, batch).await?;
168        batch_count += 1;
169    }
170    pending_messages.clear();
171    Ok(batch_count)
172}
173
174async fn send_serialized_batch(
175    writer: &mut WsWriter,
176    batch: &[Envelope],
177) -> Result<(), WebSocketCloseCode> {
178    let frame = serde_json::to_string(batch).map_err(|_error| WebSocketCloseCode::Error)?;
179    writer
180        .send(Message::Text(frame.into()))
181        .await
182        .map_err(|_error| WebSocketCloseCode::Error)
183}
184
185async fn with_outbound_write_timeout<T>(
186    operation: impl Future<Output = Result<T, WebSocketCloseCode>>,
187) -> Result<T, WebSocketCloseCode> {
188    match timeout(OUTBOUND_WRITE_TIMEOUT, operation).await {
189        Ok(result) => result,
190        Err(_elapsed) => Err(WebSocketCloseCode::Error),
191    }
192}
193
194#[cfg(test)]
195#[path = "TESTS/io.rs"]
196mod tests;
o_sfu/runtime/websocket_server/io.rs

o_sfu/runtime/websocket_server/
io.rs
