test_fork_core/
fork.rs

1// Copyright (C) 2025-2026 Daniel Mueller <deso@posteo.net>
2// SPDX-License-Identifier: (Apache-2.0 OR MIT)
3
4//-
5// Copyright 2018 Jason Lingle
6//
7// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
8// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
9// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
10// option. This file may not be copied, modified, or distributed
11// except according to those terms.
12
13use std::env;
14use std::io::Read;
15use std::io::Write as _;
16use std::net::TcpListener;
17use std::net::TcpStream;
18use std::panic;
19use std::process;
20use std::process::Child;
21use std::process::Command;
22use std::process::ExitCode;
23use std::process::Stdio;
24use std::process::Termination;
25
26use crate::cmdline;
27use crate::error::Result;
28
29
30const OCCURS_ENV: &str = "TEST_FORK_OCCURS";
31const OCCURS_TERM_LENGTH: usize = 17; /* ':' plus 16 hexits */
32
33
34fn supervise_child(child: Child) {
35    let output = child.wait_with_output().expect("failed to wait for child");
36    assert!(
37        output.status.success(),
38        "child exited unsuccessfully with {}",
39        output.status,
40    );
41
42    // Make sure to forward output we captured to our own output, using
43    // print! and eprint! macros, which hook into the test output
44    // capture mechanism, to mimic default behavior.
45
46    if !output.stdout.is_empty() {
47        let s = String::from_utf8_lossy(&output.stdout);
48        print!("{s}");
49    }
50    if !output.stderr.is_empty() {
51        let s = String::from_utf8_lossy(&output.stderr);
52        eprint!("{s}");
53    }
54}
55
56
57/// Simulate a process fork.
58///
59/// Since this is not a true process fork, the calling code must be structured
60/// to ensure that the child process, upon starting from the same entry point,
61/// also reaches this same `fork()` call. Recursive forks are supported; the
62/// child branch is taken from all child processes of the fork even if it is
63/// not directly the child of a particular branch. However, encountering the
64/// same fork point more than once in a single execution sequence of a child
65/// process is not (e.g., putting this call in a recursive function) and
66/// results in unspecified behaviour.
67///
68/// `fork_id` is a unique identifier identifying this particular fork location.
69/// This *must* be stable across processes of the same executable; pointers are
70/// not suitable stable, and string constants may not be suitably unique. The
71/// [`fork_id!()`] macro is the recommended way to supply this
72/// parameter.
73///
74/// `test_name` must exactly match the full path of the test function being
75/// run.
76///
77/// If `test` panics, the child process exits with a failure code immediately
78/// rather than let the panic propagate out of the `fork()` call.
79///
80/// ## Panics
81///
82/// Panics if the environment indicates that there are already at least 16
83/// levels of fork nesting.
84///
85/// Panics if `std::env::current_exe()` fails to determine the path to
86/// the current executable.
87///
88/// Panics if any argument to the current process is not valid UTF-8.
89pub fn fork<F, T>(fork_id: &str, test_name: &str, test: F) -> Result<()>
90where
91    // NB: We use `Fn` here, because `FnMut` and `FnOnce` would allow
92    //     for modification of captured variables, but that will not
93    //     work across process boundaries.
94    F: Fn() -> T,
95    T: Termination,
96{
97    fn no_configure_child(_child: &mut Command) {}
98
99    fork_int(
100        test_name,
101        fork_id,
102        no_configure_child,
103        supervise_child,
104        test,
105    )
106}
107
108/// Simulate a process fork.
109///
110/// This function is similar to [`fork`], except that it allows for data
111/// exchange with the child process.
112#[expect(clippy::panic_in_result_fn, clippy::unwrap_in_result)]
113pub fn fork_in_out<F, T>(fork_id: &str, test_name: &str, test: F, data: &mut [u8]) -> Result<()>
114where
115    F: Fn(&mut [u8]) -> T,
116    T: Termination,
117{
118    let listener = TcpListener::bind("127.0.0.1:0").expect("failed to bind TCP socket");
119    let addr = listener.local_addr().unwrap();
120    let data_len = data.len();
121
122    fork_int(
123        test_name,
124        fork_id,
125        |cmd| {
126            cmd.env(fork_id, addr.to_string());
127        },
128        |child| {
129            let (mut stream, _addr) = listener
130                .accept()
131                .expect("failed to listen for child connection");
132            let () = stream
133                .write_all(data)
134                .expect("failed to send data to child");
135            let () = stream
136                .read_exact(data)
137                .expect("failed to receive data from child");
138            supervise_child(child)
139        },
140        || {
141            let addr = env::var(fork_id).unwrap_or_else(|err| {
142                panic!("failed to retrieve {fork_id} environment variable: {err}")
143            });
144            let mut stream =
145                TcpStream::connect(addr).expect("failed to establish connection with parent");
146
147            let mut data = Vec::with_capacity(data_len);
148            // SAFETY: The `Vec` contains `data_len` `u8` values, which
149            //         are valid for any bit pattern, so we can safely
150            //         adjust the length.
151            let () = unsafe { data.set_len(data_len) };
152
153            let () = stream
154                .read_exact(&mut data)
155                .expect("failed to receive data from parent");
156            let status = test(&mut data);
157            let () = stream
158                .write_all(&data)
159                .expect("failed to send data to parent");
160            status
161        },
162    )
163}
164
165pub(crate) fn fork_int<M, P, C, R, T>(
166    test_name: &str,
167    fork_id: &str,
168    process_modifier: M,
169    in_parent: P,
170    in_child: C,
171) -> Result<R>
172where
173    M: FnOnce(&mut process::Command),
174    P: FnOnce(Child) -> R,
175    T: Termination,
176    C: FnOnce() -> T,
177{
178    // Erase the generics so we don't instantiate the actual implementation for
179    // every single test
180    let mut return_value = None;
181    let mut process_modifier = Some(process_modifier);
182    let mut in_parent = Some(in_parent);
183    let mut in_child = Some(in_child);
184
185    fork_impl(
186        test_name,
187        fork_id,
188        &mut |cmd| process_modifier.take().unwrap()(cmd),
189        &mut |child| return_value = Some(in_parent.take().unwrap()(child)),
190        &mut || in_child.take().unwrap()(),
191    )
192    .map(|()| return_value.unwrap())
193}
194
195#[expect(clippy::panic_in_result_fn, clippy::unwrap_in_result)]
196fn fork_impl<T: Termination>(
197    test_name: &str,
198    fork_id: &str,
199    process_modifier: &mut dyn FnMut(&mut process::Command),
200    in_parent: &mut dyn FnMut(Child),
201    in_child: &mut dyn FnMut() -> T,
202) -> Result<()> {
203    let mut occurs = env::var(OCCURS_ENV).unwrap_or_else(|_| String::new());
204    if occurs.contains(fork_id) {
205        match panic::catch_unwind(panic::AssertUnwindSafe(in_child)) {
206            Ok(test_result) => {
207                let rc = if test_result.report() == ExitCode::SUCCESS {
208                    0
209                } else {
210                    70
211                };
212                process::exit(rc)
213            }
214            // Assume that the default panic handler already printed something
215            //
216            // We don't use process::abort() since it produces core dumps on
217            // some systems and isn't something more special than a normal
218            // panic.
219            Err(_) => process::exit(70 /* EX_SOFTWARE */),
220        }
221    } else {
222        // Prevent misconfiguration creating a fork bomb
223        if occurs.len() > 16 * OCCURS_TERM_LENGTH {
224            panic!("test-fork: Not forking due to >=16 levels of recursion");
225        }
226
227        occurs.push_str(fork_id);
228        let mut command =
229            process::Command::new(env::current_exe().expect("current_exe() failed, cannot fork"));
230        command
231            .args(cmdline::strip_cmdline(env::args())?)
232            .args(cmdline::RUN_TEST_ARGS)
233            .arg(test_name)
234            .env(OCCURS_ENV, &occurs)
235            .stdin(Stdio::null())
236            .stdout(Stdio::piped())
237            .stderr(Stdio::piped());
238        process_modifier(&mut command);
239
240        let child = command.spawn()?;
241        let () = in_parent(child);
242
243        Ok(())
244    }
245}
246
247
248#[cfg(test)]
249mod test {
250    use super::*;
251
252
253    fn wait_for_child_output(child: Child) -> String {
254        let output = child.wait_with_output().expect("failed to wait for child");
255        assert!(output.status.success());
256        let output = String::from_utf8(output.stdout).unwrap();
257        output
258    }
259
260    #[test]
261    fn fork_basically_works() {
262        fork_int(
263            "fork::test::fork_basically_works",
264            fork_id!(),
265            |_| (),
266            supervise_child,
267            || println!("hello from child"),
268        )
269        .unwrap()
270    }
271
272    #[test]
273    fn child_output_captured_and_repeated() {
274        let output = fork_int(
275            "fork::test::child_output_captured_and_repeated",
276            fork_id!(),
277            |_| (),
278            wait_for_child_output,
279            || {
280                fork_int(
281                    "fork::test::child_output_captured_and_repeated",
282                    fork_id!(),
283                    |_| (),
284                    supervise_child,
285                    || println!("hello from child"),
286                )
287                .unwrap()
288            },
289        )
290        .unwrap();
291        assert!(output.contains("hello from child"));
292    }
293
294    #[test]
295    fn child_aborted_if_panics() {
296        let status = fork_int::<_, _, _, _, ()>(
297            "fork::test::child_aborted_if_panics",
298            fork_id!(),
299            |_| (),
300            |mut child| child.wait().unwrap(),
301            || panic!("testing a panic, nothing to see here"),
302        )
303        .unwrap();
304        assert_eq!(70, status.code().unwrap());
305    }
306
307    /// Check that we can exchange data with the child process.
308    #[test]
309    fn data_exchange() {
310        let mut data = [1, 2, 3, 4, 5];
311
312        let () = fork_in_out(
313            fork_id!(),
314            "fork::test::data_exchange",
315            |data| {
316                assert_eq!(data.len(), 5);
317                let () = data.iter_mut().for_each(|x| *x += 1);
318            },
319            data.as_mut_slice(),
320        )
321        .unwrap();
322
323        assert_eq!(data, [2, 3, 4, 5, 6]);
324    }
325}
test_fork_core/fork.rs

test_fork_core/
fork.rs
