1use std::{collections::BTreeMap, mem::replace};
9
10use serde::{Deserialize, Serialize};
11
12mod connection_lifecycle;
13mod outbound_batch;
14mod request_flow;
15mod request_tracker;
16mod server_events;
17mod sticky_replay;
18mod timers;
19
20use outbound_batch::{FlushMode, OutboundBatcher};
21use request_tracker::RequestTracker;
22use sticky_replay::StickyReplayState;
23use timers::RequestTimeoutId;
24
25use crate::{
26 shared::{
27 AvailableFeatures, DownloadStates, JsonPayload, RecordingState, RecordingStateUpdate,
28 StreamType, UserId, UserInfo,
29 },
30 signaling::{
31 AuthPayload, ClientBroadcastPayload, ClientEnvelope, ClientMessage, Envelope,
32 MAX_ENVELOPE_BATCH_LEN, NegotiationUploadSlot, PeerSnapshot, RecordingOptions, RequestId,
33 ServerEnvelope, StreamIntentPayload, SubscribePayload, TrackBinding, WebSocketCloseCode,
34 WelcomePayload, decode_envelope_batch,
35 },
36 wire::ServerMessage,
37};
38
39pub const RECOVERY_TIMER_ID: u32 = 1;
41const BATCH_FLUSH_TIMER_ID: u32 = 2;
42const INITIAL_RECOVERY_DELAY_MS: u32 = 1_000;
43const MAX_RECOVERY_DELAY_MS: u32 = 30_000;
44const BATCH_FLUSH_DELAY_MS: u32 = 100;
45const REQUEST_TIMEOUT_MS: u32 = 5_000;
46const MAX_OUTBOUND_BATCH_LEN: usize = 16;
47
48#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
52#[serde(tag = "kind", rename_all = "camelCase")]
53pub enum Command {
54 SendWebSocket {
56 frame: String,
57 },
58 ApplyNegotiation {
60 #[serde(rename = "requestId")]
61 request_id: RequestId,
62 #[serde(rename = "negotiationKind")]
63 kind: NegotiationKind,
64 sdp: String,
65 #[serde(rename = "uploadSlots")]
66 upload_slots: Vec<NegotiationUploadSlot>,
67 },
68 ClosePeerConnection,
69 CloseWebSocket {
70 code: u16,
71 },
72 EmitStateChange {
75 state: ConnectionState,
76 cause: Option<String>,
77 },
78 SetAvailableFeatures {
79 features: AvailableFeatures,
80 },
81 SetRecordingState {
82 state: RecordingState,
83 },
84 #[serde(rename = "emitUpdate")]
86 EmitEvent {
87 #[serde(
88 rename = "update",
89 serialize_with = "crate::host_bridge::serialize_protocol_event"
90 )]
91 event: ProtocolEvent,
92 },
93 BeginPendingRequest {
94 request: PendingRequest,
95 },
96 CompletePendingRequest {
98 #[serde(rename = "requestId")]
99 request_id: RequestId,
100 #[serde(rename = "timeoutTimerId")]
101 timeout_timer_id: u32,
102 ok: bool,
103 },
104 ScheduleTimer {
107 id: u32,
108 ms: u32,
109 },
110 CancelTimer {
111 id: u32,
112 },
113 Connect {
115 url: String,
116 },
117}
118
119pub(crate) type Commands = Vec<Command>;
120
121#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
122#[serde(rename_all = "snake_case")]
123pub enum ConnectionState {
124 Disconnected,
125 Connecting,
126 Authenticated,
127 Connected,
128 Recovering,
129 Closed,
130}
131
132#[derive(Debug, Clone, PartialEq, Eq)]
133pub enum ProtocolEvent {
134 PeerSnapshot {
135 peers: Vec<PeerSnapshot>,
136 },
137 TrackSnapshot {
138 bindings: Vec<TrackBinding>,
139 },
140 PeerInfo {
141 user_id: UserId,
142 info: UserInfo,
143 },
144 PeerLeft {
145 user_id: UserId,
146 },
147 Broadcast {
148 sender_id: UserId,
149 message: JsonPayload,
150 },
151 RecordingStateChanged {
152 state: RecordingStateUpdate,
153 },
154}
155
156#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
157#[serde(rename_all = "camelCase")]
158pub enum NegotiationKind {
159 Offer,
160 Renegotiate,
161}
162
163#[derive(Debug, Clone, Copy, PartialEq, Eq)]
164pub(crate) enum PendingRequestKind {
165 StartRecording,
166 StopRecording,
167}
168
169#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
170#[serde(rename_all = "camelCase")]
171pub struct PendingRequest {
172 pub request_id: RequestId,
173 pub timeout_timer_id: u32,
174 pub timeout_ms: u32,
175}
176
177#[derive(Debug, Clone, PartialEq, Eq)]
178struct ConnectContext {
179 url: String,
180 jwt: String,
181 room: Option<String>,
182}
183
184#[derive(Debug, Clone, PartialEq, Eq)]
185struct PendingNegotiation {
186 request_id: RequestId,
187 kind: NegotiationKind,
188}
189
190#[derive(Debug, Clone, PartialEq, Eq)]
191enum ProtocolPhase {
192 Disconnected,
193 Connecting,
194 Authenticated(NegotiationSlot),
195 Connected(NegotiationSlot),
196 Recovering,
197 Closed,
198}
199
200impl ProtocolPhase {
201 const fn connection_state(&self) -> ConnectionState {
202 match self {
203 Self::Disconnected => ConnectionState::Disconnected,
204 Self::Connecting => ConnectionState::Connecting,
205 Self::Authenticated(_) => ConnectionState::Authenticated,
206 Self::Connected(_) => ConnectionState::Connected,
207 Self::Recovering => ConnectionState::Recovering,
208 Self::Closed => ConnectionState::Closed,
209 }
210 }
211
212 const fn is_awaiting_welcome(&self) -> bool {
213 matches!(self, Self::Connecting | Self::Recovering)
214 }
215
216 fn apply_lifecycle_state(&mut self, state: ConnectionState) {
217 if self.connection_state() == state {
218 return;
219 }
220 let current = replace(self, Self::Disconnected);
221 *self = match (current, state) {
222 (Self::Authenticated(slot), ConnectionState::Connected) => Self::Connected(slot),
223 (_, ConnectionState::Disconnected) => Self::Disconnected,
224 (_, ConnectionState::Connecting) => Self::Connecting,
225 (_, ConnectionState::Authenticated) => Self::Authenticated(NegotiationSlot::Idle),
226 (_, ConnectionState::Connected) => Self::Connected(NegotiationSlot::Idle),
227 (_, ConnectionState::Recovering) => Self::Recovering,
228 (_, ConnectionState::Closed) => Self::Closed,
229 };
230 }
231
232 const fn can_send_client_messages(&self) -> bool {
233 matches!(self, Self::Authenticated(_) | Self::Connected(_))
234 }
235
236 const fn can_enter_connected(&self) -> bool {
237 matches!(self, Self::Authenticated(NegotiationSlot::Idle))
238 }
239
240 fn accept_negotiation(
241 &mut self,
242 request_id: &RequestId,
243 kind: NegotiationKind,
244 ) -> Result<(), NegotiationRejection> {
245 match (self, kind) {
246 (Self::Authenticated(slot), NegotiationKind::Offer)
247 | (Self::Connected(slot), NegotiationKind::Renegotiate) => {
248 slot.accept(request_id, kind)
249 }
250 (Self::Authenticated(_), NegotiationKind::Renegotiate)
251 | (Self::Connected(_), NegotiationKind::Offer) => {
252 Err(NegotiationRejection::ProtocolError)
253 }
254 (
255 Self::Disconnected | Self::Connecting | Self::Recovering | Self::Closed,
256 NegotiationKind::Offer | NegotiationKind::Renegotiate,
257 ) => Err(NegotiationRejection::Ignored),
258 }
259 }
260
261 fn resolve_negotiation(&mut self, request_id: &RequestId, kind: NegotiationKind) -> bool {
262 match self {
263 Self::Authenticated(slot) | Self::Connected(slot) => slot.resolve(request_id, kind),
264 Self::Disconnected | Self::Connecting | Self::Recovering | Self::Closed => false,
265 }
266 }
267}
268
269#[derive(Debug, Clone, PartialEq, Eq)]
270enum NegotiationSlot {
271 Idle,
272 WaitingForAnswer(PendingNegotiation),
273}
274
275impl NegotiationSlot {
276 fn accept(
277 &mut self,
278 request_id: &RequestId,
279 kind: NegotiationKind,
280 ) -> Result<(), NegotiationRejection> {
281 match self {
282 Self::Idle => {
283 *self = Self::WaitingForAnswer(PendingNegotiation {
284 request_id: request_id.clone(),
285 kind,
286 });
287 Ok(())
288 }
289 Self::WaitingForAnswer(_) => Err(NegotiationRejection::ProtocolError),
290 }
291 }
292
293 fn resolve(&mut self, request_id: &RequestId, kind: NegotiationKind) -> bool {
294 let Self::WaitingForAnswer(pending) = self else {
295 return false;
296 };
297 if pending.request_id != *request_id || pending.kind != kind {
298 return false;
299 }
300 *self = Self::Idle;
301 true
302 }
303}
304
305#[derive(Debug, Clone, Copy, PartialEq, Eq)]
306pub(super) enum NegotiationRejection {
307 Ignored,
308 ProtocolError,
309}
310
311#[derive(Debug, Clone, PartialEq, Eq)]
316pub struct ProtocolCore {
317 phase: ProtocolPhase,
319 track_bindings: BTreeMap<String, TrackBinding>,
324 sticky_replay: StickyReplayState,
332 connect_context: Option<ConnectContext>,
338 recovery_delay_ms: u32,
344 outbound_batch: OutboundBatcher,
351 request_tracker: RequestTracker,
357}
358
359impl Default for ProtocolCore {
360 fn default() -> Self {
361 Self::new()
362 }
363}
364
365impl ProtocolCore {
366 #[must_use]
372 pub fn new() -> Self {
373 Self {
374 phase: ProtocolPhase::Disconnected,
375 track_bindings: BTreeMap::new(),
376 sticky_replay: StickyReplayState::new(),
377 connect_context: None,
378 recovery_delay_ms: INITIAL_RECOVERY_DELAY_MS,
379 outbound_batch: OutboundBatcher::new(),
380 request_tracker: RequestTracker::new(),
381 }
382 }
383
384 #[must_use]
385 pub const fn state(&self) -> ConnectionState {
386 self.phase.connection_state()
387 }
388
389 pub fn connect(
400 &mut self,
401 url: impl Into<String>,
402 jwt: impl Into<String>,
403 room: Option<String>,
404 ) -> Vec<Command> {
405 connection_lifecycle::connect(self, url.into(), jwt.into(), room)
406 }
407
408 pub fn on_ws_open(&mut self) -> Vec<Command> {
413 if !matches!(
414 self.phase.connection_state(),
415 ConnectionState::Connecting | ConnectionState::Recovering
416 ) {
417 return Vec::new();
418 }
419 let Some(connect_context) = self.connect_context.as_ref() else {
420 return Vec::new();
421 };
422 self.enqueue_client_message(
423 ClientMessage::Auth(AuthPayload {
424 jwt: connect_context.jwt.clone(),
425 channel: connect_context.room.clone(),
426 }),
427 FlushMode::Immediate,
428 )
429 }
430
431 pub fn on_ws_message(&mut self, frame: &str) -> Vec<Command> {
437 let Ok(batch) = decode_envelope_batch(frame, MAX_ENVELOPE_BATCH_LEN) else {
438 return close_for_protocol_error();
439 };
440 let Ok(envelopes) = batch
441 .into_iter()
442 .map(ServerEnvelope::decode)
443 .collect::<Result<Vec<_>, _>>()
444 else {
445 return close_for_protocol_error();
446 };
447 let mut commands = Vec::new();
448 for envelope in envelopes {
449 match envelope {
450 ServerEnvelope::Message(message) => {
451 if self.phase.is_awaiting_welcome()
452 && !matches!(message, ServerMessage::Welcome(_))
453 {
454 return close_for_protocol_error();
455 }
456 commands.extend(server_events::handle_server_message(self, message));
457 }
458 ServerEnvelope::Request {
459 request_id,
460 request,
461 } => {
462 commands.extend(request_flow::handle_server_request(
463 self, request_id, request,
464 ));
465 }
466 ServerEnvelope::Response {
467 response_to,
468 response,
469 } => {
470 commands.extend(request_flow::handle_server_response(
471 self,
472 &response_to,
473 response,
474 ));
475 }
476 }
477 }
478 commands
479 }
480
481 fn accept_welcome(&mut self, payload: WelcomePayload) -> Commands {
482 if !matches!(
483 self.phase.connection_state(),
484 ConnectionState::Connecting | ConnectionState::Recovering
485 ) {
486 return Vec::new();
487 }
488 let WelcomePayload {
489 features,
490 recording,
491 peers,
492 } = payload;
493 self.recovery_delay_ms = INITIAL_RECOVERY_DELAY_MS;
494 self.phase
495 .apply_lifecycle_state(ConnectionState::Authenticated);
496
497 let mut commands = vec![
498 Command::SetAvailableFeatures { features },
499 Command::SetRecordingState { state: recording },
500 Command::EmitStateChange {
501 state: self.phase.connection_state(),
502 cause: None,
503 },
504 ];
505 if !peers.is_empty() {
506 commands.push(Command::EmitEvent {
507 event: ProtocolEvent::PeerSnapshot { peers },
508 });
509 }
510 commands.extend(self.replay_session_state());
511 commands
512 }
513
514 pub fn on_transport_ready(&mut self) -> Vec<Command> {
520 if !self.phase.can_enter_connected() {
521 return Vec::new();
522 }
523 self.phase.apply_lifecycle_state(ConnectionState::Connected);
524 let mut commands = vec![Command::EmitStateChange {
525 state: self.state(),
526 cause: None,
527 }];
528 commands.extend(self.replay_publication_state());
529 commands
530 }
531
532 pub fn publish(&mut self, stream_type: StreamType, active: bool) -> Vec<Command> {
537 self.sticky_replay.set_publish_active(stream_type, active);
538 if !matches!(&self.phase, ProtocolPhase::Connected(_)) {
539 return Vec::new();
540 }
541 let message = if active {
542 ClientMessage::Publish(StreamIntentPayload { stream_type })
543 } else {
544 ClientMessage::Unpublish(StreamIntentPayload { stream_type })
545 };
546 self.enqueue_client_message(message, FlushMode::Batched)
547 }
548
549 pub fn subscribe(&mut self, user_id: UserId, states: DownloadStates) -> Vec<Command> {
555 self.sticky_replay
556 .remember_subscription_states(&user_id, &states);
557 if !self.can_send_client_messages() {
558 return Vec::new();
559 }
560 self.enqueue_client_message(
561 ClientMessage::Subscribe(SubscribePayload { user_id, states }),
562 FlushMode::Batched,
563 )
564 }
565
566 pub fn update_info(&mut self, info: UserInfo) -> Vec<Command> {
572 self.sticky_replay.remember_info(&info);
573 if !self.can_send_client_messages() {
574 return Vec::new();
575 }
576 self.enqueue_client_message(ClientMessage::Info(info), FlushMode::Batched)
577 }
578
579 pub fn broadcast(&mut self, message: JsonPayload) -> Vec<Command> {
585 if !self.can_send_client_messages() {
586 return Vec::new();
587 }
588 self.enqueue_client_message(
589 ClientMessage::Broadcast(ClientBroadcastPayload { message }),
590 FlushMode::Batched,
591 )
592 }
593
594 pub fn start_recording(&mut self, options: RecordingOptions) -> Vec<Command> {
595 request_flow::start_recording(self, options)
596 }
597 pub fn stop_recording(&mut self) -> Vec<Command> {
598 request_flow::stop_recording(self)
599 }
600
601 pub fn submit_negotiation_answer(
607 &mut self,
608 request_id: &RequestId,
609 kind: NegotiationKind,
610 sdp: impl Into<String>,
611 ) -> Vec<Command> {
612 request_flow::submit_negotiation_answer(self, request_id, kind, sdp)
613 }
614
615 pub fn disconnect(&mut self) -> Vec<Command> {
616 connection_lifecycle::disconnect(self)
617 }
618
619 pub fn on_ws_close(&mut self, code: u16) -> Vec<Command> {
620 connection_lifecycle::on_ws_close(self, code)
621 }
622
623 pub fn on_timer(&mut self, timer_id: u32) -> Vec<Command> {
629 if timer_id == RECOVERY_TIMER_ID {
630 return connection_lifecycle::handle_recovery_timer(self);
631 }
632 if timer_id == BATCH_FLUSH_TIMER_ID {
633 return self.flush_pending_batch(false);
634 }
635 if let Some(commands) = RequestTimeoutId::try_from_raw(timer_id)
636 .and_then(|timeout_id| self.request_tracker.resolve_timeout(timeout_id))
637 {
638 return commands;
639 }
640 Vec::new()
641 }
642
643 fn enqueue_envelope(&mut self, envelope: Envelope, mode: FlushMode) -> Commands {
644 self.outbound_batch.enqueue(envelope, mode)
645 }
646
647 fn enqueue_client_message(&mut self, message: ClientMessage, mode: FlushMode) -> Commands {
648 let Some(envelope) = ClientEnvelope::Message(message).into_envelope().ok() else {
649 return Vec::new();
650 };
651 self.enqueue_envelope(envelope, mode)
652 }
653
654 fn flush_pending_batch(&mut self, cancel_timer: bool) -> Commands {
655 self.outbound_batch.flush(cancel_timer)
656 }
657
658 fn clear_runtime_state(&mut self) {
659 self.track_bindings.clear();
660 self.outbound_batch.clear();
661 self.request_tracker.clear();
662 }
663
664 fn teardown_runtime_state(&mut self) -> Commands {
670 let mut commands = self.outbound_batch.discard_pending();
671 commands.extend(self.request_tracker.fail_all());
672 if !self.track_bindings.is_empty() {
673 self.track_bindings.clear();
674 commands.push(Command::EmitEvent {
675 event: ProtocolEvent::TrackSnapshot {
676 bindings: Vec::new(),
677 },
678 });
679 }
680 commands
681 }
682
683 fn clear_sticky_state(&mut self) {
684 self.sticky_replay.clear();
685 }
686
687 fn replay_session_state(&mut self) -> Commands {
689 if !self.can_send_client_messages() {
690 return Vec::new();
691 }
692 let Some(replay_batch) = self.sticky_replay.replay_session_batch() else {
693 return Vec::new();
694 };
695
696 self.outbound_batch.extend(replay_batch);
697 self.flush_pending_batch(true)
698 }
699
700 fn replay_publication_state(&mut self) -> Commands {
702 if !self.can_send_client_messages() {
703 return Vec::new();
704 }
705
706 let mut replay_batch = Vec::new();
707 for stream_type in self.sticky_replay.active_publications() {
708 let Some(envelope) =
709 ClientEnvelope::Message(ClientMessage::Publish(StreamIntentPayload {
710 stream_type,
711 }))
712 .into_envelope()
713 .ok()
714 else {
715 continue;
716 };
717 replay_batch.push(envelope);
718 }
719 if replay_batch.is_empty() {
720 return Vec::new();
721 }
722
723 self.outbound_batch.extend(replay_batch);
724 self.flush_pending_batch(true)
725 }
726
727 fn can_send_client_messages(&self) -> bool {
728 self.phase.can_send_client_messages()
729 }
730}
731
732fn empty_features() -> AvailableFeatures {
733 AvailableFeatures {
734 rtc: false,
735 transcription: false,
736 audio_recording: false,
737 video_recording: false,
738 }
739}
740
741fn close_for_protocol_error() -> Commands {
742 vec![Command::CloseWebSocket {
743 code: u16::from(WebSocketCloseCode::ProtocolError),
744 }]
745}
746
747fn next_recovery_delay(current_delay_ms: u32) -> u32 {
752 current_delay_ms
753 .saturating_mul(3)
754 .checked_div(2)
755 .unwrap_or(MAX_RECOVERY_DELAY_MS)
756 .min(MAX_RECOVERY_DELAY_MS)
757}
758
759#[cfg(test)]
760#[path = "core/TESTS/mod.rs"]
761mod tests;