Skip to main content

sbe_proxy/
server.rs

1use std::{
2    collections::BTreeSet,
3    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
4    sync::Arc,
5    time::Duration,
6};
7
8use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64_STANDARD};
9use tokio::{
10    io::{
11        AsyncBufRead, AsyncBufReadExt, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt,
12        BufReader,
13    },
14    net::{TcpListener, TcpStream, lookup_host},
15    sync::{Semaphore, watch},
16    task::JoinSet,
17    time::{Instant, timeout, timeout_at},
18};
19use tracing::{debug, info, warn};
20
21use crate::{allowlist::DomainAllowlist, error::ProxyError};
22
23/// Resource and destination limits for the trusted local proxy.
24#[derive(Debug, Clone)]
25pub struct ProxyConfig {
26    pub max_request_line_bytes: usize,
27    pub max_header_line_bytes: usize,
28    pub max_header_bytes: usize,
29    pub max_headers: usize,
30    pub max_connections: usize,
31    pub max_resolved_addresses: usize,
32    pub header_timeout: Duration,
33    pub dns_timeout: Duration,
34    pub connect_timeout: Duration,
35    pub idle_timeout: Duration,
36    pub tunnel_timeout: Duration,
37    pub allowed_ports: BTreeSet<u16>,
38    /// Tests and explicitly trusted local-development profiles may opt in.
39    pub allow_special_addresses: bool,
40}
41
42impl Default for ProxyConfig {
43    fn default() -> Self {
44        Self {
45            max_request_line_bytes: 8 * 1024,
46            max_header_line_bytes: 8 * 1024,
47            max_header_bytes: 32 * 1024,
48            max_headers: 64,
49            max_connections: 64,
50            max_resolved_addresses: 32,
51            header_timeout: Duration::from_secs(10),
52            dns_timeout: Duration::from_secs(10),
53            connect_timeout: Duration::from_secs(15),
54            idle_timeout: Duration::from_secs(60),
55            tunnel_timeout: Duration::from_secs(30 * 60),
56            allowed_ports: BTreeSet::from([443]),
57            allow_special_addresses: false,
58        }
59    }
60}
61
62/// Per-run endpoint information injected into the sandbox environment.
63#[derive(Clone)]
64pub struct ProxyEndpoint {
65    pub port: u16,
66    token: String,
67}
68
69impl std::fmt::Debug for ProxyEndpoint {
70    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
71        formatter
72            .debug_struct("ProxyEndpoint")
73            .field("port", &self.port)
74            .field("token", &"<redacted>")
75            .finish()
76    }
77}
78
79impl ProxyEndpoint {
80    pub fn url(&self) -> String {
81        format!("http://sbe:{}@127.0.0.1:{}", self.token, self.port)
82    }
83
84    pub fn authorization_header(&self) -> String {
85        format!(
86            "Basic {}",
87            BASE64_STANDARD.encode(format!("sbe:{}", self.token))
88        )
89    }
90
91    fn java_tool_options(&self, agent_path: &str, temp_path: &str) -> String {
92        format!(
93            "-javaagent:{agent_path} \
94             -Djava.io.tmpdir={temp_path} \
95             -Dhttp.proxyHost=127.0.0.1 -Dhttp.proxyPort={} \
96             -Dhttp.proxyProtocol=http \
97             -Dhttps.proxyHost=127.0.0.1 -Dhttps.proxyPort={} \
98             -Dhttps.proxyProtocol=http \
99             -Dhttp.nonProxyHosts= \
100             -Djdk.http.auth.tunneling.disabledSchemes=",
101            self.port, self.port
102        )
103    }
104
105    /// Runtime variables for JVM clients. The one-time token stays outside
106    /// `JAVA_TOOL_OPTIONS`, which the JVM prints to stderr on every launch.
107    pub fn java_environment(
108        &self,
109        agent_path: &str,
110        temp_path: &str,
111    ) -> [(&'static str, String); 2] {
112        [
113            (
114                "JAVA_TOOL_OPTIONS",
115                self.java_tool_options(agent_path, temp_path),
116            ),
117            ("SBE_PROXY_TOKEN", self.token.clone()),
118        ]
119    }
120}
121
122/// A bounded, authenticated, domain-filtering HTTP CONNECT proxy.
123pub struct ProxyServer {
124    listener: TcpListener,
125    allowlist: Arc<DomainAllowlist>,
126    config: Arc<ProxyConfig>,
127    authorization: Arc<str>,
128    shutdown_rx: watch::Receiver<bool>,
129}
130
131impl ProxyServer {
132    pub async fn bind(
133        allowlist: DomainAllowlist,
134        shutdown_rx: watch::Receiver<bool>,
135    ) -> Result<(Self, ProxyEndpoint), ProxyError> {
136        Self::bind_with_config(allowlist, shutdown_rx, ProxyConfig::default()).await
137    }
138
139    pub async fn bind_with_config(
140        allowlist: DomainAllowlist,
141        shutdown_rx: watch::Receiver<bool>,
142        config: ProxyConfig,
143    ) -> Result<(Self, ProxyEndpoint), ProxyError> {
144        if config.max_connections == 0
145            || config.max_request_line_bytes == 0
146            || config.max_header_line_bytes == 0
147            || config.max_header_bytes == 0
148            || config.max_headers == 0
149            || config.max_resolved_addresses == 0
150            || config.allowed_ports.is_empty()
151            || config.header_timeout.is_zero()
152            || config.dns_timeout.is_zero()
153            || config.connect_timeout.is_zero()
154            || config.idle_timeout.is_zero()
155            || config.tunnel_timeout.is_zero()
156        {
157            return Err(ProxyError::Config(
158                "proxy limits and allowed ports must be non-zero".to_owned(),
159            ));
160        }
161        let listener = TcpListener::bind("127.0.0.1:0")
162            .await
163            .map_err(ProxyError::Bind)?;
164        let port = listener.local_addr().map_err(ProxyError::Bind)?.port();
165        let mut random = [0_u8; 32];
166        getrandom::fill(&mut random)
167            .map_err(|e| ProxyError::Config(format!("proxy authentication RNG failed: {e}")))?;
168        let endpoint = ProxyEndpoint {
169            port,
170            token: hex(&random),
171        };
172        let authorization: Arc<str> = endpoint.authorization_header().into();
173        info!(port, "sbe proxy listening");
174        Ok((
175            Self {
176                listener,
177                allowlist: Arc::new(allowlist),
178                config: Arc::new(config),
179                authorization,
180                shutdown_rx,
181            },
182            endpoint,
183        ))
184    }
185
186    /// Run until shutdown and keep every connection task inside this lifecycle.
187    pub async fn run(self) -> Result<(), ProxyError> {
188        let mut shutdown = self.shutdown_rx;
189        let semaphore = Arc::new(Semaphore::new(self.config.max_connections));
190        let mut tasks = JoinSet::new();
191        loop {
192            tokio::select! {
193                result = self.listener.accept() => {
194                    let (stream, addr) = result.map_err(ProxyError::Accept)?;
195                    let Ok(permit) = Arc::clone(&semaphore).try_acquire_owned() else {
196                        warn!(client = %addr, "proxy connection limit reached");
197                        drop(stream);
198                        continue;
199                    };
200                    let allowlist = Arc::clone(&self.allowlist);
201                    let config = Arc::clone(&self.config);
202                    let authorization = Arc::clone(&self.authorization);
203                    tasks.spawn(async move {
204                        let _permit = permit;
205                        if let Err(error) = handle_connection(
206                            stream, addr, &allowlist, &config, &authorization,
207                        ).await {
208                            debug!(%error, client = %addr, "proxy connection error");
209                        }
210                    });
211                }
212                Some(result) = tasks.join_next(), if !tasks.is_empty() => {
213                    if let Err(error) = result {
214                        warn!(%error, "proxy connection task failed");
215                    }
216                }
217                changed = shutdown.changed() => {
218                    if changed.is_err() || *shutdown.borrow() {
219                        break;
220                    }
221                }
222            }
223        }
224        tasks.shutdown().await;
225        info!("sbe proxy shut down");
226        Ok(())
227    }
228}
229
230async fn handle_connection(
231    client: TcpStream,
232    addr: SocketAddr,
233    allowlist: &DomainAllowlist,
234    config: &ProxyConfig,
235    expected_authorization: &str,
236) -> Result<(), ProxyError> {
237    let mut client = BufReader::new(client);
238    let header_deadline = Instant::now() + config.header_timeout;
239    let request_line = timeout_at(
240        header_deadline,
241        read_bounded_line(&mut client, config.max_request_line_bytes),
242    )
243    .await
244    .map_err(|_| ProxyError::Timeout("request line"))??;
245    if request_line.is_empty() {
246        return Ok(());
247    }
248    let request = std::str::from_utf8(&request_line)
249        .map_err(|_| ProxyError::Protocol("request line is not UTF-8".to_owned()))?
250        .trim_end_matches(['\r', '\n']);
251    let mut parts = request.split_ascii_whitespace();
252    let method = parts.next().unwrap_or_default();
253    let target = parts.next().unwrap_or_default();
254    let version = parts.next().unwrap_or_default();
255    if parts.next().is_some()
256        || !method.eq_ignore_ascii_case("CONNECT")
257        || target.is_empty()
258        || !matches!(version, "HTTP/1.0" | "HTTP/1.1")
259    {
260        write_response(&mut client, 400, "Bad Request").await?;
261        return Err(ProxyError::Protocol(
262            "invalid CONNECT request line".to_owned(),
263        ));
264    }
265
266    let mut header_bytes = 0_usize;
267    let mut header_count = 0_usize;
268    let mut authorized = false;
269    loop {
270        let remaining = config.max_header_bytes.saturating_sub(header_bytes);
271        if remaining == 0 {
272            write_response(&mut client, 431, "Request Header Fields Too Large").await?;
273            return Err(ProxyError::Limit("total header bytes"));
274        }
275        let line = timeout_at(
276            header_deadline,
277            read_bounded_line(&mut client, remaining.min(config.max_header_line_bytes)),
278        )
279        .await
280        .map_err(|_| ProxyError::Timeout("request headers"))??;
281        header_bytes = header_bytes.saturating_add(line.len());
282        if line == b"\r\n" || line == b"\n" || line.is_empty() {
283            break;
284        }
285        header_count += 1;
286        if header_count > config.max_headers {
287            write_response(&mut client, 431, "Request Header Fields Too Large").await?;
288            return Err(ProxyError::Limit("header count"));
289        }
290        let line = std::str::from_utf8(&line)
291            .map_err(|_| ProxyError::Protocol("header is not UTF-8".to_owned()))?;
292        let (name, value) = line
293            .trim_end_matches(['\r', '\n'])
294            .split_once(':')
295            .ok_or_else(|| ProxyError::Protocol("malformed header".to_owned()))?;
296        if name.eq_ignore_ascii_case("Proxy-Authorization")
297            && constant_time_eq(value.trim().as_bytes(), expected_authorization.as_bytes())
298        {
299            authorized = true;
300        }
301    }
302    if !authorized {
303        client.write_all(
304            b"HTTP/1.1 407 Proxy Authentication Required\r\nProxy-Authenticate: Basic realm=\"sbe\"\r\nConnection: close\r\n\r\n"
305        ).await?;
306        return Err(ProxyError::Unauthorized);
307    }
308
309    let (host, port) = parse_authority(target)?;
310    if !config.allowed_ports.contains(&port) {
311        write_response(&mut client, 403, "Forbidden").await?;
312        return Err(ProxyError::Destination(format!(
313            "port {port} is not authorized"
314        )));
315    }
316    if !allowlist.is_allowed(&host) {
317        write_response(&mut client, 403, "Forbidden").await?;
318        warn!(%host, port, client = %addr, "blocked non-allowlisted domain");
319        return Err(ProxyError::Destination(format!(
320            "domain '{host}' is not authorized"
321        )));
322    }
323
324    let resolved = timeout(config.dns_timeout, lookup_host((host.as_str(), port)))
325        .await
326        .map_err(|_| ProxyError::Timeout("DNS resolution"))??;
327    let mut addresses = BTreeSet::new();
328    for address in resolved {
329        addresses.insert(address);
330        if addresses.len() > config.max_resolved_addresses {
331            write_response(&mut client, 502, "Bad Gateway").await?;
332            return Err(ProxyError::Limit("resolved address count"));
333        }
334    }
335    if addresses.is_empty() {
336        return Err(ProxyError::Destination(format!(
337            "domain '{host}' resolved to no addresses"
338        )));
339    }
340    if !config.allow_special_addresses && addresses.iter().any(|a| !is_global(a.ip())) {
341        write_response(&mut client, 403, "Forbidden").await?;
342        return Err(ProxyError::Destination(format!(
343            "domain '{host}' resolved to a special-use address"
344        )));
345    }
346    let mut upstream = timeout(
347        config.connect_timeout,
348        connect_to_validated_address(&host, port, &addresses),
349    )
350    .await
351    .map_err(|_| ProxyError::Timeout("upstream connection"))??;
352    client
353        .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n")
354        .await?;
355    client.flush().await?;
356    debug!(%host, port, client = %addr, "proxy tunnel established");
357
358    let buffered = client.buffer().to_vec();
359    if !buffered.is_empty() {
360        timeout(config.idle_timeout, upstream.write_all(&buffered))
361            .await
362            .map_err(|_| ProxyError::Timeout("initial tunnel write"))??;
363    }
364    let client = client.into_inner();
365    let (client_read, client_write) = client.into_split();
366    let (upstream_read, upstream_write) = upstream.into_split();
367    let tunnel = relay_tunnel(
368        client_read,
369        client_write,
370        upstream_read,
371        upstream_write,
372        config.idle_timeout,
373    );
374    timeout(config.tunnel_timeout, tunnel)
375        .await
376        .map_err(|_| ProxyError::Timeout("maximum tunnel lifetime"))??;
377    Ok(())
378}
379
380async fn relay_tunnel<CR, CW, UR, UW>(
381    client_read: CR,
382    client_write: CW,
383    upstream_read: UR,
384    upstream_write: UW,
385    idle_timeout: Duration,
386) -> Result<(), ProxyError>
387where
388    CR: AsyncRead + Unpin,
389    CW: AsyncWrite + Unpin,
390    UR: AsyncRead + Unpin,
391    UW: AsyncWrite + Unpin,
392{
393    let (activity_tx, activity_rx) = watch::channel(Instant::now());
394    tokio::try_join!(
395        copy_with_tunnel_activity(client_read, upstream_write, activity_tx.clone()),
396        copy_with_tunnel_activity(upstream_read, client_write, activity_tx),
397        enforce_tunnel_idle_timeout(activity_rx, idle_timeout),
398    )?;
399    Ok(())
400}
401
402async fn connect_to_validated_address(
403    host: &str,
404    port: u16,
405    addresses: &BTreeSet<SocketAddr>,
406) -> Result<TcpStream, ProxyError> {
407    let mut last_error = None;
408    for &address in addresses {
409        match TcpStream::connect(address).await {
410            Ok(stream) => return Ok(stream),
411            Err(error) => last_error = Some(error),
412        }
413    }
414    Err(ProxyError::UpstreamConnect {
415        host: host.to_owned(),
416        port,
417        source: last_error.unwrap_or_else(|| std::io::Error::other("no usable address")),
418    })
419}
420
421async fn copy_with_tunnel_activity<R, W>(
422    mut reader: R,
423    mut writer: W,
424    activity: watch::Sender<Instant>,
425) -> Result<u64, ProxyError>
426where
427    R: AsyncRead + Unpin,
428    W: AsyncWrite + Unpin,
429{
430    let mut copied = 0_u64;
431    let mut buffer = [0_u8; 16 * 1024];
432    loop {
433        let count = reader.read(&mut buffer).await?;
434        if count == 0 {
435            writer.shutdown().await?;
436            return Ok(copied);
437        }
438        activity.send_replace(Instant::now());
439        let mut written = 0;
440        while written < count {
441            let amount = writer.write(&buffer[written..count]).await?;
442            if amount == 0 {
443                return Err(std::io::Error::from(std::io::ErrorKind::WriteZero).into());
444            }
445            written += amount;
446            activity.send_replace(Instant::now());
447        }
448        copied = copied.saturating_add(count as u64);
449    }
450}
451
452/// Expire a tunnel only when neither direction has transferred bytes for the
453/// configured interval. Each copy task updates the same watch channel, so a
454/// continuous download keeps an otherwise quiet upload direction alive.
455async fn enforce_tunnel_idle_timeout(
456    mut activity: watch::Receiver<Instant>,
457    idle_timeout: Duration,
458) -> Result<(), ProxyError> {
459    loop {
460        let deadline = *activity.borrow_and_update() + idle_timeout;
461        match timeout_at(deadline, activity.changed()).await {
462            Ok(Ok(())) => continue,
463            Ok(Err(_closed)) => return Ok(()),
464            Err(_elapsed) => {
465                if Instant::now().saturating_duration_since(*activity.borrow()) >= idle_timeout {
466                    return Err(ProxyError::Timeout("tunnel idle period"));
467                }
468            }
469        }
470    }
471}
472
473async fn read_bounded_line<R>(reader: &mut R, maximum: usize) -> Result<Vec<u8>, ProxyError>
474where
475    R: AsyncBufRead + Unpin,
476{
477    let mut output = Vec::new();
478    loop {
479        let available = reader.fill_buf().await?;
480        if available.is_empty() {
481            return Ok(output);
482        }
483        let consumed = available
484            .iter()
485            .position(|byte| *byte == b'\n')
486            .map_or(available.len(), |index| index + 1);
487        if output.len().saturating_add(consumed) > maximum {
488            return Err(ProxyError::Limit("line bytes"));
489        }
490        output.extend_from_slice(&available[..consumed]);
491        reader.consume(consumed);
492        if output.last() == Some(&b'\n') {
493            return Ok(output);
494        }
495    }
496}
497
498fn parse_authority(target: &str) -> Result<(String, u16), ProxyError> {
499    if target.starts_with('[') || target.parse::<IpAddr>().is_ok() {
500        return Err(ProxyError::Destination(
501            "IP literals are not authorized".to_owned(),
502        ));
503    }
504    let (raw_host, raw_port) = target
505        .rsplit_once(':')
506        .ok_or_else(|| ProxyError::Protocol("CONNECT target has no port".to_owned()))?;
507    if raw_host.is_empty() || raw_host.contains(':') {
508        return Err(ProxyError::Protocol(
509            "CONNECT target has an invalid host".to_owned(),
510        ));
511    }
512    let port = raw_port
513        .parse::<u16>()
514        .map_err(|_| ProxyError::Protocol("CONNECT target has an invalid port".to_owned()))?;
515    if port == 0 {
516        return Err(ProxyError::Protocol(
517            "CONNECT target port must be non-zero".to_owned(),
518        ));
519    }
520    let host = idna::domain_to_ascii(raw_host.trim_end_matches('.'))
521        .map_err(|_| ProxyError::Protocol("CONNECT target has invalid IDNA".to_owned()))?
522        .to_ascii_lowercase();
523    validate_dns_name(&host)?;
524    Ok((host, port))
525}
526
527fn validate_dns_name(host: &str) -> Result<(), ProxyError> {
528    if host.is_empty() || host.len() > 253 || host.parse::<IpAddr>().is_ok() {
529        return Err(ProxyError::Protocol("invalid DNS hostname".to_owned()));
530    }
531    for label in host.split('.') {
532        if label.is_empty()
533            || label.len() > 63
534            || label.starts_with('-')
535            || label.ends_with('-')
536            || !label
537                .bytes()
538                .all(|b| b.is_ascii_alphanumeric() || b == b'-')
539        {
540            return Err(ProxyError::Protocol("invalid DNS hostname".to_owned()));
541        }
542    }
543    Ok(())
544}
545
546fn is_global(ip: IpAddr) -> bool {
547    match ip {
548        IpAddr::V4(ip) => is_global_v4(ip),
549        IpAddr::V6(ip) => is_global_v6(ip),
550    }
551}
552
553fn is_global_v4(ip: Ipv4Addr) -> bool {
554    let [a, b, c, _] = ip.octets();
555    !(ip.is_private()
556        || ip.is_loopback()
557        || ip.is_link_local()
558        || ip.is_broadcast()
559        || ip.is_documentation()
560        || ip.is_unspecified()
561        || ip.is_multicast()
562        || a == 0
563        || a >= 240
564        || (a == 100 && (64..=127).contains(&b))
565        || (a == 192 && b == 0 && c == 0)
566        || (a == 192 && b == 88 && c == 99)
567        || (a == 198 && (18..=19).contains(&b)))
568}
569
570fn is_global_v6(ip: Ipv6Addr) -> bool {
571    let segments = ip.segments();
572    if let Some(mapped) = ip.to_ipv4_mapped() {
573        return is_global_v4(mapped);
574    }
575    !(ip.is_unspecified()
576        || ip.is_loopback()
577        || ip.is_multicast()
578        || (segments[0] & 0xfe00) == 0xfc00
579        || (segments[0] & 0xffc0) == 0xfe80
580        || (segments[0] & 0xffc0) == 0xfec0
581        // Local-use translation, discard-only, and IETF protocol-assignment
582        // blocks are not ordinary globally routable destinations.
583        || (segments[0] == 0x0064 && segments[1] == 0xff9b && segments[2] == 1)
584        || (segments[0] == 0x0100
585            && segments[1] == 0
586            && segments[2] == 0
587            && segments[3] == 0)
588        || (segments[0] == 0x2001 && segments[1] <= 0x01ff)
589        || (segments[0] == 0x2001 && segments[1] == 0x0db8)
590        || (segments[0] == 0x3fff && segments[1] < 0x1000)
591        || segments[0] == 0x5f00)
592}
593
594async fn write_response<W: AsyncWrite + Unpin>(
595    stream: &mut W,
596    status: u16,
597    reason: &str,
598) -> Result<(), ProxyError> {
599    stream
600        .write_all(format!("HTTP/1.1 {status} {reason}\r\nConnection: close\r\n\r\n").as_bytes())
601        .await?;
602    Ok(())
603}
604
605fn constant_time_eq(left: &[u8], right: &[u8]) -> bool {
606    if left.len() != right.len() {
607        return false;
608    }
609    let mut difference = 0_u8;
610    for (&left, &right) in left.iter().zip(right) {
611        difference |= left ^ right;
612    }
613    difference == 0
614}
615
616fn hex(bytes: &[u8]) -> String {
617    const HEX: &[u8; 16] = b"0123456789abcdef";
618    let mut output = String::with_capacity(bytes.len() * 2);
619    for byte in bytes {
620        output.push(HEX[(byte >> 4) as usize] as char);
621        output.push(HEX[(byte & 0x0f) as usize] as char);
622    }
623    output
624}
625
626#[cfg(test)]
627mod tests {
628    use super::*;
629
630    #[test]
631    fn parses_and_canonicalizes_authority() {
632        assert_eq!(
633            parse_authority("Registry.NPMJS.org.:443").unwrap(),
634            ("registry.npmjs.org".to_owned(), 443)
635        );
636    }
637
638    #[test]
639    fn endpoint_debug_output_redacts_authentication_token() {
640        let endpoint = ProxyEndpoint {
641            port: 12345,
642            token: "sentinel-token".to_owned(),
643        };
644        let output = format!("{endpoint:?}");
645        assert!(output.contains("<redacted>"));
646        assert!(!output.contains("sentinel-token"));
647    }
648
649    #[test]
650    fn java_environment_uses_agent_without_logging_token() {
651        let endpoint = ProxyEndpoint {
652            port: 12345,
653            token: "sentinel-token".to_owned(),
654        };
655        let environment = endpoint.java_environment("/private/sbe/proxy-agent.jar", "/private/sbe");
656        let options = &environment[0].1;
657
658        assert!(options.contains("-javaagent:/private/sbe/proxy-agent.jar"));
659        assert!(options.contains("-Djava.io.tmpdir=/private/sbe"));
660        assert!(options.contains("-Dhttp.proxyProtocol=http"));
661        assert!(options.contains("-Dhttps.proxyProtocol=http"));
662        assert!(options.contains("-Djdk.http.auth.tunneling.disabledSchemes="));
663        assert!(!options.contains("sentinel-token"));
664        assert_eq!(
665            environment[1],
666            ("SBE_PROXY_TOKEN", "sentinel-token".to_owned())
667        );
668    }
669
670    #[test]
671    fn malformed_or_literal_authorities_never_panic() {
672        for target in ["]:443", "[::1]:443", "127.0.0.1:443", "a:b:443", "x:0", "x"] {
673            assert!(parse_authority(target).is_err(), "accepted {target}");
674        }
675    }
676
677    #[test]
678    fn special_addresses_are_not_global() {
679        for address in [
680            "127.0.0.1",
681            "10.0.0.1",
682            "169.254.1.1",
683            "192.0.2.1",
684            "::1",
685            "fe80::1",
686            "fc00::1",
687            "2001:db8::1",
688            "64:ff9b:1::1",
689            "100::1",
690            "2001:20::1",
691            "3fff::1",
692            "5f00::1",
693            "fec0::1",
694        ] {
695            assert!(!is_global(address.parse().unwrap()), "accepted {address}");
696        }
697        assert!(is_global("1.1.1.1".parse().unwrap()));
698        assert!(is_global("2606:4700:4700::1111".parse().unwrap()));
699    }
700
701    #[tokio::test]
702    async fn bounded_line_reader_rejects_oversized_input() {
703        let (mut writer, reader) = tokio::io::duplex(128);
704        tokio::spawn(async move {
705            writer.write_all(&[b'a'; 65]).await.unwrap();
706            writer.write_all(b"\n").await.unwrap();
707        });
708        let mut reader = BufReader::new(reader);
709        assert!(matches!(
710            read_bounded_line(&mut reader, 64).await,
711            Err(ProxyError::Limit(_))
712        ));
713    }
714
715    #[tokio::test]
716    async fn tunnel_enforces_aggregate_idle_timeout() {
717        let (activity_tx, activity_rx) = watch::channel(Instant::now());
718        let result = enforce_tunnel_idle_timeout(activity_rx, Duration::from_millis(10)).await;
719        drop(activity_tx);
720        assert!(matches!(
721            result,
722            Err(ProxyError::Timeout("tunnel idle period"))
723        ));
724    }
725
726    #[tokio::test]
727    async fn activity_in_either_direction_resets_tunnel_idle_timeout() {
728        let (activity_tx, activity_rx) = watch::channel(Instant::now());
729        let monitor = tokio::spawn(enforce_tunnel_idle_timeout(
730            activity_rx,
731            Duration::from_millis(40),
732        ));
733
734        for _ in 0..3 {
735            tokio::time::sleep(Duration::from_millis(25)).await;
736            activity_tx.send_replace(Instant::now());
737        }
738        drop(activity_tx);
739
740        assert!(monitor.await.unwrap().is_ok());
741    }
742
743    #[tokio::test]
744    async fn continuous_one_way_transfer_keeps_quiet_direction_alive() {
745        let (client_app, relay_client) = tokio::io::duplex(64);
746        let (upstream_app, relay_upstream) = tokio::io::duplex(64);
747        let (client_read, client_write) = tokio::io::split(relay_client);
748        let (upstream_read, upstream_write) = tokio::io::split(relay_upstream);
749        let relay = tokio::spawn(relay_tunnel(
750            client_read,
751            client_write,
752            upstream_read,
753            upstream_write,
754            Duration::from_millis(200),
755        ));
756        let (mut client_output, mut client_input) = tokio::io::split(client_app);
757        let (upstream_output, mut upstream_input) = tokio::io::split(upstream_app);
758
759        for byte in b"download" {
760            tokio::time::sleep(Duration::from_millis(40)).await;
761            upstream_input.write_all(&[*byte]).await.unwrap();
762        }
763        upstream_input.shutdown().await.unwrap();
764        client_input.shutdown().await.unwrap();
765
766        let mut received = Vec::new();
767        client_output.read_to_end(&mut received).await.unwrap();
768        assert_eq!(received, b"download");
769        let relay_result = relay.await.unwrap();
770        assert!(relay_result.is_ok(), "{relay_result:?}");
771        drop(upstream_output);
772    }
773}