eggress_testkit/
pproxy_oracle.rs

1use std::net::SocketAddr;
2use std::process::{Child, Command, Stdio};
3use std::sync::{Arc, Mutex};
4use std::time::{Duration, Instant};
5
6use thiserror::Error;
7use tokio::net::TcpStream;
8
9#[derive(Debug, Error)]
10pub enum PproxyOracleError {
11    #[error("io error: {0}")]
12    Io(#[from] std::io::Error),
13
14    #[error("process failed: {0}")]
15    ProcessFailed(String),
16
17    #[error("startup timeout after {0:?}")]
18    StartupTimeout(Duration),
19
20    #[error("not ready: {0}")]
21    NotReady(String),
22
23    #[error("version mismatch: expected {expected}, got {actual}")]
24    VersionMismatch { expected: String, actual: String },
25
26    #[error("version detection failed: {0}")]
27    VersionDetectionFailed(String),
28}
29
30#[derive(Debug, Clone)]
31pub struct OracleConfig {
32    pub python_binary: String,
33    pub pproxy_version: String,
34    pub startup_timeout: Duration,
35    pub shutdown_timeout: Duration,
36    pub io_timeout: Duration,
37}
38
39impl Default for OracleConfig {
40    fn default() -> Self {
41        Self {
42            python_binary: "python3".to_string(),
43            pproxy_version: "2.7.9".to_string(),
44            startup_timeout: Duration::from_secs(15),
45            shutdown_timeout: Duration::from_secs(5),
46            io_timeout: Duration::from_secs(3),
47        }
48    }
49}
50
51pub struct PproxyProcess {
52    child: Option<Child>,
53    bound_addr: SocketAddr,
54    stdout_buf: Arc<Mutex<Vec<u8>>>,
55    stderr_buf: Arc<Mutex<Vec<u8>>>,
56    #[allow(dead_code)]
57    work_dir: Option<tempfile::TempDir>,
58}
59
60impl PproxyProcess {
61    pub async fn start(config: &OracleConfig, args: &[String]) -> Result<Self, PproxyOracleError> {
62        let work_dir = tempfile::TempDir::new().map_err(PproxyOracleError::Io)?;
63
64        let stdout_buf = Arc::new(Mutex::new(Vec::new()));
65        let stderr_buf = Arc::new(Mutex::new(Vec::new()));
66
67        let stdout_clone = Arc::clone(&stdout_buf);
68        let stderr_clone = Arc::clone(&stderr_buf);
69
70        let mut child = Command::new(&config.python_binary)
71            .arg("-m")
72            .arg("pproxy")
73            .args(args)
74            .current_dir(work_dir.path())
75            .stdout(Stdio::piped())
76            .stderr(Stdio::piped())
77            .spawn()
78            .map_err(PproxyOracleError::Io)?;
79
80        let child_stdout = child.stdout.take().expect("stdout piped");
81        let child_stderr = child.stderr.take().expect("stderr piped");
82
83        std::thread::spawn(move || {
84            let mut reader = child_stderr;
85            let mut tmp = [0u8; 4096];
86            loop {
87                match std::io::Read::read(&mut reader, &mut tmp) {
88                    Ok(0) => break,
89                    Ok(n) => {
90                        if let Ok(mut guard) = stderr_clone.lock() {
91                            guard.extend_from_slice(&tmp[..n]);
92                        }
93                    }
94                    Err(_) => break,
95                }
96            }
97        });
98
99        std::thread::spawn(move || {
100            let mut reader = child_stdout;
101            let mut tmp = [0u8; 4096];
102            loop {
103                match std::io::Read::read(&mut reader, &mut tmp) {
104                    Ok(0) => break,
105                    Ok(n) => {
106                        if let Ok(mut guard) = stdout_clone.lock() {
107                            guard.extend_from_slice(&tmp[..n]);
108                        }
109                    }
110                    Err(_) => break,
111                }
112            }
113        });
114
115        let bound_addr = if let Some(addr) = parse_listen_addr_from_args(args) {
116            addr
117        } else {
118            wait_for_output_ready(&stderr_buf, &stdout_buf, config.startup_timeout).await?
119        };
120
121        let proc = Self {
122            child: Some(child),
123            bound_addr,
124            stdout_buf,
125            stderr_buf,
126            work_dir: Some(work_dir),
127        };
128
129        proc.wait_ready(config).await?;
130
131        Ok(proc)
132    }
133
134    pub async fn wait_ready(&self, config: &OracleConfig) -> Result<(), PproxyOracleError> {
135        let start = Instant::now();
136        let interval = Duration::from_millis(100);
137
138        loop {
139            match TcpStream::connect(self.bound_addr).await {
140                Ok(_) => return Ok(()),
141                Err(_) if start.elapsed() < config.startup_timeout => {
142                    tokio::time::sleep(interval).await;
143                }
144                Err(e) => {
145                    return Err(PproxyOracleError::NotReady(format!(
146                        "tcp connect to {} failed: {}",
147                        self.bound_addr, e
148                    )));
149                }
150            }
151        }
152    }
153
154    pub fn shutdown(&mut self) {
155        if let Some(ref mut child) = self.child {
156            let _ = child.kill();
157            let _ = child.wait();
158        }
159        self.child = None;
160    }
161
162    pub fn stdout(&self) -> Vec<u8> {
163        self.stdout_buf
164            .lock()
165            .map(|g| g.clone())
166            .unwrap_or_default()
167    }
168
169    pub fn stderr(&self) -> Vec<u8> {
170        self.stderr_buf
171            .lock()
172            .map(|g| g.clone())
173            .unwrap_or_default()
174    }
175
176    pub fn bound_addr(&self) -> SocketAddr {
177        self.bound_addr
178    }
179
180    pub fn redacted_stderr(&self) -> String {
181        let raw = self.stderr();
182        redact_credentials(&raw)
183    }
184}
185
186impl Drop for PproxyProcess {
187    fn drop(&mut self) {
188        self.shutdown();
189    }
190}
191
192pub fn redact_credentials(data: &[u8]) -> String {
193    let text = String::from_utf8_lossy(data);
194    redact_uri_credentials(&text)
195}
196
197/// Scrub proxy credentials from free-form oracle transcript text.
198///
199/// Delegates to the canonical tolerant redactor
200/// ([`eggress_uri::redact_proxy_uri`]) so every credential-bearing scheme is
201/// covered without maintaining a parallel scheme whitelist here. The match is
202/// deliberately generic: any `scheme://...@...` userinfo is stripped while
203/// credential-free text passes through unchanged.
204fn redact_uri_credentials(text: &str) -> String {
205    // Fast path: skip allocation-free scan when no URI-like content exists.
206    if !text.contains("://") && !text.contains('@') {
207        return text.to_string();
208    }
209    // The canonical redactor handles one URI-like string; transcripts may
210    // embed several across whitespace, so apply it per whitespace-separated
211    // token to avoid one malformed span suppressing the rest.
212    // (`split_inclusive` keeps each delimiter attached to its token, so
213    // concatenation reconstructs the input exactly.)
214    let mut out = String::with_capacity(text.len());
215    for token in text.split_inclusive(char::is_whitespace) {
216        out.push_str(&eggress_uri::redact_proxy_uri(token));
217    }
218    out
219}
220
221async fn wait_for_output_ready(
222    stderr_buf: &Arc<Mutex<Vec<u8>>>,
223    stdout_buf: &Arc<Mutex<Vec<u8>>>,
224    timeout: Duration,
225) -> Result<SocketAddr, PproxyOracleError> {
226    let start = Instant::now();
227    let interval = Duration::from_millis(100);
228
229    loop {
230        {
231            let guard = stderr_buf.lock().map_err(|e| {
232                PproxyOracleError::ProcessFailed(format!("stderr lock poisoned: {}", e))
233            })?;
234            let text = String::from_utf8_lossy(&guard);
235            if let Some(addr) = parse_bound_addr(&text) {
236                return Ok(addr);
237            }
238        }
239        {
240            let guard = stdout_buf.lock().map_err(|e| {
241                PproxyOracleError::ProcessFailed(format!("stdout lock poisoned: {}", e))
242            })?;
243            let text = String::from_utf8_lossy(&guard);
244            if let Some(addr) = parse_bound_addr(&text) {
245                return Ok(addr);
246            }
247        }
248
249        if start.elapsed() >= timeout {
250            let stderr_text = stderr_buf
251                .lock()
252                .map(|g| String::from_utf8_lossy(&g).trim().to_string())
253                .unwrap_or_default();
254            let stdout_text = stdout_buf
255                .lock()
256                .map(|g| String::from_utf8_lossy(&g).trim().to_string())
257                .unwrap_or_default();
258            return Err(PproxyOracleError::ProcessFailed(format!(
259                "startup timeout after {:?}, stderr: [{}], stdout: [{}]",
260                timeout, stderr_text, stdout_text
261            )));
262        }
263
264        tokio::time::sleep(interval).await;
265    }
266}
267
268fn parse_listen_addr_from_args(args: &[String]) -> Option<SocketAddr> {
269    let mut i = 0;
270    while i < args.len() {
271        if (args[i] == "-l" || args[i] == "--listen") && i + 1 < args.len() {
272            let uri = &args[i + 1];
273            let host_port = if let Some(at_pos) = uri.rfind('@') {
274                &uri[at_pos + 1..]
275            } else {
276                uri.as_str()
277            };
278            let stripped = host_port
279                .strip_prefix("socks5://")
280                .or_else(|| host_port.strip_prefix("http://"))
281                .or_else(|| host_port.strip_prefix("socks4://"))
282                .or_else(|| host_port.strip_prefix("socks4a://"))
283                .or_else(|| host_port.strip_prefix("ss://"))
284                .or_else(|| host_port.strip_prefix("trojan://"))
285                .or_else(|| host_port.strip_prefix("direct://"))
286                .unwrap_or(host_port);
287            if let Some(colon_pos) = stripped.rfind(':') {
288                let port_str = &stripped[colon_pos + 1..];
289                if let Ok(port) = port_str.parse::<u16>() {
290                    if port > 0 {
291                        let host_str = &stripped[..colon_pos];
292                        let host = if host_str.is_empty() || host_str == "127.0.0.1" {
293                            std::net::IpAddr::V4(std::net::Ipv4Addr::LOCALHOST)
294                        } else if host_str == "0.0.0.0" {
295                            std::net::IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED)
296                        } else {
297                            host_str.parse().ok()?
298                        };
299                        return Some(SocketAddr::new(host, port));
300                    }
301                }
302            }
303            return None;
304        }
305        i += 1;
306    }
307    None
308}
309
310fn parse_bound_addr(text: &str) -> Option<SocketAddr> {
311    // (needle in oracle output, resulting IP). IPv6 literals are bracketed
312    // exactly as they appear in pproxy's startup banner.
313    const PATTERNS: [(&str, std::net::IpAddr); 4] = [
314        (
315            "127.0.0.1:",
316            std::net::IpAddr::V4(std::net::Ipv4Addr::LOCALHOST),
317        ),
318        (
319            "0.0.0.0:",
320            std::net::IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED),
321        ),
322        (
323            "[::1]:",
324            std::net::IpAddr::V6(std::net::Ipv6Addr::LOCALHOST),
325        ),
326        (
327            "[::]:",
328            std::net::IpAddr::V6(std::net::Ipv6Addr::UNSPECIFIED),
329        ),
330    ];
331
332    for line in text.lines() {
333        let line = line.trim();
334        if line.is_empty() {
335            continue;
336        }
337
338        for (needle, ip) in PATTERNS {
339            if let Some(idx) = line.find(needle) {
340                let port_str = &line[idx + needle.len()..];
341                let port_str = port_str
342                    .chars()
343                    .take_while(|c| c.is_ascii_digit())
344                    .collect::<String>();
345                if let Ok(port) = port_str.parse::<u16>() {
346                    if port > 0 {
347                        return Some(SocketAddr::new(ip, port));
348                    }
349                }
350            }
351        }
352    }
353
354    None
355}
356
357pub async fn verify_pproxy_version(config: &OracleConfig) -> Result<String, PproxyOracleError> {
358    let output = Command::new(&config.python_binary)
359        .args([
360            "-c",
361            "import pproxy; print(getattr(pproxy, '__version__', 'unknown'))",
362        ])
363        .stdout(Stdio::piped())
364        .stderr(Stdio::piped())
365        .output()
366        .map_err(PproxyOracleError::Io)?;
367
368    if !output.status.success() {
369        let stderr = String::from_utf8_lossy(&output.stderr);
370        return Err(PproxyOracleError::VersionDetectionFailed(
371            stderr.trim().to_string(),
372        ));
373    }
374
375    let version = String::from_utf8_lossy(&output.stdout).trim().to_string();
376
377    Ok(version)
378}
379
380pub async fn assert_pproxy_version(config: &OracleConfig) -> Result<(), PproxyOracleError> {
381    let actual = verify_pproxy_version(config).await?;
382    if actual != config.pproxy_version && actual != "unknown" {
383        return Err(PproxyOracleError::VersionMismatch {
384            expected: config.pproxy_version.clone(),
385            actual,
386        });
387    }
388    Ok(())
389}
390
391#[cfg(test)]
392mod tests {
393    use super::*;
394
395    fn require_pproxy() -> bool {
396        std::process::Command::new("python3")
397            .args(["-c", "import pproxy"])
398            .stdout(Stdio::null())
399            .stderr(Stdio::null())
400            .status()
401            .map(|s| s.success())
402            .unwrap_or(false)
403    }
404
405    #[tokio::test]
406    #[ignore]
407    async fn test_process_guard_drop_kills_child() {
408        if !require_pproxy() {
409            eprintln!("pproxy not available, skipping");
410            return;
411        }
412
413        let config = OracleConfig::default();
414        let port = crate::get_free_port().await;
415        let listen = format!("socks5://127.0.0.1:{}", port);
416        let args = vec![
417            "-l".to_string(),
418            listen,
419            "-r".to_string(),
420            "direct".to_string(),
421        ];
422
423        let proc = PproxyProcess::start(&config, &args).await.unwrap();
424        let addr = proc.bound_addr();
425        assert_eq!(addr.port(), port, "should parse port from args");
426
427        drop(proc);
428
429        tokio::time::sleep(Duration::from_millis(200)).await;
430
431        let result = TcpStream::connect(addr).await;
432        assert!(result.is_err(), "process should be dead after drop");
433    }
434
435    #[tokio::test]
436    #[ignore]
437    async fn test_readiness_probe() {
438        if !require_pproxy() {
439            eprintln!("pproxy not available, skipping");
440            return;
441        }
442
443        let config = OracleConfig::default();
444        let port = crate::get_free_port().await;
445        let listen = format!("socks5://127.0.0.1:{}", port);
446        let args = vec![
447            "-l".to_string(),
448            listen,
449            "-r".to_string(),
450            "direct".to_string(),
451        ];
452
453        let proc = PproxyProcess::start(&config, &args).await.unwrap();
454
455        let result = TcpStream::connect(proc.bound_addr()).await;
456        assert!(result.is_ok(), "process should be ready after start");
457
458        drop(proc);
459    }
460
461    #[test]
462    fn test_log_redaction() {
463        let input = b"socks5://user:secret123@127.0.0.1:1080\nhttp://admin:pw0rd@0.0.0.0:8080\n";
464        let redacted = redact_credentials(input);
465
466        assert!(!redacted.contains("secret123"));
467        assert!(!redacted.contains("pw0rd"));
468        assert!(!redacted.contains("user:"), "username must be redacted");
469        assert!(!redacted.contains("admin:"), "username must be redacted");
470        assert!(redacted.contains("****@127.0.0.1:1080"));
471        assert!(redacted.contains("****@0.0.0.0:8080"));
472    }
473
474    #[test]
475    fn test_log_redaction_no_credentials() {
476        let input = b"socks5://127.0.0.1:1080\nlistening on port 8080\n";
477        let redacted = redact_credentials(input);
478
479        assert_eq!(redacted, String::from_utf8_lossy(input));
480    }
481
482    #[test]
483    fn test_parse_bound_addr() {
484        assert_eq!(
485            parse_bound_addr("Listen: socks5://127.0.0.1:9090"),
486            Some(SocketAddr::new(
487                std::net::IpAddr::V4(std::net::Ipv4Addr::LOCALHOST),
488                9090
489            ))
490        );
491
492        assert_eq!(
493            parse_bound_addr("Listen: http://0.0.0.0:8080"),
494            Some(SocketAddr::new(
495                std::net::IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED),
496                8080
497            ))
498        );
499
500        assert_eq!(parse_bound_addr("no address here"), None);
501        assert_eq!(parse_bound_addr(""), None);
502
503        assert_eq!(
504            parse_bound_addr("Listen: socks5://[::1]:9090"),
505            Some(SocketAddr::new(
506                std::net::IpAddr::V6(std::net::Ipv6Addr::LOCALHOST),
507                9090
508            ))
509        );
510
511        assert_eq!(
512            parse_bound_addr("Listen: http://[::]:8080"),
513            Some(SocketAddr::new(
514                std::net::IpAddr::V6(std::net::Ipv6Addr::UNSPECIFIED),
515                8080
516            ))
517        );
518    }
519
520    #[test]
521    fn test_redact_uri_credentials() {
522        assert_eq!(
523            redact_uri_credentials("socks5://user:pass@host:1080"),
524            "socks5://****@host:1080"
525        );
526
527        assert_eq!(
528            redact_uri_credentials("http://admin:secret@proxy:8080"),
529            "http://****@proxy:8080"
530        );
531
532        assert_eq!(
533            redact_uri_credentials("socks5://127.0.0.1:1080"),
534            "socks5://127.0.0.1:1080"
535        );
536
537        // Schemes outside the old hardcoded whitelist are covered by the
538        // canonical redactor.
539        assert_eq!(
540            redact_uri_credentials("ssh://deploy:s3cret@build.internal:22"),
541            "ssh://****@build.internal:22"
542        );
543        assert_eq!(
544            redact_uri_credentials("h3://user:pass@[2001:db8::1]:443"),
545            "h3://****@[2001:db8::1]:443"
546        );
547
548        // Multi-token transcripts redact each credentialed token.
549        assert_eq!(
550            redact_uri_credentials("dial socks5://a:b@h1:1080 then http://c:d@h2:8080 ok"),
551            "dial socks5://****@h1:1080 then http://****@h2:8080 ok"
552        );
553    }
554
555    #[tokio::test]
556    async fn test_version_detection() {
557        if !require_pproxy() {
558            eprintln!("pproxy not available, skipping");
559            return;
560        }
561
562        let config = OracleConfig::default();
563        let result = verify_pproxy_version(&config).await;
564        assert!(
565            result.is_ok(),
566            "version detection should succeed: {:?}",
567            result.err()
568        );
569    }
570}
eggress_testkit/pproxy_oracle.rs

eggress_testkit/
pproxy_oracle.rs
