tenant.openagents/omega

1//! Per-connection logic: parse the first request, decide allowed/denied,
2//! either pump bytes through to the upstream or close (with an explanatory
3//! 511 for policy denials).
4//!
5//! Both CONNECT and HTTP forward go through here. After the policy decision,
6//! the TCP connection is pinned to the approved destination — directly, or
7//! via a CONNECT tunnel through the upstream proxy — and everything becomes
8//! opaque byte pumping. We don't parse anything else (chunked encoding,
9//! keep-alive HTTP requests after the first, etc.); because the whole
10//! connection can only reach the one approved host, later requests the
11//! client sends on it cannot escape the policy decision. Per-TCP-connection
12//! event granularity, by design.
13
14use crate::pinned_host::{PinnedHost, PinnedHostError};
15use crate::proxy::{
16    DenyReason, ProxyEvent, RequestMethod, RequestOutcome, RuntimeState, UpstreamProxy,
17};
18use anyhow::{Context, Result, anyhow, bail};
19use base64::Engine as _;
20use std::io::{Read, Write};
21use std::net::{IpAddr, Shutdown, SocketAddr, TcpStream, ToSocketAddrs as _};
22#[cfg(unix)]
23use std::os::unix::net::UnixStream;
24use std::sync::Arc;
25use std::thread;
26use std::time::{Duration, Instant};
27use url::Url;
28
29/// Buffer size for each direction of bidir copy. 64 KiB is the sweet spot
30/// for most networks — large enough to keep the pipe full, small enough
31/// not to balloon memory under many concurrent connections.
32const PUMP_BUFFER_SIZE: usize = 64 * 1024;
33
34/// Cap on request/response header bytes. The proxy runs inside the editor
35/// process and its sole client is model-driven code — exactly the party the
36/// sandbox distrusts — so an unbounded header read would let a malicious
37/// command balloon the editor's memory. 64 KiB is far beyond what real HTTP
38/// clients send.
39const MAX_HEADER_BYTES: usize = 64 * 1024;
40
41/// How long to wait for the client's request headers. Pooled connections
42/// that never send a request get closed; well-behaved clients retry on a
43/// fresh connection. Cleared before the pump phase so long-lived idle
44/// tunnels (long polls, slow downloads) are unaffected.
45const HEADER_READ_TIMEOUT: Duration = Duration::from_secs(60);
46
47/// Timeout for outbound TCP connects (direct or to the upstream proxy),
48/// so a black-holed destination doesn't pin a connection thread for the
49/// OS default (~75s or more).
50const CONNECT_TIMEOUT: Duration = Duration::from_secs(30);
51
52/// How long to wait for the upstream proxy's CONNECT response.
53const UPSTREAM_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(30);
54
55/// Top-level entry from the listener thread. Owns the client connection
56/// and the runtime state; emits events; never returns errors that escape
57/// the connection (those are logged at the listener level).
58pub(crate) enum ClientStream {
59    Tcp(TcpStream),
60    #[cfg(unix)]
61    Unix(UnixStream),
62}
63
64impl ClientStream {
65    fn try_clone(&self) -> std::io::Result<Self> {
66        match self {
67            Self::Tcp(stream) => stream.try_clone().map(Self::Tcp),
68            #[cfg(unix)]
69            Self::Unix(stream) => stream.try_clone().map(Self::Unix),
70        }
71    }
72
73    fn set_read_timeout(&self, duration: Option<Duration>) -> std::io::Result<()> {
74        match self {
75            Self::Tcp(stream) => stream.set_read_timeout(duration),
76            #[cfg(unix)]
77            Self::Unix(stream) => stream.set_read_timeout(duration),
78        }
79    }
80}
81
82impl Read for ClientStream {
83    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
84        match self {
85            Self::Tcp(stream) => stream.read(buf),
86            #[cfg(unix)]
87            Self::Unix(stream) => stream.read(buf),
88        }
89    }
90}
91
92impl Write for ClientStream {
93    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
94        match self {
95            Self::Tcp(stream) => stream.write(buf),
96            #[cfg(unix)]
97            Self::Unix(stream) => stream.write(buf),
98        }
99    }
100
101    fn flush(&mut self) -> std::io::Result<()> {
102        match self {
103            Self::Tcp(stream) => stream.flush(),
104            #[cfg(unix)]
105            Self::Unix(stream) => stream.flush(),
106        }
107    }
108}
109
110trait StreamIo: Read + Write + Send + 'static {
111    fn shutdown(&self, how: Shutdown) -> std::io::Result<()>;
112}
113
114impl StreamIo for ClientStream {
115    fn shutdown(&self, how: Shutdown) -> std::io::Result<()> {
116        match self {
117            Self::Tcp(stream) => stream.shutdown(how),
118            #[cfg(unix)]
119            Self::Unix(stream) => stream.shutdown(how),
120        }
121    }
122}
123
124impl StreamIo for TcpStream {
125    fn shutdown(&self, how: Shutdown) -> std::io::Result<()> {
126        TcpStream::shutdown(self, how)
127    }
128}
129
130pub(crate) fn handle(client: ClientStream, state: Arc<RuntimeState>) -> Result<()> {
131    if let ClientStream::Tcp(stream) = &client
132        && let Err(error) = stream.set_nodelay(true)
133    {
134        log::debug!("[http_proxy] failed to set TCP_NODELAY on client socket: {error}");
135    }
136
137    let mut client = client;
138    if let Err(error) = client.set_read_timeout(Some(HEADER_READ_TIMEOUT)) {
139        log::debug!("[http_proxy] failed to set client header read timeout: {error}");
140    }
141
142    // Read until end-of-headers (`\r\n\r\n`). The first request determines
143    // the destination for this entire TCP connection.
144    let (header_buf, header_end) = read_request_headers(&mut client)?;
145
146    // The header timeout must not apply to the pump phase. (Socket options
147    // are shared with the `try_clone` handles used by the pump.)
148    if let Err(error) = client.set_read_timeout(None) {
149        log::debug!("[http_proxy] failed to clear client read timeout: {error}");
150    }
151
152    let request = ParsedRequest::parse(&header_buf[..header_end])?;
153
154    // Bytes after the request headers — the start of an HTTP request body
155    // for forwarding cases, or (for unusually eager clients) bytes pipelined
156    // ahead of the CONNECT response. Replayed to the upstream either way.
157    let leftover_body = header_buf[header_end..].to_vec();
158
159    match request {
160        ParsedRequest::Connect { host, port } => {
161            handle_connect(client, host, port, leftover_body, state)
162        }
163        ParsedRequest::Http {
164            method,
165            host,
166            port,
167            request_bytes,
168        } => handle_http_forward(
169            client,
170            method,
171            host,
172            port,
173            request_bytes,
174            leftover_body,
175            state,
176        ),
177    }
178}
179
180/// Reads request headers (until `\r\n\r\n`) from the client. Returns the
181/// full buffer (which may include some bytes after the headers) and the
182/// offset where the headers ended. Capped at [`MAX_HEADER_BYTES`].
183fn read_request_headers(client: &mut ClientStream) -> Result<(Vec<u8>, usize)> {
184    let mut buf = Vec::with_capacity(4096);
185    let mut tmp = [0u8; 4096];
186    let mut searched = 0usize;
187    loop {
188        let n = client.read(&mut tmp)?;
189        if n == 0 {
190            bail!("client closed before sending complete request headers");
191        }
192        buf.extend_from_slice(&tmp[..n]);
193        // Only scan the new bytes (plus 3 bytes of overlap for a delimiter
194        // straddling the read boundary), keeping the search linear overall.
195        let scan_start = searched.saturating_sub(3);
196        if let Some(end) = find_double_crlf(&buf[scan_start..]) {
197            return Ok((buf, scan_start + end));
198        }
199        searched = buf.len();
200        if buf.len() > MAX_HEADER_BYTES {
201            bail!("request headers exceed {MAX_HEADER_BYTES} bytes");
202        }
203    }
204}
205
206fn find_double_crlf(buf: &[u8]) -> Option<usize> {
207    buf.windows(4).position(|w| w == b"\r\n\r\n").map(|i| i + 4)
208}
209
210/// Parsed first request from the client.
211enum ParsedRequest {
212    Connect {
213        host: String,
214        port: u16,
215    },
216    Http {
217        method: String,
218        host: String,
219        port: u16,
220        /// Request bytes (line + headers) to forward to the origin.
221        /// Absolute-form requests are rewritten to origin-form (see
222        /// [`build_origin_form_request`]); origin-form requests pass
223        /// through verbatim.
224        request_bytes: Vec<u8>,
225    },
226}
227
228impl ParsedRequest {
229    fn parse(headers: &[u8]) -> Result<Self> {
230        let mut header_storage = [httparse::EMPTY_HEADER; 64];
231        let mut req = httparse::Request::new(&mut header_storage);
232        let status = req.parse(headers).context("malformed HTTP request")?;
233        if !status.is_complete() {
234            bail!("incomplete HTTP request after \\r\\n\\r\\n boundary");
235        }
236
237        let method = req
238            .method
239            .ok_or_else(|| anyhow!("missing HTTP method"))?
240            .to_string();
241        let target = req.path.ok_or_else(|| anyhow!("missing request target"))?;
242
243        if method.eq_ignore_ascii_case("CONNECT") {
244            let (host, port) = parse_authority_form(target)?;
245            Ok(ParsedRequest::Connect { host, port })
246        } else if target.starts_with("http://") || target.starts_with("https://") {
247            let (host, port, url) = parse_absolute_form_target(target)?;
248            let request_bytes = build_origin_form_request(&method, &url, req.version, req.headers);
249            Ok(ParsedRequest::Http {
250                method,
251                host,
252                port,
253                request_bytes,
254            })
255        } else {
256            // Origin-form request: destination comes from the Host: header,
257            // and the bytes are already in the form an origin server expects,
258            // so forward them verbatim.
259            let host_hdr = req
260                .headers
261                .iter()
262                .find(|h| h.name.eq_ignore_ascii_case("host"))
263                .ok_or_else(|| anyhow!("origin-form request missing Host: header"))?;
264            let value =
265                std::str::from_utf8(host_hdr.value).context("Host: header is not valid UTF-8")?;
266            let (host, port) = parse_host_header(value)?;
267            Ok(ParsedRequest::Http {
268                method,
269                host,
270                port,
271                request_bytes: headers.to_vec(),
272            })
273        }
274    }
275}
276
277/// Parse `host:port` (CONNECT authority-form). Port is required, per RFC 9112
278/// §3.2.3.
279///
280/// The port is split off manually rather than via `Url::parse`, because `Url`
281/// elides scheme-default ports — `host:80` parsed with an `http://` prefix
282/// becomes indistinguishable from a missing port, and `CONNECT host:80` is
283/// legitimate (e.g. `ws://` through a proxy). `Url` still canonicalizes the
284/// host part (bracketed IPv6, IDN-to-punycode, lowercasing).
285fn parse_authority_form(input: &str) -> Result<(String, u16)> {
286    let (host_part, port_part) = input
287        .rsplit_once(':')
288        .ok_or_else(|| anyhow!("CONNECT target '{input}' must include a port"))?;
289    let port: u16 = port_part
290        .parse()
291        .with_context(|| format!("CONNECT target '{input}' has an invalid port"))?;
292    let parsed = Url::parse(&format!("http://{host_part}"))
293        .with_context(|| format!("parsing CONNECT target '{input}'"))?;
294    let host = parsed
295        .host_str()
296        .ok_or_else(|| anyhow!("CONNECT target '{input}' has no host"))?
297        .to_string();
298    Ok((host, port))
299}
300
301/// Parse an absolute-form HTTP request target like `http://foo.com/path`.
302/// Returns the parsed URL too, for the origin-form rewrite.
303fn parse_absolute_form_target(target: &str) -> Result<(String, u16, Url)> {
304    let parsed =
305        Url::parse(target).with_context(|| format!("parsing absolute-form target '{target}'"))?;
306    let host = parsed
307        .host_str()
308        .ok_or_else(|| anyhow!("absolute-form target '{target}' has no host"))?
309        .to_string();
310    // `port_or_known_default` covers `http://foo.com` (→ 80) without forcing
311    // the agent to spell it.
312    let port = parsed
313        .port_or_known_default()
314        .ok_or_else(|| anyhow!("absolute-form target '{target}' has no port"))?;
315    Ok((host, port, parsed))
316}
317
318/// Rewrite an absolute-form proxy request into origin-form for the origin
319/// server.
320///
321/// RFC 9112 §3.2.2 requires origin servers to accept absolute-form, but many
322/// real servers and frameworks don't handle it; every production proxy
323/// rewrites, and so do we. The `Host` header is regenerated from the URL
324/// (per the same section, a proxy must use the URI host and ignore a
325/// mismatched `Host`), and `Proxy-*` headers — which are addressed to us and
326/// may carry credentials (`Proxy-Authorization`) — are stripped rather than
327/// leaked to the origin.
328fn build_origin_form_request(
329    method: &str,
330    url: &Url,
331    version: Option<u8>,
332    headers: &[httparse::Header],
333) -> Vec<u8> {
334    let mut target = url.path().to_string();
335    if let Some(query) = url.query() {
336        target.push('?');
337        target.push_str(query);
338    }
339    // `host_str` keeps brackets on IPv6 literals, which is what the Host
340    // header wants. `url.port()` is None for scheme-default ports, matching
341    // the convention of omitting them.
342    let host_value = match (url.host_str(), url.port()) {
343        (Some(host), Some(port)) => format!("{host}:{port}"),
344        (Some(host), None) => host.to_string(),
345        // Unreachable in practice: callers parsed the host already.
346        (None, _) => String::new(),
347    };
348    let minor_version = version.unwrap_or(1);
349
350    let mut out = Vec::with_capacity(256);
351    out.extend_from_slice(format!("{method} {target} HTTP/1.{minor_version}\r\n").as_bytes());
352    out.extend_from_slice(format!("Host: {host_value}\r\n").as_bytes());
353    for header in headers {
354        let name = header.name;
355        if name.eq_ignore_ascii_case("host")
356            || name
357                .get(.."proxy-".len())
358                .is_some_and(|prefix| prefix.eq_ignore_ascii_case("proxy-"))
359        {
360            continue;
361        }
362        out.extend_from_slice(name.as_bytes());
363        out.extend_from_slice(b": ");
364        out.extend_from_slice(header.value);
365        out.extend_from_slice(b"\r\n");
366    }
367    out.extend_from_slice(b"\r\n");
368    out
369}
370
371/// Parse a `Host:` header value into `(host, port)`. Default port is 80
372/// since this is only called for HTTP forward, never CONNECT.
373fn parse_host_header(value: &str) -> Result<(String, u16)> {
374    let value = value.trim();
375    if value.is_empty() {
376        bail!("empty Host header");
377    }
378    let parsed = Url::parse(&format!("http://{value}"))
379        .with_context(|| format!("parsing Host header '{value}'"))?;
380    let host = parsed
381        .host_str()
382        .ok_or_else(|| anyhow!("Host header '{value}' has no host"))?
383        .to_string();
384    let port = parsed.port().unwrap_or(80);
385    Ok((host, port))
386}
387
388/// Normalize a hostname for allowlist matching. Strips brackets from IPv6
389/// literals and a single trailing dot. Lowercasing happens inside the
390/// allowlist matcher.
391fn normalize_host(host: &str) -> String {
392    let stripped = host
393        .strip_prefix('[')
394        .and_then(|s| s.strip_suffix(']'))
395        .unwrap_or(host);
396    stripped.trim_end_matches('.').to_string()
397}
398
399/// Whether a hostname is an IP literal, per our policy.
400fn is_ip_literal(host: &str) -> bool {
401    let stripped = host
402        .strip_prefix('[')
403        .and_then(|s| s.strip_suffix(']'))
404        .unwrap_or(host);
405    stripped.parse::<IpAddr>().is_ok()
406}
407
408/// The policy check shared by CONNECT and HTTP forward: IP-literal targets
409/// and hosts outside the allowlist are denied. Both checks are skipped when
410/// the allowlist allows any host — that grant means unrestricted egress,
411/// including IP literals (matching the pre-allowlist `allow_network`
412/// behavior).
413fn policy_denial(host: &str, port: u16, state: &RuntimeState) -> Option<DenyReason> {
414    if state.allowlist.allows_any() {
415        return None;
416    }
417    if is_ip_literal(host) {
418        return Some(DenyReason::IpLiteralRejected {
419            target: format!("{host}:{port}"),
420        });
421    }
422    if !state.allowlist.allows(host) {
423        return Some(DenyReason::HostNotInAllowlist {
424            host: host.to_string(),
425        });
426    }
427    None
428}
429
430/// How an approved request will reach its destination.
431enum Route {
432    /// Connect directly to one of the addresses pinned when the host was
433    /// resolved and vetted.
434    Direct(PinnedHost),
435    /// Tunnel through the upstream proxy with a CONNECT handshake.
436    ViaUpstream(UpstreamProxy),
437}
438
439enum RouteFailure {
440    /// Policy denial — respond with 511.
441    Denied(DenyReason),
442    /// Network-level failure — close silently, per "no proxy here".
443    Error(anyhow::Error),
444}
445
446/// Decide how to reach `host:port`, resolving and vetting addresses for
447/// direct connections.
448///
449/// Resolution happens here, in the unsandboxed editor process, so this is
450/// also where we keep allowlisted hostnames from smuggling the sandboxed
451/// command onto the local machine or local network: a hostname whose DNS
452/// points into loopback / private / link-local space (DNS rebinding) is
453/// denied, and the connection later uses the vetted addresses rather than
454/// re-resolving. The filter is skipped when the allowlist allows any host,
455/// since that grant means unrestricted egress. Upstream-proxied destinations
456/// aren't resolved locally at all — the upstream does its own resolution
457/// inside the user's trusted network.
458fn plan_route(host: &str, port: u16, state: &RuntimeState) -> Result<Route, RouteFailure> {
459    if let Some(upstream) = &state.upstream
460        && !upstream.bypasses(host, port)
461    {
462        return Ok(Route::ViaUpstream(upstream.clone()));
463    }
464
465    match PinnedHost::resolve_for_allowlist(host, port, &state.allowlist) {
466        Ok(pinned) => Ok(Route::Direct(pinned)),
467        Err(PinnedHostError::AllAddressesForbidden { host }) => {
468            Err(RouteFailure::Denied(DenyReason::ResolvedToForbiddenIp {
469                host,
470            }))
471        }
472        Err(error) => Err(RouteFailure::Error(anyhow!(error))),
473    }
474}
475
476fn handle_connect(
477    mut client: ClientStream,
478    host: String,
479    port: u16,
480    leftover_body: Vec<u8>,
481    state: Arc<RuntimeState>,
482) -> Result<()> {
483    let normalized = normalize_host(&host);
484
485    if let Some(reason) = policy_denial(&normalized, port, &state) {
486        return deny_request(
487            &mut client,
488            &state,
489            normalized,
490            port,
491            RequestMethod::Connect,
492            reason,
493        );
494    }
495
496    let route = match plan_route(&normalized, port, &state) {
497        Ok(route) => route,
498        Err(RouteFailure::Denied(reason)) => {
499            return deny_request(
500                &mut client,
501                &state,
502                normalized,
503                port,
504                RequestMethod::Connect,
505                reason,
506            );
507        }
508        Err(RouteFailure::Error(error)) => {
509            log::debug!("[http_proxy] routing failed for CONNECT {normalized}:{port}: {error:#}");
510            // Per "no proxy here" — close abruptly. Client sees a connection drop.
511            return Ok(());
512        }
513    };
514
515    emit(
516        &state,
517        ProxyEvent::RequestAttempt {
518            host: normalized.clone(),
519            port,
520            method: RequestMethod::Connect,
521            outcome: RequestOutcome::Allowed,
522        },
523    );
524
525    let (mut upstream, upstream_leftover) = match open_route(&route, &normalized, port) {
526        Ok(opened) => opened,
527        Err(error) => {
528            log::debug!(
529                "[http_proxy] upstream open failed for CONNECT {normalized}:{port}: {error:#}"
530            );
531            return Ok(());
532        }
533    };
534
535    if let Err(error) = upstream.set_nodelay(true) {
536        log::debug!("[http_proxy] failed to set TCP_NODELAY on upstream socket: {error}");
537    }
538
539    // Tell the client the tunnel is up, then replay anything the upstream
540    // sent past its CONNECT response and anything the client pipelined ahead
541    // of ours.
542    client.write_all(b"HTTP/1.1 200 Connection established\r\n\r\n")?;
543    if !upstream_leftover.is_empty() {
544        client.write_all(&upstream_leftover)?;
545    }
546    if !leftover_body.is_empty() {
547        upstream.write_all(&leftover_body)?;
548    }
549
550    let started = Instant::now();
551    let (pumped_to_remote, pumped_from_remote) = pump_bidir(client, upstream);
552
553    emit(
554        &state,
555        ProxyEvent::RequestCompleted {
556            host: normalized,
557            port,
558            method: RequestMethod::Connect,
559            bytes_to_remote: pumped_to_remote + leftover_body.len() as u64,
560            bytes_from_remote: pumped_from_remote + upstream_leftover.len() as u64,
561            duration_ms: started.elapsed().as_millis() as u64,
562        },
563    );
564
565    Ok(())
566}
567
568fn handle_http_forward(
569    mut client: ClientStream,
570    method: String,
571    host: String,
572    port: u16,
573    request_bytes: Vec<u8>,
574    leftover_body: Vec<u8>,
575    state: Arc<RuntimeState>,
576) -> Result<()> {
577    let normalized = normalize_host(&host);
578
579    if let Some(reason) = policy_denial(&normalized, port, &state) {
580        return deny_request(
581            &mut client,
582            &state,
583            normalized,
584            port,
585            RequestMethod::Http(method),
586            reason,
587        );
588    }
589
590    let route = match plan_route(&normalized, port, &state) {
591        Ok(route) => route,
592        Err(RouteFailure::Denied(reason)) => {
593            return deny_request(
594                &mut client,
595                &state,
596                normalized,
597                port,
598                RequestMethod::Http(method),
599                reason,
600            );
601        }
602        Err(RouteFailure::Error(error)) => {
603            log::debug!("[http_proxy] routing failed for {method} {normalized}:{port}: {error:#}");
604            return Ok(());
605        }
606    };
607
608    emit(
609        &state,
610        ProxyEvent::RequestAttempt {
611            host: normalized.clone(),
612            port,
613            method: RequestMethod::Http(method.clone()),
614            outcome: RequestOutcome::Allowed,
615        },
616    );
617
618    let (mut upstream, upstream_leftover) = match open_route(&route, &normalized, port) {
619        Ok(opened) => opened,
620        Err(error) => {
621            log::debug!(
622                "[http_proxy] upstream open failed for {method} {normalized}:{port}: {error:#}"
623            );
624            return Ok(());
625        }
626    };
627
628    if let Err(error) = upstream.set_nodelay(true) {
629        log::debug!("[http_proxy] failed to set TCP_NODELAY on upstream socket: {error}");
630    }
631
632    if !upstream_leftover.is_empty() {
633        client.write_all(&upstream_leftover)?;
634    }
635    upstream.write_all(&request_bytes)?;
636    if !leftover_body.is_empty() {
637        upstream.write_all(&leftover_body)?;
638    }
639
640    let started = Instant::now();
641    let (pumped_to_remote, pumped_from_remote) = pump_bidir(client, upstream);
642    let to_remote = pumped_to_remote + request_bytes.len() as u64 + leftover_body.len() as u64;
643
644    emit(
645        &state,
646        ProxyEvent::RequestCompleted {
647            host: normalized,
648            port,
649            method: RequestMethod::Http(method),
650            bytes_to_remote: to_remote,
651            bytes_from_remote: pumped_from_remote + upstream_leftover.len() as u64,
652            duration_ms: started.elapsed().as_millis() as u64,
653        },
654    );
655
656    Ok(())
657}
658
659/// Open the connection that will carry this request's bytes to the origin —
660/// a direct TCP connection to a vetted address, or a CONNECT tunnel through
661/// the upstream proxy. Returns the stream plus any bytes the upstream sent
662/// past its CONNECT response (rare; replayed to the client by the caller).
663///
664/// HTTP forward also goes through a CONNECT tunnel when chaining: handing
665/// the upstream a routable absolute-form byte stream would let later
666/// keep-alive requests on this connection name a different (unapproved)
667/// host and have the upstream route them there. A tunnel pins the whole
668/// connection to the approved `host:port` — and shares the upstream auth
669/// handshake with the CONNECT path.
670fn open_route(route: &Route, host: &str, port: u16) -> Result<(TcpStream, Vec<u8>)> {
671    match route {
672        Route::Direct(pinned) => {
673            // Connect to the addresses pinned when the host was vetted, never
674            // re-resolving `host` here — that re-resolution is exactly the
675            // DNS-rebinding window `PinnedHost` exists to close.
676            let addrs: Vec<SocketAddr> = pinned.socket_addrs().collect();
677            Ok((connect_to_any(&addrs, host, port)?, Vec::new()))
678        }
679        Route::ViaUpstream(upstream) => connect_via_upstream(host, port, upstream),
680    }
681}
682
683/// Connect to the first address that accepts, with a per-attempt timeout.
684fn connect_to_any(addrs: &[SocketAddr], host: &str, port: u16) -> Result<TcpStream> {
685    let mut last_error = None;
686    for addr in addrs {
687        match TcpStream::connect_timeout(addr, CONNECT_TIMEOUT) {
688            Ok(stream) => return Ok(stream),
689            Err(error) => last_error = Some(error),
690        }
691    }
692    match last_error {
693        Some(error) => Err(anyhow!("connect to {host}:{port}: {error}")),
694        None => Err(anyhow!("no addresses to connect to for {host}:{port}")),
695    }
696}
697
698/// Open a TCP connection to the upstream proxy and complete a CONNECT
699/// handshake to (host, port). Returns once the upstream has confirmed `200`,
700/// along with any bytes it sent past its response headers.
701fn connect_via_upstream(
702    host: &str,
703    port: u16,
704    upstream: &UpstreamProxy,
705) -> Result<(TcpStream, Vec<u8>)> {
706    let addrs: Vec<SocketAddr> = (upstream.host.as_str(), upstream.port)
707        .to_socket_addrs()
708        .with_context(|| format!("resolving upstream proxy {upstream}"))?
709        .collect();
710    let mut stream = connect_to_any(&addrs, &upstream.host, upstream.port)
711        .with_context(|| format!("connect to upstream proxy {upstream}"))?;
712
713    if let Err(error) = stream.set_read_timeout(Some(UPSTREAM_HANDSHAKE_TIMEOUT)) {
714        log::debug!("[http_proxy] failed to set upstream handshake timeout: {error}");
715    }
716
717    let auth_header = upstream.auth.as_ref().map(|auth| {
718        let creds = format!("{}:{}", auth.user, auth.password);
719        format!(
720            "Proxy-Authorization: Basic {}\r\n",
721            base64::engine::general_purpose::STANDARD.encode(creds.as_bytes())
722        )
723    });
724
725    let request = format!(
726        "CONNECT {host}:{port} HTTP/1.1\r\n\
727         Host: {host}:{port}\r\n\
728         {}\
729         \r\n",
730        auth_header.unwrap_or_default()
731    );
732    stream.write_all(request.as_bytes())?;
733
734    // Read upstream's response status line + headers (up to \r\n\r\n).
735    let mut buf = Vec::with_capacity(512);
736    let mut tmp = [0u8; 512];
737    let header_end = loop {
738        let n = stream.read(&mut tmp)?;
739        if n == 0 {
740            bail!("upstream closed before responding to CONNECT");
741        }
742        buf.extend_from_slice(&tmp[..n]);
743        if let Some(end) = find_double_crlf(&buf) {
744            break end;
745        }
746        if buf.len() > MAX_HEADER_BYTES {
747            bail!("upstream CONNECT response headers exceed {MAX_HEADER_BYTES} bytes");
748        }
749    };
750
751    let mut header_storage = [httparse::EMPTY_HEADER; 16];
752    let mut response = httparse::Response::new(&mut header_storage);
753    response
754        .parse(&buf[..header_end])
755        .context("malformed upstream CONNECT response")?;
756    let status = response
757        .code
758        .ok_or_else(|| anyhow!("upstream CONNECT response missing status code"))?;
759    if status != 200 {
760        let reason = response.reason.unwrap_or("");
761        bail!("upstream CONNECT refused: HTTP {status} {reason}");
762    }
763
764    if let Err(error) = stream.set_read_timeout(None) {
765        log::debug!("[http_proxy] failed to clear upstream handshake timeout: {error}");
766    }
767
768    // Bytes past the response headers already belong to the tunnel (an
769    // origin could in principle speak first); hand them back for replay.
770    Ok((stream, buf[header_end..].to_vec()))
771}
772
773/// Send a 511 response with `Via` and `Proxy-Status` headers and an
774/// explanatory body. Closes the connection afterwards.
775fn deny_request(
776    client: &mut ClientStream,
777    state: &RuntimeState,
778    host: String,
779    port: u16,
780    method: RequestMethod,
781    reason: DenyReason,
782) -> Result<()> {
783    emit(
784        state,
785        ProxyEvent::RequestAttempt {
786            host,
787            port,
788            method,
789            outcome: RequestOutcome::Denied {
790                reason: reason.clone(),
791            },
792        },
793    );
794
795    let body = format!(
796        "Request blocked by the Omega sandbox network policy.\n\n  \
797         Reason: {}\n\n  \
798         This is not a network or server failure — it's a policy decision.\n  \
799         To proceed, ask the user to approve the host on the next terminal call.\n",
800        reason.human_explanation()
801    );
802    let response = format!(
803        "HTTP/1.1 511 Network Authentication Required\r\n\
804         Via: 1.1 omega-sandbox-proxy\r\n\
805         Proxy-Status: omega-sandbox-proxy; error={}; details=\"{}\"\r\n\
806         Content-Type: text/plain; charset=utf-8\r\n\
807         Content-Length: {}\r\n\
808         Connection: close\r\n\r\n{body}",
809        reason.proxy_status_error(),
810        proxy_status_details(&reason),
811        body.len(),
812    );
813    client.write_all(response.as_bytes())?;
814    Ok(())
815}
816
817fn proxy_status_details(reason: &DenyReason) -> String {
818    reason
819        .human_explanation()
820        .replace(['\r', '\n'], " ")
821        .replace('"', "'")
822}
823
824fn emit(state: &RuntimeState, event: ProxyEvent) {
825    // Unbounded send is sync and never blocks. Only fails if the receiver
826    // has been dropped, which we silently ignore — events are diagnostic,
827    // not load-bearing.
828    let _ = state.events.unbounded_send(event);
829}
830
831/// Bidirectional byte pump.
832///
833/// The client→remote direction runs on a spawned thread; remote→client runs
834/// on the current (connection) thread, halving the thread count per
835/// connection. Each direction reads from one socket and writes to the other
836/// in a tight loop with a fixed-size buffer. When one direction reaches EOF
837/// (peer closed write-half), the receiving side shuts down the other side's
838/// write-half so the partner eventually sees EOF too. This mirrors what
839/// `tokio::io::copy_bidirectional` does, with explicit thread join.
840///
841/// Returns `(client→remote bytes, remote→client bytes)`. Errors are
842/// swallowed — partial transfer is fine, the caller just emits whatever
843/// totals we got.
844fn pump_bidir(client: ClientStream, upstream: TcpStream) -> (u64, u64) {
845    // Two clones per direction so each side owns the half it touches.
846    // `try_clone` dups the underlying fd, so reads/writes on the two
847    // halves don't contend for the same kernel state.
848    let client_read = match client.try_clone() {
849        Ok(s) => s,
850        Err(e) => {
851            log::debug!("[http_proxy] failed to clone client socket for bidir pump: {e}");
852            return (0, 0);
853        }
854    };
855    let upstream_read = match upstream.try_clone() {
856        Ok(s) => s,
857        Err(e) => {
858            log::debug!("[http_proxy] failed to clone upstream socket for bidir pump: {e}");
859            return (0, 0);
860        }
861    };
862    let client_write = client;
863    let upstream_write = upstream;
864
865    let to_remote_handle = match thread::Builder::new()
866        .name("http-proxy-pump-out".to_string())
867        .stack_size(128 * 1024)
868        .spawn(move || copy_one_way(client_read, upstream_write))
869    {
870        Ok(handle) => handle,
871        Err(error) => {
872            // Returning drops all stream handles, closing both sockets.
873            log::warn!("[http_proxy] failed to spawn pump thread: {error}");
874            return (0, 0);
875        }
876    };
877    let from_remote = copy_one_way(upstream_read, client_write);
878    let to_remote = to_remote_handle.join().unwrap_or_else(|_| {
879        log::warn!("[http_proxy] pump thread panicked");
880        0
881    });
882    (to_remote, from_remote)
883}
884
885/// Copy bytes from `from` to `to` until EOF on the read side or write
886/// failure. Half-closes `to` for writes when done so the partner sees EOF.
887fn copy_one_way(mut from: impl Read, mut to: impl StreamIo) -> u64 {
888    let mut total = 0u64;
889    let mut buf = vec![0u8; PUMP_BUFFER_SIZE];
890    loop {
891        let n = match from.read(&mut buf) {
892            Ok(0) => break,
893            Ok(n) => n,
894            Err(_) => break,
895        };
896        if to.write_all(&buf[..n]).is_err() {
897            break;
898        }
899        total += n as u64;
900    }
901    // Half-close the write side. Mirrors what tokio's copy_bidirectional
902    // does on EOF: the partner's read returns EOF eventually, completing
903    // the other direction.
904    let _ = to.shutdown(Shutdown::Write);
905    total
906}
907
908#[cfg(test)]
909mod tests {
910    use super::*;
911    use crate::allowlist::{Allowlist, HostPattern};
912    use std::sync::atomic::AtomicUsize;
913
914    #[test]
915    fn plan_route_denies_host_resolving_to_loopback() {
916        // `localhost` resolves to loopback without real DNS, standing in for
917        // any allowlisted hostname whose DNS points at the local machine
918        // (DNS rebinding).
919        let state = runtime_state(Allowlist::from_patterns([
920            HostPattern::parse("github.com").unwrap()
921        ]));
922        match plan_route("localhost", 80, &state) {
923            Err(RouteFailure::Denied(DenyReason::ResolvedToForbiddenIp { host })) => {
924                assert_eq!(host, "localhost");
925            }
926            Ok(_) => panic!("expected denial, got a route"),
927            Err(RouteFailure::Denied(reason)) => panic!("unexpected deny reason: {reason:?}"),
928            Err(RouteFailure::Error(error)) => panic!("expected denial, got error: {error}"),
929        }
930    }
931
932    #[test]
933    fn plan_route_allows_loopback_when_allowlist_allows_any() {
934        let state = runtime_state(Allowlist::any());
935        match plan_route("localhost", 80, &state) {
936            Ok(Route::Direct(pinned)) => {
937                let addrs: Vec<SocketAddr> = pinned.socket_addrs().collect();
938                assert!(!addrs.is_empty());
939                assert!(addrs.iter().all(|addr| addr.ip().is_loopback()));
940            }
941            Ok(Route::ViaUpstream(_)) => panic!("expected direct route"),
942            Err(RouteFailure::Denied(reason)) => panic!("unexpected denial: {reason:?}"),
943            Err(RouteFailure::Error(error)) => panic!("unexpected error: {error}"),
944        }
945    }
946
947    fn runtime_state(allowlist: Allowlist) -> RuntimeState {
948        let (events, _receiver) = futures::channel::mpsc::unbounded();
949        RuntimeState {
950            allowlist,
951            upstream: None,
952            events,
953            active_connections: AtomicUsize::new(0),
954        }
955    }
956
957    #[test]
958    fn parse_authority_form_basic() {
959        let (h, p) = parse_authority_form("github.com:443").unwrap();
960        assert_eq!(h, "github.com");
961        assert_eq!(p, 443);
962    }
963
964    #[test]
965    fn parse_authority_form_ipv6() {
966        let (h, p) = parse_authority_form("[::1]:443").unwrap();
967        assert_eq!(h, "[::1]");
968        assert_eq!(p, 443);
969    }
970
971    #[test]
972    fn parse_authority_form_accepts_scheme_default_ports() {
973        // Regression test: `Url::parse` elides scheme-default ports, so a
974        // naive URL round-trip would reject `host:80` as "missing a port".
975        let (h, p) = parse_authority_form("example.com:80").unwrap();
976        assert_eq!(h, "example.com");
977        assert_eq!(p, 80);
978    }
979
980    #[test]
981    fn parse_authority_form_requires_port() {
982        assert!(parse_authority_form("github.com").is_err());
983        assert!(parse_authority_form("[::1]").is_err());
984    }
985
986    #[test]
987    fn parse_absolute_form_basic() {
988        let (h, p, _) = parse_absolute_form_target("http://example.com/path").unwrap();
989        assert_eq!(h, "example.com");
990        assert_eq!(p, 80);
991    }
992
993    #[test]
994    fn parse_absolute_form_with_port() {
995        let (h, p, _) = parse_absolute_form_target("http://example.com:8080/").unwrap();
996        assert_eq!(h, "example.com");
997        assert_eq!(p, 8080);
998    }
999
1000    #[test]
1001    fn host_header_default_port() {
1002        let (h, p) = parse_host_header("example.com").unwrap();
1003        assert_eq!(h, "example.com");
1004        assert_eq!(p, 80);
1005    }
1006
1007    #[test]
1008    fn host_header_explicit_port() {
1009        let (h, p) = parse_host_header("example.com:8080").unwrap();
1010        assert_eq!(h, "example.com");
1011        assert_eq!(p, 8080);
1012    }
1013
1014    #[test]
1015    fn host_header_ipv6() {
1016        let (h, p) = parse_host_header("[::1]:443").unwrap();
1017        assert_eq!(h, "[::1]");
1018        assert_eq!(p, 443);
1019    }
1020
1021    #[test]
1022    fn detects_ip_literals() {
1023        assert!(is_ip_literal("1.2.3.4"));
1024        assert!(is_ip_literal("[::1]"));
1025        assert!(is_ip_literal("::1"));
1026        assert!(!is_ip_literal("github.com"));
1027        assert!(!is_ip_literal("localhost"));
1028    }
1029
1030    // Forbidden-IP range coverage lives with the logic in `pinned_host.rs`.
1031
1032    #[test]
1033    fn parsed_request_recognizes_connect() {
1034        let req = b"CONNECT example.com:443 HTTP/1.1\r\nHost: example.com:443\r\n\r\n";
1035        match ParsedRequest::parse(req).unwrap() {
1036            ParsedRequest::Connect { host, port } => {
1037                assert_eq!(host, "example.com");
1038                assert_eq!(port, 443);
1039            }
1040            ParsedRequest::Http { .. } => panic!("expected Connect"),
1041        }
1042    }
1043
1044    #[test]
1045    fn parsed_request_recognizes_http_absolute_form() {
1046        let req = b"GET http://example.com/foo HTTP/1.1\r\nHost: example.com\r\n\r\n";
1047        match ParsedRequest::parse(req).unwrap() {
1048            ParsedRequest::Http {
1049                method, host, port, ..
1050            } => {
1051                assert_eq!(method, "GET");
1052                assert_eq!(host, "example.com");
1053                assert_eq!(port, 80);
1054            }
1055            ParsedRequest::Connect { .. } => panic!("expected Http"),
1056        }
1057    }
1058
1059    #[test]
1060    fn absolute_form_is_rewritten_to_origin_form() {
1061        let req = b"GET http://example.com/foo?q=1 HTTP/1.1\r\n\
1062            Host: wrong.example\r\n\
1063            Proxy-Connection: keep-alive\r\n\
1064            Proxy-Authorization: Basic c2VjcmV0\r\n\
1065            User-Agent: test\r\n\r\n";
1066        match ParsedRequest::parse(req).unwrap() {
1067            ParsedRequest::Http { request_bytes, .. } => {
1068                let text = String::from_utf8(request_bytes).unwrap();
1069                assert!(text.starts_with("GET /foo?q=1 HTTP/1.1\r\n"), "{text}");
1070                // Host is regenerated from the URI, not trusted from the
1071                // (mismatched) header.
1072                assert!(text.contains("Host: example.com\r\n"), "{text}");
1073                assert!(!text.contains("wrong.example"), "{text}");
1074                // Proxy-* headers are addressed to us (and may carry
1075                // credentials); they must not leak to the origin.
1076                assert!(!text.to_ascii_lowercase().contains("proxy-"), "{text}");
1077                assert!(text.contains("User-Agent: test\r\n"), "{text}");
1078                assert!(text.ends_with("\r\n\r\n"), "{text}");
1079            }
1080            ParsedRequest::Connect { .. } => panic!("expected Http"),
1081        }
1082    }
1083
1084    #[test]
1085    fn absolute_form_rewrite_keeps_non_default_port_in_host_header() {
1086        let req = b"GET http://example.com:8080/ HTTP/1.1\r\nHost: example.com:8080\r\n\r\n";
1087        match ParsedRequest::parse(req).unwrap() {
1088            ParsedRequest::Http { request_bytes, .. } => {
1089                let text = String::from_utf8(request_bytes).unwrap();
1090                assert!(text.starts_with("GET / HTTP/1.1\r\n"), "{text}");
1091                assert!(text.contains("Host: example.com:8080\r\n"), "{text}");
1092            }
1093            ParsedRequest::Connect { .. } => panic!("expected Http"),
1094        }
1095    }
1096
1097    #[test]
1098    fn parsed_request_recognizes_http_origin_form_via_host_header() {
1099        let req = b"GET /foo HTTP/1.1\r\nHost: example.com:8080\r\n\r\n";
1100        match ParsedRequest::parse(req).unwrap() {
1101            ParsedRequest::Http {
1102                host,
1103                port,
1104                request_bytes,
1105                ..
1106            } => {
1107                assert_eq!(host, "example.com");
1108                assert_eq!(port, 8080);
1109                // Origin-form bytes pass through verbatim.
1110                assert_eq!(request_bytes, req.to_vec());
1111            }
1112            ParsedRequest::Connect { .. } => panic!("expected Http"),
1113        }
1114    }
1115}
1116