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
197fn redact_uri_credentials(text: &str) -> String {
198    let mut result = text.to_string();
199
200    let patterns = [
201        "socks4://",
202        "socks4a://",
203        "socks5://",
204        "http://",
205        "https://",
206        "ss://",
207        "trojan://",
208    ];
209
210    for scheme in &patterns {
211        let mut offset = 0;
212        while let Some(scheme_pos) = result[offset..].find(scheme) {
213            let abs_pos = offset + scheme_pos;
214            let after_scheme = abs_pos + scheme.len();
215            if let Some(at_rel) = result[after_scheme..].find('@') {
216                let cred_start = after_scheme;
217                let cred_end = after_scheme + at_rel;
218                let colon_pos = result[cred_start..cred_end].find(':');
219                if colon_pos.is_some() {
220                    let rest = result[cred_end + 1..].to_string();
221                    // Match the `****:****@` redaction convention used by the
222                    // other layers (embed, pproxy-compat, python bindings) so
223                    // differential transcripts never carry proxy usernames.
224                    const REDACTED_PREFIX: &str = "****:****@";
225                    let redacted = format!("{REDACTED_PREFIX}{rest}");
226                    result = format!("{}{}", &result[..cred_start], redacted);
227                    offset = cred_start + REDACTED_PREFIX.len();
228                } else {
229                    offset = cred_end;
230                }
231            } else {
232                break;
233            }
234        }
235    }
236
237    result
238}
239
240async fn wait_for_output_ready(
241    stderr_buf: &Arc<Mutex<Vec<u8>>>,
242    stdout_buf: &Arc<Mutex<Vec<u8>>>,
243    timeout: Duration,
244) -> Result<SocketAddr, PproxyOracleError> {
245    let start = Instant::now();
246    let interval = Duration::from_millis(100);
247
248    loop {
249        {
250            let guard = stderr_buf.lock().map_err(|e| {
251                PproxyOracleError::ProcessFailed(format!("stderr lock poisoned: {}", e))
252            })?;
253            let text = String::from_utf8_lossy(&guard);
254            if let Some(addr) = parse_bound_addr(&text) {
255                return Ok(addr);
256            }
257        }
258        {
259            let guard = stdout_buf.lock().map_err(|e| {
260                PproxyOracleError::ProcessFailed(format!("stdout lock poisoned: {}", e))
261            })?;
262            let text = String::from_utf8_lossy(&guard);
263            if let Some(addr) = parse_bound_addr(&text) {
264                return Ok(addr);
265            }
266        }
267
268        if start.elapsed() >= timeout {
269            let stderr_text = stderr_buf
270                .lock()
271                .map(|g| String::from_utf8_lossy(&g).trim().to_string())
272                .unwrap_or_default();
273            let stdout_text = stdout_buf
274                .lock()
275                .map(|g| String::from_utf8_lossy(&g).trim().to_string())
276                .unwrap_or_default();
277            return Err(PproxyOracleError::ProcessFailed(format!(
278                "startup timeout after {:?}, stderr: [{}], stdout: [{}]",
279                timeout, stderr_text, stdout_text
280            )));
281        }
282
283        tokio::time::sleep(interval).await;
284    }
285}
286
287fn parse_listen_addr_from_args(args: &[String]) -> Option<SocketAddr> {
288    let mut i = 0;
289    while i < args.len() {
290        if (args[i] == "-l" || args[i] == "--listen") && i + 1 < args.len() {
291            let uri = &args[i + 1];
292            let host_port = if let Some(at_pos) = uri.rfind('@') {
293                &uri[at_pos + 1..]
294            } else {
295                uri.as_str()
296            };
297            let stripped = host_port
298                .strip_prefix("socks5://")
299                .or_else(|| host_port.strip_prefix("http://"))
300                .or_else(|| host_port.strip_prefix("socks4://"))
301                .or_else(|| host_port.strip_prefix("socks4a://"))
302                .or_else(|| host_port.strip_prefix("ss://"))
303                .or_else(|| host_port.strip_prefix("trojan://"))
304                .or_else(|| host_port.strip_prefix("direct://"))
305                .unwrap_or(host_port);
306            if let Some(colon_pos) = stripped.rfind(':') {
307                let port_str = &stripped[colon_pos + 1..];
308                if let Ok(port) = port_str.parse::<u16>() {
309                    if port > 0 {
310                        let host_str = &stripped[..colon_pos];
311                        let host = if host_str.is_empty() || host_str == "127.0.0.1" {
312                            std::net::IpAddr::V4(std::net::Ipv4Addr::LOCALHOST)
313                        } else if host_str == "0.0.0.0" {
314                            std::net::IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED)
315                        } else {
316                            host_str.parse().ok()?
317                        };
318                        return Some(SocketAddr::new(host, port));
319                    }
320                }
321            }
322            return None;
323        }
324        i += 1;
325    }
326    None
327}
328
329fn parse_bound_addr(text: &str) -> Option<SocketAddr> {
330    // (needle in oracle output, resulting IP). IPv6 literals are bracketed
331    // exactly as they appear in pproxy's startup banner.
332    const PATTERNS: [(&str, std::net::IpAddr); 4] = [
333        (
334            "127.0.0.1:",
335            std::net::IpAddr::V4(std::net::Ipv4Addr::LOCALHOST),
336        ),
337        (
338            "0.0.0.0:",
339            std::net::IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED),
340        ),
341        (
342            "[::1]:",
343            std::net::IpAddr::V6(std::net::Ipv6Addr::LOCALHOST),
344        ),
345        (
346            "[::]:",
347            std::net::IpAddr::V6(std::net::Ipv6Addr::UNSPECIFIED),
348        ),
349    ];
350
351    for line in text.lines() {
352        let line = line.trim();
353        if line.is_empty() {
354            continue;
355        }
356
357        for (needle, ip) in PATTERNS {
358            if let Some(idx) = line.find(needle) {
359                let port_str = &line[idx + needle.len()..];
360                let port_str = port_str
361                    .chars()
362                    .take_while(|c| c.is_ascii_digit())
363                    .collect::<String>();
364                if let Ok(port) = port_str.parse::<u16>() {
365                    if port > 0 {
366                        return Some(SocketAddr::new(ip, port));
367                    }
368                }
369            }
370        }
371    }
372
373    None
374}
375
376pub async fn verify_pproxy_version(config: &OracleConfig) -> Result<String, PproxyOracleError> {
377    let output = Command::new(&config.python_binary)
378        .args([
379            "-c",
380            "import pproxy; print(getattr(pproxy, '__version__', 'unknown'))",
381        ])
382        .stdout(Stdio::piped())
383        .stderr(Stdio::piped())
384        .output()
385        .map_err(PproxyOracleError::Io)?;
386
387    if !output.status.success() {
388        let stderr = String::from_utf8_lossy(&output.stderr);
389        return Err(PproxyOracleError::VersionDetectionFailed(
390            stderr.trim().to_string(),
391        ));
392    }
393
394    let version = String::from_utf8_lossy(&output.stdout).trim().to_string();
395
396    Ok(version)
397}
398
399pub async fn assert_pproxy_version(config: &OracleConfig) -> Result<(), PproxyOracleError> {
400    let actual = verify_pproxy_version(config).await?;
401    if actual != config.pproxy_version && actual != "unknown" {
402        return Err(PproxyOracleError::VersionMismatch {
403            expected: config.pproxy_version.clone(),
404            actual,
405        });
406    }
407    Ok(())
408}
409
410#[cfg(test)]
411mod tests {
412    use super::*;
413
414    fn require_pproxy() -> bool {
415        std::process::Command::new("python3")
416            .args(["-c", "import pproxy"])
417            .stdout(Stdio::null())
418            .stderr(Stdio::null())
419            .status()
420            .map(|s| s.success())
421            .unwrap_or(false)
422    }
423
424    #[tokio::test]
425    #[ignore]
426    async fn test_process_guard_drop_kills_child() {
427        if !require_pproxy() {
428            eprintln!("pproxy not available, skipping");
429            return;
430        }
431
432        let config = OracleConfig::default();
433        let port = crate::get_free_port().await;
434        let listen = format!("socks5://127.0.0.1:{}", port);
435        let args = vec![
436            "-l".to_string(),
437            listen,
438            "-r".to_string(),
439            "direct".to_string(),
440        ];
441
442        let proc = PproxyProcess::start(&config, &args).await.unwrap();
443        let addr = proc.bound_addr();
444        assert_eq!(addr.port(), port, "should parse port from args");
445
446        drop(proc);
447
448        tokio::time::sleep(Duration::from_millis(200)).await;
449
450        let result = TcpStream::connect(addr).await;
451        assert!(result.is_err(), "process should be dead after drop");
452    }
453
454    #[tokio::test]
455    #[ignore]
456    async fn test_readiness_probe() {
457        if !require_pproxy() {
458            eprintln!("pproxy not available, skipping");
459            return;
460        }
461
462        let config = OracleConfig::default();
463        let port = crate::get_free_port().await;
464        let listen = format!("socks5://127.0.0.1:{}", port);
465        let args = vec![
466            "-l".to_string(),
467            listen,
468            "-r".to_string(),
469            "direct".to_string(),
470        ];
471
472        let proc = PproxyProcess::start(&config, &args).await.unwrap();
473
474        let result = TcpStream::connect(proc.bound_addr()).await;
475        assert!(result.is_ok(), "process should be ready after start");
476
477        drop(proc);
478    }
479
480    #[test]
481    fn test_log_redaction() {
482        let input = b"socks5://user:secret123@127.0.0.1:1080\nhttp://admin:pw0rd@0.0.0.0:8080\n";
483        let redacted = redact_credentials(input);
484
485        assert!(!redacted.contains("secret123"));
486        assert!(!redacted.contains("pw0rd"));
487        assert!(!redacted.contains("user:"), "username must be redacted");
488        assert!(!redacted.contains("admin:"), "username must be redacted");
489        assert!(redacted.contains("****:****@127.0.0.1:1080"));
490        assert!(redacted.contains("****:****@0.0.0.0:8080"));
491    }
492
493    #[test]
494    fn test_log_redaction_no_credentials() {
495        let input = b"socks5://127.0.0.1:1080\nlistening on port 8080\n";
496        let redacted = redact_credentials(input);
497
498        assert_eq!(redacted, String::from_utf8_lossy(input));
499    }
500
501    #[test]
502    fn test_parse_bound_addr() {
503        assert_eq!(
504            parse_bound_addr("Listen: socks5://127.0.0.1:9090"),
505            Some(SocketAddr::new(
506                std::net::IpAddr::V4(std::net::Ipv4Addr::LOCALHOST),
507                9090
508            ))
509        );
510
511        assert_eq!(
512            parse_bound_addr("Listen: http://0.0.0.0:8080"),
513            Some(SocketAddr::new(
514                std::net::IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED),
515                8080
516            ))
517        );
518
519        assert_eq!(parse_bound_addr("no address here"), None);
520        assert_eq!(parse_bound_addr(""), None);
521
522        assert_eq!(
523            parse_bound_addr("Listen: socks5://[::1]:9090"),
524            Some(SocketAddr::new(
525                std::net::IpAddr::V6(std::net::Ipv6Addr::LOCALHOST),
526                9090
527            ))
528        );
529
530        assert_eq!(
531            parse_bound_addr("Listen: http://[::]:8080"),
532            Some(SocketAddr::new(
533                std::net::IpAddr::V6(std::net::Ipv6Addr::UNSPECIFIED),
534                8080
535            ))
536        );
537    }
538
539    #[test]
540    fn test_redact_uri_credentials() {
541        assert_eq!(
542            redact_uri_credentials("socks5://user:pass@host:1080"),
543            "socks5://****:****@host:1080"
544        );
545
546        assert_eq!(
547            redact_uri_credentials("http://admin:secret@proxy:8080"),
548            "http://****:****@proxy:8080"
549        );
550
551        assert_eq!(
552            redact_uri_credentials("socks5://127.0.0.1:1080"),
553            "socks5://127.0.0.1:1080"
554        );
555    }
556
557    #[tokio::test]
558    async fn test_version_detection() {
559        if !require_pproxy() {
560            eprintln!("pproxy not available, skipping");
561            return;
562        }
563
564        let config = OracleConfig::default();
565        let result = verify_pproxy_version(&config).await;
566        assert!(
567            result.is_ok(),
568            "version detection should succeed: {:?}",
569            result.err()
570        );
571    }
572}
eggress_testkit/pproxy_oracle.rs

eggress_testkit/
pproxy_oracle.rs
