girt 0.3.4

//! A single async owner for nonblocking pipes and the unreaped SSH process group.
use std::future::Future;
use std::io::{Read, Write};
use std::os::fd::{AsFd, AsRawFd};
use std::os::unix::process::CommandExt;
use std::process::{Child, ChildStderr, ChildStdin, ChildStdout, Command, ExitStatus, Stdio};
use std::time::{Duration, Instant};

use rustix::fs::{OFlags, fcntl_getfl, fcntl_setfl};
use rustix::process::{Pid, Signal, WaitId, WaitIdOptions, kill_process_group, waitid};
use tokio::io::unix::AsyncFd;

use super::terminal::Lease;
use super::{ForegroundTerminal, SshError, TransportControl};

pub(crate) struct Session {
    // Declared first so cleanup kills the process before closing registered pipes.
    process: Process,
    input: Option<AsyncFd<ChildStdin>>,
    output: AsyncFd<ChildStdout>,
    diagnostics: AsyncFd<ChildStderr>,
    diagnostics_open: bool,
}
struct Process {
    child: Child,
    status: Option<ExitStatus>,
    terminal: Option<Lease>,
}
impl Process {
    fn pid(&self) -> Pid {
        Pid::from_raw(self.child.id() as i32).expect("positive child PID")
    }
    fn kill_group(&self) {
        let _ = kill_process_group(self.pid(), Signal::KILL);
    }

    // Every controlled pipe wait observes job-control state, not only the final exit wait.
    fn observe(&mut self) -> Result<bool, SshError> {
        if self.status.is_some() {
            return Ok(false);
        }
        let pid = self.pid();
        let mut options = WaitIdOptions::EXITED | WaitIdOptions::NOHANG | WaitIdOptions::NOWAIT;
        if self.terminal.is_some() {
            options |= WaitIdOptions::STOPPED | WaitIdOptions::CONTINUED;
        }
        let event = match waitid(WaitId::Pid(pid), options) {
            Ok(event) => event,
            Err(rustix::io::Errno::INTR) => None,
            Err(error) => return Err(std::io::Error::from(error).into()),
        };
        if let Some(event) = event {
            if event.stopped() || event.continued() {
                // Consume only job-control notifications. An exit racing this call remains
                // unreaped, reserving the leader PID until group cleanup and tty restoration.
                let options =
                    WaitIdOptions::STOPPED | WaitIdOptions::CONTINUED | WaitIdOptions::NOHANG;
                #[cfg(test)]
                if event.continued()
                    && terminal_tests::RAPID_STOP.swap(false, std::sync::atomic::Ordering::Relaxed)
                {
                    terminal_tests::replace_continued(pid);
                }
                let consumed = match waitid(WaitId::Pid(pid), options) {
                    Ok(event) => event,
                    Err(rustix::io::Errno::INTR) => None,
                    Err(error) => return Err(std::io::Error::from(error).into()),
                };
                if let Some(event) = consumed
                    && let Some(terminal) = &mut self.terminal
                {
                    if let Some(signal) = event.stopping_signal() {
                        terminal.stopped(pid, signal)?;
                    } else {
                        terminal.continued(pid)?;
                    }
                }
            } else {
                self.kill_group();
                let restored = self.terminal.as_mut().map_or(Ok(()), |terminal| {
                    terminal.restore(event.exit_status() == Some(0))
                });
                let status = self.child.wait()?;
                self.status = Some(status);
                // Preserve an unsuccessful child exit as the primary failure. Drop retries any
                // incomplete restoration before releasing the terminal lease.
                if status.success() {
                    restored?;
                }
                return Ok(false);
            }
        }
        self.terminal
            .as_mut()
            .map_or(Ok(false), |terminal| terminal.check(pid))
    }

    fn exit_result(&self) -> Result<(), SshError> {
        let status = self.status.expect("observed exit");
        if status.success() {
            Ok(())
        } else {
            Err(SshError::Exit(status.code()))
        }
    }
}
impl Drop for Process {
    fn drop(&mut self) {
        if self.status.is_none() {
            self.kill_group();
            let _ = self.child.kill();
            if let Some(terminal) = &mut self.terminal {
                let _ = terminal.restore(false);
            }
            let _ = self.child.wait();
        }
        // Release ownership only after group cleanup, terminal restoration, and leader reaping.
        self.terminal.take();
    }
}
impl Session {
    #[cfg(test)]
    pub(crate) fn spawn(command: &mut Command) -> Result<Self, SshError> {
        Self::spawn_with_terminal(command, None)
    }

    pub(crate) fn spawn_with_terminal(
        command: &mut Command,
        terminal: Option<&ForegroundTerminal>,
    ) -> Result<Self, SshError> {
        let terminal = terminal.map(ForegroundTerminal::acquire).transpose()?;
        if let Some(terminal) = &terminal {
            terminal.prepare_command(command)?;
        }
        let child = command
            .process_group(0)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()?;
        let mut process = Process {
            child,
            status: None,
            terminal,
        };
        let input = register(process.child.stdin.take().expect("piped stdin"))?;
        let output = register(process.child.stdout.take().expect("piped stdout"))?;
        let diagnostics = register(process.child.stderr.take().expect("piped stderr"))?;
        Ok(Self {
            process,
            input: Some(input),
            output,
            diagnostics,
            diagnostics_open: true,
        })
    }

    #[cfg(test)]
    pub(crate) async fn advertise(
        &mut self,
        limit: usize,
        control: TransportControl<'_>,
    ) -> Result<Vec<u8>, SshError> {
        self.advertise_with_diagnostics(limit, control, &mut |_| {})
            .await
    }

    pub(crate) async fn advertise_with_diagnostics(
        &mut self,
        limit: usize,
        control: TransportControl<'_>,
        diagnostics: &mut impl FnMut(&[u8]),
    ) -> Result<Vec<u8>, SshError> {
        let result = controlled(
            advertisement(&mut self.output, limit),
            &mut self.process,
            &mut self.diagnostics,
            &mut self.diagnostics_open,
            control,
            diagnostics,
        )
        .await;
        // EOF before an advertisement often means host-key/authentication failure. Observe the
        // process under the same deadline so it is reported as transport failure, not Git refusal.
        if matches!(result, Err(SshError::Protocol("truncated advertisement"))) {
            self.input.take();
            self.finish(control, diagnostics).await?;
        }
        result
    }

    #[cfg(test)]
    pub(crate) async fn exchange(
        &mut self,
        request: &[u8],
        pack: &[u8],
        limit: usize,
        control: TransportControl<'_>,
    ) -> (Vec<u8>, Result<(), SshError>, usize) {
        self.exchange_with_diagnostics(request, pack, limit, control, &mut |_| {})
            .await
    }

    pub(crate) async fn exchange_with_diagnostics(
        &mut self,
        request: &[u8],
        pack: &[u8],
        limit: usize,
        control: TransportControl<'_>,
        diagnostics: &mut impl FnMut(&[u8]),
    ) -> (Vec<u8>, Result<(), SshError>, usize) {
        self.exchange_observed(
            request,
            crate::fetch::PackInput::Slice(pack),
            limit,
            control,
            diagnostics,
            &mut |_| {},
        )
        .await
    }

    // Observes the retained prefix after each bounded read; observers neither own nor copy it.
    pub(crate) async fn exchange_observed(
        &mut self,
        request: &[u8],
        pack: crate::fetch::PackInput<'_>,
        limit: usize,
        control: TransportControl<'_>,
        diagnostics: &mut impl FnMut(&[u8]),
        observe: &mut impl FnMut(&[u8]),
    ) -> (Vec<u8>, Result<(), SshError>, usize) {
        let mut body = Vec::new();
        let (result, written) = self
            .exchange_streaming_pack(request, pack, limit, control, diagnostics, &mut |chunk| {
                body.extend_from_slice(chunk);
                observe(&body);
                Ok(())
            })
            .await;
        (body, result, written)
    }

    pub(crate) async fn exchange_streaming(
        &mut self,
        request: &[u8],
        pack: &[u8],
        limit: usize,
        control: TransportControl<'_>,
        diagnostics: &mut impl FnMut(&[u8]),
        receive: &mut impl FnMut(&[u8]) -> Result<(), SshError>,
    ) -> (Result<(), SshError>, usize) {
        self.exchange_streaming_pack(
            request,
            crate::fetch::PackInput::Slice(pack),
            limit,
            control,
            diagnostics,
            receive,
        )
        .await
    }

    async fn exchange_streaming_pack(
        &mut self,
        request: &[u8],
        pack: crate::fetch::PackInput<'_>,
        limit: usize,
        control: TransportControl<'_>,
        diagnostics: &mut impl FnMut(&[u8]),
        receive: &mut impl FnMut(&[u8]) -> Result<(), SshError>,
    ) -> (Result<(), SshError>, usize) {
        let mut written = 0;
        let input = self.input.take().expect("one request per session");
        let exchange = async {
            // Drain response concurrently with upload: a rejecting server can fill stdout while
            // no longer consuming input. Retain received status bytes even when writing fails.
            tokio::try_join!(
                write_request(input, request, pack, &mut written),
                read_body(&mut self.output, limit, receive)
            )?;
            Ok(())
        };
        let mut result = controlled(
            exchange,
            &mut self.process,
            &mut self.diagnostics,
            &mut self.diagnostics_open,
            control,
            diagnostics,
        )
        .await;
        if result.is_ok() {
            result = self.finish(control, diagnostics).await;
        }
        (result, written)
    }

    async fn finish(
        &mut self,
        control: TransportControl<'_>,
        diagnostics: &mut impl FnMut(&[u8]),
    ) -> Result<(), SshError> {
        let result = loop {
            if let Err(error) = self.process.observe() {
                break Err(error);
            }
            if self.process.status.is_some() {
                break self.process.exit_result();
            }
            let tick = async {
                tokio::time::sleep(Duration::from_millis(20)).await;
                Ok(())
            };
            if let Err(error) = controlled(
                tick,
                &mut self.process,
                &mut self.diagnostics,
                &mut self.diagnostics_open,
                control,
                diagnostics,
            )
            .await
            {
                break Err(error);
            }
        };
        // Exit observation kills owned descendants before reaping and draining final diagnostics.
        if self.process.status.is_some() {
            let drained = drain_diagnostics(
                &mut self.diagnostics,
                &mut self.diagnostics_open,
                control,
                diagnostics,
            )
            .await;
            return result.and(drained);
        }
        result
    }
}

fn register<T: AsFd + AsRawFd>(pipe: T) -> Result<AsyncFd<T>, SshError> {
    let flags = fcntl_getfl(&pipe).map_err(std::io::Error::from)?;
    fcntl_setfl(&pipe, flags | OFlags::NONBLOCK).map_err(std::io::Error::from)?;
    Ok(AsyncFd::new(pipe)?)
}
async fn read<T: Read + AsRawFd>(
    pipe: &mut AsyncFd<T>,
    bytes: &mut [u8],
) -> Result<usize, SshError> {
    loop {
        let mut ready = pipe.readable_mut().await?;
        match ready.try_io(|fd| fd.get_mut().read(bytes)) {
            Ok(result) => return result.map_err(Into::into),
            Err(_) => continue,
        }
    }
}
async fn exact<T: Read + AsRawFd>(
    pipe: &mut AsyncFd<T>,
    mut bytes: &mut [u8],
) -> Result<(), SshError> {
    while !bytes.is_empty() {
        let count = read(pipe, bytes).await?;
        if count == 0 {
            return Err(SshError::Protocol("truncated advertisement"));
        }
        bytes = &mut bytes[count..];
    }
    Ok(())
}
async fn advertisement(
    output: &mut AsyncFd<ChildStdout>,
    limit: usize,
) -> Result<Vec<u8>, SshError> {
    let mut bytes = Vec::new();
    loop {
        if limit.saturating_sub(bytes.len()) < 4 {
            return Err(SshError::Limit);
        }
        let mut header = [0; 4];
        exact(output, &mut header).await?;
        if !header.iter().all(u8::is_ascii_hexdigit) {
            return Err(SshError::Protocol("pkt-line header"));
        }
        let length = usize::from_str_radix(std::str::from_utf8(&header).unwrap(), 16).unwrap();
        bytes.extend_from_slice(&header);
        if length == 0 {
            return Ok(bytes);
        }
        if !(4..=65520).contains(&length) {
            return Err(SshError::Protocol("pkt-line length"));
        }
        let payload = length - 4;
        if payload > limit.saturating_sub(bytes.len()) {
            return Err(SshError::Limit);
        }
        let offset = bytes.len();
        bytes.resize(offset + payload, 0);
        exact(output, &mut bytes[offset..]).await?;
        tokio::task::yield_now().await;
    }
}
async fn write_request(
    mut input: AsyncFd<ChildStdin>,
    request: &[u8],
    pack: crate::fetch::PackInput<'_>,
    written: &mut usize,
) -> Result<(), SshError> {
    use tokio::io::AsyncReadExt as _;
    // Open before sending commands so an unreadable spool cannot cause a partial upload.
    let mut reader = pack.async_reader()?;
    write_bytes(&mut input, request, written).await?;
    let mut buffer = vec![0; 64 * 1024];
    loop {
        let n = reader.read(&mut buffer).await?;
        if n == 0 {
            break;
        }
        write_bytes(&mut input, &buffer[..n], written).await?;
    }
    Ok(())
}

async fn write_bytes(
    input: &mut AsyncFd<ChildStdin>,
    mut bytes: &[u8],
    written: &mut usize,
) -> Result<(), SshError> {
    while !bytes.is_empty() {
        let mut ready = input.writable_mut().await?;
        if let Ok(result) = ready.try_io(|fd| fd.get_mut().write(bytes)) {
            let count = result?;
            if count == 0 {
                return Err(SshError::Io(std::io::ErrorKind::WriteZero));
            }
            *written += count;
            bytes = &bytes[count..];
            tokio::task::yield_now().await;
        }
    }
    Ok(())
}
async fn read_body(
    output: &mut AsyncFd<ChildStdout>,
    limit: usize,
    receive: &mut impl FnMut(&[u8]) -> Result<(), SshError>,
) -> Result<(), SshError> {
    let mut bytes = [0; 8192];
    let mut total = 0;
    loop {
        let count = read(output, &mut bytes).await?;
        if count == 0 {
            return Ok(());
        }
        let keep = count.min(limit.saturating_sub(total));
        receive(&bytes[..keep])?;
        total += keep;
        if keep != count {
            return Err(SshError::Limit);
        }
        // AsyncFd readiness can remain immediately ready for a prolific peer. Yield explicitly
        // so the outer cancellation/deadline/diagnostic loop and other runtime tasks get polled.
        tokio::task::yield_now().await;
    }
}
async fn controlled<T>(
    future: impl Future<Output = Result<T, SshError>>,
    process: &mut Process,
    diagnostics: &mut AsyncFd<ChildStderr>,
    diagnostics_open: &mut bool,
    control: TransportControl<'_>,
    diagnostic_callback: &mut impl FnMut(&[u8]),
) -> Result<T, SshError> {
    tokio::pin!(future);
    let mut bytes = [0; 8192];
    loop {
        control.check()?;
        let suspended = if process.terminal.is_some() {
            process.observe()?
        } else {
            false
        };
        let interval = control.deadline.map_or(Duration::from_millis(20), |end| {
            end.saturating_duration_since(Instant::now())
                .min(Duration::from_millis(20))
        });
        tokio::select! {
            biased;
            _ = tokio::time::sleep(interval) => {},
            result = &mut future, if !suspended => return result,
            result = read(diagnostics, &mut bytes), if *diagnostics_open => {
                let count = result?;
                if count == 0 { *diagnostics_open = false; }
                else { diagnostic_callback(&bytes[..count]); }
                tokio::task::yield_now().await;
            },
        }
    }
}

async fn drain_diagnostics(
    diagnostics: &mut AsyncFd<ChildStderr>,
    open: &mut bool,
    control: TransportControl<'_>,
    callback: &mut impl FnMut(&[u8]),
) -> Result<(), SshError> {
    let mut bytes = [0; 8192];
    while *open {
        control.check()?;
        tokio::select! {
            biased;
            _ = tokio::time::sleep(Duration::from_millis(20)) => {},
            result = read(diagnostics, &mut bytes) => {
                let count = result?;
                if count == 0 { *open = false; }
                else { callback(&bytes[..count]); }
                tokio::task::yield_now().await;
            }
        }
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use std::sync::atomic::AtomicBool;

    use super::*;

    fn runtime() -> tokio::runtime::Runtime {
        tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .unwrap()
    }
    fn assert_reaped(pid: Pid) {
        assert!(matches!(
            waitid(
                WaitId::Pid(pid),
                WaitIdOptions::EXITED | WaitIdOptions::NOHANG
            ),
            Err(rustix::io::Errno::CHILD)
        ));
    }
    #[test]
    fn dropping_pending_future_kills_and_reaps_the_child() {
        let rt = runtime();
        let cancel = AtomicBool::new(false);
        let pid = rt.block_on(async {
            let mut session =
                Session::spawn(Command::new("/bin/sh").args(["-c", "sleep 10"])).unwrap();
            let pid = session.process.pid();
            let operation =
                async move { session.advertise(100, TransportControl::new(&cancel)).await };
            assert!(
                tokio::time::timeout(Duration::from_millis(50), operation)
                    .await
                    .is_err()
            );
            pid
        });
        assert_reaped(pid);
    }
    #[test]
    fn deadline_cleans_up_child_with_blocked_stdout_and_stderr() {
        let rt = runtime();
        let cancel = AtomicBool::new(false);
        let pid = rt.block_on(async {
            let mut command = Command::new("python3");
            command.args(["-c", "import os,time\nos.write(1,b'0000')\nwhile True: os.write(2,b'x'*8192); os.write(1,b'x'*8192)"]);
            let mut session = Session::spawn(&mut command).unwrap();
            let pid = session.process.pid();
            // Interpreter startup is outside the short blocked-pipe deadline.
            let startup = TransportControl { cancel: &cancel, deadline: Some(Instant::now()+Duration::from_secs(5)) };
            session.advertise(4, startup).await.unwrap();
            let control = TransportControl { cancel: &cancel, deadline: Some(Instant::now()+Duration::from_millis(200)) };
            let (_, result, _) = session.exchange(b"0000", &[], usize::MAX, control).await;
            assert!(matches!(result, Err(SshError::Deadline)));
            pid
        });
        assert_reaped(pid);
    }
    #[test]
    fn successful_exit_kills_pipe_holding_descendants_only() {
        let rt = runtime();
        let cancel = AtomicBool::new(false);
        let mut unrelated = Command::new("sleep").arg("10").spawn().unwrap();
        rt.block_on(async {
            let mut session = Session::spawn(
                Command::new("/bin/sh").args(["-c", "sleep 10 & printf 0000; exit 0"]),
            )
            .unwrap();
            let control = TransportControl {
                cancel: &cancel,
                deadline: Some(Instant::now() + Duration::from_secs(2)),
            };
            session.advertise(4, control).await.unwrap();
            session.finish(control, &mut |_| {}).await.unwrap();
            let count =
                tokio::time::timeout(Duration::from_secs(2), read(&mut session.output, &mut [0]))
                    .await
                    .unwrap()
                    .unwrap();
            assert_eq!(count, 0);
        });
        assert!(unrelated.try_wait().unwrap().is_none());
        unrelated.kill().unwrap();
        unrelated.wait().unwrap();
    }
    #[test]
    fn already_cancelled_exchange_attempts_no_write() {
        let rt = runtime();
        rt.block_on(async {
            let mut session =
                Session::spawn(Command::new("/bin/sh").args(["-c", "sleep 10"])).unwrap();
            let cancel = AtomicBool::new(true);
            let (_, result, written) = session
                .exchange(b"commands", &[], 100, TransportControl::new(&cancel))
                .await;
            assert_eq!(written, 0);
            assert!(matches!(result, Err(SshError::Cancelled)));
        });
    }
}

#[cfg(test)]
#[path = "terminal_tests.rs"]
mod terminal_tests;