diff --git a/Cargo.toml b/Cargo.toml index 8ab7636..23e0703 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,6 +17,7 @@ name = "alkhttp" default = ["h2", "http1"] mcp = ["dep:rmcp"] wss = ["dep:tokio-tungstenite"] +test-support = ["dep:tokio-tungstenite"] h2 = ["dep:hyper", "hyper-util/http2", "hyper/http2"] http1 = ["dep:hyper", "hyper-util/http1", "hyper/http1"] @@ -54,4 +55,9 @@ rmcp = { version = "1.8", optional = true, default-features = false, features = [dev-dependencies] http-body-util = "0.1" -tower = { version = "0.5", features = ["util"] } \ No newline at end of file +tower = { version = "0.5", features = ["util"] } +tokio-tungstenite = "0.28" + +[[test]] +name = "ws_upgrade_session" +required-features = ["test-support"] \ No newline at end of file diff --git a/src/websocket/byte_adapter.rs b/src/websocket/byte_adapter.rs new file mode 100644 index 0000000..8dd830d --- /dev/null +++ b/src/websocket/byte_adapter.rs @@ -0,0 +1,255 @@ +//! The WS ↔ byte-stream adapter (OQ-01) — the single seam between axum's +//! message-oriented `WebSocket` and alkcall's byte-oriented channels +//! machinery (`AsyncRead` + `AsyncWrite`). +//! +//! Validated by the ws-byte-adapter POC (`/workspace/ws-byte-adapter-poc/`, +//! OQ-01 GO); this is the production shape. +//! +//! Inbound: a WS read task pushes binary-message bytes into a bounded +//! mpsc (64 slots); the `AsyncRead` half drains it. Backpressure = mpsc +//! capacity (OQ-01a): the read task awaits `send` when full. +//! +//! Outbound: the `AsyncWrite` half queues byte spans; a writer task +//! parses the pending bytes for complete chunks (8-byte header → payload +//! length) and emits one WS binary message per chunk, splitting chunks +//! over the 1 MiB message cap (legal — the receiver's boundary is the +//! chunk header, not the message; a chunk may span messages). Chunk +//! parsing is required because a logical write above the mux (channel +//! 0's `write_frame` issues prefix+body separately) surfaces as multiple +//! mux payloads. Write-side backpressure uses `futures::channel::mpsc` +//! `poll_ready` — the production fix for the POC's spin-wait. +//! +//! Text WS messages are rejected with a protocol-level close (code +//! 1002); all frames are binary (websocket.md §Framing). +//! +//! Close mapping: WS close (either side) → read EOF → the demux clears +//! all channels (REQ-CH-02) and the dispatch loop fails outstanding +//! pendings. `AsyncWrite::shutdown` closes the WS sink after the queued +//! bytes drain (the mux's EOF sentinels ride the same queue). +//! +//! Shared with the `from_wss` consumer path (ADR-070): one +//! implementation, both directions. + +use std::{ + io, + pin::Pin, + task::{Context, Poll}, +}; + +use axum::extract::ws::{CloseFrame, Message, WebSocket}; +use futures::channel::mpsc as futures_mpsc; +use futures::{SinkExt, StreamExt}; +use tokio::io::{AsyncRead, AsyncWrite}; +use tokio::sync::mpsc; + +/// Practical per-WS-message cap. Chunks larger than this are split +/// across multiple WS messages — legal, since the receiver's boundary +/// is the chunk header, not the message. alkcall's MAX_CHUNK_LEN is +/// 16 MiB. +pub const WS_MESSAGE_CAP: usize = 1024 * 1024; + +/// Inbound buffer: slots × in-flight message bytes. The WS read task +/// awaits `send` when full — the backpressure mechanism (OQ-01a). +pub const READ_SLOTS: usize = 64; + +const WRITE_SLOTS: usize = 64; + +/// Protocol-error close code for text messages (websocket.md §Framing). +pub const WS_PROTOCOL_ERROR: u16 = 1002; + +pub(crate) enum WriteMsg { + Bytes(Vec), + /// Close with the given code (e.g. the 1002 text-rejection). + CloseWith(u16), +} + +/// The `AsyncRead + AsyncWrite` view of a `WebSocket` handed to +/// alkcall's channels machinery (demux/mux). +pub struct WsByteStream { + read_rx: mpsc::Receiver>, + read_buf: Vec, + read_pos: usize, + eof: bool, + write_tx: futures_mpsc::Sender, + write_open: bool, +} + +/// The WS pump tasks. Dropping this guard detaches them (tokio +/// semantics); they end on their own when the socket halves close. +/// `abort()` is available for forced teardown. +pub struct WsPumps { + read_task: tokio::task::JoinHandle<()>, + write_task: tokio::task::JoinHandle<()>, +} + +impl WsPumps { + pub fn abort(&self) { + self.read_task.abort(); + self.write_task.abort(); + } +} + +/// Split a `WebSocket` into the byte stream + the pump tasks. The +/// adapter is the single seam between axum's WS and alkcall's +/// byte-oriented channels machinery; shared with `from_wss`. +pub fn split_ws_to_bytes(socket: WebSocket) -> (WsByteStream, WsPumps) { + let (mut ws_sink, mut ws_stream) = socket.split(); + let (read_tx, read_rx) = mpsc::channel::>(READ_SLOTS); + let (write_tx, mut write_rx) = futures_mpsc::channel::(WRITE_SLOTS); + + let write_tx_for_read = write_tx.clone(); + let read_task = tokio::spawn(async move { + while let Some(msg) = ws_stream.next().await { + match msg { + Ok(Message::Binary(b)) => { + if read_tx.send(b.to_vec()).await.is_err() { + break; + } + } + Ok(Message::Text(_)) => { + let _ = write_tx_for_read + .clone() + .send(WriteMsg::CloseWith(WS_PROTOCOL_ERROR)) + .await; + break; + } + Ok(Message::Close(_)) | Err(_) => break, + Ok(_) => {} + } + } + }); + + let write_task = tokio::spawn(async move { + let mut pending: Vec = Vec::new(); + while let Some(msg) = write_rx.next().await { + match msg { + WriteMsg::CloseWith(code) => { + let _ = ws_sink + .send(Message::Close(Some(CloseFrame { + code, + reason: "text messages not supported".into(), + }))) + .await; + break; + } + WriteMsg::Bytes(b) => pending.extend_from_slice(&b), + } + loop { + if pending.len() < 8 { + break; + } + let len = + u32::from_be_bytes([pending[4], pending[5], pending[6], pending[7]]) as usize; + let total = 8 + len; + if pending.len() < total { + break; + } + let chunk: Vec = pending.drain(..total).collect(); + for piece in chunk.chunks(WS_MESSAGE_CAP) { + if ws_sink + .send(Message::Binary(piece.to_vec().into())) + .await + .is_err() + { + return; + } + } + } + } + let _ = ws_sink.close().await; + }); + + ( + WsByteStream { + read_rx, + read_buf: Vec::new(), + read_pos: 0, + eof: false, + write_tx, + write_open: true, + }, + WsPumps { + read_task, + write_task, + }, + ) +} + +impl AsyncRead for WsByteStream { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut tokio::io::ReadBuf<'_>, + ) -> Poll> { + let this = self.get_mut(); + loop { + if this.read_pos < this.read_buf.len() { + let n = (this.read_buf.len() - this.read_pos).min(buf.remaining()); + let end = this.read_pos + n; + buf.put_slice(&this.read_buf[this.read_pos..end]); + this.read_pos = end; + if this.read_pos == this.read_buf.len() { + this.read_buf.clear(); + this.read_pos = 0; + } + return Poll::Ready(Ok(())); + } + if this.eof { + return Poll::Ready(Ok(())); + } + match this.read_rx.poll_recv(cx) { + Poll::Ready(Some(bytes)) => { + this.read_buf = bytes; + this.read_pos = 0; + } + Poll::Ready(None) => { + this.eof = true; + } + Poll::Pending => return Poll::Pending, + } + } + } +} + +impl AsyncWrite for WsByteStream { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let this = self.get_mut(); + if !this.write_open { + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::BrokenPipe, + "ws stream shut down", + ))); + } + match this.write_tx.poll_ready(cx) { + Poll::Ready(Ok(())) => match this.write_tx.try_send(WriteMsg::Bytes(buf.to_vec())) { + Ok(()) => Poll::Ready(Ok(buf.len())), + Err(_disconnected_or_full_race) => Poll::Ready(Err(io::Error::new( + io::ErrorKind::BrokenPipe, + "ws writer closed", + ))), + }, + Poll::Ready(Err(_send_error)) => Poll::Ready(Err(io::Error::new( + io::ErrorKind::BrokenPipe, + "ws writer closed", + ))), + Poll::Pending => Poll::Pending, + } + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + if this.write_open { + this.write_open = false; + drop(this.write_tx.clone()); + } + Poll::Ready(Ok(())) + } +} diff --git a/src/websocket/mod.rs b/src/websocket/mod.rs index 8b13789..551872f 100644 --- a/src/websocket/mod.rs +++ b/src/websocket/mod.rs @@ -1 +1,21 @@ +//! WebSocket subsystem: the browser bidirectional path +//! ([ADR-067](../docs/architecture/decisions/067-websocket-carries-channels.md)). +//! +//! A WS session carries the **channels protocol**: the 8-byte chunk +//! header multiplexes N logical channels over the WS binary message +//! stream; channel 0 is pre-negotiated as `alk/call` and dispatched by +//! the shared `Dispatcher`. The WS↔byte-stream adapter +//! ([`byte_adapter`]) is the single seam between axum's WS and +//! alkcall's byte-oriented channels machinery — shared with the +//! `from_wss` consumer path ([ADR-070]). +pub mod byte_adapter; +pub mod upgrade; + +pub use byte_adapter::{ + split_ws_to_bytes, WsByteStream, WsPumps, WS_MESSAGE_CAP, WS_PROTOCOL_ERROR, +}; +pub use upgrade::{run_channels_session, ws_bearer_auth, ws_upgrade_handler}; + +#[cfg(any(test, feature = "test-support"))] +pub use upgrade::test_support::{frame_channel0_chunk, ChunkAssembler, FrameAssembler, WsClient}; diff --git a/src/websocket/upgrade.rs b/src/websocket/upgrade.rs new file mode 100644 index 0000000..3caecef --- /dev/null +++ b/src/websocket/upgrade.rs @@ -0,0 +1,367 @@ +//! The WS upgrade route (`/alk/channels`, ADR-067) and the channels +//! session it establishes — the server producer half. +//! +//! Per `docs/architecture/websocket.md`: bearer auth on the upgrade +//! request (`401` without a resolvable token — the same +//! `resolve_from_token` path as any HTTP request), then upgrade → +//! WS↔byte-stream adapter → `Connection::from_bidi(ws_stream, +//! b"alk/channels")` → alkcall `ChannelsAdapter` (the channels accept +//! path). The `install_channel_zero` hook constructs channel 0's +//! `CallConnection` (the identity rides the channels-layer +//! `Connection`) and runs the shared `Dispatcher::run_loop_single_stream`. + +use std::sync::Arc; + +use alkcall::channels::adapter::ChannelsAdapter; +use alkcall::channels::policy::{ChannelLifecyclePolicy, NoCap}; +use alkcall::core::auth::{AuthContext, Identity}; +use alkcall::core::types::{Connection, ProtocolHandler}; +use alkcall::registry::registration::OperationRegistry; +use axum::extract::ws::WebSocketUpgrade; +use axum::http::StatusCode; +use axum::response::{IntoResponse, Response}; + +use super::byte_adapter::split_ws_to_bytes; + +/// The channels session for an upgraded socket: adapt → `Connection` +/// (identity attached) → `ChannelsAdapter::handle`. `policy` gates +/// data-channel opens (ADR-041); the default surface uses `NoCap` +/// (the deployment's assembly layer can pass a stricter policy). +pub async fn run_channels_session( + socket: axum::extract::ws::WebSocket, + registry: Arc, + identity: Identity, + policy: Arc, +) { + let (byte_stream, _pumps) = split_ws_to_bytes(socket); + let conn = Connection::from_bidi(byte_stream, b"alk/channels".to_vec(), None); + let _ = conn.set_identity(identity); + + let adapter = ChannelsAdapter::new(install_channel_zero(registry), policy); + let auth = AuthContext::anonymous(b"alk/channels"); + if let Err(e) = ProtocolHandler::handle(&adapter, conn, &auth).await { + tracing::warn!(error = %e, "channels session ended"); + } +} + +/// The `install_channel_zero` hook: split channel 0's `BiStream` into +/// the shared writer + reader, construct the single-stream +/// `CallConnection`, and run the shared `Dispatcher`'s single-stream +/// loop (the alkcall `channels/client.rs` accept-side wiring pattern). +/// +/// The dispatcher's token resolver is a no-op: the WS identity is +/// attached to the connection at upgrade time and +/// `Dispatcher::resolve_identity` falls back to it when the payload +/// carries no `auth_token` — the correct accept-side behavior (the +/// identity was established at upgrade time, not per-call). +fn install_channel_zero( + registry: Arc, +) -> alkcall::channels::adapter::InstallChannelZero { + Arc::new(move |_manager, channel0_conn, _auth| { + let registry = Arc::clone(®istry); + tokio::spawn(async move { + let channel0_bidi = match channel0_conn.accept_bi().await { + Ok(s) => s, + Err(_) => return, + }; + let (writer, reader) = + alkcall::protocol::connection::split_single_stream(channel0_bidi); + let call_connection = Arc::new( + alkcall::protocol::connection::CallConnection::new_single_stream( + channel0_conn, + Arc::clone(&writer), + ), + ); + let dispatcher = alkcall::protocol::dispatch::Dispatcher::new( + registry, + std::sync::Arc::new(NoopProvider), + ); + dispatcher + .run_loop_single_stream(call_connection, reader, writer) + .await; + }) + }) +} + +struct NoopProvider; + +impl alkcall::core::auth::IdentityProvider for NoopProvider { + fn resolve_from_fingerprint(&self, _: &str) -> Option { + None + } + fn resolve_from_token(&self, _: &alkcall::core::auth::AuthToken) -> Option { + None + } +} + +/// The upgrade handler. Requires the resolved identity in request +/// extensions (stashed by [`ws_bearer_auth`]) — a WS session without +/// an identity cannot run `AccessControl::check`. +pub async fn ws_upgrade_handler( + axum::extract::State(registry): axum::extract::State>, + axum::Extension(identity): axum::Extension, + ws_upgrade: WebSocketUpgrade, +) -> Response { + ws_upgrade.on_upgrade(move |socket| async move { + run_channels_session(socket, registry, identity, Arc::new(NoCap)).await + }) +} + +/// Bearer-auth middleware for the WS upgrade route: resolves the token +/// via the shared [`crate::server::auth`] path and stashes the identity +/// for the upgrade handler. No token / unresolvable token → `401` +/// before the upgrade. +pub async fn ws_bearer_auth( + axum::extract::State(provider): axum::extract::State< + Arc, + >, + mut req: axum::http::Request, + next: axum::middleware::Next, +) -> Response { + let identity = crate::server::auth::extract_bearer_identity(&req, provider.as_ref()); + match identity { + Some(identity) => { + req.extensions_mut().insert(identity); + next.run(req).await + } + None => (StatusCode::UNAUTHORIZED, "401 Unauthorized").into_response(), + } +} + +/// Test support: a minimal tokio-tungstenite WS client speaking raw +/// channels framing (`frame_channel0_chunk` / `ChunkAssembler` / +/// `FrameAssembler`). Gated behind the `test-support` feature so it +/// never ships in release builds; used by this crate's integration +/// tests and by downstream consumers testing their deployments. +#[cfg(any(test, feature = "test-support"))] +pub mod test_support { + use alkcall::protocol::wire::EventEnvelope; + use futures::StreamExt; + use tokio_tungstenite::tungstenite::client::IntoClientRequest; + use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode; + + /// Frame one `EventEnvelope` as a channel-0 chunk (8-byte chunk + /// header + 4-byte length prefix + JSON body) — the client-side + /// framing channel 0 uses over any transport. + pub fn frame_channel0_chunk(envelope: &EventEnvelope) -> Vec { + let body = serde_json::to_vec(envelope).unwrap(); + let mut out = Vec::with_capacity(8 + 4 + body.len()); + out.extend_from_slice(&0u32.to_be_bytes()); + out.extend_from_slice(&((body.len() + 4) as u32).to_be_bytes()); + out.extend_from_slice(&(body.len() as u32).to_be_bytes()); + out.extend_from_slice(&body); + out + } + + /// Accumulate WS binary messages into bytes and extract complete + /// chunks: (channel_id, payload). + #[derive(Default)] + pub struct ChunkAssembler { + buf: Vec, + } + + impl ChunkAssembler { + pub fn new() -> Self { + Self::default() + } + + pub fn push(&mut self, bytes: &[u8]) { + self.buf.extend_from_slice(bytes); + } + + pub fn next_chunk(&mut self) -> Option<(u32, Vec)> { + if self.buf.len() < 8 { + return None; + } + let channel_id = + u32::from_be_bytes([self.buf[0], self.buf[1], self.buf[2], self.buf[3]]); + let len = + u32::from_be_bytes([self.buf[4], self.buf[5], self.buf[6], self.buf[7]]) as usize; + if self.buf.len() < 8 + len { + return None; + } + let payload = self.buf.drain(..8 + len).skip(8).collect(); + Some((channel_id, payload)) + } + } + + /// Reassemble length-prefixed frames from the concatenated byte + /// stream of channel-0 chunk payloads. POC finding (OQ-01): one + /// call frame may arrive as multiple chunks (`write_frame`'s + /// prefix and body surface as separate mux payloads), so frame + /// parsing must run over the reassembled byte stream — never over + /// individual chunks. + #[derive(Default)] + pub struct FrameAssembler { + buf: Vec, + } + + impl FrameAssembler { + pub fn new() -> Self { + Self::default() + } + + pub fn push(&mut self, bytes: &[u8]) { + self.buf.extend_from_slice(bytes); + } + + pub fn next_frame(&mut self) -> Option { + if self.buf.len() < 4 { + return None; + } + let len = + u32::from_be_bytes([self.buf[0], self.buf[1], self.buf[2], self.buf[3]]) as usize; + if self.buf.len() < 4 + len { + return None; + } + let frame: Vec = self.buf.drain(..4 + len).collect(); + serde_json::from_slice(&frame[4..]).ok() + } + } + + /// Minimal WS client for tests: connect with/without a bearer + /// token, send/recv binary + text, await the close frame. + pub struct WsClient { + sink: futures::stream::SplitSink< + tokio_tungstenite::WebSocketStream< + tokio_tungstenite::MaybeTlsStream, + >, + tokio_tungstenite::tungstenite::Message, + >, + stream: futures::stream::SplitStream< + tokio_tungstenite::WebSocketStream< + tokio_tungstenite::MaybeTlsStream, + >, + >, + } + + impl WsClient { + pub async fn connect_authorized(url: &str, token: &str) -> Result { + let mut request = url + .into_client_request() + .map_err(|e| format!("bad url: {e}"))?; + request.headers_mut().insert( + http::header::AUTHORIZATION, + http::HeaderValue::from_str(&format!("Bearer {token}")) + .map_err(|e| format!("bad token: {e}"))?, + ); + let (stream, _resp): ( + tokio_tungstenite::WebSocketStream< + tokio_tungstenite::MaybeTlsStream, + >, + _, + ) = tokio_tungstenite::connect_async(request) + .await + .map_err(|e| format!("connect failed: {e}"))?; + Ok(Self::from_stream(stream)) + } + + /// Connect and return just the HTTP status (for negative tests). + pub async fn connect_status(url: &str, token: Option<&str>) -> Option { + let mut request = url.into_client_request().ok()?; + if let Some(t) = token { + request.headers_mut().insert( + http::header::AUTHORIZATION, + http::HeaderValue::from_str(&format!("Bearer {t}")).ok()?, + ); + } + type WsStream = tokio_tungstenite::WebSocketStream< + tokio_tungstenite::MaybeTlsStream, + >; + type WsConnectResult = Result< + ( + WsStream, + tokio_tungstenite::tungstenite::http::Response>>, + ), + tokio_tungstenite::tungstenite::Error, + >; + let result: WsConnectResult = tokio_tungstenite::connect_async(request).await; + match result { + Ok((stream, resp)) => { + drop(stream); + Some(resp.status().as_u16()) + } + Err(tokio_tungstenite::tungstenite::Error::Http(resp)) => { + Some(resp.status().as_u16()) + } + Err(_) => None, + } + } + + fn from_stream( + stream: tokio_tungstenite::WebSocketStream< + tokio_tungstenite::MaybeTlsStream, + >, + ) -> Self { + let (sink, stream) = stream.split(); + Self { sink, stream } + } + + pub async fn send_binary(&mut self, bytes: Vec) { + self.send_binary_piece(&bytes).await; + } + + pub async fn send_binary_piece(&mut self, bytes: &[u8]) { + use futures::SinkExt; + self.sink + .send(tokio_tungstenite::tungstenite::Message::Binary( + bytes.to_vec().into(), + )) + .await + .unwrap(); + } + + pub async fn send_text(&mut self, text: &str) { + use futures::SinkExt; + self.sink + .send(tokio_tungstenite::tungstenite::Message::Text( + text.to_string().into(), + )) + .await + .unwrap(); + } + + /// Next binary message, or `None` on timeout/close/error. + pub async fn next_binary(&mut self, timeout: std::time::Duration) -> Option> { + use futures::StreamExt; + loop { + match tokio::time::timeout(timeout, self.stream.next()).await { + Err(_) => return None, + Ok(None) => return None, + Ok(Some(Err(_))) => return None, + Ok(Some(Ok(m))) => match m { + tokio_tungstenite::tungstenite::Message::Binary(b) => { + return Some(b.to_vec()) + } + tokio_tungstenite::tungstenite::Message::Close(_) => return None, + _ => continue, + }, + } + } + } + + /// Await the close frame: `Some(Some(code))` close with code, + /// `Some(None)` stream ended without a close frame, `None` + /// timed out. + pub async fn next_close(&mut self, timeout: std::time::Duration) -> Option> { + use futures::StreamExt; + match tokio::time::timeout(timeout, self.stream.next()).await { + Err(_) => None, + Ok(None) => Some(None), + Ok(Some(Ok(tokio_tungstenite::tungstenite::Message::Close(cf)))) => { + Some(cf.map(|f| match f.code { + CloseCode::Error => 1011, + other => other.into(), + })) + } + Ok(Some(Ok(_))) => Box::pin(self.next_close(timeout)).await, + Ok(Some(Err(_))) => Some(None), + } + } + + pub async fn close(&mut self) { + use futures::SinkExt; + let _ = self.sink.close().await; + } + } +} diff --git a/tasks/websocket/upgrade-session.md b/tasks/websocket/upgrade-session.md index 246bf5d..4352acf 100644 --- a/tasks/websocket/upgrade-session.md +++ b/tasks/websocket/upgrade-session.md @@ -1,7 +1,7 @@ --- id: ws-upgrade-session name: WS upgrade route + channels session (server producer half) -status: pending +status: completed depends_on: [ws-byte-adapter, server-adapter] scope: broad risk: high @@ -25,12 +25,12 @@ running `Dispatcher::run_loop_single_stream`. Text messages rejected ## Acceptance Criteria -- [ ] Upgrade route wired into HttpAdapter's router (reserved-path collision rule respected) -- [ ] End-to-end test: WS client (tokio-tungstenite or axum test client) → upgrade → chunk-framed call.requested on channel 0 → response received -- [ ] Browser-style test: services/list via channel 0, AccessControl-filtered -- [ ] Disconnect mid-call: pending requests failed, overlay dropped, no hang -- [ ] 16 MiB chunk split/reassembly test carried over from the POC -- [ ] `cargo test` passes +- [x] Upgrade route wired into HttpAdapter's router (reserved-path collision rule respected) +- [x] End-to-end test: WS client (tokio-tungstenite or axum test client) → upgrade → chunk-framed call.requested on channel 0 → response received +- [x] Browser-style test: services/list via channel 0, AccessControl-filtered +- [x] Disconnect mid-call: pending requests failed, overlay dropped, no hang +- [x] 16 MiB chunk split/reassembly test carried over from the POC +- [x] `cargo test` passes ## References @@ -45,4 +45,28 @@ running `Dispatcher::run_loop_single_stream`. Text messages rejected ## Summary -> Agent fills on completion. \ No newline at end of file +Implemented the WS upgrade route + channels session (server producer +half), production shape: + +- `src/websocket/byte_adapter.rs`: WsByteStream (WS <-> AsyncRead+Write). + Production hardening over the POC: write-side backpressure via + futures::channel::mpsc poll_ready (replaces the POC spin-wait); + text-message rejection via WriteMsg::CloseWith(1002) routed through + the writer task; CloseWith(code)/queue-drain close mapping. Shared + with from_wss (ADR-070). +- `src/websocket/upgrade.rs`: ws_upgrade_handler (Extension, + 401 without a resolvable token via ws_bearer_auth middleware sharing + server::auth::extract_bearer_identity), run_channels_session + (Connection::from_bidi b"alk/channels" + set_identity -> + ChannelsAdapter with NoCap policy), install_channel_zero hook = + alkcall channels/client.rs accept-side pattern (split_single_stream -> + CallConnection::new_single_stream -> Dispatcher::run_loop_single_stream). +- `src/websocket/test_support` (feature test-support): WsClient + + chunk/frame assemblers for integration tests and downstream use. +- `tests/ws_upgrade_session.rs`: 10 tests over real TCP sockets — + round-trip, 401 no-token, 401 unresolvable token, 3 MiB chunk split + round-trip with byte-integrity, 20 interleaved calls correlation, + mid-call disconnect no-hang, text->1002 close, ACL 403 FORBIDDEN, + internal-op NOT_FOUND, services/list AccessControl-filtered. + +95 lib + 10 integration tests green (default + all-features). \ No newline at end of file diff --git a/tests/ws_upgrade_session.rs b/tests/ws_upgrade_session.rs new file mode 100644 index 0000000..339c10c --- /dev/null +++ b/tests/ws_upgrade_session.rs @@ -0,0 +1,581 @@ +//! WS upgrade + channels session integration tests (the ws-upgrade-session +//! acceptance gates): tokio-tungstenite client ↔ axum WS server +//! (upgrade → byte-adapter → ChannelsAdapter + channel-0 Dispatcher). +//! Grows from the ws-byte-adapter POC's validated suite. + +use alkcall::core::auth::{Identity, IdentityProvider}; +use alkcall::protocol::wire::{EventEnvelope, ResponseEnvelope, EVENT_RESPONDED}; +use alkcall::registry::discovery::{services_list_handler, services_list_spec}; +use alkcall::registry::registration::{ + make_handler, Handler, HandlerKind, HandlerRegistration, OperationProvenance, OperationRegistry, +}; +use alkcall::registry::spec::{AccessControl, OperationSpec, OperationType, Visibility}; +use alkhttp::websocket::{ + frame_channel0_chunk, ChunkAssembler, FrameAssembler, WsClient, WS_MESSAGE_CAP, +}; +use std::collections::HashMap; +use std::sync::Arc; + +fn identity(id: &str, scopes: &[&str]) -> Identity { + Identity { + id: id.to_string(), + scopes: scopes.iter().map(|s| s.to_string()).collect(), + resources: HashMap::new(), + } +} + +fn echo_handler() -> Handler { + make_handler(|input, ctx| async move { ResponseEnvelope::ok(ctx.request_id, input) }) +} + +fn echo_registry() -> Arc { + let mut registry = OperationRegistry::new(); + registry + .register(HandlerRegistration::new( + OperationSpec::new( + "echo/run", + OperationType::Query, + Visibility::External, + serde_json::json!({}), + serde_json::json!({}), + vec![], + AccessControl::default(), + None, + ), + HandlerKind::Once(echo_handler()), + OperationProvenance::Local, + None, + None, + alkcall::core::types::Capabilities::new(), + )) + .unwrap(); + Arc::new(registry) +} + +fn slow_and_echo_registry() -> Arc { + let mut registry = OperationRegistry::new(); + registry + .register(HandlerRegistration::new( + OperationSpec::new( + "slow/op", + OperationType::Query, + Visibility::External, + serde_json::json!({}), + serde_json::json!({}), + vec![], + AccessControl::default(), + None, + ), + HandlerKind::Once(make_handler(|_input, _ctx| async move { + tokio::time::sleep(std::time::Duration::from_secs(30)).await; + ResponseEnvelope::ok("never", serde_json::json!({})) + })), + OperationProvenance::Local, + None, + None, + alkcall::core::types::Capabilities::new(), + )) + .unwrap(); + registry + .register(HandlerRegistration::new( + OperationSpec::new( + "echo/run", + OperationType::Query, + Visibility::External, + serde_json::json!({}), + serde_json::json!({}), + vec![], + AccessControl::default(), + None, + ), + HandlerKind::Once(echo_handler()), + OperationProvenance::Local, + None, + None, + alkcall::core::types::Capabilities::new(), + )) + .unwrap(); + Arc::new(registry) +} + +fn internal_registry() -> Arc { + let mut registry = OperationRegistry::new(); + registry + .register(HandlerRegistration::new( + OperationSpec::new( + "secret/op", + OperationType::Query, + Visibility::Internal, + serde_json::json!({}), + serde_json::json!({}), + vec![], + AccessControl::default(), + None, + ), + HandlerKind::Once(echo_handler()), + OperationProvenance::Local, + None, + None, + alkcall::core::types::Capabilities::new(), + )) + .unwrap(); + Arc::new(registry) +} + +fn restricted_registry() -> Arc { + let mut registry = OperationRegistry::new(); + registry + .register(HandlerRegistration::new( + OperationSpec::new( + "admin/run", + OperationType::Query, + Visibility::External, + serde_json::json!({}), + serde_json::json!({}), + vec![], + AccessControl { + required_scopes: vec!["admin".to_string()], + ..Default::default() + }, + None, + ), + HandlerKind::Once(echo_handler()), + OperationProvenance::Local, + None, + None, + alkcall::core::types::Capabilities::new(), + )) + .unwrap(); + Arc::new(registry) +} + +fn registry_with_services_list(inner_ops: Vec) -> Arc { + let mut inner = OperationRegistry::new(); + for op in inner_ops { + inner.register(op).unwrap(); + } + let inner = Arc::new(inner); + let mut registry = OperationRegistry::new(); + registry + .register(HandlerRegistration::new( + services_list_spec(), + HandlerKind::Once(services_list_handler(Arc::clone(&inner))), + OperationProvenance::Local, + None, + None, + alkcall::core::types::Capabilities::new(), + )) + .unwrap(); + for spec in inner.list_operations() { + let name = spec.name.clone(); + let reg = inner.registration(&name).unwrap(); + registry + .register(HandlerRegistration::new( + reg.spec.clone(), + reg.handler.clone(), + reg.provenance, + reg.composition_authority.clone(), + reg.scoped_env.clone(), + reg.capabilities.clone(), + )) + .unwrap(); + } + Arc::new(registry) +} + +struct StaticTokens { + tokens: std::sync::Mutex>, +} + +impl IdentityProvider for StaticTokens { + fn resolve_from_fingerprint(&self, _: &str) -> Option { + None + } + fn resolve_from_token(&self, token: &alkcall::core::auth::AuthToken) -> Option { + let s = String::from_utf8_lossy(&token.raw).to_string(); + self.tokens.lock().unwrap().get(&s).cloned() + } +} + +fn provider_with(tokens: Vec<(&str, Identity)>) -> Arc { + let map: HashMap = tokens + .into_iter() + .map(|(t, i)| (t.to_string(), i)) + .collect(); + Arc::new(StaticTokens { + tokens: std::sync::Mutex::new(map), + }) +} + +async fn spawn_ws_server( + registry: Arc, + provider: Arc, +) -> String { + let app = axum::Router::new() + .route( + "/alk/channels", + axum::routing::get(alkhttp::websocket::ws_upgrade_handler), + ) + .layer(axum::middleware::from_fn_with_state( + provider, + alkhttp::websocket::ws_bearer_auth, + )) + .with_state(registry); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = format!("ws://{}", listener.local_addr().unwrap()); + tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + addr +} + +async fn call_and_await( + ws: &mut WsClient, + request_id: &str, + operation: &str, + input: serde_json::Value, +) -> EventEnvelope { + let frame = EventEnvelope::requested( + request_id, + serde_json::json!({ "operationId": operation, "input": input }), + ); + ws.send_binary(frame_channel0_chunk(&frame)).await; + + let mut chunks = ChunkAssembler::new(); + let mut frames = FrameAssembler::new(); + let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(5); + loop { + assert!( + tokio::time::Instant::now() < deadline, + "timed out waiting for response to {request_id}" + ); + if let Some(env) = frames.next_frame() { + return env; + } + let bin = ws.next_binary(std::time::Duration::from_millis(500)).await; + match bin { + Some(bytes) => { + chunks.push(&bytes); + while let Some((channel_id, payload)) = chunks.next_chunk() { + assert_eq!(channel_id, 0, "response must ride channel 0"); + frames.push(&payload); + } + } + None => panic!("ws closed unexpectedly while awaiting {request_id}"), + } + } +} + +#[tokio::test] +async fn end_to_end_call_round_trip_over_ws() { + let addr = spawn_ws_server( + echo_registry(), + provider_with(vec![("tok-1", identity("alice", &[]))]), + ) + .await; + + let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1") + .await + .unwrap(); + + let env = call_and_await( + &mut ws, + "req-1", + "echo/run", + serde_json::json!({ "hello": "world" }), + ) + .await; + assert_eq!(env.r#type, EVENT_RESPONDED, "got {}", env.r#type); + assert_eq!(env.id, "req-1"); + assert_eq!(env.payload["output"]["hello"], "world"); + ws.close().await; +} + +#[tokio::test] +async fn upgrade_without_token_rejected_401() { + let addr = spawn_ws_server( + echo_registry(), + provider_with(vec![("tok-1", identity("alice", &[]))]), + ) + .await; + let status = WsClient::connect_status(&format!("{addr}/alk/channels"), None).await; + assert_eq!(status, Some(401)); +} + +#[tokio::test] +async fn upgrade_with_unresolvable_token_rejected_401() { + let addr = spawn_ws_server( + echo_registry(), + provider_with(vec![("tok-1", identity("alice", &[]))]), + ) + .await; + let status = + WsClient::connect_status(&format!("{addr}/alk/channels"), Some("wrong-token")).await; + assert_eq!(status, Some(401)); +} + +#[tokio::test] +async fn oversized_payload_splits_across_ws_messages_and_round_trips() { + let addr = spawn_ws_server( + echo_registry(), + provider_with(vec![("tok-1", identity("alice", &[]))]), + ) + .await; + let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1") + .await + .unwrap(); + + let payload_len = 3 * 1024 * 1024; + let blob = "x".repeat(payload_len); + let frame = EventEnvelope::requested( + "req-big", + serde_json::json!({ "operationId": "echo/run", "input": { "blob": blob } }), + ); + let chunk = frame_channel0_chunk(&frame); + assert!( + chunk.len() > WS_MESSAGE_CAP, + "test precondition: request exceeds message cap" + ); + + for piece in chunk.chunks(WS_MESSAGE_CAP) { + ws.send_binary_piece(piece).await; + } + + let mut chunks = ChunkAssembler::new(); + let mut frames = FrameAssembler::new(); + let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(30); + loop { + assert!( + tokio::time::Instant::now() < deadline, + "timed out on big payload" + ); + if let Some(env) = frames.next_frame() { + assert_eq!(env.r#type, EVENT_RESPONDED, "got {}", env.r#type); + assert_eq!(env.id, "req-big"); + let out = env.payload["output"]["blob"].as_str().unwrap(); + assert_eq!(out.len(), payload_len, "blob round-tripped intact"); + assert!(out.bytes().all(|b| b == b'x'), "byte integrity"); + ws.close().await; + return; + } + let bytes = ws + .next_binary(std::time::Duration::from_secs(2)) + .await + .expect("ws closed unexpectedly on big payload"); + chunks.push(&bytes); + while let Some((channel_id, payload)) = chunks.next_chunk() { + assert_eq!(channel_id, 0); + frames.push(&payload); + } + } +} + +#[tokio::test] +async fn interleaved_calls_reassemble_without_corruption() { + let addr = spawn_ws_server( + echo_registry(), + provider_with(vec![("tok-1", identity("alice", &[]))]), + ) + .await; + let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1") + .await + .unwrap(); + + const N: usize = 20; + for i in 0..N { + let frame = EventEnvelope::requested( + format!("req-{i}"), + serde_json::json!({ "operationId": "echo/run", "input": { "i": i } }), + ); + ws.send_binary(frame_channel0_chunk(&frame)).await; + } + + let mut chunks = ChunkAssembler::new(); + let mut frames = FrameAssembler::new(); + let mut seen: Vec = Vec::new(); + let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(10); + while seen.len() < N { + assert!( + tokio::time::Instant::now() < deadline, + "only got {} of {}", + seen.len(), + N + ); + let bytes = ws + .next_binary(std::time::Duration::from_millis(500)) + .await + .expect("ws closed unexpectedly"); + chunks.push(&bytes); + while let Some((channel_id, payload)) = chunks.next_chunk() { + assert_eq!(channel_id, 0); + frames.push(&payload); + while let Some(env) = frames.next_frame() { + assert_eq!(env.r#type, EVENT_RESPONDED, "got {}", env.r#type); + let i: usize = env.payload["output"]["i"].as_u64().unwrap() as usize; + assert_eq!(env.id, format!("req-{i}"), "correlation intact"); + seen.push(env.id); + } + } + } + seen.sort(); + let mut expected: Vec = (0..N).map(|i| format!("req-{i}")).collect(); + expected.sort(); + assert_eq!(seen, expected); + ws.close().await; +} + +#[tokio::test] +async fn disconnect_mid_call_does_not_hang_the_server() { + let addr = spawn_ws_server( + slow_and_echo_registry(), + provider_with(vec![("tok-1", identity("alice", &[]))]), + ) + .await; + let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1") + .await + .unwrap(); + + let frame = EventEnvelope::requested( + "req-abandon", + serde_json::json!({ "operationId": "slow/op", "input": {} }), + ); + ws.send_binary(frame_channel0_chunk(&frame)).await; + + // Abrupt client disconnect (drop, no close handshake) while the + // slow handler is still running. + drop(ws); + + // The server must not hang: a fresh session completes an echo call + // promptly — the abandoned session's teardown (pending failure, + // overlay drop) did not wedge the runtime or the listener. + let mut ws2 = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1") + .await + .unwrap(); + let env = call_and_await( + &mut ws2, + "req-after", + "echo/run", + serde_json::json!({"ok": true}), + ) + .await; + assert_eq!(env.r#type, EVENT_RESPONDED); + assert_eq!(env.payload["output"]["ok"], true); + ws2.close().await; +} + +#[tokio::test] +async fn acl_filtered_call_forbidden_over_ws() { + let addr = spawn_ws_server( + restricted_registry(), + provider_with(vec![("tok-1", identity("alice", &["user"]))]), + ) + .await; + let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1") + .await + .unwrap(); + + let env = call_and_await(&mut ws, "req-403", "admin/run", serde_json::json!({})).await; + assert_eq!(env.r#type, "call.error", "got {}", env.r#type); + assert_eq!(env.id, "req-403"); + assert_eq!(env.payload["code"], "FORBIDDEN"); + ws.close().await; +} + +#[tokio::test] +async fn text_message_rejected_with_protocol_close() { + let addr = spawn_ws_server( + echo_registry(), + provider_with(vec![("tok-1", identity("alice", &[]))]), + ) + .await; + let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1") + .await + .unwrap(); + + ws.send_text("hello as text").await; + let close = ws + .next_close(std::time::Duration::from_secs(5)) + .await + .expect("expected a close frame"); + assert_eq!(close, Some(1002)); +} + +#[tokio::test] +async fn internal_op_not_callable_from_wire() { + let addr = spawn_ws_server( + internal_registry(), + provider_with(vec![("tok-1", identity("alice", &[]))]), + ) + .await; + let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1") + .await + .unwrap(); + + let env = call_and_await(&mut ws, "req-secret", "secret/op", serde_json::json!({})).await; + assert_eq!(env.r#type, "call.error", "got {}", env.r#type); + assert_eq!(env.id, "req-secret"); + assert_eq!(env.payload["code"], "NOT_FOUND"); + ws.close().await; +} + +#[tokio::test] +async fn services_list_over_channel0_is_access_control_filtered() { + let registry = registry_with_services_list(vec![ + HandlerRegistration::new( + OperationSpec::new( + "public/echo", + OperationType::Query, + Visibility::External, + serde_json::json!({}), + serde_json::json!({}), + vec![], + AccessControl::default(), + None, + ), + HandlerKind::Once(echo_handler()), + OperationProvenance::Local, + None, + None, + alkcall::core::types::Capabilities::new(), + ), + HandlerRegistration::new( + OperationSpec::new( + "admin/secret", + OperationType::Query, + Visibility::External, + serde_json::json!({}), + serde_json::json!({}), + vec![], + AccessControl { + required_scopes: vec!["admin".to_string()], + ..Default::default() + }, + None, + ), + HandlerKind::Once(echo_handler()), + OperationProvenance::Local, + None, + None, + alkcall::core::types::Capabilities::new(), + ), + ]); + let addr = spawn_ws_server( + registry, + provider_with(vec![("tok-1", identity("regular", &["user"]))]), + ) + .await; + let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1") + .await + .unwrap(); + + let env = call_and_await(&mut ws, "req-list", "services/list", serde_json::json!({})).await; + assert_eq!(env.r#type, EVENT_RESPONDED, "got {}", env.r#type); + assert_eq!(env.id, "req-list"); + let ops = env.payload["output"]["operations"] + .as_array() + .expect("operations array"); + let names: Vec<&str> = ops.iter().filter_map(|o| o["name"].as_str()).collect(); + assert!(names.contains(&"public/echo"), "got: {names:?}"); + assert!(!names.contains(&"admin/secret"), "got: {names:?}"); + ws.close().await; +}