mod envelope;
mod forwarders;
use std::collections::HashMap;
use std::time::Duration;
use actix_web::{web, HttpRequest, Responder};
use actix_ws::Message;
use futures::StreamExt;
use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream;
use tokio_stream::StreamMap;
use serde::Deserialize;
use self::envelope::{
decode_client_frame, Channel, ClientFrame, Encoding, OutFrame, SUBPROTOCOL_JSON,
SUBPROTOCOL_MSGPACK,
};
use self::forwarders::{spawn_agent_forwarder, spawn_feed_forwarder, OutboundTx};
use crate::app_state::AppState;
use crate::handlers::agent::events::MAX_BATCH_MS;
use crate::handlers::agent::stop::cancel_session;
const OUTBOUND_BUFFER: usize = 64;
const PING_INTERVAL: Duration = Duration::from_secs(15);
const AUTH_DEADLINE: Duration = Duration::from_secs(10);
const AUTH_DEADLINE_OVERRIDE_ENV: &str = "BAMBOO_WS_AUTH_DEADLINE_MS";
fn auth_deadline() -> Duration {
std::env::var(AUTH_DEADLINE_OVERRIDE_ENV)
.ok()
.and_then(|v| v.parse::<u64>().ok())
.map(Duration::from_millis)
.unwrap_or(AUTH_DEADLINE)
}
#[derive(Debug, PartialEq, Eq)]
enum UnauthorizedAction {
VerifyHello,
TokenlessHelloNoop,
Ignore,
}
impl UnauthorizedAction {
fn classify(frame: &ClientFrame) -> Self {
match frame {
ClientFrame::Hello {
device_id: Some(_),
token: Some(_),
} => UnauthorizedAction::VerifyHello,
ClientFrame::Hello { .. } => UnauthorizedAction::TokenlessHelloNoop,
_ => UnauthorizedAction::Ignore,
}
}
}
#[derive(Debug, PartialEq, Eq)]
enum GateOutcome {
Handled,
Close,
Dispatch,
}
async fn apply_auth_gate(
state: &web::Data<AppState>,
frame: &ClientFrame,
authorized: &mut bool,
) -> GateOutcome {
if *authorized {
return GateOutcome::Dispatch;
}
match UnauthorizedAction::classify(frame) {
UnauthorizedAction::VerifyHello => {
let ClientFrame::Hello {
device_id: Some(device_id),
token: Some(token),
} = frame
else {
unreachable!("VerifyHello implies a credentialed Hello");
};
let config = state.config.read().await.clone();
if crate::handlers::settings::verify_device_token(&config, device_id, token) {
*authorized = true;
tracing::debug!("ws_v2: hello verified for device {device_id}; authorized");
GateOutcome::Handled
} else {
tracing::warn!(
"ws_v2: hello rejected — invalid device credential for {device_id}; closing"
);
GateOutcome::Close
}
}
UnauthorizedAction::TokenlessHelloNoop => {
tracing::debug!("ws_v2: token-less hello while unauthorized — not granting access");
GateOutcome::Handled
}
UnauthorizedAction::Ignore => {
tracing::debug!("ws_v2: ignoring frame on unauthorized connection");
GateOutcome::Handled
}
}
}
#[derive(Debug, Default, Deserialize)]
pub struct StreamQuery {
#[serde(default)]
pub batch_ms: u64,
}
fn negotiate_encoding(req: &HttpRequest) -> (Encoding, Option<&'static str>) {
let offered = req
.headers()
.get(actix_web::http::header::SEC_WEBSOCKET_PROTOCOL)
.and_then(|v| v.to_str().ok())
.unwrap_or("");
let mut has_msgpack = false;
let mut has_json = false;
for tok in offered.split(',') {
match tok.trim() {
SUBPROTOCOL_MSGPACK => has_msgpack = true,
SUBPROTOCOL_JSON => has_json = true,
_ => {}
}
}
if has_msgpack {
(Encoding::Msgpack, Some(SUBPROTOCOL_MSGPACK))
} else if has_json {
(Encoding::Json, Some(SUBPROTOCOL_JSON))
} else {
(Encoding::Json, None)
}
}
pub async fn handler(
state: web::Data<AppState>,
query: web::Query<StreamQuery>,
req: HttpRequest,
body: web::Payload,
) -> actix_web::Result<impl Responder> {
let batch_ms = query.batch_ms.min(MAX_BATCH_MS);
let (encoding, selected_subprotocol) = negotiate_encoding(&req);
let pre_authorized = {
let config = state.config.read().await.clone();
crate::handlers::settings::request_is_authorized(&req, &config)
};
let (mut response, session, msg_stream) = actix_ws::handle(&req, body)?;
if let Some(proto) = selected_subprotocol {
response.headers_mut().insert(
actix_web::http::header::SEC_WEBSOCKET_PROTOCOL,
actix_web::http::header::HeaderValue::from_static(proto),
);
}
actix_web::rt::spawn(drive(
state,
session,
msg_stream,
batch_ms,
pre_authorized,
encoding,
));
Ok(response)
}
async fn drive(
state: web::Data<AppState>,
mut session: actix_ws::Session,
mut msg_stream: actix_ws::MessageStream,
batch_ms: u64,
pre_authorized: bool,
encoding: Encoding,
) {
let mut forwarders: HashMap<String, tokio::task::JoinHandle<()>> = HashMap::new();
let mut queues: StreamMap<String, ReceiverStream<OutFrame>> = StreamMap::new();
let mut ping = tokio::time::interval(PING_INTERVAL);
ping.tick().await;
let mut authorized = pre_authorized;
let auth_deadline = tokio::time::sleep(auth_deadline());
tokio::pin!(auth_deadline);
loop {
tokio::select! {
_ = &mut auth_deadline, if !authorized => {
tracing::debug!("ws_v2: closing unauthorized connection after auth deadline");
break;
}
Some((_ch, frame)) = queues.next() => {
let write = match frame {
OutFrame::Text(s) => session.text(s).await,
OutFrame::Binary(b) => session.binary(b).await,
};
if write.is_err() {
break;
}
}
_ = ping.tick() => {
if session.ping(b"").await.is_err() {
break;
}
}
msg = msg_stream.next() => {
match msg {
Some(Ok(Message::Text(text))) if encoding == Encoding::Json => {
let keep_open = handle_client_bytes(
&state, &mut forwarders, &mut queues,
batch_ms, encoding, text.as_bytes(), &mut authorized,
)
.await;
if !keep_open {
break;
}
}
Some(Ok(Message::Binary(bytes))) if encoding == Encoding::Msgpack => {
let keep_open = handle_client_bytes(
&state, &mut forwarders, &mut queues,
batch_ms, encoding, &bytes, &mut authorized,
)
.await;
if !keep_open {
break;
}
}
Some(Ok(Message::Ping(bytes))) => {
if session.pong(&bytes).await.is_err() {
break;
}
}
Some(Ok(Message::Text(_)))
| Some(Ok(Message::Binary(_)))
| Some(Ok(Message::Pong(_)))
| Some(Ok(Message::Continuation(_)))
| Some(Ok(Message::Nop)) => {}
Some(Ok(Message::Close(_))) | None => break,
Some(Err(e)) => {
tracing::debug!("ws_v2: message stream error: {e}");
break;
}
}
}
else => break,
}
}
for (_ch, handle) in forwarders.drain() {
handle.abort();
}
queues.clear();
let _ = session.close(None).await;
}
async fn handle_client_bytes(
state: &web::Data<AppState>,
forwarders: &mut HashMap<String, tokio::task::JoinHandle<()>>,
queues: &mut StreamMap<String, ReceiverStream<OutFrame>>,
batch_ms: u64,
encoding: Encoding,
bytes: &[u8],
authorized: &mut bool,
) -> bool {
let frame: ClientFrame = match decode_client_frame(encoding, bytes) {
Ok(f) => f,
Err(e) => {
tracing::debug!("ws_v2: ignoring malformed client frame: {e}");
return true;
}
};
handle_client_frame(
state, forwarders, queues, batch_ms, encoding, frame, authorized,
)
.await
}
async fn handle_client_frame(
state: &web::Data<AppState>,
forwarders: &mut HashMap<String, tokio::task::JoinHandle<()>>,
queues: &mut StreamMap<String, ReceiverStream<OutFrame>>,
batch_ms: u64,
encoding: Encoding,
frame: ClientFrame,
authorized: &mut bool,
) -> bool {
match apply_auth_gate(state, &frame, authorized).await {
GateOutcome::Handled => return true,
GateOutcome::Close => return false,
GateOutcome::Dispatch => {}
}
match frame {
ClientFrame::Hello { device_id, token } => {
match (device_id, token) {
(Some(device_id), Some(token)) => {
let config = state.config.read().await.clone();
if crate::handlers::settings::verify_device_token(&config, &device_id, &token) {
tracing::debug!("ws_v2: hello verified for device {device_id}");
} else {
tracing::warn!(
"ws_v2: hello rejected — invalid device credential for {device_id}; closing"
);
return false;
}
}
_ => {
tracing::debug!("ws_v2: token-less hello on authorized connection (no-op)");
}
}
}
ClientFrame::Subscribe { ch, since } => {
subscribe(state, forwarders, queues, batch_ms, encoding, &ch, since).await;
}
ClientFrame::Unsubscribe { ch } => {
if let Some(handle) = forwarders.remove(&ch) {
handle.abort();
queues.remove(&ch);
tracing::debug!("ws_v2: unsubscribed {ch}");
}
}
ClientFrame::Stop { session_id } => {
let cancelled = cancel_session(state, &session_id).await;
tracing::debug!("ws_v2: stop {session_id} -> cancelled={cancelled}");
}
ClientFrame::Unknown => {
tracing::debug!("ws_v2: ignoring unknown client frame type");
}
}
true
}
async fn subscribe(
state: &web::Data<AppState>,
forwarders: &mut HashMap<String, tokio::task::JoinHandle<()>>,
queues: &mut StreamMap<String, ReceiverStream<OutFrame>>,
batch_ms: u64,
encoding: Encoding,
ch: &str,
since: Option<u64>,
) {
let Some(channel) = Channel::parse(ch) else {
tracing::debug!("ws_v2: ignoring subscribe to unknown channel {ch}");
return;
};
if let Some(old) = forwarders.remove(ch) {
old.abort();
}
queues.remove(ch);
let (out_tx, out_rx) = mpsc::channel::<OutFrame>(OUTBOUND_BUFFER);
let out_tx: OutboundTx = out_tx;
let handle = match channel {
Channel::Feed => {
let receiver = state.account_sink.subscribe();
let latest_at_start = state.account_sink.latest_seq();
let events_dir = state.account_sink.events_dir().to_path_buf();
let since = since.unwrap_or(0);
spawn_feed_forwarder(
out_tx,
encoding,
receiver,
events_dir,
since,
latest_at_start,
)
}
Channel::Agent(sid) => {
if state.session_store.get_index_entry(&sid).await.is_none() {
tracing::debug!("ws_v2: ignoring subscribe to unknown session {sid}");
return;
}
let sender = state.get_session_event_sender(&sid).await;
let receiver = sender.subscribe();
state.ensure_notification_relay(&sid, sender.clone());
let runner_snapshot = {
let runners = state.agent_runners.read().await;
runners.get(&sid).cloned()
};
let budget_event_to_replay = runner_snapshot
.as_ref()
.and_then(|runner| runner.last_budget_event.clone());
let critical_events_to_replay: Vec<_> = runner_snapshot
.as_ref()
.map(|runner| runner.last_critical_events.clone())
.unwrap_or_default();
spawn_agent_forwarder(
state.clone(),
sid.clone(),
out_tx,
encoding,
ch.to_string(),
receiver,
budget_event_to_replay,
critical_events_to_replay,
batch_ms,
)
}
};
forwarders.insert(ch.to_string(), handle);
queues.insert(ch.to_string(), ReceiverStream::new(out_rx));
tracing::debug!("ws_v2: subscribed {ch} (since={since:?})");
}
#[cfg(test)]
mod tests {
use super::*;
use crate::app_state::AppState;
use bamboo_config::{AccessControlConfig, DeviceCredential};
use tempfile::tempdir;
fn hello(device_id: Option<&str>, token: Option<&str>) -> ClientFrame {
ClientFrame::Hello {
device_id: device_id.map(str::to_string),
token: token.map(str::to_string),
}
}
fn negotiate_for(header: Option<&str>) -> (Encoding, Option<&'static str>) {
let mut req = actix_web::test::TestRequest::default();
if let Some(h) = header {
req = req.insert_header((actix_web::http::header::SEC_WEBSOCKET_PROTOCOL, h));
}
negotiate_encoding(&req.to_http_request())
}
#[test]
fn negotiate_encoding_branches() {
assert_eq!(negotiate_for(None), (Encoding::Json, None));
assert_eq!(negotiate_for(Some("")), (Encoding::Json, None));
assert_eq!(
negotiate_for(Some("bamboo.v2")),
(Encoding::Json, Some(SUBPROTOCOL_JSON))
);
assert_eq!(
negotiate_for(Some("bamboo.v2.msgpack")),
(Encoding::Msgpack, Some(SUBPROTOCOL_MSGPACK))
);
assert_eq!(
negotiate_for(Some("bamboo.v2.msgpack, bamboo.v2")),
(Encoding::Msgpack, Some(SUBPROTOCOL_MSGPACK))
);
assert_eq!(
negotiate_for(Some(" bamboo.v2 , bamboo.v2.msgpack ")),
(Encoding::Msgpack, Some(SUBPROTOCOL_MSGPACK))
);
assert_eq!(
negotiate_for(Some("some.other.proto")),
(Encoding::Json, None)
);
}
#[test]
fn classify_unauthorized_frame_actions() {
assert_eq!(
UnauthorizedAction::classify(&hello(Some("d"), Some("t"))),
UnauthorizedAction::VerifyHello
);
assert_eq!(
UnauthorizedAction::classify(&hello(None, None)),
UnauthorizedAction::TokenlessHelloNoop
);
assert_eq!(
UnauthorizedAction::classify(&hello(Some("d"), None)),
UnauthorizedAction::TokenlessHelloNoop
);
assert_eq!(
UnauthorizedAction::classify(&hello(None, Some("t"))),
UnauthorizedAction::TokenlessHelloNoop
);
assert_eq!(
UnauthorizedAction::classify(&ClientFrame::Subscribe {
ch: "feed".into(),
since: None
}),
UnauthorizedAction::Ignore
);
assert_eq!(
UnauthorizedAction::classify(&ClientFrame::Unsubscribe { ch: "feed".into() }),
UnauthorizedAction::Ignore
);
assert_eq!(
UnauthorizedAction::classify(&ClientFrame::Stop {
session_id: "s".into()
}),
UnauthorizedAction::Ignore
);
assert_eq!(
UnauthorizedAction::classify(&ClientFrame::Unknown),
UnauthorizedAction::Ignore
);
}
async fn app_state_with_device() -> (web::Data<AppState>, DeviceCredential, String) {
let dir = tempdir().unwrap();
let state = web::Data::new(AppState::new(dir.path().to_path_buf()).await.unwrap());
let (cred, token) = crate::handlers::settings::issue_device_token("test-device");
{
let mut config = state.config.write().await;
config.access_control = Some(AccessControlConfig {
password_enabled: false,
password_hash: None,
password_salt: None,
updated_at: None,
devices: vec![cred.clone()],
});
}
(state, cred, token)
}
#[actix_web::test]
async fn subscribe_while_unauthorized_is_ignored_and_stays_unauthorized() {
let (state, _cred, _token) = app_state_with_device().await;
let mut authorized = false;
let frame = ClientFrame::Subscribe {
ch: "feed".into(),
since: None,
};
let outcome = apply_auth_gate(&state, &frame, &mut authorized).await;
assert_eq!(outcome, GateOutcome::Handled);
assert!(!authorized, "a subscribe must never authorize a connection");
}
#[actix_web::test]
async fn stop_while_unauthorized_is_ignored() {
let (state, _cred, _token) = app_state_with_device().await;
let mut authorized = false;
let frame = ClientFrame::Stop {
session_id: "sess".into(),
};
let outcome = apply_auth_gate(&state, &frame, &mut authorized).await;
assert_eq!(outcome, GateOutcome::Handled);
assert!(!authorized);
}
#[actix_web::test]
async fn valid_hello_authorizes() {
let (state, cred, token) = app_state_with_device().await;
let mut authorized = false;
let frame = hello(Some(&cred.device_id), Some(&token));
let outcome = apply_auth_gate(&state, &frame, &mut authorized).await;
assert_eq!(outcome, GateOutcome::Handled);
assert!(authorized, "a valid hello must authorize the connection");
}
#[actix_web::test]
async fn invalid_hello_closes_and_does_not_authorize() {
let (state, cred, _token) = app_state_with_device().await;
let mut authorized = false;
let frame = hello(Some(&cred.device_id), Some("bd1_wrongwrongwrong"));
let outcome = apply_auth_gate(&state, &frame, &mut authorized).await;
assert_eq!(outcome, GateOutcome::Close);
assert!(!authorized, "an invalid hello must never authorize");
}
#[actix_web::test]
async fn tokenless_hello_does_not_authorize_unauthorized_connection() {
let (state, _cred, _token) = app_state_with_device().await;
let mut authorized = false;
let frame = hello(None, None);
let outcome = apply_auth_gate(&state, &frame, &mut authorized).await;
assert_eq!(outcome, GateOutcome::Handled);
assert!(
!authorized,
"a token-less hello must NEVER authorize a non-pre-authorized connection"
);
}
#[actix_web::test]
async fn pre_authorized_connection_dispatches_subscribe() {
let (state, _cred, _token) = app_state_with_device().await;
let mut authorized = true;
let frame = ClientFrame::Subscribe {
ch: "feed".into(),
since: None,
};
let outcome = apply_auth_gate(&state, &frame, &mut authorized).await;
assert_eq!(outcome, GateOutcome::Dispatch);
assert!(authorized);
}
}