diff --git a/CHANGELOG.md b/CHANGELOG.md index 45da1625..f8317099 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,15 @@ # Changelog +## Unreleased + +- Hardened malformed protocol framing, credential comparisons, replay rejection, and secret-safe diagnostics. +- Preserved handshake early data and response ordering in TCP, TUN, Hysteria2, and TUIC forwarding; bounded initial writes and flushes. +- Added `check` and `version` aliases, nonzero validation failures, last-good configuration preparation, and explicit signal handling. SIGHUP reloads even with `--no-reload`; orderly signal shutdown exits with status 0. +- Corrected Snell salts to 16 bytes for all ciphers. ChaCha20 and AES-256 connections to older Shoes versions using 32-byte salts are not wire-compatible; upgrade both ends together. +- Hysteria2 now reports actual connection outcomes, emits bracketed IPv6 addresses, and advertises `Hysteria-CC-RX: auto`. Compatible clients select automatic congestion control (BBR in sing-box), even when a client bandwidth is configured. +- Unsupported VMess cipher names, including `auto`, are rejected during validation instead of panicking at startup. Use `any` for automatic cipher selection. +- Improved QUIC socket portability, SOCKS UDP compatibility with NAT clients, and synchronous mobile lifecycle safety. + ## v0.3.2 ### Fixes diff --git a/CONFIG.md b/CONFIG.md index 62ec2ff1..8c4ec1a2 100644 --- a/CONFIG.md +++ b/CONFIG.md @@ -144,6 +144,8 @@ protocol: **Note:** VMess AEAD mode is always enabled. The legacy `force_aead` field is deprecated and non-AEAD mode is no longer supported. +Replay protection uses a process-wide cache of at most 65,536 Auth-IDs, shared across credentials, listeners, and reloads. Entries remain protected through their timestamp's 120-second validity window; future-dated IDs can occupy a slot for about four minutes. A full cache rejects new authentication rather than evicting live entries. Heavy traffic from one credential can therefore exhaust capacity for other VMess users in the same process. A backward system-clock step also pauses authentication until the clock catches up. + ### VLESS ```yaml protocol: @@ -778,7 +780,19 @@ client_chain: ### Hot Reloading -Configuration changes are automatically detected and applied without restarting. Disable with `--no-reload` flag. +Configuration changes are automatically detected. `--no-reload` disables file watching; on Unix, SIGHUP still requests a reload. + +On Linux, if directory permissions prevent watching a readable configuration's parent, startup falls back to direct file watches and logs a warning. In-place edits remain detectable, but atomic replacements may be missed; use SIGHUP or grant directory read permission when replacing files. + +Replacement configurations, certificates, and DNS resolvers are prepared while the current listeners continue serving. A preparation failure leaves the current configuration running. Once preparation succeeds, listeners are replaced. A failure to activate replacements exits with status 1; there is no rollback after listeners are retired. Existing TCP connections may drain for `global_limits.reload_grace_secs` during successful reloads; QUIC connections are closed. + +QUIC listener sockets must finish closing before replacements bind. Closing endpoints get up to one second for best-effort close delivery, then their drivers are stopped locally so peer-controlled timers cannot delay reloads. A five-second socket-release guard remains; exceeding it indicates a local cleanup failure and exits with an error rather than announcing a partially ready generation. + +Candidate QUIC buffer sizes are used during DNS preparation, but `global_limits` admission caps change only at activation. DNS bootstrap over HTTP/3 or QUIC proxy chains therefore needs headroom under the live cap. A cap-increasing reload may require a restart if preparation cannot fit. Lowering a cap does not revoke existing reservations, including those made during preparation; new admissions are blocked until usage falls sufficiently. At cold startup, preparation uses the initial unlimited admission caps. + +On macOS, edit the configured file path or use SIGHUP when writing through a hard-link alias in another directory. Automatic detection of such alias writes is backend-dependent. System-volume directory symlinks are resolved without watching their immutable links. + +SIGINT/SIGTERM (or Ctrl-C on Windows) stop listeners and exit with status 0. Process shutdown does not wait for the TCP reload grace period. ### mTLS (Mutual TLS) @@ -813,14 +827,18 @@ shoes [OPTIONS] [config.yaml...] OPTIONS: -t, --threads NUM Worker threads (default: CPU count) - -d, --dry-run Parse config and exit - --no-reload Disable hot-reloading + -d, --dry-run Validate config and certificates without starting listeners + --no-reload Disable automatic reloads on file changes COMMANDS: + check Alias for --dry-run + version Alias for --version generate-reality-keypair Generate Reality X25519 keypair generate-shadowsocks-2022-password Generate Shadowsocks 2022 password ``` +`check` and `--dry-run` exit with status 0 for valid configuration and 1 for errors. They do not open listeners, create a file watcher, or perform DNS bootstrap networking. Runtime failures such as unavailable bind addresses remain possible after a successful check. + ## Tips ### Generate Keys diff --git a/Cargo.lock b/Cargo.lock index 6ea9139e..15da12e6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2136,11 +2136,22 @@ dependencies = [ "rcgen", "rustls", "rustls-pemfile", + "serde_json", "tempfile", "tokio", "tokio-rustls", ] +[[package]] +name = "signal-hook-registry" +version = "1.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" +dependencies = [ + "errno", + "libc", +] + [[package]] name = "simd_cesu8" version = "1.2.0" @@ -2372,6 +2383,7 @@ dependencies = [ "libc", "mio", "pin-project-lite", + "signal-hook-registry", "socket2", "tokio-macros", "windows-sys 0.61.2", @@ -2407,6 +2419,7 @@ dependencies = [ "bytes", "futures-core", "futures-sink", + "futures-util", "libc", "pin-project-lite", "slab", diff --git a/Cargo.toml b/Cargo.toml index 3b88f8c3..b9932d98 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,6 +35,10 @@ external-network-tests = ["integration-tests"] privileged-tests = ["integration-tests"] stress-tests = ["integration-tests"] +[[test]] +name = "process_contract" +path = "tests/process_contract.rs" + [[test]] name = "integration" path = "tests/integration.rs" @@ -108,10 +112,11 @@ tokio = { version = "*", features = [ "macros", "net", "rt-multi-thread", + "signal", "sync", "time", ] } -tokio-util = { version = "*", features = ["time", "io"] } +tokio-util = { version = "*", features = ["time", "io", "rt"] } tun = { version = "*", features = ["async"] } url = "*" webpki-roots = { version = "*" } diff --git a/README.md b/README.md index 420f092e..4dc3a16b 100644 --- a/README.md +++ b/README.md @@ -67,10 +67,12 @@ shoes [OPTIONS] [config.yaml...] OPTIONS: -t, --threads NUM Set the number of worker threads (default: CPU count) - -d, --dry-run Parse the config and exit + -d, --dry-run Validate the config and certificates, then exit --no-reload Disable automatic config reloading on file changes COMMANDS: + check Alias for --dry-run + version Alias for --version generate-reality-keypair Generate a new Reality X25519 keypair generate-shadowsocks-2022-password Generate a Shadowsocks password ``` diff --git a/src/address.rs b/src/address.rs index bd4fd9bc..63b06623 100644 --- a/src/address.rs +++ b/src/address.rs @@ -171,7 +171,6 @@ impl NetLocation { Ok(Self { address, port }) } - #[cfg(test)] pub fn from_ip_addr(ip: IpAddr, port: u16) -> Self { let address = match ip { IpAddr::V4(addr) => Address::Ipv4(addr), diff --git a/src/anytls/anytls_stream.rs b/src/anytls/anytls_stream.rs index b499ccc9..371daf26 100644 --- a/src/anytls/anytls_stream.rs +++ b/src/anytls/anytls_stream.rs @@ -239,14 +239,6 @@ impl AsyncRead for AnyTlsStream { if buf.remaining() == 0 { return Poll::Ready(Ok(())); } - // Check if stream/session is closed - if self.stream_closed { - return Poll::Ready(Err(io::Error::new( - io::ErrorKind::BrokenPipe, - "stream closed", - ))); - } - // Check EOF with empty buffer let remaining_in_buffer = self.read_buffer.len() - self.read_offset; if self.eof && remaining_in_buffer == 0 { @@ -368,6 +360,7 @@ impl AsyncWrite for AnyTlsStream { return Poll::Ready(Ok(())); } if self.session_closed.load(Ordering::Relaxed) { + self.data_rx.close(); self.stream_closed = true; return Poll::Ready(Ok(())); } @@ -383,6 +376,7 @@ impl AsyncWrite for AnyTlsStream { self.fin_sent = true; } std::task::ready!(self.data_tx.poll_flush(cx))?; + self.data_rx.close(); self.stream_closed = true; Poll::Ready(Ok(())) } @@ -420,6 +414,24 @@ mod tests { use super::*; use tokio::io::{AsyncReadExt, AsyncWriteExt}; + #[tokio::test] + async fn shutdown_keeps_buffered_and_queued_response_bytes_readable() { + let (outgoing, _outgoing_rx) = mpsc::channel(STREAM_CHANNEL_BUFFER); + let (incoming, incoming_rx) = mpsc::channel(STREAM_CHANNEL_BUFFER); + let mut stream = + AnyTlsStream::new(1, incoming_rx, outgoing, Arc::new(AtomicBool::new(false))); + incoming.send(Bytes::from_static(b"first")).await.unwrap(); + incoming.send(Bytes::from_static(b"second")).await.unwrap(); + let mut first = [0; 1]; + stream.read_exact(&mut first).await.unwrap(); + stream.shutdown().await.unwrap(); + assert!(stream.write_all(b"closed").await.is_err()); + let mut rest = Vec::new(); + stream.read_to_end(&mut rest).await.unwrap(); + assert_eq!(first, *b"f"); + assert_eq!(rest, b"irstsecond"); + } + #[tokio::test] async fn test_stream_write() { let (data_tx, mut data_rx) = mpsc::channel(STREAM_CHANNEL_BUFFER); diff --git a/src/config/pem.rs b/src/config/pem.rs index 4659972f..002c16cd 100644 --- a/src/config/pem.rs +++ b/src/config/pem.rs @@ -258,6 +258,16 @@ fn gather_pem_file_paths_from_server_proxy( { process_pem_path(&mut handshake.cert, known_pem_paths, unknown_pem_paths); process_pem_path(&mut handshake.key, known_pem_paths, unknown_pem_paths); + } else if let ShadowTlsServerHandshakeConfig::Remote(handshake) = + &mut shadowtls_config.handshake + { + for hop in handshake.client_chain.iter_mut() { + gather_pem_file_paths_from_chain_hop( + hop, + known_pem_paths, + unknown_pem_paths, + ); + } } // Recurse into inner protocol gather_pem_file_paths_from_server_proxy( @@ -273,6 +283,9 @@ fn gather_pem_file_paths_from_server_proxy( // Process Reality targets for (sni, reality_config) in reality_targets.iter_mut() { + for hop in reality_config.dest_client_chain.iter_mut() { + gather_pem_file_paths_from_chain_hop(hop, known_pem_paths, unknown_pem_paths); + } // Validate Vision configuration validate_vision_protocol( reality_config.vision, diff --git a/src/config/validate.rs b/src/config/validate.rs index ffe8b3a7..1d59e531 100644 --- a/src/config/validate.rs +++ b/src/config/validate.rs @@ -4,7 +4,7 @@ use std::collections::{HashMap, HashSet}; use crate::address::NetLocationMask; use crate::dns::ParsedDnsUrl; -use crate::option_util::{NoneOrSome, OneOrSome}; +use crate::option_util::{NoneOrOne, NoneOrSome, OneOrSome}; use crate::reality::{decode_private_key, decode_short_id}; use crate::thread_util::get_num_threads; use crate::uuid_util::parse_uuid; @@ -642,6 +642,12 @@ fn validate_server_config( rule_groups: &HashMap>, named_pems: &HashMap, ) -> std::io::Result<()> { + if server_config.transport == Transport::Udp { + return Err(std::io::Error::new( + std::io::ErrorKind::Unsupported, + "UDP listener transport is not supported", + )); + } // First handle QUIC settings certificates if let Some(ref mut quic_settings) = server_config.quic_settings { embed_pem_from_map(&mut quic_settings.cert, named_pems); @@ -649,6 +655,11 @@ fn validate_server_config( for cert in quic_settings.client_ca_certs.iter_mut() { embed_pem_from_map(cert, named_pems); } + validate_server_certificates( + &quic_settings.cert, + &quic_settings.key, + &quic_settings.client_ca_certs, + )?; } if server_config.transport != Transport::Tcp && server_config.tcp_settings.is_some() { return Err(std::io::Error::new( @@ -858,24 +869,29 @@ fn validate_client_vision_protocol( /// Recursive validation of client proxy config structure (Vision rules, etc.) fn validate_client_proxy_structure(config: &ClientProxyConfig) -> std::io::Result<()> { - match config { + let inner = match config { ClientProxyConfig::Tls(tls_config) => { validate_client_vision_protocol(tls_config.vision, &tls_config.protocol, "TLS")?; - validate_client_proxy_structure(&tls_config.protocol)?; + Some(tls_config.protocol.as_ref()) } ClientProxyConfig::Reality { vision, protocol, .. } => { validate_client_vision_protocol(*vision, protocol, "Reality")?; - validate_client_proxy_structure(protocol)?; + Some(protocol.as_ref()) } - ClientProxyConfig::ShadowTls { protocol, .. } => { - validate_client_proxy_structure(protocol)?; + ClientProxyConfig::ShadowTls { protocol, .. } => Some(protocol.as_ref()), + ClientProxyConfig::Websocket(ws_config) => Some(ws_config.protocol.as_ref()), + _ => None, + }; + if let Some(inner) = inner { + if inner.is_direct() { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "direct cannot be used as an inner proxy protocol", + )); } - ClientProxyConfig::Websocket(ws_config) => { - validate_client_proxy_structure(&ws_config.protocol)?; - } - _ => {} + validate_client_proxy_structure(inner)?; } Ok(()) } @@ -915,6 +931,7 @@ fn validate_client_config( )); } validate_server_fingerprints(server_fingerprints)?; + validate_client_certificates(cert.as_deref(), key.as_deref())?; } #[cfg(not(any(target_os = "android", target_os = "fuchsia", target_os = "linux")))] @@ -925,7 +942,12 @@ fn validate_client_config( )); } - validate_client_proxy_config(&mut client_config.protocol, named_pems)?; + validate_client_proxy_structure(&client_config.protocol)?; + validate_client_proxy_config( + &mut client_config.protocol, + named_pems, + client_config.address.address().hostname(), + )?; Ok(()) } @@ -951,13 +973,42 @@ fn validate_server_fingerprints( fn validate_client_proxy_config( client_proxy_config: &mut ClientProxyConfig, named_pems: &HashMap, + default_sni: Option<&str>, ) -> std::io::Result<()> { - validate_client_proxy_structure(client_proxy_config)?; - match client_proxy_config { - ClientProxyConfig::Reality { - short_id, protocol, .. + ClientProxyConfig::Vmess { + user_id, cipher, .. } => { + parse_uuid(user_id)?; + validate_vmess_cipher(cipher)?; + } + ClientProxyConfig::Vless { user_id, .. } => { + parse_uuid(user_id)?; + } + ClientProxyConfig::Snell { + config: super::types::ShadowsocksConfig::Aead2022 { .. }, + .. + } => { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "Snell does not support shadowsocks 2022 ciphers", + )); + } + ClientProxyConfig::Reality { + public_key, + short_id, + sni_hostname, + protocol, + .. + } => { + crate::reality::decode_public_key(public_key).map_err(std::io::Error::other)?; + let sni = sni_hostname.as_deref().or(default_sni).ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "REALITY client requires sni_hostname", + ) + })?; + crate::reality::parse_server_name(sni)?; validate_reality_client_short_id(short_id)?; if short_id == DEFAULT_REALITY_SHORT_ID { @@ -967,10 +1018,18 @@ fn validate_client_proxy_config( ); } - validate_client_proxy_config(protocol, named_pems)?; + validate_client_proxy_config(protocol, named_pems, None)?; } ClientProxyConfig::Tls(tls_config) => { + let sni = match &tls_config.sni_hostname { + NoneOrOne::Unspecified => default_sni, + NoneOrOne::None => None, + NoneOrOne::One(sni) => Some(sni.as_str()), + }; + if let Some(sni) = sni { + validate_sni(sni)?; + } embed_optional_pem_from_map(&mut tls_config.cert, named_pems); embed_optional_pem_from_map(&mut tls_config.key, named_pems); @@ -981,16 +1040,32 @@ fn validate_client_proxy_config( )); } validate_server_fingerprints(&mut tls_config.server_fingerprints)?; + validate_client_certificates(tls_config.cert.as_deref(), tls_config.key.as_deref())?; - validate_client_proxy_config(&mut tls_config.protocol, named_pems)?; + validate_client_proxy_config(&mut tls_config.protocol, named_pems, None)?; } - ClientProxyConfig::ShadowTls { protocol, .. } => { - validate_client_proxy_config(protocol, named_pems)?; + ClientProxyConfig::ShadowTls { + protocol, + sni_hostname, + .. + } => { + if let Some(sni) = sni_hostname.as_deref().or(default_sni) { + validate_sni(sni)?; + } + validate_client_proxy_config(protocol, named_pems, None)?; } ClientProxyConfig::Websocket(ws_config) => { - validate_client_proxy_config(&mut ws_config.protocol, named_pems)?; + validate_client_proxy_config(&mut ws_config.protocol, named_pems, None)?; + } + + ClientProxyConfig::Anytls { + padding_scheme: Some(lines), + .. + } => { + crate::anytls::PaddingFactory::new(lines.join("\n").as_bytes()) + .map_err(std::io::Error::other)?; } _ => {} @@ -998,6 +1073,47 @@ fn validate_client_proxy_config( Ok(()) } +fn validate_sni(sni: &str) -> std::io::Result<()> { + rustls::pki_types::ServerName::try_from(sni) + .map(|_| ()) + .map_err(|_| { + std::io::Error::new(std::io::ErrorKind::InvalidInput, "invalid TLS server name") + }) +} + +fn validate_server_certificates( + cert: &str, + key: &str, + ca: &NoneOrSome, +) -> std::io::Result<()> { + crate::rustls_config_util::try_create_server_config( + cert.as_bytes(), + key.as_bytes(), + ca.iter().map(|pem| pem.as_bytes().to_vec()).collect(), + &[], + &[], + ) + .map(|_| ()) +} + +fn validate_client_certificates(cert: Option<&str>, key: Option<&str>) -> std::io::Result<()> { + if let (Some(cert), Some(key)) = (cert, key) { + crate::rustls_config_util::try_create_client_config( + false, + vec![], + vec![], + true, + Some((key.as_bytes().to_vec(), cert.as_bytes().to_vec())), + false, + )?; + } + Ok(()) +} + +fn validate_vmess_cipher(cipher: &str) -> std::io::Result<()> { + crate::vmess::DataCipher::try_from(cipher).map(|_| ()) +} + fn validate_server_proxy_config( server_proxy_config: &mut ServerProxyConfig, client_groups: &HashMap>, @@ -1006,6 +1122,13 @@ fn validate_server_proxy_config( inside_tls_or_reality: bool, ) -> std::io::Result<()> { match server_proxy_config { + ServerProxyConfig::Anytls { + padding_scheme: Some(lines), + .. + } => { + crate::anytls::PaddingFactory::new(lines.join("\n").as_bytes()) + .map_err(std::io::Error::other)?; + } ServerProxyConfig::Naiveproxy { .. } if !inside_tls_or_reality => { return Err(std::io::Error::new( std::io::ErrorKind::InvalidInput, @@ -1016,8 +1139,11 @@ fn validate_server_proxy_config( ServerProxyConfig::Vless { user_id, .. } => { parse_uuid(user_id)?; } - ServerProxyConfig::Vmess { user_id, .. } => { + ServerProxyConfig::Vmess { + user_id, cipher, .. + } => { parse_uuid(user_id)?; + validate_vmess_cipher(cipher)?; } ServerProxyConfig::Tls { tls_targets, @@ -1042,6 +1168,11 @@ fn validate_server_proxy_config( for cert in tls_server_config.client_ca_certs.iter_mut() { embed_pem_from_map(cert, named_pems); } + validate_server_certificates( + &tls_server_config.cert, + &tls_server_config.key, + &tls_server_config.client_ca_certs, + )?; let TlsServerConfig { ref mut protocol, @@ -1076,6 +1207,12 @@ fn validate_server_proxy_config( for cert in tls_server_config.client_ca_certs.iter_mut() { embed_pem_from_map(cert, named_pems); } + validate_server_certificates( + &tls_server_config.cert, + &tls_server_config.key, + &tls_server_config.client_ca_certs, + )?; + validate_client_fingerprints(&mut tls_server_config.client_fingerprints)?; let TlsServerConfig { ref mut protocol, @@ -1119,7 +1256,22 @@ fn validate_server_proxy_config( if let ShadowTlsServerHandshakeConfig::Local(local_handshake) = handshake { embed_pem_from_map(&mut local_handshake.cert, named_pems); embed_pem_from_map(&mut local_handshake.key, named_pems); + for cert in local_handshake.client_ca_certs.iter_mut() { + embed_pem_from_map(cert, named_pems); + } + validate_server_certificates( + &local_handshake.cert, + &local_handshake.key, + &local_handshake.client_ca_certs, + )?; validate_client_fingerprints(&mut local_handshake.client_fingerprints)?; + } else if let ShadowTlsServerHandshakeConfig::Remote(remote_handshake) = handshake { + validate_client_chain( + &mut remote_handshake.client_chain, + 0, + client_groups, + named_pems, + )?; } validate_server_proxy_config( @@ -1162,6 +1314,14 @@ fn validate_server_proxy_config( validate_reality_private_key(&reality_config.private_key, sni_hostname)?; validate_reality_server_short_ids(&reality_config.short_ids, sni_hostname)?; + if !reality_config.dest_client_chain.is_empty() { + let mut hops = OneOrSome::Some( + std::mem::take(&mut reality_config.dest_client_chain).into_vec(), + ); + validate_client_chain(&mut hops, 0, client_groups, named_pems)?; + reality_config.dest_client_chain = NoneOrSome::Some(hops.into_vec()); + } + validate_server_proxy_config( &mut reality_config.protocol, client_groups, @@ -1236,6 +1396,9 @@ fn validate_server_proxy_config( "Snell does not support shadowsocks 2022 ciphers", )); } + ServerProxyConfig::Snell { cipher, .. } => { + crate::shadowsocks::ShadowsocksCipher::try_from(cipher.as_str())?; + } _ => (), } Ok(()) @@ -1321,22 +1484,27 @@ fn validate_rule_config( )); } - // Validate each chain for (chain_index, chain) in client_chains.iter_mut().enumerate() { - // First validate all hops in this chain - for hop in chain.hops.iter_mut() { - validate_client_chain_hop(hop, client_groups, named_pems)?; - } - // Then expand group references to inline configs - expand_client_chain(&mut chain.hops, client_groups)?; - // Validate that direct connectors only appear at hop 0 - validate_direct_connector_positions(&chain.hops, chain_index)?; + validate_client_chain(&mut chain.hops, chain_index, client_groups, named_pems)?; } } Ok(()) } +fn validate_client_chain( + hops: &mut OneOrSome, + chain_index: usize, + client_groups: &HashMap>, + named_pems: &HashMap, +) -> std::io::Result<()> { + for hop in hops.iter_mut() { + validate_client_chain_hop(hop, client_groups, named_pems)?; + } + expand_client_chain(hops, client_groups)?; + validate_direct_connector_positions(hops, chain_index) +} + /// Validates that direct connectors only appear at hop 0. /// /// Direct connectors can only be used as the first hop in a chain because they @@ -1506,6 +1674,74 @@ mod tests { use crate::config::pem::convert_cert_paths; use crate::dns::IpStrategy; + #[test] + fn nested_structure_errors_precede_protocol_field_errors() { + let mut config = serde_json::from_value(serde_json::json!({ + "address": "example.com:443", + "protocol": { + "type": "reality", "public_key": "invalid", + "protocol": {"type": "websocket", "protocol": {"type": "direct"}} + } + })) + .unwrap(); + let error = validate_client_config(&mut config, &HashMap::new()).unwrap_err(); + assert_eq!(error.kind(), std::io::ErrorKind::InvalidInput); + assert_eq!( + error.to_string(), + "direct cannot be used as an inner proxy protocol" + ); + } + + #[test] + fn tls_sni_distinguishes_inheritance_disable_and_override() { + for (sni, valid) in [ + (NoneOrOne::Unspecified, false), + (NoneOrOne::None, true), + (NoneOrOne::One("example.com".to_owned()), true), + (NoneOrOne::One("invalid name".to_owned()), false), + ] { + let mut protocol = serde_json::from_value(serde_json::json!({ + "type": "tls", "protocol": {"type": "socks"} + })) + .unwrap(); + let ClientProxyConfig::Tls(ref mut tls) = protocol else { + unreachable!() + }; + tls.sni_hostname = sni; + assert_eq!( + validate_client_proxy_config(&mut protocol, &HashMap::new(), Some("invalid name")) + .is_ok(), + valid, + ); + } + } + + #[test] + fn reality_sni_defaults_only_at_the_outer_client_layer() { + let protocol = serde_json::json!({ + "type": "reality", "public_key": "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA", + "protocol": {"type": "socks"} + }); + for (address, protocol, valid) in [ + ("example.com:443", protocol.clone(), true), + ("127.0.0.1:443", protocol.clone(), false), + ( + "example.com:443", + serde_json::json!({"type": "websocket", "protocol": protocol}), + false, + ), + ] { + let mut config: ClientConfig = serde_json::from_value(serde_json::json!({ + "address": address, "protocol": protocol + })) + .unwrap(); + assert_eq!( + validate_client_config(&mut config, &HashMap::new()).is_ok(), + valid + ); + } + } + async fn validate_configs_test(configs: Vec) -> std::io::Result> { let (converted_configs, _) = convert_cert_paths(configs).await?; let validated = create_server_configs(converted_configs)?; @@ -1695,8 +1931,11 @@ mod tests { let cert_dir = test_dir.path().join("certs"); tokio::fs::create_dir_all(&cert_dir).await.unwrap(); - let test_cert = "-----BEGIN CERTIFICATE-----\nTEST CERT CONTENT\n-----END CERTIFICATE-----"; - let test_key = "-----BEGIN PRIVATE KEY-----\nTEST KEY CONTENT\n-----END PRIVATE KEY-----"; + let certificate = rcgen::generate_simple_self_signed(vec!["localhost".into()]).unwrap(); + let cert_pem = certificate.cert.pem(); + let key_pem = certificate.signing_key.serialize_pem(); + let test_cert = cert_pem.as_str(); + let test_key = key_pem.as_str(); let cert_files = vec![ ("quic.crt", test_cert), @@ -1738,7 +1977,7 @@ mod tests { - matching_path: "/ws" protocol: type: vmess - cipher: auto + cipher: any user_id: "123e4567-e89b-42d3-a456-426614174000" shadowtls_targets: "shadow.com": diff --git a/src/dns/builder.rs b/src/dns/builder.rs index 5961db06..d18cd9f9 100644 --- a/src/dns/builder.rs +++ b/src/dns/builder.rs @@ -5,7 +5,7 @@ use std::time::Duration; use rustc_hash::FxHashMap; -use crate::config::{DnsConfig, ExpandedDnsGroup, ExpandedDnsSpec}; +use crate::config::{DnsConfig, ExpandedDnsGroup, ExpandedDnsSpec, GlobalLimits}; use crate::dns::composite_resolver::CompositeResolver; use crate::dns::hickory_resolver::{HickoryResolver, HickoryResolverOptions}; use crate::dns::parsed::{ParsedDnsServer, ParsedDnsServerEntry, ParsedDnsUrl}; @@ -14,10 +14,11 @@ use crate::resolver::{ CachingNativeResolver, NativeResolver, RefreshPolicy, RefreshingResolver, Resolver, ResolverFactory, TimeoutResolver, }; -use crate::tcp::chain_builder::{build_client_chain_group, build_direct_chain_group}; +use crate::tcp::chain_builder::{build_client_chain_group_with_limits, build_direct_chain_group}; /// Registry of resolved DNS groups with lazy default resolver. pub struct DnsRegistry { + limits: GlobalLimits, groups: FxHashMap>, /// Default resolver, created lazily only if needed. default_resolver: Option>, @@ -27,6 +28,7 @@ impl DnsRegistry { /// Creates a new empty registry. pub fn new() -> Self { Self { + limits: crate::resources::limits(), groups: FxHashMap::default(), default_resolver: None, } @@ -380,13 +382,21 @@ fn server_to_ns_config( /// Build a resolver from parsed DNS server entries. /// When all entries are hickory-backed with compatible settings, pools them /// into a single hickory resolver instead of using CompositeResolver. +#[cfg(test)] pub fn build_resolver(entries: Vec) -> std::io::Result> { + build_resolver_with_limits(entries, crate::resources::limits()) +} + +fn build_resolver_with_limits( + entries: Vec, + limits: GlobalLimits, +) -> std::io::Result> { if entries.is_empty() { return Err(std::io::Error::other("no DNS servers configured")); } // Try to pool all hickory-backed entries into one resolver. - if let Some(pooled) = try_build_hickory_pool(&entries)? { + if let Some(pooled) = try_build_hickory_pool(&entries, limits)? { return Ok(pooled); } @@ -401,6 +411,7 @@ pub fn build_resolver(entries: Vec) -> std::io::Result 0).then(|| Duration::from_secs(timeout_secs as u64)), connect_timeout: Duration::from_secs(entry.connect_timeout_secs as u64), attempts: entry.attempts, + quic_memory_bytes: limits.quic_dns_memory_bytes, }; let resolver: Arc = match entry.server { @@ -431,6 +442,7 @@ pub fn build_resolver(entries: Vec) -> std::io::Result std::io::Result>> { if entries.is_empty() || entries.len() < 2 { return Ok(None); @@ -468,6 +480,7 @@ fn try_build_hickory_pool( request_timeout: (timeout_secs > 0).then(|| Duration::from_secs(timeout_secs as u64)), connect_timeout: Duration::from_secs(first.connect_timeout_secs as u64), attempts: first.attempts, + quic_memory_bytes: limits.quic_dns_memory_bytes, }; let descriptions: Vec = entries.iter().map(|e| format!("{:?}", e.server)).collect(); @@ -491,8 +504,19 @@ fn try_build_hickory_pool( /// - Builds client chain groups from expanded client chains /// - Resolves hostnames in DNS URLs using bootstrap resolvers /// - Creates HickoryResolver instances +#[allow(dead_code)] // Public library entry point; the binary prepares with explicit limits. pub async fn build_dns_registry(groups: Vec) -> std::io::Result { - let mut registry = DnsRegistry::new(); + build_dns_registry_with_limits(groups, crate::resources::limits()).await +} + +pub async fn build_dns_registry_with_limits( + groups: Vec, + limits: GlobalLimits, +) -> std::io::Result { + let mut registry = DnsRegistry { + limits, + ..DnsRegistry::new() + }; for group in groups { let resolver = build_resolver_from_specs(&group.specs, ®istry, &group.name).await?; @@ -546,7 +570,7 @@ async fn build_resolver_from_specs( let refreshing = RefreshingResolver::new(factory, policy, description).await?; Ok(Arc::new(refreshing)) } else { - build_resolver(entries) + build_resolver_with_limits(entries, registry.limits) } } @@ -570,7 +594,11 @@ async fn build_entry_and_plan( } else { NoneOrSome::Some(spec.client_chains.clone()) }; - Arc::new(build_client_chain_group(chains, chain_resolver)) + Arc::new(build_client_chain_group_with_limits( + chains, + chain_resolver, + registry.limits, + )?) }; // Build or get bootstrap resolver @@ -604,7 +632,7 @@ async fn build_entry_and_plan( 5, // Default connect timeout for bootstrap 2, // Default attempts for bootstrap ); - build_resolver(vec![bootstrap_entry])? + build_resolver_with_limits(vec![bootstrap_entry], registry.limits)? } } None => Arc::new(NativeResolver::new()), @@ -616,6 +644,7 @@ async fn build_entry_and_plan( request_timeout: (timeout_secs > 0).then(|| Duration::from_secs(timeout_secs as u64)), connect_timeout: Duration::from_secs(spec.connect_timeout_secs as u64), attempts: spec.attempts, + quic_memory_bytes: registry.limits.quic_dns_memory_bytes, }; // Build plan for hickory-backed resolvers (not system) @@ -723,6 +752,134 @@ mod tests { } } + #[tokio::test] + async fn dns_bootstrap_uses_live_quic_admission_before_activation() { + if std::env::var_os("SHOES_DNS_ADMISSION_TEST_CHILD").is_none() { + let status = std::process::Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "dns::builder::tests::dns_bootstrap_uses_live_quic_admission_before_activation", + ]) + .env("SHOES_DNS_ADMISSION_TEST_CHILD", "1") + .status() + .unwrap(); + assert!(status.success()); + return; + } + let blackhole = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let address = blackhole.local_addr().unwrap(); + let live = GlobalLimits { + quic_memory_bytes: Some(1), + ..Default::default() + }; + let candidate = GlobalLimits { + quic_memory_bytes: Some(64 << 20), + ..live + }; + crate::resources::configure(live).unwrap(); + for server in [ + format!("url: h3://{address}/dns-query"), + format!( + "url: tcp://127.0.0.1\n client_chain:\n address: '{address}'\n transport: quic\n protocol: {{type: socks}}" + ), + ] { + let config = format!( + "- dns_group: bootstrap\n dns_servers:\n - {server}\n timeout_secs: 1\n attempts: 1\n- dns_group: candidate\n dns_servers:\n - url: https://candidate.test/dns-query\n bootstrap_url: bootstrap\n" + ); + let validated = + crate::config::create_server_configs(serde_yaml::from_str(&config).unwrap()) + .unwrap(); + let rejected = crate::resources::snapshot().quic_buffer_bytes.rejected; + assert!( + tokio::time::timeout( + Duration::from_secs(5), + build_dns_registry_with_limits(validated.dns_groups, candidate) + ) + .await + .unwrap() + .is_err() + ); + assert!(crate::resources::snapshot().quic_buffer_bytes.rejected > rejected); + assert_eq!(crate::resources::snapshot().quic_buffer_bytes.active, 0); + assert_eq!( + crate::resources::limits().quic_memory_bytes, + live.quic_memory_bytes + ); + assert!(blackhole.try_recv(&mut [0; 1500]).is_err()); + } + } + + #[tokio::test] + async fn prepared_dns_chains_use_candidate_quic_windows() { + if std::env::var_os("SHOES_DNS_CANDIDATE_TEST_CHILD").is_none() { + let status = std::process::Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "dns::builder::tests::prepared_dns_chains_use_candidate_quic_windows", + ]) + .env("SHOES_DNS_CANDIDATE_TEST_CHILD", "1") + .status() + .unwrap(); + assert!(status.success()); + return; + } + let blackhole = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let config = format!( + "- dns_group: candidate\n dns_servers:\n - url: tcp://127.0.0.1\n client_chain:\n address: '{}'\n transport: quic\n protocol: {{type: socks}}\n", + blackhole.local_addr().unwrap() + ); + crate::resources::configure(GlobalLimits { + quic_memory_bytes: Some(1 << 20), + ..Default::default() + }) + .unwrap(); + for window in [65536, 131072] { + let previous = crate::resources::limits(); + let candidate = GlobalLimits { + quic_receive_window: window, + quic_send_window: window, + quic_stream_window: 16384, + ..previous + }; + let configs = + crate::config::create_server_configs(serde_yaml::from_str(&config).unwrap()) + .unwrap(); + let registry = build_dns_registry_with_limits(configs.dns_groups, candidate) + .await + .unwrap(); + assert_eq!( + crate::resources::limits().quic_receive_window, + previous.quic_receive_window + ); + let resolver = registry.get_by_name("candidate").unwrap(); + { + let target = + crate::address::NetLocation::from_str("candidate.test:80", None).unwrap(); + let lookup = resolver.resolve_location(&target); + tokio::pin!(lookup); + let mut packet = [0; 1500]; + tokio::select! { + result = &mut lookup => panic!("lookup completed before QUIC packet: {result:?}"), + result = tokio::time::timeout(std::time::Duration::from_secs(2), blackhole.recv(&mut packet)) => { result.unwrap().unwrap(); } + } + assert_eq!( + crate::resources::snapshot().quic_buffer_bytes.active, + window * 2 + (512 << 10) + ); + } + drop((resolver, registry)); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while crate::resources::snapshot().quic_buffer_bytes.active != 0 { + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + crate::resources::configure(candidate).unwrap(); + while blackhole.try_recv(&mut [0; 1500]).is_ok() {} + } + } + #[test] fn test_compatible_servers_are_pooled() { let (chain, bootstrap) = shared_test_deps(); diff --git a/src/dns/hickory_resolver.rs b/src/dns/hickory_resolver.rs index 6cf55cb4..d5d9ac88 100644 --- a/src/dns/hickory_resolver.rs +++ b/src/dns/hickory_resolver.rs @@ -22,6 +22,7 @@ use crate::resolver::Resolver as ShoesResolver; /// Tuning options for hickory-backed resolvers. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct HickoryResolverOptions { + pub quic_memory_bytes: usize, pub ip_strategy: IpStrategy, /// Per-request timeout passed to hickory's ResolverOpts.timeout. /// None means use hickory's default. @@ -35,6 +36,7 @@ pub struct HickoryResolverOptions { impl Default for HickoryResolverOptions { fn default() -> Self { Self { + quic_memory_bytes: crate::resources::limits().quic_dns_memory_bytes, ip_strategy: IpStrategy::default(), request_timeout: Some(Duration::from_secs(5)), connect_timeout: Duration::from_secs(5), @@ -212,7 +214,8 @@ impl HickoryResolver { ) -> std::io::Result { let config = ResolverConfig::from_parts(None, vec![], ns_configs); let provider = - ProxyRuntimeProvider::with_bootstrap(chain_group, bootstrap, options.connect_timeout); + ProxyRuntimeProvider::with_bootstrap(chain_group, bootstrap, options.connect_timeout) + .with_quic_memory(options.quic_memory_bytes); let mut builder = Resolver::builder_with_config(config, provider); let resolver_opts = builder.options_mut(); @@ -332,6 +335,7 @@ mod tests { request_timeout: Some(Duration::from_secs(3)), connect_timeout: Duration::from_secs(1), attempts: 1, + ..Default::default() }; assert_eq!(opts.ip_strategy, IpStrategy::Ipv4Only); assert_eq!(opts.request_timeout, Some(Duration::from_secs(3))); diff --git a/src/dns/mod.rs b/src/dns/mod.rs index 9fdaf430..9107533d 100644 --- a/src/dns/mod.rs +++ b/src/dns/mod.rs @@ -16,5 +16,6 @@ mod hickory_resolver; mod parsed; mod proxy_runtime; -pub use builder::build_dns_registry; +#[allow(unused_imports)] +pub use builder::{DnsRegistry, build_dns_registry, build_dns_registry_with_limits}; pub use parsed::{IpStrategy, ParsedDnsUrl}; diff --git a/src/dns/proxy_runtime.rs b/src/dns/proxy_runtime.rs index 7f8f6ac7..4927e68f 100644 --- a/src/dns/proxy_runtime.rs +++ b/src/dns/proxy_runtime.rs @@ -44,6 +44,7 @@ impl ProxyRuntimeProvider { let bind_interface = chain_group.get_bind_interface().map(ToString::to_string); let quic_binder = ProxyQuicBinder { bind_interface: bind_interface.clone(), + memory_bytes: crate::resources::limits().quic_dns_memory_bytes, }; Self { chain_group, @@ -53,6 +54,11 @@ impl ProxyRuntimeProvider { connect_timeout, } } + + pub fn with_quic_memory(mut self, bytes: usize) -> Self { + self.quic_binder.memory_bytes = bytes; + self + } } /// Spawn handle for tokio runtime. @@ -160,6 +166,7 @@ impl RuntimeProvider for ProxyRuntimeProvider { #[derive(Clone)] struct ProxyQuicBinder { bind_interface: Option, + memory_bytes: usize, } impl QuicSocketBinder for ProxyQuicBinder { @@ -168,8 +175,8 @@ impl QuicSocketBinder for ProxyQuicBinder { local_addr: SocketAddr, _server_addr: SocketAddr, ) -> Result, io::Error> { - let memory = - crate::resources::try_dns_quic_memory().ok_or_else(crate::resources::exhausted)?; + let memory = crate::resources::try_quic_memory(self.memory_bytes) + .ok_or_else(|| crate::resources::quic_memory_exhausted(self.memory_bytes))?; let socket = crate::socket_util::new_outbound_socket2_udp_socket( local_addr.is_ipv6(), self.bind_interface.clone(), @@ -315,6 +322,68 @@ mod tests { true } + #[tokio::test] + async fn quic_binder_keeps_candidate_reservation_without_publishing_limits() { + if run_quic_test_in_child( + "quic_binder_keeps_candidate_reservation_without_publishing_limits", + ) { + return; + } + let resolver = Arc::new(NativeResolver::new()); + let provider = ProxyRuntimeProvider::with_bootstrap( + Arc::new(build_direct_chain_group(resolver.clone())), + resolver, + TEST_CONNECT_TIMEOUT, + ) + .with_quic_memory(1 << 20); + assert_eq!(crate::resources::limits().quic_dns_memory_bytes, 16 << 20); + let binder = provider.quic_binder().unwrap(); + let bind = || { + binder.bind_quic( + "0.0.0.0:0".parse().unwrap(), + "127.0.0.1:443".parse().unwrap(), + ) + }; + crate::resources::configure(crate::config::GlobalLimits { + quic_memory_bytes: Some(512 << 10), + ..Default::default() + }) + .unwrap(); + let error = bind().unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::ConnectionRefused); + assert!( + error + .to_string() + .contains("requested 1048576 bytes, active 0 bytes, live cap 524288") + ); + crate::resources::configure(crate::config::GlobalLimits { + quic_memory_bytes: Some(2 << 20), + ..Default::default() + }) + .unwrap(); + let socket = bind().unwrap(); + assert_eq!( + crate::resources::snapshot().quic_buffer_bytes.active, + 1 << 20 + ); + let second = bind().unwrap(); + crate::resources::configure(crate::config::GlobalLimits { + quic_memory_bytes: Some(1 << 20), + ..Default::default() + }) + .unwrap(); + assert_eq!( + crate::resources::snapshot().quic_buffer_bytes.active, + 2 << 20 + ); + assert!(bind().is_err()); + drop(second); + assert!(bind().is_err()); + drop(socket); + assert_eq!(crate::resources::snapshot().quic_buffer_bytes.active, 0); + assert!(bind().is_ok()); + } + #[tokio::test] async fn quic_binder_reserves_shared_memory_until_last_socket_drop() { if run_quic_test_in_child("quic_binder_reserves_shared_memory_until_last_socket_drop") { @@ -322,6 +391,7 @@ mod tests { } let binder = ProxyQuicBinder { bind_interface: None, + memory_bytes: crate::resources::limits().quic_dns_memory_bytes, }; let local_addr = "0.0.0.0:0".parse().unwrap(); let server_addr = "127.0.0.1:443".parse().unwrap(); diff --git a/src/ffi/android.rs b/src/ffi/android.rs index 2ac862f8..f322d107 100644 --- a/src/ffi/android.rs +++ b/src/ffi/android.rs @@ -179,6 +179,9 @@ pub extern "system" fn Java_com_shoesproxy_ShoesNative_start<'local>( config_yaml: JString<'local>, protect_callback: JObject<'local>, ) -> jlong { + if !common::is_lifecycle_thread() { + return -1; + } let _lifecycle = common::SERVICE_LIFECYCLE.lock(); if TUN_SERVICE .get() diff --git a/src/ffi/common.rs b/src/ffi/common.rs index a4deb08b..f838e730 100644 --- a/src/ffi/common.rs +++ b/src/ffi/common.rs @@ -125,7 +125,10 @@ pub fn flush_log_file() { /// Stop the TUN service and wait for shutdown. /// /// This is the common shutdown logic used by both iOS and Android. -pub fn stop_service() { +pub fn stop_service() -> bool { + if !is_lifecycle_thread() { + return false; + } let _lifecycle = SERVICE_LIFECYCLE.lock(); info!("Stopping TUN service"); @@ -157,6 +160,16 @@ pub fn stop_service() { } info!("TUN service stop completed"); + true +} + +pub fn is_lifecycle_thread() -> bool { + if tokio::runtime::Handle::try_current().is_ok() { + log::error!("Synchronous service lifecycle calls require a thread outside Tokio"); + false + } else { + true + } } /// Check if the TUN service is running. @@ -276,6 +289,19 @@ mod tests { use std::sync::mpsc; use std::time::{Duration, Instant}; + #[test] + fn lifecycle_requires_a_non_runtime_thread() { + assert!(is_lifecycle_thread()); + let runtime = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + runtime.block_on(async { + assert!(!is_lifecycle_thread()); + assert!(!stop_service()); + }); + assert!(is_lifecycle_thread()); + } + fn start_test_service( config: String, protector: Arc, diff --git a/src/ffi/ios.rs b/src/ffi/ios.rs index 48fc8825..978255b1 100644 --- a/src/ffi/ios.rs +++ b/src/ffi/ios.rs @@ -47,9 +47,9 @@ struct IosSocketProtector; impl crate::tun::SocketProtector for IosSocketProtector { fn protect(&self, fd: std::os::unix::io::RawFd) -> std::io::Result<()> { - let callback_guard = PROTECT_CALLBACK.get_or_init(|| Mutex::new(None)).lock(); + let callback = *PROTECT_CALLBACK.get_or_init(|| Mutex::new(None)).lock(); - if let Some(callback) = *callback_guard { + if let Some(callback) = callback { if callback(fd as c_int) { Ok(()) } else { @@ -64,6 +64,24 @@ impl crate::tun::SocketProtector for IosSocketProtector { } } +#[cfg(test)] +mod tests { + use super::*; + use crate::tun::SocketProtector; + + #[test] + fn protection_callback_runs_without_holding_its_slot() { + extern "C" fn callback(_: c_int) -> bool { + PROTECT_CALLBACK.get().unwrap().try_lock().is_some() + } + let slot = PROTECT_CALLBACK.get_or_init(|| Mutex::new(None)); + let previous = slot.lock().replace(callback); + let result = IosSocketProtector.protect(-1); + *slot.lock() = previous; + result.unwrap(); + } +} + /// Initialize the shoes library. /// /// # Arguments @@ -128,6 +146,9 @@ pub unsafe extern "C" fn shoes_start( config_yaml: *const c_char, protect_callback: ProtectSocketCallback, ) -> c_long { + if !common::is_lifecycle_thread() { + return -1; + } let _lifecycle = common::SERVICE_LIFECYCLE.lock(); if TUN_SERVICE .get() @@ -195,7 +216,9 @@ pub unsafe extern "C" fn shoes_start( /// * `handle` - Handle returned by shoes_start (currently unused, we use global state) #[unsafe(no_mangle)] pub extern "C" fn shoes_stop(_handle: c_long) { - common::stop_service(); + if !common::stop_service() { + return; + } if let Some(callback) = PROTECT_CALLBACK.get() { let mut guard = callback.lock(); diff --git a/src/http_handler.rs b/src/http_handler.rs index b326e9fd..af0d2018 100644 --- a/src/http_handler.rs +++ b/src/http_handler.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use async_trait::async_trait; use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64}; use log::debug; +use subtle::ConstantTimeEq; use tokio::io::AsyncWriteExt; use crate::address::{Address, NetLocation, ResolvedLocation}; @@ -18,6 +19,51 @@ const PROXY_AUTH_HEADER_PREFIX: &str = "proxy-authorization: basic "; const CONNECTION_HEADER_PREFIX: &str = "connection: "; const PROXY_CONNECTION_HEADER_PREFIX: &str = "proxy-connection: "; +#[cfg(test)] +mod request_tests { + use super::*; + + #[tokio::test] + async fn rejected_credentials_are_not_logged() { + let capture = crate::logging::test_capture::Capture::new(); + for request in ["CONNECT example.com:443", "GET http://example.com/"] { + let (mut client, server) = tokio::io::duplex(1024); + client.write_all(format!("{request} HTTP/1.1\r\nProxy-Authorization: Basic rejected-secret-token\r\n\r\n").as_bytes()).await.unwrap(); + let handler = HttpTcpServerHandler::new( + Some(("user".into(), "password".into())), + Arc::new(ClientProxySelector::new(vec![])), + ); + let error = handler + .setup_server_stream(Box::new(server)) + .await + .err() + .unwrap(); + assert!(!error.to_string().contains("rejected-secret-token")); + } + let logs = capture.finish(); + assert!(logs.contains("authentication")); + assert!(!logs.contains("rejected-secret-token")); + } + + #[tokio::test] + async fn non_ascii_auth_header_does_not_panic_at_prefix_boundary() { + let (mut client, server) = tokio::io::duplex(1024); + let header = format!( + "{}\u{e9}invalid", + "x".repeat(PROXY_AUTH_HEADER_PREFIX.len() - 1) + ); + client + .write_all(format!("CONNECT example.com:443 HTTP/1.1\r\n{header}\r\n\r\n").as_bytes()) + .await + .unwrap(); + let handler = HttpTcpServerHandler::new( + Some(("user".into(), "password".into())), + Arc::new(ClientProxySelector::new(vec![])), + ); + assert!(handler.setup_server_stream(Box::new(server)).await.is_err()); + } +} + fn create_http_auth_token(username: &str, password: &str) -> String { BASE64.encode(format!("{username}:{password}")) } @@ -106,14 +152,14 @@ pub async fn setup_http_server_stream_inner( } if need_auth && line.len() > PROXY_AUTH_HEADER_PREFIX.len() + 1 - && line[0..PROXY_AUTH_HEADER_PREFIX.len()].to_ascii_lowercase() - == PROXY_AUTH_HEADER_PREFIX + && line.as_bytes()[0..PROXY_AUTH_HEADER_PREFIX.len()] + .eq_ignore_ascii_case(PROXY_AUTH_HEADER_PREFIX.as_bytes()) { - if &line[PROXY_AUTH_HEADER_PREFIX.len()..] != auth_token.unwrap() { - debug!( - "Received incorrect HTTP CONNECT authentication: {}", - &line[PROXY_AUTH_HEADER_PREFIX.len()..] - ); + if !bool::from( + line.as_bytes()[PROXY_AUTH_HEADER_PREFIX.len()..] + .ct_eq(auth_token.unwrap().as_bytes()), + ) { + debug!("Received incorrect HTTP CONNECT authentication"); return Err(std::io::Error::new( std::io::ErrorKind::InvalidInput, "Incorrect HTTP CONNECT authentication", @@ -122,7 +168,7 @@ pub async fn setup_http_server_stream_inner( need_auth = false; continue; } - debug!("Ignored HTTP CONNECT request header: {line}"); + debug!("Ignored HTTP CONNECT request header"); } if need_auth { @@ -203,11 +249,11 @@ pub async fn setup_http_server_stream_inner( && lowercase_line.starts_with(PROXY_AUTH_HEADER_PREFIX) { if need_auth { - if &line[PROXY_AUTH_HEADER_PREFIX.len()..] != auth_token.unwrap() { - debug!( - "Received incorrect HTTP GET authentication: {}", - &line[PROXY_AUTH_HEADER_PREFIX.len()..] - ); + if !bool::from( + line.as_bytes()[PROXY_AUTH_HEADER_PREFIX.len()..] + .ct_eq(auth_token.unwrap().as_bytes()), + ) { + debug!("Received incorrect HTTP GET authentication"); return Err(std::io::Error::new( std::io::ErrorKind::InvalidInput, "Incorrect HTTP GET authentication", diff --git a/src/hysteria2_server.rs b/src/hysteria2_server.rs index d2d7f9a1..009e6f7d 100644 --- a/src/hysteria2_server.rs +++ b/src/hysteria2_server.rs @@ -9,6 +9,7 @@ use log::{debug, error, warn}; use rand::distr::Alphanumeric; use rand::{Rng, RngExt}; use rustc_hash::FxHashMap; +use subtle::ConstantTimeEq; use tokio::io::AsyncWriteExt; use tokio::task::{JoinHandle, JoinSet}; use tokio::time::timeout; @@ -22,7 +23,7 @@ const AUTH_TIMEOUT: Duration = Duration::from_secs(3); /// Per official hysteria reference: https://github.com/apernet/hysteria/blob/master/core/server/server.go#L20 const CLOSE_ERR_CODE_OK: u32 = 0x100; // HTTP3 ErrCodeNoError -use crate::address::NetLocation; +use crate::address::{Address, NetLocation}; use crate::async_stream::AsyncStream; use crate::client_proxy_selector::{ClientProxySelector, ConnectDecision}; use crate::copy_bidirectional::copy_bidirectional_with_sizes; @@ -30,9 +31,10 @@ use crate::quic_stream::QuicStream; use crate::resolver::Resolver; use crate::routing::udp_relay::UdpRelay; use crate::stream_reader::StreamReader; +use crate::tcp::tcp_handler::TcpClientSetupResult; use crate::tcp::tcp_server::setup_client_tcp_stream; use crate::udp_fragments::UdpFragments; -use crate::util::allocate_vec; +use crate::util::{timeout_stream_setup, write_all}; async fn process_connection( client_proxy_selector: Arc, @@ -168,7 +170,7 @@ fn validate_auth_request(req: http::Request, password: &str) -> std::io::R let auth_str = auth_value .to_str() .map_err(|e| std::io::Error::other(format!("invalid auth header value: {e}")))?; - if auth_str != password { + if !bool::from(auth_str.as_bytes().ct_eq(password.as_bytes())) { return Err(std::io::Error::other("incorrect auth password")); } @@ -204,7 +206,7 @@ async fn auth_connection( let resp = http::Response::builder() .status(http::status::StatusCode::from_u16(233).unwrap()) .header("Hysteria-UDP", if udp_enabled { "true" } else { "false" }) - .header("Hysteria-CC-RX", "0") + .header("Hysteria-CC-RX", "auto") .header("Hysteria-Padding", generate_ascii_string()) .body(()) .unwrap(); @@ -249,21 +251,17 @@ struct UdpSession { send_socket: Arc, pinned_location: Option, cancel_token: CancellationToken, - task: Option, + task: tokio::task::AbortHandle, } impl Drop for UdpSession { fn drop(&mut self) { self.cancel_token.cancel(); - if let Some(task) = &self.task { - task.abort(); - } + self.task.abort(); } } impl UdpSession { - // TODO: remove this function completely and inline? - #[allow(clippy::too_many_arguments)] fn start( session_id: u32, connection: quinn::Connection, @@ -275,31 +273,29 @@ impl UdpSession { // Create a child token so this session is cancelled when the parent (connection) is cancelled let session_cancel_token = parent_cancel_token.child_token(); - let mut session = UdpSession { - send_socket: client_socket.clone(), + let send_socket = client_socket.clone(); + let cancel_token = session_cancel_token.clone(); + let task = tokio::spawn(async move { + if let Err(e) = run_udp_remote_to_local_loop( + session_id, + connection, + client_socket, + override_local_write_location, + session_cancel_token, + ) + .await + { + error!("UDP remote-to-local write loop ended with error: {e}"); + } + }) + .abort_handle(); + + Self { + send_socket, pinned_location, - cancel_token: session_cancel_token.clone(), - task: None, - }; - - session.task = Some( - tokio::spawn(async move { - if let Err(e) = run_udp_remote_to_local_loop( - session_id, - connection, - client_socket, - override_local_write_location, - session_cancel_token, - ) - .await - { - error!("UDP remote-to-local write loop ended with error: {e}"); - } - }) - .abort_handle(), - ); - - session + cancel_token, + task, + } } } @@ -316,7 +312,7 @@ async fn run_udp_remote_to_local_loop( let original_address_bytes: Option<(Bytes, Bytes)> = match override_local_write_address { Some(a) => { - let address_bytes: Bytes = a.to_string().into_bytes().into(); + let address_bytes: Bytes = wire_address(&a).into_bytes().into(); let address_len = address_bytes.len(); let address_len_bytes = encode_varint(address_len as u64)?; Some((address_bytes, address_len_bytes.into())) @@ -436,7 +432,8 @@ fn parse_udp_packet(data: &[u8]) -> std::io::Result> { } let end = 8 + length + address_len as usize; let address = data.get(8 + length..end).ok_or_else(invalid)?; - let location = NetLocation::from_str(str::from_utf8(address).map_err(|_| invalid())?, None)?; + let location = + NetLocation::from_authority(str::from_utf8(address).map_err(|_| invalid())?, None)?; Ok(UdpPacket { session_id: u32::from_be_bytes(data[..4].try_into().unwrap()), packet_id: u16::from_be_bytes(data[4..6].try_into().unwrap()), @@ -608,8 +605,7 @@ async fn run_tcp_loop( const FRAME_TYPE_TCP_REQUEST: u64 = 0x401; async fn handle_tcp_header( - send: &mut quinn::SendStream, - recv: &mut quinn::RecvStream, + recv: &mut (impl tokio::io::AsyncRead + Unpin), ) -> std::io::Result<(NetLocation, StreamReader)> { let mut stream_reader = StreamReader::new_with_buffer_size(8192); @@ -631,7 +627,7 @@ async fn handle_tcp_header( let address_bytes = stream_reader.read_slice(recv, address_len as usize).await?; let address = std::str::from_utf8(address_bytes) .map_err(|e| std::io::Error::other(format!("invalid address encoding: {e}")))?; - let remote_location = NetLocation::from_str(address, None)?; + let remote_location = NetLocation::from_authority(address, None)?; let padding_len = read_varint(recv, &mut stream_reader).await?; if padding_len > 4096 { @@ -639,105 +635,110 @@ async fn handle_tcp_header( } stream_reader.read_slice(recv, padding_len as usize).await?; - let response_bytes = { - // [uint8] Status (0x00 = OK, 0x01 = Error) - // [varint] Message length - // [bytes] Message string - // [varint] Padding length - // [bytes] Random padding - - let mut rng = rand::rng(); - - // only use the lower 6 bits so that the varint always fits in a single u8 - let padding_len = rng.random_range(0..=63); - - // first 3 bytes of status = 0x0, message length = 0, padding length - let mut response_bytes = allocate_vec(3 + (padding_len as usize)); - response_bytes[0] = 0; - response_bytes[1] = 0; - response_bytes[2] = padding_len; - rng.fill_bytes(&mut response_bytes[3..]); - - response_bytes - }; - - let len = response_bytes.len(); - let mut i = 0; - while i < len { - let count = send - .write(&response_bytes[i..len]) - .await - .map_err(|e| std::io::Error::other(format!("H3 stream write failed: {e}")))?; - i += count; - } - Ok((remote_location, stream_reader)) } +fn wire_address(location: &NetLocation) -> String { + match location.address() { + Address::Ipv6(ip) => format!("[{ip}]:{}", location.port()), + _ => location.to_string(), + } +} + +fn tcp_response(success: bool, message: &str) -> Vec { + assert!(message.len() <= 2048); + let mut response = vec![u8::from(!success)]; + response.extend_from_slice(&encode_varint(message.len() as u64).unwrap()); + response.extend_from_slice(message.as_bytes()); + let padding_len = rand::rng().random_range(0..=63); + response.push(padding_len); + let start = response.len(); + response.resize(start + padding_len as usize, 0); + rand::rng().fill_bytes(&mut response[start..]); + response +} + +async fn complete_tcp_setup( + server: &mut Box, + setup: impl std::future::Future>>, +) -> std::io::Result>> { + match timeout_stream_setup(setup).await { + Ok(Some(TcpClientSetupResult { + client_stream, + early_data, + })) => { + timeout_stream_setup(async { + write_all(server, &tcp_response(true, "")).await?; + if let Some(data) = early_data { + write_all(server, &data).await?; + } + server.flush().await + }) + .await?; + Ok(Some(client_stream)) + } + result => { + let message = match &result { + Ok(None) => "connection blocked", + Err(error) if error.kind() == std::io::ErrorKind::PermissionDenied => { + "connection blocked" + } + Err(error) if error.kind() == std::io::ErrorKind::TimedOut => { + "connection timed out" + } + _ => "connection failed", + }; + let _ = timeout(crate::util::SHUTDOWN_TIMEOUT, async { + write_all(server, &tcp_response(false, message)).await?; + server.flush().await + }) + .await; + crate::util::shutdown_stream(server).await; + result.map(|_| None) + } + } +} + async fn process_tcp_stream( client_proxy_selector: Arc, resolver: Arc, mut send: quinn::SendStream, mut recv: quinn::RecvStream, ) -> std::io::Result<()> { - let (remote_location, stream_reader) = match handle_tcp_header(&mut send, &mut recv).await { - Ok(res) => res, - Err(e) => { - let _ = send.shutdown().await; - return Err(e); - } - }; + let (remote_location, stream_reader) = + match timeout_stream_setup(handle_tcp_header(&mut recv)).await { + Ok(res) => res, + Err(e) => { + crate::util::shutdown_stream(&mut send).await; + return Err(e); + } + }; let mut server_stream: Box = Box::new(QuicStream::from(send, recv)); - let setup_client_stream_future = timeout( - Duration::from_secs(60), - setup_client_tcp_stream( - &mut server_stream, - client_proxy_selector, - resolver, - remote_location.clone(), - ), - ); - - let mut client_stream = match setup_client_stream_future.await { - Ok(Ok(Some(s))) => s, - Ok(Ok(None)) => { - // Must have been blocked. - crate::util::shutdown_stream(&mut server_stream).await; - return Ok(()); - } - Ok(Err(e)) => { - crate::util::shutdown_stream(&mut server_stream).await; - return Err(std::io::Error::new( - e.kind(), - format!("failed to setup client stream to {remote_location}: {e}"), - )); - } - Err(elapsed) => { - crate::util::shutdown_stream(&mut server_stream).await; - return Err(std::io::Error::new( - std::io::ErrorKind::TimedOut, - format!("client setup to {remote_location} timed out: {elapsed}"), - )); - } + let Some(mut client_stream) = complete_tcp_setup( + &mut server_stream, + setup_client_tcp_stream(client_proxy_selector, resolver, remote_location.clone()), + ) + .await + .map_err(|error| { + std::io::Error::new( + error.kind(), + format!("TCP setup for {remote_location} failed: {error}"), + ) + })? + else { + return Ok(()); }; let unparsed_data = stream_reader.unparsed_data(); - let client_requires_flush = if unparsed_data.is_empty() { - false - } else { - let len = unparsed_data.len(); - let mut i = 0; - while i < len { - let count = client_stream - .write(&unparsed_data[i..len]) - .await - .map_err(|e| std::io::Error::other(format!("H3 stream write failed: {e}")))?; - i += count; - } - true - }; + if !unparsed_data.is_empty() { + timeout_stream_setup(async { + write_all(&mut client_stream, unparsed_data).await?; + client_stream.flush().await + }) + .await?; + } drop(stream_reader); // Use 32KB buffers to match hysteria2/sing-box reference implementations @@ -746,7 +747,7 @@ async fn process_tcp_stream( &mut client_stream, // no need to flush even through we wrote this response since it's quic false, - client_requires_flush, + false, 32768, 32768, ) @@ -783,7 +784,7 @@ fn encode_varint(value: u64) -> std::io::Result> { } async fn read_varint( - recv: &mut quinn::RecvStream, + recv: &mut (impl tokio::io::AsyncRead + Unpin), stream_reader: &mut StreamReader, ) -> std::io::Result { let first_byte = stream_reader.read_u8(recv).await?; @@ -926,3 +927,109 @@ mod datagram_tests { } } } + +#[cfg(test)] +mod tcp_tests { + use super::*; + use tokio::io::AsyncReadExt; + + #[tokio::test] + async fn success_follows_dial_and_precedes_target_bytes() { + let (server, mut peer) = tokio::io::duplex(512); + let (target, _target_peer) = tokio::io::duplex(512); + let (ready, dial) = tokio::sync::oneshot::channel(); + let task = tokio::spawn(async move { + let mut server: Box = Box::new(server); + let _target = complete_tcp_setup(&mut server, async { + dial.await.unwrap(); + Ok(Some(TcpClientSetupResult { + client_stream: Box::new(target), + early_data: Some(b"greeting".to_vec()), + })) + }) + .await + .unwrap(); + server.write_all(b"later").await.unwrap(); + }); + assert!( + timeout(Duration::from_millis(10), peer.read_u8()) + .await + .is_err() + ); + ready.send(()).unwrap(); + let mut wire = Vec::new(); + peer.read_to_end(&mut wire).await.unwrap(); + task.await.unwrap(); + assert_eq!(&wire[..2], &[0, 0]); + assert_eq!(&wire[3 + wire[2] as usize..], b"greetinglater"); + } + + #[tokio::test(start_paused = true)] + async fn failed_dials_report_failure_without_exposing_internal_errors() { + for (kind, blocked) in [ + (std::io::ErrorKind::PermissionDenied, true), + (std::io::ErrorKind::PermissionDenied, false), + (std::io::ErrorKind::ConnectionRefused, false), + (std::io::ErrorKind::TimedOut, false), + ] { + let (server, mut peer) = tokio::io::duplex(512); + let mut server: Box = Box::new(server); + let result = complete_tcp_setup(&mut server, async { + if blocked { + Ok(None) + } else if kind == std::io::ErrorKind::TimedOut { + std::future::pending().await + } else { + Err(std::io::Error::new(kind, "private target details")) + } + }) + .await; + if blocked { + assert!(result.unwrap().is_none()); + } else { + assert_eq!(result.err().unwrap().kind(), kind); + } + let mut wire = Vec::new(); + peer.read_to_end(&mut wire).await.unwrap(); + assert_eq!(wire[0], 1); + let message_len = wire[1] as usize; + assert!(message_len > 0); + assert!(!String::from_utf8_lossy(&wire).contains("private target details")); + assert_eq!(wire.len(), 3 + message_len + wire[2 + message_len] as usize); + } + } + + #[tokio::test(start_paused = true)] + async fn stalled_success_response_has_a_deadline() { + let (server, _peer) = tokio::io::duplex(1); + let (target, _target_peer) = tokio::io::duplex(1); + let mut server: Box = Box::new(server); + let result = complete_tcp_setup(&mut server, async { + Ok(Some(TcpClientSetupResult { + client_stream: Box::new(target), + early_data: None, + })) + }) + .await; + assert_eq!(result.err().unwrap().kind(), std::io::ErrorKind::TimedOut); + } + + #[tokio::test] + async fn ipv6_request_preserves_coalesced_payload() { + let location = NetLocation::new(Address::Ipv6(std::net::Ipv6Addr::LOCALHOST), 443); + let address = wire_address(&location); + assert_eq!(address, "[::1]:443"); + let mut request = encode_varint(FRAME_TYPE_TCP_REQUEST).unwrap().to_vec(); + request.extend_from_slice(&encode_varint(address.len() as u64).unwrap()); + request.extend_from_slice(address.as_bytes()); + request.extend_from_slice(&[0]); + request.extend_from_slice(b"request payload"); + let (parsed, reader) = handle_tcp_header(&mut request.as_slice()).await.unwrap(); + assert_eq!(parsed, location); + assert_eq!(reader.unparsed_data(), b"request payload"); + let mut packet = vec![0, 0, 0, 1, 0, 0, 0, 1, address.len() as u8]; + packet.extend_from_slice(address.as_bytes()); + packet.extend_from_slice(b"udp"); + assert_eq!(parse_udp_packet(&packet).unwrap().location, location); + } +} diff --git a/src/listener_tasks.rs b/src/listener_tasks.rs index 07282f54..997cc980 100644 --- a/src/listener_tasks.rs +++ b/src/listener_tasks.rs @@ -1,8 +1,43 @@ use std::future::Future; +use std::io; +use std::sync::Arc; use std::time::Duration; +use parking_lot::Mutex; +use tokio::sync::oneshot; use tokio::task::JoinSet; +tokio::task_local! { + static QUIC_RETIREMENTS: QuicRetirements; +} + +/// Collects listener socket releases without including outbound QUIC used by draining TCP tasks. +#[derive(Clone, Default)] +pub(crate) struct QuicRetirements(Arc>>>); + +impl QuicRetirements { + pub async fn track(&self, startup: F) -> F::Output { + QUIC_RETIREMENTS.scope(self.clone(), startup).await + } + + pub async fn wait(&self) -> io::Result<()> { + let releases = std::mem::take(&mut *self.0.lock()); + tokio::time::timeout(Duration::from_secs(5), async { + for released in releases { + // Sender destruction, rather than a sent value, confirms socket release. + let _ = released.await; + } + }) + .await + .map_err(|_| { + io::Error::new( + io::ErrorKind::TimedOut, + "QUIC listener sockets did not retire", + ) + }) + } +} + pub(crate) struct QuicListener(pub crate::quic_endpoint::QuicEndpoint); impl QuicListener { @@ -12,21 +47,29 @@ impl QuicListener { count: usize, memory_bytes: usize, ) -> std::io::Result> { + let reuse_port = cfg!(all( + unix, + not(any(target_os = "solaris", target_os = "illumos")) + )); + let count = if reuse_port { count } else { count.min(1) }; (0..count) .map(|_| { let socket = crate::socket_util::new_socket2_udp_socket_with_buffer_size( address.is_ipv6(), None, Some(address), - true, + reuse_port, Some(crate::resources::limits().quic_socket_buffer), )?; - crate::quic_endpoint::QuicEndpoint::new( - Some(config.clone()), + let (endpoint, released) = crate::quic_endpoint::QuicEndpoint::listen( + config.clone(), socket.into(), memory_bytes, - ) - .map(Self) + )?; + let _ = QUIC_RETIREMENTS.try_with(|retirements| { + retirements.0.lock().push(released); + }); + Ok(Self(endpoint)) }) .collect() } @@ -119,8 +162,61 @@ impl Drop for ListenerTasks { #[cfg(test)] mod tests { use super::*; - use std::sync::Arc; - use tokio::sync::oneshot; + + #[tokio::test(start_paused = true)] + async fn retirement_timeout_is_an_error() { + let retirements = QuicRetirements::default(); + let (_socket, released) = oneshot::channel(); + retirements.0.lock().push(released); + assert_eq!( + retirements.wait().await.unwrap_err().kind(), + io::ErrorKind::TimedOut + ); + } + + #[tokio::test] + async fn retirement_does_not_wait_for_a_slow_peer_close_timer() { + let cert = rcgen::generate_simple_self_signed(vec!["localhost".into()]).unwrap(); + let mut roots = rustls::RootCertStore::empty(); + roots.add(cert.cert.der().clone()).unwrap(); + let mut config = quinn::ServerConfig::with_single_cert( + vec![cert.cert.der().clone()], + rustls::pki_types::PrivatePkcs8KeyDer::from(cert.signing_key.serialize_der()).into(), + ) + .unwrap(); + Arc::get_mut(&mut config.transport) + .unwrap() + .initial_rtt(Duration::from_secs(2)); + let retirements = QuicRetirements::default(); + let listener = retirements + .track(async { + QuicListener::bind_all("127.0.0.1:0".parse().unwrap(), config, 1, 1) + .unwrap() + .pop() + .unwrap() + }) + .await; + let address = listener.local_addr().unwrap(); + let mut client = quinn::Endpoint::client("127.0.0.1:0".parse().unwrap()).unwrap(); + client.set_default_client_config( + quinn::ClientConfig::with_root_certificates(Arc::new(roots)).unwrap(), + ); + let connecting = client.connect(address, "localhost").unwrap(); + let incoming = tokio::time::timeout(Duration::from_secs(2), listener.accept()) + .await + .unwrap() + .unwrap(); + client.close(0u32.into(), b"abandon handshake"); + drop((connecting, client)); + let pending = incoming.accept().unwrap(); + drop(pending); + drop(listener); + tokio::time::timeout(Duration::from_secs(3), retirements.wait()) + .await + .unwrap() + .unwrap(); + drop(std::net::UdpSocket::bind(address).unwrap()); + } #[tokio::test(start_paused = true)] async fn reload_preserves_active_work_then_cancels_stalled_connections() { @@ -174,6 +270,18 @@ mod tests { .unwrap(); let memory_bytes = crate::resources::configure_quic(&mut server_config, 100, 0); let client_config = quinn::ClientConfig::with_root_certificates(Arc::new(roots)).unwrap(); + #[cfg(windows)] + assert_eq!( + QuicListener::bind_all( + "0.0.0.0:0".parse().unwrap(), + server_config.clone(), + 4, + memory_bytes, + ) + .unwrap() + .len(), + 1, + ); let occupied = std::net::UdpSocket::bind("0.0.0.0:0").unwrap(); assert!( QuicListener::bind_all( @@ -190,7 +298,10 @@ mod tests { .pop() .unwrap() }; - let listener = create_listener("0.0.0.0:0".parse().unwrap()); + let retirements = QuicRetirements::default(); + let listener = retirements + .track(async { create_listener("0.0.0.0:0".parse().unwrap()) }) + .await; let bind_address = listener.local_addr().unwrap(); let target = std::net::SocketAddr::from(([127, 0, 0, 1], bind_address.port())); let create_client = || { @@ -232,8 +343,9 @@ mod tests { )); drop(server_conn); drop(client_conn); - // Match the config watcher's debounce before rebinding the listener. - tokio::time::sleep(Duration::from_secs(3)).await; + retirements.wait().await.unwrap(); + // A non-reuse socket proves retirement released the kernel port, not just Quinn state. + drop(std::net::UdpSocket::bind(bind_address).unwrap()); let listener = create_listener(bind_address); for _ in 0..12 { let client = create_client(); diff --git a/src/logging.rs b/src/logging.rs index 3dba8d85..270a1860 100644 --- a/src/logging.rs +++ b/src/logging.rs @@ -429,3 +429,53 @@ mod tests { assert!(!logger.matches(Level::Error, "shoes")); } } +#[cfg(test)] +pub(crate) mod test_capture { + use std::cell::RefCell; + use std::sync::Once; + + thread_local! { + static RECORDS: RefCell> = const { RefCell::new(None) }; + } + + struct Logger; + + impl log::Log for Logger { + fn enabled(&self, _: &log::Metadata<'_>) -> bool { + true + } + fn log(&self, record: &log::Record<'_>) { + use std::fmt::Write; + RECORDS.with(|records| { + if let Some(output) = records.borrow_mut().as_mut() { + let _ = writeln!(output, "{}", record.args()); + } + }); + } + fn flush(&self) {} + } + + pub struct Capture; + + impl Capture { + pub fn new() -> Self { + static INIT: Once = Once::new(); + INIT.call_once(|| { + log::set_logger(&Logger).unwrap(); + log::set_max_level(log::LevelFilter::Trace); + }); + RECORDS.with(|records| assert!(records.replace(Some(String::new())).is_none())); + Self + } + + pub fn finish(self) -> String { + RECORDS.with(|records| records.take().unwrap()) + } + } + + impl Drop for Capture { + fn drop(&mut self) { + RECORDS.with(|records| records.take()); + } + } +} diff --git a/src/main.rs b/src/main.rs index e6879461..5aa98f74 100644 --- a/src/main.rs +++ b/src/main.rs @@ -19,6 +19,7 @@ mod naiveproxy; mod option_util; mod port_forward_handler; mod prepend_stream; +mod process; mod quic_endpoint; mod quic_server; mod quic_stream; @@ -63,45 +64,14 @@ use tikv_jemallocator::Jemalloc; #[global_allocator] static GLOBAL: Jemalloc = Jemalloc; -use std::path::Path; - use aws_lc_rs::rand::{SecureRandom, SystemRandom}; use base64::engine::{Engine as _, general_purpose::STANDARD}; use log::debug; -use notify::{Event, EventKind, RecommendedWatcher, RecursiveMode, Watcher}; -use tcp_server::start_servers; use tokio::runtime::Builder; -use tokio::sync::mpsc::{Receiver, channel}; use crate::reality::generate_keypair; use crate::shadowsocks::ShadowsocksCipher; use crate::thread_util::set_num_threads; -use tcp::*; - -#[derive(Debug)] -struct ConfigChanged; - -fn start_notify_thread(config_paths: Vec) -> (RecommendedWatcher, Receiver) { - let (tx, rx) = channel(1); - - let mut watcher = notify::recommended_watcher(move |res: notify::Result| match res { - Ok(event) => { - if matches!(event.kind, EventKind::Modify(..)) { - let _ = tx.try_send(ConfigChanged {}); - } - } - Err(e) => println!("watch error: {e:?}"), - }) - .unwrap(); - - for config_path in config_paths { - watcher - .watch(Path::new(&config_path), RecursiveMode::NonRecursive) - .unwrap(); - } - - (watcher, rx) -} fn print_usage_and_exit(arg0: String) { eprintln!("{arg0} [OPTIONS] [config.yaml...]"); @@ -126,6 +96,10 @@ fn print_usage_and_exit(arg0: String) { ); eprintln!(); eprintln!("COMMANDS:"); + eprintln!(" check Validate configuration and exit"); + eprintln!(" version Print version information"); + eprintln!(" SIGHUP reloads configuration even with --no-reload."); + eprintln!(" SIGINT/SIGTERM stop listeners and exit without waiting for TCP reload drains."); eprintln!( " generate-reality-keypair Generate a new Reality X25519 keypair" ); @@ -139,6 +113,11 @@ fn print_usage_and_exit(arg0: String) { fn main() { let mut args: Vec = std::env::args().collect(); let arg0 = args.remove(0); + if args.first().is_some_and(|arg| arg == "check") { + args[0] = "--dry-run".into(); + } else if args.first().is_some_and(|arg| arg == "version") { + args[0] = "--version".into(); + } let mut num_threads = 0usize; let mut dry_run = false; let mut no_reload = false; @@ -326,113 +305,14 @@ fn main() { .build() .expect("Could not build tokio runtime"); - runtime.block_on(async move { - let mut reload_state = if no_reload { - None - } else { - let (watcher, rx) = start_notify_thread(args.clone()); - Some((watcher, rx)) - }; - - loop { - let configs = match config::load_configs(&args).await { - Ok(c) => c, - Err(e) => { - eprintln!("Failed to load server configs: {e}\n"); - print_usage_and_exit(arg0); - return; - } - }; - - let (configs, load_file_count) = match config::convert_cert_paths(configs).await { - Ok(c) => c, - Err(e) => { - eprintln!("Failed to load cert files: {e}\n"); - print_usage_and_exit(arg0); - return; - } - }; - - if load_file_count > 0 { - println!("Loaded {load_file_count} certs/keys from files"); - } - - if dry_run { - if let Err(e) = config::create_server_configs(configs) { - eprintln!("Dry run failed, could not create server configs: {e}\n"); - } else { - println!("Finishing dry run, config parsed successfully."); - } - return; - } - - let mut join_handles = vec![]; - - let server_configs = match config::create_server_configs(configs) { - Ok(c) => c, - Err(e) => { - eprintln!("Failed to create server configs: {e}\n"); - print_usage_and_exit(arg0); - return; - } - }; - - let config::ValidatedConfigs { - configs: server_configs, - dns_groups, - global_limits, - } = server_configs; - resources::configure(global_limits).expect("validated global limits"); - let _resource_reporter = resources::ResourceReporter::start(); - - // Build DNS registry from expanded groups (async - resolves hostnames) - let mut dns_registry = match dns::build_dns_registry(dns_groups).await { - Ok(r) => r, - Err(e) => { - eprintln!("Failed to build DNS registry: {e}\n"); - print_usage_and_exit(arg0); - return; - } - }; - - println!("\nStarting {} server(s)..", server_configs.len()); - - for server_config in server_configs { - // Get the resolver for this server from the registry - let dns_ref = match &server_config { - config::Config::Server(s) => s.dns.as_ref(), - config::Config::TunServer(t) => t.dns.as_ref(), - _ => None, - }; - let resolver = dns_registry.get_for_server(dns_ref); - join_handles.extend(start_servers(server_config, resolver).await.unwrap()); - } - - match reload_state.as_mut() { - Some((_watcher, rx)) => { - rx.recv().await.unwrap(); - - println!("Configs changed, restarting servers in 3 seconds.."); - - for join_handle in &join_handles { - join_handle.abort(); - } - for join_handle in join_handles { - let _ = join_handle.await; - } - - tokio::time::sleep(std::time::Duration::from_secs(3)).await; - - // Remove any extra events - while rx.try_recv().is_ok() {} - } - None => { - // No reload mode - wait forever - // TODO: signal handling? - futures::future::pending::<()>().await; - unreachable!(); - } - } + let code = match runtime.block_on(process::run(args, dry_run, no_reload)) { + Ok(()) => 0, + Err(error) => { + eprintln!("shoes: {error}"); + 1 } - }); + }; + runtime.shutdown_timeout(std::time::Duration::from_secs(5)); + log::logger().flush(); + std::process::exit(code); } diff --git a/src/process.rs b/src/process.rs new file mode 100644 index 00000000..8566260f --- /dev/null +++ b/src/process.rs @@ -0,0 +1,530 @@ +use std::collections::HashSet; +use std::io; +use std::path::{Component, Path, PathBuf}; +use std::time::Duration; + +use notify::{EventKind, RecommendedWatcher, RecursiveMode, Watcher}; +use tokio::sync::mpsc; +use tokio::task::JoinHandle; + +use crate::address::NetLocation; +use crate::config::{self, BindLocation, Config, ValidatedConfigs}; +use crate::dns::{self, DnsRegistry}; +use crate::option_util::OneOrSome; +use crate::resources; + +async fn load_validated(paths: &[String]) -> io::Result { + let configs = config::load_configs(&paths.to_vec()).await?; + let (configs, _) = config::convert_cert_paths(configs).await?; + config::create_server_configs(configs) +} + +struct PreparedServers { + validated: ValidatedConfigs, + dns: DnsRegistry, +} + +async fn prepare(paths: &[String]) -> io::Result { + let mut validated = load_validated(paths).await?; + for config in &mut validated.configs { + let Config::Server(server) = config else { + continue; + }; + let BindLocation::Address(addresses) = &server.bind_location else { + continue; + }; + let addresses = addresses.clone(); + let resolved = tokio::task::spawn_blocking(move || { + let mut resolved = Vec::new(); + for address in addresses.iter() { + let socket_addrs = address.to_socket_addrs()?; + // NetLocation has no scope ID; retain scoped literals in their original form. + if socket_addrs.iter().any( + |addr| matches!(addr, std::net::SocketAddr::V6(addr) if addr.scope_id() != 0), + ) { + resolved.push(address.clone()); + continue; + } + for socket_addr in socket_addrs { + resolved.push( + NetLocation::from_ip_addr(socket_addr.ip(), socket_addr.port()).into(), + ); + } + } + Ok::<_, io::Error>(resolved) + }) + .await + .map_err(io::Error::other)??; + if resolved.is_empty() { + return Err(io::Error::new( + io::ErrorKind::AddrNotAvailable, + "listener address resolved to no addresses", + )); + } + server.bind_location = BindLocation::Address(OneOrSome::Some(resolved)); + } + let dns = dns::build_dns_registry_with_limits( + std::mem::take(&mut validated.dns_groups), + validated.global_limits, + ) + .await + .map_err(|error| { + let live = resources::limits().quic_memory_bytes; + let candidate = validated.global_limits.quic_memory_bytes; + if live == candidate { + error + } else { + io::Error::new(error.kind(), format!( + "{error}; DNS preparation uses the live QUIC cap {live:?}, not the candidate cap {candidate:?}" + )) + } + })?; + Ok(PreparedServers { validated, dns }) +} + +#[derive(Default)] +struct RunningServers { + listeners: Vec>, + reporter: Option, + quic_retirements: crate::listener_tasks::QuicRetirements, +} + +impl RunningServers { + async fn stop(&mut self) { + for task in &self.listeners { + task.abort(); + } + for task in self.listeners.drain(..) { + let _ = task.await; + } + self.reporter = None; + } + + async fn start(&mut self, prepared: &mut PreparedServers) -> io::Result<()> { + resources::configure(prepared.validated.global_limits)?; + self.reporter = Some(resources::ResourceReporter::start()); + for config in &prepared.validated.configs { + let dns_ref = match config { + Config::Server(server) => server.dns.as_ref(), + Config::TunServer(tun) => tun.dns.as_ref(), + _ => None, + }; + let resolver = prepared.dns.get_for_server(dns_ref); + let startup = crate::tcp::tcp_server::start_servers(config.clone(), resolver); + match self.quic_retirements.track(startup).await { + Ok(tasks) => self.listeners.extend(tasks), + Err(error) => { + self.stop().await; + return Err(error); + } + } + } + println!("Servers ready"); + Ok(()) + } +} + +impl Drop for RunningServers { + fn drop(&mut self) { + for task in &self.listeners { + task.abort(); + } + } +} + +fn watch(paths: &[String], tx: mpsc::Sender<()>) -> io::Result { + let mut files = HashSet::new(); + let mut required_directories = HashSet::new(); + for path in paths { + let path = std::path::absolute(path)?; + if let Some(parent) = path.parent() { + required_directories.insert(parent.to_path_buf()); + } + if let Ok(canonical) = std::fs::canonicalize(&path) { + if let Some(parent) = canonical.parent() { + required_directories.insert(parent.to_path_buf()); + } + files.insert(canonical); + } + files.insert(path.clone()); + resolve_watch_path(&path, &mut files, &mut 0)?; + } + let required_directories: HashSet<_> = required_directories + .iter() + .map(|path| watch_directory(path)) + .collect(); + let directories: HashSet<_> = files + .iter() + .filter_map(|path| path.parent()?.ancestors().find(|parent| parent.is_dir())) + .map(watch_directory) + .collect(); + let direct_files: HashSet<_> = files + .iter() + .filter_map(|path| std::fs::canonicalize(path).ok()) + .filter(|path| path.is_file()) + .collect(); + files.extend(direct_files.iter().cloned()); + let mut watcher = + notify::recommended_watcher(move |result: notify::Result| match result { + Ok(event) + if matches!( + event.kind, + EventKind::Create(_) | EventKind::Modify(_) | EventKind::Remove(_) + ) && event.paths.iter().any(|path| files.contains(path)) => + { + let _ = tx.try_send(()); + } + Err(error) => log::warn!("Config watcher failed: {error}"), + _ => {} + }) + .map_err(io::Error::other)?; + for directory in directories { + if let Err(error) = watcher.watch(&directory, RecursiveMode::NonRecursive) { + if required_directories.contains(&directory) { + if !matches!(&error.kind, notify::ErrorKind::Io(error) if error.kind() == io::ErrorKind::PermissionDenied) + { + return Err(io::Error::other(error)); + } + log::warn!( + "Could not watch config directory {}: {error}; using file watches, atomic replacements may be missed", + directory.display() + ); + } else { + log::warn!( + "Could not watch intermediate config directory {}: {error}", + directory.display() + ); + } + } + } + for file in direct_files { + if let Err(error) = watcher.watch(&file, RecursiveMode::NonRecursive) { + let missing = match &error.kind { + notify::ErrorKind::PathNotFound => true, + notify::ErrorKind::Io(error) => error.kind() == io::ErrorKind::NotFound, + _ => false, + }; + if !missing { + log::warn!("Could not watch config file {}: {error}", file.display()); + } + } + } + Ok(watcher) +} + +fn watch_directory(path: &Path) -> PathBuf { + // inotify keeps one event path per inode; register directory aliases only once. + #[cfg(unix)] + if let Ok(canonical) = std::fs::canonicalize(path) { + return canonical; + } + path.to_path_buf() +} + +fn resolve_watch_path( + path: &Path, + files: &mut HashSet, + symlinks: &mut usize, +) -> io::Result { + let mut resolved = PathBuf::new(); + for component in path.components() { + if component == Component::ParentDir { + resolved.pop(); + continue; + } + resolved.push(component); + match std::fs::symlink_metadata(&resolved) { + Ok(metadata) if metadata.file_type().is_symlink() => { + *symlinks += 1; + if *symlinks > 40 { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "too many config symlinks", + )); + } + let target = resolved + .parent() + .unwrap() + .join(std::fs::read_link(&resolved)?); + if cfg!(target_os = "macos") && resolved.parent() == Some(Path::new("/")) { + // System-volume links cannot be retargeted; tracking them subscribes + // FSEvents to the entire filesystem even with NonRecursive watches. + resolved = resolve_watch_path(&target, &mut HashSet::new(), symlinks)?; + } else { + files.insert(resolved.clone()); + resolved = resolve_watch_path(&target, files, symlinks)?; + } + } + // Retain missing components so creating a target or its parent triggers another reload. + Err(error) if error.kind() == io::ErrorKind::NotFound => { + files.insert(resolved.clone()); + } + Err(error) => return Err(error), + Ok(_) => {} + } + } + files.insert(resolved.clone()); + Ok(resolved) +} + +#[cfg(unix)] +struct ShutdownSignals { + interrupt: tokio::signal::unix::Signal, + terminate: tokio::signal::unix::Signal, +} + +#[cfg(unix)] +impl ShutdownSignals { + fn new() -> io::Result { + use tokio::signal::unix::{SignalKind, signal}; + Ok(Self { + interrupt: signal(SignalKind::interrupt())?, + terminate: signal(SignalKind::terminate())?, + }) + } + + async fn recv(&mut self) { + tokio::select! { + _ = self.interrupt.recv() => {}, + _ = self.terminate.recv() => {}, + } + } +} + +#[cfg(windows)] +struct ShutdownSignals(tokio::signal::windows::CtrlC); + +#[cfg(windows)] +impl ShutdownSignals { + fn new() -> io::Result { + tokio::signal::windows::ctrl_c().map(Self) + } + async fn recv(&mut self) { + self.0.recv().await; + } +} + +struct ReloadSignal { + #[cfg(unix)] + signal: tokio::signal::unix::Signal, +} + +impl ReloadSignal { + fn new() -> io::Result { + Ok(Self { + #[cfg(unix)] + signal: tokio::signal::unix::signal(tokio::signal::unix::SignalKind::hangup())?, + }) + } + + async fn recv(&mut self) { + #[cfg(unix)] + { + self.signal.recv().await; + } + #[cfg(not(unix))] + std::future::pending::<()>().await; + } +} + +pub async fn run(paths: Vec, dry_run: bool, no_reload: bool) -> io::Result<()> { + if dry_run { + load_validated(&paths).await?; + println!("Configuration valid"); + return Ok(()); + } + let mut shutdown = ShutdownSignals::new()?; + let mut reload = ReloadSignal::new()?; + let (change_tx, mut changes) = mpsc::channel(1); + let mut _watcher = if no_reload { + None + } else { + Some(watch(&paths, change_tx.clone())?) + }; + let mut prepared = tokio::select! { + biased; + _ = shutdown.recv() => return Ok(()), + result = prepare(&paths) => result?, + }; + let mut running = RunningServers::default(); + tokio::select! { + biased; + _ = shutdown.recv() => { running.stop().await; return Ok(()); } + result = running.start(&mut prepared) => result?, + } + loop { + let debounce = tokio::select! { + biased; + _ = shutdown.recv() => { running.stop().await; return Ok(()); } + _ = reload.recv() => false, + _ = changes.recv() => true, + }; + if debounce { + tokio::select! { + biased; + _ = shutdown.recv() => { running.stop().await; return Ok(()); } + _ = reload.recv() => {} + _ = tokio::time::sleep(Duration::from_secs(3)) => {} + } + } + while changes.try_recv().is_ok() {} + if !no_reload { + // A symlink may now point outside the directories watched by the previous generation. + match watch(&paths, change_tx.clone()) { + Ok(watcher) => _watcher = Some(watcher), + Err(error) => log::warn!("Could not refresh config watches: {error}"), + } + } + let candidate = tokio::select! { + biased; + _ = shutdown.recv() => { running.stop().await; return Ok(()); } + result = prepare(&paths) => result, + }; + let mut candidate = match candidate { + Ok(candidate) => candidate, + Err(error) => { + eprintln!("Reload rejected; current servers retained: {error}"); + continue; + } + }; + running.stop().await; + tokio::select! { + biased; + _ = shutdown.recv() => return Ok(()), + result = running.quic_retirements.wait() => result?, + } + let deadline = tokio::time::Instant::now() + Duration::from_secs(3); + loop { + let result = tokio::select! { + biased; + _ = shutdown.recv() => { running.stop().await; return Ok(()); } + result = running.start(&mut candidate) => result, + }; + match result { + Ok(()) => break, + Err(error) + if error.kind() == io::ErrorKind::AddrInUse + && tokio::time::Instant::now() < deadline => + { + tokio::select! { + biased; + _ = shutdown.recv() => return Ok(()), + result = running.quic_retirements.wait() => result?, + } + tokio::select! { + biased; + _ = shutdown.recv() => return Ok(()), + _ = tokio::time::sleep(Duration::from_millis(50)) => {}, + } + } + Err(error) => return Err(error), + } + } + drop(std::mem::replace(&mut prepared, candidate)); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[cfg(unix)] + #[tokio::test] + async fn directory_aliases_preserve_missing_target_notifications() { + let directory = tempfile::tempdir().unwrap(); + let root = directory.path(); + let target = root.join("target"); + std::fs::create_dir(&target).unwrap(); + std::os::unix::fs::symlink(&target, root.join("alias")).unwrap(); + let existing = root.join("alias/existing.yaml"); + std::fs::write(&existing, "existing").unwrap(); + let missing = target.join("missing.yaml"); + let (tx, mut events) = mpsc::channel(1); + let _watcher = watch( + &[ + existing.to_str().unwrap().to_owned(), + missing.to_str().unwrap().to_owned(), + ], + tx, + ) + .unwrap(); + std::fs::write(&missing, "created").unwrap(); + tokio::time::timeout(Duration::from_secs(3), events.recv()) + .await + .unwrap() + .unwrap(); + } + + #[cfg(target_os = "macos")] + #[test] + fn system_symlinks_do_not_add_root_directory_watches() { + let mut files = HashSet::new(); + resolve_watch_path(Path::new("/var/db/config.yaml"), &mut files, &mut 0).unwrap(); + assert!(files.contains(Path::new("/private/var/db/config.yaml"))); + assert!(!files.contains(Path::new("/var"))); + assert!(!files.contains(Path::new("/private/var"))); + } + + #[tokio::test] + async fn preparation_preserves_ipv6_scope_ids_and_port_ranges() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.yaml"); + std::fs::write( + &path, + "- address: '[fe80::1%2]:8080-8081'\n protocol: {type: http}\n", + ) + .unwrap(); + let prepared = prepare(&[path.to_str().unwrap().to_owned()]).await.unwrap(); + let Config::Server(server) = &prepared.validated.configs[0] else { + panic!() + }; + let BindLocation::Address(addresses) = &server.bind_location else { + panic!() + }; + let resolved: Vec<_> = addresses + .iter() + .flat_map(|address| address.to_socket_addrs().unwrap()) + .collect(); + assert_eq!( + resolved, + [ + "[fe80::1%2]:8080".parse().unwrap(), + "[fe80::1%2]:8081".parse().unwrap() + ] + ); + } + + #[cfg(unix)] + #[test] + fn watch_paths_include_missing_targets_and_symlinked_parents() { + let directory = tempfile::tempdir().unwrap(); + let root = directory.path().canonicalize().unwrap(); + std::fs::create_dir(root.join("targets")).unwrap(); + std::os::unix::fs::symlink("targets", root.join("alias")).unwrap(); + std::os::unix::fs::symlink("alias/missing/config.yaml", root.join("config.yaml")).unwrap(); + let mut files = HashSet::new(); + resolve_watch_path(&root.join("config.yaml"), &mut files, &mut 0).unwrap(); + for path in [ + "config.yaml", + "alias", + "targets/missing", + "targets/missing/config.yaml", + ] { + assert!(files.contains(&root.join(path)), "missing {path}"); + } + } + + #[cfg(unix)] + #[test] + fn watch_paths_reject_symlink_cycles() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.yaml"); + std::os::unix::fs::symlink("config.yaml", &path).unwrap(); + assert_eq!( + resolve_watch_path(&path, &mut HashSet::new(), &mut 0) + .unwrap_err() + .kind(), + io::ErrorKind::InvalidInput + ); + } +} diff --git a/src/quic_endpoint.rs b/src/quic_endpoint.rs index df7365de..f438e5c5 100644 --- a/src/quic_endpoint.rs +++ b/src/quic_endpoint.rs @@ -1,12 +1,62 @@ use std::cell::RefCell; +use std::future::Future; use std::io::{self, IoSliceMut}; use std::net::SocketAddr; use std::pin::Pin; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use std::task::{Context, Poll}; +use std::time::{Duration, Instant}; use crate::resources::BudgetPermit; use quinn::{AsyncUdpSocket, Runtime, UdpPoller}; +use tokio::sync::oneshot; +use tokio_util::sync::{CancellationToken, DropGuard}; +use tokio_util::task::AbortOnDropHandle; + +const LISTENER_CLOSE_GRACE: Duration = Duration::from_secs(1); + +#[derive(Debug)] +struct ListenerRuntime { + endpoint_started: AtomicBool, + retired: CancellationToken, +} + +impl Runtime for ListenerRuntime { + fn new_timer(&self, deadline: Instant) -> Pin> { + quinn::TokioRuntime.new_timer(deadline) + } + + fn spawn(&self, future: Pin + Send>>) { + // Quinn 0.11 spawns the endpoint driver synchronously before exposing the + // endpoint. Later spawns are connection drivers; recheck this on upgrades. + if self.endpoint_started.swap(true, Ordering::Relaxed) { + quinn::TokioRuntime.spawn(future); + return; + } + let mut driver = AbortOnDropHandle::new(tokio::spawn(future)); + let retired = self.retired.clone(); + tokio::spawn(async move { + tokio::select! { + _ = &mut driver => {}, + _ = async { + retired.cancelled().await; + tokio::time::sleep(LISTENER_CLOSE_GRACE).await; + } => {}, + } + // Dropping the endpoint driver closes its connection-event channels. + // Connection drivers then terminate and wake application waiters. + }); + } + + fn wrap_udp_socket(&self, socket: std::net::UdpSocket) -> io::Result> { + quinn::TokioRuntime.wrap_udp_socket(socket) + } + + fn now(&self) -> Instant { + quinn::TokioRuntime.now() + } +} tokio::task_local! { static CONNECTION_MEMORY: RefCell>; @@ -19,6 +69,7 @@ pub(crate) struct QuicEndpoint { inner: quinn::Endpoint, // Transport windows belong to this endpoint, not the latest config generation. memory_bytes: usize, + _driver_retirement: Option, } impl QuicEndpoint { @@ -27,19 +78,50 @@ impl QuicEndpoint { socket: std::net::UdpSocket, memory_bytes: usize, ) -> io::Result { + Self::with_release(config, socket, memory_bytes, None) + } + + pub fn listen( + config: quinn::ServerConfig, + socket: std::net::UdpSocket, + memory_bytes: usize, + ) -> io::Result<(Self, oneshot::Receiver<()>)> { + let (release, released) = oneshot::channel(); + let endpoint = Self::with_release(Some(config), socket, memory_bytes, Some(release))?; + Ok((endpoint, released)) + } + + fn with_release( + config: Option, + socket: std::net::UdpSocket, + memory_bytes: usize, + release: Option>, + ) -> io::Result { + let (runtime, retirement): (Arc, _) = if release.is_some() { + let retired = CancellationToken::new(); + let runtime = ListenerRuntime { + endpoint_started: AtomicBool::new(false), + retired: retired.clone(), + }; + (Arc::new(runtime), Some(retired.drop_guard())) + } else { + (Arc::new(quinn::TokioRuntime), None) + }; let socket = MemorySocket { inner: quinn::TokioRuntime.wrap_udp_socket(socket)?, endpoint_memory: None, + _release: release, }; quinn::Endpoint::new_with_abstract_socket( quinn::EndpointConfig::default(), config, Arc::new(socket), - Arc::new(quinn::TokioRuntime), + runtime, ) .map(|inner| Self { inner, memory_bytes, + _driver_retirement: retirement, }) } @@ -49,7 +131,7 @@ impl QuicEndpoint { pub fn connect(&self, address: SocketAddr, name: &str) -> io::Result { let memory = crate::resources::try_quic_memory(self.memory_bytes) - .ok_or_else(crate::resources::exhausted)?; + .ok_or_else(|| crate::resources::quic_memory_exhausted(self.memory_bytes))?; self.connect_with_memory(address, name, memory) } @@ -102,8 +184,9 @@ impl Incoming { pub fn accept(self) -> io::Result { let Some(memory) = crate::resources::try_quic_memory(self.memory_bytes) else { + let error = crate::resources::quic_memory_exhausted(self.memory_bytes); self.refuse(); - return Err(crate::resources::exhausted()); + return Err(error); }; self.accept_with_memory(memory) } @@ -124,6 +207,7 @@ pub(crate) fn socket_with_memory( Ok(Arc::new(MemorySocket { inner: quinn::TokioRuntime.wrap_udp_socket(socket)?, endpoint_memory: Some(memory), + _release: None, })) } @@ -131,6 +215,8 @@ pub(crate) fn socket_with_memory( struct MemorySocket { inner: Arc, endpoint_memory: Option, + // Field order signals retirement only after the native socket has been dropped. + _release: Option>, } impl AsyncUdpSocket for MemorySocket { @@ -150,6 +236,7 @@ impl AsyncUdpSocket for MemorySocket { Box::pin(MemoryPoller { inner: self.inner.clone().create_io_poller(), _memory: memory, + _socket: self, }) } @@ -187,6 +274,8 @@ impl AsyncUdpSocket for MemorySocket { struct MemoryPoller { inner: Pin>, _memory: BudgetPermit, + // The delegated poller owns a socket reference and must drop before the release signal. + _socket: Arc, } impl UdpPoller for MemoryPoller { @@ -205,6 +294,137 @@ mod tests { use bytes::Bytes; use tokio::time::timeout; + #[tokio::test(start_paused = true)] + async fn listener_runtime_cancels_only_the_endpoint_driver_after_grace() { + let retired = CancellationToken::new(); + let runtime = ListenerRuntime { + endpoint_started: AtomicBool::new(false), + retired: retired.clone(), + }; + let endpoint = Arc::new(()); + let owned_endpoint = endpoint.clone(); + runtime.spawn(Box::pin(async move { + let _owned = owned_endpoint; + std::future::pending::<()>().await; + })); + let connection = Arc::new(()); + let owned_connection = connection.clone(); + let (finish, finished) = oneshot::channel(); + runtime.spawn(Box::pin(async move { + let _owned = owned_connection; + let _ = finished.await; + })); + tokio::task::yield_now().await; + retired.cancel(); + tokio::task::yield_now().await; + tokio::time::advance(LISTENER_CLOSE_GRACE - Duration::from_millis(1)).await; + assert_eq!(Arc::strong_count(&endpoint), 2); + tokio::time::advance(Duration::from_millis(2)).await; + for _ in 0..4 { + tokio::task::yield_now().await; + } + assert_eq!(Arc::strong_count(&endpoint), 1); + assert_eq!(Arc::strong_count(&connection), 2); + finish.send(()).unwrap(); + tokio::task::yield_now().await; + assert_eq!(Arc::strong_count(&connection), 1); + } + + #[tokio::test] + async fn listener_runtime_reaps_a_completed_driver_and_its_watcher() { + let metrics = tokio::runtime::Handle::current().metrics(); + let baseline = metrics.num_alive_tasks(); + let runtime = ListenerRuntime { + endpoint_started: AtomicBool::new(false), + retired: CancellationToken::new(), + }; + runtime.spawn(Box::pin(async {})); + for _ in 0..4 { + tokio::task::yield_now().await; + } + assert_eq!(metrics.num_alive_tasks(), baseline); + } + + #[tokio::test] + async fn endpoint_cancellation_wakes_retained_handles_before_releasing_memory() { + let (client_config, server_config) = configs(); + let (server, mut released) = QuicEndpoint::listen( + server_config, + std::net::UdpSocket::bind("127.0.0.1:0").unwrap(), + 1, + ) + .unwrap(); + let address = server.local_addr().unwrap(); + let mut client = quinn::Endpoint::client("127.0.0.1:0".parse().unwrap()).unwrap(); + client.set_default_client_config(client_config); + let budget = Arc::new(Budget::new(Some(1))); + let (client_conn, server_conn) = tokio::join!( + async { client.connect(address, "localhost").unwrap().await.unwrap() }, + async { + server + .accept() + .await + .unwrap() + .accept_with_memory(budget.acquire(1).unwrap()) + .unwrap() + .await + .unwrap() + }, + ); + let (mut client_send, _client_recv) = client_conn.open_bi().await.unwrap(); + client_send.write_all(b"x").await.unwrap(); + let (mut server_send, mut server_recv) = server_conn.accept_bi().await.unwrap(); + server_recv.read_exact(&mut [0]).await.unwrap(); + // No close() here: channel termination must wake handles on its own. + drop(server); + timeout(Duration::from_secs(3), server_conn.closed()) + .await + .unwrap(); + assert!(server_send.write_all(b"x").await.is_err()); + assert!(server_recv.read(&mut [0]).await.is_err()); + assert_eq!( + released.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + ); + assert_eq!(budget.snapshot().active, 1); + drop((server_conn, server_send, server_recv)); + assert!( + timeout(Duration::from_secs(1), released) + .await + .unwrap() + .is_err() + ); + assert_eq!(budget.snapshot().active, 0); + drop(std::net::UdpSocket::bind(address).unwrap()); + client.close(0u32.into(), b"done"); + } + + #[tokio::test] + async fn socket_release_waits_for_delegated_poller() { + let socket = std::net::UdpSocket::bind("127.0.0.1:0").unwrap(); + let address = socket.local_addr().unwrap(); + let (release, mut released) = oneshot::channel(); + let socket = Arc::new(MemorySocket { + inner: quinn::TokioRuntime.wrap_udp_socket(socket).unwrap(), + endpoint_memory: None, + _release: Some(release), + }); + let budget = Arc::new(Budget::new(Some(1))); + let poller = CONNECTION_MEMORY.sync_scope(RefCell::new(budget.acquire(1)), || { + socket.clone().create_io_poller() + }); + drop(socket); + assert_eq!( + released.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + ); + assert!(std::net::UdpSocket::bind(address).is_err()); + drop(poller); + assert!(released.await.is_err()); + assert_eq!(budget.available_permits(), 1); + drop(std::net::UdpSocket::bind(address).unwrap()); + } + async fn wait_for_slot(budget: &Arc) -> BudgetPermit { loop { if let Some(permit) = budget.acquire(1) { diff --git a/src/quic_server.rs b/src/quic_server.rs index aaf77567..a7e4ac77 100644 --- a/src/quic_server.rs +++ b/src/quic_server.rs @@ -10,7 +10,7 @@ use crate::config::{ }; use crate::quic_stream::QuicStream; use crate::resolver::Resolver; -use crate::rustls_config_util::create_server_config; +use crate::rustls_config_util::try_create_server_config; use crate::tcp::tcp_client_handler_factory::create_tcp_client_proxy_selector; use crate::tcp::tcp_handler::TcpServerHandler; use crate::tcp::tcp_server_handler_factory::create_tcp_server_handler; @@ -174,13 +174,13 @@ pub async fn start_quic_servers( processed_ca_certs.push(cert.as_bytes().to_vec()); } - let server_config = Arc::new(create_server_config( + let server_config = Arc::new(try_create_server_config( &cert_bytes, &key_bytes, processed_ca_certs, &alpn_protocols.into_vec(), &client_fingerprints.into_vec(), - )); + )?); let quic_server_config: quinn::crypto::rustls::QuicServerConfig = server_config .try_into() @@ -191,7 +191,7 @@ pub async fn start_quic_servers( let client_proxy_selector = Arc::new(create_tcp_client_proxy_selector( rules.clone(), resolver.clone(), - )); + )?); let mut handles = vec![]; @@ -248,20 +248,24 @@ pub async fn start_quic_servers( tcp_protocol => { // Shares protocol state across ports without reusing an interface-specific UDP bind IP. let mut handlers: HashMap> = HashMap::new(); - - for bind_address in bind_addresses.into_iter() { - let tcp_handler: Arc = handlers - .entry(bind_address.ip()) - .or_insert_with(|| { + // Construct every handler before spawning listeners so errors cannot detach tasks. + for bind_address in &bind_addresses { + if let std::collections::hash_map::Entry::Vacant(entry) = + handlers.entry(bind_address.ip()) + { + entry.insert( create_tcp_server_handler( tcp_protocol.clone(), &client_proxy_selector, &resolver, Some(bind_address.ip()), - ) - .into() - }) - .clone(); + )? + .into(), + ); + } + } + for bind_address in bind_addresses.into_iter() { + let tcp_handler = handlers[&bind_address.ip()].clone(); let quic_server_config = quic_server_config.clone(); let resolver = resolver.clone(); let result = start_quic_server( diff --git a/src/reality/common.rs b/src/reality/common.rs index 77a186ed..01538d46 100644 --- a/src/reality/common.rs +++ b/src/reality/common.rs @@ -63,6 +63,16 @@ pub const TLS_RECORD_HEADER_SIZE: usize = 5; /// Maximum TLS record size (ciphertext + header) pub const TLS_MAX_RECORD_SIZE: usize = MAX_TLS_CIPHERTEXT_LEN + TLS_RECORD_HEADER_SIZE; +pub(crate) fn checked_record_size(payload_len: usize) -> std::io::Result { + if payload_len > MAX_TLS_CIPHERTEXT_LEN { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "TLS record length exceeds maximum", + )); + } + Ok(TLS_RECORD_HEADER_SIZE + payload_len) +} + /// Buffer capacity for ciphertext read (2x TLS max record for safety) pub const CIPHERTEXT_READ_BUF_CAPACITY: usize = TLS_MAX_RECORD_SIZE * 2; @@ -158,6 +168,23 @@ pub fn strip_content_type_with_padding(plaintext: &mut Vec) -> io::Result io::Resu "REALITY CLIENT: HMAC verification - ed25519_pubkey={:02x?}", pubkey_data ); - log::debug!( - "REALITY CLIENT: HMAC verification - expected_sig={:02x?}", - expected_signature - ); log::debug!( "REALITY CLIENT: HMAC verification - actual_sig={:02x?}", signature diff --git a/src/reality/reality_server_connection.rs b/src/reality/reality_server_connection.rs index 836c1d7b..2976fa13 100644 --- a/src/reality/reality_server_connection.rs +++ b/src/reality/reality_server_connection.rs @@ -379,13 +379,10 @@ impl RealityServerConnection { }); if !short_id_ok { - log::warn!( - "REALITY: Client short_id {:02x?} not in configured list", - client_short_id - ); + log::warn!("REALITY: Client short_id not in configured list"); return Err(io::Error::new( io::ErrorKind::PermissionDenied, - format!("Invalid short_id: {:02x?}", client_short_id), + "Invalid short_id", )); } @@ -454,8 +451,7 @@ impl RealityServerConnection { } log::debug!( - "REALITY: Client authentication successful - short_id: {:02x?}, version: {:?}, timestamp: {}", - client_short_id, + "REALITY: Client authentication successful - version: {:?}, timestamp: {}", client_version, client_timestamp ); @@ -738,15 +734,26 @@ impl RealityServerConnection { .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "Buffer too short"))? as usize; - // Need complete ChangeCipherSpec record - if self.ciphertext_read_buf.len() < TLS_RECORD_HEADER_SIZE + ccs_len { + let record_size = super::common::checked_record_size(ccs_len)?; + if ccs_len != 1 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "invalid ChangeCipherSpec length", + )); + } + if self.ciphertext_read_buf.len() < record_size { return Ok(false); // Need more data } + if self.ciphertext_read_buf[TLS_RECORD_HEADER_SIZE] != 1 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "invalid ChangeCipherSpec payload", + )); + } // Skip ChangeCipherSpec (compatibility message) log::debug!("REALITY: Skipping ChangeCipherSpec (compatibility message)"); - self.ciphertext_read_buf - .consume(TLS_RECORD_HEADER_SIZE + ccs_len); + self.ciphertext_read_buf.consume(record_size); // Check if we have the next record header if self.ciphertext_read_buf.len() < TLS_RECORD_HEADER_SIZE { @@ -762,7 +769,7 @@ impl RealityServerConnection { as usize; // Check if we have the complete record - let total_record_len = TLS_RECORD_HEADER_SIZE + record_len; + let total_record_len = super::common::checked_record_size(record_len)?; if self.ciphertext_read_buf.len() < total_record_len { return Ok(false); // Need more data } @@ -903,7 +910,7 @@ impl RealityServerConnection { as usize; // Check if we have the complete record - let total_record_len = TLS_RECORD_HEADER_SIZE + record_len; + let total_record_len = super::common::checked_record_size(record_len)?; if self.ciphertext_read_buf.len() < total_record_len { break; // Need more data } @@ -1112,6 +1119,38 @@ pub fn feed_reality_server_connection( mod tests { use super::*; + #[test] + fn fragmented_change_cipher_spec_rejects_invalid_length_before_buffering_payload() { + for length in [0u16, 2, 65535] { + let mut connection = RealityServerConnection::new(RealityServerConfig { + private_key: [0; 32], + short_ids: vec![[0; 8]], + dest: NetLocation::new(Address::UNSPECIFIED, 443), + max_time_diff: None, + min_client_version: None, + max_client_version: None, + cipher_suites: vec![], + }) + .unwrap(); + let header = [ + CONTENT_TYPE_CHANGE_CIPHER_SPEC, + 3, + 3, + (length >> 8) as u8, + length as u8, + ]; + for byte in &header[..4] { + connection.read_tls(&mut &[*byte][..]).unwrap(); + assert!(!connection.process_client_finished().unwrap()); + } + connection.read_tls(&mut &header[4..]).unwrap(); + assert_eq!( + connection.process_client_finished().unwrap_err().kind(), + io::ErrorKind::InvalidData + ); + } + } + #[test] fn test_reality_server_connection_creation() { let config = RealityServerConfig { diff --git a/src/reality/reality_util.rs b/src/reality/reality_util.rs index eaa9ec1b..829b18e2 100644 --- a/src/reality/reality_util.rs +++ b/src/reality/reality_util.rs @@ -5,6 +5,18 @@ use rand::Rng; use super::reality_cipher_suite::CipherSuite; use crate::buf_reader::BufReader; +pub(crate) fn parse_server_name( + name: &str, +) -> std::io::Result> { + match rustls::pki_types::ServerName::try_from(name.to_owned()) { + Ok(name @ rustls::pki_types::ServerName::DnsName(_)) => Ok(name), + _ => Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "REALITY requires a valid DNS server name", + )), + } +} + /// Decodes a base64url-encoded public key pub fn decode_public_key(encoded: &str) -> Result<[u8; 32], std::io::Error> { let decoded = URL_SAFE_NO_PAD.decode(encoded).map_err(|e| { diff --git a/src/resources.rs b/src/resources.rs index 399deef6..d96c94ab 100644 --- a/src/resources.rs +++ b/src/resources.rs @@ -165,9 +165,18 @@ pub(crate) fn try_quic_memory(bytes: usize) -> Option { QUIC_BYTES.acquire(bytes) } -pub(crate) fn try_dns_quic_memory() -> Option { - let limits = LIMITS.read(); - QUIC_BYTES.acquire(limits.quic_dns_memory_bytes) +pub(crate) fn quic_memory_exhausted(bytes: usize) -> std::io::Error { + let state = QUIC_BYTES.0.lock(); + let limit = state + .limit + .map_or_else(|| "unlimited".to_owned(), |limit| limit.to_string()); + std::io::Error::new( + std::io::ErrorKind::ConnectionRefused, + format!( + "QUIC memory budget exhausted: requested {bytes} bytes, active {} bytes, live cap {limit}", + state.active, + ), + ) } pub(crate) fn configure_quic( @@ -200,7 +209,13 @@ pub(crate) fn configure_quic( /// Returns the per-connection allowance from the same limits used to configure the windows. pub(crate) fn configure_quic_transport(transport: &mut quinn::TransportConfig) -> usize { - let limits = limits(); + configure_quic_transport_with_limits(transport, limits()) +} + +pub(crate) fn configure_quic_transport_with_limits( + transport: &mut quinn::TransportConfig, + limits: GlobalLimits, +) -> usize { transport .receive_window((limits.quic_receive_window as u32).into()) .send_window(limits.quic_send_window as u64) diff --git a/src/rustls_config_util.rs b/src/rustls_config_util.rs index e9fea190..3c9532f3 100644 --- a/src/rustls_config_util.rs +++ b/src/rustls_config_util.rs @@ -4,14 +4,14 @@ use std::sync::OnceLock; use rustls::pki_types::pem::PemObject; -pub fn create_client_config( +pub fn try_create_client_config( verify_webpki: bool, server_fingerprints: Vec, alpn_protocols: Vec, enable_sni: bool, client_key_and_cert: Option<(Vec, Vec)>, tls13_only: bool, -) -> rustls::ClientConfig { +) -> std::io::Result { let builder = rustls::ClientConfig::builder_with_provider(get_crypto_provider()); let builder = if tls13_only { builder @@ -33,7 +33,7 @@ pub fn create_client_config( .dangerous() .with_custom_certificate_verifier(Arc::new(ServerFingerprintVerifier { supported_algs: get_supported_algorithms(), - server_fingerprints: process_fingerprints(&server_fingerprints).unwrap(), + server_fingerprints: process_fingerprints(&server_fingerprints)?, webpki_verifier: Some(Arc::into_inner(webpki_verifier).unwrap()), })) } else { @@ -44,7 +44,7 @@ pub fn create_client_config( .dangerous() .with_custom_certificate_verifier(Arc::new(ServerFingerprintVerifier { supported_algs: get_supported_algorithms(), - server_fingerprints: process_fingerprints(&server_fingerprints).unwrap(), + server_fingerprints: process_fingerprints(&server_fingerprints)?, webpki_verifier: None, })) } else { @@ -57,13 +57,14 @@ pub fn create_client_config( Some((key_bytes, cert_bytes)) => { // Parse all certificates from the PEM file (client cert + intermediates if any) let certs: Vec<_> = rustls::pki_types::CertificateDer::pem_slice_iter(&cert_bytes) - .map(|cert| cert.unwrap().into_owned()) - .collect(); + .collect::, _>>() + .map_err(std::io::Error::other)?; - let privkey = rustls::pki_types::PrivateKeyDer::from_pem_slice(&key_bytes).unwrap(); + let privkey = rustls::pki_types::PrivateKeyDer::from_pem_slice(&key_bytes) + .map_err(std::io::Error::other)?; builder .with_client_auth_cert(certs, privkey) - .expect("Could not parse client certificate") + .map_err(std::io::Error::other)? } None => builder.with_no_client_auth(), }; @@ -74,7 +75,7 @@ pub fn create_client_config( .collect(); config.enable_sni = enable_sni; - config + Ok(config) } #[derive(Debug)] @@ -229,24 +230,25 @@ pub fn create_dns_client_config() -> rustls::ClientConfig { .with_no_client_auth() } -pub fn create_server_config( +pub fn try_create_server_config( cert_bytes: &[u8], key_bytes: &[u8], ca_cert_bytes: Vec>, alpn_protocols: &[String], client_fingerprints: &[String], -) -> rustls::ServerConfig { +) -> std::io::Result { // Parse all certificates from the PEM file (server cert + intermediates) let certs: Vec<_> = rustls::pki_types::CertificateDer::pem_slice_iter(cert_bytes) - .map(|cert| cert.unwrap().into_owned()) - .collect(); + .collect::, _>>() + .map_err(std::io::Error::other)?; log::debug!( "TLS server config: loaded {} certificate(s) in chain", certs.len() ); - let privkey = rustls::pki_types::PrivateKeyDer::from_pem_slice(key_bytes).unwrap(); + let privkey = rustls::pki_types::PrivateKeyDer::from_pem_slice(key_bytes) + .map_err(std::io::Error::other)?; let webpki_verifier = if ca_cert_bytes.is_empty() { None @@ -254,16 +256,16 @@ pub fn create_server_config( let mut store = rustls::RootCertStore::empty(); for ca_cert in ca_cert_bytes.into_iter() { let ca_cert = rustls::pki_types::CertificateDer::from_pem_slice(&ca_cert) - .unwrap() + .map_err(std::io::Error::other)? .into_owned(); - store.add(ca_cert).unwrap(); + store.add(ca_cert).map_err(std::io::Error::other)?; } let verifier = rustls::server::WebPkiClientVerifier::builder_with_provider( Arc::new(store), get_crypto_provider(), ) .build() - .unwrap(); + .map_err(std::io::Error::other)?; Some(verifier) }; @@ -278,12 +280,12 @@ pub fn create_server_config( builder.with_client_cert_verifier(Arc::new(ClientFingerprintVerifier { supported_algs: get_supported_algorithms(), webpki_verifier, - client_fingerprints: process_fingerprints(client_fingerprints).unwrap(), + client_fingerprints: process_fingerprints(client_fingerprints)?, })) }; let mut config = builder .with_single_cert(certs, privkey) - .expect("bad certificate/key"); + .map_err(std::io::Error::other)?; config.alpn_protocols = alpn_protocols .iter() @@ -294,7 +296,7 @@ pub fn create_server_config( config.max_early_data_size = u32::MAX; config.ignore_client_order = true; - config + Ok(config) } pub fn process_fingerprints(client_fingerprints: &[String]) -> std::io::Result>> { @@ -304,10 +306,12 @@ pub fn process_fingerprints(client_fingerprints: &[String]) -> std::io::Result std::io::Result bool { + let Some((tag, payload)) = record.split_at_checked(4) else { + return false; + }; + let mut candidate = self.clone(); + candidate.update(payload); + bool::from(candidate.finalized_digest().ct_eq(tag)) + } + pub fn digest(&self) -> [u8; 4] { let tag = self.context.clone().sign(); let mut out = [0u8; 4]; @@ -28,3 +39,28 @@ impl ShadowTlsHmac { out } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn record_verification_is_bounded_and_does_not_advance_the_mac() { + let key = aws_lc_rs::hmac::Key::new(aws_lc_rs::hmac::HMAC_SHA256, b"test"); + let mac = ShadowTlsHmac::new(&key); + for size in 0..4 { + assert!(!mac.authenticates_record(&[0; 4][..size])); + } + let mut signed = mac.clone(); + signed.update(b"payload"); + let mut record = signed.finalized_digest().to_vec(); + record.extend_from_slice(b"payload"); + assert!(mac.authenticates_record(&record)); + assert!(mac.authenticates_record(&record)); + for index in 0..record.len() { + record[index] ^= 1; + assert!(!mac.authenticates_record(&record)); + record[index] ^= 1; + } + } +} diff --git a/src/shadow_tls/shadow_tls_server_handler.rs b/src/shadow_tls/shadow_tls_server_handler.rs index 01d1bcc9..ecb1feba 100644 --- a/src/shadow_tls/shadow_tls_server_handler.rs +++ b/src/shadow_tls/shadow_tls_server_handler.rs @@ -7,6 +7,7 @@ use std::fmt::Debug; use std::io::Cursor; use std::sync::Arc; +use subtle::ConstantTimeEq; use tokio::io::AsyncWriteExt; @@ -167,7 +168,7 @@ fn validate_shadowtls_client_hello( hmac.update(&[0; 4]); hmac.update(&client_hello_frame[digest.client_hello_digest_end_index..]); - if digest.client_hello_digest != hmac.finalized_digest() { + if !bool::from(digest.client_hello_digest.ct_eq(&hmac.finalized_digest())) { return Err(std::io::Error::new( std::io::ErrorKind::PermissionDenied, "HMAC tag mismatch", @@ -867,38 +868,34 @@ async fn setup_remote_handshake( format!("failed to read TLS payload from client during handshake (size {client_payload_size}): {e}") ))?; - if client_content_type == CONTENT_TYPE_APPLICATION_DATA { - let mut tmp_hmac = hmac_client_data.clone(); - tmp_hmac.update(&client_payload_bytes[4..]); + if client_content_type == CONTENT_TYPE_APPLICATION_DATA + && hmac_client_data.authenticates_record(client_payload_bytes) { + let initial_client_data = &client_payload_bytes[4..]; - if tmp_hmac.finalized_digest() == client_payload_bytes[..4] { - let initial_client_data = &client_payload_bytes[4..]; + hmac_client_data.update(initial_client_data); + hmac_client_data.update(&hmac_client_data.digest()); - hmac_client_data.update(initial_client_data); - hmac_client_data.update(&hmac_client_data.digest()); + let _ = client_stream.shutdown().await; - let _ = client_stream.shutdown().await; + let mut shadow_tls_stream = ShadowTlsStream::new( + server_stream, + initial_client_data, + hmac_client_data, + hmac_server_data, + None, + ).map_err(|e| std::io::Error::other( + format!("failed to create ShadowTlsStream: {e}") + ))?; - let mut shadow_tls_stream = ShadowTlsStream::new( - server_stream, - initial_client_data, - hmac_client_data, - hmac_server_data, - None, - ).map_err(|e| std::io::Error::other( - format!("failed to create ShadowTlsStream: {e}") - ))?; - - let unparsed_data = client_reader.unparsed_data(); - if !unparsed_data.is_empty() { - shadow_tls_stream.feed_initial_read_data(unparsed_data) - .map_err(|e| std::io::Error::other( - format!("failed to feed initial data to ShadowTlsStream: {e}") - ))?; - } - - return Ok(shadow_tls_stream); + let unparsed_data = client_reader.unparsed_data(); + if !unparsed_data.is_empty() { + shadow_tls_stream.feed_initial_read_data(unparsed_data) + .map_err(|e| std::io::Error::other( + format!("failed to feed initial data to ShadowTlsStream: {e}") + ))?; } + + return Ok(shadow_tls_stream); } client_frame.extend_from_slice(client_payload_bytes); @@ -1113,32 +1110,29 @@ async fn setup_local_handshake( .read_slice(&mut server_stream, client_payload_size as usize) .await?; - if client_content_type == CONTENT_TYPE_APPLICATION_DATA { - let mut tmp_hmac = hmac_client_data.clone(); - tmp_hmac.update(&client_payload_bytes[4..]); + if client_content_type == CONTENT_TYPE_APPLICATION_DATA + && hmac_client_data.authenticates_record(client_payload_bytes) + { + let initial_client_data = &client_payload_bytes[4..]; - if tmp_hmac.finalized_digest() == client_payload_bytes[..4] { - let initial_client_data = &client_payload_bytes[4..]; + hmac_client_data.update(initial_client_data); + hmac_client_data.update(&hmac_client_data.digest()); - hmac_client_data.update(initial_client_data); - hmac_client_data.update(&hmac_client_data.digest()); + let mut shadow_tls_stream = ShadowTlsStream::new( + server_stream, + initial_client_data, + hmac_client_data, + hmac_server_data, + None, + )?; - let mut shadow_tls_stream = ShadowTlsStream::new( - server_stream, - initial_client_data, - hmac_client_data, - hmac_server_data, - None, - )?; - - // Feeds any leftover data from the reader to the stream. - let leftover = client_reader.unparsed_data(); - if !leftover.is_empty() { - shadow_tls_stream.feed_initial_read_data(leftover)?; - } - - return Ok(shadow_tls_stream); + // Feeds any leftover data from the reader to the stream. + let leftover = client_reader.unparsed_data(); + if !leftover.is_empty() { + shadow_tls_stream.feed_initial_read_data(leftover)?; } + + return Ok(shadow_tls_stream); } feed_rustls_server_connection(&mut server_connection, &client_header_bytes)?; diff --git a/src/shadow_tls/shadow_tls_stream.rs b/src/shadow_tls/shadow_tls_stream.rs index 81cf88a4..1f7044c0 100644 --- a/src/shadow_tls/shadow_tls_stream.rs +++ b/src/shadow_tls/shadow_tls_stream.rs @@ -1,5 +1,6 @@ use std::pin::Pin; use std::task::{Context, Poll}; +use subtle::ConstantTimeEq; use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; @@ -195,7 +196,7 @@ impl ShadowTlsStream { if let Some(ref mut handshake_hmac) = self.handshake_hmac { handshake_hmac.update(payload); let expected_digest = handshake_hmac.digest(); - if received_digest == expected_digest { + if bool::from(received_digest.ct_eq(&expected_digest)) { if total_len < self.unprocessed_end_offset { self.unprocessed_buf .copy_within(total_len..self.unprocessed_end_offset, 0); @@ -213,7 +214,7 @@ impl ShadowTlsStream { self.read_hmac.update(payload); let expected_digest = self.read_hmac.digest(); - if received_digest != expected_digest { + if !bool::from(received_digest.ct_eq(&expected_digest)) { return Err(std::io::Error::new( std::io::ErrorKind::InvalidData, "HMAC verification failed", diff --git a/src/shadowsocks/shadowsocks_stream.rs b/src/shadowsocks/shadowsocks_stream.rs index 8e4a3cbd..b5ae4525 100644 --- a/src/shadowsocks/shadowsocks_stream.rs +++ b/src/shadowsocks/shadowsocks_stream.rs @@ -88,6 +88,68 @@ enum DecryptState { const METADATA_SIZE: usize = 2 + (2 * TAG_LEN); +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn timestamp_window_is_symmetric_and_inclusive() { + for timestamp in [970, 1000, 1030] { + assert!(validate_timestamp(timestamp, 1000).is_ok()); + } + for timestamp in [969, 1031, u64::MAX] { + assert!(validate_timestamp(timestamp, 1000).is_err()); + } + assert!(validate_timestamp(0, 0).is_ok()); + } + + #[tokio::test] + async fn full_output_buffer_preserves_the_decrypted_length_and_nonce() { + let algorithm = &aws_lc_rs::aead::AES_128_GCM; + let (transport, _) = tokio::io::duplex(1); + let mut stream = ShadowsocksStream::new( + Box::new(transport), + ShadowsocksStreamType::Aead, + algorithm, + 16, + Arc::new(Box::new(super::super::default_key::DefaultKey::new( + "test", 16, + ))), + None, + ); + let mut seal = SealingKey::new( + UnboundKey::new(algorithm, &[7; 16]).unwrap(), + IncreasingSequence::new(), + ); + stream.opening_key = Some(OpeningKey::new( + UnboundKey::new(algorithm, &[7; 16]).unwrap(), + IncreasingSequence::new(), + )); + let mut packet = Vec::new(); + for plaintext in [&[0, 3][..], b"abc"] { + let mut chunk = plaintext.to_vec(); + seal.seal_in_place_append_tag(Aad::empty(), &mut chunk) + .unwrap(); + packet.extend(chunk); + } + stream.unprocessed_buf[..packet.len()].copy_from_slice(&packet); + stream.unprocessed_end_offset = packet.len(); + stream.processed_end_offset = stream.processed_buf.len(); + assert!(matches!( + stream.try_decrypt().unwrap(), + DecryptState::BufferFull + )); + assert_eq!(stream.unprocessed_pending_len, Some(3)); + stream.processed_end_offset = 0; + assert!(matches!( + stream.try_decrypt().unwrap(), + DecryptState::Success + )); + assert_eq!(&stream.processed_buf[..3], b"abc"); + assert_eq!(stream.unprocessed_pending_len, None); + } +} + impl ShadowsocksStream { pub fn new( stream: Box, @@ -222,6 +284,8 @@ impl ShadowsocksStream { } if self.processed_end_offset + data_len_no_tag > self.processed_buf.len() { + // The length nonce has already advanced; retry the payload, not its header. + self.unprocessed_pending_len = Some(data_len_no_tag); return Ok(DecryptState::BufferFull); } @@ -449,21 +513,7 @@ impl ShadowsocksStream { let timestamp_bytes = &self.unprocessed_buf[self.salt_len + 1..self.salt_len + 9]; let timestamp_secs = u64::from_be_bytes(timestamp_bytes.try_into().unwrap()); let current_time_secs = current_time_secs(); - if current_time_secs >= timestamp_secs { - if current_time_secs - timestamp_secs > 30 { - return Err(std::io::Error::other( - "timestamp is greater than 30 seconds", - )); - } - } else { - // Make sure times aren't too far in the future. - if timestamp_secs - current_time_secs > 2 { - return Err(std::io::Error::other(format!( - "timestamp is {} seconds in the future", - timestamp_secs - current_time_secs - ))); - } - } + validate_timestamp(timestamp_secs, current_time_secs)?; let decrypt_iv = &self.unprocessed_buf[0..self.salt_len]; if let Some(salt_checker) = &self.salt_checker @@ -516,21 +566,7 @@ impl ShadowsocksStream { let timestamp_bytes = &self.unprocessed_buf[self.salt_len + 1..self.salt_len + 9]; let timestamp_secs = u64::from_be_bytes(timestamp_bytes.try_into().unwrap()); let current_time_secs = current_time_secs(); - if current_time_secs >= timestamp_secs { - if current_time_secs - timestamp_secs > 30 { - return Err(std::io::Error::other( - "timestamp is greater than 30 seconds", - )); - } - } else { - // Make sure times aren't too far in the future. - if timestamp_secs - current_time_secs > 2 { - return Err(std::io::Error::other(format!( - "timestamp is {} seconds in the future", - timestamp_secs - current_time_secs - ))); - } - } + validate_timestamp(timestamp_secs, current_time_secs)?; if let Some(salt_checker) = &self.salt_checker { let decrypt_iv = &self.unprocessed_buf[0..self.salt_len]; @@ -929,6 +965,16 @@ fn current_time_secs() -> u64 { SystemTime::UNIX_EPOCH.elapsed().unwrap().as_secs() } +fn validate_timestamp(timestamp: u64, now: u64) -> std::io::Result<()> { + if timestamp.abs_diff(now) > 30 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "timestamp outside the 30-second window", + )); + } + Ok(()) +} + #[cfg(test)] mod flush_tests { use super::*; diff --git a/src/shadowsocks/shadowsocks_tcp_handler.rs b/src/shadowsocks/shadowsocks_tcp_handler.rs index 35e5629d..629ce35f 100644 --- a/src/shadowsocks/shadowsocks_tcp_handler.rs +++ b/src/shadowsocks/shadowsocks_tcp_handler.rs @@ -76,7 +76,7 @@ impl ShadowsocksTcpHandler { cipher, key, aead2022: false, - salt_checker: None, + salt_checker: Some(Arc::new(Mutex::new(TimedSaltChecker::new(60)))), udp_enabled, proxy_selector: None, resolver: None, @@ -99,7 +99,7 @@ impl ShadowsocksTcpHandler { cipher, key, aead2022: true, - salt_checker: Some(Arc::new(Mutex::new(TimedSaltChecker::new(60)))), + salt_checker: Some(Arc::new(Mutex::new(TimedSaltChecker::new(61)))), udp_enabled, proxy_selector: Some(proxy_selector), resolver: Some(resolver), @@ -120,7 +120,7 @@ impl ShadowsocksTcpHandler { cipher, key, aead2022: true, - salt_checker: Some(Arc::new(Mutex::new(TimedSaltChecker::new(60)))), + salt_checker: Some(Arc::new(Mutex::new(TimedSaltChecker::new(61)))), udp_enabled, proxy_selector: None, resolver: None, @@ -402,3 +402,56 @@ impl TcpClientHandler for ShadowsocksTcpHandler { Ok(Box::new(message_stream)) } } + +#[cfg(test)] +mod replay_tests { + use super::*; + use tokio::io::AsyncReadExt; + + #[tokio::test] + async fn classic_client_rejects_authenticated_response_replays() { + let handler = ShadowsocksTcpHandler::new_client( + ShadowsocksCipher::try_from("aes-128-gcm").unwrap(), + "test", + false, + ); + let (transport, mut wire_reader) = tokio::io::duplex(4096); + let mut response = ShadowsocksStream::new( + Box::new(transport), + ShadowsocksStreamType::Aead, + handler.cipher.algorithm(), + handler.cipher.salt_len(), + handler.key.clone(), + None, + ); + response.write_all(b"hello").await.unwrap(); + response.flush().await.unwrap(); + let mut wire = vec![0; handler.cipher.salt_len() + 2 + 16 + 5 + 16]; + wire_reader.read_exact(&mut wire).await.unwrap(); + + for attempt in 0..3 { + let (transport, mut peer) = tokio::io::duplex(4096); + let mut client = handler + .setup_client_tcp_stream( + Box::new(transport), + NetLocation::from_str("127.0.0.1:80", None).unwrap().into(), + ) + .await + .unwrap() + .client_stream; + let mut packet = wire.clone(); + if attempt == 0 { + *packet.last_mut().unwrap() ^= 1; + } + peer.write_all(&packet).await.unwrap(); + let mut plaintext = [0; 5]; + let result = client.read_exact(&mut plaintext).await; + if attempt == 1 { + result.unwrap(); + assert_eq!(&plaintext, b"hello"); + } else { + assert!(result.is_err()); + } + } + } +} diff --git a/src/shadowsocks/timed_salt_checker.rs b/src/shadowsocks/timed_salt_checker.rs index da51ec0c..9ceac899 100644 --- a/src/shadowsocks/timed_salt_checker.rs +++ b/src/shadowsocks/timed_salt_checker.rs @@ -104,6 +104,16 @@ impl SaltChecker for TimedSaltChecker { mod tests { use super::*; + #[tokio::test(start_paused = true)] + async fn ss2022_future_timestamp_is_protected_at_its_last_valid_second() { + let mut checker = TimedSaltChecker::new(61); + assert!(checker.insert_and_check(b"future salt")); + tokio::time::advance(Duration::from_secs(60)).await; + assert!(!checker.insert_and_check(b"future salt")); + tokio::time::advance(Duration::from_secs(1)).await; + assert!(checker.insert_and_check(b"future salt")); + } + #[tokio::test(start_paused = true)] async fn capacity_is_replay_safe_and_idle_entries_expire() { let mut checker = TimedSaltChecker::with_capacity(60, 2); diff --git a/src/snell/snell_handler.rs b/src/snell/snell_handler.rs index 659c176a..ab6fa44a 100644 --- a/src/snell/snell_handler.rs +++ b/src/snell/snell_handler.rs @@ -1,15 +1,17 @@ +use std::pin::Pin; use std::sync::Arc; +use std::task::{Context, Poll}; use argon2::{Config as Argon2Config, ThreadMode, Variant, Version}; use async_trait::async_trait; use log::debug; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWriteExt, ReadBuf}; use super::snell_fixed_target_stream::SnellFixedTargetStream; use super::snell_udp_stream::{SnellUdpClientStream, SnellUdpStream}; use crate::address::{Address, NetLocation, ResolvedLocation}; -use crate::async_stream::AsyncMessageStream; use crate::async_stream::AsyncStream; +use crate::async_stream::{AsyncMessageStream, AsyncReadMessage}; use crate::client_proxy_selector::ClientProxySelector; use crate::h2mux::{MUX_DESTINATION_HOST, MUX_DESTINATION_PORT, handle_h2mux_session}; use crate::resolver::Resolver; @@ -22,6 +24,21 @@ use crate::tcp::tcp_handler::{ }; use crate::util::write_all; +const SNELL_SALT_LEN: usize = 16; + +// Handshake reads must not coalesce the following UDP packet into the header buffer. +struct HeaderReader<'a>(&'a mut ShadowsocksStream); + +impl AsyncRead for HeaderReader<'_> { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + Pin::new(&mut *self.get_mut().0).poll_read_message(cx, buf) + } +} + // Snell protocol Argon2 parameters // ref: https://github.com/icpz/open-snell/blob/master/components/aead/cipher.go#L48 const SNELL_ARGON2_CONFIG: Argon2Config<'static> = Argon2Config { @@ -107,21 +124,22 @@ impl TcpServerHandler for SnellServerHandler { server_stream, ShadowsocksStreamType::Aead, self.cipher.algorithm(), - self.cipher.salt_len(), + SNELL_SALT_LEN, self.key.clone(), None, ); let mut stream_reader = StreamReader::new_with_buffer_size(400); + let mut header_stream = HeaderReader(&mut server_stream); - let version = stream_reader.read_u8(&mut server_stream).await?; + let version = stream_reader.read_u8(&mut header_stream).await?; if version != 1 { return Err(std::io::Error::other(format!( "unexpected snell version: {version}" ))); } - let command_type = stream_reader.read_u8(&mut server_stream).await?; + let command_type = stream_reader.read_u8(&mut header_stream).await?; let is_udp = match command_type { 0 => { // Ping command @@ -148,18 +166,18 @@ impl TcpServerHandler for SnellServerHandler { } }; - let client_id_len = stream_reader.read_u8(&mut server_stream).await?; + let client_id_len = stream_reader.read_u8(&mut header_stream).await?; if client_id_len > 0 { stream_reader - .read_slice(&mut server_stream, client_id_len as usize) + .read_slice(&mut header_stream, client_id_len as usize) .await?; } if !is_udp { - let hostname_len = stream_reader.read_u8(&mut server_stream).await? as usize; + let hostname_len = stream_reader.read_u8(&mut header_stream).await? as usize; let hostname_and_port_bytes = stream_reader - .read_slice(&mut server_stream, hostname_len + 2) + .read_slice(&mut header_stream, hostname_len + 2) .await?; let hostname_str = match std::str::from_utf8(&hostname_and_port_bytes[0..hostname_len]) @@ -267,7 +285,7 @@ impl TcpClientHandler for SnellClientHandler { client_stream, ShadowsocksStreamType::Aead, self.cipher.algorithm(), - self.cipher.salt_len(), + SNELL_SALT_LEN, self.key.clone(), None, )); @@ -336,7 +354,7 @@ impl TcpClientHandler for SnellClientHandler { client_stream, ShadowsocksStreamType::Aead, self.cipher.algorithm(), - self.cipher.salt_len(), + SNELL_SALT_LEN, self.key.clone(), None, ); diff --git a/src/socket_util.rs b/src/socket_util.rs index 8ba9791e..f861b305 100644 --- a/src/socket_util.rs +++ b/src/socket_util.rs @@ -10,6 +10,90 @@ use std::path::Path; use socket2::{Domain, Protocol, SockAddr, Socket, Type}; +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn hostname_socket_fallback_is_limited_to_family_errors() { + #[cfg(unix)] + let unavailable = libc::EAFNOSUPPORT; + #[cfg(windows)] + let unavailable = 10047; + let mut attempts = Vec::new(); + let socket = prefer_ipv6_socket(|ipv6| { + attempts.push(ipv6); + if ipv6 { + Err(std::io::Error::from_raw_os_error(unavailable)) + } else { + new_socket2_udp_socket(false, None, Some("0.0.0.0:0".parse().unwrap()), false) + } + }) + .unwrap(); + assert_eq!(attempts, [true, false]); + assert!(socket.local_addr().unwrap().as_socket().unwrap().is_ipv4()); + + attempts.clear(); + let error = prefer_ipv6_socket(|ipv6| { + attempts.push(ipv6); + Err(std::io::Error::new( + std::io::ErrorKind::PermissionDenied, + "denied", + )) + }) + .unwrap_err(); + assert_eq!(attempts, [true]); + assert_eq!(error.kind(), std::io::ErrorKind::PermissionDenied); + } + + #[tokio::test] + async fn hostname_socket_prefers_dual_stack_and_enforces_interface_binding() { + let ipv6_available = + match new_socket2_udp_socket(true, None, Some("[::]:0".parse().unwrap()), false) { + Ok(_) => true, + Err(error) if ipv6_unavailable(&error) => false, + Err(error) => panic!("IPv6 socket probe failed: {error}"), + }; + let socket = new_hostname_udp_socket(None).unwrap(); + assert_eq!(socket.local_addr().unwrap().is_ipv6(), ipv6_available); + assert!(new_hostname_udp_socket(Some("shoes-missing-interface".into())).is_err()); + } + + #[tokio::test] + async fn ipv6_udp_listener_accepts_both_address_families() { + let socket = + match new_socket2_udp_socket(true, None, Some("[::]:0".parse().unwrap()), false) { + Ok(socket) => socket, + Err(error) if ipv6_unavailable(&error) => return, + Err(error) => panic!("IPv6 socket creation failed: {error}"), + }; + assert!(!socket.only_v6().unwrap()); + let socket = into_tokio_udp_socket(socket).unwrap(); + let port = socket.local_addr().unwrap().port(); + for (bind, destination) in [ + ("127.0.0.1:0", format!("127.0.0.1:{port}")), + ("[::1]:0", format!("[::1]:{port}")), + ] { + tokio::time::timeout(std::time::Duration::from_secs(2), async { + let peer = match tokio::net::UdpSocket::bind(bind).await { + Ok(peer) => peer, + Err(error) if bind == "[::1]:0" && ipv6_unavailable(&error) => return, + Err(error) => panic!("UDP peer bind failed: {error}"), + }; + peer.send_to(b"request", destination).await.unwrap(); + let mut buf = [0; 16]; + let (len, sender) = socket.recv_from(&mut buf).await.unwrap(); + assert_eq!(&buf[..len], b"request"); + socket.send_to(b"reply", sender).await.unwrap(); + let len = peer.recv(&mut buf).await.unwrap(); + assert_eq!(&buf[..len], b"reply"); + }) + .await + .unwrap(); + } + } +} + #[cfg(all(test, unix, feature = "ffi"))] mod protection_tests { use super::*; @@ -39,6 +123,7 @@ mod protection_tests { ))); assert!(new_tcp_socket(None, false).is_err()); assert!(new_udp_socket(false, None).is_err()); + assert!(new_hostname_udp_socket(None).is_err()); assert!( new_outbound_socket2_udp_socket(false, None, Some("0.0.0.0:0".parse().unwrap())) .is_err() @@ -63,6 +148,59 @@ pub fn new_udp_socket( into_tokio_udp_socket(socket) } +pub fn new_hostname_udp_socket( + bind_interface: Option, +) -> std::io::Result { + let socket = prefer_ipv6_socket(|ipv6| { + new_socket2_udp_socket(ipv6, None, Some(get_unspecified_socket_addr(ipv6)), false) + })?; + bind_udp_interface(&socket, bind_interface.as_deref())?; + #[cfg(any(target_os = "android", target_os = "ios", all(unix, feature = "ffi")))] + crate::tun::protect_socket(socket.as_raw_fd())?; + into_tokio_udp_socket(socket) +} + +fn prefer_ipv6_socket( + mut create: impl FnMut(bool) -> std::io::Result, +) -> std::io::Result { + match create(true) { + Err(error) if ipv6_unavailable(&error) => create(false), + result => result, + } +} + +fn ipv6_unavailable(error: &std::io::Error) -> bool { + #[cfg(unix)] + let codes = [ + libc::EAFNOSUPPORT, + libc::EPROTONOSUPPORT, + libc::ENOPROTOOPT, + libc::EADDRNOTAVAIL, + ]; + // WSAENOPROTOOPT, WSAEPROTONOSUPPORT, WSAEAFNOSUPPORT, WSAEADDRNOTAVAIL. + #[cfg(windows)] + let codes = [10042, 10043, 10047, 10049]; + error + .raw_os_error() + .is_some_and(|code| codes.contains(&code)) +} + +fn bind_udp_interface(socket: &Socket, interface: Option<&str>) -> std::io::Result<()> { + if let Some(interface) = interface { + #[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))] + socket.bind_device(Some(interface.as_bytes()))?; + #[cfg(not(any(target_os = "android", target_os = "fuchsia", target_os = "linux")))] + { + let _ = (socket, interface); + return Err(std::io::Error::new( + std::io::ErrorKind::Unsupported, + "binding UDP sockets to an interface is unsupported", + )); + } + } + Ok(()) +} + pub fn new_outbound_socket2_udp_socket( is_ipv6: bool, bind_interface: Option, @@ -101,6 +239,10 @@ pub fn new_socket2_udp_socket_with_buffer_size( let domain = if is_ipv6 { Domain::IPV6 } else { Domain::IPV4 }; let socket = Socket::new(domain, Type::DGRAM, Some(Protocol::UDP))?; + if is_ipv6 { + socket.set_only_v6(false)?; + } + socket.set_nonblocking(true)?; // Set socket buffer sizes if specified. @@ -116,17 +258,13 @@ pub fn new_socket2_udp_socket_with_buffer_size( socket.set_reuse_port(true)?; #[cfg(any(not(unix), target_os = "solaris", target_os = "illumos"))] - panic!("Cannot support reuse sockets"); + return Err(std::io::Error::new( + std::io::ErrorKind::Unsupported, + "UDP port reuse is unsupported on this platform", + )); } - if let Some(ref interface) = bind_interface { - #[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))] - socket.bind_device(Some(interface.as_bytes()))?; - - // This should be handled during config validation. - #[cfg(not(any(target_os = "android", target_os = "fuchsia", target_os = "linux")))] - panic!("Could not bind to device, unsupported platform.") - } + bind_udp_interface(&socket, bind_interface.as_deref())?; if let Some(bind_address) = bind_address { socket.bind(&SockAddr::from(bind_address))?; diff --git a/src/socks5_udp_relay.rs b/src/socks5_udp_relay.rs index 22713e12..2551aa33 100644 --- a/src/socks5_udp_relay.rs +++ b/src/socks5_udp_relay.rs @@ -187,7 +187,7 @@ pub fn encode_socks5_udp_packet(source: &NetLocation, payload: &[u8], packet: &m /// and sends it to the client. pub struct Socks5UdpRelayStream { socket: Arc, - /// The client's UDP address, learned from the first received packet. + /// The client's UDP address, learned from the first valid, non-empty packet. client_addr: Option, /// Receiver for incoming packets from the socket reader task. receiver: mpsc::Receiver<(Box<[u8]>, SocketAddr)>, @@ -251,22 +251,17 @@ impl AsyncReadTargetedMessage for Socks5UdpRelayStream { match Pin::new(&mut this.receiver).poll_recv(cx) { Poll::Ready(Some((packet, from_addr))) => { - // Learn/verify client address - if let Some(expected) = this.client_addr { - if from_addr != expected { - // Packet from unexpected source, ignore - log::debug!( - "SOCKS5 UDP relay: ignoring packet from {} (expected {})", - from_addr, - expected - ); - cx.waker().wake_by_ref(); - return Poll::Pending; - } - } else { - // Learn client address from first packet - log::debug!("SOCKS5 UDP relay: learned client address: {}", from_addr); - this.client_addr = Some(from_addr); + if let Some(expected) = this.client_addr + && from_addr != expected + { + // Packet from unexpected source, ignore + log::debug!( + "SOCKS5 UDP relay: ignoring packet from {} (expected {})", + from_addr, + expected + ); + cx.waker().wake_by_ref(); + return Poll::Pending; } match parse_socks5_udp_packet(&packet) { @@ -292,6 +287,7 @@ impl AsyncReadTargetedMessage for Socks5UdpRelayStream { ), ))); } + this.client_addr = Some(from_addr); buf.put_slice(payload); Poll::Ready(Ok(target)) } @@ -410,6 +406,47 @@ impl Drop for Socks5UdpRelayStream { mod tests { use super::*; + #[tokio::test] + async fn only_a_valid_nonempty_packet_can_claim_an_association() { + for invalid in [vec![0; 2], vec![0, 0, 0, 1, 127, 0, 0, 1, 0, 53]] { + let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let mut relay = Socks5UdpRelayStream::new(socket); + let (tx, rx) = mpsc::channel(8); + relay.receiver = rx; + let client: SocketAddr = "127.0.0.1:1234".parse().unwrap(); + let foreign: SocketAddr = "127.0.0.2:5678".parse().unwrap(); + let target = NetLocation::from_str("127.0.0.1:53", None).unwrap(); + let mut packet = Vec::new(); + encode_socks5_udp_packet(&target, b"valid", &mut packet); + tx.send((invalid.into_boxed_slice(), foreign)) + .await + .unwrap(); + tx.send((packet.clone().into_boxed_slice(), client)) + .await + .unwrap(); + let mut output = [0; 64]; + let mut output = ReadBuf::new(&mut output); + let read_target = std::future::poll_fn(|cx| { + Pin::new(&mut relay).poll_read_targeted_message(cx, &mut output) + }) + .await + .unwrap(); + assert_eq!(read_target, target); + assert_eq!(output.filled(), b"valid"); + assert_eq!(relay.client_addr, Some(client)); + tx.send((packet.clone().into_boxed_slice(), foreign)) + .await + .unwrap(); + let mut cx = Context::from_waker(std::task::Waker::noop()); + assert!( + Pin::new(&mut relay) + .poll_read_targeted_message(&mut cx, &mut output) + .is_pending() + ); + assert_eq!(relay.client_addr, Some(client)); + } + } + #[test] fn reply_storage_reuses_capacity_without_retaining_previous_packet() { let mut packet = Vec::with_capacity(1024); diff --git a/src/socks_handler.rs b/src/socks_handler.rs index 92e7741b..ebefff40 100644 --- a/src/socks_handler.rs +++ b/src/socks_handler.rs @@ -373,9 +373,8 @@ async fn handle_udp_associate( mut server_stream: Box, stream_reader: &mut StreamReader, ) -> std::io::Result { - // Read client's hint address (DST.ADDR:DST.PORT) - we ignore this per RFC + // NAT and some clients make this hint differ from the actual UDP source. let _client_hint = read_location(&mut server_stream, stream_reader).await?; - log::debug!("SOCKS5 UDP ASSOCIATE: client hint = {:?}", _client_hint); // Uses 2MB buffer to prevent packet drops during bursts. const UDP_BUFFER_SIZE: usize = 2 * 1024 * 1024; diff --git a/src/tcp/chain_builder.rs b/src/tcp/chain_builder.rs index 5c185df3..2ea0c369 100644 --- a/src/tcp/chain_builder.rs +++ b/src/tcp/chain_builder.rs @@ -16,41 +16,54 @@ use crate::tcp::socket_connector_impl::SocketConnectorImpl; /// Creates InitialHopEntry (socket + optional proxy paired) from hop 0. /// Creates ProxyConnectors for subsequent hops (1+). /// `protocol: direct` at hop 0 creates InitialHopEntry::Direct. +#[cfg(test)] pub fn build_client_proxy_chain( client_chain: crate::option_util::OneOrSome, resolver: Arc, ) -> ClientProxyChain { + try_build_client_proxy_chain(client_chain, resolver).expect("Failed to create client chain") +} + +pub fn try_build_client_proxy_chain( + client_chain: crate::option_util::OneOrSome, + resolver: Arc, +) -> std::io::Result { + build_client_proxy_chain_with_limits(client_chain, resolver, crate::resources::limits()) +} + +fn build_client_proxy_chain_with_limits( + client_chain: crate::option_util::OneOrSome, + resolver: Arc, + limits: crate::config::GlobalLimits, +) -> std::io::Result { let hops: Vec> = client_chain .into_vec() .into_iter() - .map(|hop| match hop { - ClientChainHop::Single(selection) => match selection { - ConfigSelection::Config(config) => vec![config], - ConfigSelection::GroupName(group_name) => { - panic!( - "Group reference '{}' was not resolved during config validation.", - group_name - ); - } - }, - ClientChainHop::Pool(selections) => selections - .into_vec() + .map(|hop| { + let selections = match hop { + ClientChainHop::Single(selection) => vec![selection], + ClientChainHop::Pool(selections) => selections.into_vec(), + }; + selections .into_iter() - .flat_map(|selection| match selection { - ConfigSelection::Config(config) => vec![config], - ConfigSelection::GroupName(group_name) => { - panic!( - "Group reference '{}' was not resolved during config validation.", - group_name - ); - } + .map(|selection| match selection { + ConfigSelection::Config(config) => Ok(config), + ConfigSelection::GroupName(group_name) => Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!( + "Group reference '{group_name}' was not resolved during config validation." + ), + )), }) - .collect(), + .collect::>() }) - .collect(); + .collect::>()?; if hops.is_empty() { - panic!("Client chain must have at least one hop"); + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "Client chain must have at least one hop", + )); } // Build initial hop entries from hop 0. @@ -61,22 +74,26 @@ pub fn build_client_proxy_chain( // Find the first proxy address for QUIC socket configuration let target_address = find_first_proxy_address(&hops, config); - let socket = SocketConnectorImpl::from_config(config, target_address) - .map(|s| Box::new(s) as Box) - .expect("Failed to create SocketConnector"); + let socket = + SocketConnectorImpl::from_config_with_limits(config, target_address, limits) + .map(|s| Box::new(s) as Box)?; if config.protocol.is_direct() { // Direct: socket only, no proxy - InitialHopEntry::Direct(socket) + Ok(InitialHopEntry::Direct(socket)) } else { // Proxy: socket + proxy paired - let proxy = ProxyConnectorImpl::from_config(config.clone(), resolver.clone()) + let proxy = ProxyConnectorImpl::from_config(config.clone(), resolver.clone())? .map(|p| Box::new(p) as Box) - .expect("Failed to create ProxyConnector for non-direct config"); - InitialHopEntry::Proxy { socket, proxy } + .ok_or_else(|| { + std::io::Error::other( + "Missing proxy connector for non-direct configuration", + ) + })?; + Ok(InitialHopEntry::Proxy { socket, proxy }) } }) - .collect(); + .collect::>()?; // Build proxy connectors for subsequent hops (1+) let subsequent_hops: Vec>> = hops @@ -90,22 +107,24 @@ pub fn build_client_proxy_chain( .map(|config| { // Subsequent hops MUST NOT have direct protocol if config.protocol.is_direct() { - panic!( - "protocol: direct is only valid at hop 0. Found direct at hop {} with address {}", - hop_index, - config.address - ); + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!( + "protocol: direct is only valid at hop 0. Found direct at hop {} with address {}", + hop_index, config.address + ), + )); } - ProxyConnectorImpl::from_config(config, resolver.clone()) + ProxyConnectorImpl::from_config(config, resolver.clone())? .map(|p| Box::new(p) as Box) - .expect("Failed to create ProxyConnector for subsequent hop") + .ok_or_else(|| std::io::Error::other("Missing proxy connector for subsequent hop")) }) - .collect() + .collect::>() }) - .collect(); + .collect::>()?; - ClientProxyChain::new(initial_hop, subsequent_hops) + Ok(ClientProxyChain::new(initial_hop, subsequent_hops)) } /// Find the first proxy address in the chain (for socket connector target). @@ -141,22 +160,39 @@ pub fn build_client_chain_group( client_chains: crate::option_util::NoneOrSome, resolver: Arc, ) -> ClientChainGroup { + try_build_client_chain_group(client_chains, resolver) + .expect("Failed to create client chain group") +} + +pub fn try_build_client_chain_group( + client_chains: crate::option_util::NoneOrSome, + resolver: Arc, +) -> std::io::Result { + build_client_chain_group_with_limits(client_chains, resolver, crate::resources::limits()) +} + +pub fn build_client_chain_group_with_limits( + client_chains: crate::option_util::NoneOrSome, + resolver: Arc, + limits: crate::config::GlobalLimits, +) -> std::io::Result { let chains: Vec = if client_chains.is_empty() { - vec![build_client_proxy_chain( + vec![build_client_proxy_chain_with_limits( crate::option_util::OneOrSome::One(ClientChainHop::Single(ConfigSelection::Config( ClientConfig::default(), ))), resolver, - )] + limits, + )?] } else { client_chains .into_vec() .into_iter() - .map(|chain| build_client_proxy_chain(chain.hops, resolver.clone())) - .collect() + .map(|chain| build_client_proxy_chain_with_limits(chain.hops, resolver.clone(), limits)) + .collect::>()? }; - ClientChainGroup::new(chains) + Ok(ClientChainGroup::new(chains)) } #[cfg(test)] @@ -168,6 +204,17 @@ mod tests { use crate::resolver::NativeResolver; use std::net::{IpAddr, Ipv4Addr}; + #[cfg(target_os = "linux")] + #[tokio::test] + async fn socket_initialization_failure_is_returned() { + let config: ClientConfig = serde_yaml::from_str("address: '127.0.0.1:443'\ntransport: quic\nbind_interface: shoes-missing-interface\nprotocol: {type: socks}\n").unwrap(); + let result = try_build_client_proxy_chain( + OneOrSome::One(ClientChainHop::Single(ConfigSelection::Config(config))), + mock_resolver(), + ); + assert!(result.is_err()); + } + fn mock_resolver() -> Arc { Arc::new(NativeResolver::new()) } @@ -284,21 +331,23 @@ mod tests { } #[test] - #[should_panic(expected = "protocol: direct is only valid at hop 0")] - fn test_direct_at_hop1_panics() { - build_client_proxy_chain( + fn test_direct_at_hop1_returns_error() { + let result = try_build_client_proxy_chain( OneOrSome::Some(vec![ ClientChainHop::Single(ConfigSelection::Config(socks_config(1080))), ClientChainHop::Single(ConfigSelection::Config(direct_config())), ]), mock_resolver(), ); + assert_eq!( + result.err().unwrap().kind(), + std::io::ErrorKind::InvalidInput + ); } #[test] - #[should_panic(expected = "protocol: direct is only valid at hop 0")] - fn test_direct_in_pool_at_hop1_panics() { - build_client_proxy_chain( + fn test_direct_in_pool_at_hop1_returns_error() { + let result = try_build_client_proxy_chain( OneOrSome::Some(vec![ ClientChainHop::Single(ConfigSelection::Config(socks_config(1080))), ClientChainHop::Pool(OneOrSome::Some(vec![ @@ -308,19 +357,36 @@ mod tests { ]), mock_resolver(), ); + assert_eq!( + result.err().unwrap().kind(), + std::io::ErrorKind::InvalidInput + ); } #[test] - #[should_panic(expected = "was not resolved during config validation")] - fn test_unresolved_group_reference_panics() { - build_client_proxy_chain( - OneOrSome::One(ClientChainHop::Single(ConfigSelection::GroupName( - "unresolved_group".to_string(), - ))), - mock_resolver(), + fn test_empty_chain_returns_error() { + let result = try_build_client_proxy_chain(OneOrSome::Some(vec![]), mock_resolver()); + assert_eq!( + result.err().unwrap().kind(), + std::io::ErrorKind::InvalidInput ); } + #[test] + fn test_unresolved_group_reference_returns_error() { + let selection = ConfigSelection::GroupName("unresolved_group".to_string()); + for hop in [ + ClientChainHop::Single(selection.clone()), + ClientChainHop::Pool(OneOrSome::One(selection)), + ] { + let result = try_build_client_proxy_chain(OneOrSome::One(hop), mock_resolver()); + assert_eq!( + result.err().unwrap().kind(), + std::io::ErrorKind::InvalidInput + ); + } + } + #[test] fn test_find_first_proxy_address_direct_only() { let direct = direct_config(); diff --git a/src/tcp/proxy_connector_impl.rs b/src/tcp/proxy_connector_impl.rs index 862e062f..cd62288b 100644 --- a/src/tcp/proxy_connector_impl.rs +++ b/src/tcp/proxy_connector_impl.rs @@ -34,21 +34,24 @@ impl ProxyConnectorImpl { /// Create a ProxyConnector from a ClientConfig's protocol-related fields. /// /// Returns None for direct protocol (direct has no ProxyConnector). - pub fn from_config(config: ClientConfig, resolver: Arc) -> Option { + pub fn from_config( + config: ClientConfig, + resolver: Arc, + ) -> std::io::Result> { if config.protocol.is_direct() { - return None; + return Ok(None); } let default_sni_hostname = config.address.address().hostname().map(ToString::to_string); - Some(Self { + Ok(Some(Self { location: config.address, client_handler: create_tcp_client_handler( config.protocol, default_sni_hostname, resolver, - ), - }) + )?, + })) } } @@ -241,11 +244,36 @@ mod tests { let _ = peer.await; } + #[test] + fn malformed_protocol_settings_return_errors_without_validation() { + for protocol in [ + serde_json::json!({"type": "reality", "public_key": "invalid", "sni_hostname": "example.com", "protocol": {"type": "socks"}}), + serde_json::json!({"type": "reality", "public_key": "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA", "protocol": {"type": "socks"}}), + serde_json::json!({"type": "anytls", "password": "test", "padding_scheme": ["stop=invalid"]}), + serde_json::json!({"type": "reality", "public_key": "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA", "sni_hostname": "127.0.0.1", "protocol": {"type": "socks"}}), + serde_json::json!({"type": "reality", "public_key": "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA", "sni_hostname": "::1", "protocol": {"type": "socks"}}), + serde_json::json!({"type": "reality", "public_key": "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA", "sni_hostname": "example.com", "protocol": {"type": "direct"}}), + serde_json::json!({"type": "tls", "protocol": {"type": "direct"}}), + serde_json::json!({"type": "shadowtls", "password": "test", "protocol": {"type": "direct"}}), + serde_json::json!({"type": "websocket", "protocol": {"type": "direct"}}), + ] { + let config = serde_json::from_value(serde_json::json!({ + "address": "127.0.0.1:443", "protocol": protocol + })) + .unwrap(); + assert!(ProxyConnectorImpl::from_config(config, mock_resolver()).is_err()); + } + } + #[test] fn test_from_direct_config_returns_none() { let config = ClientConfig::default(); assert!(config.protocol.is_direct()); - assert!(ProxyConnectorImpl::from_config(config, mock_resolver()).is_none()); + assert!( + ProxyConnectorImpl::from_config(config, mock_resolver()) + .unwrap() + .is_none() + ); } #[test] @@ -258,7 +286,7 @@ mod tests { }, ..Default::default() }; - let connector = ProxyConnectorImpl::from_config(config, mock_resolver()); + let connector = ProxyConnectorImpl::from_config(config, mock_resolver()).unwrap(); assert!(connector.is_some()); let connector = connector.unwrap(); assert_eq!(connector.proxy_location().port(), 1080); diff --git a/src/tcp/socket_connector_impl.rs b/src/tcp/socket_connector_impl.rs index 68a472b3..e74f84ae 100644 --- a/src/tcp/socket_connector_impl.rs +++ b/src/tcp/socket_connector_impl.rs @@ -21,7 +21,7 @@ use crate::config::{ClientConfig, ClientQuicConfig, Transport}; use crate::quic_endpoint::QuicEndpoint; use crate::quic_stream::QuicStream; use crate::resolver::{Resolver, resolve_addresses, resolve_location}; -use crate::rustls_config_util::create_client_config; +use crate::rustls_config_util::try_create_client_config; use crate::socket_util::{new_tcp_socket, new_udp_socket, set_tcp_keepalive}; use crate::thread_util::get_num_threads; @@ -140,11 +140,20 @@ impl SocketConnectorImpl { /// Pass None for direct protocol (QUIC is not supported for direct). /// /// # Returns - /// None if QUIC endpoint creation fails. + /// Returns any socket or QUIC endpoint initialization error. + #[cfg(test)] pub fn from_config( config: &ClientConfig, target_address: Option<&NetLocation>, - ) -> Option { + ) -> std::io::Result { + Self::from_config_with_limits(config, target_address, crate::resources::limits()) + } + + pub fn from_config_with_limits( + config: &ClientConfig, + target_address: Option<&NetLocation>, + limits: crate::config::GlobalLimits, + ) -> std::io::Result { let bind_interface = config.bind_interface.clone().into_option(); let default_sni_hostname = @@ -168,9 +177,12 @@ impl SocketConnectorImpl { } Transport::Quic => { // QUIC requires a target address for endpoint creation - let target_address = target_address.expect( - "QUIC transport requires target_address (direct protocol should use TCP)", - ); + let target_address = target_address.ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "QUIC transport requires a target address", + ) + })?; let ClientQuicConfig { verify, @@ -207,27 +219,29 @@ impl SocketConnectorImpl { (key_bytes, cert_bytes) }); - let rustls_client_config = create_client_config( + let rustls_client_config = try_create_client_config( verify, server_fingerprints.into_vec(), alpn_protocols.into_vec(), sni_hostname.is_some(), key_and_cert_bytes, false, // tls13_only - QUIC enforces TLS 1.3 anyway - ); + )?; let quic_client_config = quinn::crypto::rustls::QuicClientConfig::with_initial( Arc::new(rustls_client_config), tls13_suite.quic_suite().unwrap(), ) - .unwrap(); + .map_err(std::io::Error::other)?; let mut quinn_client_config = quinn::ClientConfig::new(Arc::new(quic_client_config)); let mut transport_config = quinn::TransportConfig::default(); - let memory_bytes = - crate::resources::configure_quic_transport(&mut transport_config); + let memory_bytes = crate::resources::configure_quic_transport_with_limits( + &mut transport_config, + limits, + ); transport_config .max_concurrent_bidi_streams(0_u32.into()) .max_concurrent_uni_streams(0_u8.into()) @@ -240,19 +254,14 @@ impl SocketConnectorImpl { let mut endpoints = Vec::with_capacity(endpoints_len); for _ in 0..endpoints_len { - let udp_socket = match new_udp_socket( - target_address.address().is_ipv6(), - bind_interface.clone(), - ) { - Ok(s) => s, - Err(e) => { - error!("Failed to bind new UDP socket for QUIC: {e}"); - return None; - } + let udp_socket = if target_address.address().hostname().is_some() { + crate::socket_util::new_hostname_udp_socket(bind_interface.clone())? + } else { + new_udp_socket(target_address.address().is_ipv6(), bind_interface.clone())? }; - let udp_socket = udp_socket.into_std().unwrap(); + let udp_socket = udp_socket.into_std()?; - let mut endpoint = QuicEndpoint::new(None, udp_socket, memory_bytes).unwrap(); + let mut endpoint = QuicEndpoint::new(None, udp_socket, memory_bytes)?; endpoint.set_default_client_config(quinn_client_config.clone()); endpoints.push(Arc::new(endpoint)); } @@ -265,7 +274,7 @@ impl SocketConnectorImpl { } }; - Some(Self { + Ok(Self { bind_interface, transport, }) @@ -290,7 +299,7 @@ impl SocketConnector for SocketConnectorImpl { resolver: &Arc, address: &ResolvedLocation, ) -> std::io::Result> { - let target_addrs = match address.resolved_addr() { + let mut target_addrs = match address.resolved_addr() { Some(r) => vec![r], None => resolve_addresses(resolver, address.location()).await?, }; @@ -319,6 +328,10 @@ impl SocketConnector for SocketConnectorImpl { next_endpoint_index, sni_hostname, } => { + // Preserve IPv4 preference for hostnames without excluding IPv6-only targets. + if address.address().hostname().is_some() { + target_addrs.sort_by_key(SocketAddr::is_ipv6); + } let domain = match sni_hostname { Some(s) => s.as_str(), None => address.address().hostname().unwrap_or("example.com"), @@ -493,6 +506,75 @@ mod tests { use std::time::Duration; use tokio::time::Instant; + #[derive(Debug)] + struct FixedResolver(Vec); + + impl Resolver for FixedResolver { + fn resolve_location( + &self, + _location: &NetLocation, + ) -> Pin>> + Send>> + { + let addresses = self.0.clone(); + Box::pin(async move { Ok(addresses) }) + } + } + + #[tokio::test] + async fn hostname_quic_prefers_ipv4_without_waiting_for_stalled_ipv6() { + use tokio::io::AsyncWriteExt; + let Ok(blackhole) = UdpSocket::bind("[::1]:0").await else { + return; + }; + let stalled = blackhole.local_addr().unwrap(); + let certificate = rcgen::generate_simple_self_signed(vec!["localhost".into()]).unwrap(); + let server_config = quinn::ServerConfig::with_single_cert( + vec![certificate.cert.der().clone()], + rustls::pki_types::PrivatePkcs8KeyDer::from(certificate.signing_key.serialize_der()) + .into(), + ) + .unwrap(); + for bind in ["127.0.0.1:0", "[::1]:0"] { + let server = + quinn::Endpoint::server(server_config.clone(), bind.parse().unwrap()).unwrap(); + let server_addr = server.local_addr().unwrap(); + let addresses = if server_addr.is_ipv4() { + vec![stalled, server_addr] + } else { + vec![server_addr] + }; + let resolver: Arc = Arc::new(FixedResolver(addresses)); + let config: ClientConfig = serde_yaml::from_str("address: 'localhost:443'\ntransport: quic\nquic_settings: {verify: false}\nprotocol: {type: socks}\n").unwrap(); + let connector = + SocketConnectorImpl::from_config(&config, Some(&config.address)).unwrap(); + tokio::time::timeout(Duration::from_secs(3), async { + let (client, received) = tokio::join!( + async { + let mut stream = connector + .connect(&resolver, &ResolvedLocation::new(config.address.clone())) + .await + .unwrap(); + stream.write_all(b"ping").await.unwrap(); + stream.flush().await.unwrap(); + stream + }, + async { + let connection = server.accept().await.unwrap().await.unwrap(); + let (_send, mut recv) = connection.accept_bi().await.unwrap(); + let mut data = [0; 4]; + recv.read_exact(&mut data).await.unwrap(); + assert_eq!(&data, b"ping"); + connection + }, + ); + drop((client, received)); + }) + .await + .unwrap(); + server.close(0u32.into(), b"done"); + } + } + fn addresses(count: u16) -> Vec { (1..=count) .map(|port| SocketAddr::from(([127, 0, 0, 1], port))) @@ -693,7 +775,7 @@ mod tests { fn test_from_config_direct_protocol() { let config = ClientConfig::default(); // default is direct protocol let connector = SocketConnectorImpl::from_config(&config, None); - assert!(connector.is_some()); + assert!(connector.is_ok()); assert!(matches!( connector.unwrap().transport, TransportConfig::Tcp { .. } diff --git a/src/tcp/tcp_client_handler_factory.rs b/src/tcp/tcp_client_handler_factory.rs index 4c9029f3..43602974 100644 --- a/src/tcp/tcp_client_handler_factory.rs +++ b/src/tcp/tcp_client_handler_factory.rs @@ -15,12 +15,12 @@ use crate::http_handler::HttpTcpClientHandler; use crate::naiveproxy::NaiveProxyTcpClientHandler; use crate::port_forward_handler::PortForwardClientHandler; use crate::resolver::Resolver; -use crate::rustls_config_util::create_client_config; +use crate::rustls_config_util::try_create_client_config; use crate::shadow_tls::ShadowTlsClientHandler; use crate::shadowsocks::ShadowsocksTcpHandler; use crate::snell::snell_handler::SnellClientHandler; use crate::socks_handler::SocksTcpClientHandler; -use crate::tcp::chain_builder::build_client_chain_group; +use crate::tcp::chain_builder::try_build_client_chain_group; use crate::tcp::tcp_handler::TcpClientHandler; use crate::tls_client_handler::TlsClientHandler; use crate::trojan_handler::TrojanTcpHandler; @@ -43,10 +43,13 @@ pub fn create_tcp_client_handler( client_proxy_config: ClientProxyConfig, default_sni_hostname: Option, resolver: Arc, -) -> Box { - match client_proxy_config { +) -> std::io::Result> { + Ok(match client_proxy_config { ClientProxyConfig::Direct => { - panic!("Tried to create a direct tcp client handler"); + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "direct cannot be used as an inner proxy protocol", + )); } ClientProxyConfig::Http { username, @@ -154,17 +157,19 @@ pub fn create_tcp_client_handler( (key_bytes, cert_bytes) }); - let client_config = Arc::new(create_client_config( + let client_config = Arc::new(try_create_client_config( verify, server_fingerprints.into_vec(), alpn_protocols.into_vec(), sni_hostname.is_some(), key_and_cert_bytes, false, // tls13_only - )); + )?); let server_name = match sni_hostname { - Some(s) => rustls::pki_types::ServerName::try_from(s).unwrap(), + Some(s) => { + rustls::pki_types::ServerName::try_from(s).map_err(std::io::Error::other)? + } // This is unused, since enable_sni is false, but connect_with still requires a // parameter. None => "example.com".try_into().unwrap(), @@ -191,7 +196,7 @@ pub fn create_tcp_client_handler( *udp_enabled, )) } else { - let handler = create_tcp_client_handler(*protocol, None, resolver.clone()); + let handler = create_tcp_client_handler(*protocol, None, resolver.clone())?; Box::new(TlsClientHandler::new( client_config, @@ -209,24 +214,17 @@ pub fn create_tcp_client_handler( vision, protocol, } => { - // Decode public key from base64url let public_key_bytes = - crate::reality::decode_public_key(&public_key).expect("Invalid REALITY public key"); - - // Decode short ID from hex string + crate::reality::decode_public_key(&public_key).map_err(std::io::Error::other)?; let short_id_bytes = - crate::reality::decode_short_id(&short_id).expect("Invalid REALITY short_id"); - - // Determine SNI hostname - let sni_hostname = sni_hostname.or(default_sni_hostname.clone()); - let server_name = match sni_hostname { - Some(s) => rustls::pki_types::ServerName::try_from(s) - .unwrap() - .to_owned(), - None => { - panic!("REALITY client requires sni_hostname to be specified"); - } + crate::reality::decode_short_id(&short_id).map_err(std::io::Error::other)?; + let Some(sni_hostname) = sni_hostname.or(default_sni_hostname) else { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "REALITY client requires sni_hostname", + )); }; + let server_name = crate::reality::parse_server_name(&sni_hostname)?; let cipher_suites = cipher_suites.into_vec(); @@ -253,7 +251,7 @@ pub fn create_tcp_client_handler( ), ) } else { - let inner_handler = create_tcp_client_handler(*protocol, None, resolver.clone()); + let inner_handler = create_tcp_client_handler(*protocol, None, resolver.clone())?; Box::new(crate::reality_client_handler::RealityClientHandler::new( public_key_bytes, short_id_bytes, @@ -272,7 +270,9 @@ pub fn create_tcp_client_handler( let enable_sni = sni_hostname.is_some(); let server_name = match sni_hostname { - Some(s) => rustls::pki_types::ServerName::try_from(s).unwrap(), + Some(s) => { + rustls::pki_types::ServerName::try_from(s).map_err(std::io::Error::other)? + } None => "example.com".try_into().unwrap(), // Fallback }; @@ -281,16 +281,16 @@ pub fn create_tcp_client_handler( // an HMAC tag, and rustls doesn't validate session_id echo for TLS 1.3 ServerHello. // TLS 1.2 would fail anyway (no supported_versions extension), but restricting // here provides defense in depth and fails fast at the TLS level. - let client_config = Arc::new(create_client_config( + let client_config = Arc::new(try_create_client_config( false, // No WebPKI verification needed for ShadowTLS Vec::new(), // No fingerprints Vec::new(), // No ALPN enable_sni, // Enable SNI if hostname provided None, // No client cert true, // tls13_only - required for ShadowTLS v3 - )); + )?); - let handler = create_tcp_client_handler(*protocol, None, resolver.clone()); + let handler = create_tcp_client_handler(*protocol, None, resolver.clone())?; Box::new(ShadowTlsClientHandler::new( password, @@ -306,7 +306,7 @@ pub fn create_tcp_client_handler( h2mux, } => { let handler: Box = - Box::new(VmessTcpClientHandler::new(&cipher, &user_id, udp_enabled)); + Box::new(VmessTcpClientHandler::new(&cipher, &user_id, udp_enabled)?); if let Some(h2mux_config) = h2mux { Box::new(H2MuxClientHandler::new( Arc::from(handler), @@ -324,7 +324,7 @@ pub fn create_tcp_client_handler( protocol, } = websocket_client_config; - let handler = create_tcp_client_handler(*protocol, None, resolver.clone()); + let handler = create_tcp_client_handler(*protocol, None, resolver.clone())?; Box::new(WebsocketTcpClientHandler::new( matching_path, @@ -342,10 +342,7 @@ pub fn create_tcp_client_handler( let padding = match padding_scheme { Some(lines) => { let scheme = lines.join("\n"); - Arc::new( - PaddingFactory::new(scheme.as_bytes()) - .expect("Invalid padding scheme in AnyTLS config"), - ) + Arc::new(PaddingFactory::new(scheme.as_bytes()).map_err(std::io::Error::other)?) } None => PaddingFactory::default_factory(), }; @@ -358,13 +355,13 @@ pub fn create_tcp_client_handler( } => Box::new(NaiveProxyTcpClientHandler::new( &username, &password, padding, )), - } + }) } pub fn create_tcp_client_proxy_selector( rules: Vec, resolver: Arc, -) -> ClientProxySelector { +) -> std::io::Result { let rules = rules .into_iter() .map(|rule_config| { @@ -374,13 +371,14 @@ pub fn create_tcp_client_proxy_selector( override_address, client_chains, } => { - let chain_group = build_client_chain_group(client_chains, resolver.clone()); + let chain_group = + try_build_client_chain_group(client_chains, resolver.clone())?; ConnectAction::new_allow(override_address, chain_group) } RuleActionConfig::Block => ConnectAction::new_block(), }; - ConnectRule::new(masks.into_vec(), connect_action) + Ok(ConnectRule::new(masks.into_vec(), connect_action)) }) - .collect::>(); - ClientProxySelector::new(rules) + .collect::>>()?; + Ok(ClientProxySelector::new(rules)) } diff --git a/src/tcp/tcp_server.rs b/src/tcp/tcp_server.rs index 518a448e..4d73054b 100644 --- a/src/tcp/tcp_server.rs +++ b/src/tcp/tcp_server.rs @@ -1,6 +1,5 @@ use std::collections::HashMap; -use std::net::{IpAddr, SocketAddr}; -use std::path::PathBuf; +use std::net::IpAddr; use std::sync::Arc; use std::time::Duration; @@ -28,14 +27,13 @@ use crate::tun::start_tun_server; use crate::util::write_all; async fn run_tcp_server( - bind_address: SocketAddr, + listener: tokio::net::TcpListener, tcp_config: TcpConfig, resolver: Arc, server_handler: Arc, ) -> std::io::Result<()> { let TcpConfig { no_delay } = tcp_config; - let listener = new_tcp_listener(bind_address, 4096, None)?; let mut tasks = crate::listener_tasks::ListenerTasks::new(); loop { @@ -82,19 +80,10 @@ async fn run_tcp_server( #[cfg(target_family = "unix")] async fn run_unix_server( - path_buf: PathBuf, + listener: tokio::net::UnixListener, resolver: Arc, server_handler: Arc, ) -> std::io::Result<()> { - if tokio::fs::symlink_metadata(&path_buf).await.is_ok() { - println!( - "WARNING: replacing file at socket path {}", - path_buf.display() - ); - let _ = tokio::fs::remove_file(&path_buf).await; - } - - let listener = crate::socket_util::new_unix_listener(path_buf, 4096)?; let mut tasks = crate::listener_tasks::ListenerTasks::new(); loop { @@ -179,15 +168,13 @@ where } => { let setup_client_stream_future = timeout( Duration::from_secs(60), - setup_client_tcp_stream( - &mut server_stream, - proxy_selector, - resolver, - remote_location.clone(), - ), + setup_client_tcp_stream(proxy_selector, resolver, remote_location.clone()), ); - let mut client_stream = match setup_client_stream_future.await { + let TcpClientSetupResult { + mut client_stream, + early_data, + } = match setup_client_stream_future.await { Ok(Ok(Some(s))) => s, Ok(Ok(None)) => { // Must have been blocked. @@ -210,27 +197,27 @@ where } }; - if let Some(data) = connection_success_response { - write_all(&mut server_stream, &data).await?; - // server_need_initial_flush should be set to true by the handler if - // it's needed. - } - - let client_need_initial_flush = match initial_remote_data { - Some(data) => { - write_all(&mut client_stream, &data).await?; - true + crate::util::timeout_stream_setup(async { + let flush_server = server_need_initial_flush || early_data.is_some(); + if let Some(data) = connection_success_response { + write_all(&mut server_stream, &data).await?; } - None => false, - }; + if let Some(data) = early_data { + write_all(&mut server_stream, &data).await?; + } + if flush_server { + server_stream.flush().await?; + } + if let Some(data) = initial_remote_data { + write_all(&mut client_stream, &data).await?; + client_stream.flush().await?; + } + Ok(()) + }) + .await?; - let copy_result = copy_bidirectional( - &mut server_stream, - &mut client_stream, - server_need_initial_flush, - client_need_initial_flush, - ) - .await; + let copy_result = + copy_bidirectional(&mut server_stream, &mut client_stream, false, false).await; futures::join!( crate::util::shutdown_stream(&mut server_stream), @@ -317,11 +304,10 @@ where } pub async fn setup_client_tcp_stream( - server_stream: &mut Box, client_proxy_selector: Arc, resolver: Arc, remote_location: NetLocation, -) -> std::io::Result>> { +) -> std::io::Result> { let action = client_proxy_selector .judge(remote_location.into(), &resolver) .await?; @@ -330,19 +316,10 @@ pub async fn setup_client_tcp_stream( ConnectDecision::Allow { chain_group, remote_location, - } => { - let TcpClientSetupResult { - client_stream, - early_data, - } = chain_group.connect_tcp(remote_location, &resolver).await?; - - if let Some(data) = early_data { - server_stream.write_all(&data).await?; - server_stream.flush().await?; - } - - Ok(Some(client_stream)) - } + } => chain_group + .connect_tcp(remote_location, &resolver) + .await + .map(Some), ConnectDecision::Block => Ok(None), } } @@ -398,38 +375,17 @@ async fn start_tcp_or_quic_servers( config: ServerConfig, resolver: Arc, ) -> std::io::Result>> { - let mut join_handles = Vec::with_capacity(3); - - match config.transport { - Transport::Tcp => match start_tcp_servers(config.clone(), resolver).await { - Ok(handles) => { - join_handles.extend(handles); - } - Err(e) => { - for join_handle in join_handles { - join_handle.abort(); - } - return Err(e); - } - }, - Transport::Quic => match start_quic_servers(config.clone(), resolver).await { - Ok(handles) => { - join_handles.extend(handles); - } - Err(e) => { - for join_handle in join_handles { - join_handle.abort(); - } - return Err(e); - } - }, + let bind_location = config.bind_location.to_string(); + let join_handles = match config.transport { + Transport::Tcp => start_tcp_servers(config, resolver).await?, + Transport::Quic => start_quic_servers(config, resolver).await?, Transport::Udp => todo!(), - } + }; if join_handles.is_empty() { return Err(std::io::Error::other(format!( "failed to start servers at {}", - config.bind_location + bind_location ))); } @@ -459,7 +415,7 @@ async fn start_tcp_servers( let client_proxy_selector = Arc::new(create_tcp_client_proxy_selector( rules.clone(), resolver.clone(), - )); + )?); let mut handles = vec![]; @@ -467,41 +423,56 @@ async fn start_tcp_servers( BindLocation::Address(addresses) => { // Shares protocol state across ports without reusing an interface-specific UDP bind IP. let mut handlers: HashMap> = HashMap::new(); + let mut listeners = Vec::new(); for address in addresses.into_vec() { for socket_addr in address.to_socket_addrs()? { - let tcp_handler = handlers - .entry(socket_addr.ip()) - .or_insert_with(|| { + let listener = new_tcp_listener(socket_addr, 4096, None)?; + let tcp_handler = match handlers.entry(socket_addr.ip()) { + std::collections::hash_map::Entry::Occupied(entry) => entry.into_mut(), + std::collections::hash_map::Entry::Vacant(entry) => entry.insert( create_tcp_server_handler( protocol.clone(), &client_proxy_selector, &resolver, Some(socket_addr.ip()), - ) - .into() - }) - .clone(); - let tcp_config = tcp_config.clone(); - let resolver = resolver.clone(); - let handle = tokio::spawn(async move { - run_tcp_server(socket_addr, tcp_config, resolver, tcp_handler) - .await - .unwrap(); - }); - handles.push(handle); + )? + .into(), + ), + } + .clone(); + listeners.push((listener, tcp_handler)); } } + for (listener, tcp_handler) in listeners { + let tcp_config = tcp_config.clone(); + let resolver = resolver.clone(); + handles.push(tokio::spawn(async move { + if let Err(error) = + run_tcp_server(listener, tcp_config, resolver, tcp_handler).await + { + error!("TCP listener stopped: {error}"); + } + })); + } } BindLocation::Path(path_buf) => { #[cfg(target_family = "unix")] { + if tokio::fs::symlink_metadata(&path_buf).await.is_ok() { + println!( + "WARNING: replacing file at socket path {}", + path_buf.display() + ); + tokio::fs::remove_file(&path_buf).await?; + } + let listener = crate::socket_util::new_unix_listener(path_buf, 4096)?; let tcp_handler: Arc = - create_tcp_server_handler(protocol, &client_proxy_selector, &resolver, None) + create_tcp_server_handler(protocol, &client_proxy_selector, &resolver, None)? .into(); let handle = tokio::spawn(async move { - run_unix_server(path_buf, resolver, tcp_handler) - .await - .unwrap(); + if let Err(error) = run_unix_server(listener, resolver, tcp_handler).await { + error!("Unix listener stopped: {error}"); + } }); handles.push(handle); } @@ -523,6 +494,7 @@ mod lifetime_tests { use super::*; use async_trait::async_trait; use std::io; + use std::net::SocketAddr; use std::pin::Pin; use std::sync::atomic::{AtomicUsize, Ordering}; use std::task::{Context, Poll}; @@ -535,6 +507,224 @@ mod lifetime_tests { use crate::client_proxy_selector::{ConnectAction, ConnectRule}; use crate::tcp::socket_connector::SocketConnector; + #[derive(Debug)] + struct Handler(Arc); + #[async_trait] + impl TcpServerHandler for Handler { + async fn setup_server_stream( + &self, + stream: Box, + ) -> io::Result { + Ok(TcpServerSetupResult::TcpForward { + remote_location: NetLocation::from_str("127.0.0.1:80", None)?, + stream, + need_initial_flush: true, + proxy_selector: self.0.clone(), + connection_success_response: Some(b"response".to_vec().into_boxed_slice()), + initial_remote_data: Some(b"request".to_vec().into_boxed_slice()), + }) + } + } + + #[derive(Debug)] + struct ReadyConnector(std::sync::Mutex>); + + #[derive(Debug)] + struct FlushStream { + inner: tokio::io::DuplexStream, + stalled: bool, + flushes: Arc, + } + + impl tokio::io::AsyncRead for FlushStream { + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + Pin::new(&mut self.inner).poll_read(cx, buf) + } + } + + impl tokio::io::AsyncWrite for FlushStream { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + Pin::new(&mut self.inner).poll_write(cx, buf) + } + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.flushes.fetch_add(1, Ordering::Relaxed); + if self.stalled { + Poll::Pending + } else { + Pin::new(&mut self.inner).poll_flush(cx) + } + } + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.inner).poll_shutdown(cx) + } + } + + impl AsyncPing for FlushStream { + fn supports_ping(&self) -> bool { + false + } + fn poll_write_ping(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(false)) + } + } + impl AsyncStream for FlushStream {} + + #[async_trait] + impl SocketConnector for ReadyConnector { + async fn connect( + &self, + _: &Arc, + _: &crate::address::ResolvedLocation, + ) -> io::Result> { + Ok(Box::new(self.0.lock().unwrap().take().unwrap())) + } + async fn connect_udp_bidirectional( + &self, + _: &Arc, + _: crate::address::ResolvedLocation, + ) -> io::Result> { + unreachable!() + } + fn bind_interface(&self) -> Option<&str> { + None + } + } + + #[tokio::test(start_paused = true)] + async fn initial_writes_and_flushes_have_deadlines() { + use tokio::io::AsyncReadExt; + for (inbound_capacity, outbound_capacity, stall_response, stall_request) in [ + (1, 64, false, false), + (64, 1, false, false), + (64, 64, true, false), + (64, 64, false, true), + ] { + let (inbound, mut peer) = tokio::io::duplex(inbound_capacity); + let (outbound, mut target) = tokio::io::duplex(outbound_capacity); + let response_flushes = Arc::new(AtomicUsize::new(0)); + let request_flushes = Arc::new(AtomicUsize::new(0)); + let inbound = FlushStream { + inner: inbound, + stalled: stall_response, + flushes: response_flushes.clone(), + }; + let outbound = FlushStream { + inner: outbound, + stalled: stall_request, + flushes: request_flushes.clone(), + }; + let chain = ClientProxyChain::new( + vec![InitialHopEntry::Direct(Box::new(ReadyConnector( + std::sync::Mutex::new(Some(outbound)), + )))], + vec![], + ); + let selector = Arc::new(ClientProxySelector::new(vec![ConnectRule::new( + vec![crate::address::NetLocationMask::ANY], + ConnectAction::new_allow(None, ClientChainGroup::new(vec![chain])), + )])); + let started = tokio::time::Instant::now(); + let error = timeout( + Duration::from_secs(61), + process_stream( + inbound, + Arc::new(Handler(selector)), + Arc::new(PendingResolver), + ), + ) + .await + .expect("initial writes retained the stream indefinitely") + .unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::TimedOut); + assert_eq!(started.elapsed(), Duration::from_secs(60)); + if stall_response { + assert!(response_flushes.load(Ordering::Relaxed) > 0); + } + if stall_request { + assert!(request_flushes.load(Ordering::Relaxed) > 0); + } + peer.read_to_end(&mut Vec::new()).await.unwrap(); + target.read_to_end(&mut Vec::new()).await.unwrap(); + } + } + + #[tokio::test] + async fn protocol_response_precedes_chained_greeting() { + use tokio::io::AsyncReadExt; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let rule = serde_yaml::from_str(&format!( + "masks: '0.0.0.0/0'\nclient_proxy:\n address: '{}'\n protocol: {{type: socks}}\n", + listener.local_addr().unwrap() + )) + .unwrap(); + let resolver: Arc = Arc::new(crate::resolver::NativeResolver::new()); + let handler = Arc::new(Handler(Arc::new( + create_tcp_client_proxy_selector(vec![rule], resolver.clone()).unwrap(), + ))); + let (server, mut peer) = tokio::io::duplex(512); + let mut tasks = tokio::task::JoinSet::new(); + tasks.spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + socket.read_exact(&mut [0; 3]).await.unwrap(); + socket.write_all(&[5, 0]).await.unwrap(); + socket.read_exact(&mut [0; 10]).await.unwrap(); + socket + .write_all(b"\x05\x00\x00\x01\x00\x00\x00\x00\x00\x00greeting") + .await + .unwrap(); + let mut request = [0; 7]; + socket.read_exact(&mut request).await.unwrap(); + assert_eq!(&request, b"request"); + socket.write_all(b"later").await.unwrap(); + }); + tasks.spawn(async move { + process_stream(server, handler, resolver).await.unwrap(); + }); + let mut response = [0; 21]; + timeout(Duration::from_secs(3), peer.read_exact(&mut response)) + .await + .unwrap() + .unwrap(); + assert_eq!(&response, b"responsegreetinglater"); + drop(peer); + while let Some(result) = timeout(Duration::from_secs(3), tasks.join_next()) + .await + .unwrap() + { + result.unwrap(); + } + } + + #[tokio::test] + async fn tcp_bind_failure_releases_all_prepared_listeners() { + let first = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let first_addr = first.local_addr().unwrap(); + let occupied = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let config = format!( + "- address: ['{first_addr}', '{}']\n protocol:\n type: http\n", + occupied.local_addr().unwrap() + ); + let configs = + crate::config::create_server_configs(serde_yaml::from_str(&config).unwrap()).unwrap(); + drop(first); + let error = start_servers( + configs.configs.into_iter().next().unwrap(), + Arc::new(crate::resolver::NativeResolver::new()), + ) + .await + .unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::AddrInUse); + std::net::TcpListener::bind(first_addr).expect("earlier bind leaked after startup failed"); + } + #[derive(Debug)] struct StalledMessageShutdown { _marker: Arc<()>, diff --git a/src/tcp/tcp_server_handler_factory.rs b/src/tcp/tcp_server_handler_factory.rs index 4de6f163..cdcc0b81 100644 --- a/src/tcp/tcp_server_handler_factory.rs +++ b/src/tcp/tcp_server_handler_factory.rs @@ -19,12 +19,12 @@ use crate::option_util::OneOrSome; use crate::port_forward_handler::PortForwardServerHandler; use crate::reality::RealityServerTarget; use crate::resolver::Resolver; -use crate::rustls_config_util::create_server_config; +use crate::rustls_config_util::try_create_server_config; use crate::shadow_tls::{ShadowTlsServerTarget, ShadowTlsServerTargetHandshake}; use crate::shadowsocks::ShadowsocksTcpHandler; use crate::snell::snell_handler::SnellServerHandler; use crate::socks_handler::SocksTcpServerHandler; -use crate::tcp::chain_builder::build_client_proxy_chain; +use crate::tcp::chain_builder::try_build_client_proxy_chain; use crate::tcp::tcp_handler::TcpServerHandler; use crate::tls_server_handler::NaiveConfig; use crate::tls_server_handler::{ @@ -64,8 +64,8 @@ pub fn create_tcp_server_handler( client_proxy_selector: &Arc, resolver: &Arc, bind_ip: Option, -) -> Box { - match server_proxy_config { +) -> std::io::Result> { + Ok(match server_proxy_config { ServerProxyConfig::Http { username, password } => Box::new(HttpTcpServerHandler::new( create_auth_credentials(username, password), client_proxy_selector.clone(), @@ -164,44 +164,46 @@ pub fn create_tcp_server_handler( let mut all_targets = tls_targets .into_iter() .map(|(sni, config)| { - ( + Ok(( sni, - create_tls_server_target(config, client_proxy_selector, resolver, bind_ip), - ) + create_tls_server_target(config, client_proxy_selector, resolver, bind_ip)?, + )) }) - .collect::>(); - let default_tls_target = default_tls_target.map(|config| { - create_tls_server_target(*config, client_proxy_selector, resolver, bind_ip) - }); + .collect::>>()?; + let default_tls_target = default_tls_target + .map(|config| { + create_tls_server_target(*config, client_proxy_selector, resolver, bind_ip) + }) + .transpose()?; let shadowtls_targets = shadowtls_targets .into_iter() .map(|(sni, config)| { - ( + Ok(( sni, create_shadow_tls_server_target( config, client_proxy_selector, resolver, bind_ip, - ), - ) + )?, + )) }) - .collect::>(); + .collect::>>()?; all_targets.extend(shadowtls_targets); let reality_server_targets = reality_targets .into_iter() .map(|(sni, config)| { - ( + Ok(( sni, create_reality_server_target( config, client_proxy_selector, resolver, bind_ip, - ), - ) + )?, + )) }) - .collect::>(); + .collect::>>()?; all_targets.extend(reality_server_targets); Box::new(TlsServerHandler::new( all_targets, @@ -220,7 +222,7 @@ pub fn create_tcp_server_handler( udp_enabled, client_proxy_selector.clone(), resolver.clone(), - )), + )?), ServerProxyConfig::Websocket { targets } => { let server_targets: Vec = targets .into_vec() @@ -228,7 +230,7 @@ pub fn create_tcp_server_handler( .map(|config| { create_websocket_server_target(config, client_proxy_selector, resolver, bind_ip) }) - .collect::>(); + .collect::>>()?; Box::new(WebsocketTcpServerHandler::new(server_targets)) } ServerProxyConfig::PortForward { targets } => { @@ -252,10 +254,7 @@ pub fn create_tcp_server_handler( let padding = if let Some(scheme_lines) = padding_scheme { let scheme_str = scheme_lines.join("\n"); - Arc::new( - PaddingFactory::new(scheme_str.as_bytes()) - .expect("Invalid padding scheme (should be validated during config load)"), - ) + Arc::new(PaddingFactory::new(scheme_str.as_bytes()).map_err(std::io::Error::other)?) } else { PaddingFactory::default_factory() }; @@ -281,7 +280,7 @@ pub fn create_tcp_server_handler( unknown_config => { panic!("Unsupported TCP proxy config: {unknown_config:?}") } - } + }) } fn create_tls_server_target( @@ -289,7 +288,7 @@ fn create_tls_server_target( client_proxy_selector: &Arc, resolver: &Arc, bind_ip: Option, -) -> TlsServerTarget { +) -> std::io::Result { let TlsServerConfig { cert, key, @@ -326,20 +325,20 @@ fn create_tls_server_target( alpn_protocols.into_vec() }; - let server_config = Arc::new(create_server_config( + let server_config = Arc::new(try_create_server_config( &cert_bytes, &key_bytes, client_ca_certs, &effective_alpn, &client_fingerprints.into_vec(), - )); + )?); // Compute effective selector: if override_rules exist, create new selector; otherwise use parent's let effective_selector = if !override_rules.is_empty() { let rules = override_rules .map(ConfigSelection::unwrap_config) .into_vec(); - Arc::new(create_tcp_client_proxy_selector(rules, resolver.clone())) + Arc::new(create_tcp_client_proxy_selector(rules, resolver.clone())?) } else { client_proxy_selector.clone() }; @@ -385,15 +384,15 @@ fn create_tls_server_target( unreachable!("Vision requires VLESS (should be validated during config load)") } } else { - let handler = create_tcp_server_handler(protocol, &effective_selector, resolver, bind_ip); + let handler = create_tcp_server_handler(protocol, &effective_selector, resolver, bind_ip)?; InnerProtocol::Normal(handler) }; - TlsServerTarget::Tls { + Ok(TlsServerTarget::Tls { server_config, effective_selector, inner_protocol, - } + }) } fn create_shadow_tls_server_target( @@ -401,7 +400,7 @@ fn create_shadow_tls_server_target( client_proxy_selector: &Arc, resolver: &Arc, bind_ip: Option, -) -> TlsServerTarget { +) -> std::io::Result { let ShadowTlsServerConfig { password, handshake, @@ -421,20 +420,21 @@ fn create_shadow_tls_server_target( .map(|cert| cert.as_bytes().to_vec()) .collect(); - let server_config = Arc::new(create_server_config( + let server_config = Arc::new(try_create_server_config( &cert_bytes, &key_bytes, client_ca_certs, &handshake.alpn_protocols.into_vec(), &handshake.client_fingerprints.into_vec(), - )); + )?); ShadowTlsServerTargetHandshake::new_local(server_config) } ShadowTlsServerHandshakeConfig::Remote(handshake) => { // Build ClientProxyChain from client_chain // client_chain is guaranteed to be non-empty (defaults to direct hop) - let client_chain = build_client_proxy_chain(handshake.client_chain, resolver.clone()); + let client_chain = + try_build_client_proxy_chain(handshake.client_chain, resolver.clone())?; ShadowTlsServerTargetHandshake::new_remote(handshake.address, client_chain) } }; @@ -444,18 +444,18 @@ fn create_shadow_tls_server_target( let rules = override_rules .map(ConfigSelection::unwrap_config) .into_vec(); - Arc::new(create_tcp_client_proxy_selector(rules, resolver.clone())) + Arc::new(create_tcp_client_proxy_selector(rules, resolver.clone())?) } else { client_proxy_selector.clone() }; - let handler = create_tcp_server_handler(protocol, &effective_selector, resolver, bind_ip); + let handler = create_tcp_server_handler(protocol, &effective_selector, resolver, bind_ip)?; - TlsServerTarget::ShadowTls(ShadowTlsServerTarget::new( + Ok(TlsServerTarget::ShadowTls(ShadowTlsServerTarget::new( password, target_handshake, handler, - )) + ))) } fn create_reality_server_target( @@ -463,7 +463,7 @@ fn create_reality_server_target( client_proxy_selector: &Arc, resolver: &Arc, bind_ip: Option, -) -> TlsServerTarget { +) -> std::io::Result { let RealityServerConfig { private_key, short_ids, @@ -498,7 +498,7 @@ fn create_reality_server_target( let rules = override_rules .map(ConfigSelection::unwrap_config) .into_vec(); - Arc::new(create_tcp_client_proxy_selector(rules, resolver.clone())) + Arc::new(create_tcp_client_proxy_selector(rules, resolver.clone())?) } else { client_proxy_selector.clone() }; @@ -544,7 +544,7 @@ fn create_reality_server_target( unreachable!("Vision requires VLESS (should be validated during config load)") } } else { - let handler = create_tcp_server_handler(protocol, &effective_selector, resolver, bind_ip); + let handler = create_tcp_server_handler(protocol, &effective_selector, resolver, bind_ip)?; InnerProtocol::Normal(handler) }; @@ -553,23 +553,18 @@ fn create_reality_server_target( let hops = dest_client_chain.into_vec(); if hops.is_empty() { // Default to direct connection - build_client_proxy_chain( + try_build_client_proxy_chain( OneOrSome::One(ClientChainHop::Single(ConfigSelection::Config( ClientConfig::default(), ))), resolver.clone(), ) - } else if hops.len() == 1 { - build_client_proxy_chain( - OneOrSome::One(hops.into_iter().next().unwrap()), - resolver.clone(), - ) } else { - build_client_proxy_chain(OneOrSome::Some(hops), resolver.clone()) + try_build_client_proxy_chain(OneOrSome::Some(hops), resolver.clone()) } - }; + }?; - TlsServerTarget::Reality(RealityServerTarget { + Ok(TlsServerTarget::Reality(RealityServerTarget { private_key: private_key_bytes, short_ids: short_id_bytes, dest, @@ -580,7 +575,7 @@ fn create_reality_server_target( effective_selector, inner_protocol, dest_client_chain, - }) + })) } fn create_websocket_server_target( @@ -588,7 +583,7 @@ fn create_websocket_server_target( client_proxy_selector: &Arc, resolver: &Arc, bind_ip: Option, -) -> WebsocketServerTarget { +) -> std::io::Result { let WebsocketServerConfig { matching_path, matching_headers, @@ -611,17 +606,17 @@ fn create_websocket_server_target( let rules = override_rules .map(ConfigSelection::unwrap_config) .into_vec(); - Arc::new(create_tcp_client_proxy_selector(rules, resolver.clone())) + Arc::new(create_tcp_client_proxy_selector(rules, resolver.clone())?) } else { client_proxy_selector.clone() }; - let handler = create_tcp_server_handler(protocol, &effective_selector, resolver, bind_ip); + let handler = create_tcp_server_handler(protocol, &effective_selector, resolver, bind_ip)?; - WebsocketServerTarget { + Ok(WebsocketServerTarget { matching_path, matching_headers, ping_type, handler, - } + }) } diff --git a/src/tuic_server.rs b/src/tuic_server.rs index 510f5357..01053f34 100644 --- a/src/tuic_server.rs +++ b/src/tuic_server.rs @@ -2,6 +2,8 @@ use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr}; use std::str; use std::sync::Arc; use std::time::Duration; +use subtle::ConstantTimeEq; +use tokio::io::AsyncWriteExt; use bytes::{Bytes, BytesMut}; use dashmap::DashMap; @@ -233,18 +235,13 @@ async fn auth_connection( continue; } - let specified_uuid = stream_reader.read_slice(&mut recv_stream, 16).await?; - if specified_uuid != uuid { + let credentials = stream_reader.read_slice(&mut recv_stream, 48).await?; + let authenticated = + credentials[..16].ct_eq(uuid) & credentials[16..].ct_eq(&expected_token_bytes); + if !bool::from(authenticated) { return Err(std::io::Error::new( std::io::ErrorKind::PermissionDenied, - "incorrect uuid", - )); - } - let token_bytes = stream_reader.read_slice(&mut recv_stream, 32).await?; - if token_bytes != expected_token_bytes { - return Err(std::io::Error::new( - std::io::ErrorKind::PermissionDenied, - "incorrect token", + "incorrect credentials", )); } @@ -439,15 +436,13 @@ async fn process_tcp_stream( let mut server_stream: Box = Box::new(QuicStream::from(send, recv)); let setup_client_stream_future = timeout( Duration::from_secs(60), - setup_client_tcp_stream( - &mut server_stream, - client_proxy_selector, - resolver, - remote_location.clone(), - ), + setup_client_tcp_stream(client_proxy_selector, resolver, remote_location.clone()), ); - let mut client_stream = match setup_client_stream_future.await { + let crate::tcp::tcp_handler::TcpClientSetupResult { + mut client_stream, + early_data, + } = match setup_client_stream_future.await { Ok(Ok(Some(s))) => s, Ok(Ok(None)) => { // Must have been blocked. @@ -470,13 +465,19 @@ async fn process_tcp_stream( } }; - let unparsed_data = stream_reader.unparsed_data(); - let client_requires_flush = if unparsed_data.is_empty() { - false - } else { - write_all(&mut client_stream, unparsed_data).await?; - true - }; + crate::util::timeout_stream_setup(async { + if let Some(data) = early_data { + write_all(&mut server_stream, &data).await?; + server_stream.flush().await?; + } + let unparsed_data = stream_reader.unparsed_data(); + if !unparsed_data.is_empty() { + write_all(&mut client_stream, unparsed_data).await?; + client_stream.flush().await?; + } + Ok(()) + }) + .await?; drop(stream_reader); // Use 32KB buffers to match reference implementations @@ -484,7 +485,7 @@ async fn process_tcp_stream( &mut server_stream, &mut client_stream, false, // no need to flush since it's QUIC - client_requires_flush, + false, 32768, 32768, ) diff --git a/src/tun/mod.rs b/src/tun/mod.rs index 07bf506f..88c1564a 100644 --- a/src/tun/mod.rs +++ b/src/tun/mod.rs @@ -89,11 +89,22 @@ use udp_manager::TunUdpManager; /// 3. The stack thread reads packets directly from TUN using poll() /// 4. Handles TCP connections through the proxy chain /// 5. Handles UDP packets through tokio (forwarded from stack thread) +#[allow(dead_code)] pub async fn run_tun_server( + config: TunServerConfig, + proxy_selector: Arc, + resolver: Arc, + shutdown_rx: oneshot::Receiver<()>, +) -> std::io::Result<()> { + run_tun_server_inner(config, proxy_selector, resolver, shutdown_rx, &mut None).await +} + +async fn run_tun_server_inner( config: TunServerConfig, proxy_selector: Arc, resolver: Arc, mut shutdown_rx: oneshot::Receiver<()>, + ready: &mut Option>>, ) -> std::io::Result<()> { config.resource_limits.validate()?; info!( @@ -190,6 +201,9 @@ pub async fn run_tun_server( } info!("TUN server started successfully"); + if let Some(ready) = ready.take() { + let _ = ready.send(Ok(())); + } // Wait for shutdown signal or stack thread exit tokio::select! { @@ -248,21 +262,25 @@ async fn handle_tcp_connection( remote_location.location() ); - match tokio::time::timeout( - std::time::Duration::from_secs(30), - chain_group.connect_tcp(remote_location.clone(), &resolver), - ) + match tokio::time::timeout(std::time::Duration::from_secs(30), async { + let setup = chain_group + .connect_tcp(remote_location.clone(), &resolver) + .await?; + if let Some(data) = setup.early_data { + tokio::io::AsyncWriteExt::write_all(&mut connection, &data).await?; + } + Ok::<_, std::io::Error>(setup.client_stream) + }) .await .map_err(|_| { std::io::Error::new(std::io::ErrorKind::TimedOut, "TUN TCP setup timed out") })? { - Ok(setup_result) => { + Ok(mut remote) => { debug!( "TCP: connected to {}, starting bidirectional copy", remote_location.location() ); - let mut remote = setup_result.client_stream; let result = tokio::io::copy_bidirectional(&mut connection, &mut remote).await; match result { @@ -332,23 +350,45 @@ pub async fn start_tun_server( resolver: std::sync::Arc, ) -> std::io::Result> { let (shutdown_tx, shutdown_rx) = oneshot::channel(); + let (ready_tx, ready_rx) = oneshot::channel(); let handle = tokio::spawn(async move { let _keep_alive = shutdown_tx; - if let Err(e) = run_tun_from_config(config, shutdown_rx, true, resolver).await { - warn!("TUN server error: {}", e); + let mut ready = Some(ready_tx); + if let Err(error) = + run_tun_from_config_inner(config, shutdown_rx, true, resolver, &mut ready).await + { + if let Some(ready) = ready.take() { + let _ = ready.send(Err(error)); + } else { + warn!("TUN server error: {error}"); + } } }); - - Ok(handle) + let handle = tokio_util::task::AbortOnDropHandle::new(handle); + ready_rx + .await + .map_err(|_| std::io::Error::other("TUN task exited before startup completed"))??; + Ok(handle.detach()) } /// Run TUN server from config with external shutdown control. +#[allow(dead_code)] pub async fn run_tun_from_config( config: TunConfig, shutdown_rx: tokio::sync::oneshot::Receiver<()>, close_fd_on_drop: bool, resolver: Arc, +) -> std::io::Result<()> { + run_tun_from_config_inner(config, shutdown_rx, close_fd_on_drop, resolver, &mut None).await +} + +async fn run_tun_from_config_inner( + config: TunConfig, + shutdown_rx: tokio::sync::oneshot::Receiver<()>, + close_fd_on_drop: bool, + resolver: Arc, + ready: &mut Option>>, ) -> std::io::Result<()> { let mut tun_server_config = TunServerConfig::new() .mtu(config.mtu) @@ -384,13 +424,15 @@ pub async fn run_tun_from_config( } let rules = config.rules.map(ConfigSelection::unwrap_config).into_vec(); - let client_proxy_selector = Arc::new(create_tcp_client_proxy_selector(rules, resolver.clone())); + let client_proxy_selector = + Arc::new(create_tcp_client_proxy_selector(rules, resolver.clone())?); - run_tun_server( + run_tun_server_inner( tun_server_config, client_proxy_selector, resolver, shutdown_rx, + ready, ) .await } @@ -400,6 +442,96 @@ mod tests { use super::*; use std::os::fd::AsRawFd; + #[tokio::test] + async fn startup_acknowledges_initialization_and_returns_errors() { + use std::os::fd::IntoRawFd; + let resolver: Arc = Arc::new(crate::resolver::NativeResolver::new()); + let invalid = serde_yaml::from_str("resource_limits: {tcp_buffer_size: 0}").unwrap(); + assert_eq!( + start_tun_server(invalid, resolver.clone()) + .await + .unwrap_err() + .kind(), + std::io::ErrorKind::InvalidInput + ); + let (_peer, device) = std::os::unix::net::UnixDatagram::pair().unwrap(); + let fd = device.into_raw_fd(); + let config = serde_yaml::from_str(&format!("device_fd: {fd}")).unwrap(); + let task = start_tun_server(config, resolver).await.unwrap(); + assert!(!task.is_finished()); + assert!(unsafe { libc::fcntl(fd, libc::F_GETFD) } >= 0); + task.abort(); + let _ = task.await; + } + + #[tokio::test] + async fn proxy_handshake_bytes_reach_the_tun_connection_before_relay_data() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let proxy = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let rule: crate::config::RuleConfig = serde_yaml::from_str(&format!( + "masks: '0.0.0.0/0'\nclient_proxy:\n address: '{}'\n protocol:\n type: socks\n", + proxy.local_addr().unwrap() + )) + .unwrap(); + let resolver: Arc = Arc::new(crate::resolver::NativeResolver::new()); + let selector = + Arc::new(create_tcp_client_proxy_selector(vec![rule], resolver.clone()).unwrap()); + let control = Arc::new(tcp_conn::TcpConnectionControl::new(1024, 1024)); + let (wake, _wake_rx) = wake::Wake::new().unwrap(); + let connection = tcp_conn::TcpConnection::new(control.clone(), wake); + let (release, released) = tokio::sync::oneshot::channel(); + let mut tasks = JoinSet::new(); + tasks.spawn(async move { + let (mut socket, _) = proxy.accept().await.unwrap(); + let mut greeting = [0; 3]; + socket.read_exact(&mut greeting).await.unwrap(); + socket.write_all(&[5, 0]).await.unwrap(); + let mut request = [0; 10]; + socket.read_exact(&mut request).await.unwrap(); + socket + .write_all(b"\x05\x00\x00\x01\x00\x00\x00\x00\x00\x00greeting") + .await + .unwrap(); + released.await.unwrap(); + socket.write_all(b"later").await.unwrap(); + std::future::pending::<()>().await; + }); + tasks.spawn(async move { + handle_tcp_connection( + connection, + NetLocation::from_str("127.0.0.1:80", None).unwrap(), + selector, + resolver, + ) + .await + .unwrap(); + }); + let mut received = Vec::new(); + receive_tun_bytes(&control, &mut received, b"greeting").await; + release.send(()).unwrap(); + receive_tun_bytes(&control, &mut received, b"greetinglater").await; + tasks.shutdown().await; + } + + async fn receive_tun_bytes( + control: &tcp_conn::TcpConnectionControl, + received: &mut Vec, + expected: &[u8], + ) { + tokio::time::timeout(std::time::Duration::from_secs(3), async { + while received.len() < expected.len() { + let mut bytes = [0; 64]; + let count = control.dequeue_send_data(&mut bytes); + received.extend_from_slice(&bytes[..count]); + control.wake_sender(); + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + assert_eq!(received, expected); + } + #[tokio::test] async fn borrowed_tun_fd_remains_open_after_stop() { let (_peer, client) = std::os::unix::net::UnixStream::pair().unwrap(); diff --git a/src/tun/udp_tests.rs b/src/tun/udp_tests.rs index 54741107..316e7b81 100644 --- a/src/tun/udp_tests.rs +++ b/src/tun/udp_tests.rs @@ -91,10 +91,13 @@ async fn roundtrips(ipv6: bool, framed: bool) { .close_fd_on_drop(false) .packet_information(framed); let resolver: Arc = Arc::new(crate::resolver::NativeResolver::new()); - let selector = Arc::new(create_tcp_client_proxy_selector( - vec![crate::config::RuleConfig::default()], - resolver.clone(), - )); + let selector = Arc::new( + create_tcp_client_proxy_selector( + vec![crate::config::RuleConfig::default()], + resolver.clone(), + ) + .unwrap(), + ); let (shutdown, shutdown_rx) = oneshot::channel(); let mut tasks = JoinSet::new(); tasks.spawn(run_tun_server(config, selector, resolver, shutdown_rx)); diff --git a/src/vless/vless_util.rs b/src/vless/vless_util.rs index 990295d1..1dd4e476 100644 --- a/src/vless/vless_util.rs +++ b/src/vless/vless_util.rs @@ -22,12 +22,10 @@ pub async fn parse_addons_from_reader( .read_slice(stream, addon_length as usize) .await?; - log::debug!( - "Parsing addons: length={}, bytes={:?}", - addon_length, - addon_bytes - ); + parse_addons(addon_bytes) +} +fn parse_addons(addon_bytes: &[u8]) -> std::io::Result { // Parse protobuf-encoded addons // Format: field_tag length data [field_tag length data ...] // Field 1 = flow (string) @@ -38,13 +36,13 @@ pub async fn parse_addons_from_reader( while addon_cursor < addon_bytes.len() { // Read field tag - let field_tag = addon_bytes[addon_cursor]; - addon_cursor += 1; + let (field_tag, bytes_used) = read_varint(&addon_bytes[addon_cursor..])?; + addon_cursor += bytes_used; let field_number = field_tag >> 3; let wire_type = field_tag & 0x07; - if wire_type != 2 { + if wire_type != 2 || field_number == 0 || field_number > 0x1fff_ffff { return Err(std::io::Error::other(format!( "Unexpected wire type {} for field {}", wire_type, field_number @@ -61,7 +59,7 @@ pub async fn parse_addons_from_reader( addon_cursor += bytes_used; // Validate field_length is within bounds - if addon_cursor + field_length as usize > addon_bytes.len() { + if field_length > (addon_bytes.len() - addon_cursor) as u64 { return Err(std::io::Error::other(format!( "Field {} length {} exceeds remaining addon bytes (cursor: {}, total: {})", field_number, @@ -168,21 +166,23 @@ pub fn vision_flow_addon_data() -> &'static [u8] { } fn read_varint(data: &[u8]) -> std::io::Result<(u64, usize)> { - let mut cursor = 0usize; - let mut length = 0u64; - loop { - let byte = data[cursor]; - if (byte & 0b10000000) != 0 { - length = (length << 8) | ((byte ^ 0b10000000) as u64); - } else { - length = (length << 8) | (byte as u64); - return Ok((length, cursor + 1)); + let mut value = 0; + for (index, &byte) in data.iter().take(10).enumerate() { + if index == 9 && byte > 1 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "Varint overflow", + )); } - if cursor == 7 || cursor == data.len() { - return Err(std::io::Error::other("Varint is too long")); + value |= u64::from(byte & 0x7f) << (index * 7); + if byte & 0x80 == 0 { + return Ok((value, index + 1)); } - cursor += 1; } + Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + "Truncated varint", + )) } /// Encode a flow string as protobuf addon data @@ -212,3 +212,38 @@ fn encode_flow_addon(flow: &str) -> std::io::Result> { Ok(result) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn protobuf_varints_are_bounded_little_endian() { + assert_eq!(read_varint(&[0xac, 2]).unwrap(), (300, 2)); + let mut maximum = vec![0xff; 9]; + maximum.push(1); + assert_eq!(read_varint(&maximum).unwrap(), (u64::MAX, 10)); + maximum[9] = 2; + assert!(read_varint(&maximum).is_err()); + for invalid in [&[][..], &[0x80], &[0x80; 10]] { + assert!(read_varint(invalid).is_err()); + } + } + + #[test] + fn addon_lengths_cannot_wrap_or_escape_the_header() { + let mut addon = vec![0x0a]; + addon.extend([0xff; 9]); + addon.push(1); + assert!(parse_addons(&addon).is_err()); + assert!(parse_addons(&[0x0a, 2, b'a']).is_err()); + assert!(parse_addons(&[2, 0]).is_err()); + assert_eq!( + parse_addons(vision_flow_addon_data()).unwrap(), + XTLS_VISION_FLOW + ); + let mut long = vec![0x0a, 0x80, 1]; + long.extend([b'x'; 128]); + assert_eq!(parse_addons(&long).unwrap(), "x".repeat(128)); + } +} diff --git a/src/vmess/mod.rs b/src/vmess/mod.rs index d1d6836f..d5fa6227 100644 --- a/src/vmess/mod.rs +++ b/src/vmess/mod.rs @@ -2,8 +2,10 @@ mod crc32; mod fnv1a; mod md5; mod nonce; +mod replay; mod sha2; mod typed; mod vmess_handler; mod vmess_stream; +pub(crate) use vmess_handler::DataCipher; pub use vmess_handler::{VmessTcpClientHandler, VmessTcpServerHandler}; diff --git a/src/vmess/replay.rs b/src/vmess/replay.rs new file mode 100644 index 00000000..d2430ef4 --- /dev/null +++ b/src/vmess/replay.rs @@ -0,0 +1,125 @@ +use std::collections::HashMap; +use std::io; +use std::sync::OnceLock; +use std::time::SystemTime; + +use parking_lot::Mutex; + +use super::vmess_handler::{AUTH_ID_TIME_WINDOW_SECS, unix_time_secs}; + +static AUTH_IDS: OnceLock> = OnceLock::new(); + +pub(super) fn admit(identity: [u8; 16], auth_id: [u8; 16], timestamp: u64) -> io::Result<()> { + let mut cache = AUTH_IDS + .get_or_init(|| Mutex::new(ReplayCache::new(65536))) + .lock(); + cache.admit( + identity, + auth_id, + timestamp, + unix_time_secs(SystemTime::now())?, + ) +} + +struct ReplayCache { + expires: HashMap<([u8; 16], [u8; 16]), u64>, + last_now: u64, + capacity: usize, +} + +impl ReplayCache { + fn new(capacity: usize) -> Self { + Self { + expires: HashMap::new(), + last_now: 0, + capacity, + } + } + + fn admit( + &mut self, + identity: [u8; 16], + auth_id: [u8; 16], + timestamp: u64, + now: u64, + ) -> io::Result<()> { + // Expired IDs must not become admissible again after a local clock rollback. + if now < self.last_now { + return Err(io::Error::other( + "VMess authentication paused after clock rollback", + )); + } + if now > self.last_now { + self.expires.retain(|_, expires| *expires >= now); + self.last_now = now; + } + if timestamp.abs_diff(now) > AUTH_ID_TIME_WINDOW_SECS { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "VMess authentication timestamp outside allowed window", + )); + } + let key = (identity, auth_id); + if self.expires.contains_key(&key) || self.expires.len() >= self.capacity { + return Err(io::Error::new( + io::ErrorKind::PermissionDenied, + "VMess replay or replay cache capacity exhausted", + )); + } + self.expires + .insert(key, timestamp.saturating_add(AUTH_ID_TIME_WINDOW_SECS)); + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn future_id_remains_protected_through_inclusive_expiry() { + let mut cache = ReplayCache::new(2); + assert!(cache.admit([1; 16], [2; 16], 1120, 1000).is_ok()); + for now in [1120, 1239, 1240] { + assert!(cache.admit([1; 16], [2; 16], 1120, now).is_err()); + } + assert!(cache.admit([1; 16], [2; 16], 1120, 1241).is_err()); + assert!(cache.expires.is_empty()); + assert!(cache.admit([1; 16], [2; 16], 1120, 1240).is_err()); + assert!(cache.admit([1; 16], [3; 16], 1241, 1241).is_ok()); + } + + #[test] + fn timestamp_edges_capacity_and_identities() { + let mut cache = ReplayCache::new(2); + for timestamp in [879, 1121] { + assert!(cache.admit([1; 16], [2; 16], timestamp, 1000).is_err()); + } + assert!(cache.expires.is_empty()); + assert!(cache.admit([1; 16], [2; 16], 880, 1000).is_ok()); + assert!(cache.admit([3; 16], [2; 16], 1120, 1000).is_ok()); + assert!(cache.admit([4; 16], [2; 16], 1000, 1000).is_err()); + assert!(cache.admit([1; 16], [2; 16], 880, 1000).is_err()); + assert!(cache.admit([4; 16], [2; 16], 1001, 1001).is_ok()); + assert!(cache.admit([3; 16], [2; 16], 1120, 1001).is_err()); + } + + #[test] + fn process_cache_admits_a_concurrent_identity_only_once() { + let identity = rand::random(); + let auth_id = rand::random(); + let timestamp = unix_time_secs(SystemTime::now()).unwrap(); + let admitted = std::thread::scope(|scope| { + let threads: Vec<_> = (0..8) + .map(|_| scope.spawn(|| admit(identity, auth_id, timestamp).is_ok())) + .collect(); + threads + .into_iter() + .map(|thread| thread.join().unwrap()) + .filter(|admitted| *admitted) + .count() + }); + assert_eq!(admitted, 1); + assert!(admit(identity, auth_id, timestamp).is_err()); + } +} diff --git a/src/vmess/vmess_handler.rs b/src/vmess/vmess_handler.rs index c4418eb8..3c125dd3 100644 --- a/src/vmess/vmess_handler.rs +++ b/src/vmess/vmess_handler.rs @@ -34,6 +34,20 @@ use crate::uuid_util::parse_uuid; use crate::xudp::XudpMessageStream; const TAG_LEN: usize = 16; +pub(super) const AUTH_ID_TIME_WINDOW_SECS: u64 = 120; + +pub(super) fn unix_time_secs(time: SystemTime) -> std::io::Result { + time.duration_since(SystemTime::UNIX_EPOCH) + .map(|elapsed| elapsed.as_secs()) + .map_err(|_| std::io::Error::other("system clock is before the Unix epoch")) +} + +fn random_auth_id_timestamp() -> std::io::Result { + let now = unix_time_secs(SystemTime::now())?; + Ok(rand::rng().random_range( + now.saturating_sub(AUTH_ID_TIME_WINDOW_SECS)..=now.saturating_add(AUTH_ID_TIME_WINDOW_SECS), + )) +} // VMess protocol command types const COMMAND_TCP: u8 = 1; @@ -41,24 +55,29 @@ const COMMAND_UDP: u8 = 2; const COMMAND_MUX: u8 = 3; // MUX/XUDP mode #[derive(Debug, Clone, PartialEq, Eq)] -enum DataCipher { +pub(crate) enum DataCipher { Any, Aes128Gcm, ChaCha20Poly1305, None, } -impl From<&str> for DataCipher { - fn from(name: &str) -> Self { - match name { +impl TryFrom<&str> for DataCipher { + type Error = std::io::Error; + + fn try_from(name: &str) -> Result { + Ok(match name { "" | "any" => DataCipher::Any, "aes-128-gcm" => DataCipher::Aes128Gcm, "chacha20-poly1305" | "chacha20-ietf-poly1305" => DataCipher::ChaCha20Poly1305, "none" => DataCipher::None, _ => { - panic!("Unknown cipher: {name}"); + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("unsupported VMess cipher: {name}"), + )); } - } + }) } } @@ -87,8 +106,8 @@ impl VmessTcpServerHandler { udp_enabled: bool, proxy_selector: Arc, resolver: Arc, - ) -> Self { - let mut user_id_bytes = parse_uuid(user_id).unwrap(); + ) -> std::io::Result { + let mut user_id_bytes = parse_uuid(user_id)?; user_id_bytes.extend(b"c48619fe-8f02-49e0-b9e9-edf763e17e21"); let instruction_key: [u8; 16] = compute_md5(&user_id_bytes); @@ -96,14 +115,14 @@ impl VmessTcpServerHandler { let unbound_key = UnboundCipherKey::new(&AES_128, &derived_key[0..16]).unwrap(); let aead_decrypting_key = CipherDecryptingKey::ecb(unbound_key).unwrap(); - Self { - data_cipher: cipher_name.into(), + Ok(Self { + data_cipher: DataCipher::try_from(cipher_name)?, aead_decrypting_key, instruction_key, udp_enabled, proxy_selector, resolver, - } + }) } } @@ -144,14 +163,7 @@ impl TcpServerHandler for VmessTcpServerHandler { } let time_secs = u64::from_be_bytes(aead_bytes[0..8].try_into().unwrap()); - let current_time_secs = SystemTime::UNIX_EPOCH.elapsed().unwrap().as_secs(); - let time_delta = time_secs.abs_diff(current_time_secs); - if time_delta > 120 { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!("Hash timestamp is too old ({time_secs} is {time_delta} seconds old)"), - )); - } + super::replay::admit(self.instruction_key, cert_hash, time_secs)?; let mut encrypted_payload_length = [0u8; 18]; stream_reader @@ -202,8 +214,7 @@ impl TcpServerHandler for VmessTcpServerHandler { &[b"VMess Header AEAD Nonce", &cert_hash, &nonce], ); - let mut encrypted_header = - allocate_vec(payload_length as usize + TAG_LEN).into_boxed_slice(); + let mut encrypted_header = allocate_vec(payload_length as usize + TAG_LEN); stream_reader .read_slice_into(&mut server_stream, &mut encrypted_header) @@ -222,6 +233,7 @@ impl TcpServerHandler for VmessTcpServerHandler { return Err(std::io::Error::other("failed to open encrypted header")); } + encrypted_header.truncate(payload_length as usize); let mut header_reader = AeadHeaderReader { server_stream, decrypted_header: encrypted_header, @@ -331,31 +343,22 @@ impl TcpServerHandler for VmessTcpServerHandler { }; let margin_len: u8 = fixed_header[35] >> 4; - log::info!("VMess margin_len: {}, command: {}", margin_len, command); + log::debug!("VMess margin_len: {}, command: {}", margin_len, command); if margin_len > 0 { let mut margin_bytes = allocate_vec(margin_len as usize).into_boxed_slice(); header_reader.read_slice_into(&mut margin_bytes)?; - log::info!("VMess margin_bytes: {:?}", &margin_bytes[..]); fnv_hasher.write(&margin_bytes); } let mut check_bytes = [0u8; 4]; header_reader.read_slice_into(&mut check_bytes)?; - log::info!("VMess check_bytes: {:?}", check_bytes); let expected_check_value = u32::from_be_bytes(check_bytes[0..4].try_into().unwrap()); let actual_check_value = fnv_hasher.finish(); - log::info!( - "VMess FNV1a: expected={}, actual={}", - expected_check_value, - actual_check_value - ); if expected_check_value != actual_check_value { return Err(std::io::Error::new( std::io::ErrorKind::InvalidData, - format!( - "Bad fnv1a checksum, expected {expected_check_value}, got {actual_check_value}" - ), + "Bad VMess header checksum", )); } @@ -684,14 +687,21 @@ impl TcpServerHandler for VmessTcpServerHandler { struct AeadHeaderReader { server_stream: Box, - decrypted_header: Box<[u8]>, + decrypted_header: Vec, cursor: usize, } impl AeadHeaderReader { fn read_slice_into(&mut self, data: &mut [u8]) -> std::io::Result<()> { let len = data.len(); - data.copy_from_slice(&self.decrypted_header[self.cursor..self.cursor + len]); + let chunk = self + .decrypted_header + .get(self.cursor..) + .and_then(|remaining| remaining.get(..len)) + .ok_or_else(|| { + std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "truncated VMess header") + })?; + data.copy_from_slice(chunk); self.cursor += len; Ok(()) } @@ -718,8 +728,8 @@ impl std::fmt::Debug for VmessTcpClientHandler { } impl VmessTcpClientHandler { - pub fn new(cipher_name: &str, user_id: &str, udp_enabled: bool) -> Self { - let mut user_id_bytes = parse_uuid(user_id).unwrap(); + pub fn new(cipher_name: &str, user_id: &str, udp_enabled: bool) -> std::io::Result { + let mut user_id_bytes = parse_uuid(user_id)?; user_id_bytes.extend(b"c48619fe-8f02-49e0-b9e9-edf763e17e21"); let instruction_key: [u8; 16] = compute_md5(&user_id_bytes); @@ -727,12 +737,12 @@ impl VmessTcpClientHandler { let unbound_key = UnboundCipherKey::new(&AES_128, &derived_key[0..16]).unwrap(); let aead_encrypting_key = CipherEncryptingKey::ecb(unbound_key).unwrap(); - Self { - data_cipher: cipher_name.into(), + Ok(Self { + data_cipher: DataCipher::try_from(cipher_name)?, aead_encrypting_key, instruction_key, udp_enabled, - } + }) } } @@ -745,9 +755,7 @@ impl TcpClientHandler for VmessTcpClientHandler { ) -> std::io::Result { // AEAD allows 120 second delta from the current time. // See authid.go in v2ray-core. - let random_delta: u64 = rand::rng().random_range(0..241); - let time_secs: u64 = - SystemTime::UNIX_EPOCH.elapsed().unwrap().as_secs() - 120u64 + random_delta; + let time_secs = random_auth_id_timestamp()?; let mut aead_bytes = [0u8; 16]; let time_bytes = time_secs.to_be_bytes(); @@ -1028,9 +1036,7 @@ impl VmessTcpClientHandler { // Same as TCP setup but with command=2 and is_udp=true for VmessStream. // AEAD allows 120 second delta from the current time. - let random_delta: u64 = rand::rng().random_range(0..241); - let time_secs: u64 = - SystemTime::UNIX_EPOCH.elapsed().unwrap().as_secs() - 120u64 + random_delta; + let time_secs = random_auth_id_timestamp()?; let mut aead_bytes = [0u8; 16]; let time_bytes = time_secs.to_be_bytes(); @@ -1262,3 +1268,122 @@ impl VmessTcpClientHandler { Ok(Box::new(vmess_stream)) } } + +#[cfg(test)] +mod header_tests { + use super::*; + + #[test] + fn configured_ciphers_share_one_fallible_parser() { + for (name, cipher) in [ + ("", DataCipher::Any), + ("any", DataCipher::Any), + ("none", DataCipher::None), + ("aes-128-gcm", DataCipher::Aes128Gcm), + ("chacha20-poly1305", DataCipher::ChaCha20Poly1305), + ("chacha20-ietf-poly1305", DataCipher::ChaCha20Poly1305), + ] { + assert_eq!(DataCipher::try_from(name).unwrap(), cipher); + } + let uuid = crate::uuid_util::generate_uuid(); + for name in ["auto", "AES-128-GCM", "unsupported"] { + let expected = format!("unsupported VMess cipher: {name}"); + let server = VmessTcpServerHandler::new( + name, + &uuid, + false, + Arc::new(ClientProxySelector::new(vec![])), + Arc::new(crate::resolver::NativeResolver::new()), + ); + let client = VmessTcpClientHandler::new(name, &uuid, false); + for error in [ + DataCipher::try_from(name).unwrap_err(), + server.unwrap_err(), + client.unwrap_err(), + ] { + assert_eq!(error.kind(), std::io::ErrorKind::InvalidInput); + assert_eq!(error.to_string(), expected); + } + } + } + + #[tokio::test] + async fn replay_protection_survives_handler_recreation() { + let uuid = crate::uuid_util::generate_uuid(); + let client = VmessTcpClientHandler::new("none", &uuid, false).unwrap(); + let make_auth_id = |valid_checksum: bool, timestamp: u64| { + let mut id = [0; 16]; + id[..8].copy_from_slice(×tamp.to_be_bytes()); + let checksum = super::super::crc32::crc32c(&id[..12]) ^ u32::from(!valid_checksum); + id[12..].copy_from_slice(&checksum.to_be_bytes()); + client + .aead_encrypting_key + .less_safe_encrypt(&mut id, EncryptionContext::None) + .unwrap(); + id + }; + let now = unix_time_secs(SystemTime::now()).unwrap(); + for (id, expected) in [ + (make_auth_id(false, now), std::io::ErrorKind::InvalidData), + ( + make_auth_id(true, now.saturating_sub(121)), + std::io::ErrorKind::InvalidData, + ), + ( + make_auth_id(true, now), + std::io::ErrorKind::ConnectionAborted, + ), + ( + make_auth_id(true, now), + std::io::ErrorKind::PermissionDenied, + ), + ] { + let handler = VmessTcpServerHandler::new( + "none", + &uuid, + false, + Arc::new(ClientProxySelector::new(vec![])), + Arc::new(crate::resolver::NativeResolver::new()), + ) + .unwrap(); + let (server, mut peer) = tokio::io::duplex(64); + peer.write_all(&id).await.unwrap(); + drop(peer); + let error = match handler.setup_server_stream(Box::new(server)).await { + Err(error) => error, + Ok(_) => panic!("header unexpectedly accepted"), + }; + assert_eq!(error.kind(), expected, "{error}"); + } + } + + #[tokio::test] + async fn decrypted_header_reads_are_bounded_and_non_consuming_on_error() { + for len in [0, 1, 16, 37, 38] { + let (stream, _) = tokio::io::duplex(1); + let mut reader = AeadHeaderReader { + server_stream: Box::new(stream), + decrypted_header: vec![7; len], + cursor: 0, + }; + assert_eq!(reader.read_slice_into(&mut [0; 38]).is_ok(), len == 38); + if len < 38 { + assert_eq!(reader.cursor, 0); + reader.read_slice_into(&mut vec![0; len]).unwrap(); + } + reader.read_slice_into(&mut []).unwrap(); + assert_eq!( + reader.read_slice_into(&mut [0]).unwrap_err().kind(), + std::io::ErrorKind::UnexpectedEof + ); + } + } + + #[test] + fn clock_before_epoch_is_an_error() { + assert!( + unix_time_secs(SystemTime::UNIX_EPOCH - std::time::Duration::from_secs(1)).is_err() + ); + assert_eq!(unix_time_secs(SystemTime::UNIX_EPOCH).unwrap(), 0); + } +} diff --git a/src/vmess/vmess_stream.rs b/src/vmess/vmess_stream.rs index 071a556e..7d202210 100644 --- a/src/vmess/vmess_stream.rs +++ b/src/vmess/vmess_stream.rs @@ -423,6 +423,12 @@ impl VmessStream { )); } + if data_len < padding_len { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "padding exceeds VMess chunk length", + )); + } if self.tag_len > 0 && (data_len - padding_len) < self.tag_len { return Err(std::io::Error::new( std::io::ErrorKind::InvalidData, @@ -450,7 +456,7 @@ impl VmessStream { } let processed_data_len = data_len - padding_len - self.tag_len; - if self.processed_end_offset + processed_data_len >= self.processed_buf.len() { + if self.processed_end_offset + processed_data_len > self.processed_buf.len() { self.unprocessed_pending_len = Some((padding_len, data_len)); if self.unprocessed_start_offset == self.unprocessed_end_offset { self.unprocessed_start_offset = 0; @@ -468,7 +474,7 @@ impl VmessStream { } let processed_data_len = data_len - padding_len - self.tag_len; - if self.processed_end_offset + processed_data_len >= self.processed_buf.len() { + if self.processed_end_offset + processed_data_len > self.processed_buf.len() { return Ok(DecryptState::BufferFull); } @@ -1165,6 +1171,42 @@ mod tests { use shake::digest::{ExtendableOutput, Update}; use tokio::io::AsyncReadExt; + #[tokio::test] + async fn full_output_chunk_is_accepted_with_or_without_pending_length() { + for pending in [false, true] { + let (transport, _) = tokio::io::duplex(1); + let mut stream = response_test_stream(Box::new(transport), None, false, false, None); + let len = stream.processed_buf.len(); + assert!(len <= u16::MAX as usize); + stream.unprocessed_buf[..2].copy_from_slice(&(len as u16).to_be_bytes()); + stream.unprocessed_buf[2..2 + len].fill(42); + stream.unprocessed_end_offset = 2 + len; + if pending { + stream.unprocessed_pending_len = Some((0, len)); + stream.unprocessed_start_offset = 2; + } + assert!(matches!( + stream.try_decrypt().unwrap(), + DecryptState::Success + )); + assert_eq!(&*stream.processed_buf, vec![42; len]); + } + } + + #[tokio::test] + async fn padding_larger_than_chunk_is_rejected() { + let (transport, _) = tokio::io::duplex(1); + let mut stream = response_test_stream(Box::new(transport), None, true, true, None); + let (padding, mask) = + LengthMask::new(create_shake128_reader(&[23; 16]), true).next_values(); + assert!(padding > 0); + stream.unprocessed_buf[..2].copy_from_slice(&mask.to_be_bytes()); + stream.unprocessed_end_offset = 2; + assert!( + matches!(stream.try_decrypt(), Err(error) if error.kind() == std::io::ErrorKind::InvalidData) + ); + } + fn create_shake128_reader(iv: &[u8]) -> VmessReader { let mut hasher = Shake128::default(); hasher.update(iv); diff --git a/src/websocket/websocket_handler.rs b/src/websocket/websocket_handler.rs index 1bc65b9d..27812bb7 100644 --- a/src/websocket/websocket_handler.rs +++ b/src/websocket/websocket_handler.rs @@ -35,6 +35,31 @@ impl WebsocketTcpServerHandler { } } +#[cfg(test)] +mod request_tests { + use super::*; + + #[tokio::test] + async fn malformed_request_targets_are_errors_not_panics() { + for line in [ + "GET HTTP/1.1", + "GET HTTP/1.1", + "GET / a HTTP/1.1", + "POST / HTTP/1.1", + ] { + let (mut client, server) = tokio::io::duplex(256); + client + .write_all(format!("{line}\r\n\r\n").as_bytes()) + .await + .unwrap(); + let result = WebsocketTcpServerHandler::new(vec![]) + .setup_server_stream(Box::new(server)) + .await; + assert!(result.is_err(), "{line}"); + } + } +} + #[async_trait] impl TcpServerHandler for WebsocketTcpServerHandler { async fn setup_server_stream( @@ -42,7 +67,7 @@ impl TcpServerHandler for WebsocketTcpServerHandler { mut server_stream: Box, ) -> std::io::Result { let ParsedHttpData { - mut first_line, + first_line, headers: mut request_headers, stream_reader, } = ParsedHttpData::parse(&mut server_stream).await?; @@ -59,11 +84,13 @@ impl TcpServerHandler for WebsocketTcpServerHandler { ))); } - // remove ' HTTP/1.x' - first_line.truncate(first_line.len() - 9); - - // return the path after 'GET ' - first_line.split_off(4) + let path = first_line + .get(4..first_line.len() - 9) + .filter(|path| { + !path.is_empty() && !path.bytes().any(|byte| byte.is_ascii_whitespace()) + }) + .ok_or_else(|| std::io::Error::other("invalid HTTP request target"))?; + path.to_owned() }; let websocket_key = request_headers diff --git a/tests/cases/integration/resource_lifecycle_test.rs b/tests/cases/integration/resource_lifecycle_test.rs index 80a76975..b8de9ec1 100644 --- a/tests/cases/integration/resource_lifecycle_test.rs +++ b/tests/cases/integration/resource_lifecycle_test.rs @@ -133,24 +133,26 @@ async fn tcp_reload_drains_existing_connections_within_grace() -> io::Result<()> let config = format!( "- address: '0.0.0.0:{proxy_port}'\n protocol:\n type: forward\n target: '127.0.0.1:{echo_port}'\n" ); - let (_process, mut file, _) = start_process(&config, true)?; + let (_process, file, _) = start_process(&config, true)?; ports.wait_for_all_ports().await?; let mut existing = TcpStream::connect(("127.0.0.1", proxy_port)).await?; round_trip(&mut existing, b"before reload").await?; - file.write_all(b"\n# trigger reload\n")?; - file.flush()?; - - // Observe listener retirement rather than assuming notification timing. - timeout(Duration::from_secs(3), async { - while TcpStream::connect(("127.0.0.1", proxy_port)).await.is_ok() { - sleep(Duration::from_millis(20)).await; - } - }) - .await?; + let reserved = std::net::TcpListener::bind("127.0.0.1:0")?; + let replacement_port = reserved.local_addr()?.port(); + drop(reserved); + let replacement_config = config.replace( + &format!("0.0.0.0:{proxy_port}"), + &format!("0.0.0.0:{replacement_port}"), + ); + std::fs::write( + file.path(), + format!("- global_limits: {{reload_grace_secs: 6}}\n{replacement_config}"), + )?; round_trip(&mut existing, b"during reload").await?; + // A distinct listener identifies the new generation without requiring an outage. let mut replacement = timeout(Duration::from_secs(5), async { loop { - if let Ok(stream) = TcpStream::connect(("127.0.0.1", proxy_port)).await { + if let Ok(stream) = TcpStream::connect(("127.0.0.1", replacement_port)).await { break stream; } sleep(Duration::from_millis(20)).await; @@ -160,7 +162,7 @@ async fn tcp_reload_drains_existing_connections_within_grace() -> io::Result<()> round_trip(&mut existing, b"after reload").await?; round_trip(&mut replacement, b"new generation").await?; let mut byte = [0]; - let closed = timeout(Duration::from_secs(5), existing.read(&mut byte)).await?; + let closed = timeout(Duration::from_secs(8), existing.read(&mut byte)).await?; assert!( matches!(closed, Ok(0) | Err(_)), "old generation survived its grace deadline" @@ -247,6 +249,7 @@ async fn hysteria2_control_streams_survive_small_application_caps() }) .await??; assert_eq!(response.status().as_u16(), 233, "stream cap {limit}"); + assert_eq!(response.headers()["Hysteria-CC-RX"], "auto"); } Ok(()) } @@ -324,7 +327,7 @@ async fn quic_reload_disconnects_old_connections_and_allows_reconnect() file.write_all(b"\n# trigger QUIC reload\n")?; file.flush()?; for (_, connection, _, _) in &old { - let closed = timeout(Duration::from_secs(2), connection.closed()).await?; + let closed = timeout(Duration::from_secs(6), connection.closed()).await?; assert!(matches!( closed, quinn::ConnectionError::ApplicationClosed(_) diff --git a/tests/cases/integration/socks5_udp_associate_test.rs b/tests/cases/integration/socks5_udp_associate_test.rs index aa70c833..24e74863 100644 --- a/tests/cases/integration/socks5_udp_associate_test.rs +++ b/tests/cases/integration/socks5_udp_associate_test.rs @@ -39,6 +39,48 @@ use tokio::time::{sleep, timeout}; // BASIC FUNCTIONALITY TESTS +#[tokio::test] +async fn udp_associate_accepts_private_and_cross_family_hints() +-> Result<(), Box> { + let mut ports = PortHelper::new(); + let (ip, port) = ports.get_listener_port(); + let (echo_ip, echo_port) = ports.get_port(); + let _echo = start_udp_echo_server(&echo_ip, echo_port).await?; + let (_server, _config) = start_shoes_server(&format!( + "- address: '{ip}:{port}'\n protocol: {{type: socks, udp_enabled: true}}\n" + ))?; + ports.wait_for_all_ports().await?; + let target: std::net::SocketAddr = format!("{echo_ip}:{echo_port}").parse()?; + for hint in ["[::1]:0", "10.0.0.1:0", "[fd00::1]:0", "127.0.0.2:0"] { + let mut control = TcpStream::connect((ip.as_str(), port)).await?; + control.write_all(&[5, 1, 0]).await?; + let mut auth = [0; 2]; + control.read_exact(&mut auth).await?; + assert_eq!(auth, [5, 0]); + let mut request = encode_udp_datagram(&SocksDestination::Ip(hint.parse()?), b"")?; + request[..3].copy_from_slice(&[5, 3, 0]); + control.write_all(&request).await?; + let mut reply = [0; 10]; + control.read_exact(&mut reply).await?; + assert_eq!(&reply[..4], &[5, 0, 0, 1]); + let relay = std::net::SocketAddr::from(( + [reply[4], reply[5], reply[6], reply[7]], + u16::from_be_bytes([reply[8], reply[9]]), + )); + let socket = tokio::net::UdpSocket::bind((ip.as_str(), 0)).await?; + let response = send_udp_datagram( + &socket, + relay, + &SocksDestination::Ip(target), + b"hint", + Duration::from_secs(3), + ) + .await?; + assert_eq!(response, b"hint [ECHO]"); + } + Ok(()) +} + /// Test basic UDP ASSOCIATE - single packet echo #[tokio::test] async fn test_udp_associate_basic() -> Result<(), Box> { diff --git a/tests/cases/optional_peer/snell_mihomo.rs b/tests/cases/optional_peer/snell_mihomo.rs new file mode 100644 index 00000000..858f6bca --- /dev/null +++ b/tests/cases/optional_peer/snell_mihomo.rs @@ -0,0 +1,71 @@ +use std::error::Error; +use std::time::Duration; + +use shoes_test_support as common; +use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader}; +use tokio::net::TcpStream; + +use common::process::start_mihomo_server; +use common::socks5::Socks5UdpAssociation; +use common::test_fixture::start_shoes_server; +use common::test_servers::{start_tcp_stream_echo_server, start_udp_echo_server}; + +async fn roundtrip(version: u8, cipher: &str) -> Result<(), Box> { + let mut ports = common::port_helper::PortHelper::new(); + let (server_ip, server_port) = ports.get_listener_port(); + let (client_ip, client_port) = ports.get_localhost_listener_port(); + let tcp_echo = start_tcp_stream_echo_server("127.0.0.1", 0).await?; + let udp_echo = start_udp_echo_server("127.0.0.1", 0).await?; + let _server = start_shoes_server(&format!( + "- address: '{server_ip}:{server_port}'\n protocol:\n type: snell\n cipher: {cipher}\n password: snell-peer-test\n" + ))?; + let _client = start_mihomo_server(&format!( + "mixed-port: {client_port}\nmode: rule\nproxies:\n - name: snell\n type: snell\n server: {server_ip}\n port: {server_port}\n psk: snell-peer-test\n version: {version}\n udp: {}\nrules:\n - MATCH,snell\n", version == 3 + )) + .await?; + ports.wait_for_all_ports().await?; + + tokio::time::timeout(Duration::from_secs(10), async { + let mut stream = + BufReader::new(TcpStream::connect((client_ip.as_str(), client_port)).await?); + let destination = tcp_echo.local_addr(); + stream + .write_all( + format!("CONNECT {destination} HTTP/1.1\r\nHost: {destination}\r\n\r\n").as_bytes(), + ) + .await?; + let mut line = String::new(); + stream.read_line(&mut line).await?; + assert!(line.starts_with("HTTP/1.1 200"), "{line}"); + loop { + line.clear(); + assert_ne!(stream.read_line(&mut line).await?, 0); + if line == "\r\n" { + break; + } + } + let payload = b"independent Snell framing"; + stream.write_all(payload).await?; + let mut response = vec![0; payload.len()]; + stream.read_exact(&mut response).await?; + assert_eq!(response, payload); + if version == 3 { + let association = Socks5UdpAssociation::connect(&client_ip, client_port).await?; + let response = association.send_to(udp_echo.local_addr(), payload).await?; + assert_eq!(response, [payload.as_slice(), b" [ECHO]"].concat()); + } + Ok::<_, Box>(()) + }) + .await??; + Ok(()) +} + +#[tokio::test] +async fn chacha20_tcp() -> Result<(), Box> { + roundtrip(1, "chacha20-ietf-poly1305").await +} + +#[tokio::test] +async fn aes128_tcp_and_udp() -> Result<(), Box> { + roundtrip(3, "aes-128-gcm").await +} diff --git a/tests/optional_peer.rs b/tests/optional_peer.rs index 5385d5c1..778af006 100644 --- a/tests/optional_peer.rs +++ b/tests/optional_peer.rs @@ -1,2 +1,5 @@ #[path = "cases/optional_peer/anytls_uot_mihomo.rs"] mod anytls_uot_mihomo; + +#[path = "cases/optional_peer/snell_mihomo.rs"] +mod snell_mihomo; diff --git a/tests/process_contract.rs b/tests/process_contract.rs new file mode 100644 index 00000000..f64d6a72 --- /dev/null +++ b/tests/process_contract.rs @@ -0,0 +1,696 @@ +use std::io::{self, BufRead}; +use std::net::SocketAddr; +use std::path::Path; +use std::process::{Child, Command, ExitStatus, Stdio}; +use std::time::Duration; + +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::TcpStream; +use tokio::time::timeout; + +struct Process { + child: Child, + lines: tokio::sync::mpsc::UnboundedReceiver, +} + +impl Process { + fn start(path: &Path, arguments: &[&str]) -> Self { + let mut child = Command::new(env!("CARGO_BIN_EXE_shoes")) + .args(["-t", "1"]) + .args(arguments) + .arg(path) + .env("RUST_LOG", "error") + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .unwrap(); + let (tx, lines) = tokio::sync::mpsc::unbounded_channel(); + let readers: Vec> = vec![ + Box::new(child.stdout.take().unwrap()), + Box::new(child.stderr.take().unwrap()), + ]; + for reader in readers { + let tx = tx.clone(); + std::thread::spawn(move || { + for line in io::BufReader::new(reader).lines().map_while(Result::ok) { + if tx.send(line).is_err() { + break; + } + } + }); + } + Self { child, lines } + } + + async fn wait_for(&mut self, marker: &str) { + let mut seen = String::new(); + let result = timeout(Duration::from_secs(8), async { + while let Some(line) = self.lines.recv().await { + seen.push_str(&line); + seen.push('\n'); + if line.contains(marker) { + return; + } + } + panic!("process exited waiting for {marker}: {seen}"); + }) + .await; + assert!(result.is_ok(), "timed out waiting for {marker}: {seen}"); + } + + async fn exit(&mut self) -> ExitStatus { + timeout(Duration::from_secs(8), async { + loop { + if let Some(status) = self.child.try_wait().unwrap() { + return status; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("process did not exit") + } + + #[cfg(unix)] + fn signal(&self, signal: i32) { + assert_eq!(unsafe { libc::kill(self.child.id() as i32, signal) }, 0); + } +} + +impl Drop for Process { + fn drop(&mut self) { + let _ = self.child.kill(); + let _ = self.child.wait(); + } +} + +fn available_address() -> SocketAddr { + std::net::TcpListener::bind("127.0.0.1:0") + .unwrap() + .local_addr() + .unwrap() +} + +fn forward(address: SocketAddr, target: SocketAddr) -> String { + format!( + "- global_limits: {{reload_grace_secs: 2}}\n- address: '{address}'\n protocol:\n type: forward\n target: '{target}'\n" + ) +} + +async fn roundtrip(stream: &mut TcpStream) { + timeout(Duration::from_secs(2), async { + stream.write_all(b"still serving").await.unwrap(); + let mut response = [0; 13]; + stream.read_exact(&mut response).await.unwrap(); + assert_eq!(&response, b"still serving"); + }) + .await + .unwrap(); +} + +fn invalid_protocol_configs(address: SocketAddr) -> Vec { + use serde_json::json; + let reality = json!({ + "type": "reality", "public_key": "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA", + "sni_hostname": "example.com", "protocol": {"type": "socks"} + }); + let mut bad_key = reality.clone(); + bad_key["public_key"] = json!("invalid"); + let mut bad_sni = reality.clone(); + bad_sni["sni_hostname"] = json!("invalid name"); + let mut ipv4_sni = reality.clone(); + ipv4_sni["sni_hostname"] = json!("127.0.0.1"); + let mut ipv6_sni = reality.clone(); + ipv6_sni["sni_hostname"] = json!("::1"); + let mut nested_direct = reality.clone(); + nested_direct["protocol"] = json!({"type": "direct"}); + let mut missing_sni = reality; + missing_sni.as_object_mut().unwrap().remove("sni_hostname"); + let padding = json!({"type": "anytls", "password": "test", "padding_scheme": ["stop=invalid"]}); + let protocols = [ + bad_key, + bad_sni, + missing_sni, + ipv4_sni, + ipv6_sni, + nested_direct, + padding.clone(), + json!({"type": "tls", "protocol": {"type": "direct"}}), + json!({"type": "shadowtls", "password": "test", "protocol": {"type": "direct"}}), + json!({"type": "websocket", "protocol": {"type": "direct"}}), + ]; + let mut configs: Vec<_> = protocols + .iter() + .map(|protocol| { + serde_yaml::to_string(&json!([{ + "address": address.to_string(), "protocol": {"type": "http"}, + "rules": [{"masks": "0.0.0.0/0", "client_proxy": { + "address": "127.0.0.1:443", "protocol": protocol + }}] + }])) + .unwrap() + }) + .collect(); + let mut invalid_chains: Vec<_> = protocols + .into_iter() + .map(|protocol| { + json!({ + "address": "127.0.0.1:443", "protocol": protocol + }) + }) + .collect(); + invalid_chains.extend([ + json!("missing-group"), + json!(["direct", "direct"]), + json!({"address": "127.0.0.1:443", "protocol": { + "type": "tls", "cert": "-----BEGIN CERTIFICATE-----bad", + "key": "-----BEGIN PRIVATE KEY-----bad", "protocol": {"type": "socks"} + }}), + ]); + for chain in invalid_chains { + for protocol in handshake_protocols(chain) { + configs.push( + serde_yaml::to_string(&json!([{ + "address": address.to_string(), "protocol": protocol + }])) + .unwrap(), + ); + } + } + configs.push( + serde_yaml::to_string(&json!([{"address": address.to_string(), "protocol": padding}])) + .unwrap(), + ); + configs +} + +fn handshake_protocols(chain: serde_json::Value) -> [serde_json::Value; 2] { + use serde_json::json; + [ + json!({"type": "tls", "reality_targets": {"example.com": { + "private_key": "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA", + "dest": "example.com:443", "dest_client_chain": chain, + "protocol": {"type": "socks"} + }}}), + json!({"type": "tls", "shadowtls_targets": {"example.com": { + "password": "test", "handshake": { + "address": "example.com:443", "client_chain": chain + }, "protocol": {"type": "socks"} + }}}), + ] +} + +#[tokio::test] +async fn handshake_chains_expand_groups_and_load_certificates() { + use serde_json::json; + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.yaml"); + let cert_path = directory.path().join("client.pem"); + let key_path = directory.path().join("client.key"); + let cert = rcgen::generate_simple_self_signed(vec!["localhost".into()]).unwrap(); + std::fs::write(&cert_path, cert.cert.pem()).unwrap(); + std::fs::write(&key_path, cert.signing_key.serialize_pem()).unwrap(); + let client = json!({"address": "127.0.0.1:443", "protocol": { + "type": "tls", "cert": cert_path, "key": key_path, + "protocol": {"type": "socks"} + }}); + for chain in [json!("fallback"), client.clone()] { + for protocol in handshake_protocols(chain) { + let config = json!([ + {"client_group": "fallback", "client_proxies": [client.clone(), client.clone()]}, + {"address": available_address().to_string(), "protocol": protocol} + ]); + std::fs::write(&path, serde_yaml::to_string(&config).unwrap()).unwrap(); + let output = Command::new(env!("CARGO_BIN_EXE_shoes")) + .args(["check", "-t", "1"]) + .arg(&path) + .output() + .unwrap(); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + let mut child = Process::start(&path, &["--no-reload"]); + child.wait_for("Servers ready").await; + } + } +} + +#[test] +fn check_rejects_invalid_protocol_constructor_settings() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.yaml"); + for config in invalid_protocol_configs(available_address()) { + std::fs::write(&path, config).unwrap(); + let output = Command::new(env!("CARGO_BIN_EXE_shoes")) + .args(["check", "-t", "1"]) + .arg(&path) + .output() + .unwrap(); + assert_eq!(output.status.code(), Some(1)); + let error = String::from_utf8_lossy(&output.stderr); + assert!(!error.contains("panicked"), "{error}"); + } +} + +#[tokio::test] +async fn check_alias_and_validation_exit_codes() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.yaml"); + let occupied = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let valid = format!( + "- address: '{}'\n protocol: {{type: http}}\n", + occupied.local_addr().unwrap() + ); + for option in ["--dry-run", "check"] { + for (config, expected) in [ + (valid.as_str(), 0), + ( + "- address: '127.0.0.1:1'\n transport: udp\n protocol: {type: http}\n", + 1, + ), + ("not: [valid", 1), + ] { + std::fs::write(&path, config).unwrap(); + let output = Command::new(env!("CARGO_BIN_EXE_shoes")) + .args([option, "-t", "1"]) + .arg(&path) + .output() + .unwrap(); + assert_eq!( + output.status.code(), + Some(expected), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + assert!(!String::from_utf8_lossy(&output.stderr).contains("panicked")); + } + } + std::fs::remove_file(&path).unwrap(); + let mut child = Process::start(&path, &[]); + assert_eq!(child.exit().await.code(), Some(1)); + let version = Command::new(env!("CARGO_BIN_EXE_shoes")) + .arg("version") + .output() + .unwrap(); + assert!(version.status.success()); + assert_eq!( + String::from_utf8(version.stdout).unwrap().trim(), + concat!("shoes ", env!("CARGO_PKG_VERSION")) + ); +} + +#[tokio::test] +async fn initial_partial_bind_failure_exits_and_releases_addresses() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.yaml"); + let first = available_address(); + let occupied = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + for config in [ + format!( + "- address: ['{first}', '{}']\n protocol: {{type: http}}\n", + occupied.local_addr().unwrap() + ), + format!( + "- address: '{first}'\n protocol: {{type: http}}\n- address: '{}'\n protocol: {{type: http}}\n", + occupied.local_addr().unwrap() + ), + ] { + std::fs::write(&path, config).unwrap(); + let mut child = Process::start(&path, &["--no-reload"]); + assert_eq!(child.exit().await.code(), Some(1)); + std::net::TcpListener::bind(first).unwrap(); + } +} + +#[cfg(unix)] +#[tokio::test] +async fn rejected_reload_preserves_traffic_and_valid_reload_drains_tcp() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.yaml"); + let echo = shoes_test_support::test_servers::start_tcp_stream_echo_server("127.0.0.1", 0) + .await + .unwrap(); + let address = available_address(); + let valid = forward(address, echo.local_addr()); + std::fs::write(&path, &valid).unwrap(); + let mut child = Process::start(&path, &["--no-reload"]); + child.wait_for("Servers ready").await; + let mut existing = TcpStream::connect(address).await.unwrap(); + roundtrip(&mut existing).await; + let cert = rcgen::generate_simple_self_signed(vec!["localhost".into()]).unwrap(); + let other = rcgen::generate_simple_self_signed(vec!["localhost".into()]).unwrap(); + let mismatched_identity = serde_yaml::to_string(&serde_json::json!([{ + "address": address.to_string(), + "protocol": {"type": "tls", "tls_targets": {"localhost": { + "cert": cert.cert.pem(), + "key": other.signing_key.serialize_pem(), + "protocol": {"type": "http"} + }}} + }])) + .unwrap(); + let mut invalid_configs = vec![ + "invalid: [yaml".to_owned(), + format!( + "- address: '{address}'\n protocol:\n type: tls\n tls_targets:\n localhost:\n cert: absent.pem\n key: absent.key\n protocol: {{type: http}}\n" + ), + format!( + "- address: '{address}'\n protocol:\n type: tls\n tls_targets:\n localhost:\n cert: '-----BEGIN CERTIFICATE-----bad'\n key: '-----BEGIN PRIVATE KEY-----bad'\n protocol: {{type: http}}\n" + ), + mismatched_identity, + ]; + invalid_configs.extend(invalid_protocol_configs(address)); + if cfg!(target_os = "linux") { + invalid_configs.push(format!("{valid}\n- dns_group: failing\n dns_servers:\n - url: udp://127.0.0.1\n client_chain:\n address: '127.0.0.1:443'\n transport: quic\n bind_interface: shoes-missing-interface\n protocol: {{type: socks}}\n")); + } + for invalid in invalid_configs { + std::fs::write(&path, invalid).unwrap(); + child.signal(libc::SIGHUP); + child.wait_for("Reload rejected").await; + roundtrip(&mut existing).await; + roundtrip(&mut TcpStream::connect(address).await.unwrap()).await; + } + std::fs::write(&path, valid).unwrap(); + child.signal(libc::SIGHUP); + child.wait_for("Servers ready").await; + roundtrip(&mut existing).await; + roundtrip(&mut TcpStream::connect(address).await.unwrap()).await; + let result = timeout(Duration::from_secs(4), existing.read_u8()) + .await + .unwrap(); + assert!(result.is_err()); + child.signal(libc::SIGTERM); + assert_eq!(child.exit().await.code(), Some(0)); +} + +#[cfg(unix)] +#[tokio::test] +async fn shutdown_interrupts_watcher_debounce() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.yaml"); + let config = format!( + "- address: '{}'\n protocol: {{type: http}}\n", + available_address() + ); + std::fs::write(&path, &config).unwrap(); + let mut child = Process::start(&path, &[]); + child.wait_for("Servers ready").await; + std::fs::write(&path, config).unwrap(); + tokio::time::sleep(Duration::from_millis(200)).await; + child.signal(libc::SIGTERM); + assert_eq!( + timeout(Duration::from_secs(2), child.exit()) + .await + .unwrap() + .code(), + Some(0) + ); +} + +#[cfg(unix)] +#[tokio::test] +async fn sighup_completes_pending_file_reload_without_waiting_for_debounce() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.yaml"); + let config = format!( + "- address: '{}'\n protocol: {{type: http}}\n", + available_address() + ); + std::fs::write(&path, &config).unwrap(); + let mut child = Process::start(&path, &[]); + child.wait_for("Servers ready").await; + std::fs::write(&path, config).unwrap(); + tokio::time::sleep(Duration::from_millis(200)).await; + child.signal(libc::SIGHUP); + timeout(Duration::from_secs(2), child.wait_for("Servers ready")) + .await + .unwrap(); + assert!( + timeout(Duration::from_millis(3500), async { + while let Some(line) = child.lines.recv().await { + assert!(!line.contains("Servers ready"), "duplicate reload"); + } + }) + .await + .is_err() + ); + child.signal(libc::SIGTERM); + assert_eq!(child.exit().await.code(), Some(0)); +} + +#[cfg(unix)] +#[tokio::test] +async fn atomic_replace_keeps_serving_through_debounce() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.yaml"); + let echo = shoes_test_support::test_servers::start_tcp_stream_echo_server("127.0.0.1", 0) + .await + .unwrap(); + let address = available_address(); + let valid = forward(address, echo.local_addr()); + std::fs::write(&path, &valid).unwrap(); + let mut child = Process::start(&path, &[]); + child.wait_for("Servers ready").await; + let replacement = directory.path().join("replacement.yaml"); + std::fs::write(&replacement, &valid).unwrap(); + std::fs::rename(replacement, &path).unwrap(); + tokio::time::sleep(Duration::from_millis(100)).await; + roundtrip(&mut TcpStream::connect(address).await.unwrap()).await; + child.wait_for("Servers ready").await; + roundtrip(&mut TcpStream::connect(address).await.unwrap()).await; + child.signal(libc::SIGINT); + assert_eq!(child.exit().await.code(), Some(0)); +} + +#[cfg(any(target_os = "linux", target_os = "android"))] +#[tokio::test] +async fn hard_link_updates_trigger_reload() { + let directory = tempfile::tempdir().unwrap(); + let aliases = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.yaml"); + let config = format!( + "- address: '{}'\n protocol: {{type: http}}\n", + available_address() + ); + std::fs::write(&path, &config).unwrap(); + let mut child = Process::start(&path, &[]); + child.wait_for("Servers ready").await; + for alias in [ + directory.path().join("alias.yaml"), + aliases.path().join("alias.yaml"), + ] { + std::fs::hard_link(&path, &alias).unwrap(); + std::fs::write(&alias, &config).unwrap(); + child.wait_for("Servers ready").await; + } + child.signal(libc::SIGTERM); + assert!(child.exit().await.success()); +} + +#[cfg(unix)] +#[tokio::test] +async fn quic_reload_readiness_excludes_retired_sockets() { + use std::sync::Arc; + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.yaml"); + let socket = std::net::UdpSocket::bind("127.0.0.1:0").unwrap(); + let address = socket.local_addr().unwrap(); + drop(socket); + let echo = shoes_test_support::test_servers::start_tcp_stream_echo_server("127.0.0.1", 0) + .await + .unwrap(); + let certificate = rcgen::generate_simple_self_signed(vec!["localhost".into()]).unwrap(); + let config = serde_json::json!([{ + "address": address.to_string(), "transport": "quic", + "quic_settings": {"cert": certificate.cert.pem(), "key": certificate.signing_key.serialize_pem(), "num_endpoints": 2}, + "protocol": {"type": "forward", "target": echo.local_addr().to_string()} + }]); + std::fs::write(&path, serde_yaml::to_string(&config).unwrap()).unwrap(); + let mut child = Process::start(&path, &["--no-reload"]); + child.wait_for("Servers ready").await; + let mut roots = rustls::RootCertStore::empty(); + roots.add(certificate.cert.der().clone()).unwrap(); + let config = quinn::ClientConfig::with_root_certificates(Arc::new(roots)).unwrap(); + let connect = || async { + let mut endpoint = quinn::Endpoint::client("0.0.0.0:0".parse().unwrap()).unwrap(); + endpoint.set_default_client_config(config.clone()); + let connection = timeout( + Duration::from_secs(2), + endpoint.connect(address, "localhost").unwrap(), + ) + .await + .unwrap() + .unwrap(); + let (mut send, mut recv) = connection.open_bi().await.unwrap(); + send.write_all(b"ping").await.unwrap(); + let mut response = [0; 4]; + timeout(Duration::from_secs(2), recv.read_exact(&mut response)) + .await + .unwrap() + .unwrap(); + assert_eq!(&response, b"ping"); + (endpoint, connection, send, recv) + }; + let mut old = Vec::new(); + for _ in 0..8 { + old.push(connect().await); + } + child.signal(libc::SIGHUP); + child.wait_for("Servers ready").await; + for _ in 0..24 { + let (_endpoint, connection, _send, _recv) = connect().await; + connection.close(0u32.into(), b"done"); + } + for (_, connection, _, _) in &old { + assert!(matches!( + connection.close_reason(), + Some(quinn::ConnectionError::ApplicationClosed(_)) + )); + } + child.signal(libc::SIGTERM); + assert!(child.exit().await.success()); +} + +#[cfg(unix)] +#[tokio::test] +async fn symlink_targets_remain_watched_after_retargeting() { + let directory = tempfile::tempdir().unwrap(); + let targets = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.yaml"); + let first = targets.path().join("first.yaml"); + let second = targets.path().join("second.yaml"); + let config = format!( + "- address: '{}'\n protocol: {{type: http}}\n", + available_address() + ); + std::fs::write(&first, &config).unwrap(); + std::fs::write(&second, &config).unwrap(); + std::os::unix::fs::symlink(&first, &path).unwrap(); + let mut child = Process::start(&path, &[]); + child.wait_for("Servers ready").await; + std::fs::write(&first, &config).unwrap(); + child.wait_for("Servers ready").await; + let link = directory.path().join("replacement.yaml"); + std::os::unix::fs::symlink(&second, &link).unwrap(); + std::fs::rename(link, &path).unwrap(); + child.wait_for("Servers ready").await; + std::fs::write(&second, &config).unwrap(); + child.wait_for("Servers ready").await; + std::fs::remove_file(&second).unwrap(); + child.wait_for("Reload rejected").await; + std::fs::write(&second, &config).unwrap(); + child.wait_for("Servers ready").await; + let missing = targets.path().join("missing.yaml"); + let link = directory.path().join("replacement.yaml"); + std::os::unix::fs::symlink(&missing, &link).unwrap(); + std::fs::rename(link, &path).unwrap(); + child.wait_for("Reload rejected").await; + std::fs::write(&missing, &config).unwrap(); + child.wait_for("Servers ready").await; + let missing_directory = targets.path().join("missing-directory"); + let nested = missing_directory.join("config.yaml"); + let link = directory.path().join("replacement.yaml"); + std::os::unix::fs::symlink(&nested, &link).unwrap(); + std::fs::rename(link, &path).unwrap(); + child.wait_for("Reload rejected").await; + std::fs::create_dir(&missing_directory).unwrap(); + child.wait_for("Reload rejected").await; + std::fs::write(&nested, &config).unwrap(); + child.wait_for("Servers ready").await; + child.signal(libc::SIGTERM); + assert_eq!(child.exit().await.code(), Some(0)); +} + +#[cfg(unix)] +#[tokio::test] +async fn config_watch_directories_need_only_traverse_permission() { + use std::os::unix::fs::PermissionsExt; + struct RestorePermissions(std::path::PathBuf); + impl Drop for RestorePermissions { + fn drop(&mut self) { + let _ = std::fs::set_permissions(&self.0, std::fs::Permissions::from_mode(0o700)); + } + } + if unsafe { libc::geteuid() } == 0 { + return; + } + let directory = tempfile::tempdir().unwrap(); + let private = directory.path().join("private"); + let config_dir = private.join("config"); + std::fs::create_dir_all(&config_dir).unwrap(); + let alias = directory.path().join("alias"); + std::os::unix::fs::symlink(&config_dir, &alias).unwrap(); + let path = alias.join("config.yaml"); + let config = format!( + "- address: '{}'\n protocol: {{type: http}}\n", + available_address() + ); + std::fs::write(&path, &config).unwrap(); + for restricted in [&private, &config_dir] { + std::fs::set_permissions(restricted, std::fs::Permissions::from_mode(0o111)).unwrap(); + let _restore = RestorePermissions(restricted.clone()); + assert!(std::fs::read_dir(restricted).is_err()); + let mut child = Process::start(&path, &[]); + child.wait_for("Servers ready").await; + std::fs::write(&path, &config).unwrap(); + child.wait_for("Servers ready").await; + child.signal(libc::SIGTERM); + assert!(child.exit().await.success()); + } +} + +#[cfg(unix)] +#[tokio::test] +async fn shutdown_interrupts_dns_preparation_without_retiring_listeners() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.yaml"); + let echo = shoes_test_support::test_servers::start_tcp_stream_echo_server("127.0.0.1", 0) + .await + .unwrap(); + let blackhole = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let address = available_address(); + let valid = forward(address, echo.local_addr()); + std::fs::write(&path, &valid).unwrap(); + let mut child = Process::start(&path, &["--no-reload"]); + child.wait_for("Servers ready").await; + std::fs::write(&path, format!("{valid}\n- dns_group: pending\n dns_servers:\n - url: https://pending.test/dns-query\n bootstrap_url: udp://{}\n", blackhole.local_addr().unwrap())).unwrap(); + child.signal(libc::SIGHUP); + timeout(Duration::from_secs(3), blackhole.recv(&mut [0; 512])) + .await + .unwrap() + .unwrap(); + roundtrip(&mut TcpStream::connect(address).await.unwrap()).await; + child.signal(libc::SIGTERM); + assert_eq!( + timeout(Duration::from_secs(2), child.exit()) + .await + .unwrap() + .code(), + Some(0) + ); +} + +#[cfg(unix)] +#[tokio::test] +async fn replacement_bind_failure_is_fatal_and_cleans_partial_generation() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.yaml"); + let address = available_address(); + let occupied = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let config = format!("- address: '{address}'\n protocol: {{type: http}}\n"); + std::fs::write(&path, &config).unwrap(); + let mut child = Process::start(&path, &["--no-reload"]); + child.wait_for("Servers ready").await; + std::fs::write( + &path, + format!( + "{config}- address: '{}'\n protocol: {{type: http}}\n", + occupied.local_addr().unwrap() + ), + ) + .unwrap(); + child.signal(libc::SIGHUP); + assert_eq!(child.exit().await.code(), Some(1)); + std::net::TcpListener::bind(address).unwrap(); +} diff --git a/tests/support/Cargo.toml b/tests/support/Cargo.toml index 639e9bd9..5fea3c77 100644 --- a/tests/support/Cargo.toml +++ b/tests/support/Cargo.toml @@ -16,6 +16,7 @@ rand = "*" rcgen = { version = "*", default-features = false, features = ["aws_lc_rs", "pem"] } rustls = "*" rustls-pemfile = "*" +serde_json = "*" tempfile = "*" tokio = { version = "*", features = [ "fs", diff --git a/tests/support/src/test_fixture.rs b/tests/support/src/test_fixture.rs index 71a0735e..26d25257 100644 --- a/tests/support/src/test_fixture.rs +++ b/tests/support/src/test_fixture.rs @@ -3168,6 +3168,17 @@ impl ProxyTestFixture { let proxies_clone = self.proxies.clone(); for proxy in &proxies_clone { if let Some(config) = self.generate_singbox_config(proxy)? { + let mut config: serde_json::Value = serde_json::from_str(&config)?; + if self + .local_server + .as_ref() + .is_some_and(|server| server.ip == "127.0.0.1") + { + config["dns"] = serde_json::json!({"servers": [{ + "type": "hosts", "predefined": {"localhost": ["127.0.0.1"]} + }]}); + } + let config = serde_json::to_string(&config)?; let (guard, config_file) = start_singbox_server(&config)?; guards.push(guard); config_files.push(config_file);