Skip to main content

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}