puressh 0.2.1

//! Sans-IO server transport driver.
//!
//! [`ServerDriver`] is the server-role counterpart of
//! [`ClientDriver`](super::ClientDriver): it owns the transport codec and the
//! KEX runner and sequences the server side of the SSH handshake (version
//! exchange, key exchange, re-key, EXT_INFO/PING handling), performing no I/O.
//! The frontend feeds inbound bytes, drains outbound frames, pulls
//! [`Event`]s, and ticks timers.
//!
//! As on the client, the driver does not own the connection multiplexer or the
//! authentication state machine: once the handshake completes it surfaces each
//! decoded payload as [`Event::AppData`]. The frontend runs `ServerAuth` and
//! its session/channel handlers over those payloads — so the existing blocking
//! server's handler model (and a future async one) layer on unchanged.
//!
//! The host keys and algorithm policy come from the server [`Config`], which
//! the driver holds behind an `Arc` to build KEXINIT adverts and select the
//! host key during KEX.

#![cfg(feature = "server")]

use alloc::collections::VecDeque;
use alloc::vec::Vec;
use std::sync::Arc;
use std::time::{Duration, Instant};

use purecrypto::rng::{CryptoRng, OsRng, RngCore};

use crate::error::{Error, Result};
use crate::hostkey::HostKey;
use crate::server::{Config, build_server_kexinit, pick_host_key, server_ext_info};
use crate::transport::ping::{SSH_MSG_PING, SSH_MSG_PONG, pong_for_ping};
use crate::transport::rekey::{RekeyPolicy, is_kex_msg, may_send_during_kex};
use crate::transport::{ExtInfo, KexRunner, PacketCodec, Role, VersionExchange};

use super::{
    Event, MAX_BANNER_LINE, MAX_DEFERRED_OUT_BYTES, MAX_INBOX_BYTES, SSH_MSG_EXT_INFO,
    SSH_MSG_KEXINIT, defer_inbound, generic_msg_during_initial_kex, keepalive_request,
};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Phase {
    /// Waiting for the peer's `SSH-2.0-…` identification line.
    AwaitingVersion,
    /// First key exchange in flight.
    Kex,
    /// Handshake done; post-NEWKEYS payloads surface as [`Event::AppData`].
    PostKex,
}

/// Sans-IO driver for the transport half of an SSH server connection.
///
/// See the [module docs](crate::driver) for the bytes-in / bytes-out / events
/// contract. Construct with [`ServerDriver::new`], kick off with
/// [`ServerDriver::start`], then pump.
pub struct ServerDriver {
    phase: Phase,
    codec: PacketCodec,
    runner: KexRunner,
    rng: OsRng,
    cfg: Arc<Config>,

    inbox: Vec<u8>,
    outbox: VecDeque<Vec<u8>>,
    events: VecDeque<Event>,
    /// Application packets received while a re-key was in flight (RFC 4253
    /// §7.3); replayed once NEWKEYS lands. Only ever filled *before* the
    /// peer's KEXINIT arrives — see [`Self::route_packet`].
    deferred: VecDeque<Vec<u8>>,
    /// Total payload bytes in `deferred`, capped at `MAX_DEFERRED_IN_BYTES`.
    deferred_bytes: usize,
    /// Egress held back because a key exchange is in flight (RFC 4253 §7.1).
    /// Stored as un-encoded payloads so the flush encrypts them under the
    /// *new* keys, which is what the peer expects after NEWKEYS.
    deferred_out: VecDeque<Vec<u8>>,
    deferred_out_bytes: usize,
    /// Application payloads accepted by [`Self::enqueue_payload`]; the
    /// connection loop reads it to tell whether a tick produced output.
    payloads_enqueued: u64,

    /// The peer's (client's) identification string.
    v_c: Vec<u8>,
    session_id: Vec<u8>,

    rekey_policy: RekeyPolicy,
    last_kex: Option<Instant>,

    keepalive: Option<(Duration, u32)>,
    last_activity: Instant,
    missed_keepalives: u32,
}

impl ServerDriver {
    /// Build a driver from the server [`Config`] (host keys + algorithm
    /// policy). Call [`ServerDriver::start`] before pumping to emit the
    /// version line and initial KEXINIT.
    pub fn new(cfg: Arc<Config>) -> Self {
        let mut rng = OsRng;
        let placeholder = build_server_kexinit(&mut rng, &cfg);
        Self {
            phase: Phase::AwaitingVersion,
            codec: PacketCodec::new(),
            runner: KexRunner::new(Role::Server, placeholder),
            rng,
            cfg,
            inbox: Vec::new(),
            outbox: VecDeque::new(),
            events: VecDeque::new(),
            deferred: VecDeque::new(),
            deferred_bytes: 0,
            deferred_out: VecDeque::new(),
            deferred_out_bytes: 0,
            payloads_enqueued: 0,
            v_c: Vec::new(),
            session_id: Vec::new(),
            rekey_policy: RekeyPolicy::default(),
            last_kex: None,
            keepalive: None,
            last_activity: Instant::now(),
            missed_keepalives: 0,
        }
    }

    /// Override the re-key thresholds (defaults to [`RekeyPolicy::default`]).
    pub fn set_rekey_policy(&mut self, policy: RekeyPolicy) {
        self.rekey_policy = policy;
    }

    /// Enable server keepalive (`ClientAliveInterval` / `ClientAliveCountMax`):
    /// after `interval` of inbound silence, [`handle_timeout`](Self::handle_timeout)
    /// emits a `keepalive@openssh.com` global request, failing the connection
    /// after `count_max` unanswered probes.
    pub fn set_keepalive(&mut self, interval: Duration, count_max: u32) {
        self.keepalive = Some((interval, count_max));
    }

    /// Emit the local version line and the initial KEXINIT (with the
    /// `ext-info-s` marker and queued outbound EXT_INFO). `now` seeds the
    /// keepalive clock. Call once, before pumping.
    pub fn start(&mut self, now: Instant) -> Result<()> {
        self.last_activity = now;
        self.outbox.push_back(VersionExchange::outgoing_bytes());
        let advert =
            build_server_kexinit(&mut self.rng, &self.cfg).with_ext_info_marker(Role::Server);
        self.runner = KexRunner::new(Role::Server, advert);
        // Queue our outbound EXT_INFO; the runner emits it after NEWKEYS iff
        // the client also advertised the marker.
        self.runner
            .set_outbound_ext_info(server_ext_info(&self.cfg));
        let initial = self.runner.start(&mut self.rng)?;
        for p in initial.outbound {
            self.enqueue_payload(&p)?;
        }
        Ok(())
    }

    // --- accessors ---

    /// The session identifier (KEX exchange hash `H`). Empty until handshake.
    pub fn session_id(&self) -> &[u8] {
        &self.session_id
    }

    /// The algorithms the most recent key exchange chose.
    pub fn negotiated(&self) -> Option<crate::transport::Negotiated> {
        self.runner.negotiated()
    }

    /// The peer's negotiated EXT_INFO, if any.
    pub fn peer_ext_info(&self) -> Option<&ExtInfo> {
        self.runner.peer_ext_info()
    }

    /// True once the handshake has completed (post-NEWKEYS).
    pub fn handshake_done(&self) -> bool {
        self.phase == Phase::PostKex
    }

    /// True while a key exchange is in flight.
    pub fn is_kexing(&self) -> bool {
        self.runner.is_kexing()
    }

    /// Notify the driver that userauth has just succeeded: activate
    /// `zlib@openssh.com` compression (RFC 4253 §6.2). The frontend calls this
    /// once its `ServerAuth` reports `Authenticated`, before the next packet.
    pub fn notify_auth_success(&mut self) {
        self.codec.activate_compress();
    }

    // --- pump surface ---

    /// Number of application payloads enqueued so far.
    pub fn payloads_enqueued(&self) -> u64 {
        self.payloads_enqueued
    }

    /// Encode `payload` with the current keys and queue it for transmission.
    pub fn enqueue_payload(&mut self, payload: &[u8]) -> Result<()> {
        self.payloads_enqueued = self.payloads_enqueued.wrapping_add(1);
        // RFC 4253 §7.1: between our KEXINIT and NEWKEYS only messages 1–49
        // may go out. Holding the rest back is not politeness — a strict peer
        // (OpenSSH 10.4+) aborts the connection over a single stray
        // CHANNEL_WINDOW_ADJUST emitted mid-re-key.
        let must_wait = match payload.first() {
            Some(&b) => !may_send_during_kex(b),
            None => false,
        };
        if must_wait && self.runner.is_kexing() {
            self.deferred_out_bytes += payload.len();
            if self.deferred_out_bytes > MAX_DEFERRED_OUT_BYTES {
                return Err(Error::Protocol("re-key: deferred egress buffer too large"));
            }
            self.deferred_out.push_back(payload.to_vec());
            return Ok(());
        }
        self.encode_and_queue(payload)
    }

    /// Encode `payload` under the current keys and queue the frame.
    fn encode_and_queue(&mut self, payload: &[u8]) -> Result<()> {
        let frame = self.codec.encode(payload, &mut self.rng)?;
        self.outbox.push_back(frame);
        Ok(())
    }

    /// Flush egress held back during a re-key, in the order it was queued.
    /// Runs after NEWKEYS, so each payload is encoded under the new keys.
    fn drain_deferred_out(&mut self) -> Result<()> {
        while let Some(payload) = self.deferred_out.pop_front() {
            self.encode_and_queue(&payload)?;
        }
        self.deferred_out_bytes = 0;
        Ok(())
    }

    /// Pop the next fully-encoded frame to write to the transport, if any.
    pub fn poll_transmit(&mut self) -> Option<Vec<u8>> {
        self.outbox.pop_front()
    }

    /// Pop the next high-level [`Event`], if any.
    pub fn poll_event(&mut self) -> Option<Event> {
        self.events.pop_front()
    }

    /// Feed inbound transport bytes; decode and route as many packets as are
    /// available. `now` advances the keepalive activity clock.
    pub fn handle_input(&mut self, bytes: &[u8], now: Instant) -> Result<()> {
        self.inbox.extend_from_slice(bytes);
        if self.inbox.len() > MAX_INBOX_BYTES {
            return Err(Error::Protocol("inbound buffer too large"));
        }
        if self.phase == Phase::AwaitingVersion && !self.scan_peer_version()? {
            return Ok(());
        }
        loop {
            match self.codec.decode(&self.inbox)? {
                Some((payload, consumed)) => {
                    self.inbox.drain(..consumed);
                    self.route_packet(&payload, now)?;
                }
                None => return Ok(()),
            }
        }
    }

    /// Drive re-key and keepalive timers. Returns `Err` if keepalive has gone
    /// unanswered past `count_max`.
    pub fn handle_timeout(&mut self, now: Instant) -> Result<()> {
        if self.runner.is_kexing() {
            return Ok(());
        }
        if let Some(last) = self.last_kex
            && self.rekey_policy.should_rekey(&self.codec, last, now)
        {
            self.initiate_rekey()?;
            return Ok(());
        }
        if let Some((interval, count_max)) = self.keepalive
            && now.duration_since(self.last_activity) >= interval
        {
            if self.missed_keepalives >= count_max {
                return Err(Error::Protocol("keepalive: no response from peer"));
            }
            let probe = keepalive_request();
            self.enqueue_payload(&probe)?;
            self.missed_keepalives += 1;
            self.last_activity = now;
        }
        Ok(())
    }

    /// When the frontend should next call [`handle_timeout`](Self::handle_timeout).
    pub fn next_timeout(&self) -> Option<Instant> {
        self.keepalive
            .map(|(interval, _)| self.last_activity + interval)
    }

    // --- internal routing ---

    fn scan_peer_version(&mut self) -> Result<bool> {
        let Some(pos) = self.inbox.iter().position(|&b| b == b'\n') else {
            if self.inbox.len() > MAX_BANNER_LINE {
                return Err(Error::Protocol("banner line too long"));
            }
            return Ok(false);
        };
        let line: Vec<u8> = self.inbox.drain(..=pos).collect();
        // Unlike a server, a client may send nothing before its
        // identification: anything else is refused, with OpenSSH's message.
        let parsed = line
            .starts_with(b"SSH-")
            .then(|| VersionExchange::parse_remote(&line).ok())
            .flatten();
        let Some(parsed) = parsed else {
            self.outbox
                .push_back(b"Invalid SSH identification string.\r\n".to_vec());
            return Err(Error::Protocol("invalid SSH identification string"));
        };
        self.v_c = parsed.into_bytes();
        self.phase = Phase::Kex;
        Ok(true)
    }

    fn route_packet(&mut self, payload: &[u8], now: Instant) -> Result<()> {
        self.note_activity(now);
        // Until the first NEWKEYS nothing on the wire is authenticated, so
        // only the KEX stream itself may flow. Anything the application
        // would act on is refused outright rather than buffered for replay
        // after the keys are in (which is what an on-path injector wants).
        let initial_kex = self.phase == Phase::Kex;
        match payload.first().copied() {
            Some(1) => Err(super::disconnect_error(payload)),
            // Generic transport messages during the initial KEX: fatal under
            // strict-kex, dropped otherwise. A PING gets no PONG here — the
            // peer has no business probing an unkeyed connection.
            Some(2) | Some(3) | Some(4) | Some(SSH_MSG_PING) | Some(SSH_MSG_PONG)
                if initial_kex =>
            {
                generic_msg_during_initial_kex(self.runner.strict_kex_enabled())
            }
            Some(2) | Some(3) | Some(4) => Ok(()),
            Some(SSH_MSG_PING) => {
                let pong = pong_for_ping(payload)?;
                self.enqueue_payload(&pong)
            }
            Some(SSH_MSG_PONG) => Ok(()),
            Some(SSH_MSG_EXT_INFO) => {
                if !self.runner.may_accept_ext_info() {
                    return Err(Error::Protocol("unexpected SSH_MSG_EXT_INFO"));
                }
                self.runner.handle_inbound_ext_info(payload)
            }
            Some(b) if is_kex_msg(b) => {
                if b == SSH_MSG_KEXINIT && !self.runner.is_kexing() {
                    self.initiate_rekey()?;
                }
                self.route_kex(payload)?;
                if self.runner.is_completed() {
                    if self.phase == Phase::Kex {
                        self.session_id = self
                            .runner
                            .session_id()
                            .ok_or(Error::Protocol("kex: missing session id"))?
                            .to_vec();
                        self.phase = Phase::PostKex;
                        self.events.push_back(Event::HandshakeComplete);
                    }
                    self.last_kex = Some(now);
                    self.drain_deferred_out()?;
                    self.drain_deferred()?;
                }
                Ok(())
            }
            _ => {
                if initial_kex {
                    return Err(Error::Protocol("non-KEX message before first NEWKEYS"));
                }
                if self.runner.is_kexing() {
                    // Re-key. The peer may still be sending application
                    // traffic only while it has not seen our KEXINIT — i.e.
                    // its own KEXINIT has not reached us yet. Once it has,
                    // RFC 4253 §7.1 forbids anything but KEX messages from
                    // it, so this packet is an injection or a broken peer.
                    if !self.runner.awaiting_peer_kexinit() {
                        return Err(Error::Protocol("non-KEX message during key re-exchange"));
                    }
                    return defer_inbound(&mut self.deferred, &mut self.deferred_bytes, payload);
                }
                self.runner.note_inbound_other();
                self.events.push_back(Event::AppData(payload.to_vec()));
                Ok(())
            }
        }
    }

    /// Feed one KEX packet into the runner, selecting the host key for the
    /// negotiated algorithm, and enqueue its output. (`v_s` is our own version;
    /// `v_c` is the client's, parsed during the banner scan.)
    fn route_kex(&mut self, payload: &[u8]) -> Result<()> {
        // Select the host key once the algorithms are negotiated.
        let hk_owned: Option<&(dyn HostKey + Send + Sync)> = match self.runner.negotiated() {
            Some(neg) => {
                let hk = pick_host_key(&self.cfg.host_keys, &neg.host_key);
                if hk.is_none() {
                    return Err(Error::Protocol("kex: no host key for negotiated algorithm"));
                }
                hk
            }
            None => None,
        };
        // An RSA key signs with the hash the negotiated name pins.
        let variant: Option<Box<dyn HostKey>> = match (hk_owned, self.runner.negotiated()) {
            (Some(k), Some(neg)) if k.algorithm() != neg.host_key => k.upgraded_for(&neg.host_key),
            _ => None,
        };
        let hk_ref: Option<&dyn HostKey> = match &variant {
            Some(v) => Some(v.as_ref()),
            None => hk_owned.map(|k| k as &dyn HostKey),
        };
        let v_c = self.v_c.clone();
        let v_s = crate::transport::version::LOCAL_VERSION.as_bytes().to_vec();
        let adv = self.runner.on_packet(
            &mut self.rng,
            &mut self.codec,
            payload,
            hk_ref,
            None,
            &v_c,
            &v_s,
        )?;
        for p in adv.outbound {
            self.enqueue_payload(&p)?;
        }
        Ok(())
    }

    /// Replay application packets buffered during a re-key, in arrival order.
    fn drain_deferred(&mut self) -> Result<()> {
        while !self.runner.is_kexing() {
            let Some(payload) = self.deferred.pop_front() else {
                break;
            };
            self.deferred_bytes = self.deferred_bytes.saturating_sub(payload.len());
            self.runner.note_inbound_other();
            self.events.push_back(Event::AppData(payload));
        }
        Ok(())
    }

    fn initiate_rekey(&mut self) -> Result<()> {
        let advert = build_server_kexinit(&mut self.rng, &self.cfg);
        let adv = self.runner.restart(&mut self.rng, advert)?;
        for p in adv.outbound {
            self.enqueue_payload(&p)?;
        }
        Ok(())
    }

    fn note_activity(&mut self, now: Instant) {
        self.last_activity = now;
        self.missed_keepalives = 0;
    }

    /// Test hook: the live transport codec, so a test can forge frames under
    /// the current keys (to play a misbehaving peer).
    #[cfg(test)]
    pub(crate) fn codec_mut(&mut self) -> &mut PacketCodec {
        &mut self.codec
    }

    /// Test hook: start a re-key right now, regardless of the policy.
    #[cfg(test)]
    pub(crate) fn force_rekey(&mut self) -> Result<()> {
        self.initiate_rekey()
    }
}

// `OsRng` is `CryptoRng + RngCore`; assert the bounds we rely on.
const _: fn() = || {
    fn _assert<T: CryptoRng + RngCore>() {}
    _assert::<OsRng>();
};

// End-to-end: act as the server *frontend* over a `ServerDriver` (handshake)
// plus a frontend-owned `ServerAuth` and `ConnectionState`, against the real
// blocking `Client` (connect → publickey auth → exec). This proves the
// server-side sans-IO engine symmetrically to the client driver test.
#[cfg(all(test, feature = "client"))]
mod tests {
    use super::*;
    use alloc::string::{String, ToString};
    use std::io::{Read as _, Write as _};
    use std::net::TcpListener;
    use std::sync::Arc;
    use std::thread;
    use std::time::Duration;

    use crate::auth::{AuthAttempt, AuthDecision, Authenticator, ServerAuth, ServerStep};
    use crate::channel::{ChannelEvent, ChannelOpen, ChannelRequest, ConnectionState};
    use crate::client::{Client, Config as ClientConfig};
    use crate::hostkey::{Ed25519HostKey, HostKey};
    use crate::server::{
        AuthenticatorFactory, CommandHandler, Config as ServerConfig, ExecResult, SessionEnv,
    };

    struct OneKeyAuth {
        user: String,
        blob: Vec<u8>,
    }
    impl Authenticator for OneKeyAuth {
        fn evaluate(&mut self, attempt: AuthAttempt) -> AuthDecision {
            match attempt {
                AuthAttempt::PublicKey {
                    user,
                    public_blob,
                    probe_only,
                    verified,
                    ..
                } => {
                    if user != self.user || public_blob != self.blob {
                        return AuthDecision::Reject;
                    }
                    if probe_only {
                        return AuthDecision::Accept;
                    }
                    if verified {
                        AuthDecision::Accept
                    } else {
                        AuthDecision::Reject
                    }
                }
                _ => AuthDecision::Reject,
            }
        }
    }

    struct UnusedHandler;
    impl CommandHandler for UnusedHandler {
        fn handle(&self, _u: &str, _e: &SessionEnv, _c: &str) -> ExecResult {
            ExecResult {
                stdout: Vec::new(),
                stderr: Vec::new(),
                exit_status: 0,
            }
        }
    }

    fn fresh_seed() -> [u8; 32] {
        let mut s = [0u8; 32];
        OsRng.fill_bytes(&mut s);
        s
    }

    #[test]
    fn server_driver_handshake_auth_exec_round_trip() {
        let host_seed = fresh_seed();
        let client_seed = fresh_seed();
        let client_key = Ed25519HostKey::from_seed(client_seed);
        let client_blob = client_key.public_blob();
        let user = "srv-driver-user".to_string();
        let reply = b"server driver says hi\n".to_vec();

        // Server config — only host keys + algorithms are used by the driver.
        let host_key: Box<dyn HostKey + Send + Sync> =
            Box::new(Ed25519HostKey::from_seed(host_seed));
        let factory: Arc<dyn AuthenticatorFactory> =
            Arc::new(|| Box::new(UnusedHandlerAuth) as Box<dyn Authenticator>);
        let cfg = Arc::new(ServerConfig::new(
            vec![host_key],
            factory,
            vec!["publickey"],
            Arc::new(UnusedHandler),
        ));

        let listener = TcpListener::bind("127.0.0.1:0").expect("bind");
        let addr = listener.local_addr().expect("addr");

        // Blocking client on a thread: connect, auth, exec, report output.
        let cu = user.clone();
        let client_thread = thread::spawn(move || -> ExecResultOut {
            let mut client =
                Client::connect(addr, ClientConfig::insecure()).expect("client connect");
            client
                .authenticate_publickey(&cu, Box::new(Ed25519HostKey::from_seed(client_seed)))
                .expect("auth");
            let out = client.exec("hi").expect("exec");
            ExecResultOut {
                stdout: out.stdout,
                exit: out.exit_status,
            }
        });

        // We are the server frontend.
        let (mut sock, _peer) = listener.accept().expect("accept");
        sock.set_read_timeout(Some(Duration::from_millis(50)))
            .unwrap();
        let mut driver = ServerDriver::new(cfg.clone());
        driver.start(Instant::now()).expect("start");

        let mut server_auth: Option<ServerAuth> = None;
        let mut conn = ConnectionState::new();
        let mut authed = false;
        let mut session_ch: Option<u32> = None;
        let mut responded = false;
        let mut remote_close = false;

        macro_rules! flush {
            () => {
                while let Some(frame) = driver.poll_transmit() {
                    sock.write_all(&frame).expect("write");
                }
            };
        }

        'pump: for _ in 0..200_000 {
            flush!();
            while let Some(ev) = driver.poll_event() {
                match ev {
                    Event::HandshakeComplete => {
                        let auth = ServerAuth::new(
                            driver.session_id().to_vec(),
                            vec!["publickey"],
                            Box::new(OneKeyAuth {
                                user: user.clone(),
                                blob: client_blob.clone(),
                            }),
                        );
                        server_auth = Some(auth);
                    }
                    Event::AppData(payload) if !authed => {
                        let a = server_auth.as_mut().unwrap();
                        match a.on_packet(&payload).expect("auth on_packet") {
                            ServerStep::Send(p) => driver.enqueue_payload(&p).expect("enq"),
                            ServerStep::Authenticated { payload, .. } => {
                                driver.enqueue_payload(&payload).expect("enq");
                                driver.notify_auth_success();
                                authed = true;
                            }
                            ServerStep::Disconnect(_) => panic!("auth disconnect"),
                        }
                    }
                    Event::AppData(payload) => {
                        match conn.on_packet(&payload).expect("on_packet") {
                            ChannelEvent::OpenRequest { channel, kind } => {
                                if matches!(kind, ChannelOpen::Session) {
                                    let p = conn.accept_open(channel).expect("accept");
                                    driver.enqueue_payload(&p).expect("enq");
                                    session_ch = Some(channel);
                                }
                            }
                            ChannelEvent::Request {
                                channel,
                                request,
                                want_reply,
                            } if Some(channel) == session_ch => {
                                if matches!(request, ChannelRequest::Exec { .. }) {
                                    if want_reply {
                                        let p = conn.send_request_success(channel).expect("succ");
                                        driver.enqueue_payload(&p).expect("enq");
                                    }
                                    // stdout
                                    let (p, _n) = conn.send_data(channel, &reply).expect("data");
                                    driver.enqueue_payload(&p).expect("enq");
                                    // exit-status 0
                                    let p = conn
                                        .send_request(
                                            channel,
                                            ChannelRequest::ExitStatus { code: 0 },
                                            false,
                                        )
                                        .expect("exit");
                                    driver.enqueue_payload(&p).expect("enq");
                                    // eof + close
                                    let p = conn.send_eof(channel).expect("eof");
                                    driver.enqueue_payload(&p).expect("enq");
                                    let p = conn.send_close(channel).expect("close");
                                    driver.enqueue_payload(&p).expect("enq");
                                    responded = true;
                                }
                            }
                            ChannelEvent::Close { channel } if Some(channel) == session_ch => {
                                remote_close = true;
                            }
                            _ => {}
                        }
                    }
                }
            }
            flush!();
            if responded && remote_close {
                break 'pump;
            }
            let mut buf = [0u8; 16 * 1024];
            match sock.read(&mut buf) {
                Ok(0) => break 'pump,
                Ok(n) => driver
                    .handle_input(&buf[..n], Instant::now())
                    .expect("input"),
                Err(e)
                    if e.kind() == std::io::ErrorKind::WouldBlock
                        || e.kind() == std::io::ErrorKind::TimedOut =>
                {
                    driver.handle_timeout(Instant::now()).expect("timeout");
                }
                Err(e) => panic!("read error: {e}"),
            }
        }

        let out = client_thread.join().expect("client thread");
        assert_eq!(
            out.stdout, reply,
            "exec stdout round-trips through the server driver"
        );
        assert_eq!(out.exit, Some(0));
        let _ = sock;
    }

    struct ExecResultOut {
        stdout: Vec<u8>,
        exit: Option<u32>,
    }

    // Placeholder authenticator for the server Config (the driver never uses
    // it — the test drives ServerAuth with OneKeyAuth directly).
    struct UnusedHandlerAuth;
    impl Authenticator for UnusedHandlerAuth {
        fn evaluate(&mut self, _a: AuthAttempt) -> AuthDecision {
            AuthDecision::Reject
        }
    }

    // --- initial key exchange: nothing but the KEX stream may flow ---

    /// A cleartext frame, as the wire looks before the first NEWKEYS.
    fn cleartext(payload: &[u8]) -> Vec<u8> {
        PacketCodec::new()
            .encode(payload, &mut OsRng)
            .expect("cleartext frame")
    }

    /// A client KEXINIT that agrees with the driver's default advert, with or
    /// without the client-side strict-kex marker.
    fn peer_client_kexinit(strict: bool) -> Vec<u8> {
        use crate::transport::KexInit;
        use crate::transport::kex::{KexAlgorithms, STRICT_KEX_CLIENT_MARKER};
        let kex: &[&str] = if strict {
            &["curve25519-sha256", STRICT_KEX_CLIENT_MARKER]
        } else {
            &["curve25519-sha256"]
        };
        let algs = KexAlgorithms {
            kex,
            server_host_key: &["ssh-ed25519"],
            ciphers_c2s: &["chacha20-poly1305@openssh.com"],
            ciphers_s2c: &["chacha20-poly1305@openssh.com"],
            macs_c2s: &["hmac-sha2-256"],
            macs_s2c: &["hmac-sha2-256"],
            comp_c2s: &["none"],
            comp_s2c: &["none"],
            lang_c2s: &[],
            lang_s2c: &[],
        };
        KexInit::from_algorithms(&algs, [7u8; 16]).encode()
    }

    /// A server driver that has sent its KEXINIT and consumed the peer's
    /// version line: the initial KEX is in flight, nothing decoded yet.
    fn server_in_initial_kex() -> ServerDriver {
        let host_key: Box<dyn HostKey + Send + Sync> =
            Box::new(Ed25519HostKey::from_seed(fresh_seed()));
        let factory: Arc<dyn AuthenticatorFactory> =
            Arc::new(|| Box::new(UnusedHandlerAuth) as Box<dyn Authenticator>);
        let cfg = Arc::new(ServerConfig::new(
            vec![host_key],
            factory,
            vec!["publickey"],
            Arc::new(UnusedHandler),
        ));
        let mut d = ServerDriver::new(cfg);
        d.start(Instant::now()).expect("start");
        while d.poll_transmit().is_some() {}
        d.handle_input(b"SSH-2.0-peer\r\n", Instant::now())
            .expect("version line");
        assert!(d.is_kexing() && !d.handshake_done());
        d
    }

    fn protocol_err(res: Result<()>) -> &'static str {
        match res {
            Err(Error::Protocol(msg)) => msg,
            other => panic!("expected a protocol error, got {other:?}"),
        }
    }

    /// Regression (T1/L1): an application-layer packet received during the
    /// initial (cleartext, unauthenticated) key exchange used to be queued
    /// and replayed to the application after NEWKEYS. It must be fatal,
    /// both before and after the peer's KEXINIT.
    #[test]
    fn initial_kex_rejects_application_packets() {
        // Before the peer's KEXINIT.
        let mut d = server_in_initial_kex();
        let err = protocol_err(d.handle_input(&cleartext(&[50, 0, 0, 0, 0]), Instant::now()));
        assert_eq!(err, "non-KEX message before first NEWKEYS");
        assert!(
            d.poll_event().is_none(),
            "nothing surfaced to the application"
        );

        // After the peer's KEXINIT (mid-exchange, awaiting ECDH_INIT).
        for strict in [true, false] {
            let mut d = server_in_initial_kex();
            d.handle_input(&cleartext(&peer_client_kexinit(strict)), Instant::now())
                .expect("peer kexinit");
            let userauth = [50, 0, 0, 0, 0];
            let err = protocol_err(d.handle_input(&cleartext(&userauth), Instant::now()));
            assert_eq!(err, "non-KEX message before first NEWKEYS");
            assert!(d.poll_event().is_none());
        }
    }

    /// Regression (T1/L1): `IGNORE`/`DEBUG`/`UNIMPLEMENTED`/`PING`/`PONG`
    /// during the initial KEX terminate the connection when strict-kex was
    /// negotiated (OpenSSH PROTOCOL §1.10) and are dropped otherwise.
    #[test]
    fn initial_kex_generic_messages_follow_strict_kex() {
        let ignore = [2u8, 0, 0, 0, 0];
        let unimplemented = [3u8, 0, 0, 0, 0];
        let ping = [SSH_MSG_PING, 0, 0, 0, 0];

        // Strict: fatal.
        for msg in [&ignore[..], &unimplemented[..], &ping[..]] {
            let mut d = server_in_initial_kex();
            d.handle_input(&cleartext(&peer_client_kexinit(true)), Instant::now())
                .expect("peer kexinit");
            let err = protocol_err(d.handle_input(&cleartext(msg), Instant::now()));
            assert_eq!(
                err,
                "strict-kex: unexpected message during initial key exchange"
            );
        }

        // Not strict: tolerated (and a PING draws no PONG), but application
        // packets are still refused.
        let mut d = server_in_initial_kex();
        d.handle_input(&cleartext(&peer_client_kexinit(false)), Instant::now())
            .expect("peer kexinit");
        while d.poll_transmit().is_some() {}
        for msg in [&ignore[..], &unimplemented[..], &ping[..]] {
            d.handle_input(&cleartext(msg), Instant::now())
                .expect("tolerated without strict-kex");
        }
        assert!(
            d.poll_transmit().is_none(),
            "no PONG before the connection is keyed"
        );
        assert!(d.poll_event().is_none());
        let err = protocol_err(d.handle_input(&cleartext(&[50, 0, 0, 0, 0]), Instant::now()));
        assert_eq!(err, "non-KEX message before first NEWKEYS");
    }

    /// Regression (T2, through the driver): a packet slipped in ahead of the
    /// peer's KEXINIT makes a strict-kex negotiation fail.
    #[test]
    fn initial_kex_packet_before_kexinit_fails_strict_negotiation() {
        let mut d = server_in_initial_kex();
        d.handle_input(&cleartext(&[2, 0, 0, 0, 0]), Instant::now())
            .expect("IGNORE before KEXINIT is not yet known to be strict");
        let err =
            protocol_err(d.handle_input(&cleartext(&peer_client_kexinit(true)), Instant::now()));
        assert_eq!(err, "strict-kex: KEXINIT was not the first packet");
    }
}