eggress_testkit/
eggress_runner.rs

1use std::net::SocketAddr;
2use std::path::PathBuf;
3use std::process::{Child, Command, Stdio};
4use std::time::{Duration, Instant};
5
6#[derive(Debug, thiserror::Error)]
7pub enum EggressRunnerError {
8    #[error("failed to spawn eggress process: {0}")]
9    Spawn(#[source] std::io::Error),
10    #[error("failed to write config to temp file: {0}")]
11    ConfigWrite(#[source] std::io::Error),
12    #[error("config file has no path")]
13    ConfigNoPath,
14    #[error("process exited before ready with status: {status}")]
15    EarlyExit { status: String },
16    #[error("port {0} not ready within timeout")]
17    PortNotReady(u16),
18    #[error("process has no stdout handle")]
19    NoStdout,
20    #[error("process has no stderr handle")]
21    NoStderr,
22    #[error("failed to read process output: {0}")]
23    OutputRead(#[source] std::io::Error),
24    #[error("shutdown timed out")]
25    ShutdownTimeout,
26    #[error("failed to kill process: {0}")]
27    Kill(#[source] std::io::Error),
28    #[error("failed to wait for process: {0}")]
29    Wait(#[source] std::io::Error),
30}
31
32#[derive(Debug, Clone)]
33pub struct EggressRunnerConfig {
34    pub binary_path: Option<PathBuf>,
35    pub startup_timeout: Duration,
36    pub shutdown_timeout: Duration,
37    pub io_timeout: Duration,
38}
39
40impl Default for EggressRunnerConfig {
41    fn default() -> Self {
42        Self {
43            binary_path: None,
44            startup_timeout: Duration::from_secs(10),
45            shutdown_timeout: Duration::from_secs(5),
46            io_timeout: Duration::from_secs(5),
47        }
48    }
49}
50
51pub struct EggressProcess {
52    child: Option<Child>,
53    addr: SocketAddr,
54    stdout_lines: Vec<String>,
55    stderr_lines: Vec<String>,
56    #[allow(dead_code)]
57    config_file: Option<tempfile::NamedTempFile>,
58}
59
60impl EggressProcess {
61    pub async fn start_from_toml(
62        config: &EggressRunnerConfig,
63        toml_config: &str,
64    ) -> Result<Self, EggressRunnerError> {
65        let mut f = tempfile::NamedTempFile::new().map_err(EggressRunnerError::ConfigWrite)?;
66        std::io::Write::write_all(&mut f, toml_config.as_bytes())
67            .map_err(EggressRunnerError::ConfigWrite)?;
68        std::io::Write::flush(&mut f).map_err(EggressRunnerError::ConfigWrite)?;
69        let path = f
70            .path()
71            .to_str()
72            .ok_or(EggressRunnerError::ConfigNoPath)?
73            .to_string();
74
75        let mut cmd = binary_command(config)?;
76        cmd.args(["--config", &path]);
77
78        let mut child = spawn_with_output(config, &mut cmd)?;
79        let port = extract_listen_port(toml_config).ok_or_else(|| {
80            let _ = child.kill();
81            EggressRunnerError::PortNotReady(0)
82        })?;
83
84        wait_ready(port, config.startup_timeout)
85            .await
86            .inspect_err(|_| {
87                let _ = child.kill();
88            })?;
89
90        let addr: SocketAddr = format!("127.0.0.1:{port}").parse().map_err(|_| {
91            let _ = child.kill();
92            EggressRunnerError::PortNotReady(port)
93        })?;
94
95        Ok(Self {
96            child: Some(child),
97            addr,
98            stdout_lines: Vec::new(),
99            stderr_lines: Vec::new(),
100            config_file: Some(f),
101        })
102    }
103
104    pub async fn start_from_args(
105        config: &EggressRunnerConfig,
106        args: &[&str],
107    ) -> Result<Self, EggressRunnerError> {
108        let mut cmd = binary_command(config)?;
109        cmd.args(args);
110
111        let mut child = spawn_with_output(config, &mut cmd)?;
112        let port = extract_port_from_args(args).ok_or_else(|| {
113            let _ = child.kill();
114            EggressRunnerError::PortNotReady(0)
115        })?;
116
117        wait_ready(port, config.startup_timeout)
118            .await
119            .inspect_err(|_| {
120                let _ = child.kill();
121            })?;
122
123        let addr: SocketAddr = format!("127.0.0.1:{port}").parse().map_err(|_| {
124            let _ = child.kill();
125            EggressRunnerError::PortNotReady(port)
126        })?;
127
128        Ok(Self {
129            child: Some(child),
130            addr,
131            stdout_lines: Vec::new(),
132            stderr_lines: Vec::new(),
133            config_file: None,
134        })
135    }
136
137    pub fn addr(&self) -> SocketAddr {
138        self.addr
139    }
140
141    pub async fn shutdown(&mut self) -> Result<(), EggressRunnerError> {
142        if self.child.is_none() {
143            return Ok(());
144        }
145        self.drain_output()?;
146        if let Some(ref mut child) = self.child {
147            child.kill().map_err(EggressRunnerError::Kill)?;
148            let deadline = Instant::now() + Duration::from_secs(5);
149            loop {
150                match child.try_wait() {
151                    Ok(Some(_)) => {
152                        self.child = None;
153                        return Ok(());
154                    }
155                    Ok(None) => {
156                        if Instant::now() >= deadline {
157                            let _ = child.kill();
158                            let _ = child.wait();
159                            self.child = None;
160                            return Err(EggressRunnerError::ShutdownTimeout);
161                        }
162                        tokio::time::sleep(Duration::from_millis(10)).await;
163                    }
164                    Err(e) => {
165                        self.child = None;
166                        return Err(EggressRunnerError::Wait(e));
167                    }
168                }
169            }
170        }
171        Ok(())
172    }
173
174    pub fn stdout_lines(&self) -> &[String] {
175        &self.stdout_lines
176    }
177
178    pub fn stderr_lines(&self) -> &[String] {
179        &self.stderr_lines
180    }
181
182    fn drain_output(&mut self) -> Result<(), EggressRunnerError> {
183        if let Some(ref mut child) = self.child {
184            if let Some(ref mut stdout) = child.stdout {
185                use std::io::Read;
186                let mut buf = String::new();
187                let _ = stdout.read_to_string(&mut buf);
188                self.stdout_lines.extend(buf.lines().map(String::from));
189            }
190            if let Some(ref mut stderr) = child.stderr {
191                use std::io::Read;
192                let mut buf = String::new();
193                let _ = stderr.read_to_string(&mut buf);
194                self.stderr_lines.extend(buf.lines().map(String::from));
195            }
196        }
197        Ok(())
198    }
199}
200
201impl Drop for EggressProcess {
202    fn drop(&mut self) {
203        if let Some(ref mut child) = self.child {
204            let _ = child.kill();
205            let _ = child.wait();
206        }
207    }
208}
209
210fn binary_command(config: &EggressRunnerConfig) -> Result<Command, EggressRunnerError> {
211    let cmd = match &config.binary_path {
212        Some(path) => Command::new(path),
213        None => {
214            let mut c = Command::new("cargo");
215            c.args(["run", "--bin", "eggress", "--"]);
216            c
217        }
218    };
219    Ok(cmd)
220}
221
222fn spawn_with_output(
223    _config: &EggressRunnerConfig,
224    cmd: &mut Command,
225) -> Result<Child, EggressRunnerError> {
226    cmd.stdin(Stdio::null())
227        .stdout(Stdio::piped())
228        .stderr(Stdio::piped());
229    cmd.spawn().map_err(EggressRunnerError::Spawn)
230}
231
232async fn wait_ready(port: u16, timeout: Duration) -> Result<(), EggressRunnerError> {
233    let start = Instant::now();
234    loop {
235        if start.elapsed() >= timeout {
236            return Err(EggressRunnerError::PortNotReady(port));
237        }
238        match tokio::net::TcpStream::connect(format!("127.0.0.1:{port}")).await {
239            Ok(stream) => {
240                drop(stream);
241                return Ok(());
242            }
243            Err(_) => {
244                tokio::time::sleep(Duration::from_millis(50)).await;
245            }
246        }
247    }
248}
249
250fn extract_listen_port(toml: &str) -> Option<u16> {
251    for line in toml.lines() {
252        let trimmed = line.trim();
253        if trimmed.starts_with("listen") && trimmed.contains('=') {
254            if let Some(val) = trimmed.split_once('=') {
255                let val = val.1.trim();
256                if let Some(colon) = val.rfind(':') {
257                    let port_str = &val[colon + 1..];
258                    let port_str = port_str.trim_matches(|c: char| c == '"' || c == '\'');
259                    if let Ok(port) = port_str.parse::<u16>() {
260                        return Some(port);
261                    }
262                }
263            }
264        }
265    }
266    None
267}
268
269fn extract_port_from_args(args: &[&str]) -> Option<u16> {
270    let mut i = 0;
271    while i < args.len() {
272        if args[i] == "-l" || args[i] == "--listen" {
273            if let Some(uri) = args.get(i + 1) {
274                if let Some(colon) = uri.rfind(':') {
275                    let port_str = &uri[colon + 1..];
276                    if let Ok(port) = port_str.parse::<u16>() {
277                        return Some(port);
278                    }
279                }
280            }
281        }
282        i += 1;
283    }
284    None
285}
286
287#[cfg(test)]
288mod tests {
289    use super::*;
290
291    #[test]
292    fn process_guard_drop_kills_child() {
293        let mut cmd = Command::new("sleep");
294        cmd.arg("300");
295        cmd.stdin(Stdio::null());
296        cmd.stdout(Stdio::null());
297        cmd.stderr(Stdio::null());
298        let child = cmd.spawn().expect("failed to spawn sleep");
299        let pid = child.id();
300
301        {
302            let _process = EggressProcess {
303                child: Some(child),
304                addr: "127.0.0.1:0".parse().unwrap(),
305                stdout_lines: Vec::new(),
306                stderr_lines: Vec::new(),
307                config_file: None,
308            };
309        }
310
311        std::thread::sleep(Duration::from_millis(200));
312        let result = Command::new("kill").args(["-0", &pid.to_string()]).status();
313        assert!(
314            result.is_err() || !result.unwrap().success(),
315            "process {pid} should be dead after drop"
316        );
317    }
318
319    #[tokio::test]
320    async fn start_from_toml_smoke_socks5() {
321        let port = crate::get_free_port().await;
322        let toml = format!(
323            r#"
324[listener]
325bind = "127.0.0.1:{port}"
326protocols = ["socks5"]
327"#
328        );
329        let config = EggressRunnerConfig {
330            startup_timeout: Duration::from_secs(30),
331            shutdown_timeout: Duration::from_secs(5),
332            ..Default::default()
333        };
334        let result = EggressProcess::start_from_toml(&config, &toml).await;
335        match result {
336            Ok(mut proc) => {
337                assert_eq!(proc.addr().port(), port);
338                let _ = proc.shutdown().await;
339            }
340            Err(EggressRunnerError::PortNotReady(_)) => {
341                // Expected if eggress binary is not built
342            }
343            Err(EggressRunnerError::Spawn(e)) => {
344                assert!(
345                    e.kind() == std::io::ErrorKind::NotFound
346                        || e.kind() == std::io::ErrorKind::Other,
347                    "unexpected spawn error: {e}"
348                );
349            }
350            Err(_) => {}
351        }
352    }
353}
eggress_testkit/eggress_runner.rs

eggress_testkit/
eggress_runner.rs
