Skip to main content

eggress_server/
execute.rs

1use crate::accept::{AcceptedSession, PendingHttpForward, PendingTunnel, PendingUdpAssociate};
2use crate::error::SessionOpenError;
3use crate::reply;
4use crate::ConnectionConfig;
5use eggress_core::chain::{ChainExecutor, HopHandler};
6use eggress_core::connector::DirectConnector;
7use eggress_core::relay::relay;
8use eggress_core::BoxStream;
9use eggress_core::{TargetAddr, TargetHost};
10use eggress_routing::{RouteRequest, SelectedRoute};
11use std::pin::Pin;
12use std::task::{ready, Context, Poll};
13use tokio::io::{AsyncReadExt, AsyncWriteExt};
14
15pub struct SessionReport {
16    pub protocol: Option<String>,
17    pub target: Option<String>,
18    pub route: String,
19    pub bytes_upstream: u64,
20    pub bytes_downstream: u64,
21    pub outcome: SessionOutcome,
22    pub failure: Option<FailureCategory>,
23    pub rule_id: Option<String>,
24    pub upstream_group: Option<String>,
25    pub upstream_id: Option<String>,
26    pub selection_reason: Option<eggress_routing::SelectionReason>,
27}
28
29/// Outcome of a session.
30#[derive(Debug)]
31pub enum SessionOutcome {
32    Completed,
33    ClientProtocolError,
34    AuthenticationFailed,
35    HandshakeTimedOut,
36    RouteFailed,
37    RelayFailed,
38    Cancelled,
39}
40
41/// Specific failure category for structured diagnostics and metrics.
42#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43pub enum FailureCategory {
44    Protocol,
45    Authentication,
46    HandshakeTimeout,
47    Dns,
48    ConnectionRefused,
49    NetworkUnreachable,
50    HostUnreachable,
51    RouteTimeout,
52    RouteHop,
53    UpstreamAuthentication,
54    PolicyDenied,
55    UpstreamUnavailable,
56    Relay,
57    Cancelled,
58    Internal,
59}
60
61impl SessionReport {
62    pub fn open_failed(
63        error: SessionOpenError,
64        protocol: Option<String>,
65        target: Option<String>,
66        route: String,
67    ) -> Self {
68        SessionReport {
69            protocol,
70            target,
71            route,
72            bytes_upstream: 0,
73            bytes_downstream: 0,
74            outcome: SessionOutcome::RouteFailed,
75            failure: Some(FailureCategory::from(&error)),
76            rule_id: None,
77            upstream_group: None,
78            upstream_id: None,
79            selection_reason: None,
80        }
81    }
82
83    pub fn completed(
84        protocol: Option<String>,
85        target: Option<String>,
86        route: String,
87        bytes_upstream: u64,
88        bytes_downstream: u64,
89    ) -> Self {
90        SessionReport {
91            protocol,
92            target,
93            route,
94            bytes_upstream,
95            bytes_downstream,
96            outcome: SessionOutcome::Completed,
97            failure: None,
98            rule_id: None,
99            upstream_group: None,
100            upstream_id: None,
101            selection_reason: None,
102        }
103    }
104
105    pub fn cancelled(protocol: Option<String>, target: Option<String>, route: String) -> Self {
106        SessionReport {
107            protocol,
108            target,
109            route,
110            bytes_upstream: 0,
111            bytes_downstream: 0,
112            outcome: SessionOutcome::Cancelled,
113            failure: Some(FailureCategory::Cancelled),
114            rule_id: None,
115            upstream_group: None,
116            upstream_id: None,
117            selection_reason: None,
118        }
119    }
120
121    pub fn rejected(protocol: Option<String>, target: Option<String>, rule_id: String) -> Self {
122        SessionReport {
123            protocol,
124            target,
125            route: "reject".to_string(),
126            bytes_upstream: 0,
127            bytes_downstream: 0,
128            outcome: SessionOutcome::RouteFailed,
129            failure: Some(FailureCategory::PolicyDenied),
130            rule_id: Some(rule_id),
131            upstream_group: None,
132            upstream_id: None,
133            selection_reason: None,
134        }
135    }
136}
137
138impl From<&SessionOpenError> for FailureCategory {
139    fn from(error: &SessionOpenError) -> Self {
140        match error {
141            SessionOpenError::Dns => FailureCategory::Dns,
142            SessionOpenError::Refused => FailureCategory::ConnectionRefused,
143            SessionOpenError::NetworkUnreachable => FailureCategory::NetworkUnreachable,
144            SessionOpenError::HostUnreachable => FailureCategory::HostUnreachable,
145            SessionOpenError::Timeout => FailureCategory::RouteTimeout,
146            SessionOpenError::UpstreamAuthentication => FailureCategory::UpstreamAuthentication,
147            SessionOpenError::Hop { .. } => FailureCategory::RouteHop,
148            SessionOpenError::PolicyDenied => FailureCategory::PolicyDenied,
149            SessionOpenError::UpstreamUnavailable => FailureCategory::UpstreamUnavailable,
150            SessionOpenError::Other(_) => FailureCategory::Relay,
151        }
152    }
153}
154
155impl FailureCategory {
156    pub fn from_io_error(error: &std::io::Error) -> Self {
157        match error.kind() {
158            std::io::ErrorKind::ConnectionRefused => FailureCategory::ConnectionRefused,
159            std::io::ErrorKind::ConnectionReset => FailureCategory::Relay,
160            std::io::ErrorKind::TimedOut => FailureCategory::Relay,
161            _ => FailureCategory::Relay,
162        }
163    }
164}
165
166/// Execute a session from an accepted connection.
167pub async fn execute(session: AcceptedSession, config: &ConnectionConfig) -> SessionReport {
168    match session {
169        AcceptedSession::Tunnel(pending) => {
170            let protocol = Some(match pending.protocol {
171                crate::accept::TunnelProtocol::HttpConnect => "http".to_string(),
172                crate::accept::TunnelProtocol::Http2 => "h2".to_string(),
173                crate::accept::TunnelProtocol::Http3 => "h3".to_string(),
174                crate::accept::TunnelProtocol::WebSocket => "websocket".to_string(),
175                crate::accept::TunnelProtocol::Socks4 => "socks4".to_string(),
176                crate::accept::TunnelProtocol::Socks5 => "socks5".to_string(),
177                crate::accept::TunnelProtocol::Shadowsocks => "shadowsocks".to_string(),
178                crate::accept::TunnelProtocol::ShadowsocksR => "ssr".to_string(),
179                crate::accept::TunnelProtocol::Trojan => "trojan".to_string(),
180                crate::accept::TunnelProtocol::Raw => "raw".to_string(),
181            });
182            let target = Some(pending.target.to_string());
183            execute_tunnel(pending, config, protocol, target).await
184        }
185        AcceptedSession::HttpForward(pending) => {
186            let target = Some(pending.target.to_string());
187            execute_http_forward(pending, config, target).await
188        }
189        AcceptedSession::UdpAssociate(pending) => execute_udp_associate(pending, config).await,
190        AcceptedSession::Echo(stream) => execute_echo(stream).await,
191    }
192}
193
194async fn execute_echo(mut stream: BoxStream) -> SessionReport {
195    let mut buf = [0u8; 16 * 1024];
196    let mut bytes = 0u64;
197    loop {
198        match stream.read(&mut buf).await {
199            Ok(0) => break,
200            Ok(n) => {
201                bytes += n as u64;
202                if stream.write_all(&buf[..n]).await.is_err() {
203                    break;
204                }
205            }
206            Err(_) => break,
207        }
208    }
209    SessionReport::completed(
210        Some("echo".to_string()),
211        None,
212        "echo".to_string(),
213        bytes,
214        bytes,
215    )
216}
217
218fn route_description(selected: &SelectedRoute) -> String {
219    match selected {
220        SelectedRoute::Direct {
221            selection_reason, ..
222        } => match selection_reason {
223            eggress_routing::SelectionReason::DirectFallback => "direct(fallback)".to_string(),
224            _ => "direct".to_string(),
225        },
226        SelectedRoute::Upstream {
227            upstream, group, ..
228        } => format!("upstream({}/{})", group.0, upstream),
229    }
230}
231
232fn route_metadata(
233    selected: &SelectedRoute,
234) -> (
235    Option<String>,
236    Option<String>,
237    Option<String>,
238    Option<eggress_routing::SelectionReason>,
239) {
240    match selected {
241        SelectedRoute::Direct {
242            decision,
243            selection_reason,
244        } => {
245            let rule_id = match decision {
246                eggress_routing::RouteDecision::Direct { rule, .. }
247                | eggress_routing::RouteDecision::UpstreamGroup { rule, .. }
248                | eggress_routing::RouteDecision::Reject { rule, .. } => rule.0.to_string(),
249            };
250            (Some(rule_id), None, None, Some(*selection_reason))
251        }
252        SelectedRoute::Upstream {
253            decision,
254            group,
255            upstream,
256            selection_reason,
257            ..
258        } => {
259            let rule_id = match decision {
260                eggress_routing::RouteDecision::Direct { rule, .. }
261                | eggress_routing::RouteDecision::UpstreamGroup { rule, .. }
262                | eggress_routing::RouteDecision::Reject { rule, .. } => rule.0.to_string(),
263            };
264            (
265                Some(rule_id),
266                Some(group.0.to_string()),
267                Some(upstream.to_string()),
268                Some(*selection_reason),
269            )
270        }
271    }
272}
273
274struct OpenedRoute {
275    stream: BoxStream,
276    active_lease: Option<eggress_routing::lease::ActiveLease>,
277    route_description: String,
278    rule_id: Option<String>,
279    upstream_group: Option<String>,
280    upstream_id: Option<String>,
281    selection_reason: Option<eggress_routing::SelectionReason>,
282}
283
284fn upstream_protocol_label(chain: &eggress_uri::ProxyChainSpec) -> &'static str {
285    chain
286        .hops
287        .first()
288        .and_then(|h| h.protocols.first())
289        .map(|p| match p {
290            eggress_uri::ProtocolSpec::Http => "http",
291            eggress_uri::ProtocolSpec::HttpOnly => "httponly",
292            eggress_uri::ProtocolSpec::Socks4 => "socks4",
293            eggress_uri::ProtocolSpec::Socks5 => "socks5",
294            eggress_uri::ProtocolSpec::Shadowsocks => "shadowsocks",
295            eggress_uri::ProtocolSpec::ShadowsocksR => "ssr",
296            eggress_uri::ProtocolSpec::Trojan => "trojan",
297            eggress_uri::ProtocolSpec::Http2 => "h2",
298            eggress_uri::ProtocolSpec::Http3 => "h3",
299            eggress_uri::ProtocolSpec::Quic => "quic",
300            eggress_uri::ProtocolSpec::WebSocket => "websocket",
301            eggress_uri::ProtocolSpec::Raw => "raw",
302            eggress_uri::ProtocolSpec::Ssh => "ssh",
303            eggress_uri::ProtocolSpec::Unix => "unix",
304        })
305        .unwrap_or("unknown")
306}
307
308fn failure_reason_label(error: &SessionOpenError) -> &'static str {
309    match error {
310        SessionOpenError::Dns => "dns",
311        SessionOpenError::Refused => "connection_refused",
312        SessionOpenError::NetworkUnreachable => "network_unreachable",
313        SessionOpenError::HostUnreachable => "host_unreachable",
314        SessionOpenError::Timeout => "timeout",
315        SessionOpenError::UpstreamAuthentication => "auth_failed",
316        SessionOpenError::PolicyDenied => "policy_denied",
317        SessionOpenError::Hop { .. } => "handshake",
318        SessionOpenError::UpstreamUnavailable => "upstream_unavailable",
319        SessionOpenError::Other(_) => "io",
320    }
321}
322
323async fn open_route(
324    config: &ConnectionConfig,
325    request: &RouteRequest<'_>,
326) -> Result<OpenedRoute, SessionOpenError> {
327    let selected = config.routing.route(request).map_err(|e| match e {
328        eggress_routing::RouteError::Rejected { .. } => SessionOpenError::PolicyDenied,
329        eggress_routing::RouteError::NoEligibleUpstream(_) => SessionOpenError::PolicyDenied,
330        eggress_routing::RouteError::UnknownGroup(_) => SessionOpenError::PolicyDenied,
331    })?;
332
333    let route = route_description(&selected);
334    let (rule_id, upstream_group, upstream_id, selection_reason) = route_metadata(&selected);
335
336    if let Some(metrics) = &config.metrics {
337        let rule_str = rule_id.as_deref().unwrap_or("default");
338        let action_str = match &selected {
339            SelectedRoute::Direct { .. } => "direct",
340            SelectedRoute::Upstream { .. } => "upstream",
341        };
342        metrics.record_route_decision(rule_str, action_str, "selected");
343    }
344
345    let upstream_protocol = match &selected {
346        SelectedRoute::Upstream { chain, .. } => Some(upstream_protocol_label(chain)),
347        SelectedRoute::Direct { .. } => None,
348    };
349
350    let tls_override = config.tls_client_config.as_ref();
351
352    let result = tokio::time::timeout(config.connect_timeout, async {
353        match selected {
354            SelectedRoute::Direct { .. } => {
355                let bind = config
356                    .local_bind
357                    .as_deref()
358                    .map(|v| {
359                        v.parse().map_err(|e| {
360                            SessionOpenError::Other(format!("invalid local bind '{}': {}", v, e))
361                        })
362                    })
363                    .transpose()?;
364                let stream = DirectConnector
365                    .connect_with_options(
366                        request.target,
367                        &eggress_core::connector::ConnectOptions {
368                            local_bind: bind,
369                            ..Default::default()
370                        },
371                    )
372                    .await?;
373                Ok::<_, SessionOpenError>((stream, None))
374            }
375            SelectedRoute::Upstream {
376                chain,
377                pending_lease,
378                ..
379            } => {
380                #[cfg(feature = "extended")]
381                let shadowsocks_metrics = config.shadowsocks_metrics.clone();
382                #[cfg(not(feature = "extended"))]
383                let shadowsocks_metrics = config.shadowsocks_metrics;
384                #[cfg(feature = "ssh")]
385                let executor = build_chain_executor(
386                    tls_override,
387                    shadowsocks_metrics,
388                    config.ssh_sessions.clone(),
389                );
390                #[cfg(not(feature = "ssh"))]
391                let executor = build_chain_executor(tls_override, shadowsocks_metrics);
392                let stream = executor.execute(&chain.hops, request.target).await?;
393                let active_lease = pending_lease.established();
394                Ok::<_, SessionOpenError>((stream, Some(active_lease)))
395            }
396        }
397    })
398    .await;
399
400    match result {
401        Ok(Ok((stream, active_lease))) => {
402            if let (Some(metrics), Some(protocol)) = (&config.metrics, upstream_protocol) {
403                metrics.record_upstream_open(protocol, "success");
404            }
405            Ok(OpenedRoute {
406                stream,
407                active_lease,
408                route_description: route,
409                rule_id,
410                upstream_group,
411                upstream_id,
412                selection_reason,
413            })
414        }
415        Ok(Err(e)) => {
416            if let (Some(metrics), Some(protocol)) = (&config.metrics, upstream_protocol) {
417                metrics.record_upstream_failure(protocol, failure_reason_label(&e));
418            }
419            Err(e)
420        }
421        Err(_timeout) => {
422            if let Some(metrics) = &config.metrics {
423                if let Some(protocol) = upstream_protocol {
424                    metrics.record_upstream_failure(protocol, "timeout");
425                }
426            }
427            Err(SessionOpenError::Timeout)
428        }
429    }
430}
431
432/// Execute a tunnel session: open route, send success/failure, relay.
433async fn execute_tunnel(
434    mut pending: PendingTunnel,
435    config: &ConnectionConfig,
436    protocol: Option<String>,
437    target: Option<String>,
438) -> SessionReport {
439    tracing::info!("connecting to {}", pending.target);
440
441    let request = RouteRequest {
442        target: &pending.target,
443        source: config.context.source,
444        listener: &config.context.listener,
445        inbound_protocol: match pending.protocol {
446            crate::accept::TunnelProtocol::HttpConnect => eggress_core::ProtocolId::Http,
447            crate::accept::TunnelProtocol::Http2 => eggress_core::ProtocolId::Http2,
448            crate::accept::TunnelProtocol::Http3 => eggress_core::ProtocolId::Http3,
449            crate::accept::TunnelProtocol::WebSocket => eggress_core::ProtocolId::WebSocket,
450            crate::accept::TunnelProtocol::Socks4 => eggress_core::ProtocolId::Socks4,
451            crate::accept::TunnelProtocol::Socks5 => eggress_core::ProtocolId::Socks5,
452            crate::accept::TunnelProtocol::Shadowsocks => eggress_core::ProtocolId::Shadowsocks,
453            crate::accept::TunnelProtocol::ShadowsocksR => eggress_core::ProtocolId::ShadowsocksR,
454            crate::accept::TunnelProtocol::Trojan => eggress_core::ProtocolId::Trojan,
455            crate::accept::TunnelProtocol::Raw => eggress_core::ProtocolId::Raw,
456        },
457        identity: &pending.identity,
458        transport: eggress_routing::TransportKind::Tcp,
459    };
460
461    match open_route(config, &request).await {
462        Ok(opened) => {
463            let route = opened.route_description;
464            let rule_id = opened.rule_id;
465            let upstream_group = opened.upstream_group;
466            let upstream_id = opened.upstream_id;
467            let selection_reason = opened.selection_reason;
468            let _active_lease = opened.active_lease;
469            if let Err(e) = reply::send_tunnel_success(&mut pending, None).await {
470                tracing::debug!("failed to send success reply: {e}");
471                return SessionReport {
472                    protocol,
473                    target,
474                    route,
475                    bytes_upstream: 0,
476                    bytes_downstream: 0,
477                    outcome: SessionOutcome::ClientProtocolError,
478                    failure: Some(FailureCategory::Protocol),
479                    rule_id,
480                    upstream_group,
481                    upstream_id,
482                    selection_reason,
483                };
484            }
485            let result = relay(pending.client, opened.stream).await;
486            tracing::debug!(
487                "relay complete: upstream={}B downstream={}B reason={:?}",
488                result.bytes_upstream,
489                result.bytes_downstream,
490                result.termination_reason
491            );
492            match result.termination_reason {
493                eggress_core::relay::TerminationReason::Error => SessionReport {
494                    protocol,
495                    target,
496                    route,
497                    bytes_upstream: result.bytes_upstream,
498                    bytes_downstream: result.bytes_downstream,
499                    outcome: SessionOutcome::RelayFailed,
500                    failure: Some(FailureCategory::Relay),
501                    rule_id,
502                    upstream_group,
503                    upstream_id,
504                    selection_reason,
505                },
506                _ => SessionReport {
507                    protocol,
508                    target,
509                    route,
510                    bytes_upstream: result.bytes_upstream,
511                    bytes_downstream: result.bytes_downstream,
512                    outcome: SessionOutcome::Completed,
513                    failure: None,
514                    rule_id,
515                    upstream_group,
516                    upstream_id,
517                    selection_reason,
518                },
519            }
520        }
521        Err(SessionOpenError::PolicyDenied) => {
522            let _ = reply::send_tunnel_failure(&mut pending, &SessionOpenError::PolicyDenied).await;
523            SessionReport::rejected(protocol, target, "reject".to_string())
524        }
525        Err(error) => {
526            let _ = reply::send_tunnel_failure(&mut pending, &error).await;
527            SessionReport::open_failed(error, protocol, target, "error".to_string())
528        }
529    }
530}
531
532/// Execute an HTTP forward-proxy session with persistent connection support.
533///
534/// Loops over requests on the client connection, forwarding each to the
535/// appropriate upstream. Supports HTTP/1.1 keep-alive semantics: the
536/// connection persists until the client sends `Connection: close` or the
537/// upstream signals close.
538async fn execute_http_forward(
539    pending: PendingHttpForward,
540    config: &ConnectionConfig,
541    _target: Option<String>,
542) -> SessionReport {
543    tracing::info!("forward proxy to {}", pending.target);
544
545    let mut client = pending.client;
546    let mut total_bytes_upstream: u64 = 0;
547    let mut total_bytes_downstream: u64 = 0;
548    let mut last_target: Option<String>;
549    let mut last_rule_id: Option<String> = None;
550    let mut last_upstream_group: Option<String> = None;
551    let mut last_upstream_id: Option<String> = None;
552    let mut last_selection_reason: Option<eggress_routing::SelectionReason> = None;
553    let mut last_route = String::new();
554
555    // Process the first request (already parsed in pending)
556    let mut request = pending.request;
557    let mut client_close = request.connection_close;
558
559    loop {
560        let target_addr = request.target.clone();
561        last_target = Some(target_addr.to_string());
562
563        if eggress_protocol_http::has_unsupported_expectation(&request.headers) {
564            let _ = reply::send_http_expectation_failed(&mut client).await;
565            return SessionReport {
566                protocol: None,
567                target: last_target,
568                route: last_route,
569                bytes_upstream: total_bytes_upstream,
570                bytes_downstream: total_bytes_downstream,
571                outcome: SessionOutcome::ClientProtocolError,
572                failure: Some(FailureCategory::Protocol),
573                rule_id: last_rule_id,
574                upstream_group: last_upstream_group,
575                upstream_id: last_upstream_id,
576                selection_reason: last_selection_reason,
577            };
578        }
579
580        let route_request = RouteRequest {
581            target: &target_addr,
582            source: config.context.source,
583            listener: &config.context.listener,
584            inbound_protocol: eggress_core::ProtocolId::Http,
585            identity: &pending.identity,
586            transport: eggress_routing::TransportKind::Tcp,
587        };
588
589        match open_route(config, &route_request).await {
590            Ok(mut opened) => {
591                last_route = opened.route_description;
592                last_rule_id = opened.rule_id;
593                last_upstream_group = opened.upstream_group;
594                last_upstream_id = opened.upstream_id;
595                last_selection_reason = opened.selection_reason;
596                let _active_lease = opened.active_lease;
597
598                let origin_req = eggress_protocol_http::build_origin_request(&request);
599                let head_bytes = origin_req.len() as u64;
600
601                if let Err(e) = opened.stream.write_all(origin_req.as_bytes()).await {
602                    let _ = reply::send_http_forward_failure(
603                        &mut client,
604                        &SessionOpenError::Other(e.to_string()),
605                    )
606                    .await;
607                    return SessionReport {
608                        protocol: None,
609                        target: last_target,
610                        route: last_route,
611                        bytes_upstream: total_bytes_upstream,
612                        bytes_downstream: total_bytes_downstream,
613                        outcome: SessionOutcome::RelayFailed,
614                        failure: Some(FailureCategory::Relay),
615                        rule_id: last_rule_id,
616                        upstream_group: last_upstream_group,
617                        upstream_id: last_upstream_id,
618                        selection_reason: last_selection_reason,
619                    };
620                }
621                if let Err(e) = opened.stream.flush().await {
622                    let _ = reply::send_http_forward_failure(
623                        &mut client,
624                        &SessionOpenError::Other(e.to_string()),
625                    )
626                    .await;
627                    return SessionReport {
628                        protocol: None,
629                        target: last_target,
630                        route: last_route,
631                        bytes_upstream: total_bytes_upstream + head_bytes,
632                        bytes_downstream: total_bytes_downstream,
633                        outcome: SessionOutcome::RelayFailed,
634                        failure: Some(FailureCategory::Relay),
635                        rule_id: last_rule_id,
636                        upstream_group: last_upstream_group,
637                        upstream_id: last_upstream_id,
638                        selection_reason: last_selection_reason,
639                    };
640                }
641
642                // The connect timeout ends once the upstream route has been
643                // opened. Uploading a request body is an independent stream
644                // operation and may legitimately take longer.
645                let body_result = async {
646                    let report = eggress_protocol_http::copy_request_body(
647                        &mut client,
648                        &mut opened.stream,
649                        request.body_kind(),
650                        &eggress_protocol_http::BodyCopyLimits::default(),
651                    )
652                    .await?;
653                    opened.stream.flush().await?;
654                    Ok::<_, eggress_protocol_http::HttpError>(report)
655                }
656                .await;
657                let body_report = match body_result {
658                    Ok(report) => report,
659                    Err(_) => {
660                        let _ = opened.stream.shutdown().await;
661                        let _ = client.shutdown().await;
662                        return SessionReport {
663                            protocol: None,
664                            target: last_target,
665                            route: last_route,
666                            bytes_upstream: total_bytes_upstream + head_bytes,
667                            bytes_downstream: total_bytes_downstream,
668                            outcome: SessionOutcome::ClientProtocolError,
669                            failure: Some(FailureCategory::Protocol),
670                            rule_id: last_rule_id,
671                            upstream_group: last_upstream_group,
672                            upstream_id: last_upstream_id,
673                            selection_reason: last_selection_reason,
674                        };
675                    }
676                };
677
678                total_bytes_upstream += head_bytes + body_report.wire_bytes;
679
680                let forward_result =
681                    match eggress_protocol_http::forward_response(&mut opened.stream, &mut client)
682                        .await
683                    {
684                        Ok(result) => result,
685                        Err(eggress_protocol_http::HttpError::UpgradeUnsupported) => {
686                            let _ = reply::send_http_upgrade_unsupported(&mut client).await;
687                            return SessionReport {
688                                protocol: None,
689                                target: last_target,
690                                route: last_route,
691                                bytes_upstream: total_bytes_upstream,
692                                bytes_downstream: total_bytes_downstream,
693                                outcome: SessionOutcome::RelayFailed,
694                                failure: Some(FailureCategory::Protocol),
695                                rule_id: last_rule_id,
696                                upstream_group: last_upstream_group,
697                                upstream_id: last_upstream_id,
698                                selection_reason: last_selection_reason,
699                            };
700                        }
701                        Err(_e) => {
702                            let _ = client.shutdown().await;
703                            return SessionReport {
704                                protocol: None,
705                                target: last_target,
706                                route: last_route,
707                                bytes_upstream: total_bytes_upstream,
708                                bytes_downstream: total_bytes_downstream,
709                                outcome: SessionOutcome::RelayFailed,
710                                failure: Some(FailureCategory::Relay),
711                                rule_id: last_rule_id,
712                                upstream_group: last_upstream_group,
713                                upstream_id: last_upstream_id,
714                                selection_reason: last_selection_reason,
715                            };
716                        }
717                    };
718
719                total_bytes_downstream += forward_result.report.bytes_forwarded;
720
721                // Determine whether to continue the session
722                let should_close = client_close
723                    || forward_result.client_should_close
724                    || !forward_result.upstream_alive;
725
726                if should_close {
727                    break;
728                }
729
730                // Read the next request from the client
731                match eggress_protocol_http::forward_request_stream(&mut client).await {
732                    Ok(next_request) => {
733                        client_close = next_request.connection_close;
734                        request = next_request;
735                    }
736                    Err(eggress_protocol_http::HttpError::Io(ref e))
737                        if e.kind() == std::io::ErrorKind::UnexpectedEof =>
738                    {
739                        // Client closed the connection
740                        break;
741                    }
742                    Err(_) => {
743                        // Malformed next request — close with error
744                        break;
745                    }
746                }
747            }
748            Err(SessionOpenError::PolicyDenied) => {
749                let _ =
750                    reply::send_http_forward_failure(&mut client, &SessionOpenError::PolicyDenied)
751                        .await;
752                return SessionReport::rejected(None, last_target, "reject".to_string());
753            }
754            Err(error) => {
755                let _ = reply::send_http_forward_failure(&mut client, &error).await;
756                return SessionReport::open_failed(error, None, last_target, "error".to_string());
757            }
758        }
759    }
760
761    SessionReport {
762        protocol: None,
763        target: last_target,
764        route: last_route,
765        bytes_upstream: total_bytes_upstream,
766        bytes_downstream: total_bytes_downstream,
767        outcome: SessionOutcome::Completed,
768        failure: None,
769        rule_id: last_rule_id,
770        upstream_group: last_upstream_group,
771        upstream_id: last_upstream_id,
772        selection_reason: last_selection_reason,
773    }
774}
775
776type HandshakeFuture<'a> = std::pin::Pin<
777    Box<
778        dyn std::future::Future<
779                Output = Result<BoxStream, Box<dyn std::error::Error + Send + Sync>>,
780            > + Send
781            + 'a,
782    >,
783>;
784
785async fn execute_udp_associate(
786    pending: PendingUdpAssociate,
787    config: &ConnectionConfig,
788) -> SessionReport {
789    let protocol = Some("socks5".to_string());
790
791    let udp_service = match &config.udp {
792        Some(svc) if svc.is_enabled() => svc,
793        _ => {
794            tracing::debug!("UDP ASSOCIATE rejected: UDP service not available");
795            let mut stream = pending.client;
796            let target = pending.client_hint.unwrap_or(TargetAddr {
797                host: TargetHost::Ip(std::net::IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED)),
798                port: 0,
799            });
800            let socks_addr = target_to_socks_addr(&target);
801            let _ = eggress_protocol_socks::socks5::server::send_connect_reply(
802                &mut stream,
803                eggress_protocol_socks::socks5::server::REP_NOT_ALLOWED,
804                &socks_addr,
805            )
806            .await;
807            return SessionReport {
808                protocol,
809                target: None,
810                route: "udp_associate_disabled".to_string(),
811                bytes_upstream: 0,
812                bytes_downstream: 0,
813                outcome: SessionOutcome::RouteFailed,
814                failure: Some(FailureCategory::Protocol),
815                rule_id: None,
816                upstream_group: None,
817                upstream_id: None,
818                selection_reason: None,
819            };
820        }
821    };
822
823    let client_tcp_peer = config.context.source;
824
825    let gen = config.context.generation;
826
827    let handle = match tokio::time::timeout(
828        config.connect_timeout,
829        udp_service.create_association(
830            &config.context.listener,
831            client_tcp_peer,
832            pending.identity.clone(),
833            gen,
834        ),
835    )
836    .await
837    {
838        Ok(Ok(handle)) => handle,
839        Ok(Err(e)) => {
840            tracing::debug!("UDP ASSOCIATE failed: {e}");
841            let mut stream = pending.client;
842            let target = pending.client_hint.unwrap_or(TargetAddr {
843                host: TargetHost::Ip(std::net::IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED)),
844                port: 0,
845            });
846            let socks_addr = target_to_socks_addr(&target);
847            let _ = eggress_protocol_socks::socks5::server::send_connect_reply(
848                &mut stream,
849                eggress_protocol_socks::socks5::server::REP_GENERAL_FAILURE,
850                &socks_addr,
851            )
852            .await;
853            return SessionReport {
854                protocol,
855                target: None,
856                route: "udp_associate_failed".to_string(),
857                bytes_upstream: 0,
858                bytes_downstream: 0,
859                outcome: SessionOutcome::RouteFailed,
860                failure: Some(FailureCategory::Protocol),
861                rule_id: None,
862                upstream_group: None,
863                upstream_id: None,
864                selection_reason: None,
865            };
866        }
867        Err(_) => {
868            tracing::debug!("UDP ASSOCIATE failed: timeout");
869            let mut stream = pending.client;
870            let target = pending.client_hint.unwrap_or(TargetAddr {
871                host: TargetHost::Ip(std::net::IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED)),
872                port: 0,
873            });
874            let socks_addr = target_to_socks_addr(&target);
875            let _ = eggress_protocol_socks::socks5::server::send_connect_reply(
876                &mut stream,
877                eggress_protocol_socks::socks5::server::REP_GENERAL_FAILURE,
878                &socks_addr,
879            )
880            .await;
881            return SessionReport {
882                protocol,
883                target: None,
884                route: "udp_associate_timeout".to_string(),
885                bytes_upstream: 0,
886                bytes_downstream: 0,
887                outcome: SessionOutcome::HandshakeTimedOut,
888                failure: Some(FailureCategory::RouteTimeout),
889                rule_id: None,
890                upstream_group: None,
891                upstream_id: None,
892                selection_reason: None,
893            };
894        }
895    };
896
897    let relay_ip = handle.relay_addr.ip();
898    let relay_port = handle.relay_addr.port();
899    let socks_addr = match relay_ip {
900        std::net::IpAddr::V4(ip) => {
901            eggress_protocol_socks::socks5::server::SocksAddr::IPv4(ip.octets(), relay_port)
902        }
903        std::net::IpAddr::V6(ip) => {
904            eggress_protocol_socks::socks5::server::SocksAddr::IPv6(ip.octets(), relay_port)
905        }
906    };
907
908    let mut stream = pending.client;
909    if let Err(e) =
910        eggress_protocol_socks::socks5::server::send_udp_associate_reply(&mut stream, &socks_addr)
911            .await
912    {
913        tracing::debug!("failed to send UDP ASSOCIATE reply: {e}");
914        handle.cancel.cancel();
915        return SessionReport {
916            protocol,
917            target: None,
918            route: "udp_associate_reply_failed".to_string(),
919            bytes_upstream: 0,
920            bytes_downstream: 0,
921            outcome: SessionOutcome::ClientProtocolError,
922            failure: Some(FailureCategory::Protocol),
923            rule_id: None,
924            upstream_group: None,
925            upstream_id: None,
926            selection_reason: None,
927        };
928    }
929
930    tracing::info!(
931        association_id = ?handle.id,
932        relay_addr = %handle.relay_addr,
933        "UDP ASSOCIATE established, keeping TCP control connection alive"
934    );
935
936    let mut buf = [0u8; 1];
937    tokio::select! {
938        result = stream.read_exact(&mut buf) => {
939            match result {
940                Ok(_) => {
941                    tracing::debug!(
942                        association_id = ?handle.id,
943                        "TCP control connection closed by client"
944                    );
945                }
946                Err(_) => {
947                    tracing::debug!(
948                        association_id = ?handle.id,
949                        "TCP control connection read failed"
950                    );
951                }
952            }
953        }
954        _ = handle.cancel.cancelled() => {
955            tracing::debug!(
956                association_id = ?handle.id,
957                "UDP association cancelled"
958            );
959        }
960    }
961
962    handle.cancel.cancel();
963
964    SessionReport {
965        protocol,
966        target: None,
967        route: "udp_associate".to_string(),
968        bytes_upstream: 0,
969        bytes_downstream: 0,
970        outcome: SessionOutcome::Completed,
971        failure: None,
972        rule_id: None,
973        upstream_group: None,
974        upstream_id: None,
975        selection_reason: None,
976    }
977}
978
979pub fn build_chain_executor(
980    tls_override: Option<&std::sync::Arc<rustls::ClientConfig>>,
981    #[cfg(feature = "extended")] shadowsocks_metrics: Option<
982        std::sync::Arc<eggress_protocol_shadowsocks::ShadowsocksMetrics>,
983    >,
984    #[cfg(not(feature = "extended"))] _shadowsocks_metrics: Option<()>,
985    #[cfg(feature = "ssh")] ssh_sessions: Option<
986        std::sync::Arc<eggress_transport_ssh::SshSessionCache>,
987    >,
988) -> ChainExecutor {
989    // Build shared TLS client config for upstream hops
990    let shared_tls_config = match tls_override {
991        Some(config) => Some(config.clone()),
992        None => {
993            let builder = eggress_transport_tls::TlsClientConfigBuilder::new();
994            match builder.with_system_roots().and_then(|b| b.build()) {
995                Ok(config) => Some(config),
996                Err(e) => {
997                    tracing::warn!("failed to build shared TLS config: {e}");
998                    None
999                }
1000            }
1001        }
1002    };
1003
1004    #[cfg(feature = "extended")]
1005    let shared_tls_config_arc = shared_tls_config.clone();
1006    #[cfg(not(feature = "extended"))]
1007    let _shared_tls_config_arc = shared_tls_config.clone();
1008
1009    // Per-hop `?insecure` requires an insecure verifier. Build it only when
1010    // the `insecure-tls` feature is available; otherwise per-hop insecure hops
1011    // will be rejected in `ChainExecutor::validate_chain` / `execute` with an
1012    // explicit error. The transport's `with_insecure` is feature-gated, so
1013    // `cargo test` without the feature intentionally leaves this as `None`.
1014    #[cfg(feature = "insecure-tls")]
1015    let insecure_shared_tls_config: Option<std::sync::Arc<rustls::ClientConfig>> =
1016        if tls_override.is_some() {
1017            None
1018        } else {
1019            let builder = eggress_transport_tls::TlsClientConfigBuilder::new();
1020            match builder
1021                .with_system_roots()
1022                .map(|b| b.with_insecure())
1023                .and_then(|b| b.build())
1024            {
1025                Ok(cfg) => Some(cfg),
1026                Err(e) => {
1027                    tracing::debug!("failed to build insecure TLS config: {e}");
1028                    None
1029                }
1030            }
1031        };
1032    #[cfg(not(feature = "insecure-tls"))]
1033    let insecure_shared_tls_config: Option<std::sync::Arc<rustls::ClientConfig>> = None;
1034
1035    let mut handlers: Vec<Box<dyn HopHandler>> = vec![
1036        Box::new(HttpHopHandler),
1037        Box::new(HttpOnlyHopHandler),
1038        Box::new(Socks5HopHandler),
1039        Box::new(Socks4HopHandler),
1040    ];
1041
1042    #[cfg(feature = "extended")]
1043    {
1044        handlers.push(Box::new(ShadowsocksHopHandler {
1045            metrics: shadowsocks_metrics,
1046        }));
1047        handlers.push(Box::new(TrojanHopHandler {
1048            tls_config: shared_tls_config_arc.clone(),
1049            insecure_tls_config: insecure_shared_tls_config.clone(),
1050            tls_override: tls_override.cloned(),
1051        }));
1052        handlers.push(Box::new(WebSocketHopHandler));
1053    }
1054
1055    #[cfg(feature = "pproxy-legacy")]
1056    handlers.push(Box::new(ShadowsocksRHopHandler));
1057
1058    handlers.push(Box::new(RawHopHandler));
1059    handlers.push(Box::new(UnixHopHandler));
1060    #[cfg(feature = "ssh")]
1061    if let Some(sessions) = ssh_sessions {
1062        handlers.push(Box::new(SshHopHandler { sessions }));
1063    }
1064    handlers.push(Box::new(H2HopHandler));
1065
1066    #[cfg(feature = "quic")]
1067    {
1068        handlers.push(Box::new(QuicHopHandler));
1069        handlers.push(Box::new(H3HopHandler));
1070    }
1071
1072    // Pre-build TLS configs per distinct ALPN set so we don't re-read
1073    // and re-parse system roots on every handshake (O-05).
1074    let tls_wrapper_default = shared_tls_config.clone();
1075    let tls_wrapper_h2: Option<std::sync::Arc<rustls::ClientConfig>> = if tls_override.is_none() {
1076        let builder = eggress_transport_tls::TlsClientConfigBuilder::new();
1077        match builder.with_system_roots().and_then(|b| {
1078            b.with_alpn(vec![b"h2".to_vec(), b"http/1.1".to_vec()])
1079                .build()
1080        }) {
1081            Ok(cfg) => Some(cfg),
1082            Err(e) => {
1083                tracing::debug!("failed to build h2 TLS config: {e}");
1084                None
1085            }
1086        }
1087    } else {
1088        None
1089    };
1090    #[cfg(feature = "insecure-tls")]
1091    let insecure_wrapper_default = insecure_shared_tls_config.clone();
1092    #[cfg(not(feature = "insecure-tls"))]
1093    let insecure_wrapper_default: Option<std::sync::Arc<rustls::ClientConfig>> = None;
1094    #[cfg(feature = "insecure-tls")]
1095    let insecure_wrapper_h2: Option<std::sync::Arc<rustls::ClientConfig>> =
1096        if tls_override.is_none() && insecure_shared_tls_config.is_some() {
1097            let builder = eggress_transport_tls::TlsClientConfigBuilder::new();
1098            match builder
1099                .with_system_roots()
1100                .map(|b| b.with_insecure())
1101                .and_then(|b| {
1102                    b.with_alpn(vec![b"h2".to_vec(), b"http/1.1".to_vec()])
1103                        .build()
1104                }) {
1105                Ok(cfg) => Some(cfg),
1106                Err(e) => {
1107                    tracing::debug!("failed to build insecure h2 TLS config: {e}");
1108                    None
1109                }
1110            }
1111        } else {
1112            None
1113        };
1114    #[cfg(not(feature = "insecure-tls"))]
1115    let insecure_wrapper_h2: Option<std::sync::Arc<rustls::ClientConfig>> = None;
1116    fn build_alpn_config(
1117        alpn: Option<Vec<Vec<u8>>>,
1118    ) -> Result<std::sync::Arc<rustls::ClientConfig>, Box<dyn std::error::Error + Send + Sync>>
1119    {
1120        let mut builder = eggress_transport_tls::TlsClientConfigBuilder::new();
1121        builder = builder.with_system_roots()?;
1122        if let Some(protocols) = alpn {
1123            builder = builder.with_alpn(protocols);
1124        }
1125        Ok(builder.build()?)
1126    }
1127    #[cfg(feature = "insecure-tls")]
1128    fn build_insecure_alpn_config(
1129        alpn: Option<Vec<Vec<u8>>>,
1130    ) -> Result<std::sync::Arc<rustls::ClientConfig>, Box<dyn std::error::Error + Send + Sync>>
1131    {
1132        let mut builder = eggress_transport_tls::TlsClientConfigBuilder::new();
1133        builder = builder.with_system_roots()?;
1134        builder = builder.with_insecure();
1135        if let Some(protocols) = alpn {
1136            builder = builder.with_alpn(protocols);
1137        }
1138        Ok(builder.build()?)
1139    }
1140    #[cfg(not(feature = "insecure-tls"))]
1141    fn build_insecure_alpn_config(
1142        _alpn: Option<Vec<Vec<u8>>>,
1143    ) -> Result<std::sync::Arc<rustls::ClientConfig>, Box<dyn std::error::Error + Send + Sync>>
1144    {
1145        Err("insecure TLS requires the insecure-tls feature".into())
1146    }
1147    let tls_wrapper: eggress_core::chain::TlsWrapper =
1148        Box::new(move |stream, server_name, alpn, insecure| {
1149            let default = tls_wrapper_default.clone();
1150            let h2_cfg = tls_wrapper_h2.clone();
1151            let insecure_default = insecure_wrapper_default.clone();
1152            let insecure_h2_cfg = insecure_wrapper_h2.clone();
1153            Box::pin(async move {
1154                let config = if insecure {
1155                    match insecure_default.clone() {
1156                        Some(c) => {
1157                            if let Some(ref protocols) = alpn {
1158                                if c.alpn_protocols == *protocols {
1159                                    c
1160                                } else if let Some(h2) = insecure_h2_cfg.clone() {
1161                                    if *protocols == vec![b"h2".to_vec(), b"http/1.1".to_vec()] {
1162                                        h2
1163                                    } else {
1164                                        build_insecure_alpn_config(Some(protocols.clone()))?
1165                                    }
1166                                } else {
1167                                    build_insecure_alpn_config(Some(protocols.clone()))?
1168                                }
1169                            } else {
1170                                c
1171                            }
1172                        }
1173                        None => build_insecure_alpn_config(alpn)?,
1174                    }
1175                } else {
1176                    match default {
1177                        Some(c) => {
1178                            if let Some(ref protocols) = alpn {
1179                                if c.alpn_protocols == *protocols {
1180                                    c
1181                                } else if let Some(h2) = h2_cfg {
1182                                    if *protocols == vec![b"h2".to_vec(), b"http/1.1".to_vec()] {
1183                                        h2
1184                                    } else {
1185                                        build_alpn_config(Some(protocols.clone()))?
1186                                    }
1187                                } else {
1188                                    build_alpn_config(Some(protocols.clone()))?
1189                                }
1190                            } else {
1191                                c
1192                            }
1193                        }
1194                        None => build_alpn_config(alpn)?,
1195                    }
1196                };
1197                eggress_transport_tls::tls_connect(stream, config, &server_name)
1198                    .await
1199                    .map_err(|e| -> Box<dyn std::error::Error + Send + Sync> { Box::new(e) as _ })
1200            })
1201        });
1202
1203    ChainExecutor::new(handlers)
1204        .with_tls_wrapper(tls_wrapper)
1205        .with_shared_tls_config(shared_tls_config)
1206        .with_insecure_shared_tls_config(insecure_shared_tls_config)
1207}
1208
1209/// Adapts an origin-form request into the absolute-form request expected by
1210/// pproxy's `httponly` upstream mode. The adapter is deliberately limited to
1211/// the request headers; bodies are passed through unchanged.
1212struct HttpOnlyStream {
1213    inner: BoxStream,
1214    target: TargetAddr,
1215    pending: Vec<u8>,
1216    rewritten: bool,
1217}
1218
1219/// Cap on bytes buffered while the httponly upstream is stalled. Beyond this,
1220/// `poll_write` exerts backpressure instead of growing memory without bound.
1221const HTTPONLY_MAX_BUFFERED: usize = 64 * 1024;
1222
1223impl tokio::io::AsyncRead for HttpOnlyStream {
1224    fn poll_read(
1225        mut self: Pin<&mut Self>,
1226        cx: &mut Context<'_>,
1227        buf: &mut tokio::io::ReadBuf<'_>,
1228    ) -> Poll<std::io::Result<()>> {
1229        Pin::new(&mut self.inner).poll_read(cx, buf)
1230    }
1231}
1232
1233impl tokio::io::AsyncWrite for HttpOnlyStream {
1234    fn poll_write(
1235        mut self: Pin<&mut Self>,
1236        cx: &mut Context<'_>,
1237        data: &[u8],
1238    ) -> Poll<std::io::Result<usize>> {
1239        // Drain buffered bytes into the upstream first; if it is stalled,
1240        // propagating `Pending` from the flush engages relay backpressure
1241        // instead of buffering without bound.
1242        if !self.pending.is_empty() {
1243            ready!(self.as_mut().poll_flush(cx))?;
1244        }
1245        let room = HTTPONLY_MAX_BUFFERED.saturating_sub(self.pending.len());
1246        if room == 0 && !data.is_empty() {
1247            // Backpressure: pending is full after flush attempt, signal
1248            // Pending so the caller (which must use `write_all`/retry loop)
1249            // waits for the waker instead of busy-spinning on Ok(0).
1250            return Poll::Pending;
1251        }
1252        let accepted = data.len().min(room);
1253        self.pending.extend_from_slice(&data[..accepted]);
1254        // A short write is valid `AsyncWrite` behavior: callers retry with
1255        // the remainder once the buffered bytes have drained. Callers must
1256        // use `write_all` or explicit retry; a single `write` may return
1257        // short and the tail must be retried by the caller.
1258        Poll::Ready(Ok(accepted))
1259    }
1260
1261    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
1262        if !self.rewritten
1263            && !self.pending.is_empty()
1264            && self.pending.windows(4).any(|window| window == b"\r\n\r\n")
1265        {
1266            self.pending = rewrite_request_head(&self.pending, &self.target);
1267            self.rewritten = true;
1268        }
1269        // The stream type is Unpin, so plain field access keeps the borrows
1270        // of `inner` and `pending` disjoint inside the drain loop.
1271        let this = self.get_mut();
1272        while !this.pending.is_empty() {
1273            match Pin::new(&mut this.inner).poll_write(cx, &this.pending) {
1274                Poll::Ready(Ok(0)) => {
1275                    return Poll::Ready(Err(std::io::Error::new(
1276                        std::io::ErrorKind::WriteZero,
1277                        "httponly upstream accepted zero bytes",
1278                    )));
1279                }
1280                Poll::Ready(Ok(n)) => {
1281                    this.pending.drain(..n);
1282                }
1283                Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
1284                Poll::Pending => return Poll::Pending,
1285            }
1286        }
1287        Pin::new(&mut this.inner).poll_flush(cx)
1288    }
1289
1290    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
1291        ready!(self.as_mut().poll_flush(cx))?;
1292        Pin::new(&mut self.inner).poll_shutdown(cx)
1293    }
1294}
1295
1296/// Rewrites an origin-form request head into the absolute-form request
1297/// expected by pproxy's `httponly` upstream mode. Only the request line
1298/// changes; every other byte — including all header line terminators and
1299/// the body — is preserved exactly as received. Input without a complete,
1300/// rewritable request head is returned unchanged.
1301fn rewrite_request_head(data: &[u8], target: &TargetAddr) -> Vec<u8> {
1302    let Some(pos) = data.windows(4).position(|w| w == b"\r\n\r\n") else {
1303        // Head not complete yet; leave everything untouched.
1304        return data.to_vec();
1305    };
1306    let end = pos + 4;
1307    let head = &data[..end];
1308    let mut rewritten = Vec::with_capacity(data.len() + 32);
1309    if let Some(nl) = head.iter().position(|b| *b == b'\n') {
1310        let raw_first = &head[..nl];
1311        let first = raw_first.strip_suffix(b"\r").unwrap_or(raw_first);
1312        if let Some(space) = first.iter().position(|b| *b == b' ') {
1313            if let Some(second) = first[space + 1..].iter().position(|b| *b == b' ') {
1314                let method = &first[..space];
1315                let path = &first[space + 1..space + 1 + second];
1316                if path.starts_with(b"/") {
1317                    rewritten.extend_from_slice(method);
1318                    rewritten.extend_from_slice(b" http://");
1319                    rewritten.extend_from_slice(target.to_string().as_bytes());
1320                    rewritten.extend_from_slice(path);
1321                    rewritten.extend_from_slice(&first[space + 1 + second..]);
1322                    // Restore the original request-line terminator, then
1323                    // copy every remaining header byte verbatim; only the
1324                    // request line itself changes.
1325                    if raw_first.len() != first.len() {
1326                        rewritten.push(b'\r');
1327                    }
1328                    rewritten.extend_from_slice(&head[nl..]);
1329                }
1330            }
1331        }
1332    }
1333    if rewritten.is_empty() {
1334        rewritten.extend_from_slice(head);
1335    }
1336    rewritten.extend_from_slice(&data[end..]);
1337    rewritten
1338}
1339
1340struct HttpOnlyHopHandler;
1341
1342impl HopHandler for HttpOnlyHopHandler {
1343    fn protocol(&self) -> eggress_uri::ProtocolSpec {
1344        eggress_uri::ProtocolSpec::HttpOnly
1345    }
1346    fn handshake<'a>(
1347        &'a self,
1348        stream: BoxStream,
1349        target: &'a TargetAddr,
1350        _hop: &'a eggress_uri::ProxyHopSpec,
1351        _hop_index: usize,
1352    ) -> HandshakeFuture<'a> {
1353        let target = target.clone();
1354        Box::pin(async move {
1355            Ok(Box::new(HttpOnlyStream {
1356                inner: stream,
1357                target,
1358                pending: Vec::new(),
1359                rewritten: false,
1360            }) as BoxStream)
1361        })
1362    }
1363}
1364
1365struct HttpHopHandler;
1366
1367impl HopHandler for HttpHopHandler {
1368    fn protocol(&self) -> eggress_uri::ProtocolSpec {
1369        eggress_uri::ProtocolSpec::Http
1370    }
1371
1372    fn handshake<'a>(
1373        &'a self,
1374        stream: BoxStream,
1375        target: &'a TargetAddr,
1376        hop: &'a eggress_uri::ProxyHopSpec,
1377        _hop_index: usize,
1378    ) -> HandshakeFuture<'a> {
1379        let auth = hop
1380            .credentials
1381            .as_ref()
1382            .map(|c| (c.username.as_str(), c.password.as_str()));
1383        Box::pin(async move {
1384            eggress_protocol_http::http_connect(stream, target, auth, &Default::default())
1385                .await
1386                .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)
1387        })
1388    }
1389}
1390
1391struct Socks5HopHandler;
1392
1393impl HopHandler for Socks5HopHandler {
1394    fn protocol(&self) -> eggress_uri::ProtocolSpec {
1395        eggress_uri::ProtocolSpec::Socks5
1396    }
1397
1398    fn handshake<'a>(
1399        &'a self,
1400        stream: BoxStream,
1401        target: &'a TargetAddr,
1402        hop: &'a eggress_uri::ProxyHopSpec,
1403        _hop_index: usize,
1404    ) -> HandshakeFuture<'a> {
1405        let socks_addr = target_to_socks_addr(target);
1406        let auth = hop
1407            .credentials
1408            .as_ref()
1409            .map(|c| (c.username.as_str(), c.password.as_str()));
1410        Box::pin(async move {
1411            eggress_protocol_socks::socks5::client::socks5_connect(stream, &socks_addr, auth)
1412                .await
1413                .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)
1414        })
1415    }
1416}
1417
1418struct Socks4HopHandler;
1419
1420impl HopHandler for Socks4HopHandler {
1421    fn protocol(&self) -> eggress_uri::ProtocolSpec {
1422        eggress_uri::ProtocolSpec::Socks4
1423    }
1424
1425    fn handshake<'a>(
1426        &'a self,
1427        stream: BoxStream,
1428        target: &'a TargetAddr,
1429        hop: &'a eggress_uri::ProxyHopSpec,
1430        _hop_index: usize,
1431    ) -> HandshakeFuture<'a> {
1432        let user_id = hop.credentials.as_ref().map(|c| c.username.as_str());
1433        Box::pin(async move {
1434            eggress_protocol_socks::socks4_connect(stream, target, user_id)
1435                .await
1436                .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)
1437        })
1438    }
1439}
1440
1441#[cfg(feature = "extended")]
1442struct ShadowsocksHopHandler {
1443    metrics: Option<std::sync::Arc<eggress_protocol_shadowsocks::ShadowsocksMetrics>>,
1444}
1445
1446#[cfg(feature = "extended")]
1447impl HopHandler for ShadowsocksHopHandler {
1448    fn protocol(&self) -> eggress_uri::ProtocolSpec {
1449        eggress_uri::ProtocolSpec::Shadowsocks
1450    }
1451
1452    fn handshake<'a>(
1453        &'a self,
1454        stream: BoxStream,
1455        target: &'a TargetAddr,
1456        hop: &'a eggress_uri::ProxyHopSpec,
1457        _hop_index: usize,
1458    ) -> HandshakeFuture<'a> {
1459        let metrics = self.metrics.clone();
1460        Box::pin(async move {
1461            let creds = hop.credentials.as_ref().ok_or_else(|| {
1462                Box::new(eggress_protocol_shadowsocks::ShadowsocksError::Other(
1463                    "shadowsocks requires credentials (method:password)".to_string(),
1464                )) as Box<dyn std::error::Error + Send + Sync>
1465            })?;
1466
1467            match eggress_protocol_shadowsocks::CipherMethod::parse_method(&creds.username) {
1468                Ok(method) => eggress_protocol_shadowsocks::shadowsocks_connect(
1469                    stream,
1470                    target,
1471                    method,
1472                    &creds.password,
1473                    metrics,
1474                )
1475                .await
1476                .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>),
1477                Err(modern_error) => {
1478                    #[cfg(feature = "legacy-crypto")]
1479                    if let Ok(legacy_method) =
1480                        eggress_protocol_shadowsocks::legacy::LegacyMethod::parse(&creds.username)
1481                    {
1482                        return eggress_protocol_shadowsocks::legacy::legacy_connect(
1483                            stream,
1484                            target,
1485                            legacy_method,
1486                            creds.password.as_bytes(),
1487                        )
1488                        .await
1489                        .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>);
1490                    }
1491                    if let Some(m) = metrics.as_ref() {
1492                        m.record_tcp_unsupported_method_reject();
1493                    }
1494                    Err(Box::new(modern_error) as Box<dyn std::error::Error + Send + Sync>)
1495                }
1496            }
1497        })
1498    }
1499}
1500
1501#[cfg(feature = "pproxy-legacy")]
1502struct ShadowsocksRHopHandler;
1503
1504#[cfg(feature = "pproxy-legacy")]
1505impl HopHandler for ShadowsocksRHopHandler {
1506    fn protocol(&self) -> eggress_uri::ProtocolSpec {
1507        eggress_uri::ProtocolSpec::ShadowsocksR
1508    }
1509
1510    fn handshake<'a>(
1511        &'a self,
1512        stream: BoxStream,
1513        target: &'a TargetAddr,
1514        hop: &'a eggress_uri::ProxyHopSpec,
1515        _hop_index: usize,
1516    ) -> HandshakeFuture<'a> {
1517        Box::pin(async move {
1518            let plugins = eggress_protocol_shadowsocks::compat::plugin::parse_plugins(&hop.plugins)
1519                .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?;
1520            eggress_protocol_shadowsocks::compat::ssr::ssr_connect(
1521                stream,
1522                target,
1523                &eggress_protocol_shadowsocks::compat::ssr::SsrConfig {
1524                    auth_prefix: hop.auth_prefix.as_deref().map(str::as_bytes).map(Vec::from),
1525                    plugins,
1526                },
1527            )
1528            .await
1529            .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)
1530        })
1531    }
1532}
1533
1534#[cfg(feature = "extended")]
1535struct TrojanHopHandler {
1536    tls_config: Option<std::sync::Arc<rustls::ClientConfig>>,
1537    insecure_tls_config: Option<std::sync::Arc<rustls::ClientConfig>>,
1538    tls_override: Option<std::sync::Arc<rustls::ClientConfig>>,
1539}
1540
1541#[cfg(feature = "extended")]
1542impl HopHandler for TrojanHopHandler {
1543    fn protocol(&self) -> eggress_uri::ProtocolSpec {
1544        eggress_uri::ProtocolSpec::Trojan
1545    }
1546
1547    fn handshake<'a>(
1548        &'a self,
1549        stream: BoxStream,
1550        target: &'a TargetAddr,
1551        hop: &'a eggress_uri::ProxyHopSpec,
1552        _hop_index: usize,
1553    ) -> HandshakeFuture<'a> {
1554        let tls_config = self.tls_config.clone();
1555        let insecure_tls_config = self.insecure_tls_config.clone();
1556        let tls_override = self.tls_override.clone();
1557        let insecure = hop.insecure;
1558        let password = hop.credentials.as_ref().map(|c| c.password.clone());
1559        let server_name = hop
1560            .server_name
1561            .clone()
1562            .unwrap_or_else(|| hop.endpoint.host.clone());
1563        Box::pin(async move {
1564            let password = password.ok_or_else(|| {
1565                Box::new(eggress_protocol_trojan::TrojanError::Protocol(
1566                    "trojan requires credentials (password)".to_string(),
1567                )) as Box<dyn std::error::Error + Send + Sync>
1568            })?;
1569
1570            let chosen = if insecure {
1571                if let Some(ovr) = tls_override.clone() {
1572                    // Global override is already insecure in tests that set it;
1573                    // reuse it for per-hop insecure when available.
1574                    Some(ovr)
1575                } else {
1576                    insecure_tls_config.clone().or(tls_config.clone())
1577                }
1578            } else {
1579                tls_config.clone()
1580            };
1581
1582            eggress_protocol_trojan::trojan_connect(stream, target, &password, &server_name, chosen)
1583                .await
1584                .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)
1585        })
1586    }
1587}
1588
1589#[cfg(feature = "extended")]
1590struct WebSocketHopHandler;
1591
1592#[cfg(feature = "extended")]
1593impl HopHandler for WebSocketHopHandler {
1594    fn protocol(&self) -> eggress_uri::ProtocolSpec {
1595        eggress_uri::ProtocolSpec::WebSocket
1596    }
1597
1598    fn handshake<'a>(
1599        &'a self,
1600        stream: BoxStream,
1601        _target: &'a TargetAddr,
1602        hop: &'a eggress_uri::ProxyHopSpec,
1603        _hop_index: usize,
1604    ) -> HandshakeFuture<'a> {
1605        let use_tls = hop.tls;
1606        let scheme = if use_tls { "wss" } else { "ws" };
1607        let url = format!("{}://{}:{}", scheme, hop.endpoint.host, hop.endpoint.port);
1608        Box::pin(async move {
1609            let client = eggress_protocol_websocket::WebSocketTunnelClient::with_default_config();
1610            client
1611                .connect_over_stream(&url, stream)
1612                .await
1613                .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)
1614        })
1615    }
1616}
1617
1618struct RawHopHandler;
1619
1620impl HopHandler for RawHopHandler {
1621    fn protocol(&self) -> eggress_uri::ProtocolSpec {
1622        eggress_uri::ProtocolSpec::Raw
1623    }
1624
1625    fn handshake<'a>(
1626        &'a self,
1627        stream: BoxStream,
1628        _target: &'a TargetAddr,
1629        _hop: &'a eggress_uri::ProxyHopSpec,
1630        _hop_index: usize,
1631    ) -> HandshakeFuture<'a> {
1632        Box::pin(async move { Ok(stream) })
1633    }
1634}
1635
1636#[cfg(feature = "ssh")]
1637struct SshHopHandler {
1638    sessions: std::sync::Arc<eggress_transport_ssh::SshSessionCache>,
1639}
1640
1641#[cfg(feature = "ssh")]
1642impl HopHandler for SshHopHandler {
1643    fn protocol(&self) -> eggress_uri::ProtocolSpec {
1644        eggress_uri::ProtocolSpec::Ssh
1645    }
1646
1647    fn handshake<'a>(
1648        &'a self,
1649        stream: BoxStream,
1650        target: &'a TargetAddr,
1651        hop: &'a eggress_uri::ProxyHopSpec,
1652        hop_index: usize,
1653    ) -> HandshakeFuture<'a> {
1654        let sessions = self.sessions.clone();
1655        let target = target.clone();
1656        let endpoint = hop.endpoint.clone();
1657        let credentials = hop.credentials.clone();
1658        Box::pin(async move {
1659            let credentials = credentials.ok_or_else(|| {
1660                Box::new(eggress_transport_ssh::SshTransportError::MissingUsername)
1661                    as Box<dyn std::error::Error + Send + Sync>
1662            })?;
1663            if credentials.username.is_empty() {
1664                return Err(
1665                    Box::new(eggress_transport_ssh::SshTransportError::MissingUsername)
1666                        as Box<dyn std::error::Error + Send + Sync>,
1667                );
1668            }
1669            let auth = if let Some(path) = credentials.password.strip_prefix(':') {
1670                if path.is_empty() {
1671                    return Err(Box::new(
1672                        eggress_transport_ssh::SshTransportError::EmptyPrivateKeyPath,
1673                    )
1674                        as Box<dyn std::error::Error + Send + Sync>);
1675                }
1676                eggress_transport_ssh::SshAuth::PrivateKey(path.to_string())
1677            } else {
1678                eggress_transport_ssh::SshAuth::Password(credentials.password)
1679            };
1680            let key = eggress_transport_ssh::SshSessionKey {
1681                host: endpoint.host,
1682                port: endpoint.port,
1683                username: credentials.username,
1684                auth,
1685                hop_index,
1686            };
1687            // Format once (O-03); previously each branch formatted separately.
1688            let target_host = target.host.to_string();
1689            let result = if target.port == 0 {
1690                sessions.open_unix_channel(key, stream, &target_host).await
1691            } else {
1692                sessions
1693                    .open_tcp_channel(key, stream, &target_host, target.port)
1694                    .await
1695            };
1696            result.map_err(|error| Box::new(error) as Box<dyn std::error::Error + Send + Sync>)
1697        })
1698    }
1699}
1700
1701struct UnixHopHandler;
1702
1703impl HopHandler for UnixHopHandler {
1704    fn protocol(&self) -> eggress_uri::ProtocolSpec {
1705        eggress_uri::ProtocolSpec::Unix
1706    }
1707
1708    fn handshake<'a>(
1709        &'a self,
1710        stream: BoxStream,
1711        _target: &'a TargetAddr,
1712        _hop: &'a eggress_uri::ProxyHopSpec,
1713        _hop_index: usize,
1714    ) -> HandshakeFuture<'a> {
1715        Box::pin(async move { Ok(stream) })
1716    }
1717}
1718
1719struct H2HopHandler;
1720
1721#[cfg(feature = "quic")]
1722struct QuicHopHandler;
1723
1724#[cfg(feature = "quic")]
1725impl HopHandler for QuicHopHandler {
1726    fn protocol(&self) -> eggress_uri::ProtocolSpec {
1727        eggress_uri::ProtocolSpec::Quic
1728    }
1729
1730    fn open<'a>(
1731        &'a self,
1732        endpoint: &'a eggress_uri::EndpointSpec,
1733        hop: &'a eggress_uri::ProxyHopSpec,
1734        _target: &'a TargetAddr,
1735    ) -> Option<HandshakeFuture<'a>> {
1736        let endpoint = endpoint.clone();
1737        let server_name = hop
1738            .server_name
1739            .clone()
1740            .unwrap_or_else(|| endpoint.host.clone());
1741        Some(Box::pin(async move {
1742            let client = eggress_transport_quic::QuicClient::connect(
1743                &endpoint.host,
1744                endpoint.port,
1745                eggress_transport_quic::QuicClientConfig {
1746                    server_name,
1747                    insecure: hop.insecure,
1748                    alpn_protocols: Vec::new(),
1749                    ..Default::default()
1750                },
1751            )
1752            .await?;
1753            client
1754                .open_stream()
1755                .await
1756                .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)
1757        }))
1758    }
1759
1760    fn handshake<'a>(
1761        &'a self,
1762        stream: BoxStream,
1763        _target: &'a TargetAddr,
1764        _hop: &'a eggress_uri::ProxyHopSpec,
1765        _hop_index: usize,
1766    ) -> HandshakeFuture<'a> {
1767        Box::pin(async move { Ok(stream) })
1768    }
1769}
1770
1771#[cfg(feature = "quic")]
1772struct H3HopHandler;
1773
1774#[cfg(feature = "quic")]
1775impl HopHandler for H3HopHandler {
1776    fn protocol(&self) -> eggress_uri::ProtocolSpec {
1777        eggress_uri::ProtocolSpec::Http3
1778    }
1779
1780    fn open<'a>(
1781        &'a self,
1782        endpoint: &'a eggress_uri::EndpointSpec,
1783        hop: &'a eggress_uri::ProxyHopSpec,
1784        target: &'a TargetAddr,
1785    ) -> Option<HandshakeFuture<'a>> {
1786        let endpoint = endpoint.clone();
1787        let target = target.clone();
1788        let server_name = hop
1789            .server_name
1790            .clone()
1791            .unwrap_or_else(|| endpoint.host.clone());
1792        let authorization = hop
1793            .credentials
1794            .as_ref()
1795            .map(|credentials| (credentials.username.clone(), credentials.password.clone()));
1796        Some(Box::pin(async move {
1797            let client = eggress_transport_quic::QuicClient::connect(
1798                &endpoint.host,
1799                endpoint.port,
1800                eggress_transport_quic::QuicClientConfig {
1801                    server_name,
1802                    insecure: hop.insecure,
1803                    alpn_protocols: vec![b"h3".to_vec()],
1804                    ..Default::default()
1805                },
1806            )
1807            .await?;
1808            eggress_protocol_h3::H3Client::new(client, authorization)
1809                .connect(&target)
1810                .await
1811                .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)
1812        }))
1813    }
1814
1815    fn handshake<'a>(
1816        &'a self,
1817        stream: BoxStream,
1818        _target: &'a TargetAddr,
1819        _hop: &'a eggress_uri::ProxyHopSpec,
1820        _hop_index: usize,
1821    ) -> HandshakeFuture<'a> {
1822        Box::pin(async move { Ok(stream) })
1823    }
1824}
1825
1826/// Wrapper that holds an H2PoolGuard alongside the bidirectional stream,
1827/// ensuring the pooled connection is released back to the pool only when
1828/// the stream is dropped.
1829struct PooledH2Stream {
1830    inner:
1831        tokio::io::Join<eggress_protocol_http::H2StreamRead, eggress_protocol_http::H2StreamWrite>,
1832    _guard: eggress_protocol_http::H2PoolGuard,
1833}
1834
1835impl tokio::io::AsyncRead for PooledH2Stream {
1836    fn poll_read(
1837        mut self: std::pin::Pin<&mut Self>,
1838        cx: &mut std::task::Context<'_>,
1839        buf: &mut tokio::io::ReadBuf<'_>,
1840    ) -> std::task::Poll<std::io::Result<()>> {
1841        std::pin::Pin::new(&mut self.inner).poll_read(cx, buf)
1842    }
1843}
1844
1845impl tokio::io::AsyncWrite for PooledH2Stream {
1846    fn poll_write(
1847        mut self: std::pin::Pin<&mut Self>,
1848        cx: &mut std::task::Context<'_>,
1849        buf: &[u8],
1850    ) -> std::task::Poll<std::io::Result<usize>> {
1851        std::pin::Pin::new(&mut self.inner).poll_write(cx, buf)
1852    }
1853
1854    fn poll_flush(
1855        mut self: std::pin::Pin<&mut Self>,
1856        cx: &mut std::task::Context<'_>,
1857    ) -> std::task::Poll<std::io::Result<()>> {
1858        std::pin::Pin::new(&mut self.inner).poll_flush(cx)
1859    }
1860
1861    fn poll_shutdown(
1862        mut self: std::pin::Pin<&mut Self>,
1863        cx: &mut std::task::Context<'_>,
1864    ) -> std::task::Poll<std::io::Result<()>> {
1865        std::pin::Pin::new(&mut self.inner).poll_shutdown(cx)
1866    }
1867}
1868
1869impl HopHandler for H2HopHandler {
1870    fn protocol(&self) -> eggress_uri::ProtocolSpec {
1871        eggress_uri::ProtocolSpec::Http2
1872    }
1873
1874    fn handshake<'a>(
1875        &'a self,
1876        stream: BoxStream,
1877        target: &'a TargetAddr,
1878        hop: &'a eggress_uri::ProxyHopSpec,
1879        hop_index: usize,
1880    ) -> HandshakeFuture<'a> {
1881        let endpoint_host = hop.endpoint.host.clone();
1882        let endpoint_port = hop.endpoint.port;
1883        let auth = hop
1884            .credentials
1885            .as_ref()
1886            .map(|c| (c.username.clone(), c.password.clone()));
1887        let target_clone = target.clone();
1888        let pool_key = eggress_protocol_http::H2PoolKey::with_hop_index(
1889            &endpoint_host,
1890            endpoint_port,
1891            hop.tls,
1892            hop.server_name.as_deref(),
1893            auth.as_ref().map(|(u, p)| (u.as_str(), p.as_str())),
1894            hop_index,
1895        );
1896        Box::pin(async move {
1897            let stream: BoxStream = stream;
1898
1899            let auth_ref = auth.as_ref().map(|(u, p)| (u.as_str(), p.as_str()));
1900            let (send_stream, recv_stream, guard) =
1901                eggress_protocol_http::h2_connect_client_pooled(
1902                    stream,
1903                    &target_clone,
1904                    auth_ref,
1905                    &pool_key,
1906                )
1907                .await
1908                .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?;
1909
1910            let h2_write = eggress_protocol_http::H2StreamWrite::new(send_stream);
1911            let h2_read = eggress_protocol_http::H2StreamRead::new(recv_stream);
1912
1913            let pooled = PooledH2Stream {
1914                inner: tokio::io::join(h2_read, h2_write),
1915                _guard: guard,
1916            };
1917            Ok(Box::new(pooled) as BoxStream)
1918        })
1919    }
1920}
1921
1922fn target_to_socks_addr(target: &TargetAddr) -> eggress_protocol_socks::socks5::server::SocksAddr {
1923    use eggress_protocol_socks::socks5::server::SocksAddr;
1924    match &target.host {
1925        TargetHost::Ip(std::net::IpAddr::V4(ip)) => SocksAddr::IPv4(ip.octets(), target.port),
1926        TargetHost::Ip(std::net::IpAddr::V6(ip)) => SocksAddr::IPv6(ip.octets(), target.port),
1927        TargetHost::Domain(d) => SocksAddr::Domain(d.clone(), target.port),
1928    }
1929}
1930
1931#[cfg(test)]
1932mod tests {
1933    use super::*;
1934
1935    fn http_only_target() -> TargetAddr {
1936        TargetAddr {
1937            host: TargetHost::Domain("target.example".into()),
1938            port: 8080,
1939        }
1940    }
1941
1942    #[test]
1943    fn httponly_rewrite_preserves_header_terminators() {
1944        let request = b"GET /path HTTP/1.1\r\nHost: example.com\r\nX-Foo: bar\r\n\r\nbody";
1945        let rewritten = rewrite_request_head(request, &http_only_target());
1946        assert_eq!(
1947            std::str::from_utf8(&rewritten).unwrap(),
1948            "GET http://target.example:8080/path HTTP/1.1\r\n\
1949             Host: example.com\r\n\
1950             X-Foo: bar\r\n\
1951             \r\n\
1952             body"
1953        );
1954    }
1955
1956    #[test]
1957    fn httponly_rewrite_preserves_mixed_line_endings() {
1958        // Request line terminated by bare LF while the head ends with CRLF.
1959        let request = b"GET /path HTTP/1.1\nHost: example.com\r\nX-Foo: bar\r\n\r\n";
1960        let rewritten = rewrite_request_head(request, &http_only_target());
1961        assert_eq!(
1962            std::str::from_utf8(&rewritten).unwrap(),
1963            "GET http://target.example:8080/path HTTP/1.1\n\
1964             Host: example.com\r\n\
1965             X-Foo: bar\r\n\
1966             \r\n"
1967        );
1968    }
1969
1970    #[test]
1971    fn httponly_rewrite_waits_for_complete_head() {
1972        // No `\r\n\r\n` terminator yet — unchanged (flush retries later).
1973        let partial = b"GET /path HTTP/1.1\r\nHost: example.com\r\n".as_slice();
1974        assert_eq!(rewrite_request_head(partial, &http_only_target()), partial);
1975    }
1976
1977    #[test]
1978    fn httponly_rewrite_leaves_absolute_form_and_incomplete_heads_alone() {
1979        // Absolute-form request line is not origin-form; unchanged.
1980        let absolute =
1981            b"GET http://example.com/path HTTP/1.1\r\nHost: example.com\r\n\r\n".as_slice();
1982        assert_eq!(
1983            rewrite_request_head(absolute, &http_only_target()),
1984            absolute
1985        );
1986        // No complete head yet; unchanged (flush will retry later).
1987        let partial = b"GET /path HTTP/1.1\r\nHost: example.com\r\n".as_slice();
1988        assert_eq!(rewrite_request_head(partial, &http_only_target()), partial);
1989        // Empty input stays empty.
1990        assert!(rewrite_request_head(b"", &http_only_target()).is_empty());
1991    }
1992
1993    #[tokio::test]
1994    async fn httponly_stream_rewrites_once_and_drains_on_shutdown() {
1995        // A tiny duplex buffer forces the flush loop across multiple polls;
1996        // the remaining bytes still contain `\r\n\r\n`, which a second
1997        // rewrite pass would mangle ("X-Foo: /bar baz" looks like an
1998        // origin-form request line to the rewriter).
1999        let (mut peer, inner) = tokio::io::duplex(16);
2000        let mut stream = HttpOnlyStream {
2001            inner: Box::new(inner),
2002            target: http_only_target(),
2003            pending: Vec::new(),
2004            rewritten: false,
2005        };
2006        let request =
2007            b"GET /path HTTP/1.1\r\nHost: example.com\r\nX-Foo: /bar baz\r\n\r\ntail".as_slice();
2008        let expected =
2009            b"GET http://target.example:8080/path HTTP/1.1\r\nHost: example.com\r\nX-Foo: /bar baz\r\n\r\ntail";
2010        let reader = tokio::spawn(async move {
2011            let mut received = Vec::new();
2012            peer.read_to_end(&mut received).await.unwrap();
2013            received
2014        });
2015        use tokio::io::AsyncWriteExt;
2016        stream.write_all(request).await.unwrap();
2017        stream.shutdown().await.unwrap();
2018        let received = reader.await.unwrap();
2019        assert_eq!(received, expected);
2020    }
2021}