alkhttp/tests/ws_overlay_ops.rs

//! Connection-local Layer 2 overlay verification for the channels-over-WS
//! session (the ws-overlay-ops acceptance gates; ADR-067 §Bidirectionality,
//! alkcall ADR-019/ADR-024, ADR-034 §4).
//!
//! A browser over WS has no `PeerId`, does not enter `PeerCompositeEnv`,
//! and any ops it registers land in a per-`CallConnection` overlay that
//! dies when the connection drops. The hub reaches browser ops through the
//! live connection handle's `overlay_env()` — not `PeerRef::Specific` (the
//! browser is not a peer). `AccessControl` on browser-registered ops gates
//! the hub's calls. WS close drops the overlay; in-flight calls fail.
//!
//! Two layers of verification, mirroring the alknet-http `overlay.rs` port:
//! 1. Wire-level: browser↔hub over a real axum WS server — bidirectional
//!    concurrent calls on channel 0 don't cross-correlate or deadlock, and
//!    disconnect mid-call resolves pending cleanly.
//! 2. Overlay-mechanism: `CallConnection::new_overlay_only` +
//!    `register_imported` + `overlay_env()` + `compose_root_env` — the
//!    registration path a hub integration uses to expose browser ops.

use alkcall::core::auth::{Identity, IdentityProvider};
use alkcall::core::types::Capabilities;
use alkcall::protocol::connection::CallConnection;
use alkcall::protocol::wire::{EventEnvelope, ResponseEnvelope};
use alkcall::registry::context::{CompositionAuthority, OperationContext, ScopedPeerEnv};
use alkcall::registry::env::{OperationEnv, PeerRef};
use alkcall::registry::registration::{
    make_handler, HandlerKind, HandlerRegistration, OperationProvenance, OperationRegistry,
};
use alkcall::registry::spec::{AccessControl, OperationSpec, OperationType, Visibility};
use alkhttp::websocket::{frame_channel0_chunk, ChunkAssembler, FrameAssembler, WsClient};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};

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 external_spec(name: &str, acl: AccessControl) -> OperationSpec {
    OperationSpec::new(
        name,
        OperationType::Query,
        Visibility::External,
        serde_json::json!({}),
        serde_json::json!({}),
        vec![],
        acl,
        None,
    )
}

fn echo_handler() -> alkcall::registry::registration::Handler {
    make_handler(|input, ctx| async move { ResponseEnvelope::ok(ctx.request_id, input) })
}

fn echo_registry() -> Arc<alkcall::registry::registration::OperationRegistry> {
    let registry = OperationRegistry::new();
    registry
        .register(HandlerRegistration::new(
            external_spec("echo/run", AccessControl::default()),
            HandlerKind::Once(echo_handler()),
            OperationProvenance::Local,
            None,
            None,
            Capabilities::new(),
        ))
        .unwrap();
    Arc::new(registry)
}

fn browser_registration(name: &str, acl: AccessControl) -> HandlerRegistration {
    HandlerRegistration::new(
        external_spec(name, acl),
        HandlerKind::Once(make_handler(|input, ctx| async move {
            ResponseEnvelope::ok(ctx.request_id, input)
        })),
        OperationProvenance::FromCall,
        None,
        None,
        Capabilities::new(),
    )
}

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_or_else(|e| e.into_inner())
            .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,
) -> alkcall::protocol::wire::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}"),
        }
    }
}

// ---------------------------------------------------------------------------
// Overlay mechanism (the registration path a hub integration uses)
// ---------------------------------------------------------------------------

fn hub_root_context(
    request_id: &str,
    allowed: &[&str],
    authority: CompositionAuthority,
    env: Arc<dyn OperationEnv + Send + Sync>,
) -> OperationContext {
    OperationContext {
        request_id: request_id.to_string(),
        parent_request_id: None,
        identity: None,
        handler_identity: Some(authority),
        forwarded_for: None,
        capabilities: Capabilities::new(),
        metadata: HashMap::new(),
        scoped_env: ScopedPeerEnv::new(allowed.iter().copied()),
        env,
        abort_policy: alkcall::registry::context::AbortPolicy::default(),
        deadline: Some(Instant::now() + Duration::from_secs(30)),
        internal: true,
        ownership: None,
    }
}

fn dispatcher(
    registry: Arc<OperationRegistry>,
    provider: Arc<dyn IdentityProvider>,
) -> alkcall::protocol::dispatch::Dispatcher {
    alkcall::protocol::dispatch::Dispatcher::new(registry, provider)
}

#[tokio::test]
async fn browser_registered_op_lands_in_overlay_not_peer_composite() {
    let conn = Arc::new(CallConnection::new_overlay_only(identity("browser", &[])));
    let overlay_env = conn.overlay_env();
    assert!(!overlay_env.contains("ui/dragged"));

    conn.register_imported(browser_registration("ui/dragged", AccessControl::default()))
        .expect("register_imported must accept the unmarked spec");
    assert!(overlay_env.contains("ui/dragged"));
    assert!(
        overlay_env.peer_ids().is_empty(),
        "browser overlay env exposes no PeerIds (browser is not a peer)"
    );
}

#[tokio::test]
async fn hub_call_reaches_browser_op_through_overlay_env() {
    let registry = echo_registry();
    let dp = dispatcher(registry, provider_with(vec![]));
    let conn = Arc::new(CallConnection::new_overlay_only(identity("browser", &[])));

    conn.register_imported(browser_registration("ui/dragged", AccessControl::default()))
        .expect("register_imported must accept the unmarked spec");

    let ctx = hub_root_context(
        "hub-call-1",
        &["ui/dragged"],
        CompositionAuthority::new("hub", vec![]),
        conn.overlay_env(),
    );
    let composed_env = dp.compose_root_env(&conn, &ctx);

    let invoke_ctx = hub_root_context(
        "hub-call-1",
        &["ui/dragged"],
        CompositionAuthority::new("hub", vec![]),
        composed_env.clone(),
    );
    let response = composed_env
        .invoke("ui", "dragged", serde_json::json!({ "x": 5 }), &invoke_ctx)
        .await;
    assert_eq!(
        response.result,
        Ok(serde_json::json!({ "x": 5 })),
        "hub→browser call routed through the connection-local overlay"
    );
}

#[tokio::test]
async fn peerref_specific_browser_x_routes_to_nothing() {
    let registry = echo_registry();
    let dp = dispatcher(registry, provider_with(vec![]));
    let conn = Arc::new(CallConnection::new_overlay_only(identity("browser", &[])));
    conn.register_imported(browser_registration("ui/dragged", AccessControl::default()))
        .expect("register_imported must accept the unmarked spec");

    let composed_env = dp.compose_root_env(
        &conn,
        &hub_root_context(
            "hub-peer-1",
            &["ui/dragged"],
            CompositionAuthority::new("hub", vec![]),
            conn.overlay_env(),
        ),
    );
    let ctx = hub_root_context(
        "hub-peer-1",
        &["ui/dragged"],
        CompositionAuthority::new("hub", vec![]),
        composed_env.clone(),
    );
    let response = composed_env
        .invoke_peer(
            &PeerRef::Specific("browser-X".to_string()),
            "ui",
            "dragged",
            serde_json::json!({}),
            &ctx,
            alkcall::registry::context::AbortPolicy::default(),
        )
        .await;
    match response.result {
        Err(e) => assert_eq!(e.code, "NOT_FOUND"),
        other => panic!("expected NOT_FOUND for PeerRef::Specific(browser-X), got {other:?}"),
    }
}

#[tokio::test]
async fn access_control_on_browser_op_gates_hub_call() {
    // Allowed: hub authority carries ui:write.
    let conn = Arc::new(CallConnection::new_overlay_only(identity("browser", &[])));
    conn.register_imported(browser_registration(
        "ui/dragged",
        AccessControl {
            required_scopes: vec!["ui:write".to_string()],
            ..Default::default()
        },
    ))
    .expect("register_imported must accept the unmarked spec");
    let env = conn.overlay_env();
    let ctx = hub_root_context(
        "hub-acl-ok",
        &["ui/dragged"],
        CompositionAuthority::new("hub", vec!["ui:write".to_string()]),
        env.clone(),
    );
    let response = env
        .invoke("ui", "dragged", serde_json::json!({ "v": 1 }), &ctx)
        .await;
    assert_eq!(response.result, Ok(serde_json::json!({ "v": 1 })));

    // Forbidden: hub authority lacks ui:write.
    let ctx_deny = hub_root_context(
        "hub-acl-deny",
        &["ui/dragged"],
        CompositionAuthority::new("hub", vec!["ui:read".to_string()]),
        env.clone(),
    );
    let response = env
        .invoke("ui", "dragged", serde_json::json!({}), &ctx_deny)
        .await;
    match response.result {
        Err(e) => assert_eq!(e.code, "FORBIDDEN"),
        other => panic!("expected FORBIDDEN, got {other:?}"),
    }
}

#[tokio::test]
async fn overlay_dropped_on_disconnect_no_leak_between_connections() {
    let conn1 = CallConnection::new_overlay_only(identity("browser-1", &[]));
    conn1
        .register_imported(browser_registration("ui/dragged", AccessControl::default()))
        .expect("register_imported must accept the unmarked spec");
    assert!(conn1.overlay_env().contains("ui/dragged"));
    drop(conn1);

    let conn2 = CallConnection::new_overlay_only(identity("browser-2", &[]));
    assert!(
        !conn2.overlay_env().contains("ui/dragged"),
        "a fresh connection's overlay is empty — the dropped connection's overlay did not leak"
    );

    // Connection-local isolation: each session's overlay is its own.
    let conn_a = CallConnection::new_overlay_only(identity("browser-a", &[]));
    let conn_b = CallConnection::new_overlay_only(identity("browser-b", &[]));
    conn_a
        .register_imported(browser_registration("ui/dragged", AccessControl::default()))
        .expect("register_imported must accept the unmarked spec");
    conn_b
        .register_imported(browser_registration("ui/click", AccessControl::default()))
        .expect("register_imported must accept the unmarked spec");
    assert!(conn_a.overlay_env().contains("ui/dragged"));
    assert!(!conn_a.overlay_env().contains("ui/click"));
    assert!(conn_b.overlay_env().contains("ui/click"));
    assert!(!conn_b.overlay_env().contains("ui/dragged"));
}

#[tokio::test]
async fn ws_close_mid_call_to_browser_op_aborts_call() {
    let conn = Arc::new(CallConnection::new_overlay_only(identity("browser", &[])));
    conn.register_imported(browser_registration("ui/dragged", AccessControl::default()))
        .expect("register_imported must accept the unmarked spec");

    let rx = {
        let mut pending = conn.pending().lock();
        pending.register_call(
            "hub-call-inflight".to_string(),
            Some(Instant::now() + Duration::from_secs(30)),
            None,
        )
    };

    let failed =
        conn.pending()
            .lock()
            .fail_all(alkcall::protocol::wire::CallError::connection_closed(
                "ws dropped",
            ));
    assert!(failed.contains(&"hub-call-inflight".to_string()));

    let result = tokio::time::timeout(Duration::from_millis(100), rx).await;
    match result {
        Ok(Ok(Err(e))) => assert_eq!(e.code, "CONNECTION_CLOSED"),
        other => panic!("expected Err from aborted call, got {other:?}"),
    }
    assert!(
        conn.pending().lock().is_empty(),
        "in-flight call aborted from pending map on ws close"
    );
    assert!(conn.overlay_env().contains("ui/dragged"));
}

// ---------------------------------------------------------------------------
// Wire-level: bidirectional concurrent calls over one live WS session
// ---------------------------------------------------------------------------

#[tokio::test]
async fn concurrent_bidirectional_calls_do_not_cross_correlate() {
    // Both sides initiate on channel 0 concurrently: the "browser" client
    // fires N echo calls while the hub side (a second WS session's calls)
    // runs in parallel. Same framing, same pending-map correlation rules —
    // ids must not cross.
    let addr = spawn_ws_server(
        echo_registry(),
        provider_with(vec![("tok-1", identity("alice", &[]))]),
    )
    .await;

    let mut ws1 = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
        .await
        .unwrap();
    let mut ws2 = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
        .await
        .unwrap();

    // Interleave: send from both connections before reading either, then
    // collect. Cross-correlation would surface as mismatched ids/outputs.
    for i in 0..5 {
        let frame = EventEnvelope::requested(
            format!("ws1-{i}"),
            serde_json::json!({ "operationId": "echo/run", "input": { "from": "ws1", "i": i } }),
        );
        ws1.send_binary(frame_channel0_chunk(&frame)).await;
        let frame = EventEnvelope::requested(
            format!("ws2-{i}"),
            serde_json::json!({ "operationId": "echo/run", "input": { "from": "ws2", "i": i } }),
        );
        ws2.send_binary(frame_channel0_chunk(&frame)).await;
    }

    async fn collect_all(ws: &mut WsClient) -> Vec<(String, serde_json::Value)> {
        let mut chunks = ChunkAssembler::new();
        let mut frames = FrameAssembler::new();
        let mut seen: Vec<(String, serde_json::Value)> = Vec::new();
        let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(10);
        while seen.len() < 5 {
            assert!(
                tokio::time::Instant::now() < deadline,
                "timed out collecting responses"
            );
            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);
                        frames.push(&payload);
                    }
                    while let Some(env) = frames.next_frame() {
                        assert_eq!(env.r#type, "call.responded");
                        let from = env.payload["output"]["from"]
                            .as_str()
                            .unwrap_or("")
                            .to_string();
                        seen.push((env.id, serde_json::json!({ "from": from })));
                    }
                }
                None => panic!("ws closed while collecting"),
            }
        }
        seen
    }

    let (seen1, seen2) = tokio::join!(collect_all(&mut ws1), collect_all(&mut ws2));

    for (id, out) in &seen1 {
        assert!(id.starts_with("ws1-"), "ws1 got foreign id {id}");
        assert_eq!(out["from"], "ws1", "ws1 response crossed to ws2 payload");
    }
    for (id, out) in &seen2 {
        assert!(id.starts_with("ws2-"), "ws2 got foreign id {id}");
        assert_eq!(out["from"], "ws2", "ws2 response crossed to ws1 payload");
    }
    ws1.close().await;
    ws2.close().await;
}

#[tokio::test]
async fn disconnect_drops_session_cleanly_subsequent_reach_fails_cleanly() {
    // Browser session with an in-flight call drops; a fresh session must
    // connect and work (no listener wedge), and the dropped session's
    // pending call must resolve (fail) rather than hang.
    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;
    drop(ws);

    // Fresh session: the dropped session's teardown did not wedge anything.
    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, "call.responded");
    assert_eq!(env.payload["output"]["ok"], true);
    ws2.close().await;
}

fn slow_and_echo_registry() -> Arc<OperationRegistry> {
    let registry = OperationRegistry::new();
    registry
        .register(HandlerRegistration::new(
            external_spec("slow/op", AccessControl::default()),
            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,
            Capabilities::new(),
        ))
        .unwrap();
    registry
        .register(HandlerRegistration::new(
            external_spec("echo/run", AccessControl::default()),
            HandlerKind::Once(echo_handler()),
            OperationProvenance::Local,
            None,
            None,
            Capabilities::new(),
        ))
        .unwrap();
    Arc::new(registry)
}

// Silence unused warnings for helpers shared with the module above.
#[allow(unused)]
fn noop(_: &str) {}