1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
use std::{
    io,
    process::{Child, Command},
};

#[cfg(unix)]
use std::{
    io::IsTerminal,
    sync::{
        atomic::{AtomicU32, Ordering},
        Arc,
    },
};

#[cfg(unix)]
pub use foreground_pgroup::stdin_fd;

/// A simple wrapper for [`std::process::Child`]
///
/// It can only be created by [`ForegroundChild::spawn`].
///
/// # Spawn behavior
/// ## Unix
///
/// For interactive shells, the spawned child process will get its own process group id,
/// and it will be put in the foreground (by making stdin belong to the child's process group).
/// On drop, the calling process's group will become the foreground process group once again.
///
/// For non-interactive mode, processes are spawned normally without any foreground process handling.
///
/// ## Other systems
///
/// It does nothing special on non-unix systems, so `spawn` is the same as [`std::process::Command::spawn`].
pub struct ForegroundChild {
    inner: Child,
    #[cfg(unix)]
    pipeline_state: Option<Arc<(AtomicU32, AtomicU32)>>,
}

impl ForegroundChild {
    #[cfg(not(unix))]
    pub fn spawn(mut command: Command) -> io::Result<Self> {
        command.spawn().map(|child| Self { inner: child })
    }

    #[cfg(unix)]
    pub fn spawn(
        mut command: Command,
        interactive: bool,
        pipeline_state: &Arc<(AtomicU32, AtomicU32)>,
    ) -> io::Result<Self> {
        if interactive && io::stdin().is_terminal() {
            let (pgrp, pcnt) = pipeline_state.as_ref();
            let existing_pgrp = pgrp.load(Ordering::SeqCst);
            foreground_pgroup::prepare_command(&mut command, existing_pgrp);
            command
                .spawn()
                .map(|child| {
                    foreground_pgroup::set(&child, existing_pgrp);
                    let _ = pcnt.fetch_add(1, Ordering::SeqCst);
                    if existing_pgrp == 0 {
                        pgrp.store(child.id(), Ordering::SeqCst);
                    }
                    Self {
                        inner: child,
                        pipeline_state: Some(pipeline_state.clone()),
                    }
                })
                .map_err(|e| {
                    foreground_pgroup::reset();
                    e
                })
        } else {
            command.spawn().map(|child| Self {
                inner: child,
                pipeline_state: None,
            })
        }
    }
}

impl AsMut<Child> for ForegroundChild {
    fn as_mut(&mut self) -> &mut Child {
        &mut self.inner
    }
}

#[cfg(unix)]
impl Drop for ForegroundChild {
    fn drop(&mut self) {
        if let Some((pgrp, pcnt)) = self.pipeline_state.as_deref() {
            if pcnt.fetch_sub(1, Ordering::SeqCst) == 1 {
                pgrp.store(0, Ordering::SeqCst);
                foreground_pgroup::reset()
            }
        }
    }
}

// It's a simpler version of fish shell's external process handling.
#[cfg(unix)]
mod foreground_pgroup {
    use nix::{
        sys::signal::{sigaction, SaFlags, SigAction, SigHandler, SigSet, Signal},
        unistd::{self, Pid},
    };
    use std::{
        os::{
            fd::{AsFd, BorrowedFd},
            unix::prelude::CommandExt,
        },
        process::{Child, Command},
    };

    /// Alternative to having to call `std::io::stdin()` just to get the file descriptor of stdin
    ///
    /// # Safety
    /// I/O safety of reading from `STDIN_FILENO` unclear.
    ///
    /// Currently only intended to access `tcsetpgrp` and `tcgetpgrp` with the I/O safe `nix`
    /// interface.
    pub unsafe fn stdin_fd() -> impl AsFd {
        unsafe { BorrowedFd::borrow_raw(nix::libc::STDIN_FILENO) }
    }

    pub fn prepare_command(external_command: &mut Command, existing_pgrp: u32) {
        unsafe {
            // Safety:
            // POSIX only allows async-signal-safe functions to be called.
            // `sigaction` and `getpid` are async-signal-safe according to:
            // https://manpages.ubuntu.com/manpages/bionic/man7/signal-safety.7.html
            // Also, `set_foreground_pid` is async-signal-safe.
            external_command.pre_exec(move || {
                // When this callback is run, std::process has already:
                // - reset SIGPIPE to SIG_DFL

                // According to glibc's job control manual:
                // https://www.gnu.org/software/libc/manual/html_node/Launching-Jobs.html
                // This has to be done *both* in the parent and here in the child due to race conditions.
                set_foreground_pid(Pid::this(), existing_pgrp);

                // Reset signal handlers for child, sync with `terminal.rs`
                let default = SigAction::new(SigHandler::SigDfl, SaFlags::empty(), SigSet::empty());
                // SIGINT has special handling
                let _ = sigaction(Signal::SIGQUIT, &default);
                // We don't support background jobs, so keep some signals blocked for now
                // let _ = sigaction(Signal::SIGTSTP, &default);
                // let _ = sigaction(Signal::SIGTTIN, &default);
                // let _ = sigaction(Signal::SIGTTOU, &default);
                let _ = sigaction(Signal::SIGTERM, &default);

                Ok(())
            });
        }
    }

    pub fn set(process: &Child, existing_pgrp: u32) {
        set_foreground_pid(Pid::from_raw(process.id() as i32), existing_pgrp);
    }

    fn set_foreground_pid(pid: Pid, existing_pgrp: u32) {
        // Safety: needs to be async-signal-safe.
        // `setpgid` and `tcsetpgrp` are async-signal-safe.

        // `existing_pgrp` is 0 when we don't have an existing foreground process in the pipeline.
        // A pgrp of 0 means the calling process's pid for `setpgid`. But not for `tcsetpgrp`.
        let pgrp = if existing_pgrp == 0 {
            pid
        } else {
            Pid::from_raw(existing_pgrp as i32)
        };
        let _ = unistd::setpgid(pid, pgrp);
        let _ = unistd::tcsetpgrp(unsafe { stdin_fd() }, pgrp);
    }

    /// Reset the foreground process group to the shell
    pub fn reset() {
        if let Err(e) = unistd::tcsetpgrp(unsafe { stdin_fd() }, unistd::getpgrp()) {
            eprintln!("ERROR: reset foreground id failed, tcsetpgrp result: {e:?}");
        }
    }
}