puressh 0.1.7

//! `ProxyCommand` transport: run the SSH session over a spawned helper
//! process's stdio instead of a direct socket.
//!
//! `ssh_config(5)`'s `ProxyCommand` names a shell command whose stdin/stdout
//! become the byte transport to the server (e.g. `nc -X connect -x
//! proxy:3128 %h %p`). puressh spawns it under `/bin/sh -c` with piped
//! stdin/stdout and inherited stderr, then drives a [`Client`] over the
//! pipes via [`Client::connect_via`].
//!
//! [`Client`]: crate::client::Client
//! [`Client::connect_via`]: crate::client::Client::connect_via

#![cfg(all(unix, feature = "client"))]

use std::io::{self, Read, Write};
use std::os::fd::AsFd;
use std::process::{Child, ChildStdin, ChildStdout, Command, Stdio};
use std::time::{Duration, Instant};

use nix::fcntl::{FcntlArg, OFlag, fcntl};
use nix::poll::{PollFd, PollFlags, PollTimeout, poll};

use crate::client::Transport;
use crate::error::{Error, Result};

/// Is `host` safe to substitute into a shell command line?
///
/// Mirrors OpenSSH's `valid_hostname()` (the CVE-2023-51385 fix): only
/// ASCII alphanumerics, `.`, `-` and `_`, plus the IPv6-literal
/// characters `:`, `[`, `]` and `%` (zone id), are allowed; the name must
/// be non-empty and must not start with `-` (so it can never be parsed as
/// an option by the helper it is handed to).
pub fn valid_hostname(host: &str) -> bool {
    if host.is_empty() || host.starts_with('-') {
        return false;
    }
    host.bytes().all(|b| {
        b.is_ascii_alphanumeric() || matches!(b, b'.' | b'-' | b'_' | b':' | b'[' | b']' | b'%')
    })
}

/// Is `user` safe to substitute into a shell command line?
///
/// Mirrors OpenSSH's `valid_ruser()` (a strict superset of its deny-list):
/// rejects a leading `-`, any whitespace or control character, and the
/// shell metacharacters `` ' ` \ " ; & | ( ) { } < > $ ``. Any other
/// character is accepted, as OpenSSH allows locale-specific login names.
pub fn valid_ruser(user: &str) -> bool {
    if user.starts_with('-') {
        return false;
    }
    !user.chars().any(|c| {
        c.is_whitespace()
            || c.is_control()
            || matches!(
                c,
                '\'' | '`' | '\\' | '"' | ';' | '&' | '|' | '(' | ')' | '{' | '}' | '<' | '>' | '$'
            )
    })
}

/// [`expand_tokens`] for strings that will be handed to `/bin/sh -c`
/// (`ProxyCommand`): the `%h` host and `%r` user are validated with
/// [`valid_hostname`] / [`valid_ruser`] before being spliced in, so a
/// hostname such as `` `id` `` or `-oProxyCommand=...` coming from an
/// untrusted source (a `.git/config`, a URL, a `ProxyJump` chain) cannot
/// inject shell syntax or options into the helper command. Only tokens
/// that actually appear in `s` are checked. Returns [`Error::Config`]
/// on an invalid value.
pub fn expand_tokens_checked(s: &str, host: &str, port: u16, user: &str) -> Result<String> {
    if uses_token(s, 'h') && !valid_hostname(host) {
        return Err(Error::Config(
            "ProxyCommand: hostname contains characters that are not permitted \
             in a shell command",
        ));
    }
    if uses_token(s, 'r') && !valid_ruser(user) {
        return Err(Error::Config(
            "ProxyCommand: remote user name contains characters that are not \
             permitted in a shell command",
        ));
    }
    Ok(expand_tokens(s, host, port, user))
}

/// Does `s` contain the `%<tok>` sequence (with `%%` escaping honoured)?
fn uses_token(s: &str, tok: char) -> bool {
    let mut chars = s.chars();
    while let Some(c) = chars.next() {
        if c == '%' && chars.next() == Some(tok) {
            return true;
        }
    }
    false
}

/// Expand OpenSSH-style `%`-tokens in a `ProxyCommand` / `ProxyJump` host
/// string.
///
/// Supported tokens:
/// - `%h` → target host
/// - `%p` → target port
/// - `%r` → remote (login) user
/// - `%%` → a literal `%`
///
/// Any other `%x` sequence is passed through verbatim (the `%` and the
/// following char are both kept), matching OpenSSH's lenient handling of
/// unknown tokens in these particular directives.
///
/// This performs **no validation** of `host` / `user` and is suitable for
/// values that are not interpreted by a shell (e.g. `ControlPath`). For a
/// string destined for `/bin/sh -c` use [`expand_tokens_checked`].
pub fn expand_tokens(s: &str, host: &str, port: u16, user: &str) -> String {
    let mut out = String::with_capacity(s.len());
    let mut chars = s.chars();
    while let Some(c) = chars.next() {
        if c != '%' {
            out.push(c);
            continue;
        }
        match chars.next() {
            Some('h') => out.push_str(host),
            Some('p') => {
                use core::fmt::Write as _;
                let _ = write!(out, "{port}");
            }
            Some('r') => out.push_str(user),
            Some('%') => out.push('%'),
            Some(other) => {
                out.push('%');
                out.push(other);
            }
            None => out.push('%'),
        }
    }
    out
}

/// A [`Transport`] backed by a spawned `ProxyCommand` helper process.
///
/// Writes go to the child's stdin; reads come from its stdout. The child's
/// stderr is inherited so any diagnostics surface on the terminal. The
/// child is killed (and reaped) on drop.
///
/// # Read timeouts over a pipe
///
/// A plain anonymous pipe has no `SO_RCVTIMEO`-style knob, so the serve /
/// forwarding poll loops (which expect [`Transport::set_read_timeout`] to
/// produce a periodic [`io::ErrorKind::WouldBlock`] tick) historically could
/// not run over a `ProxyCommand`. We give them one anyway: the child's stdout
/// fd is switched to `O_NONBLOCK`, and [`Transport::set_read_timeout`] records
/// a deadline. With a timeout set, [`Read::read`] uses `poll(2)` to wait up to
/// the deadline for readability and translates the empty-poll case into
/// `WouldBlock` — exactly the signal the serve loop already handles for a
/// `TcpStream` read timeout. With no timeout the read blocks (via an infinite
/// `poll`) so the blocking KEX / userauth phases behave as before.
pub struct ProcTransport {
    child: Child,
    stdin: ChildStdin,
    stdout: ChildStdout,
    /// Current read timeout. `None` ⇒ block until data (or EOF). `Some(d)` ⇒
    /// each `read` waits at most `d` for readability, then returns
    /// `WouldBlock`. The stdout fd is always `O_NONBLOCK`; this field only
    /// decides how long `poll(2)` waits before giving up.
    read_timeout: Option<Duration>,
}

impl ProcTransport {
    /// Spawn `command` under `/bin/sh -c` with piped stdin/stdout and
    /// inherited stderr. `command` should already have its `%`-tokens
    /// expanded (see [`expand_tokens`]). Returns an error if the helper
    /// cannot be spawned — there is intentionally **no** fallback to a
    /// direct connection.
    pub fn spawn(command: &str) -> io::Result<Self> {
        let mut child = Command::new("/bin/sh")
            .arg("-c")
            .arg(command)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::inherit())
            .spawn()?;
        let stdin = child
            .stdin
            .take()
            .ok_or_else(|| io::Error::other("ProxyCommand: child stdin not captured"))?;
        let stdout = child
            .stdout
            .take()
            .ok_or_else(|| io::Error::other("ProxyCommand: child stdout not captured"))?;
        // Put the read side into non-blocking mode up front: every read goes
        // through `poll(2)` + a non-blocking `read`, regardless of whether a
        // timeout is currently armed (a `None` timeout polls with no deadline).
        set_nonblocking(&stdout)?;
        Ok(Self {
            child,
            stdin,
            stdout,
            read_timeout: None,
        })
    }
}

/// Flip `O_NONBLOCK` on for `fd` via `fcntl(F_GETFL)` + `fcntl(F_SETFL)`.
fn set_nonblocking<F: AsFd>(fd: &F) -> io::Result<()> {
    let borrowed = fd.as_fd();
    let cur = fcntl(borrowed, FcntlArg::F_GETFL).map_err(io::Error::from)?;
    let flags = OFlag::from_bits_truncate(cur) | OFlag::O_NONBLOCK;
    fcntl(borrowed, FcntlArg::F_SETFL(flags)).map_err(io::Error::from)?;
    Ok(())
}

impl Read for ProcTransport {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        // The fd is O_NONBLOCK, so a bare `read` returns `WouldBlock` instead
        // of parking. We wrap it in `poll(2)` to get the wait semantics the
        // serve loop expects:
        //   * `read_timeout = None`  → poll with no deadline (block for data).
        //   * `read_timeout = Some`  → poll up to the deadline, surfacing
        //     `WouldBlock` if nothing arrives so the caller can tick.
        let deadline = self.read_timeout.map(|d| Instant::now() + d);
        loop {
            // Compute the remaining poll budget for this iteration.
            let timeout: PollTimeout = match deadline {
                None => PollTimeout::NONE,
                Some(end) => {
                    let now = Instant::now();
                    if now >= end {
                        return Err(io::Error::from(io::ErrorKind::WouldBlock));
                    }
                    let remaining = end - now;
                    // Clamp to u16 ms (PollTimeout's largest finite value is
                    // ~65s; our timeouts are tens of ms, so saturation is fine).
                    let ms = remaining.as_millis().min(u16::MAX as u128) as u16;
                    PollTimeout::from(ms)
                }
            };

            let mut fds = [PollFd::new(self.stdout.as_fd(), PollFlags::POLLIN)];
            match poll(&mut fds, timeout) {
                Ok(0) => {
                    // poll timed out without readiness → no data yet.
                    return Err(io::Error::from(io::ErrorKind::WouldBlock));
                }
                Ok(_) => {
                    // Readable (or hangup): attempt the read. On a real pipe a
                    // POLLHUP fd reads as EOF (0), which we surface as Ok(0).
                    match self.stdout.read(buf) {
                        // Spurious wakeup (poll said readable but the read
                        // raced to empty) — loop and re-poll within budget.
                        Err(e) if e.kind() == io::ErrorKind::WouldBlock => continue,
                        other => return other,
                    }
                }
                Err(nix::errno::Errno::EINTR) => continue,
                Err(e) => return Err(io::Error::from(e)),
            }
        }
    }
}

impl Write for ProcTransport {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.stdin.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.stdin.flush()
    }
}

impl Transport for ProcTransport {
    /// Arm (or clear) a read timeout. The child stdout fd is already
    /// `O_NONBLOCK`; this just records the deadline [`Read::read`] enforces
    /// via `poll(2)`. `Some(d)` makes reads return [`io::ErrorKind::WouldBlock`]
    /// after `d` of no data (the periodic tick the serve / forwarding loops
    /// rely on); `None` reverts to a blocking read.
    fn set_read_timeout(&mut self, t: Option<Duration>) -> io::Result<()> {
        self.read_timeout = t;
        Ok(())
    }
}

impl Drop for ProcTransport {
    fn drop(&mut self) {
        // Best-effort: kill the helper and reap it so we don't leave a
        // zombie. Ignore errors — the child may have already exited.
        let _ = self.child.kill();
        let _ = self.child.wait();
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn expand_basic_tokens() {
        assert_eq!(
            expand_tokens("nc %h %p", "example.com", 2222, "alice"),
            "nc example.com 2222"
        );
        assert_eq!(expand_tokens("login=%r", "h", 22, "bob"), "login=bob");
        assert_eq!(expand_tokens("100%%done", "h", 22, "u"), "100%done");
    }

    #[test]
    fn valid_hostname_matches_openssh_rules() {
        for ok in [
            "example.com",
            "10.0.0.1",
            "[::1]",
            "::1",
            "fe80::1%eth0",
            "host_1",
            "a-b.c",
            "LOCALHOST",
        ] {
            assert!(valid_hostname(ok), "{ok:?} should be accepted");
        }
        for bad in [
            "",
            "`id`",
            "$(id)",
            "a;b",
            "-oProxyCommand=x",
            "-",
            "a b",
            "a|b",
            "a&b",
            "a'b",
            "a\"b",
            "a\\b",
            "a/b",
            "a\nb",
            "a\tb",
            "hôst",
            "a>b",
            "a<b",
            "a(b)",
        ] {
            assert!(!valid_hostname(bad), "{bad:?} should be rejected");
        }
    }

    #[test]
    fn valid_ruser_matches_openssh_rules() {
        for ok in ["alice", "bob.smith", "user_1", "u-1", "", "李雷", "a@b"] {
            assert!(valid_ruser(ok), "{ok:?} should be accepted");
        }
        for bad in [
            "-bob", "a b", "a;b", "$(id)", "`id`", "a|b", "a&b", "a'b", "a\"b", "a\\b", "a(b",
            "a)b", "a<b", "a>b", "a\nb", "a\x1bb",
        ] {
            assert!(!valid_ruser(bad), "{bad:?} should be rejected");
        }
    }

    #[test]
    fn expand_tokens_checked_rejects_shell_metacharacters() {
        // Injection via the host or user reaching a `%h` / `%r` token fails
        // closed with a config error rather than reaching `/bin/sh -c`.
        for host in ["`id`", "$(id)", "a;b", "-oProxyCommand=x"] {
            let err = expand_tokens_checked("nc %h %p", host, 22, "u")
                .err()
                .unwrap_or_else(|| panic!("{host:?} must be rejected"));
            assert!(matches!(err, Error::Config(_)));
        }
        for user in ["-x", "a;b", "$(id)"] {
            assert!(expand_tokens_checked("ssh -l %r jump", "h", 22, user).is_err());
        }
        // Valid values expand exactly like the unchecked variant.
        for host in ["example.com", "10.0.0.1", "[::1]", "host_1"] {
            assert_eq!(
                expand_tokens_checked("nc %h %p", host, 2222, "u").unwrap(),
                format!("nc {host} 2222")
            );
        }
        // Values are only checked when the corresponding token is used —
        // a command that never interpolates the host cannot be injected
        // through it — and `%%h` is an escaped literal, not a token.
        assert_eq!(
            expand_tokens_checked("nc proxy 3128", "`id`", 22, "$(id)").unwrap(),
            "nc proxy 3128"
        );
        assert_eq!(
            expand_tokens_checked("echo 100%%h", "`id`", 22, "u").unwrap(),
            "echo 100%h"
        );
    }

    #[test]
    fn expand_unknown_token_passthrough() {
        // %z isn't a recognised token → kept verbatim.
        assert_eq!(expand_tokens("a%zb", "h", 22, "u"), "a%zb");
        // trailing lone % is kept.
        assert_eq!(expand_tokens("end%", "h", 22, "u"), "end%");
    }

    #[test]
    fn spawn_failure_is_strict_error() {
        // A command that cannot possibly run still *spawns* /bin/sh (sh
        // exists), so to exercise the strict no-fallback path we instead
        // verify that a command exiting immediately yields EOF on read
        // rather than silently succeeding. The genuine spawn-failure path
        // (e.g. /bin/sh missing) is exercised by the binary's error
        // propagation; here we just confirm the transport surfaces the
        // child's premature exit as a clean EOF.
        let mut t = ProcTransport::spawn("exit 0").expect("sh -c spawns");
        let mut buf = [0u8; 16];
        let n = t.read(&mut buf).expect("read after child exit");
        assert_eq!(n, 0, "closed child stdout should read as EOF");
    }

    #[test]
    fn round_trip_through_cat() {
        // `cat` echoes stdin to stdout — a trivial loopback transport.
        let mut t = ProcTransport::spawn("cat").expect("spawn cat");
        t.write_all(b"hello pipe").expect("write");
        t.flush().expect("flush");
        let mut buf = [0u8; 10];
        t.read_exact(&mut buf).expect("read echo");
        assert_eq!(&buf, b"hello pipe");
    }

    #[test]
    fn read_timeout_ticks_as_wouldblock() {
        // `sleep` produces no output, so a read with a short timeout must
        // surface WouldBlock (the serve-loop tick) rather than blocking.
        let mut t = ProcTransport::spawn("sleep 5").expect("spawn sleep");
        t.set_read_timeout(Some(Duration::from_millis(50)))
            .expect("set timeout");
        let mut buf = [0u8; 16];
        let start = Instant::now();
        let err = t.read(&mut buf).expect_err("should time out");
        assert_eq!(err.kind(), io::ErrorKind::WouldBlock);
        // It should have actually waited ~the timeout, not busy-returned.
        assert!(
            start.elapsed() >= Duration::from_millis(40),
            "poll should have waited out the deadline, took {:?}",
            start.elapsed()
        );
    }

    #[test]
    fn read_timeout_still_delivers_data() {
        // With a timeout armed, data that *does* arrive must still be read.
        let mut t = ProcTransport::spawn("printf abc; sleep 5").expect("spawn");
        t.set_read_timeout(Some(Duration::from_millis(500)))
            .expect("set timeout");
        let mut buf = [0u8; 3];
        t.read_exact(&mut buf).expect("read the printf output");
        assert_eq!(&buf, b"abc");
    }

    #[test]
    fn blocking_read_after_clearing_timeout() {
        // `set_read_timeout(None)` reverts to a blocking poll: a delayed
        // write is still delivered (no premature WouldBlock).
        let mut t = ProcTransport::spawn("sleep 0.1; printf ok").expect("spawn");
        t.set_read_timeout(None).expect("clear timeout");
        let mut buf = [0u8; 2];
        t.read_exact(&mut buf)
            .expect("blocking read waits for data");
        assert_eq!(&buf, b"ok");
    }
}