alkssh/tests/spawn_return.rs

//! The spawn-and-return mechanics proof (`w2-spawn-return`) — the
//! structural constraint (ADR-006 decision 4) pinned as tests, over
//! the real wire.
//!
//! russh at this pin runs server handler callbacks inline inside the
//! session's packet loop; a blocking handler deadlocks against the
//! peer's follow-up bytes (POC-3 proved by construction with
//! receive-pack). The crate's plumbing
//! ([`alkssh::spawn_serving_path`],
//! [`alkssh::report_exit_status_then_close`] — the `#[doc(hidden)]`
//! contract-test re-exports from `src/serve.rs`) makes the escape
//! mechanical: the handler hook clones `session.handle()`, spawns the
//! serving future, and returns.
//!
//! These rows compose exactly that plumbing end to end: a minimal
//! proof handler stands in for the crate's own auth/dispatch hooks
//! (their correctness is `w2-auth-bridge`/`w2-dispatch-decline-core`
//! contract territory; the `dispatch` seam stubs now spawn through the
//! same plumbing, so the crate's own rows exercise the spawn too —
//! `tests/dispatch_core.rs`). What is under test here is the
//! *mechanics*: the spawn-then-return, the cloned-handle reporting,
//! and the wire ordering — `exit-status` before close, arriving after
//! the handler already returned, while the session's packet loop keeps
//! routing the connection's follow-up traffic.

#![allow(dead_code)]

mod common;

use std::sync::Arc;
use std::time::Instant;

use common::fixture;
use common::harness::{
    HarnessClientHandler, authenticate_publickey, client_config, next_message, open_session_channel,
};
use parking_lot::Mutex;
use russh::keys::Algorithm;
use russh::server::{Auth, ChannelOpenHandle, Session};
use russh::{Channel, ChannelId, ChannelMsg, Disconnect, Error, MethodKind, MethodSet};

use alkssh::Limits;

/// The exit status the proof stub reports — an arbitrary distinct
/// value the rows assert verbatim (`exit-status(7)`).
const PROOF_EXIT: u32 = 7;

/// The command the rows exec — anything; the proof handler ignores it.
const PROOF_COMMAND: &str = "mechanics-proof-serving-path";

/// The async completion gate the spawned serving paths hold on —
/// released only after the rows have proven the handler returned.
/// `watch`-based (no lost-wakeup race: a late subscriber still sees
/// the release).
fn serving_gate() -> tokio::sync::watch::Sender<bool> {
    tokio::sync::watch::channel(false).0
}

/// The spawned task's wait: blocks until (or returns if) released —
/// `borrow_and_update` marks the release seen so the wakeup is never
/// lost or replayed.
async fn gate_wait(rx: &mut tokio::sync::watch::Receiver<bool>) {
    if !*rx.borrow_and_update() {
        rx.changed().await.expect("the gate lives for the row");
    }
}

/// What the proof handler records — the timestamps the rows compare.
#[derive(Default)]
struct ProofRecord {
    handler_returned_at: Option<Instant>,
    exit_reported_at: Option<Instant>,
}

type Shared = Arc<Mutex<ProofRecord>>;

/// The proof handler (server side): mechanics-only — it authenticates
/// the fixture alpha by the crate's one fingerprint path, accepts
/// session channels, and serves every exec by *the crate's own
/// spawn-and-return plumbing*: a serving path gated on the row's
/// release, reporting `exit-status(7)` then close through the cloned
/// handle.
struct ProofServer {
    gate: Arc<tokio::sync::watch::Sender<bool>>,
    record: Shared,
}

impl russh::server::Handler for ProofServer {
    type Error = Error;

    async fn auth_publickey_offered(
        &mut self,
        _user: &str,
        _offer: russh::server::PublicKeyOffer<'_>,
    ) -> Result<Auth, Self::Error> {
        Ok(Auth::Accept)
    }

    /// The mechanics rows' own fixture check — fingerprint equality
    /// against fixture alpha through the crate's fingerprint path
    /// (not the bridge contract; that is `w2-auth-bridge`'s suite).
    async fn auth_publickey(
        &mut self,
        _user: &str,
        key: russh::server::VerifiedPublicKey<'_>,
    ) -> Result<Auth, Self::Error> {
        let accepted = *key.algorithm == Algorithm::Ed25519
            && alkssh::fingerprint(key.public_key) == fixture("alpha").fingerprint;
        Ok(if accepted {
            Auth::Accept
        } else {
            Auth::reject()
        })
    }

    async fn channel_open_session(
        &mut self,
        _channel: Channel<russh::server::Msg>,
        reply: ChannelOpenHandle,
        _session: &mut Session,
    ) -> Result<(), Self::Error> {
        reply.accept().await;
        Ok(())
    }

    /// The hook under proof: clone the handle, spawn the serving path
    /// (gated — it stays pending past this method's return), record
    /// the handler's return instant, and go.
    async fn exec_request(
        &mut self,
        channel: ChannelId,
        _data: &[u8],
        session: &mut Session,
    ) -> Result<(), Self::Error> {
        let handle = session.handle();
        let record = Arc::clone(&self.record);
        let mut gate_rx = self.gate.subscribe();

        alkssh::spawn_serving_path(async move {
            gate_wait(&mut gate_rx).await;
            alkssh::report_exit_status_then_close(&handle, channel, PROOF_EXIT).await;
            record.lock().exit_reported_at = Some(Instant::now());
        });

        self.record.lock().handler_returned_at = Some(Instant::now());
        Ok(())
    }
}

/// One mounted mechanics-proof session: the duplex pair, the spawned
/// server task (russh's own stack over the proof handler), and the
/// connected client. Transport and posture mirror the loopback
/// harness (its recorded decisions); the handler is intentionally not
/// the crate's `SessionHandler` — the crate stack under a custom
/// serving hook is the proof's subject.
struct ProofSession {
    client: russh::client::Handle<HarnessClientHandler>,
    gate: Arc<tokio::sync::watch::Sender<bool>>,
    record: Shared,
    server_task: tokio::task::JoinHandle<Result<(), Error>>,
}

/// Mounts one proof session: the russh server stack over the proof
/// handler, the harness client against it, the fixture alpha
/// authenticated.
async fn mount() -> ProofSession {
    let gate = Arc::new(serving_gate());
    let record: Shared = Arc::new(Mutex::new(ProofRecord::default()));

    let config = Arc::new(russh::server::Config {
        methods: MethodSet::from(&[MethodKind::PublicKey][..]),
        keys: vec![fixture("bravo").private_key()],
        inactivity_timeout: Some(Limits::default().inactivity_timeout),
        ..russh::server::Config::default()
    });

    let (client_io, server_io) = tokio::io::duplex(common::harness::DUPLEX_CAPACITY);
    let server = ProofServer {
        gate: Arc::clone(&gate),
        record: Arc::clone(&record),
    };
    let server_task = tokio::spawn(async move {
        let running = russh::server::run_stream(config, server_io, server)
            .await
            .expect("the proof server mounts");
        running.await
    });

    let policy = Arc::new(common::MemHostKeyPolicy::with_default(true));
    let mut client = russh::client::connect_stream(
        client_config(Limits::default()),
        client_io,
        HarnessClientHandler::new(Arc::clone(&policy)),
    )
    .await
    .expect("the client mounts against the proof server");

    let result = authenticate_publickey(&mut client, "mechanics-proof", fixture("alpha")).await;
    assert!(
        result.success(),
        "the fixture identity authenticates before every mechanics row"
    );

    ProofSession {
        client,
        gate,
        record,
        server_task,
    }
}

impl ProofSession {
    /// Releases the gate — the spawned serving paths complete
    /// asynchronously from here on.
    fn release_gate(&self) {
        let _ = self.gate.send(true);
    }

    /// Ends a row cleanly: disconnect, wait for the server task, and
    /// require a clean session end (the harness posture).
    async fn unmount(self) {
        self.client
            .disconnect(Disconnect::ByApplication, "row done", "en")
            .await
            .expect("client disconnects cleanly");
        let outcome = self.server_task.await.expect("server task ends");
        outcome.expect("clean session end");
    }

    /// Waits (bounded) until the exec hook has provably run and
    /// returned — the client's `exec` returns when the request is on
    /// the wire; the loop's processing of it is observed here.
    async fn exec_hook_has_returned(&self) {
        tokio::time::timeout(std::time::Duration::from_secs(5), async {
            while self.record.lock().handler_returned_at.is_none() {
                tokio::time::sleep(std::time::Duration::from_millis(2)).await;
            }
        })
        .await
        .expect("the exec hook returns promptly");
    }

    /// The handler provably returned while the spawned serving path
    /// was still pending — the handler-returns-fast assertion.
    fn assert_handler_returned_spawn_still_pending(&self, context: &str) {
        let record = self.record.lock();
        assert!(
            record.handler_returned_at.is_some(),
            "{context}: the exec hook must already have returned"
        );
        assert!(
            record.exit_reported_at.is_none(),
            "{context}: the spawned serving path is still gated — nothing reported yet"
        );
    }
}

// ---------------------------------------------------------------------------
// Rows.
// ---------------------------------------------------------------------------

/// - *Contract stamp:* ADR-006 decision 4 (the structural constraint:
///   every serving path is spawned — the handler completes while the
///   spawned task is still running) + mapping.md §The structural
///   constraint: spawn-and-return (the follow-up-bytes deadlock
///   inverted: the connection's loop keeps routing while a serving
///   path is pending).
#[tokio::test]
async fn handler_returns_fast_while_the_spawned_serving_path_is_pending() {
    let session = mount().await;

    let mut channel = open_session_channel(&session.client)
        .await
        .expect("the session channel opens");
    channel
        .exec(false, PROOF_COMMAND)
        .await
        .expect("the exec request rides the wire");
    session.exec_hook_has_returned().await;
    session.assert_handler_returned_spawn_still_pending("after the exec");

    // The follow-up-bytes inversion — POC-3's deadlock shape: the peer
    // sends immediate follow-up traffic and the loop routes it while
    // the first serving path is still pending.
    let follow_up = open_session_channel(&session.client)
        .await
        .expect("the packet loop routes while the spawn is pending");
    follow_up
        .exec(false, PROOF_COMMAND)
        .await
        .expect("a second exec dispatches while the first serving path is pending");
    session.assert_handler_returned_spawn_still_pending("after the follow-up traffic");

    session.release_gate();

    let first = next_message(&mut channel)
        .await
        .expect("the released serving path reports");
    assert!(
        matches!(
            first,
            ChannelMsg::ExitStatus {
                exit_status: PROOF_EXIT
            }
        ),
        "the spawned stub reports exit-status({PROOF_EXIT}), got {first:?}"
    );

    session.unmount().await;
}

/// - *Contract stamp:* mapping.md §exit_status row (sent **before**
///   close, from the spawned task via `Session::handle()`) + ADR-006
///   decision 4 (exit/status reporting rides the cloned handle).
#[tokio::test]
async fn spawned_stub_reports_exit_status_before_close_after_the_handler_returned() {
    let session = mount().await;

    let mut channel = open_session_channel(&session.client)
        .await
        .expect("the session channel opens");
    let handler_returned_at = {
        channel
            .exec(false, PROOF_COMMAND)
            .await
            .expect("the exec request rides the wire");
        session.exec_hook_has_returned().await;
        let handler_returned_at = session
            .record
            .lock()
            .handler_returned_at
            .expect("the exec hook returned before this row observes anything");
        session.release_gate();
        handler_returned_at
    };

    // Strict wire order, message by message: exit-status first, close
    // after — the ordering contract verbatim.
    let first = next_message(&mut channel)
        .await
        .expect("the released serving path reports");
    assert!(
        matches!(
            first,
            ChannelMsg::ExitStatus {
                exit_status: PROOF_EXIT
            }
        ),
        "exit-status precedes close, got {first:?}"
    );
    let exit_reported_at = session.record.lock().exit_reported_at;

    let second = next_message(&mut channel)
        .await
        .expect("the channel closes after its exit report");
    assert!(
        matches!(second, ChannelMsg::Close),
        "close follows the exit report, got {second:?}"
    );

    // The literal criterion: the client observes exit-status that left
    // the server *after* the handler had already returned.
    let exit_reported_at = exit_reported_at.expect("the report timestamped");
    assert!(
        exit_reported_at > handler_returned_at,
        "the spawned task reported exit-status after the handler returned \
         (handler {handler_returned_at:?} vs report {exit_reported_at:?})"
    );

    session.unmount().await;
}