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:
@@ -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(®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<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;
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user