alkssh/tests/common/tty_doubles.rs

//! The `session`-seam test doubles (`w3-tty-session`): the alkssh-side
//! **TtyBackend shim double**: alktty's `MockBackend`/`testing.rs` are
//! `pub(crate)` + `cfg(test)` and never importable, so the suite
//! hand-rolls the shim per the `TtyHandle` field shapes: a backend
//! that records the `TtyParams` it was allocated, consumes stdin
//! (recording bytes, the EOF shutdown and the abort kill), emits
//! scripted stdout/stderr chunks and a scripted exit code, and
//! surfaces a control handle recording resizes/signals. Plus the
//! and the ownership witness double: the identity that reaches
//! `drive_session_pre_negotiated`'s ADR-050 ownership check is the
//! identity ride's live witness (the task's "identity-landing item").
//!
//! Every stream/future here is test-scripted and instantly-completed
//! or never-completing; no real process, no real PTY (the real-tool
//! smoke composes alktty's own `local` backend instead, dev-dep only).

use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use alkcall::core::auth::Identity;
use alkcall::core::ownership::OwnershipProvider;
use alktty::backend::{TtyBackend, TtyControl, TtyError, TtyHandle, TtyParams};
use async_trait::async_trait;
use bytes::Bytes;
use tokio::io::AsyncWrite;
use tokio_stream::wrappers::ReceiverStream;

/// The shim's shared observation surface; every row's assertions
/// come off this (`Record TtyParams`, consume stdin/EOF/resizes/
/// signals, the abort kill witness).
#[derive(Debug, Default)]
pub struct TtyShimObservation {
    /// The `TtyParams` of every allocation, in order.
    pub allocations: parking_lot::Mutex<Vec<TtyParams>>,
    /// Bytes written to the shim's stdin, concatenated per allocation.
    pub stdin: parking_lot::Mutex<Vec<u8>>,
    /// Whether a stdin shutdown (the EOF translation's destination)
    /// was observed (one flag per shutdown, but one allocation per
    /// row in practice).
    pub stdin_eof: parking_lot::Mutex<Vec<bool>>,
    /// Resizes from the control handle (`cols`, `rows`).
    pub resizes: parking_lot::Mutex<Vec<(u16, u16)>>,
    /// Signals from the control handle, the name strings verbatim.
    pub signals: parking_lot::Mutex<Vec<String>>,
    /// The abort kill witness: pushed once per `exit_code` future
    /// dropped *unresolved*; alktty's kill-on-drop contract
    /// (ADR-056) observed through the shim.
    pub kills: Arc<parking_lot::Mutex<Vec<bool>>>,
}

/// The shim's scripted per-allocation behavior: what the shim emits
/// and what exit code it resolves with.
#[derive(Debug, Clone, Default)]
pub struct ShimScript {
    /// stdout chunks, in emission order, then end-of-stream.
    pub stdout: Vec<Vec<u8>>,
    /// stderr chunks, in emission order, then end-of-stream (the shim
    /// carries the stderr stream only when these are non-empty *and*
    /// the allocation is pipe-mode).
    pub stderr: Vec<Vec<u8>>,
    /// The exit code: `Ok(code)` resolves instantly; `None` never
    /// resolves (the cancel row's never-ending session).
    pub exit: Option<i32>,
}

/// The TtyBackend shim double, a hand-rolled minimal backend per the
/// `TtyHandle` field shapes (see the module doc). Clone-able (the
/// observation + script are `Arc`s).
#[derive(Clone, Default)]
pub struct TtyShim {
    observation: Arc<TtyShimObservation>,
    script: Arc<parking_lot::Mutex<ShimScript>>,
}

impl TtyShim {
    /// A shim with an empty script (allocates, pumps nothing, resolves
    /// exit-status 0 by default).
    pub fn new() -> Self {
        Self::default()
    }

    /// Sets the scripted behavior for subsequent allocations.
    pub fn set_script(&self, script: ShimScript) {
        *self.script.lock() = script;
    }

    /// The observation surface.
    pub fn observation(&self) -> &Arc<TtyShimObservation> {
        &self.observation
    }
}

#[async_trait]
impl TtyBackend for TtyShim {
    /// Records the params verbatim and hands back the scripted
    /// `TtyHandle`.
    async fn allocate(&self, params: &TtyParams) -> Result<TtyHandle, TtyError> {
        self.observation.allocations.lock().push(params.clone());
        let script = self.script.lock().clone();

        let stdin = ShimStdin {
            observation: Arc::clone(&self.observation),
        };

        let stdout = scripted_stream(&script.stdout);
        let stderr = if script.terminal_is_pty(params) {
            None
        } else {
            Some(scripted_stream(&script.stderr))
        };

        let exit_code: alktty::backend::BoxFuture<Result<i32, TtyError>> = Box::pin(GuardedExit {
            inner: ShimExit { code: script.exit },
            resolved: false,
            kills: Arc::clone(&self.observation.kills),
        });

        let control = Some(alktty::backend::TtyControlHandle::new(Arc::new(
            ShimControl {
                observation: Arc::clone(&self.observation),
            },
        )));

        Ok(TtyHandle {
            stdin: Box::new(stdin),
            stdout,
            stderr,
            exit_code,
            control,
        })
    }

    /// The ADR-050 check input (every allocation targets this shim
    /// resource, so the ownership witness row (the identity surface)
    /// always evaluates.
    fn resource_id(&self, _params: &TtyParams) -> Option<(&'static str, String)> {
        Some(("shim", "demo".to_owned()))
    }
}

impl ShimScript {
    /// Whether the allocation is PTY-mode (the TtyParams's terminal is
    /// not the NegotiateRequest.tty this module speaks; the adapter's
    /// `From<NegotiateRequest>` conversion is what the seam consumes;
    /// the shim sees the converted `TtyParams`).
    fn terminal_is_pty(&self, params: &TtyParams) -> bool {
        params.terminal.is_some()
    }
}

/// A scripted stream: the chunks in order, then end. Instantly
/// drained (the alktty pumps consume; the exit chunk arrives after,
/// since the adapter's exit-chunk-is-last invariant orders everything).
fn scripted_stream(chunks: &[Vec<u8>]) -> Pin<Box<dyn tokio_stream::Stream<Item = Bytes> + Send>> {
    let (tx, rx) = tokio::sync::mpsc::channel(chunks.len().max(1) + 1);
    for chunk in chunks {
        let _ = tx.try_send(Bytes::copy_from_slice(chunk));
    }
    drop(tx);
    Box::pin(ReceiverStream::new(rx))
}

/// The shim's stdin, an `AsyncWrite` recording bytes and the EOF
/// shutdown (`poll_shutdown` = the zero-length sentinel's effect at
/// the backend).
struct ShimStdin {
    observation: Arc<TtyShimObservation>,
}

impl AsyncWrite for ShimStdin {
    fn poll_write(
        self: Pin<&mut Self>,
        _cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<std::io::Result<usize>> {
        self.observation.stdin.lock().extend_from_slice(buf);
        Poll::Ready(Ok(buf.len()))
    }

    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        Poll::Ready(Ok(()))
    }

    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        self.observation.stdin_eof.lock().push(true);
        Poll::Ready(Ok(()))
    }
}

/// The scripted exit-code future: `Some(code)` resolves instantly,
/// `None` pends forever (the cancel row's never-ending session).
struct ShimExit {
    code: Option<i32>,
}

impl Future for ShimExit {
    type Output = Result<i32, TtyError>;

    fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
        match self.code {
            Some(code) => Poll::Ready(Ok(code)),
            None => Poll::Pending,
        }
    }
}

/// Wrap [`ShimExit`] with alktty's ADR-056 kill-on-drop contract
/// witness: an unresolved drop records the kill; a resolved drop (the
/// future completed) must not.
struct GuardedExit {
    inner: ShimExit,
    resolved: bool,
    kills: Arc<parking_lot::Mutex<Vec<bool>>>,
}

impl Unpin for GuardedExit {}

impl Future for GuardedExit {
    type Output = Result<i32, TtyError>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let me = self.get_mut();
        let out = Pin::new(&mut me.inner).poll(cx);
        if matches!(&out, Poll::Ready(Ok(_))) {
            me.resolved = true;
        }
        out
    }
}

impl Drop for GuardedExit {
    fn drop(&mut self) {
        if !self.resolved {
            self.kills.lock().push(true);
        }
    }
}

/// The shim's control surface; it records resizes and signals verbatim.
struct ShimControl {
    observation: Arc<TtyShimObservation>,
}

impl TtyControl for ShimControl {
    fn resize(&self, cols: u16, rows: u16, _pixel_width: u16, _pixel_height: u16) {
        self.observation.resizes.lock().push((cols, rows));
    }

    fn signal(&self, name: &str) {
        self.observation.signals.lock().push(name.to_owned());
    }
}

/// A shim-backed `SessionSeam` in the `session`-seam shape the
/// adapters take: the `backends` map holds the shim under `"shim"`, the
/// selector names it, ownership unset (the ADR-050 check skips).
pub fn tty_seam(shim: &Arc<TtyShim>) -> alkssh::SessionSeam {
    let mut backends = HashMap::new();
    backends.insert("shim".to_owned(), Arc::clone(shim) as Arc<dyn TtyBackend>);
    alkssh::SessionSeam {
        backends: Arc::new(backends),
        backend: "shim".to_owned(),
        ownership: None,
    }
}

/// The ownership witness double ([`alkcall::core::ownership::
/// OwnershipProvider`]): records every identity evaluation and rules
/// by its configured verdict (the deny shape produces the negotiation
/// error frame on the wire, the bridge's `negotiation-failed`
/// decline).
#[derive(Debug)]
pub struct MemOwnership {
    pub owns: bool,
    /// The identity ids that reached `owns`, the identity ride's
    /// live witness.
    pub ids: parking_lot::Mutex<Vec<String>>,
}

impl MemOwnership {
    pub fn new(owns: bool) -> Self {
        Self {
            owns,
            ids: parking_lot::Mutex::new(Vec::new()),
        }
    }
}

impl OwnershipProvider for MemOwnership {
    fn owns(
        &self,
        identity: &Identity,
        _resource_type: &str,
        _resource_id: &str,
        _action: &str,
    ) -> bool {
        self.ids.lock().push(identity.id.clone());
        self.owns
    }

    fn owned_resources(&self, _identity: &Identity, _resource_type: &str) -> Vec<String> {
        Vec::new()
    }

    fn owns_any(&self, _identity: &Identity, _resource_type: &str) -> bool {
        self.owns
    }
}

/// A shim-backed seam with the ownership provider present, the
/// witness rows' builder.
pub fn tty_seam_with_ownership(
    shim: &Arc<TtyShim>,
    ownership: Arc<MemOwnership>,
) -> alkssh::SessionSeam {
    let mut backends = HashMap::new();
    backends.insert("shim".to_owned(), Arc::clone(shim) as Arc<dyn TtyBackend>);
    alkssh::SessionSeam {
        backends: Arc::new(backends),
        backend: "shim".to_owned(),
        ownership: Some(ownership),
    }
}

/// A seam whose selector names no backend in its own map; it renders the
/// `unknown_backend` negotiation error frame, driven from config.
pub fn misconfigured_seam() -> alkssh::SessionSeam {
    alkssh::SessionSeam {
        backends: Arc::new(HashMap::new()),
        backend: "absent".to_owned(),
        ownership: None,
    }
}