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#[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 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#[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 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
122pub 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 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
452async 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 || (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}