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,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<u8>),
|
||||
/// 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<Vec<u8>>,
|
||||
read_buf: Vec<u8>,
|
||||
read_pos: usize,
|
||||
eof: bool,
|
||||
write_tx: futures_mpsc::Sender<WriteMsg>,
|
||||
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::<Vec<u8>>(READ_SLOTS);
|
||||
let (write_tx, mut write_rx) = futures_mpsc::channel::<WriteMsg>(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<u8> = 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<u8> = 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<io::Result<()>> {
|
||||
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<io::Result<usize>> {
|
||||
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<io::Result<()>> {
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
|
||||
fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
||||
let this = self.get_mut();
|
||||
if this.write_open {
|
||||
this.write_open = false;
|
||||
drop(this.write_tx.clone());
|
||||
}
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
}
|
||||
@@ -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};
|
||||
|
||||
@@ -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