sqz_engine/
exit_marker.rs

1//! Exit-status trailer protocol between hook-rewritten commands and
2//! `sqz compress`.
3//!
4//! A pipeline exits with the status of its last command, so
5//! `cmd 2>&1 | sqz compress` would report sqz's status, not cmd's. The
6//! rewrite sends the real status through the pipe as a trailer line
7//! that sqz strips and exits with:
8//!
9//! ```sh
10//! { ( cmd
11//! ); printf '\n__SQZ_EXIT_%d__\n' "$?"; } 2>&1 | sqz compress --cmd cmd
12//! ```
13//!
14//! Plain POSIX, so it works in bash/zsh/dash/ash/ksh where `pipefail`
15//! (rejected by dash/ash) and `PIPESTATUS` (bash-only) would not.
16//! PowerShell gets its own form. Input without a trailer keeps the old
17//! exit-0 behavior, so manual pipes and legacy shell functions are
18//! unaffected.
19
20/// Build the POSIX form of a rewritten command that carries its exit
21/// status to sqz. `escaped_label` must already be shell-escaped. The
22/// newline before `)` ends a trailing comment or heredoc. The `2>&1` sits
23/// on the group, so a redirect of the command's own, such as
24/// `2>/dev/null` or `2> err.log`, still wins.
25pub fn posix_rewrite(command: &str, escaped_label: &str) -> String {
26    format!(
27        "{{ ( {command}\n); printf '\\n__SQZ_EXIT_%d__\\n' \"$?\"; }} 2>&1 | sqz compress --cmd {escaped_label}"
28    )
29}
30
31/// The same form for a list or pipeline. The group's `2>&1` covers every
32/// command in it.
33pub fn posix_rewrite_list(body: &str, escaped_label: &str) -> String {
34    posix_rewrite(body, escaped_label)
35}
36
37/// The ungrouped form for bash and zsh: each marked command gets its own
38/// pipe, and `pipefail` makes each pipe exit with the command's status, so
39/// `&&` and `||` behave as before and builtins run in the shell itself.
40/// The shell must not keep state between commands, or `pipefail` leaks.
41pub fn posix_flat_rewrite(command: &str, cmds: &[crate::shell_scan::FlatCmd]) -> String {
42    let command = command.trim_end();
43    let mut out = String::from("set -o pipefail; ");
44    let mut at = 0;
45    for cmd in cmds {
46        out.push_str(&command[at..cmd.end]);
47        if let Some(label) = &cmd.label {
48            if cmd.merge_stderr {
49                out.push_str(" 2>&1");
50            }
51            out.push_str(" | sqz compress --cmd ");
52            if label.contains(' ') {
53                out.push('\'');
54                out.push_str(label);
55                out.push('\'');
56            } else {
57                out.push_str(label);
58            }
59        }
60        at = cmd.end;
61    }
62    out.push_str(&command[at..]);
63    out
64}
65
66/// Build the PowerShell form. Native commands set `$LASTEXITCODE`;
67/// cmdlets set only `$?`, so fall back to that. The `$(...)`
68/// subexpression keeps the trailing `__` out of the variable name.
69pub fn powershell_rewrite(command: &str, escaped_label: &str) -> String {
70    format!(
71        "& {{ {command} 2>&1; \"__SQZ_EXIT_$(if ($null -ne $LASTEXITCODE) {{ $LASTEXITCODE }} elseif ($?) {{ 0 }} else {{ 1 }})__\" }} | sqz compress --cmd {escaped_label}"
72    )
73}
74
75/// Split a trailing `__SQZ_EXIT_<n>__` marker off `input`.
76///
77/// Returns the input with the marker (and the separator newline the
78/// producer added) removed, plus the status to exit with. When no marker
79/// is present the input is returned unchanged with `None`.
80///
81/// The match is deliberately strict: the marker must be the final
82/// non-empty line, consist of nothing but the literal frame, an optional
83/// `-` and 1 to 10 digits, and sit at the very end of the stream. A
84/// command whose real output happens to contain the string mid-stream is
85/// left alone.
86pub fn strip_exit_marker(input: &str) -> (String, Option<i32>) {
87    let (head, code) = strip_exit_marker_bytes(input.as_bytes());
88    // `head` is a prefix of `input` that ends before an ASCII newline.
89    (input[..head.len()].to_string(), code)
90}
91
92/// [`strip_exit_marker`] for output that isn't valid UTF-8.
93pub fn strip_exit_marker_bytes(input: &[u8]) -> (&[u8], Option<i32>) {
94    const PREFIX: &[u8] = b"__SQZ_EXIT_";
95    const SUFFIX: &[u8] = b"__";
96
97    // Peel trailing newlines (the producer emits a trailing \n; PowerShell
98    // may emit \r\n) to expose the marker line.
99    let mut end = input.len();
100    while end > 0 && matches!(input[end - 1], b'\n' | b'\r') {
101        end -= 1;
102    }
103    let trimmed = &input[..end];
104    let line_start = trimmed.iter().rposition(|&b| b == b'\n').map_or(0, |i| i + 1);
105    let last_line = &trimmed[line_start..];
106
107    let Some(digits) = last_line.strip_prefix(PREFIX).and_then(|r| r.strip_suffix(SUFFIX)) else {
108        return (input, None);
109    };
110    let unsigned = digits.strip_prefix(b"-").unwrap_or(digits);
111    if unsigned.is_empty() || unsigned.len() > 10 || !unsigned.iter().all(u8::is_ascii_digit) {
112        return (input, None);
113    }
114    let Some(code) = std::str::from_utf8(digits).ok().and_then(|d| d.parse::<i64>().ok()) else {
115        return (input, None);
116    };
117
118    // Drop the marker line plus the single separator newline the producer
119    // prepended (`printf '\n…'`), restoring the original bytes exactly.
120    let mut head = &trimmed[..line_start];
121    head = head
122        .strip_suffix(b"\r\n")
123        .or_else(|| head.strip_suffix(b"\n"))
124        .unwrap_or(head);
125    (head, Some(exit_status(code)))
126}
127
128/// The status to exit with: all 32 bits on Windows, the low byte on Unix, never 0 for a failure.
129fn exit_status(code: i64) -> i32 {
130    let kept = if cfg!(windows) { code as i32 } else { (code & 0xFF) as i32 };
131    if kept == 0 && code != 0 {
132        1
133    } else {
134        kept
135    }
136}
137
138// ── Tests ─────────────────────────────────────────────────────────────────
139
140#[cfg(test)]
141mod tests {
142    use super::*;
143
144    #[test]
145    fn strips_marker_and_restores_exact_bytes() {
146        let (out, code) = strip_exit_marker("hello\nworld\n\n__SQZ_EXIT_2__\n");
147        assert_eq!(out, "hello\nworld\n");
148        assert_eq!(code, Some(2));
149    }
150
151    #[test]
152    fn output_without_trailing_newline() {
153        let (out, code) = strip_exit_marker("no newline\n__SQZ_EXIT_42__\n");
154        assert_eq!(out, "no newline");
155        assert_eq!(code, Some(42));
156    }
157
158    #[test]
159    fn empty_output_just_marker() {
160        let (out, code) = strip_exit_marker("\n__SQZ_EXIT_0__\n");
161        assert_eq!(out, "");
162        assert_eq!(code, Some(0));
163    }
164
165    #[test]
166    fn crlf_from_powershell() {
167        let (out, code) = strip_exit_marker("line one\r\n__SQZ_EXIT_1__\r\n");
168        assert_eq!(out, "line one");
169        assert_eq!(code, Some(1));
170    }
171
172    #[test]
173    fn no_marker_passes_through() {
174        let input = "plain output\nwith lines\n";
175        let (out, code) = strip_exit_marker(input);
176        assert_eq!(out, input);
177        assert_eq!(code, None);
178    }
179
180    #[test]
181    fn marker_mid_stream_is_not_a_trailer() {
182        let input = "__SQZ_EXIT_7__\nreal output after\n";
183        let (out, code) = strip_exit_marker(input);
184        assert_eq!(out, input);
185        assert_eq!(code, None);
186    }
187
188    #[test]
189    fn malformed_markers_left_alone() {
190        for input in [
191            "x\n__SQZ_EXIT___\n",
192            "x\n__SQZ_EXIT_12a__\n",
193            "x\n__SQZ_EXIT_12345678901__\n",
194            "x\n__SQZ_EXIT_-__\n",
195            "x\n__SQZ_EXIT_--5__\n",
196            "x\n__SQZ_EXIT_+5__\n",
197            "x\n__SQZ_EXIT_5_\n",
198            "x\nSQZ_EXIT_5__\n",
199        ] {
200            let (out, code) = strip_exit_marker(input);
201            assert_eq!(out, input, "must not strip: {input:?}");
202            assert_eq!(code, None);
203        }
204    }
205
206    #[test]
207    fn negative_and_wide_codes_still_fail() {
208        let (out, code) = strip_exit_marker("fatal: crashed\n\n__SQZ_EXIT_-1073741819__\n");
209        assert_eq!(out, "fatal: crashed\n");
210        assert_eq!(code, Some(if cfg!(windows) { -1073741819 } else { 5 }));
211        let (_, code) = strip_exit_marker("x\n__SQZ_EXIT_3221225477__\n");
212        assert_eq!(code, Some(if cfg!(windows) { -1073741819 } else { 5 }));
213        let (_, code) = strip_exit_marker("x\n__SQZ_EXIT_256__\n");
214        assert_eq!(code, Some(if cfg!(windows) { 256 } else { 1 }));
215        assert_eq!(strip_exit_marker("x\n__SQZ_EXIT_4294967296__\n").1, Some(1));
216        assert_eq!(strip_exit_marker("x\n__SQZ_EXIT_-0__\n").1, Some(0));
217        let (_, code) = strip_exit_marker_bytes(b"\xff\n__SQZ_EXIT_-1__\n");
218        assert_eq!(code, Some(if cfg!(windows) { -1 } else { 255 }));
219    }
220
221    #[test]
222    fn posix_rewrite_shape() {
223        let r = posix_rewrite("git status", "'git status'");
224        assert_eq!(
225            r,
226            "{ ( git status\n); printf '\\n__SQZ_EXIT_%d__\\n' \"$?\"; } 2>&1 | sqz compress --cmd 'git status'"
227        );
228    }
229
230    #[test]
231    fn strips_marker_from_binary_output() {
232        let input = b"\xff\xfe\x00raw\n\n__SQZ_EXIT_3__\n";
233        assert_eq!(strip_exit_marker_bytes(input), (&b"\xff\xfe\x00raw\n"[..], Some(3)));
234        let plain = b"\xff no marker";
235        assert_eq!(strip_exit_marker_bytes(plain), (&plain[..], None));
236    }
237
238    #[test]
239    fn list_rewrite_shape() {
240        assert_eq!(
241            posix_rewrite_list("ls # all", "'sh: ls # all'"),
242            "{ ( ls # all\n); printf '\\n__SQZ_EXIT_%d__\\n' \"$?\"; } 2>&1 | sqz compress --cmd 'sh: ls # all'"
243        );
244    }
245
246    #[test]
247    fn flat_rewrite_shape() {
248        let cmd = "cd sub && cargo test 2>&1; git log -3 # recent\n";
249        let parts = crate::shell_scan::flat(cmd).unwrap();
250        assert_eq!(
251            posix_flat_rewrite(cmd, &parts),
252            "set -o pipefail; cd sub && cargo test 2>&1 | sqz compress --cmd 'cargo test'; \
253             git log -3 2>&1 | sqz compress --cmd 'git log -3' # recent"
254        );
255        let parts = crate::shell_scan::flat("ls").unwrap();
256        assert_eq!(posix_flat_rewrite("ls", &parts), "set -o pipefail; ls 2>&1 | sqz compress --cmd ls");
257    }
258
259    #[test]
260    fn round_trip_posix_semantics() {
261        // Simulate what the shell group produces for a failing command.
262        let produced = format!("ls: cannot access '/nope': No such file or directory\n\n__SQZ_EXIT_{}__\n", 2);
263        let (out, code) = strip_exit_marker(&produced);
264        assert_eq!(out, "ls: cannot access '/nope': No such file or directory\n");
265        assert_eq!(code, Some(2));
266    }
267}
sqz_engine/exit_marker.rs

sqz_engine/
exit_marker.rs
