wrapcli/
lib.rs

1//! For docs look at [WrapCLI Book](https://vi-is-ramen.github.io/wrapcli/)
2
3use regex::Regex;
4use std::io::{self, BufRead, BufReader, Write};
5use std::process::{Child, Command, ExitStatus, Stdio};
6use std::thread;
7
8/// Wrapper configuration.
9#[derive(Debug, Clone)]
10pub struct WrapConfig {
11    /// Original binary/program name to determine what replace to.
12    pub orig_name: String,
13    /// Fake name of program.
14    pub fake_name: String,
15    /// Fake version of program.
16    pub fake_ver: String,
17    /// Save original program's version in parens or not?
18    pub save_orig: bool,
19}
20
21/// Result of running a command with wrapper.
22#[derive(Debug)]
23pub struct WrapResult {
24    /// Exit status of called binary.
25    pub status: ExitStatus,
26    /// Stdout stream of called binary.
27    pub stdout: Vec<u8>,
28    /// Stderr stream of called binary.
29    pub stderr: Vec<u8>,
30}
31
32struct Rewriter {
33    re_version: Regex,
34    re_usage: Regex,
35    cfg: WrapConfig,
36}
37
38impl Rewriter {
39    fn new(cfg: WrapConfig) -> Self {
40        let re_version = Regex::new(&format!(
41            r"(?m)^([ \t]*){}([ \t]+)([^ \t\r\n]+)([^\r\n]*?)(\r?\n)?$",
42            regex::escape(&cfg.orig_name)
43        ))
44        .expect("invalid version regex");
45
46        let re_usage = Regex::new(&format!(
47            r"(?im)(^|[^\w])([Uu]sage:?[ \t]+){}(\b)",
48            regex::escape(&cfg.orig_name)
49        ))
50        .expect("invalid usage regex");
51
52        Self {
53            re_version,
54            re_usage,
55            cfg,
56        }
57    }
58
59    fn rewrite_first(&self, line: &str) -> String {
60        if let Some(caps) = self.re_version.captures(line) {
61            let indent = caps.get(1).map(|m| m.as_str()).unwrap_or("");
62            let sep = caps.get(2).map(|m| m.as_str()).unwrap_or(" ");
63            let orig_ver = caps.get(3).map(|m| m.as_str()).unwrap_or("");
64            let rest = caps.get(4).map(|m| m.as_str()).unwrap_or("");
65            let eol = caps.get(5).map(|m| m.as_str()).unwrap_or("");
66
67            let mut out = String::with_capacity(line.len() + 32);
68            out.push_str(indent);
69            out.push_str(&self.cfg.fake_name);
70            out.push_str(sep);
71            out.push_str(&self.cfg.fake_ver);
72
73            if self.cfg.save_orig {
74                out.push_str(" (on ");
75                out.push_str(&self.cfg.orig_name);
76                out.push(' ');
77                out.push_str(orig_ver);
78                out.push(')');
79            }
80
81            out.push_str(rest);
82            out.push_str(eol);
83            return out;
84        }
85
86        self.rewrite_usage(line)
87    }
88
89    fn rewrite_usage(&self, line: &str) -> String {
90        self.re_usage
91            .replace_all(line, |caps: &regex::Captures| {
92                let prefix = caps.get(1).map(|m| m.as_str()).unwrap_or("");
93                let usage = caps.get(2).map(|m| m.as_str()).unwrap_or("");
94                let boundary = caps.get(3).map(|m| m.as_str()).unwrap_or("");
95                format!("{}{}{}{}", prefix, usage, self.cfg.fake_name, boundary)
96            })
97            .into_owned()
98    }
99}
100
101/// Run wrap configuration with given CLI arguments.
102pub fn run_streaming<S: AsRef<str>>(
103    cfg: &WrapConfig,
104    args: impl IntoIterator<Item = S>,
105) -> io::Result<ExitStatus> {
106    let mut child = spawn(cfg, args)?;
107
108    let child_stdout = child.stdout.take().expect("stdout was piped");
109    let child_stderr = child.stderr.take().expect("stderr was piped");
110
111    let rewriter_out = Rewriter::new(cfg.clone());
112    let rewriter_err = Rewriter::new(cfg.clone());
113
114    let t_out = thread::spawn(move || -> io::Result<()> {
115        let mut stdout = io::stdout().lock();
116        let rdr = BufReader::new(child_stdout);
117        stream_and_rewrite(rdr, &mut stdout, &rewriter_out)
118    });
119
120    let t_err = thread::spawn(move || -> io::Result<()> {
121        let mut stderr = io::stderr().lock();
122        let rdr = BufReader::new(child_stderr);
123        stream_and_rewrite(rdr, &mut stderr, &rewriter_err)
124    });
125
126    let status = child.wait()?;
127
128    t_out.join().unwrap()?;
129    t_err.join().unwrap()?;
130
131    Ok(status)
132}
133
134/// Run wrap configuration with given CLI arguments and capture child process' stdout/err.
135pub fn run_capture<S: AsRef<str>>(
136    cfg: &WrapConfig,
137    args: impl IntoIterator<Item = S>,
138) -> io::Result<WrapResult> {
139    let mut child = spawn(cfg, args)?;
140
141    let child_stdout = child.stdout.take().expect("stdout was piped");
142    let child_stderr = child.stderr.take().expect("stderr was piped");
143
144    let rewriter_out = Rewriter::new(cfg.clone());
145    let rewriter_err = Rewriter::new(cfg.clone());
146
147    let t_out = thread::spawn(move || -> io::Result<Vec<u8>> {
148        let mut buf: Vec<u8> = Vec::new();
149        let rdr = BufReader::new(child_stdout);
150        stream_and_rewrite(rdr, &mut buf, &rewriter_out)?;
151        Ok(buf)
152    });
153
154    let t_err = thread::spawn(move || -> io::Result<Vec<u8>> {
155        let mut buf: Vec<u8> = Vec::new();
156        let rdr = BufReader::new(child_stderr);
157        stream_and_rewrite(rdr, &mut buf, &rewriter_err)?;
158        Ok(buf)
159    });
160
161    let status = child.wait()?;
162    let stdout = t_out.join().unwrap()?;
163    let stderr = t_err.join().unwrap()?;
164
165    Ok(WrapResult {
166        status,
167        stdout,
168        stderr,
169    })
170}
171
172fn spawn<S: AsRef<str>>(cfg: &WrapConfig, args: impl IntoIterator<Item = S>) -> io::Result<Child> {
173    Command::new(&cfg.orig_name)
174        .args(args.into_iter().map(|s| s.as_ref().to_owned()))
175        .stdin(Stdio::inherit())
176        .stdout(Stdio::piped())
177        .stderr(Stdio::piped())
178        .spawn()
179}
180
181fn stream_and_rewrite<R: BufRead, W: Write>(
182    mut rdr: R,
183    mut out: W,
184    rewriter: &Rewriter,
185) -> io::Result<()> {
186    let mut first_line_done = false;
187    let mut line_buf = String::new();
188
189    loop {
190        line_buf.clear();
191        let n = rdr.read_line(&mut line_buf)?;
192        if n == 0 {
193            break;
194        }
195
196        let rewritten = if !first_line_done {
197            first_line_done = true;
198            rewriter.rewrite_first(&line_buf)
199        } else {
200            rewriter.rewrite_usage(&line_buf)
201        };
202
203        out.write_all(rewritten.as_bytes())?;
204        out.flush()?;
205    }
206
207    Ok(())
208}
209
210#[cfg(test)]
211mod tests {
212    use super::*;
213
214    fn cfg(save_orig: bool) -> WrapConfig {
215        WrapConfig {
216            orig_name: "rustc".into(),
217            fake_name: "dustc".into(),
218            fake_ver: "2.0.0".into(),
219            save_orig,
220        }
221    }
222
223    #[test]
224    fn rewrites_version_line_with_save_orig() {
225        let c = cfg(true);
226        let rewriter = Rewriter::new(c);
227
228        let line = "rustc 1.76.0 (07dca489a 2024-02-04)\n";
229        let out = rewriter.rewrite_first(line);
230        assert_eq!(
231            out,
232            "dustc 2.0.0 (on rustc 1.76.0) (07dca489a 2024-02-04)\n"
233        );
234    }
235
236    #[test]
237    fn rewrites_version_line_without_save_orig() {
238        let c = cfg(false);
239        let rewriter = Rewriter::new(c);
240
241        let line = "rustc 1.76.0 (07dca489a 2024-02-04)\n";
242        let out = rewriter.rewrite_first(line);
243        assert_eq!(out, "dustc 2.0.0 (07dca489a 2024-02-04)\n");
244    }
245
246    #[test]
247    fn rewrites_usage_line() {
248        let c = cfg(false);
249        let rewriter = Rewriter::new(c);
250
251        let line = "Usage: rustc [OPTIONS] INPUT\n";
252        let out = rewriter.rewrite_first(line);
253        assert_eq!(out, "Usage: dustc [OPTIONS] INPUT\n");
254    }
255
256    #[test]
257    fn rewrites_usage_case_insensitive() {
258        let c = cfg(false);
259        let rewriter = Rewriter::new(c);
260
261        let line = "usage: rustc [OPTIONS] INPUT\n";
262        let out = rewriter.rewrite_first(line);
263        assert_eq!(out, "usage: dustc [OPTIONS] INPUT\n");
264    }
265
266    #[test]
267    fn rewrites_usage_in_other_lines() {
268        let c = cfg(false);
269        let rewriter = Rewriter::new(c);
270
271        let line = "Some text. Usage: rustc [OPTIONS]\n";
272        let out = rewriter.rewrite_usage(line);
273        assert_eq!(out, "Some text. Usage: dustc [OPTIONS]\n");
274    }
275
276    #[test]
277    fn does_not_touch_unrelated_occurrences() {
278        let c = cfg(true);
279        let rewriter = Rewriter::new(c);
280
281        let line = "error: aborting due to rustc internal error\n";
282        let out = rewriter.rewrite_first(line);
283        assert_eq!(out, line);
284
285        let out2 = rewriter.rewrite_usage(line);
286        assert_eq!(out2, line);
287    }
288
289    #[test]
290    fn handles_crlf() {
291        let c = cfg(false);
292        let rewriter = Rewriter::new(c);
293
294        let line = "rustc 1.76.0\r\n";
295        let out = rewriter.rewrite_first(line);
296        assert_eq!(out, "dustc 2.0.0\r\n");
297    }
298
299    #[test]
300    fn preserves_indentation() {
301        let c = cfg(false);
302        let rewriter = Rewriter::new(c);
303
304        let line = "  Usage: rustc [OPTIONS]\n";
305        let out = rewriter.rewrite_usage(line);
306        assert_eq!(out, "  Usage: dustc [OPTIONS]\n");
307    }
308}
wrapcli/lib.rs

wrapcli/
lib.rs
