tenant.openagents/omega

1//! The sans-IO proxy handshake state machine.
2
3use std::collections::VecDeque;
4use std::net::{IpAddr, SocketAddr};
5
6use base64::Engine as _;
7
8use crate::{Credentials, MAX_HTTP_RESPONSE_LENGTH, ProxyError, ProxyScheme, ProxySpec, Target};
9
10const MAX_HTTP_RESPONSE_HEADERS: usize = 64;
11
12const SOCKS4_VERSION: u8 = 0x04;
13const SOCKS5_VERSION: u8 = 0x05;
14const SOCKS_COMMAND_CONNECT: u8 = 0x01;
15const SOCKS4_REPLY_GRANTED: u8 = 0x5A;
16const SOCKS5_METHOD_NO_AUTHENTICATION: u8 = 0x00;
17const SOCKS5_METHOD_USERNAME_PASSWORD: u8 = 0x02;
18const SOCKS5_ADDRESS_IPV4: u8 = 0x01;
19const SOCKS5_ADDRESS_DOMAIN: u8 = 0x03;
20const SOCKS5_ADDRESS_IPV6: u8 = 0x04;
21
22/// A proxy handshake in progress.
23///
24/// Sans-IO: the handshake never performs I/O itself. Drive it by calling
25/// [`advance`](Self::advance) with bytes received from the proxy and acting
26/// on the returned [`Step`]s until it reports [`Step::Done`].
27pub struct Handshake {
28    exchanges: VecDeque<Exchange>,
29    input: Vec<u8>,
30}
31
32/// What the caller must do next to make progress on the handshake.
33#[derive(Debug, PartialEq, Eq)]
34pub enum Step {
35    /// Write these bytes to the proxy.
36    Send(Vec<u8>),
37    /// Read more bytes from the proxy and pass them to the next `advance`.
38    NeedMoreInput,
39    /// The tunnel is established. `leftover` holds any bytes the proxy sent
40    /// past the end of the handshake; they belong to the tunneled stream and
41    /// must be consumed before reading from the transport again.
42    Done { leftover: Vec<u8> },
43}
44
45enum Exchange {
46    Send(Vec<u8>),
47    Expect(Expectation),
48}
49
50enum Expectation {
51    HttpResponseHead,
52    Socks5MethodSelection { offered: u8 },
53    Socks5AuthReply,
54    Socks5ConnectReply,
55    Socks4ConnectReply,
56}
57
58impl Handshake {
59    /// Plans the handshake for tunneling to `target` through the proxy
60    /// described by `spec`.
61    ///
62    /// Passing a [`Target::Domain`] delegates DNS resolution to the proxy
63    /// regardless of scheme (SOCKS4 proxies are sent the 4a extension);
64    /// resolve the host locally and pass a [`Target::Address`] when
65    /// [`ProxySpec::remote_dns`] is false, to follow the proxy URL's intent.
66    pub fn new(spec: &ProxySpec, target: &Target) -> Result<Self, ProxyError> {
67        // The domain is spliced into protocol messages (the CONNECT request
68        // line, NUL-terminated SOCKS4a fields), so control bytes would let a
69        // hostile target string forge protocol structure. No valid host name
70        // contains them; fail closed rather than trusting the caller to have
71        // validated.
72        if let Target::Domain(domain, _) = target
73            && domain
74                .bytes()
75                .any(|byte| byte.is_ascii_control() || byte == 0x7F)
76        {
77            return Err(ProxyError::DomainContainsControlBytes);
78        }
79
80        let mut exchanges = VecDeque::new();
81        match spec.scheme {
82            ProxyScheme::Http { .. } => {
83                exchanges.push_back(Exchange::Send(http_connect_request(
84                    target,
85                    spec.credentials.as_ref(),
86                )));
87                exchanges.push_back(Exchange::Expect(Expectation::HttpResponseHead));
88            }
89            ProxyScheme::Socks4 { .. } => {
90                exchanges.push_back(Exchange::Send(socks4_connect_request(
91                    target,
92                    spec.credentials.as_ref(),
93                )?));
94                exchanges.push_back(Exchange::Expect(Expectation::Socks4ConnectReply));
95            }
96            ProxyScheme::Socks5 { .. } => {
97                let method = if spec.credentials.is_some() {
98                    SOCKS5_METHOD_USERNAME_PASSWORD
99                } else {
100                    SOCKS5_METHOD_NO_AUTHENTICATION
101                };
102                exchanges.push_back(Exchange::Send(vec![SOCKS5_VERSION, 1, method]));
103                exchanges.push_back(Exchange::Expect(Expectation::Socks5MethodSelection {
104                    offered: method,
105                }));
106                if let Some(credentials) = &spec.credentials {
107                    exchanges.push_back(Exchange::Send(socks5_auth_request(credentials)?));
108                    exchanges.push_back(Exchange::Expect(Expectation::Socks5AuthReply));
109                }
110                exchanges.push_back(Exchange::Send(socks5_connect_request(target)?));
111                exchanges.push_back(Exchange::Expect(Expectation::Socks5ConnectReply));
112            }
113        }
114        Ok(Self {
115            exchanges,
116            input: Vec::new(),
117        })
118    }
119
120    /// Advances the state machine with bytes newly received from the proxy
121    /// (empty on the first call, and whenever there is nothing new to feed).
122    ///
123    /// Callers must not feed the same bytes twice: partial input is buffered
124    /// internally until it can be parsed.
125    pub fn advance(&mut self, received: &[u8]) -> Result<Step, ProxyError> {
126        self.input.extend_from_slice(received);
127        loop {
128            match self.exchanges.pop_front() {
129                None => {
130                    return Ok(Step::Done {
131                        leftover: std::mem::take(&mut self.input),
132                    });
133                }
134                Some(Exchange::Send(bytes)) => return Ok(Step::Send(bytes)),
135                Some(Exchange::Expect(expectation)) => {
136                    match try_parse(&expectation, &self.input)? {
137                        Some(consumed) => {
138                            self.input.drain(..consumed);
139                        }
140                        None => {
141                            self.exchanges.push_front(Exchange::Expect(expectation));
142                            return Ok(Step::NeedMoreInput);
143                        }
144                    }
145                }
146            }
147        }
148    }
149}
150
151/// Attempts to parse one expected proxy message from the buffered input.
152/// Returns the number of bytes consumed, or `None` when more input is needed.
153fn try_parse(expectation: &Expectation, input: &[u8]) -> Result<Option<usize>, ProxyError> {
154    match expectation {
155        Expectation::HttpResponseHead => parse_http_response_head(input),
156        Expectation::Socks5MethodSelection { offered } => {
157            let Some(reply) = input.get(..2) else {
158                return Ok(None);
159            };
160            if reply[0] != SOCKS5_VERSION {
161                return Err(ProxyError::UnexpectedSocksVersion(reply[0]));
162            }
163            if reply[1] != *offered {
164                return Err(ProxyError::AuthMethodRejected { selected: reply[1] });
165            }
166            Ok(Some(2))
167        }
168        Expectation::Socks5AuthReply => {
169            let Some(reply) = input.get(..2) else {
170                return Ok(None);
171            };
172            // RFC 1929: any non-zero status is failure.
173            if reply[1] != 0x00 {
174                return Err(ProxyError::AuthenticationFailed);
175            }
176            Ok(Some(2))
177        }
178        Expectation::Socks5ConnectReply => {
179            let Some(header) = input.get(..4) else {
180                return Ok(None);
181            };
182            if header[0] != SOCKS5_VERSION {
183                return Err(ProxyError::UnexpectedSocksVersion(header[0]));
184            }
185            if header[1] != 0x00 {
186                return Err(ProxyError::Socks5ConnectRefused(header[1]));
187            }
188            let address_length = match header[3] {
189                SOCKS5_ADDRESS_IPV4 => 4,
190                SOCKS5_ADDRESS_IPV6 => 16,
191                SOCKS5_ADDRESS_DOMAIN => match input.get(4) {
192                    Some(length) => 1 + *length as usize,
193                    None => return Ok(None),
194                },
195                other => return Err(ProxyError::UnknownAddressType(other)),
196            };
197            // The reply's bound address and port only matter for BIND, so
198            // they are consumed and discarded.
199            let total = 4 + address_length + 2;
200            if input.len() < total {
201                return Ok(None);
202            }
203            Ok(Some(total))
204        }
205        Expectation::Socks4ConnectReply => {
206            if input.len() < 8 {
207                return Ok(None);
208            }
209            if input[1] != SOCKS4_REPLY_GRANTED {
210                return Err(ProxyError::Socks4ConnectRefused(input[1]));
211            }
212            Ok(Some(8))
213        }
214    }
215}
216
217fn parse_http_response_head(input: &[u8]) -> Result<Option<usize>, ProxyError> {
218    let mut headers = [httparse::EMPTY_HEADER; MAX_HTTP_RESPONSE_HEADERS];
219    let mut response = httparse::Response::new(&mut headers);
220    match response.parse(input) {
221        Ok(httparse::Status::Complete(consumed)) => {
222            // The cap applies to complete heads too: a driver that reads in
223            // large chunks could otherwise hand a well-formed multi-megabyte
224            // head to a single `advance` and have it accepted.
225            if consumed > MAX_HTTP_RESPONSE_LENGTH {
226                return Err(ProxyError::HttpResponseTooLarge);
227            }
228            let code = response.code.ok_or_else(|| {
229                ProxyError::MalformedHttpResponse("missing status code".to_string())
230            })?;
231            if !(200..300).contains(&code) {
232                return Err(ProxyError::HttpConnectRefused(code));
233            }
234            Ok(Some(consumed))
235        }
236        Ok(httparse::Status::Partial) => {
237            if input.len() > MAX_HTTP_RESPONSE_LENGTH {
238                Err(ProxyError::HttpResponseTooLarge)
239            } else {
240                Ok(None)
241            }
242        }
243        Err(error) => Err(ProxyError::MalformedHttpResponse(error.to_string())),
244    }
245}
246
247fn http_connect_request(target: &Target, credentials: Option<&Credentials>) -> Vec<u8> {
248    let host = match target {
249        Target::Domain(domain, port) => format!("{domain}:{port}"),
250        // `SocketAddr`'s `Display` brackets IPv6 addresses, matching the
251        // `host:port` form CONNECT requires.
252        Target::Address(address) => address.to_string(),
253    };
254    let mut request =
255        format!("CONNECT {host} HTTP/1.1\r\nHost: {host}\r\nProxy-Connection: Keep-Alive\r\n");
256    if let Some(Credentials { username, password }) = credentials {
257        let encoded = base64::prelude::BASE64_STANDARD.encode(format!("{username}:{password}"));
258        request.push_str(&format!("Proxy-Authorization: Basic {encoded}\r\n"));
259    }
260    request.push_str("\r\n");
261    request.into_bytes()
262}
263
264fn socks5_auth_request(credentials: &Credentials) -> Result<Vec<u8>, ProxyError> {
265    let username = credentials.username.as_bytes();
266    let password = credentials.password.as_bytes();
267    if username.len() > 255 || password.len() > 255 {
268        return Err(ProxyError::CredentialsTooLong);
269    }
270    let mut request = Vec::with_capacity(3 + username.len() + password.len());
271    // RFC 1929 subnegotiation version.
272    request.push(0x01);
273    request.push(username.len() as u8);
274    request.extend_from_slice(username);
275    request.push(password.len() as u8);
276    request.extend_from_slice(password);
277    Ok(request)
278}
279
280fn socks5_connect_request(target: &Target) -> Result<Vec<u8>, ProxyError> {
281    let mut request = vec![SOCKS5_VERSION, SOCKS_COMMAND_CONNECT, 0x00];
282    let port = match target {
283        Target::Domain(domain, port) => {
284            let domain = domain.as_bytes();
285            if domain.len() > 255 {
286                return Err(ProxyError::DomainTooLong);
287            }
288            request.push(SOCKS5_ADDRESS_DOMAIN);
289            request.push(domain.len() as u8);
290            request.extend_from_slice(domain);
291            *port
292        }
293        Target::Address(address) => {
294            match address.ip() {
295                IpAddr::V4(ip) => {
296                    request.push(SOCKS5_ADDRESS_IPV4);
297                    request.extend_from_slice(&ip.octets());
298                }
299                IpAddr::V6(ip) => {
300                    request.push(SOCKS5_ADDRESS_IPV6);
301                    request.extend_from_slice(&ip.octets());
302                }
303            }
304            address.port()
305        }
306    };
307    request.extend_from_slice(&port.to_be_bytes());
308    Ok(request)
309}
310
311fn socks4_connect_request(
312    target: &Target,
313    credentials: Option<&Credentials>,
314) -> Result<Vec<u8>, ProxyError> {
315    let user_id = credentials
316        .map(|credentials| credentials.username.as_str())
317        .unwrap_or("");
318    // The user id field is NUL-terminated on the wire, so an embedded NUL
319    // would truncate it and shift the fields that follow.
320    if user_id.as_bytes().contains(&0x00) {
321        return Err(ProxyError::Socks4UserIdContainsNul);
322    }
323    let mut request = vec![SOCKS4_VERSION, SOCKS_COMMAND_CONNECT];
324    match target {
325        Target::Address(SocketAddr::V4(address)) => {
326            request.extend_from_slice(&address.port().to_be_bytes());
327            request.extend_from_slice(&address.ip().octets());
328            request.extend_from_slice(user_id.as_bytes());
329            request.push(0x00);
330        }
331        Target::Address(SocketAddr::V6(_)) => return Err(ProxyError::Socks4Ipv4Only),
332        Target::Domain(domain, port) => {
333            // SOCKS4a: the invalid destination address 0.0.0.1 signals that
334            // a host name follows the user id.
335            request.extend_from_slice(&port.to_be_bytes());
336            request.extend_from_slice(&[0, 0, 0, 1]);
337            request.extend_from_slice(user_id.as_bytes());
338            request.push(0x00);
339            request.extend_from_slice(domain.as_bytes());
340            request.push(0x00);
341        }
342    }
343    Ok(request)
344}
345
346#[cfg(test)]
347mod tests {
348    use proptest::prelude::*;
349
350    use super::*;
351
352    #[test]
353    fn http_connect_completes_on_200_and_preserves_leftover() {
354        let server_bytes = b"HTTP/1.1 200 Connection Established\r\n\r\ntunnel-bytes";
355        let (sends, leftover) = run(
356            "http://proxy:8080",
357            &domain_target(),
358            server_bytes,
359            &[server_bytes.len()],
360        )
361        .unwrap();
362
363        assert_eq!(sends.len(), 1);
364        let request = String::from_utf8(sends[0].clone()).unwrap();
365        assert!(
366            request.starts_with("CONNECT cloud.example.com:443 HTTP/1.1\r\n"),
367            "unexpected request: {request:?}"
368        );
369        assert!(request.contains("Host: cloud.example.com:443\r\n"));
370        assert!(request.ends_with("\r\n\r\n"));
371        assert_eq!(leftover, b"tunnel-bytes");
372    }
373
374    #[test]
375    fn http_connect_sends_basic_authorization() {
376        let (sends, _) = run(
377            "http://user:pass@proxy:8080",
378            &domain_target(),
379            b"HTTP/1.1 200 OK\r\n\r\n",
380            &[64],
381        )
382        .unwrap();
383
384        let request = String::from_utf8(sends[0].clone()).unwrap();
385        // "dXNlcjpwYXNz" is base64("user:pass").
386        assert!(
387            request.contains("Proxy-Authorization: Basic dXNlcjpwYXNz\r\n"),
388            "unexpected request: {request:?}"
389        );
390    }
391
392    #[test]
393    fn http_connect_fails_on_error_status() {
394        let error = run(
395            "http://proxy:8080",
396            &domain_target(),
397            b"HTTP/1.1 407 Proxy Authentication Required\r\nProxy-Authenticate: Basic\r\n\r\n",
398            &[64],
399        )
400        .unwrap_err();
401        assert!(matches!(error, ProxyError::HttpConnectRefused(407)));
402    }
403
404    #[test]
405    fn http_connect_rejects_oversized_response() {
406        // A single never-ending header keeps the response incomplete (and
407        // under httparse's header-count limit) until the size cap trips.
408        let mut server_bytes = b"HTTP/1.1 200 OK\r\nX-Padding: ".to_vec();
409        server_bytes.resize(MAX_HTTP_RESPONSE_LENGTH + 1024, b'a');
410        let error = run("http://proxy:8080", &domain_target(), &server_bytes, &[64]).unwrap_err();
411        assert!(matches!(error, ProxyError::HttpResponseTooLarge));
412    }
413
414    #[test]
415    fn http_connect_rejects_oversized_complete_response_in_one_chunk() {
416        // A *well-formed* head over the cap, delivered whole in a single
417        // `advance`, must be rejected too: the cap can't depend on the
418        // driver's read chunking.
419        let mut server_bytes = b"HTTP/1.1 200 OK\r\nX-Padding: ".to_vec();
420        server_bytes.resize(MAX_HTTP_RESPONSE_LENGTH + 1024, b'a');
421        server_bytes.extend_from_slice(b"\r\n\r\n");
422        let error = run(
423            "http://proxy:8080",
424            &domain_target(),
425            &server_bytes,
426            &[server_bytes.len()],
427        )
428        .unwrap_err();
429        assert!(matches!(error, ProxyError::HttpResponseTooLarge));
430    }
431
432    #[test]
433    fn rejects_domains_containing_control_bytes() {
434        for scheme in ["http", "socks4a", "socks5h"] {
435            for domain in ["evil.com\r\nX-Injected: x", "evil.com\0.example.com"] {
436                let spec = spec(&format!("{scheme}://proxy:1080"));
437                let target = Target::Domain(domain.to_string(), 443);
438                let Err(error) = Handshake::new(&spec, &target) else {
439                    panic!("expected {scheme} handshake construction to fail for {domain:?}");
440                };
441                assert!(matches!(error, ProxyError::DomainContainsControlBytes));
442            }
443        }
444    }
445
446    #[test]
447    fn socks4_rejects_user_id_containing_nul() {
448        let spec = spec("socks4://user%00id@proxy:1080");
449        let target = Target::Address("192.0.2.10:443".parse().unwrap());
450        let Err(error) = Handshake::new(&spec, &target) else {
451            panic!("expected handshake construction to fail");
452        };
453        assert!(matches!(error, ProxyError::Socks4UserIdContainsNul));
454    }
455
456    #[test]
457    fn socks5_connect_with_domain_target() {
458        let mut server_bytes = vec![0x05, 0x00];
459        server_bytes.extend_from_slice(&[0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0]);
460        server_bytes.extend_from_slice(b"tunnel");
461        let (sends, leftover) = run(
462            "socks5h://proxy:1080",
463            &domain_target(),
464            &server_bytes,
465            &[server_bytes.len()],
466        )
467        .unwrap();
468
469        assert_eq!(sends[0], vec![0x05, 0x01, 0x00]);
470        let mut expected_connect = vec![0x05, 0x01, 0x00, 0x03, 17];
471        expected_connect.extend_from_slice(b"cloud.example.com");
472        expected_connect.extend_from_slice(&443u16.to_be_bytes());
473        assert_eq!(sends[1], expected_connect);
474        assert_eq!(leftover, b"tunnel");
475    }
476
477    #[test]
478    fn socks5_connect_with_username_password() {
479        let mut server_bytes = vec![0x05, 0x02];
480        server_bytes.extend_from_slice(&[0x01, 0x00]);
481        server_bytes.extend_from_slice(&[0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0]);
482        let (sends, _) = run(
483            "socks5h://user:pass@proxy:1080",
484            &domain_target(),
485            &server_bytes,
486            &[server_bytes.len()],
487        )
488        .unwrap();
489
490        assert_eq!(sends[0], vec![0x05, 0x01, 0x02]);
491        assert_eq!(sends[1], b"\x01\x04user\x04pass".to_vec());
492        assert_eq!(sends.len(), 3);
493    }
494
495    #[test]
496    fn socks5_reports_connection_refusal() {
497        let mut server_bytes = vec![0x05, 0x00];
498        server_bytes.extend_from_slice(&[0x05, 0x02, 0x00, 0x01, 0, 0, 0, 0, 0, 0]);
499        let error = run(
500            "socks5h://proxy:1080",
501            &domain_target(),
502            &server_bytes,
503            &[server_bytes.len()],
504        )
505        .unwrap_err();
506        assert!(matches!(error, ProxyError::Socks5ConnectRefused(0x02)));
507    }
508
509    #[test]
510    fn socks5_reports_rejected_authentication_method() {
511        let error = run(
512            "socks5h://proxy:1080",
513            &domain_target(),
514            &[0x05, 0xFF],
515            &[2],
516        )
517        .unwrap_err();
518        assert!(matches!(
519            error,
520            ProxyError::AuthMethodRejected { selected: 0xFF }
521        ));
522    }
523
524    #[test]
525    fn socks5_encodes_ip_targets() {
526        let mut server_bytes = vec![0x05, 0x00];
527        server_bytes.extend_from_slice(&[0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0]);
528
529        let target = Target::Address("192.0.2.10:443".parse().unwrap());
530        let (sends, _) = run("socks5://proxy:1080", &target, &server_bytes, &[64]).unwrap();
531        assert_eq!(
532            sends[1],
533            vec![0x05, 0x01, 0x00, 0x01, 192, 0, 2, 10, 0x01, 0xBB]
534        );
535
536        let target = Target::Address("[2001:db8::1]:443".parse().unwrap());
537        let (sends, _) = run("socks5://proxy:1080", &target, &server_bytes, &[64]).unwrap();
538        assert_eq!(sends[1][3], SOCKS5_ADDRESS_IPV6);
539        assert_eq!(sends[1].len(), 4 + 16 + 2);
540    }
541
542    #[test]
543    fn socks5_reply_with_domain_bound_address_is_consumed() {
544        let mut server_bytes = vec![0x05, 0x00];
545        server_bytes.extend_from_slice(&[0x05, 0x00, 0x00, 0x03, 4]);
546        server_bytes.extend_from_slice(b"host");
547        server_bytes.extend_from_slice(&[0, 80]);
548        server_bytes.extend_from_slice(b"rest");
549        let (_, leftover) = run(
550            "socks5h://proxy:1080",
551            &domain_target(),
552            &server_bytes,
553            &[server_bytes.len()],
554        )
555        .unwrap();
556        assert_eq!(leftover, b"rest");
557    }
558
559    #[test]
560    fn socks4_connect_with_ipv4_target() {
561        let target = Target::Address("192.0.2.10:443".parse().unwrap());
562        let (sends, leftover) = run(
563            "socks4://proxy:1080",
564            &target,
565            &[0x00, 0x5A, 0, 0, 0, 0, 0, 0],
566            &[8],
567        )
568        .unwrap();
569        assert_eq!(sends[0], vec![0x04, 0x01, 0x01, 0xBB, 192, 0, 2, 10, 0x00]);
570        assert_eq!(leftover, b"");
571    }
572
573    #[test]
574    fn socks4a_connect_with_domain_target_and_user_id() {
575        let (sends, _) = run(
576            "socks4a://userid@proxy:1080",
577            &domain_target(),
578            &[0x00, 0x5A, 0, 0, 0, 0, 0, 0],
579            &[8],
580        )
581        .unwrap();
582        let mut expected = vec![0x04, 0x01, 0x01, 0xBB, 0, 0, 0, 1];
583        expected.extend_from_slice(b"userid\x00");
584        expected.extend_from_slice(b"cloud.example.com\x00");
585        assert_eq!(sends[0], expected);
586    }
587
588    #[test]
589    fn socks4_rejects_ipv6_targets() {
590        let spec = spec("socks4://proxy:1080");
591        let target = Target::Address("[2001:db8::1]:443".parse().unwrap());
592        let Err(error) = Handshake::new(&spec, &target) else {
593            panic!("expected handshake construction to fail");
594        };
595        assert!(matches!(error, ProxyError::Socks4Ipv4Only));
596    }
597
598    #[test]
599    fn socks4_reports_rejection() {
600        let target = Target::Address("192.0.2.10:443".parse().unwrap());
601        let error = run(
602            "socks4://proxy:1080",
603            &target,
604            &[0x00, 0x5B, 0, 0, 0, 0, 0, 0],
605            &[8],
606        )
607        .unwrap_err();
608        assert!(matches!(error, ProxyError::Socks4ConnectRefused(0x5B)));
609    }
610
611    proptest! {
612        /// The outcome of a handshake must not depend on how the proxy's
613        /// bytes are chunked across reads.
614        #[test]
615        fn chunking_does_not_change_the_outcome(
616            chunk_sizes in proptest::collection::vec(1usize..16, 1..32),
617        ) {
618            for (url, server_bytes) in transcripts() {
619                let baseline = run(url, &domain_target(), &server_bytes, &[server_bytes.len()]);
620                let chunked = run(url, &domain_target(), &server_bytes, &chunk_sizes);
621                match (baseline, chunked) {
622                    (Ok(expected), Ok(actual)) => prop_assert_eq!(expected, actual),
623                    (Err(_), Err(_)) => {}
624                    (baseline, chunked) => {
625                        return Err(TestCaseError::fail(format!(
626                            "outcomes diverged: baseline {baseline:?} vs chunked {chunked:?}"
627                        )));
628                    }
629                }
630            }
631        }
632    }
633
634    /// Scripted proxy conversations used by the chunking property test, each
635    /// with trailing bytes to pin leftover handling.
636    fn transcripts() -> Vec<(&'static str, Vec<u8>)> {
637        let mut socks5 = vec![0x05, 0x02];
638        socks5.extend_from_slice(&[0x01, 0x00]);
639        socks5.extend_from_slice(&[0x05, 0x00, 0x00, 0x03, 4]);
640        socks5.extend_from_slice(b"host");
641        socks5.extend_from_slice(&[0, 80]);
642        socks5.extend_from_slice(b"leftover-bytes");
643
644        vec![
645            (
646                "http://user:pass@proxy:8080",
647                b"HTTP/1.1 200 Connection Established\r\nVia: test\r\n\r\nleftover-bytes".to_vec(),
648            ),
649            ("socks5h://user:pass@proxy:1080", socks5),
650            (
651                "socks4a://userid@proxy:1080",
652                b"\x00\x5A\x00\x00\x00\x00\x00\x00leftover-bytes".to_vec(),
653            ),
654        ]
655    }
656
657    fn spec(url: &str) -> ProxySpec {
658        ProxySpec::parse(&url.parse().unwrap()).unwrap()
659    }
660
661    fn domain_target() -> Target {
662        Target::Domain("cloud.example.com".to_string(), 443)
663    }
664
665    /// Drives a handshake to completion against scripted proxy bytes,
666    /// feeding input in the given chunk sizes (cycled). Returns everything
667    /// the client sent and the leftover bytes reported at completion.
668    fn run(
669        url: &str,
670        target: &Target,
671        server_bytes: &[u8],
672        chunk_sizes: &[usize],
673    ) -> Result<(Vec<Vec<u8>>, Vec<u8>), ProxyError> {
674        let mut handshake = Handshake::new(&spec(url), target)?;
675        let mut sends = Vec::new();
676        let mut remaining = server_bytes;
677        let mut sizes = chunk_sizes.iter().copied().cycle();
678        let mut pending: &[u8] = &[];
679        loop {
680            let step = handshake.advance(pending)?;
681            pending = &[];
682            match step {
683                Step::Send(bytes) => sends.push(bytes),
684                Step::NeedMoreInput => {
685                    assert!(
686                        !remaining.is_empty(),
687                        "handshake wants more input than the transcript provides"
688                    );
689                    let size = sizes
690                        .next()
691                        .expect("chunk_sizes must not be empty")
692                        .clamp(1, remaining.len());
693                    let (chunk, rest) = remaining.split_at(size);
694                    pending = chunk;
695                    remaining = rest;
696                }
697                Step::Done { leftover } => {
698                    let mut leftover = leftover;
699                    // Any transcript bytes not yet fed to the machine are
700                    // also tunnel bytes; account for them so results don't
701                    // depend on chunking.
702                    leftover.extend_from_slice(remaining);
703                    return Ok((sends, leftover));
704                }
705            }
706        }
707    }
708}
709