bamboo_server/handlers/agent/ws_v2/mod.rs
1//! v2-P1 unified WebSocket multiplex: `GET /v2/stream`.
2//!
3//! One WebSocket replaces the two v1 SSE streams (`GET /events/{id}` +
4//! `GET /stream`) plus a minimal `stop` control uplink, so a mobile client uses
5//! ONE connection instead of two SSE + scattered POSTs. The v1 SSE/REST
6//! endpoints stay UNCHANGED (dual-track per `docs/api-v2-transport.md` §8.1).
7//!
8//! ## Channels (server→client)
9//! - `feed` — the account `ChangeEvent` stream (reuses `plan_replay` + journal).
10//! - `agent.{session_id}` — a per-session `AgentEvent` stream (reuses the v1
11//! critical-event replay + the v1 `Coalescer` for token batching).
12//!
13//! ## Control (client→server, P1)
14//! - `stop` only — cancels a session (reuses the v1 stop discipline).
15//!
16//! ## Design
17//! Each subscribed channel runs its OWN forwarder task with its OWN broadcast
18//! receiver (per-channel lag independence; §10-Q3) AND its OWN bounded outbound
19//! `mpsc` queue. The driver holds a `StreamMap<channel, ReceiverStream>` and
20//! drains every per-channel queue with a fair (randomized-start) merge, so a burst
21//! on one channel can no longer head-of-line another at the socket (RFC §10-Q3 —
22//! see `forwarders.rs` for the honest fairness guarantee). A per-channel
23//! `JoinHandle` lets `unsubscribe`/teardown abort exactly that forwarder and drop
24//! its queue, leaving no orphaned broadcast reader.
25//!
26//! ## Encoding (v2-P3, #181)
27//! The wire encoding is negotiated ONCE at the upgrade from the offered
28//! `Sec-WebSocket-Protocol`: `bamboo.v2.msgpack` selects binary MessagePack,
29//! anything else (or nothing) stays JSON text (the default — desktop /
30//! debuggability). The SAME envelope schema is carried either way; only the
31//! serialization + WS frame type (Text vs Binary) differs. See `envelope::Encoding`.
32//!
33//! ## DEFERRED (later slices; see PR / §5.3)
34//! - `execute` / `approve` over control (handler refactors) — clients keep REST.
35//!
36//! ## Auth (v2-P2, #181 / #189)
37//! `/v2/stream` is on the public route whitelist, so the upgrade OPENS without
38//! a middleware credential — this is the only way a browser device-token client
39//! (which cannot set `Authorization`/`X-Device-Id` headers on a WS upgrade) can
40//! present its token, via the `hello` frame. The handler is therefore the
41//! AUTHORITATIVE gate:
42//!
43//! - `pre_authorized` is computed from the SAME allow-decision the middleware
44//! uses for every other route (`request_is_authorized`): a local bypass, a
45//! verified password cookie (cookies ARE sent on the upgrade), or a header
46//! device token. These connections stay frictionless — exactly as before.
47//! - While NOT authorized, the ONLY frame that does anything is `hello` carrying
48//! a VALID `device_id` + `token`. A `subscribe`/`unsubscribe`/`stop` received
49//! before authorization is IGNORED — no forwarder, no cancel, no channel data.
50//! - A token-less `hello` NEVER authorizes a connection that wasn't already
51//! pre-authorized.
52//! - An unauthenticated socket that never sends a valid `hello` within
53//! [`AUTH_DEADLINE`] is CLOSED, so it cannot linger.
54//! - The token is NEVER logged.
55
56mod envelope;
57mod forwarders;
58
59use std::collections::HashMap;
60use std::time::Duration;
61
62use actix_web::{web, HttpRequest, Responder};
63use actix_ws::Message;
64use futures::StreamExt;
65use tokio::sync::mpsc;
66use tokio_stream::wrappers::ReceiverStream;
67use tokio_stream::StreamMap;
68
69use serde::Deserialize;
70
71use self::envelope::{
72 decode_client_frame, Channel, ClientFrame, Encoding, OutFrame, SUBPROTOCOL_JSON,
73 SUBPROTOCOL_MSGPACK,
74};
75use self::forwarders::{spawn_agent_forwarder, spawn_feed_forwarder, OutboundTx};
76use crate::app_state::AppState;
77use crate::handlers::agent::events::MAX_BATCH_MS;
78use crate::handlers::agent::stop::cancel_session;
79
80/// Bound on EACH per-channel outbound mpsc (RFC §10-Q3). Caps how far ahead one
81/// channel's forwarder may run before the WS writer applies backpressure to that
82/// channel ALONE — a slow socket no longer lets a bursting channel starve the
83/// queue of another, because each channel owns its own bounded buffer and the
84/// driver drains them with a fair merge.
85const OUTBOUND_BUFFER: usize = 64;
86
87/// WS ping interval (~15s), mirroring the v1 SSE `[KEEPALIVE]` cadence.
88const PING_INTERVAL: Duration = Duration::from_secs(15);
89
90/// How long an UNAUTHORIZED connection may stay open before it must present a
91/// valid `hello` device token. A connection that opens the (now public) upgrade
92/// without a header/cookie credential and never authenticates is CLOSED when
93/// this elapses, so an unauthenticated socket can never linger (#189). Once the
94/// connection is authorized this deadline is disarmed.
95const AUTH_DEADLINE: Duration = Duration::from_secs(10);
96
97/// Env var the live integration test sets to SHORTEN the auth deadline so the
98/// "closed when no hello arrives" path is asserted in milliseconds, not 10s. It
99/// is read once per connection in [`drive`]. Production NEVER sets it, so the
100/// 10s [`AUTH_DEADLINE`] default is unchanged.
101const AUTH_DEADLINE_OVERRIDE_ENV: &str = "BAMBOO_WS_AUTH_DEADLINE_MS";
102
103/// The effective unauthorized deadline: the [`AUTH_DEADLINE`] default unless
104/// [`AUTH_DEADLINE_OVERRIDE_ENV`] is set to a valid millisecond count (test-only).
105fn auth_deadline() -> Duration {
106 std::env::var(AUTH_DEADLINE_OVERRIDE_ENV)
107 .ok()
108 .and_then(|v| v.parse::<u64>().ok())
109 .map(Duration::from_millis)
110 .unwrap_or(AUTH_DEADLINE)
111}
112
113/// The verdict for one client frame while a connection is NOT yet authorized.
114///
115/// Pure decision over the parsed frame (no `AppState`), so the auth-gating
116/// contract is unit-testable without a live WS driver: while unauthorized the
117/// ONLY frame that can change anything is a `hello`, and only a `hello` carrying
118/// a credential is even a candidate to authorize.
119#[derive(Debug, PartialEq, Eq)]
120enum UnauthorizedAction {
121 /// A `hello` carrying `device_id` + `token` — verify it; valid → authorize,
122 /// invalid → close.
123 VerifyHello,
124 /// A token-less `hello` while unauthorized — a harmless no-op that does NOT
125 /// grant access (the deadline still governs).
126 TokenlessHelloNoop,
127 /// Any other frame (`subscribe`/`unsubscribe`/`stop`/unknown) while
128 /// unauthorized — IGNORE it (serve no channel data, cancel nothing).
129 Ignore,
130}
131
132impl UnauthorizedAction {
133 /// Classify a parsed frame for an unauthorized connection.
134 fn classify(frame: &ClientFrame) -> Self {
135 match frame {
136 ClientFrame::Hello {
137 device_id: Some(_),
138 token: Some(_),
139 } => UnauthorizedAction::VerifyHello,
140 ClientFrame::Hello { .. } => UnauthorizedAction::TokenlessHelloNoop,
141 _ => UnauthorizedAction::Ignore,
142 }
143 }
144}
145
146/// What the driver should do with a frame after the auth gate ran.
147#[derive(Debug, PartialEq, Eq)]
148enum GateOutcome {
149 /// The frame was fully handled by the gate (the connection is/was
150 /// unauthorized, or it just authorized): keep the socket open, do NOT fall
151 /// through to channel dispatch.
152 Handled,
153 /// An invalid credential was presented: CLOSE the connection.
154 Close,
155 /// The connection is authorized and this frame should proceed to the normal
156 /// channel dispatch (subscribe/unsubscribe/stop/hello-rebind).
157 Dispatch,
158}
159
160/// The AppState-aware auth gate, factored out of `handle_client_frame` so it is
161/// unit-testable WITHOUT a live WS `Session`. It owns the entire `!authorized`
162/// decision plus the credential verification, and flips `authorized` to `true`
163/// only on a VERIFIED `hello` device token.
164///
165/// Invariants (the security review hammers these):
166/// - While `!*authorized`, a `subscribe`/`unsubscribe`/`stop` returns
167/// [`GateOutcome::Handled`] (ignored — never `Dispatch`), so no channel is
168/// served and no session cancelled before authorization.
169/// - A token-less `hello` NEVER sets `*authorized` when it was `false`.
170/// - A credentialed `hello` with a VALID token sets `*authorized = true`.
171/// - A credentialed `hello` with an INVALID token returns [`GateOutcome::Close`].
172/// - The token is NEVER logged.
173async fn apply_auth_gate(
174 state: &web::Data<AppState>,
175 frame: &ClientFrame,
176 authorized: &mut bool,
177) -> GateOutcome {
178 if *authorized {
179 return GateOutcome::Dispatch;
180 }
181
182 match UnauthorizedAction::classify(frame) {
183 UnauthorizedAction::VerifyHello => {
184 let ClientFrame::Hello {
185 device_id: Some(device_id),
186 token: Some(token),
187 } = frame
188 else {
189 unreachable!("VerifyHello implies a credentialed Hello");
190 };
191 let config = state.config.read().await.clone();
192 if crate::handlers::settings::verify_device_token(&config, device_id, token) {
193 *authorized = true;
194 // Bind device id for logging. NEVER log the token.
195 tracing::debug!("ws_v2: hello verified for device {device_id}; authorized");
196 GateOutcome::Handled
197 } else {
198 tracing::warn!(
199 "ws_v2: hello rejected — invalid device credential for {device_id}; closing"
200 );
201 GateOutcome::Close
202 }
203 }
204 UnauthorizedAction::TokenlessHelloNoop => {
205 // A token-less hello does NOT authorize a non-pre-authorized
206 // connection. Keep the socket open; the deadline still governs.
207 tracing::debug!("ws_v2: token-less hello while unauthorized — not granting access");
208 GateOutcome::Handled
209 }
210 UnauthorizedAction::Ignore => {
211 // subscribe/unsubscribe/stop/unknown before auth: serve nothing,
212 // cancel nothing. Tolerate hello-after-subscribe ordering.
213 tracing::debug!("ws_v2: ignoring frame on unauthorized connection");
214 GateOutcome::Handled
215 }
216 }
217}
218
219/// Query parameters for the `GET /v2/stream` upgrade.
220#[derive(Debug, Default, Deserialize)]
221pub struct StreamQuery {
222 /// Token-coalescing window in milliseconds for `agent.{sid}` channels.
223 /// `0` (default) = no coalescing (desktop). Mobile passes e.g. `50`.
224 #[serde(default)]
225 pub batch_ms: u64,
226}
227
228/// Negotiate the wire [`Encoding`] from the upgrade's `Sec-WebSocket-Protocol`
229/// header (v2-P3, §5.3 / §7.2). Returns the chosen encoding AND the subprotocol
230/// token the server must ECHO on the upgrade response (per RFC 6455 the server
231/// echoes the single selected subprotocol).
232///
233/// - Offers including `bamboo.v2.msgpack` → `(Msgpack, Some("bamboo.v2.msgpack"))`.
234/// - Offers including only `bamboo.v2` → `(Json, Some("bamboo.v2"))`.
235/// - No (recognized) subprotocol offered → `(Json, None)` — today's behavior is
236/// preserved byte-for-byte for clients that offer nothing.
237///
238/// `bamboo.v2.msgpack` wins if BOTH are offered (the client opted into binary).
239fn negotiate_encoding(req: &HttpRequest) -> (Encoding, Option<&'static str>) {
240 let offered = req
241 .headers()
242 .get(actix_web::http::header::SEC_WEBSOCKET_PROTOCOL)
243 .and_then(|v| v.to_str().ok())
244 .unwrap_or("");
245 // The header is a comma-separated list of subprotocol tokens.
246 let mut has_msgpack = false;
247 let mut has_json = false;
248 for tok in offered.split(',') {
249 match tok.trim() {
250 SUBPROTOCOL_MSGPACK => has_msgpack = true,
251 SUBPROTOCOL_JSON => has_json = true,
252 _ => {}
253 }
254 }
255 if has_msgpack {
256 (Encoding::Msgpack, Some(SUBPROTOCOL_MSGPACK))
257 } else if has_json {
258 (Encoding::Json, Some(SUBPROTOCOL_JSON))
259 } else {
260 (Encoding::Json, None)
261 }
262}
263
264/// `GET /v2/stream` — upgrade to the unified WS multiplex.
265///
266/// The upgrade itself is PUBLIC (whitelisted in `is_public_access_route`), so a
267/// browser device-token client can open the socket and authenticate via `hello`
268/// (it cannot set headers on a WS upgrade). Auth is enforced in `drive`: a
269/// connection that the middleware WOULD have allowed is `pre_authorized` here
270/// (local bypass / verified password cookie / header device token), via the same
271/// `request_is_authorized` allow-decision the middleware uses; everything else
272/// must present a verified `hello` before any channel is served, on a deadline.
273///
274/// The wire [`Encoding`] is negotiated from `Sec-WebSocket-Protocol` (v2-P3):
275/// `bamboo.v2.msgpack` → binary MessagePack; otherwise JSON text (default). The
276/// selected subprotocol is ECHOED on the upgrade response per RFC 6455.
277pub async fn handler(
278 state: web::Data<AppState>,
279 query: web::Query<StreamQuery>,
280 req: HttpRequest,
281 body: web::Payload,
282) -> actix_web::Result<impl Responder> {
283 // Clamp the untrusted `batch_ms` the same way the v1 SSE handler does.
284 let batch_ms = query.batch_ms.min(MAX_BATCH_MS);
285
286 // Negotiate the wire encoding from the offered subprotocols (v2-P3).
287 let (encoding, selected_subprotocol) = negotiate_encoding(&req);
288
289 // Pre-authorize exactly the connections the middleware would have allowed on
290 // a gated route: local bypass, a verified password cookie (sent on the
291 // upgrade), or a header device token. This preserves every existing client
292 // with ZERO change; a remote browser device-token client is NOT pre-auth and
293 // must send a valid `hello`.
294 let pre_authorized = {
295 let config = state.config.read().await.clone();
296 crate::handlers::settings::request_is_authorized(&req, &config)
297 };
298
299 let (mut response, session, msg_stream) = actix_ws::handle(&req, body)?;
300
301 // Echo the selected subprotocol on the upgrade RESPONSE (RFC 6455). A client
302 // that offered no recognized subprotocol gets none, keeping the legacy
303 // handshake byte-for-byte unchanged.
304 if let Some(proto) = selected_subprotocol {
305 response.headers_mut().insert(
306 actix_web::http::header::SEC_WEBSOCKET_PROTOCOL,
307 actix_web::http::header::HeaderValue::from_static(proto),
308 );
309 }
310
311 actix_web::rt::spawn(drive(
312 state,
313 session,
314 msg_stream,
315 batch_ms,
316 pre_authorized,
317 encoding,
318 ));
319
320 Ok(response)
321}
322
323/// The per-connection driver: owns the WS `session` (write) + `msg_stream`
324/// (read), the shared outbound mpsc, the per-channel forwarder handles, and the
325/// keepalive ping timer.
326async fn drive(
327 state: web::Data<AppState>,
328 mut session: actix_ws::Session,
329 mut msg_stream: actix_ws::MessageStream,
330 batch_ms: u64,
331 pre_authorized: bool,
332 encoding: Encoding,
333) {
334 // Per-channel outbound (RFC §10-Q3): every subscribed channel owns its OWN
335 // bounded queue. `forwarders` keeps the task handle (for unsubscribe/teardown
336 // abort) and `queues` keeps the matching receiver in a `StreamMap` the driver
337 // drains with a fair (randomized-start) merge, so a burst on one channel cannot
338 // head-of-line another at the socket. The two maps are keyed by the SAME
339 // channel id and mutated together (see [`subscribe`] / unsubscribe / teardown).
340 let mut forwarders: HashMap<String, tokio::task::JoinHandle<()>> = HashMap::new();
341 let mut queues: StreamMap<String, ReceiverStream<OutFrame>> = StreamMap::new();
342 let mut ping = tokio::time::interval(PING_INTERVAL);
343 ping.tick().await; // skip the immediate tick
344
345 // Authorization state (#189). Seeded from the upgrade-time decision so local
346 // / cookie / header clients are authorized immediately; everything else must
347 // present a valid `hello` before any channel is served.
348 let mut authorized = pre_authorized;
349 // The unauthorized-deadline timer. Pinned + biased to fire while `!authorized`
350 // and disarmed once the connection authorizes, so an unauthenticated socket
351 // is closed but an authorized one runs indefinitely.
352 let auth_deadline = tokio::time::sleep(auth_deadline());
353 tokio::pin!(auth_deadline);
354
355 loop {
356 tokio::select! {
357 // While unauthorized, close the socket once the deadline elapses. The
358 // `!authorized` guard disarms this arm the moment the connection
359 // authorizes (a never-completing branch is simply never selected).
360 _ = &mut auth_deadline, if !authorized => {
361 tracing::debug!("ws_v2: closing unauthorized connection after auth deadline");
362 break;
363 }
364 // Drain the per-channel queues to the WS with a fair merge. The
365 // `StreamMap` randomizes its poll start each call, so no single channel's
366 // queue is favored — a bursting channel cannot starve another at the
367 // socket. `(channel, frame)` is yielded; only the frame goes on the
368 // wire. The frame is ALREADY encoded per the connection's `Encoding`
369 // (the forwarder did the encode), so the driver only picks the WS frame
370 // type: `session.text` for a JSON `Text`, `session.binary` for a
371 // MessagePack `Binary`. If the write fails the peer is gone — break and
372 // tear down.
373 Some((_ch, frame)) = queues.next() => {
374 let write = match frame {
375 OutFrame::Text(s) => session.text(s).await,
376 OutFrame::Binary(b) => session.binary(b).await,
377 };
378 if write.is_err() {
379 break;
380 }
381 }
382 // Keepalive ping. Liveness is write-driven (a dead peer surfaces as a
383 // failed `ping`/`text` write); we do NOT track Pong arrivals or run a
384 // pong timeout — same one-directional keepalive contract as the v1 SSE
385 // stream.
386 _ = ping.tick() => {
387 if session.ping(b"").await.is_err() {
388 break;
389 }
390 }
391 // Client frames. In JSON mode the client sends TEXT frames; in msgpack
392 // mode it sends BINARY frames. A frame carrying the inbound bytes for the
393 // ACTIVE encoding is decoded + dispatched; a frame of the other kind is
394 // ignored (a Binary frame in JSON mode, a Text frame in msgpack mode),
395 // exactly as a malformed frame is — it never tears down the connection.
396 msg = msg_stream.next() => {
397 match msg {
398 // The inbound frame type that matches the active encoding.
399 Some(Ok(Message::Text(text))) if encoding == Encoding::Json => {
400 let keep_open = handle_client_bytes(
401 &state, &mut forwarders, &mut queues,
402 batch_ms, encoding, text.as_bytes(), &mut authorized,
403 )
404 .await;
405 if !keep_open {
406 break;
407 }
408 }
409 Some(Ok(Message::Binary(bytes))) if encoding == Encoding::Msgpack => {
410 let keep_open = handle_client_bytes(
411 &state, &mut forwarders, &mut queues,
412 batch_ms, encoding, &bytes, &mut authorized,
413 )
414 .await;
415 if !keep_open {
416 break;
417 }
418 }
419 Some(Ok(Message::Ping(bytes))) => {
420 if session.pong(&bytes).await.is_err() {
421 break;
422 }
423 }
424 // A frame of the WRONG kind for the active encoding (a Binary
425 // frame in JSON mode / a Text frame in msgpack mode), plus Pong /
426 // Continuation / Nop: ignore, never disconnect.
427 Some(Ok(Message::Text(_)))
428 | Some(Ok(Message::Binary(_)))
429 | Some(Ok(Message::Pong(_)))
430 | Some(Ok(Message::Continuation(_)))
431 | Some(Ok(Message::Nop)) => {}
432 Some(Ok(Message::Close(_))) | None => break,
433 Some(Err(e)) => {
434 tracing::debug!("ws_v2: message stream error: {e}");
435 break;
436 }
437 }
438 }
439 else => break,
440 }
441 }
442
443 // Clean teardown: abort every forwarder so no orphaned broadcast reader
444 // survives the connection, and drop every per-channel queue (clearing the
445 // `StreamMap` drops each receiver).
446 for (_ch, handle) in forwarders.drain() {
447 handle.abort();
448 }
449 queues.clear();
450 let _ = session.close(None).await;
451}
452
453/// Decode one inbound frame's bytes per the connection's [`Encoding`] (serde_json
454/// for `Json`, rmp-serde for `Msgpack`) and dispatch the resulting [`ClientFrame`].
455/// A malformed body logs and is ignored — it NEVER tears down the connection, in
456/// EITHER encoding (the msgpack decode error is treated exactly like a malformed
457/// JSON text frame).
458///
459/// This is the thin decode step in front of [`handle_client_frame`]; the dispatch
460/// and auth logic is encoding-agnostic (the SAME `ClientFrame` flows through both
461/// paths), so the JSON behavior is byte-for-byte unchanged.
462///
463/// Returns `false` to signal the driver to CLOSE the connection (a `hello` that
464/// presents an INVALID device credential); `true` to keep it open.
465async fn handle_client_bytes(
466 state: &web::Data<AppState>,
467 forwarders: &mut HashMap<String, tokio::task::JoinHandle<()>>,
468 queues: &mut StreamMap<String, ReceiverStream<OutFrame>>,
469 batch_ms: u64,
470 encoding: Encoding,
471 bytes: &[u8],
472 authorized: &mut bool,
473) -> bool {
474 let frame: ClientFrame = match decode_client_frame(encoding, bytes) {
475 Ok(f) => f,
476 Err(e) => {
477 tracing::debug!("ws_v2: ignoring malformed client frame: {e}");
478 return true;
479 }
480 };
481 handle_client_frame(
482 state, forwarders, queues, batch_ms, encoding, frame, authorized,
483 )
484 .await
485}
486
487/// Dispatch one decoded client frame. A malformed/unknown frame logs and is
488/// ignored — it never tears down the connection.
489///
490/// `authorized` is the per-connection auth state (#189). While it is `false` the
491/// ONLY frame that does anything is a `hello` carrying a VALID device token (it
492/// flips `authorized` to `true`); EVERY other frame — including a
493/// `subscribe`/`unsubscribe`/`stop` — is IGNORED, so a remote unauthenticated
494/// connection gets NO channel data and cancels nothing. The driver's deadline
495/// closes a socket that never authorizes.
496///
497/// Returns `false` to signal the driver to CLOSE the connection (a `hello` that
498/// presents an INVALID device credential); `true` to keep it open.
499async fn handle_client_frame(
500 state: &web::Data<AppState>,
501 forwarders: &mut HashMap<String, tokio::task::JoinHandle<()>>,
502 queues: &mut StreamMap<String, ReceiverStream<OutFrame>>,
503 batch_ms: u64,
504 encoding: Encoding,
505 frame: ClientFrame,
506 authorized: &mut bool,
507) -> bool {
508 // Auth gate (#189). Until the connection is authorized, no frame may serve a
509 // channel or cancel a session — the only frame that can change anything is a
510 // `hello` that carries a verifiable device credential.
511 match apply_auth_gate(state, &frame, authorized).await {
512 GateOutcome::Handled => return true,
513 GateOutcome::Close => return false,
514 GateOutcome::Dispatch => {}
515 }
516
517 // Authorized path: full dispatch.
518 match frame {
519 ClientFrame::Hello { device_id, token } => {
520 // Already authorized. A credentialed hello re-verifies as identity
521 // binding; an invalid one still closes. A token-less hello on an
522 // already-authorized connection is a harmless no-op.
523 match (device_id, token) {
524 (Some(device_id), Some(token)) => {
525 let config = state.config.read().await.clone();
526 if crate::handlers::settings::verify_device_token(&config, &device_id, &token) {
527 // NEVER log the token.
528 tracing::debug!("ws_v2: hello verified for device {device_id}");
529 } else {
530 tracing::warn!(
531 "ws_v2: hello rejected — invalid device credential for {device_id}; closing"
532 );
533 return false;
534 }
535 }
536 _ => {
537 tracing::debug!("ws_v2: token-less hello on authorized connection (no-op)");
538 }
539 }
540 }
541 ClientFrame::Subscribe { ch, since } => {
542 subscribe(state, forwarders, queues, batch_ms, encoding, &ch, since).await;
543 }
544 ClientFrame::Unsubscribe { ch } => {
545 if let Some(handle) = forwarders.remove(&ch) {
546 handle.abort();
547 // Drop this channel's queue too: removing it from the StreamMap
548 // drops the receiver so the (now-aborted) forwarder's sender is
549 // dead and no stale frame can leak onto the socket.
550 queues.remove(&ch);
551 tracing::debug!("ws_v2: unsubscribed {ch}");
552 }
553 }
554 ClientFrame::Stop { session_id } => {
555 let cancelled = cancel_session(state, &session_id).await;
556 tracing::debug!("ws_v2: stop {session_id} -> cancelled={cancelled}");
557 }
558 ClientFrame::Unknown => {
559 tracing::debug!("ws_v2: ignoring unknown client frame type");
560 }
561 }
562 true
563}
564
565/// Subscribe to a channel, replacing any existing forwarder for the same `ch`
566/// (a re-subscribe with a new cursor aborts the old one first, so there is never
567/// a duplicate reader on one channel).
568async fn subscribe(
569 state: &web::Data<AppState>,
570 forwarders: &mut HashMap<String, tokio::task::JoinHandle<()>>,
571 queues: &mut StreamMap<String, ReceiverStream<OutFrame>>,
572 batch_ms: u64,
573 encoding: Encoding,
574 ch: &str,
575 since: Option<u64>,
576) {
577 let Some(channel) = Channel::parse(ch) else {
578 tracing::debug!("ws_v2: ignoring subscribe to unknown channel {ch}");
579 return;
580 };
581
582 // Replace any prior forwarder for this exact channel id, dropping its queue
583 // (a re-subscribe with a new cursor must not leave the old queue behind).
584 if let Some(old) = forwarders.remove(ch) {
585 old.abort();
586 }
587 queues.remove(ch);
588
589 // This channel's OWN bounded outbound queue (RFC §10-Q3). The forwarder owns
590 // the sender; the driver drains the receiver via the fair `StreamMap` merge.
591 // Carries already-encoded [`OutFrame`]s (Text/Binary per the connection's
592 // `Encoding`).
593 let (out_tx, out_rx) = mpsc::channel::<OutFrame>(OUTBOUND_BUFFER);
594 let out_tx: OutboundTx = out_tx;
595
596 let handle = match channel {
597 Channel::Feed => {
598 // Subscribe FIRST so events written during journal replay are buffered
599 // in the ring and delivered in the live phase (no gap) — exactly the
600 // v1 SSE handoff discipline.
601 let receiver = state.account_sink.subscribe();
602 let latest_at_start = state.account_sink.latest_seq();
603 let events_dir = state.account_sink.events_dir().to_path_buf();
604 let since = since.unwrap_or(0);
605 spawn_feed_forwarder(
606 out_tx,
607 encoding,
608 receiver,
609 events_dir,
610 since,
611 latest_at_start,
612 )
613 }
614 Channel::Agent(sid) => {
615 // The session must exist; otherwise ignore (parity with the v1
616 // events handler's 404, but here we just skip the subscribe).
617 if state.session_store.get_index_entry(&sid).await.is_none() {
618 tracing::debug!("ws_v2: ignoring subscribe to unknown session {sid}");
619 return;
620 }
621
622 let sender = state.get_session_event_sender(&sid).await;
623 let receiver = sender.subscribe();
624 // Keep the notification relay running for this session (parity with
625 // the v1 events handler).
626 state.ensure_notification_relay(&sid, sender.clone());
627
628 // Snapshot the runner for critical-event + budget replay (mirrors
629 // `events/handler.rs:80-89`).
630 let runner_snapshot = {
631 let runners = state.agent_runners.read().await;
632 runners.get(&sid).cloned()
633 };
634 let budget_event_to_replay = runner_snapshot
635 .as_ref()
636 .and_then(|runner| runner.last_budget_event.clone());
637 let critical_events_to_replay: Vec<_> = runner_snapshot
638 .as_ref()
639 .map(|runner| runner.last_critical_events.clone())
640 .unwrap_or_default();
641
642 spawn_agent_forwarder(
643 state.clone(),
644 sid.clone(),
645 out_tx,
646 encoding,
647 ch.to_string(),
648 receiver,
649 budget_event_to_replay,
650 critical_events_to_replay,
651 batch_ms,
652 )
653 }
654 };
655
656 forwarders.insert(ch.to_string(), handle);
657 // Register this channel's queue receiver into the fair-merge drain. Keyed by
658 // the SAME channel id as `forwarders`, so unsubscribe/teardown drops both.
659 queues.insert(ch.to_string(), ReceiverStream::new(out_rx));
660 tracing::debug!("ws_v2: subscribed {ch} (since={since:?})");
661}
662
663#[cfg(test)]
664mod tests {
665 use super::*;
666 use crate::app_state::AppState;
667 use bamboo_config::{AccessControlConfig, DeviceCredential};
668 use tempfile::tempdir;
669
670 fn hello(device_id: Option<&str>, token: Option<&str>) -> ClientFrame {
671 ClientFrame::Hello {
672 device_id: device_id.map(str::to_string),
673 token: token.map(str::to_string),
674 }
675 }
676
677 // ── Subprotocol negotiation (v2-P3) ───────────────────────────────────────
678
679 fn negotiate_for(header: Option<&str>) -> (Encoding, Option<&'static str>) {
680 let mut req = actix_web::test::TestRequest::default();
681 if let Some(h) = header {
682 req = req.insert_header((actix_web::http::header::SEC_WEBSOCKET_PROTOCOL, h));
683 }
684 negotiate_encoding(&req.to_http_request())
685 }
686
687 #[test]
688 fn negotiate_encoding_branches() {
689 // No header / empty → JSON default, NO echo (unchanged for old clients).
690 assert_eq!(negotiate_for(None), (Encoding::Json, None));
691 assert_eq!(negotiate_for(Some("")), (Encoding::Json, None));
692 // Explicit JSON subprotocol → JSON, echo it.
693 assert_eq!(
694 negotiate_for(Some("bamboo.v2")),
695 (Encoding::Json, Some(SUBPROTOCOL_JSON))
696 );
697 // Msgpack offered → Msgpack, echo it.
698 assert_eq!(
699 negotiate_for(Some("bamboo.v2.msgpack")),
700 (Encoding::Msgpack, Some(SUBPROTOCOL_MSGPACK))
701 );
702 // Multi-offer: msgpack preferred regardless of order, and whitespace is trimmed.
703 assert_eq!(
704 negotiate_for(Some("bamboo.v2.msgpack, bamboo.v2")),
705 (Encoding::Msgpack, Some(SUBPROTOCOL_MSGPACK))
706 );
707 assert_eq!(
708 negotiate_for(Some(" bamboo.v2 , bamboo.v2.msgpack ")),
709 (Encoding::Msgpack, Some(SUBPROTOCOL_MSGPACK))
710 );
711 // Unknown-only offer → JSON, and NO bogus echo (echoing a non-offered
712 // subprotocol would violate RFC 6455).
713 assert_eq!(
714 negotiate_for(Some("some.other.proto")),
715 (Encoding::Json, None)
716 );
717 }
718
719 // ── Pure classification (no AppState) ─────────────────────────────────────
720
721 #[test]
722 fn classify_unauthorized_frame_actions() {
723 // A credentialed hello is the only authorize candidate.
724 assert_eq!(
725 UnauthorizedAction::classify(&hello(Some("d"), Some("t"))),
726 UnauthorizedAction::VerifyHello
727 );
728 // A token-less hello (any missing field) is a no-op, never an authorizer.
729 assert_eq!(
730 UnauthorizedAction::classify(&hello(None, None)),
731 UnauthorizedAction::TokenlessHelloNoop
732 );
733 assert_eq!(
734 UnauthorizedAction::classify(&hello(Some("d"), None)),
735 UnauthorizedAction::TokenlessHelloNoop
736 );
737 assert_eq!(
738 UnauthorizedAction::classify(&hello(None, Some("t"))),
739 UnauthorizedAction::TokenlessHelloNoop
740 );
741 // Every channel-touching / control frame is IGNORED while unauthorized.
742 assert_eq!(
743 UnauthorizedAction::classify(&ClientFrame::Subscribe {
744 ch: "feed".into(),
745 since: None
746 }),
747 UnauthorizedAction::Ignore
748 );
749 assert_eq!(
750 UnauthorizedAction::classify(&ClientFrame::Unsubscribe { ch: "feed".into() }),
751 UnauthorizedAction::Ignore
752 );
753 assert_eq!(
754 UnauthorizedAction::classify(&ClientFrame::Stop {
755 session_id: "s".into()
756 }),
757 UnauthorizedAction::Ignore
758 );
759 assert_eq!(
760 UnauthorizedAction::classify(&ClientFrame::Unknown),
761 UnauthorizedAction::Ignore
762 );
763 }
764
765 // ── AppState-aware gate (no live Session) ─────────────────────────────────
766
767 async fn app_state_with_device() -> (web::Data<AppState>, DeviceCredential, String) {
768 let dir = tempdir().unwrap();
769 let state = web::Data::new(AppState::new(dir.path().to_path_buf()).await.unwrap());
770 let (cred, token) = crate::handlers::settings::issue_device_token("test-device");
771 {
772 let mut config = state.config.write().await;
773 config.access_control = Some(AccessControlConfig {
774 password_enabled: false,
775 password_hash: None,
776 password_salt: None,
777 updated_at: None,
778 devices: vec![cred.clone()],
779 });
780 }
781 (state, cred, token)
782 }
783
784 #[actix_web::test]
785 async fn subscribe_while_unauthorized_is_ignored_and_stays_unauthorized() {
786 let (state, _cred, _token) = app_state_with_device().await;
787 let mut authorized = false;
788 let frame = ClientFrame::Subscribe {
789 ch: "feed".into(),
790 since: None,
791 };
792 let outcome = apply_auth_gate(&state, &frame, &mut authorized).await;
793 // The driver keeps the socket open but does NOT dispatch (no forwarder).
794 assert_eq!(outcome, GateOutcome::Handled);
795 assert!(!authorized, "a subscribe must never authorize a connection");
796 }
797
798 #[actix_web::test]
799 async fn stop_while_unauthorized_is_ignored() {
800 let (state, _cred, _token) = app_state_with_device().await;
801 let mut authorized = false;
802 let frame = ClientFrame::Stop {
803 session_id: "sess".into(),
804 };
805 let outcome = apply_auth_gate(&state, &frame, &mut authorized).await;
806 assert_eq!(outcome, GateOutcome::Handled);
807 assert!(!authorized);
808 }
809
810 #[actix_web::test]
811 async fn valid_hello_authorizes() {
812 let (state, cred, token) = app_state_with_device().await;
813 let mut authorized = false;
814 let frame = hello(Some(&cred.device_id), Some(&token));
815 let outcome = apply_auth_gate(&state, &frame, &mut authorized).await;
816 assert_eq!(outcome, GateOutcome::Handled);
817 assert!(authorized, "a valid hello must authorize the connection");
818 }
819
820 #[actix_web::test]
821 async fn invalid_hello_closes_and_does_not_authorize() {
822 let (state, cred, _token) = app_state_with_device().await;
823 let mut authorized = false;
824 let frame = hello(Some(&cred.device_id), Some("bd1_wrongwrongwrong"));
825 let outcome = apply_auth_gate(&state, &frame, &mut authorized).await;
826 assert_eq!(outcome, GateOutcome::Close);
827 assert!(!authorized, "an invalid hello must never authorize");
828 }
829
830 #[actix_web::test]
831 async fn tokenless_hello_does_not_authorize_unauthorized_connection() {
832 let (state, _cred, _token) = app_state_with_device().await;
833 let mut authorized = false;
834 let frame = hello(None, None);
835 let outcome = apply_auth_gate(&state, &frame, &mut authorized).await;
836 assert_eq!(outcome, GateOutcome::Handled);
837 assert!(
838 !authorized,
839 "a token-less hello must NEVER authorize a non-pre-authorized connection"
840 );
841 }
842
843 #[actix_web::test]
844 async fn pre_authorized_connection_dispatches_subscribe() {
845 let (state, _cred, _token) = app_state_with_device().await;
846 // Pre-authorized (local / cookie / header equivalent): the gate passes
847 // every frame straight through to channel dispatch.
848 let mut authorized = true;
849 let frame = ClientFrame::Subscribe {
850 ch: "feed".into(),
851 since: None,
852 };
853 let outcome = apply_auth_gate(&state, &frame, &mut authorized).await;
854 assert_eq!(outcome, GateOutcome::Dispatch);
855 assert!(authorized);
856 }
857}