feat(websocket): WS upgrade route + channels session (server producer half)

- src/websocket/byte_adapter.rs: production WsByteStream from the POC —
  inbound bounded mpsc (64 slots, backpressure), outbound chunk parser
  emitting one WS message per chunk with 1 MiB split; write-side
  backpressure now uses futures mpsc poll_ready (POC spin-wait fixed);
  text messages closed with 1002; close mapping per websocket.md
- src/websocket/upgrade.rs: /alk/channels upgrade route — bearer auth
  (401 unresolvable), identity attached to the channels Connection,
  ChannelsAdapter + install_channel_zero running
  Dispatcher::run_loop_single_stream
- test_support module (feature test-support): WsClient, chunk/frame
  assemblers; shared with from_wss consumer path (ADR-070)
- tests/ws_upgrade_session.rs: 10 integration tests — call round-trip,
  services/list ACL-filtered, 3 MiB split, interleaved calls, ACL 403,
  internal-op NOT_FOUND, text->1002 close, disconnect mid-call no-hang

Verified: cargo test (95), cargo test --all-features (95+10),
clippy -D warnings (default + all-features), fmt.
This commit is contained in:
2026-08-28 08:47:13 +00:00
parent 0c1de05f85
commit 4ba9b652b3
6 changed files with 1262 additions and 9 deletions
+367
View File
@@ -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<OperationRegistry>,
identity: Identity,
policy: Arc<dyn ChannelLifecyclePolicy>,
) {
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<OperationRegistry>,
) -> alkcall::channels::adapter::InstallChannelZero {
Arc::new(move |_manager, channel0_conn, _auth| {
let registry = Arc::clone(&registry);
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<Identity> {
None
}
fn resolve_from_token(&self, _: &alkcall::core::auth::AuthToken) -> Option<Identity> {
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<Arc<OperationRegistry>>,
axum::Extension(identity): axum::Extension<Identity>,
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<dyn alkcall::core::auth::IdentityProvider>,
>,
mut req: axum::http::Request<axum::body::Body>,
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<u8> {
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<u8>,
}
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<u8>)> {
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<u8>,
}
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<EventEnvelope> {
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<u8> = 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::net::TcpStream>,
>,
tokio_tungstenite::tungstenite::Message,
>,
stream: futures::stream::SplitStream<
tokio_tungstenite::WebSocketStream<
tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
>,
>,
}
impl WsClient {
pub async fn connect_authorized(url: &str, token: &str) -> Result<Self, String> {
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::net::TcpStream>,
>,
_,
) = 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<u16> {
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<tokio::net::TcpStream>,
>;
type WsConnectResult = Result<
(
WsStream,
tokio_tungstenite::tungstenite::http::Response<Option<Vec<u8>>>,
),
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<tokio::net::TcpStream>,
>,
) -> Self {
let (sink, stream) = stream.split();
Self { sink, stream }
}
pub async fn send_binary(&mut self, bytes: Vec<u8>) {
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<Vec<u8>> {
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<Option<u16>> {
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;
}
}
}