551 lines
21 KiB
Rust
551 lines
21 KiB
Rust
//! 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::collections::HashMap;
|
|
use std::sync::atomic::{AtomicU64, Ordering};
|
|
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 parking_lot::Mutex;
|
|
|
|
use super::byte_adapter::{
|
|
split_ws_to_bytes, WsPumps, INBOUND_WS_FRAME_CAP, INBOUND_WS_MESSAGE_CAP,
|
|
};
|
|
|
|
/// Registry of live WS session pump handles (WS-08). The upgrade
|
|
/// handler retains each session's [`WsPumps`] handle here for the
|
|
/// session's lifetime so a stuck session is evictable in-crate:
|
|
/// [`WsSessions::abort`] forces the pumps' tasks to end (the socket
|
|
/// halves close and the channels session unwinds). The handle for a
|
|
/// session is removed when the session task finishes (self-removing
|
|
/// guard), so the registry only holds live sessions. Assembly layers
|
|
/// share one instance via [`crate::server::state::RouterState`] (the
|
|
/// upgrade handler registers against it) or per-route request
|
|
/// extensions (an extension clone takes precedence).
|
|
///
|
|
/// Un-registered (default) — the upgrade runs fine and simply keeps no
|
|
/// eviction lever, matching the pre-WS-08 behavior.
|
|
#[derive(Clone, Default)]
|
|
pub struct WsSessions {
|
|
counter: Arc<AtomicU64>,
|
|
sessions: Arc<Mutex<HashMap<u64, Arc<WsPumps>>>>,
|
|
}
|
|
|
|
/// The upgrade handler's state slice: what it needs beyond the request
|
|
/// itself. Axum lifts it via `FromRef` from either full router state —
|
|
/// the adapter's [`crate::server::state::RouterState`] (carrying the
|
|
/// shared [`WsSessions`] instance) or a bare `Arc<OperationRegistry>`
|
|
/// (custom upgrade routes / integration tests get a handler-private
|
|
/// registry; eviction still works in-crate, just not shared).
|
|
#[derive(Clone)]
|
|
pub struct SessionState {
|
|
registry: Arc<OperationRegistry>,
|
|
sessions: Arc<WsSessions>,
|
|
}
|
|
|
|
impl SessionState {
|
|
/// From the plain registry state (custom upgrade routes /
|
|
/// integration tests): sessions default to a fresh [`WsSessions`]
|
|
/// private to the handler (eviction still functional in-crate but
|
|
/// not shared with the assembly layer).
|
|
pub(crate) fn from_registry(registry: &Arc<OperationRegistry>) -> Self {
|
|
Self {
|
|
registry: Arc::clone(registry),
|
|
sessions: Arc::new(WsSessions::new()),
|
|
}
|
|
}
|
|
|
|
/// From the adapter's router state: the shared [`WsSessions`]
|
|
/// instance the assembly layer can hold for eviction.
|
|
pub(crate) fn new(registry: Arc<OperationRegistry>, sessions: Arc<WsSessions>) -> Self {
|
|
Self { registry, sessions }
|
|
}
|
|
|
|
pub(crate) fn registry(&self) -> &Arc<OperationRegistry> {
|
|
&self.registry
|
|
}
|
|
|
|
pub(crate) fn sessions(&self) -> &Arc<WsSessions> {
|
|
&self.sessions
|
|
}
|
|
}
|
|
|
|
impl axum::extract::FromRef<SessionState> for Arc<OperationRegistry> {
|
|
fn from_ref(state: &SessionState) -> Self {
|
|
Arc::clone(&state.registry)
|
|
}
|
|
}
|
|
|
|
/// `FromRef` chain: a bare `Arc<OperationRegistry>` router state lifts
|
|
/// into the handler's [`SessionState`]; a full [`RouterState`] carries
|
|
/// the shared [`WsSessions`] instance and lifts through its own impl.
|
|
impl axum::extract::FromRef<Arc<OperationRegistry>> for SessionState {
|
|
fn from_ref(registry: &Arc<OperationRegistry>) -> Self {
|
|
SessionState::from_registry(registry)
|
|
}
|
|
}
|
|
|
|
impl WsSessions {
|
|
pub fn new() -> Self {
|
|
Self::default()
|
|
}
|
|
|
|
/// Abort every live session's pump tasks (forced teardown).
|
|
pub fn abort(&self) {
|
|
for (_, pumps) in self.sessions.lock().drain() {
|
|
pumps.abort();
|
|
}
|
|
}
|
|
|
|
/// Number of live sessions currently tracked.
|
|
pub fn len(&self) -> usize {
|
|
self.sessions.lock().len()
|
|
}
|
|
|
|
pub fn is_empty(&self) -> bool {
|
|
self.sessions.lock().is_empty()
|
|
}
|
|
|
|
fn insert(&self, pumps: Arc<WsPumps>) -> u64 {
|
|
let id = self.counter.fetch_add(1, Ordering::Relaxed);
|
|
self.sessions.lock().insert(id, pumps);
|
|
id
|
|
}
|
|
|
|
fn remove(&self, id: u64) {
|
|
self.sessions.lock().remove(&id);
|
|
}
|
|
}
|
|
|
|
/// The channels session for an upgraded socket: adapt → `Connection`
|
|
/// (identity attached) → `ChannelsAdapter::handle`. `policy` gates
|
|
/// data-channel opens (ADR-041). When `sessions` is `Some`, the
|
|
/// session's pump handle is registered for the session's lifetime —
|
|
/// the WS-08 eviction lever ([`WsSessions::abort`]).
|
|
pub async fn run_channels_session(
|
|
socket: axum::extract::ws::WebSocket,
|
|
registry: Arc<OperationRegistry>,
|
|
identity: Identity,
|
|
policy: Arc<dyn ChannelLifecyclePolicy>,
|
|
sessions: Option<WsSessions>,
|
|
) {
|
|
let (byte_stream, pumps) = split_ws_to_bytes(socket);
|
|
let pumps = Arc::new(pumps);
|
|
|
|
let _guard = sessions.map(|sessions| {
|
|
let id = sessions.insert(Arc::clone(&pumps));
|
|
SessionGuard { sessions, id }
|
|
});
|
|
|
|
let conn = Connection::from_bidi(byte_stream, b"alk/channels".to_vec(), None);
|
|
let _ = conn.set_identity(identity.clone());
|
|
|
|
let adapter = ChannelsAdapter::new(install_channel_zero(registry), policy);
|
|
let auth = AuthContext {
|
|
identity: Some(identity),
|
|
alpn: b"alk/channels".to_vec(),
|
|
remote_addr: None,
|
|
tls_client_fingerprint: None,
|
|
};
|
|
if let Err(e) = ProtocolHandler::handle(&adapter, conn, &auth).await {
|
|
tracing::warn!(error = %e, "channels session ended");
|
|
}
|
|
}
|
|
|
|
struct SessionGuard {
|
|
sessions: WsSessions,
|
|
id: u64,
|
|
}
|
|
|
|
impl Drop for SessionGuard {
|
|
fn drop(&mut self) {
|
|
self.sessions.remove(self.id);
|
|
}
|
|
}
|
|
|
|
/// 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 {
|
|
// The WS identity rides the upgrade request; propagate it to
|
|
// channel 0's `CallConnection` so the dispatcher's
|
|
// `resolve_identity` (and thus `AccessControl::check` on
|
|
// `services/list` and every operation) sees it. Without this
|
|
// the freshly constructed channel-0 `Connection` carries no
|
|
// identity and all ACL-restricted ops look unauthenticated.
|
|
if let Some(identity) = auth.identity.clone() {
|
|
let _ = channel0_conn.set_identity(identity);
|
|
}
|
|
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;
|
|
})
|
|
})
|
|
}
|
|
|
|
#[allow(dead_code)]
|
|
pub(crate) fn adapter_install_channel_zero(
|
|
registry: Arc<OperationRegistry>,
|
|
) -> alkcall::channels::adapter::InstallChannelZero {
|
|
install_channel_zero(registry)
|
|
}
|
|
|
|
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
|
|
}
|
|
}
|
|
|
|
/// Extension wrapper for the channels-policy injection point (SRV-10):
|
|
/// a deployment inserts `ChannelsPolicy(Arc<dyn ChannelLifecyclePolicy>)`
|
|
/// into the request extensions (a route layer on the WS route) to gate
|
|
/// data-channel opens per identity (ADR-041). Without the extension the
|
|
/// upgrade defaults to [`NoCap`] — the crate default for the built-in
|
|
/// surface (POC/trusted-peer semantics); assembly layers that build
|
|
/// their own upgrade route pass a stricter policy directly to
|
|
/// [`run_channels_session`].
|
|
#[derive(Clone)]
|
|
pub struct ChannelsPolicy(pub Arc<dyn ChannelLifecyclePolicy>);
|
|
|
|
/// 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`.
|
|
///
|
|
/// The channel lifecycle policy comes from the
|
|
/// [`ChannelsPolicy`] request extension when present, else `NoCap`.
|
|
/// The session pump handles are retained in the [`WsSessions`]
|
|
/// registry — the shared instance from the router state
|
|
/// (`RouterState::ws_sessions`) unless a request extension carries
|
|
/// one, so a stuck session stays evictable (WS-08).
|
|
pub async fn ws_upgrade_handler(
|
|
sessions: Option<axum::Extension<WsSessions>>,
|
|
axum::extract::State(state): axum::extract::State<SessionState>,
|
|
axum::Extension(identity): axum::Extension<Identity>,
|
|
policy: Option<axum::Extension<ChannelsPolicy>>,
|
|
ws_upgrade: WebSocketUpgrade,
|
|
) -> Response {
|
|
let sessions = Some(
|
|
sessions
|
|
.map(|axum::Extension(s)| s)
|
|
.unwrap_or_else(|| WsSessions::clone(state.sessions())),
|
|
);
|
|
let policy = policy
|
|
.map(|axum::Extension(p)| p.0)
|
|
.unwrap_or_else(|| Arc::new(NoCap));
|
|
let registry = Arc::clone(state.registry());
|
|
ws_upgrade
|
|
.max_frame_size(INBOUND_WS_FRAME_CAP)
|
|
.max_message_size(INBOUND_WS_MESSAGE_CAP)
|
|
.on_upgrade(move |socket| async move {
|
|
run_channels_session(socket, registry, identity, policy, sessions).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;
|
|
}
|
|
}
|
|
}
|