691 lines
22 KiB
Rust
691 lines
22 KiB
Rust
//! WS upgrade + channels session integration tests (the ws-upgrade-session
|
|
//! acceptance gates): tokio-tungstenite client ↔ axum WS server
|
|
//! (upgrade → byte-adapter → ChannelsAdapter + channel-0 Dispatcher).
|
|
//! Grows from the ws-byte-adapter POC's validated suite.
|
|
|
|
use alkcall::core::auth::{Identity, IdentityProvider};
|
|
use alkcall::protocol::wire::{EventEnvelope, ResponseEnvelope, EVENT_RESPONDED};
|
|
use alkcall::registry::discovery::{services_list_handler, services_list_spec};
|
|
use alkcall::registry::registration::{
|
|
make_handler, Handler, HandlerKind, HandlerRegistration, OperationProvenance, OperationRegistry,
|
|
};
|
|
use alkcall::registry::spec::{AccessControl, OperationSpec, OperationType, Visibility};
|
|
use alkhttp::websocket::{
|
|
frame_channel0_chunk, ChunkAssembler, FrameAssembler, WsClient, INBOUND_WS_MESSAGE_CAP,
|
|
WS_MESSAGE_CAP,
|
|
};
|
|
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
|
|
fn identity(id: &str, scopes: &[&str]) -> Identity {
|
|
Identity {
|
|
id: id.to_string(),
|
|
scopes: scopes.iter().map(|s| s.to_string()).collect(),
|
|
resources: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
fn echo_handler() -> Handler {
|
|
make_handler(|input, ctx| async move { ResponseEnvelope::ok(ctx.request_id, input) })
|
|
}
|
|
|
|
fn echo_registry() -> Arc<OperationRegistry> {
|
|
let mut registry = OperationRegistry::new();
|
|
registry
|
|
.register(HandlerRegistration::new(
|
|
OperationSpec::new(
|
|
"echo/run",
|
|
OperationType::Query,
|
|
Visibility::External,
|
|
serde_json::json!({}),
|
|
serde_json::json!({}),
|
|
vec![],
|
|
AccessControl::default(),
|
|
None,
|
|
),
|
|
HandlerKind::Once(echo_handler()),
|
|
OperationProvenance::Local,
|
|
None,
|
|
None,
|
|
alkcall::core::types::Capabilities::new(),
|
|
))
|
|
.unwrap();
|
|
Arc::new(registry)
|
|
}
|
|
|
|
fn slow_and_echo_registry() -> Arc<OperationRegistry> {
|
|
let mut registry = OperationRegistry::new();
|
|
registry
|
|
.register(HandlerRegistration::new(
|
|
OperationSpec::new(
|
|
"slow/op",
|
|
OperationType::Query,
|
|
Visibility::External,
|
|
serde_json::json!({}),
|
|
serde_json::json!({}),
|
|
vec![],
|
|
AccessControl::default(),
|
|
None,
|
|
),
|
|
HandlerKind::Once(make_handler(|_input, _ctx| async move {
|
|
tokio::time::sleep(std::time::Duration::from_secs(30)).await;
|
|
ResponseEnvelope::ok("never", serde_json::json!({}))
|
|
})),
|
|
OperationProvenance::Local,
|
|
None,
|
|
None,
|
|
alkcall::core::types::Capabilities::new(),
|
|
))
|
|
.unwrap();
|
|
registry
|
|
.register(HandlerRegistration::new(
|
|
OperationSpec::new(
|
|
"echo/run",
|
|
OperationType::Query,
|
|
Visibility::External,
|
|
serde_json::json!({}),
|
|
serde_json::json!({}),
|
|
vec![],
|
|
AccessControl::default(),
|
|
None,
|
|
),
|
|
HandlerKind::Once(echo_handler()),
|
|
OperationProvenance::Local,
|
|
None,
|
|
None,
|
|
alkcall::core::types::Capabilities::new(),
|
|
))
|
|
.unwrap();
|
|
Arc::new(registry)
|
|
}
|
|
|
|
fn internal_registry() -> Arc<OperationRegistry> {
|
|
let mut registry = OperationRegistry::new();
|
|
registry
|
|
.register(HandlerRegistration::new(
|
|
OperationSpec::new(
|
|
"secret/op",
|
|
OperationType::Query,
|
|
Visibility::Internal,
|
|
serde_json::json!({}),
|
|
serde_json::json!({}),
|
|
vec![],
|
|
AccessControl::default(),
|
|
None,
|
|
),
|
|
HandlerKind::Once(echo_handler()),
|
|
OperationProvenance::Local,
|
|
None,
|
|
None,
|
|
alkcall::core::types::Capabilities::new(),
|
|
))
|
|
.unwrap();
|
|
Arc::new(registry)
|
|
}
|
|
|
|
fn restricted_registry() -> Arc<OperationRegistry> {
|
|
let mut registry = OperationRegistry::new();
|
|
registry
|
|
.register(HandlerRegistration::new(
|
|
OperationSpec::new(
|
|
"admin/run",
|
|
OperationType::Query,
|
|
Visibility::External,
|
|
serde_json::json!({}),
|
|
serde_json::json!({}),
|
|
vec![],
|
|
AccessControl {
|
|
required_scopes: vec!["admin".to_string()],
|
|
..Default::default()
|
|
},
|
|
None,
|
|
),
|
|
HandlerKind::Once(echo_handler()),
|
|
OperationProvenance::Local,
|
|
None,
|
|
None,
|
|
alkcall::core::types::Capabilities::new(),
|
|
))
|
|
.unwrap();
|
|
Arc::new(registry)
|
|
}
|
|
|
|
fn registry_with_services_list(inner_ops: Vec<HandlerRegistration>) -> Arc<OperationRegistry> {
|
|
let mut inner = OperationRegistry::new();
|
|
for op in inner_ops {
|
|
inner.register(op).unwrap();
|
|
}
|
|
let inner = Arc::new(inner);
|
|
let mut registry = OperationRegistry::new();
|
|
registry
|
|
.register(HandlerRegistration::new(
|
|
services_list_spec(),
|
|
HandlerKind::Once(services_list_handler(Arc::clone(&inner))),
|
|
OperationProvenance::Local,
|
|
None,
|
|
None,
|
|
alkcall::core::types::Capabilities::new(),
|
|
))
|
|
.unwrap();
|
|
for spec in inner.list_operations() {
|
|
let name = spec.name.clone();
|
|
let reg = inner.registration(&name).unwrap();
|
|
registry
|
|
.register(HandlerRegistration::new(
|
|
reg.spec.clone(),
|
|
reg.handler.clone(),
|
|
reg.provenance,
|
|
reg.composition_authority.clone(),
|
|
reg.scoped_env.clone(),
|
|
reg.capabilities.clone(),
|
|
))
|
|
.unwrap();
|
|
}
|
|
Arc::new(registry)
|
|
}
|
|
|
|
struct StaticTokens {
|
|
tokens: std::sync::Mutex<HashMap<String, Identity>>,
|
|
}
|
|
|
|
impl IdentityProvider for StaticTokens {
|
|
fn resolve_from_fingerprint(&self, _: &str) -> Option<Identity> {
|
|
None
|
|
}
|
|
fn resolve_from_token(&self, token: &alkcall::core::auth::AuthToken) -> Option<Identity> {
|
|
let s = String::from_utf8_lossy(&token.raw).to_string();
|
|
self.tokens.lock().unwrap().get(&s).cloned()
|
|
}
|
|
}
|
|
|
|
fn provider_with(tokens: Vec<(&str, Identity)>) -> Arc<dyn IdentityProvider> {
|
|
let map: HashMap<String, Identity> = tokens
|
|
.into_iter()
|
|
.map(|(t, i)| (t.to_string(), i))
|
|
.collect();
|
|
Arc::new(StaticTokens {
|
|
tokens: std::sync::Mutex::new(map),
|
|
})
|
|
}
|
|
|
|
async fn spawn_ws_server(
|
|
registry: Arc<OperationRegistry>,
|
|
provider: Arc<dyn IdentityProvider>,
|
|
) -> String {
|
|
let app = axum::Router::new()
|
|
.route(
|
|
"/alk/channels",
|
|
axum::routing::get(alkhttp::websocket::ws_upgrade_handler),
|
|
)
|
|
.layer(axum::middleware::from_fn_with_state(
|
|
provider,
|
|
alkhttp::websocket::ws_bearer_auth,
|
|
))
|
|
.with_state(registry);
|
|
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
|
let addr = format!("ws://{}", listener.local_addr().unwrap());
|
|
tokio::spawn(async move {
|
|
axum::serve(listener, app).await.unwrap();
|
|
});
|
|
addr
|
|
}
|
|
|
|
async fn call_and_await(
|
|
ws: &mut WsClient,
|
|
request_id: &str,
|
|
operation: &str,
|
|
input: serde_json::Value,
|
|
) -> EventEnvelope {
|
|
let frame = EventEnvelope::requested(
|
|
request_id,
|
|
serde_json::json!({ "operationId": operation, "input": input }),
|
|
);
|
|
ws.send_binary(frame_channel0_chunk(&frame)).await;
|
|
|
|
let mut chunks = ChunkAssembler::new();
|
|
let mut frames = FrameAssembler::new();
|
|
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(5);
|
|
loop {
|
|
assert!(
|
|
tokio::time::Instant::now() < deadline,
|
|
"timed out waiting for response to {request_id}"
|
|
);
|
|
if let Some(env) = frames.next_frame() {
|
|
return env;
|
|
}
|
|
let bin = ws.next_binary(std::time::Duration::from_millis(500)).await;
|
|
match bin {
|
|
Some(bytes) => {
|
|
chunks.push(&bytes);
|
|
while let Some((channel_id, payload)) = chunks.next_chunk() {
|
|
assert_eq!(channel_id, 0, "response must ride channel 0");
|
|
frames.push(&payload);
|
|
}
|
|
}
|
|
None => panic!("ws closed unexpectedly while awaiting {request_id}"),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn end_to_end_call_round_trip_over_ws() {
|
|
let addr = spawn_ws_server(
|
|
echo_registry(),
|
|
provider_with(vec![("tok-1", identity("alice", &[]))]),
|
|
)
|
|
.await;
|
|
|
|
let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
|
|
.await
|
|
.unwrap();
|
|
|
|
let env = call_and_await(
|
|
&mut ws,
|
|
"req-1",
|
|
"echo/run",
|
|
serde_json::json!({ "hello": "world" }),
|
|
)
|
|
.await;
|
|
assert_eq!(env.r#type, EVENT_RESPONDED, "got {}", env.r#type);
|
|
assert_eq!(env.id, "req-1");
|
|
assert_eq!(env.payload["output"]["hello"], "world");
|
|
ws.close().await;
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn upgrade_without_token_rejected_401() {
|
|
let addr = spawn_ws_server(
|
|
echo_registry(),
|
|
provider_with(vec![("tok-1", identity("alice", &[]))]),
|
|
)
|
|
.await;
|
|
let status = WsClient::connect_status(&format!("{addr}/alk/channels"), None).await;
|
|
assert_eq!(status, Some(401));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn upgrade_with_unresolvable_token_rejected_401() {
|
|
let addr = spawn_ws_server(
|
|
echo_registry(),
|
|
provider_with(vec![("tok-1", identity("alice", &[]))]),
|
|
)
|
|
.await;
|
|
let status =
|
|
WsClient::connect_status(&format!("{addr}/alk/channels"), Some("wrong-token")).await;
|
|
assert_eq!(status, Some(401));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn oversized_payload_splits_across_ws_messages_and_round_trips() {
|
|
let addr = spawn_ws_server(
|
|
echo_registry(),
|
|
provider_with(vec![("tok-1", identity("alice", &[]))]),
|
|
)
|
|
.await;
|
|
let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
|
|
.await
|
|
.unwrap();
|
|
|
|
let payload_len = 3 * 1024 * 1024;
|
|
let blob = "x".repeat(payload_len);
|
|
let frame = EventEnvelope::requested(
|
|
"req-big",
|
|
serde_json::json!({ "operationId": "echo/run", "input": { "blob": blob } }),
|
|
);
|
|
let chunk = frame_channel0_chunk(&frame);
|
|
assert!(
|
|
chunk.len() > WS_MESSAGE_CAP,
|
|
"test precondition: request exceeds message cap"
|
|
);
|
|
|
|
for piece in chunk.chunks(WS_MESSAGE_CAP) {
|
|
ws.send_binary_piece(piece).await;
|
|
}
|
|
|
|
let mut chunks = ChunkAssembler::new();
|
|
let mut frames = FrameAssembler::new();
|
|
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(30);
|
|
loop {
|
|
assert!(
|
|
tokio::time::Instant::now() < deadline,
|
|
"timed out on big payload"
|
|
);
|
|
if let Some(env) = frames.next_frame() {
|
|
assert_eq!(env.r#type, EVENT_RESPONDED, "got {}", env.r#type);
|
|
assert_eq!(env.id, "req-big");
|
|
let out = env.payload["output"]["blob"].as_str().unwrap();
|
|
assert_eq!(out.len(), payload_len, "blob round-tripped intact");
|
|
assert!(out.bytes().all(|b| b == b'x'), "byte integrity");
|
|
ws.close().await;
|
|
return;
|
|
}
|
|
let bytes = ws
|
|
.next_binary(std::time::Duration::from_secs(2))
|
|
.await
|
|
.expect("ws closed unexpectedly on big payload");
|
|
chunks.push(&bytes);
|
|
while let Some((channel_id, payload)) = chunks.next_chunk() {
|
|
assert_eq!(channel_id, 0);
|
|
frames.push(&payload);
|
|
}
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn interleaved_calls_reassemble_without_corruption() {
|
|
let addr = spawn_ws_server(
|
|
echo_registry(),
|
|
provider_with(vec![("tok-1", identity("alice", &[]))]),
|
|
)
|
|
.await;
|
|
let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
|
|
.await
|
|
.unwrap();
|
|
|
|
const N: usize = 20;
|
|
for i in 0..N {
|
|
let frame = EventEnvelope::requested(
|
|
format!("req-{i}"),
|
|
serde_json::json!({ "operationId": "echo/run", "input": { "i": i } }),
|
|
);
|
|
ws.send_binary(frame_channel0_chunk(&frame)).await;
|
|
}
|
|
|
|
let mut chunks = ChunkAssembler::new();
|
|
let mut frames = FrameAssembler::new();
|
|
let mut seen: Vec<String> = Vec::new();
|
|
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(10);
|
|
while seen.len() < N {
|
|
assert!(
|
|
tokio::time::Instant::now() < deadline,
|
|
"only got {} of {}",
|
|
seen.len(),
|
|
N
|
|
);
|
|
let bytes = ws
|
|
.next_binary(std::time::Duration::from_millis(500))
|
|
.await
|
|
.expect("ws closed unexpectedly");
|
|
chunks.push(&bytes);
|
|
while let Some((channel_id, payload)) = chunks.next_chunk() {
|
|
assert_eq!(channel_id, 0);
|
|
frames.push(&payload);
|
|
while let Some(env) = frames.next_frame() {
|
|
assert_eq!(env.r#type, EVENT_RESPONDED, "got {}", env.r#type);
|
|
let i: usize = env.payload["output"]["i"].as_u64().unwrap() as usize;
|
|
assert_eq!(env.id, format!("req-{i}"), "correlation intact");
|
|
seen.push(env.id);
|
|
}
|
|
}
|
|
}
|
|
seen.sort();
|
|
let mut expected: Vec<String> = (0..N).map(|i| format!("req-{i}")).collect();
|
|
expected.sort();
|
|
assert_eq!(seen, expected);
|
|
ws.close().await;
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn disconnect_mid_call_does_not_hang_the_server() {
|
|
let addr = spawn_ws_server(
|
|
slow_and_echo_registry(),
|
|
provider_with(vec![("tok-1", identity("alice", &[]))]),
|
|
)
|
|
.await;
|
|
let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
|
|
.await
|
|
.unwrap();
|
|
|
|
let frame = EventEnvelope::requested(
|
|
"req-abandon",
|
|
serde_json::json!({ "operationId": "slow/op", "input": {} }),
|
|
);
|
|
ws.send_binary(frame_channel0_chunk(&frame)).await;
|
|
|
|
// Abrupt client disconnect (drop, no close handshake) while the
|
|
// slow handler is still running.
|
|
drop(ws);
|
|
|
|
// The server must not hang: a fresh session completes an echo call
|
|
// promptly — the abandoned session's teardown (pending failure,
|
|
// overlay drop) did not wedge the runtime or the listener.
|
|
let mut ws2 = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
|
|
.await
|
|
.unwrap();
|
|
let env = call_and_await(
|
|
&mut ws2,
|
|
"req-after",
|
|
"echo/run",
|
|
serde_json::json!({"ok": true}),
|
|
)
|
|
.await;
|
|
assert_eq!(env.r#type, EVENT_RESPONDED);
|
|
assert_eq!(env.payload["output"]["ok"], true);
|
|
ws2.close().await;
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn acl_filtered_call_forbidden_over_ws() {
|
|
let addr = spawn_ws_server(
|
|
restricted_registry(),
|
|
provider_with(vec![("tok-1", identity("alice", &["user"]))]),
|
|
)
|
|
.await;
|
|
let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
|
|
.await
|
|
.unwrap();
|
|
|
|
let env = call_and_await(&mut ws, "req-403", "admin/run", serde_json::json!({})).await;
|
|
assert_eq!(env.r#type, "call.error", "got {}", env.r#type);
|
|
assert_eq!(env.id, "req-403");
|
|
assert_eq!(env.payload["code"], "FORBIDDEN");
|
|
ws.close().await;
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn text_message_rejected_with_protocol_close() {
|
|
let addr = spawn_ws_server(
|
|
echo_registry(),
|
|
provider_with(vec![("tok-1", identity("alice", &[]))]),
|
|
)
|
|
.await;
|
|
let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
|
|
.await
|
|
.unwrap();
|
|
|
|
ws.send_text("hello as text").await;
|
|
let close = ws
|
|
.next_close(std::time::Duration::from_secs(5))
|
|
.await
|
|
.expect("expected a close frame");
|
|
assert_eq!(close, Some(1002));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn internal_op_not_callable_from_wire() {
|
|
let addr = spawn_ws_server(
|
|
internal_registry(),
|
|
provider_with(vec![("tok-1", identity("alice", &[]))]),
|
|
)
|
|
.await;
|
|
let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
|
|
.await
|
|
.unwrap();
|
|
|
|
let env = call_and_await(&mut ws, "req-secret", "secret/op", serde_json::json!({})).await;
|
|
assert_eq!(env.r#type, "call.error", "got {}", env.r#type);
|
|
assert_eq!(env.id, "req-secret");
|
|
assert_eq!(env.payload["code"], "NOT_FOUND");
|
|
ws.close().await;
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn services_list_over_channel0_is_access_control_filtered() {
|
|
let registry = registry_with_services_list(vec![
|
|
HandlerRegistration::new(
|
|
OperationSpec::new(
|
|
"public/echo",
|
|
OperationType::Query,
|
|
Visibility::External,
|
|
serde_json::json!({}),
|
|
serde_json::json!({}),
|
|
vec![],
|
|
AccessControl::default(),
|
|
None,
|
|
),
|
|
HandlerKind::Once(echo_handler()),
|
|
OperationProvenance::Local,
|
|
None,
|
|
None,
|
|
alkcall::core::types::Capabilities::new(),
|
|
),
|
|
HandlerRegistration::new(
|
|
OperationSpec::new(
|
|
"admin/secret",
|
|
OperationType::Query,
|
|
Visibility::External,
|
|
serde_json::json!({}),
|
|
serde_json::json!({}),
|
|
vec![],
|
|
AccessControl {
|
|
required_scopes: vec!["admin".to_string()],
|
|
..Default::default()
|
|
},
|
|
None,
|
|
),
|
|
HandlerKind::Once(echo_handler()),
|
|
OperationProvenance::Local,
|
|
None,
|
|
None,
|
|
alkcall::core::types::Capabilities::new(),
|
|
),
|
|
]);
|
|
let addr = spawn_ws_server(
|
|
registry,
|
|
provider_with(vec![("tok-1", identity("regular", &["user"]))]),
|
|
)
|
|
.await;
|
|
let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
|
|
.await
|
|
.unwrap();
|
|
|
|
let env = call_and_await(&mut ws, "req-list", "services/list", serde_json::json!({})).await;
|
|
assert_eq!(env.r#type, EVENT_RESPONDED, "got {}", env.r#type);
|
|
assert_eq!(env.id, "req-list");
|
|
let ops = env.payload["output"]["operations"]
|
|
.as_array()
|
|
.expect("operations array");
|
|
let names: Vec<&str> = ops.iter().filter_map(|o| o["name"].as_str()).collect();
|
|
assert!(names.contains(&"public/echo"), "got: {names:?}");
|
|
assert!(!names.contains(&"admin/secret"), "got: {names:?}");
|
|
ws.close().await;
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn inbound_message_over_cap_fails_the_connection() {
|
|
let addr = spawn_ws_server(
|
|
echo_registry(),
|
|
provider_with(vec![("tok-1", identity("alice", &[]))]),
|
|
)
|
|
.await;
|
|
let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
|
|
.await
|
|
.unwrap();
|
|
|
|
// One WS binary message larger than INBOUND_WS_MESSAGE_CAP: the
|
|
// server's explicit max_message_size (WS-06) must fail the
|
|
// connection instead of accepting it (the byte-adapter boundary is
|
|
// the chunk header, so the adapter never dissects messages; the WS
|
|
// library cap is the enforcement point).
|
|
let oversized = vec![b'x'; INBOUND_WS_MESSAGE_CAP + 1];
|
|
ws.send_binary_piece(&oversized).await;
|
|
|
|
let close = ws
|
|
.next_close(std::time::Duration::from_secs(5))
|
|
.await
|
|
.expect("server must fail the oversized-inbound connection (WS-06)");
|
|
assert!(
|
|
close.is_some(),
|
|
"expected a close or stream end, got {close:?}"
|
|
);
|
|
}
|
|
|
|
/// WS-08 acceptance: the upgrade retains each session's pump handle in
|
|
/// the shared `WsSessions` registry; `abort()` evicts a stuck session
|
|
/// (the peer's socket closes) and the entry is removed when the
|
|
/// session task ends, so the registry only tracks live sessions.
|
|
#[tokio::test]
|
|
async fn ws_sessions_registry_tracks_and_aborts_live_sessions() {
|
|
use alkhttp::websocket::WsSessions;
|
|
|
|
let sessions = WsSessions::new();
|
|
let registry = echo_registry();
|
|
let provider = provider_with(vec![("tok-1", identity("alice", &[]))]);
|
|
|
|
let app = axum::Router::new()
|
|
.route(
|
|
"/alk/channels",
|
|
axum::routing::get(alkhttp::websocket::ws_upgrade_handler),
|
|
)
|
|
.layer(axum::middleware::from_fn_with_state(
|
|
Arc::clone(&provider),
|
|
alkhttp::websocket::ws_bearer_auth,
|
|
))
|
|
.with_state(registry)
|
|
.layer(axum::Extension(sessions.clone()));
|
|
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
|
let addr = format!("ws://{}", listener.local_addr().unwrap());
|
|
tokio::spawn(async move {
|
|
axum::serve(listener, app).await.unwrap();
|
|
});
|
|
|
|
let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
|
|
.await
|
|
.unwrap();
|
|
let env = call_and_await(&mut ws, "req-1", "echo/run", serde_json::json!({})).await;
|
|
assert_eq!(env.r#type, EVENT_RESPONDED, "session established");
|
|
|
|
// The upgrade registered the session.
|
|
let saw_session = tokio::time::timeout(std::time::Duration::from_secs(5), async {
|
|
loop {
|
|
if !sessions.is_empty() {
|
|
return true;
|
|
}
|
|
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
|
|
}
|
|
})
|
|
.await
|
|
.expect("session registered in the shared registry");
|
|
assert!(saw_session);
|
|
assert_eq!(sessions.len(), 1);
|
|
|
|
// Forced teardown: abort() closes the peer's socket.
|
|
sessions.abort();
|
|
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(10);
|
|
let close = loop {
|
|
assert!(
|
|
tokio::time::Instant::now() < deadline,
|
|
"abort never reached the peer"
|
|
);
|
|
match ws.next_close(std::time::Duration::from_millis(250)).await {
|
|
Some(x) => break Some(x),
|
|
None => continue,
|
|
}
|
|
};
|
|
assert!(close.is_some(), "peer observes the teardown: {close:?}");
|
|
|
|
// The self-removing guard drops the entry once the session task ends.
|
|
let drained = tokio::time::timeout(std::time::Duration::from_secs(5), async {
|
|
loop {
|
|
if sessions.is_empty() {
|
|
return true;
|
|
}
|
|
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
|
|
}
|
|
})
|
|
.await
|
|
.expect("session entry removed after the session task ends");
|
|
assert!(drained);
|
|
ws.close().await;
|
|
}
|