ctf_pwn/io/pipe/convert/
process.rs

1use crate::io::merge::MergedAsyncReader;
2use crate::io::Pipe;
3use std::{ffi::OsStr, process::Stdio, result};
4use thiserror::*;
5use tokio::process::*;
6
7#[derive(Error, Debug)]
8pub enum ProcessPipeError {
9    #[error("stdin not set")]
10    StdinNotSet,
11    #[error("stdout not set")]
12    StdoutNotSet,
13    #[error("stderr not set")]
14    StdErrNotSet,
15    #[error("io error")]
16    IoError(#[from] tokio::io::Error),
17}
18
19pub type ProcessPipeResult<T> = result::Result<T, ProcessPipeError>;
20
21pub type ProcessPipe = Pipe<MergedAsyncReader<ChildStdout, ChildStderr>, ChildStdin>;
22pub type StdoutPipe = Pipe<ChildStdout, ChildStdin>;
23pub type StderrPipe = Pipe<ChildStderr, ChildStdin>;
24
25impl ProcessPipe {
26    pub async fn from_app<S: AsRef<OsStr>>(program: S) -> ProcessPipeResult<Self> {
27        let command = Command::new(program);
28        Self::spawn_command(command)
29    }
30
31    pub async fn from_app_args<S: AsRef<OsStr>, I: IntoIterator<Item = S>>(
32        program: S,
33        args: I,
34    ) -> ProcessPipeResult<Self> {
35        let mut command = Command::new(program);
36        let _ = command.args(args);
37        Self::spawn_command(command)
38    }
39
40    pub fn spawn_command(mut value: Command) -> ProcessPipeResult<Self> {
41        let process = value
42            .stdout(Stdio::piped())
43            .stdin(Stdio::piped())
44            .stderr(Stdio::piped())
45            .spawn()?;
46
47        process.try_into()
48    }
49}
50
51impl StdoutPipe {
52    pub async fn from_app<S: AsRef<OsStr>>(program: S) -> ProcessPipeResult<Self> {
53        let command = Command::new(program);
54        Self::spawn_command(command)
55    }
56
57    pub async fn from_app_args<S: AsRef<OsStr>, I: IntoIterator<Item = S>>(
58        program: S,
59        args: I,
60    ) -> ProcessPipeResult<Self> {
61        let mut command = Command::new(program);
62        let _ = command.args(args);
63        Self::spawn_command(command)
64    }
65
66    pub fn spawn_command(mut value: Command) -> ProcessPipeResult<Self> {
67        let process = value.stdout(Stdio::piped()).stdin(Stdio::piped()).spawn()?;
68
69        let stdin = process.stdin.ok_or(ProcessPipeError::StdinNotSet)?;
70        let stdout = process.stdout.ok_or(ProcessPipeError::StdoutNotSet)?;
71        Ok((stdin, stdout).into())
72    }
73}
74
75impl StderrPipe {
76    pub async fn from_app<S: AsRef<OsStr>>(program: S) -> ProcessPipeResult<Self> {
77        let command = Command::new(program);
78        Self::spawn_command(command)
79    }
80
81    pub async fn from_app_args<S: AsRef<OsStr>, I: IntoIterator<Item = S>>(
82        program: S,
83        args: I,
84    ) -> ProcessPipeResult<Self> {
85        let mut command = Command::new(program);
86        let _ = command.args(args);
87        Self::spawn_command(command)
88    }
89
90    pub fn spawn_command(mut value: Command) -> ProcessPipeResult<Self> {
91        let process = value.stderr(Stdio::piped()).stdin(Stdio::piped()).spawn()?;
92
93        let stdin = process.stdin.ok_or(ProcessPipeError::StdinNotSet)?;
94        let stderr = process.stderr.ok_or(ProcessPipeError::StdErrNotSet)?;
95        Ok((stdin, stderr).into())
96    }
97}
98
99impl TryFrom<Child> for ProcessPipe {
100    type Error = ProcessPipeError;
101
102    fn try_from(value: Child) -> ProcessPipeResult<Self> {
103        let stdin = value.stdin.ok_or(ProcessPipeError::StdinNotSet)?;
104        let stdout = value.stdout.ok_or(ProcessPipeError::StdoutNotSet)?;
105        let stderr = value.stderr.ok_or(ProcessPipeError::StdErrNotSet)?;
106        Ok((stdin, stdout, stderr).into())
107    }
108}
109
110impl From<(ChildStdin, ChildStdout, ChildStderr)> for ProcessPipe {
111    fn from(value: (ChildStdin, ChildStdout, ChildStderr)) -> Self {
112        let (stdin, stdout, stderr) = value;
113        let read_stream = MergedAsyncReader {
114            stream0: stdout,
115            stream1: stderr,
116        };
117        let write_stream = stdin;
118        Pipe::new(read_stream, write_stream)
119    }
120}
121
122impl From<(ChildStdin, ChildStdout)> for StdoutPipe {
123    fn from(value: (ChildStdin, ChildStdout)) -> Self {
124        let (stdin, stdout) = value;
125        let read_stream = stdout;
126        let write_stream = stdin;
127        Pipe::new(read_stream, write_stream)
128    }
129}
130
131impl From<(ChildStdin, ChildStderr)> for StderrPipe {
132    fn from(value: (ChildStdin, ChildStderr)) -> Self {
133        let (stdin, stderr) = value;
134        let read_stream = stderr;
135        let write_stream = stdin;
136        Pipe::new(read_stream, write_stream)
137    }
138}