glm-5.3-flash
08584b229d
Coverage-confirmed dead code (every binary, zero hits): - server/state.rs: drop FromRef<RouterState> impls for Arc<OperationRegistry> and Arc<dyn IdentityProvider> — no route extracts these types; the auth middleware receives the provider directly via from_fn_with_state - gateway/dispatch.rs: drop identity_provider() and resolve_bearer() accessors; resolve_bearer's doc promised an auth hook the middleware never calls (spec/code drift). Wire-or-delete resolved to delete: bearer resolution lives in the middleware (SRV-11 single-resolve ordering), the dispatch spine only needs the per-call Option<Identity>. GatewayDispatch::new consequently takes the registry alone (GatewayState loses its unused identity_provider passthrough; dispatch.rs/to_mcp.rs tests simplified) - websocket/upgrade.rs: drop FromRef<SessionState> for Arc<OperationRegistry> — no router carries SessionState as its state type; the inverse FromRef<Arc<OperationRegistry>> for SessionState (custom upgrade routes, integration tests) remains Verification: ./scripts/verify.sh (352 passed), ./scripts/verify.sh --all-features (466 passed), clippy -D warnings, fmt --check.
719 lines
28 KiB
Rust
719 lines
28 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`.
|
|
//!
|
|
//! ## Detached task lifetime semantics (WS-10)
|
|
//!
|
|
//! Two task families spawned here are **detached by design** and
|
|
//! outlive the session task that spawned them:
|
|
//!
|
|
//! - the WS pump tasks (the `WsPumps` pair behind the byte adapter,
|
|
//! re-exported from `crate::websocket`), and
|
|
//! - the channel-0 dispatcher task spawned by `install_channel_zero`
|
|
//! per accepted channels connection.
|
|
//!
|
|
//! If the channels session task (`run_channels_session`) dies —
|
|
//! upgrade-time early return, an adapter fault, or its own task being
|
|
//! cancelled — these tasks keep running **self-healing**: each ends on
|
|
//! its own when its stream half closes (peer disconnect, peer close
|
|
//! frame, read error, or the local teardown arms —
|
|
//! `WsSessions::abort`, the idle-read timeout, `AsyncWrite::shutdown`).
|
|
//! They are never leaked unconditionally: the leak window is bounded
|
|
//! by peer behavior (a peer that holds the socket open keeps the pump
|
|
//! and dispatcher tasks alive with it) and additionally by the WS-01
|
|
//! idle-read timeout when configured, and every session's pumps stay
|
|
//! force-evictable through the WS-08 registry
|
|
//! ([`WsSessions::abort`]) for the session's whole lifetime. The
|
|
//! dispatcher task specifically ends when channel 0's `BiStream` read
|
|
//! side hits EOF (its only exit condition), i.e. when the underlying
|
|
//! WS connection ends by any of the paths above.
|
|
|
|
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_idle_with_write, 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 the adapter's `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>>>>,
|
|
}
|
|
|
|
/// Default bound on concurrent WS sessions (WS-09): the semaphore
|
|
/// initial capacity in [`SessionState`]; a deployment overrides it with
|
|
/// `HttpAdapter::with_ws_max_sessions`.
|
|
pub const DEFAULT_WS_MAX_SESSIONS: usize = 64;
|
|
|
|
/// 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 `RouterState` (carrying the shared [`WsSessions`]
|
|
/// instance and the configured session cap) or a bare
|
|
/// `Arc<OperationRegistry>` (custom upgrade routes / integration tests
|
|
/// get a handler-private registry and the default cap; eviction still
|
|
/// works in-crate, just not shared).
|
|
#[derive(Clone)]
|
|
pub struct SessionState {
|
|
registry: Arc<OperationRegistry>,
|
|
sessions: Arc<WsSessions>,
|
|
/// Session cap (WS-09): acquired post-auth, pre-upgrade; a caller
|
|
/// over the cap is rejected with 503.
|
|
session_slots: Arc<tokio::sync::Semaphore>,
|
|
/// Idle-read timeout (WS-01): `None` disables the knob.
|
|
idle_timeout: Option<std::time::Duration>,
|
|
}
|
|
|
|
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) and the default session cap.
|
|
pub(crate) fn from_registry(registry: &Arc<OperationRegistry>) -> Self {
|
|
Self {
|
|
registry: Arc::clone(registry),
|
|
sessions: Arc::new(WsSessions::new()),
|
|
session_slots: Arc::new(tokio::sync::Semaphore::new(DEFAULT_WS_MAX_SESSIONS)),
|
|
idle_timeout: Some(crate::websocket::DEFAULT_WS_IDLE_TIMEOUT),
|
|
}
|
|
}
|
|
|
|
/// From the adapter's router state: the shared [`WsSessions`]
|
|
/// instance the assembly layer can hold for eviction, and the
|
|
/// pre-built session-cap semaphore (one per `HttpAdapter`, shared
|
|
/// across requests).
|
|
pub(crate) fn new(
|
|
registry: Arc<OperationRegistry>,
|
|
sessions: Arc<WsSessions>,
|
|
session_slots: Arc<tokio::sync::Semaphore>,
|
|
idle_timeout: Option<std::time::Duration>,
|
|
) -> Self {
|
|
Self {
|
|
registry,
|
|
sessions,
|
|
session_slots,
|
|
idle_timeout,
|
|
}
|
|
}
|
|
|
|
pub(crate) fn registry(&self) -> &Arc<OperationRegistry> {
|
|
&self.registry
|
|
}
|
|
|
|
pub(crate) fn sessions(&self) -> &Arc<WsSessions> {
|
|
&self.sessions
|
|
}
|
|
|
|
pub(crate) fn idle_timeout(&self) -> Option<std::time::Duration> {
|
|
self.idle_timeout
|
|
}
|
|
}
|
|
|
|
/// `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 {
|
|
/// A fresh, empty session registry.
|
|
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()
|
|
}
|
|
|
|
/// Whether no sessions are tracked.
|
|
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`]). `idle_timeout`
|
|
/// bounds the read stall (WS-01): `None` disables the knob, `Some(d)`
|
|
/// closes the read with 1001 after `d` without inbound chunk
|
|
/// progress. `write_timeout` bounds one outbound WS send (WS-18):
|
|
/// `None` = the crate default window, `Some(d)` a deployment-set
|
|
/// window.
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub async fn run_channels_session(
|
|
socket: axum::extract::ws::WebSocket,
|
|
registry: Arc<OperationRegistry>,
|
|
identity: Identity,
|
|
policy: Arc<dyn ChannelLifecyclePolicy>,
|
|
sessions: Option<WsSessions>,
|
|
idle_timeout: Option<std::time::Duration>,
|
|
write_timeout: Option<std::time::Duration>,
|
|
) {
|
|
let (byte_stream, pumps) =
|
|
split_ws_to_bytes_idle_with_write(socket, idle_timeout, write_timeout);
|
|
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>);
|
|
|
|
/// Per-request WS pump timeouts (WS-17): a deployment inserts
|
|
/// `WsTimeouts` into the request extensions (a route layer on the WS
|
|
/// route, mirroring [`ChannelsPolicy`]) to set the pump knobs per
|
|
/// route instead of the router-state defaults. `idle` (WS-01) bounds
|
|
/// the read staleness (no completed inbound chunk for the window →
|
|
/// 1001 eviction); `write` (WS-18) bounds one outbound WS send (a peer
|
|
/// that stops reading is evicted once a single send outlasts it);
|
|
/// `None` disables a knob. Without the extension the
|
|
/// [`SessionState`] values apply — the built-in surface carries
|
|
/// `HttpAdapter::with_ws_idle_timeout` /
|
|
/// `DEFAULT_WS_IDLE_TIMEOUT` for the read side and
|
|
/// [`crate::websocket::DEFAULT_WS_WRITE_TIMEOUT`] for the write side,
|
|
/// which is also what bare-registry routes (custom upgrade routes on
|
|
/// a plain `Arc<OperationRegistry>` state) get: 60 s idle + 60 s
|
|
/// write windows, 64-session semaphore, handler-private
|
|
/// [`WsSessions`].
|
|
#[derive(Clone, Copy, Debug)]
|
|
pub struct WsTimeouts {
|
|
/// Idle-read window (WS-01); `None` disables the eviction.
|
|
pub idle: Option<std::time::Duration>,
|
|
/// Write-progress window (WS-18); `None` selects the crate
|
|
/// default (`DEFAULT_WS_WRITE_TIMEOUT`).
|
|
pub write: Option<std::time::Duration>,
|
|
}
|
|
|
|
impl Default for WsTimeouts {
|
|
fn default() -> Self {
|
|
Self {
|
|
idle: Some(crate::websocket::DEFAULT_WS_IDLE_TIMEOUT),
|
|
write: 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`.
|
|
///
|
|
/// 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).
|
|
///
|
|
/// The pump timeout knobs (WS-01 read / WS-18 write) come from the
|
|
/// [`WsTimeouts`] request extension when present, else from the
|
|
/// router state (`HttpAdapter::with_ws_idle_timeout` for the read;
|
|
/// the write side is fixed at
|
|
/// [`crate::websocket::DEFAULT_WS_WRITE_TIMEOUT`] on the built-in
|
|
/// surface). Bare-registry routes (`Arc<OperationRegistry>` state)
|
|
/// carry no configurable state — they run the documented defaults
|
|
/// above; a deployment overriding them inserts the `WsTimeouts`
|
|
/// extension on its upgrade route.
|
|
///
|
|
/// Session cap (WS-09): one semaphore permit is acquired per upgrade,
|
|
/// post-auth and pre-upgrade; when the configured cap
|
|
/// (`HttpAdapter::with_ws_max_sessions`, default
|
|
/// [`DEFAULT_WS_MAX_SESSIONS`]) is exhausted the upgrade is rejected
|
|
/// with **503 Service Unavailable** — holding the permit for the
|
|
/// session's lifetime, so ended sessions free their slot.
|
|
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>>,
|
|
timeouts: Option<axum::Extension<WsTimeouts>>,
|
|
ws_upgrade: WebSocketUpgrade,
|
|
) -> Response {
|
|
let Ok(permit) = state.session_slots.clone().try_acquire_owned() else {
|
|
return (
|
|
StatusCode::SERVICE_UNAVAILABLE,
|
|
"503 Service Unavailable: WS session cap reached",
|
|
)
|
|
.into_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());
|
|
let idle_timeout = match timeouts {
|
|
Some(axum::Extension(t)) => t.idle,
|
|
None => state.idle_timeout(),
|
|
};
|
|
let write_timeout = timeouts.and_then(|t| t.write);
|
|
ws_upgrade
|
|
.max_frame_size(INBOUND_WS_FRAME_CAP)
|
|
.max_message_size(INBOUND_WS_MESSAGE_CAP)
|
|
.on_upgrade(move |socket| async move {
|
|
let _permit = permit;
|
|
run_channels_session(
|
|
socket,
|
|
registry,
|
|
identity,
|
|
policy,
|
|
sessions,
|
|
idle_timeout,
|
|
Some(write_timeout.unwrap_or(crate::websocket::DEFAULT_WS_WRITE_TIMEOUT)),
|
|
)
|
|
.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.
|
|
///
|
|
/// # Panics
|
|
///
|
|
/// Panics only if `serde_json` cannot serialize the envelope —
|
|
/// unreachable for the acyclic wire type (no non-string map keys,
|
|
/// no untagged ambiguities), which is why this returns `Vec<u8>`
|
|
/// rather than `Result`: a test helper returning `Result` for an
|
|
/// impossible case is worse ergonomics than a documented panic
|
|
/// (review 001 HY-04, kept-as-is decision — the item ships behind
|
|
/// the opt-in `test-support` feature, the crate's documented
|
|
/// exception to no-panics-in-library-code).
|
|
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 {
|
|
/// A fresh, empty assembler.
|
|
pub fn new() -> Self {
|
|
Self::default()
|
|
}
|
|
|
|
/// Append raw bytes to the assembly buffer.
|
|
pub fn push(&mut self, bytes: &[u8]) {
|
|
self.buf.extend_from_slice(bytes);
|
|
}
|
|
|
|
/// Extract the next complete `(channel_id, payload)` chunk, if
|
|
/// a full one is buffered (8-byte header + declared payload
|
|
/// length).
|
|
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 {
|
|
/// A fresh, empty assembler.
|
|
pub fn new() -> Self {
|
|
Self::default()
|
|
}
|
|
|
|
/// Append raw bytes to the assembly buffer.
|
|
pub fn push(&mut self, bytes: &[u8]) {
|
|
self.buf.extend_from_slice(bytes);
|
|
}
|
|
|
|
/// Extract the next complete `EventEnvelope` frame, if a full
|
|
/// length-prefixed frame is buffered and parses. Unparseable
|
|
/// frames are dropped (test-only surface).
|
|
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 {
|
|
/// Connect with a `Bearer` token header; the full WS stream is
|
|
/// returned (upgrade succeeded).
|
|
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 }
|
|
}
|
|
|
|
/// Send one binary WS message.
|
|
///
|
|
/// # Panics
|
|
///
|
|
/// Panics on a socket write failure — a test client that cannot
|
|
/// send has a broken test, not a recoverable runtime state.
|
|
pub async fn send_binary(&mut self, bytes: Vec<u8>) {
|
|
self.send_binary_piece(&bytes).await;
|
|
}
|
|
|
|
/// Send one binary WS message, in pieces (for split-frame
|
|
/// tests). Same panic contract as [`Self::send_binary`].
|
|
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();
|
|
}
|
|
|
|
/// Send one text WS message. Same panic contract as
|
|
/// [`Self::send_binary`].
|
|
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),
|
|
}
|
|
}
|
|
|
|
/// Close the WS with a normal-close frame.
|
|
pub async fn close(&mut self) {
|
|
use futures::SinkExt;
|
|
let _ = self.sink.close().await;
|
|
}
|
|
}
|
|
}
|