From ced60d27e32cfbb7962d0fe91bb104cfd91cf3c7 Mon Sep 17 00:00:00 2001 From: cfal Date: Sun, 16 Nov 2025 20:20:01 +0800 Subject: [PATCH] 0.2.0 update --- CONFIG.md | 873 ++++--- Cargo.lock | 554 +++-- Cargo.toml | 18 +- README.md | 285 ++- examples/named_certificates.yaml | 41 +- examples/reality_basic.yaml | 39 + examples/reality_client.yaml | 45 + examples/reality_trojan.yaml | 56 + examples/reality_vision.yaml | 31 + examples/reality_vision_client.yaml | 31 + examples/vless_variants.yaml | 44 + src/address.rs | 65 +- src/async_stream.rs | 114 +- src/client_proxy_selector.rs | 2005 ++++++++++++++++- src/config/mod.rs | 403 +++- src/config/types.rs | 386 +++- src/copy_session_messages.rs | 287 +++ src/crypto/crypto_connection.rs | 296 +++ src/crypto/crypto_reader_writer.rs | 61 + src/crypto/crypto_tls_stream.rs | 344 +++ src/crypto/mod.rs | 13 + src/main.rs | 96 +- src/port_forward_handler.rs | 18 +- src/quic_server.rs | 43 +- src/reality/common.rs | 114 + src/reality/mod.rs | 33 + src/reality/reality_aead.rs | 365 +++ src/reality/reality_auth.rs | 552 +++++ src/reality/reality_certificate.rs | 178 ++ src/reality/reality_client_connection.rs | 865 +++++++ src/reality/reality_client_verify.rs | 314 +++ src/reality/reality_io_state.rs | 20 + src/reality/reality_reader_writer.rs | 117 + src/reality/reality_records.rs | 1594 +++++++++++++ src/reality/reality_server_connection.rs | 901 ++++++++ src/reality/reality_tls13_keys.rs | 429 ++++ src/reality/reality_tls13_messages.rs | 431 ++++ src/reality/reality_util.rs | 458 ++++ src/reality_client_handler.rs | 209 ++ src/{rustls_util.rs => rustls_config_util.rs} | 0 src/rustls_connection_util.rs | 77 + src/rustls_handshake.rs | 263 +++ src/shadow_tls/mod.rs | 4 +- src/shadow_tls/shadow_tls_client_handler.rs | 24 +- src/shadow_tls/shadow_tls_server_handler.rs | 34 +- src/shadowsocks/mod.rs | 2 + src/{ => shadowsocks}/salt_checker.rs | 0 src/shadowsocks/shadowsocks_cipher.rs | 20 +- src/shadowsocks/shadowsocks_stream.rs | 16 +- src/shadowsocks/shadowsocks_tcp_handler.rs | 85 +- src/{ => shadowsocks}/timed_salt_checker.rs | 2 +- src/slide_buffer.rs | 417 ++++ src/snell/snell_handler.rs | 4 +- src/socks_handler.rs | 79 + src/sync_adapter.rs | 56 + src/tcp/tcp_client_connector.rs | 6 +- src/tcp/tcp_handler.rs | 17 + src/tcp/tcp_handler_util.rs | 278 ++- src/tcp/tcp_server.rs | 52 + src/tls_client_handler.rs | 112 + src/tls_handler.rs | 210 -- src/tls_server_handler.rs | 397 ++++ src/trojan_handler.rs | 10 +- src/udp_session_message_stream.rs | 266 +++ src/uot/mod.rs | 41 + src/uot/uot_v1_stream.rs | 400 ++++ src/vless/mod.rs | 20 + src/vless/tls_deframer.rs | 599 +++++ src/vless/tls_fuzzy_deframer.rs | 531 +++++ src/vless/tls_handshake_util.rs | 241 ++ src/vless/vision_filter.rs | 512 +++++ src/vless/vision_pad.rs | 67 + src/vless/vision_stream.rs | 1480 ++++++++++++ src/vless/vision_unpad.rs | 547 +++++ src/vless/vless_client_handler.rs | 151 ++ src/{ => vless}/vless_message_stream.rs | 0 src/vless/vless_response_stream.rs | 137 ++ src/vless/vless_server_handler.rs | 327 +++ src/vless/vless_util.rs | 214 ++ src/vless_handler.rs | 324 --- src/vmess/md5.rs | 19 - src/vmess/typed.rs | 2 - src/vmess/vmess_handler.rs | 916 ++++---- src/vmess/vmess_stream.rs | 55 +- src/xudp/frame.rs | 846 +++++++ src/xudp/message_stream.rs | 621 +++++ src/xudp/mod.rs | 8 + 87 files changed, 21202 insertions(+), 2015 deletions(-) create mode 100644 examples/reality_basic.yaml create mode 100644 examples/reality_client.yaml create mode 100644 examples/reality_trojan.yaml create mode 100644 examples/reality_vision.yaml create mode 100644 examples/reality_vision_client.yaml create mode 100644 src/copy_session_messages.rs create mode 100644 src/crypto/crypto_connection.rs create mode 100644 src/crypto/crypto_reader_writer.rs create mode 100644 src/crypto/crypto_tls_stream.rs create mode 100644 src/crypto/mod.rs create mode 100644 src/reality/common.rs create mode 100644 src/reality/mod.rs create mode 100644 src/reality/reality_aead.rs create mode 100644 src/reality/reality_auth.rs create mode 100644 src/reality/reality_certificate.rs create mode 100644 src/reality/reality_client_connection.rs create mode 100644 src/reality/reality_client_verify.rs create mode 100644 src/reality/reality_io_state.rs create mode 100644 src/reality/reality_reader_writer.rs create mode 100644 src/reality/reality_records.rs create mode 100644 src/reality/reality_server_connection.rs create mode 100644 src/reality/reality_tls13_keys.rs create mode 100644 src/reality/reality_tls13_messages.rs create mode 100644 src/reality/reality_util.rs create mode 100644 src/reality_client_handler.rs rename src/{rustls_util.rs => rustls_config_util.rs} (100%) create mode 100644 src/rustls_connection_util.rs create mode 100644 src/rustls_handshake.rs rename src/{ => shadowsocks}/salt_checker.rs (100%) rename src/{ => shadowsocks}/timed_salt_checker.rs (97%) create mode 100644 src/slide_buffer.rs create mode 100644 src/sync_adapter.rs create mode 100644 src/tls_client_handler.rs delete mode 100644 src/tls_handler.rs create mode 100644 src/tls_server_handler.rs create mode 100644 src/udp_session_message_stream.rs create mode 100644 src/uot/mod.rs create mode 100644 src/uot/uot_v1_stream.rs create mode 100644 src/vless/mod.rs create mode 100644 src/vless/tls_deframer.rs create mode 100644 src/vless/tls_fuzzy_deframer.rs create mode 100644 src/vless/tls_handshake_util.rs create mode 100644 src/vless/vision_filter.rs create mode 100644 src/vless/vision_pad.rs create mode 100644 src/vless/vision_stream.rs create mode 100644 src/vless/vision_unpad.rs create mode 100644 src/vless/vless_client_handler.rs rename src/{ => vless}/vless_message_stream.rs (100%) create mode 100644 src/vless/vless_response_stream.rs create mode 100644 src/vless/vless_server_handler.rs create mode 100644 src/vless/vless_util.rs delete mode 100644 src/vless_handler.rs create mode 100644 src/xudp/frame.rs create mode 100644 src/xudp/message_stream.rs create mode 100644 src/xudp/mod.rs diff --git a/CONFIG.md b/CONFIG.md index 52dd19f8..12aab5a1 100644 --- a/CONFIG.md +++ b/CONFIG.md @@ -1,485 +1,658 @@ -# Configuration Guide +# Configuration Reference -shoes uses a YAML-based configuration format. Each configuration file can contain multiple entries of different types. +shoes uses YAML configuration files. Multiple configuration types can be combined in a single file or split across multiple files. -## Configuration Types +## Table of Contents +- [Configuration Structure](#configuration-structure) +- [Server Config](#server-config) +- [Server Protocols](#server-protocols) +- [Client Config](#client-config) +- [Client Protocols](#client-protocols) +- [Rules System](#rules-system) +- [Named Groups](#named-groups) +- [Named PEMs](#named-pems) +- [Advanced Features](#advanced-features) +- [Command Line](#command-line) -There are four main configuration types: +## Configuration Structure -1. Server configurations (`ServerConfig`) -2. Client proxy groups (`ClientConfigGroup`) -3. Rule groups (`RuleConfigGroup`) -4. Named PEM files (`NamedPem`) +A configuration file is a YAML array containing one or more configuration entries. Each entry can be: -## Named PEM Files - -You can define PEM files (containing certificates and/or private keys) once and reference them throughout your configuration. This is particularly useful when the same certificate/key pair is used in multiple places. - -### Named PEM -```yaml -pem: string # Name identifier for the PEM data -path: string # File path to PEM file -# OR -data: string # Inline PEM-encoded data -``` - -**Note**: A single PEM file can contain: -- Just a certificate -- Just a private key -- Both certificate and private key -- Multiple certificates (chain) -- Any combination of the above - -When referenced in a `cert` field, only the certificate portion is used. When referenced in a `key` field, only the private key portion is used. - -## Server Configuration - -A server configuration defines a proxy server instance: +- **Server Config** - Defines a proxy server instance +- **Client Config Group** - Defines reusable upstream proxy configurations +- **Rule Config Group** - Defines reusable routing rules +- **Named PEM** - Defines reusable certificate/key data ```yaml -bind_location: address | path # Network address or Unix socket path (use a file path for Unix sockets) -protocol: ServerProxyConfig # Server protocol configuration (see available types below) -transport: tcp | quic | udp # Optional, defaults to tcp; note: Unix socket binding is supported only with tcp -tcp_settings: # Optional TCP settings - no_delay: bool # Default: true -quic_settings: # Required if transport is quic - cert: string # TLS certificate (path, inline PEM data, or named PEM reference) - key: string # TLS private key (path, inline PEM data, or named PEM reference) - alpn_protocols: [string] # Optional ALPN protocols (alias: alpn_protocol) - client_fingerprints: [string] # Optional allowed client cert fingerprints (alias: client_fingerprint) - client_ca_certs: [string] # Optional client CA certificates (path, inline PEM, or named PEM reference) - num_endpoints: int # Optional; if set to 0, defaults to the number of available threads -rules: string | RuleConfig # Optional; defaults to allow-all-direct if omitted +# Server configs have 'address' or 'path' +- address: "0.0.0.0:8080" + protocol: ... + +# Client config groups have 'client_group' +- client_group: my-upstream + client_proxy: ... + +# Rule config groups have 'rule_group' +- rule_group: my-rules + rules: ... + +# Named PEMs have 'pem' +- pem: my-cert + path: /path/to/cert.pem ``` -## Protocol Types +## Server Config -### Server Protocols +```yaml +# Bind to IP address and port +address: "0.0.0.0:8080" # IPv4 +address: "[::]:8080" # IPv6 +address: "0.0.0.0:443-445" # Port range -#### HTTP/SOCKS5 Proxy +# OR bind to Unix socket (TCP only) +path: "/tmp/shoes.sock" + +# Protocol configuration (required) +protocol: ServerProxyConfig + +# Transport layer (default: tcp) +transport: tcp | quic + +# TCP settings (only when transport: tcp) +tcp_settings: + no_delay: true # Default: true + +# QUIC settings (required when transport: quic) +quic_settings: + cert: string # TLS certificate (path or named PEM) + key: string # TLS private key (path or named PEM) + alpn_protocols: [string] # Optional ALPN protocols + client_ca_certs: [string] # Optional client CA certificates + client_fingerprints: [string] # Optional client certificate fingerprints + num_endpoints: int # Optional, 0 = auto (based on thread count) + +# Routing rules (default: allow-all-direct) +rules: string | [RuleConfig] +``` + +## Server Protocols + +### HTTP ```yaml protocol: - type: http | socks - username: string? # Optional - password: string? # Optional + type: http + username: string? # Optional authentication + password: string? ``` -#### Shadowsocks +### SOCKS5 ```yaml protocol: - type: shadowsocks | ss - cipher: string # Encryption algorithm + type: socks # Aliases: socks5 + username: string? + password: string? +``` + +### Shadowsocks +```yaml +protocol: + type: shadowsocks # Aliases: ss + cipher: string # See supported ciphers below password: string + +# Supported ciphers: +# - aes-128-gcm +# - aes-256-gcm +# - chacha20-ietf-poly1305 +# - 2022-blake3-aes-128-gcm +# - 2022-blake3-aes-256-gcm +# - 2022-blake3-chacha20-ietf-poly1305 ``` -#### Snell +### VMess ```yaml protocol: - type: snell - cipher: string # Encryption algorithm - password: string - udp_enabled: bool # Optional, defaults to true - udp_num_sockets: int # Optional, defaults to 1 + type: vmess + cipher: string # aes-128-gcm, chacha20-poly1305, none + user_id: string # UUID + udp_enabled: true # Default: true (enables XUDP) ``` -#### VLESS +### VLESS ```yaml protocol: type: vless - user_id: string # UUID + user_id: string # UUID + udp_enabled: true # Default: true (enables XUDP) ``` -#### Trojan +### Trojan ```yaml protocol: type: trojan password: string - shadowsocks: # Optional additional encryption + shadowsocks: # Optional encryption layer cipher: string password: string ``` -#### VMess +### Snell v3 ```yaml protocol: - type: vmess - cipher: string - user_id: string # UUID - force_aead: bool # Default: true - udp_enabled: bool # Default: true + type: snell + cipher: string # aes-128-gcm, aes-256-gcm, chacha20-ietf-poly1305 + password: string + udp_enabled: true # Default: true + udp_num_sockets: 1 # Default: 1, sockets per UDP session ``` -#### TLS Server +### TLS Server ```yaml protocol: type: tls - sni_targets: # Map of SNI hostnames to TLS configs + + # Standard TLS targets (by SNI) + tls_targets: # Aliases: sni_targets, targets "example.com": - cert: string # Certificate (path, inline PEM data, or named PEM reference) - key: string # Private key (path, inline PEM data, or named PEM reference) - alpn_protocols: [string] # Optional ALPN protocols - client_fingerprints: [string] # Optional allowed client fingerprints - client_ca_certs: [string] # Optional client CA certificates - protocol: ServerProxyConfig # Inner protocol configuration - override_rules: string | [RuleConfig] # Optional override rules - default_target: # Optional default configuration + cert: string # Certificate (path or named PEM) + key: string # Private key (path or named PEM) + alpn_protocols: [string] # Optional ALPN + client_ca_certs: [string] # Optional client CA certs + client_fingerprints: [string] # Optional client cert fingerprints + vision: false # Enable Vision (requires VLESS inner protocol) + protocol: ServerProxyConfig + override_rules: [RuleConfig] # Optional rule override + + # Default TLS target (for unmatched/no SNI) + default_tls_target: # Aliases: default_target cert: string key: string - client_ca_certs: [string] - protocol: ServerProxyConfig - override_rules: string | [RuleConfig] - tls_buffer_size: int # Optional TLS buffer size (minimum 16384) -``` + # ... same fields as tls_targets -#### ShadowTLS v3 (in-process handshake) -```yaml -protocol: - type: shadowtls - sni_targets: # Map of SNI hostnames to ShadowTLS v3 configs + # Reality targets (by SNI) + reality_targets: + "www.cloudflare.com": + private_key: string # X25519 private key (base64url) + short_ids: [string] # Valid client IDs (hex, 0-16 chars) + dest: string # Fallback destination (e.g., "example.com:443") + max_time_diff: 60000 # Max timestamp diff in ms (default: 60000) + vision: false # Enable Vision (requires VLESS inner protocol) + protocol: ServerProxyConfig + override_rules: [RuleConfig] + + # ShadowTLS v3 targets (by SNI) + shadowtls_targets: "example.com": - password: string # ShadowTLS password - handshake: # TLS handshake configuration; can be defined in two ways: - cert: string # Local handshake: certificate (path, inline PEM, or named PEM reference) - key: string # Local handshake: private key (path, inline PEM, or named PEM reference) - alpn_protocols: [string] # Optional ALPN protocols - client_fingerprints: [string] # Optional allowed client fingerprints - protocol: ServerProxyConfig # Inner protocol configuration - override_rules: string | [RuleConfig] # Optional override rules - default_target: # Optional default configuration for ShadowTLS v3 - password: string - handshake: { ... } # Handshake configuration (local or remote) - protocol: ServerProxyConfig - override_rules: string | [RuleConfig] + password: string + handshake: + # Local handshake (with own certificate): + cert: string + key: string + alpn_protocols: [string] + client_ca_certs: [string] + client_fingerprints: [string] + # OR Remote handshake (proxy to real server): + address: string # e.g., "google.com:443" + client_proxies: [ClientConfig] # Optional proxies for handshake + protocol: ServerProxyConfig + override_rules: [RuleConfig] + + # Buffer size for TLS (optional, min 16384) + tls_buffer_size: int ``` -#### ShadowTLS v3 (remote handshake server) +### WebSocket ```yaml protocol: - type: shadowtls - sni_targets: # Map of SNI hostnames to ShadowTLS v3 configs - "example.com": - password: string # ShadowTLS password - handshake: # TLS handshake configuration; can be defined in two ways: - address: example.com:443 # Remote handshake: address of handshake server - client_proxies: ClientConfig # Remote handshake: client proxy configuration for handshake server - protocol: ServerProxyConfig # Inner protocol configuration - override_rules: string | [RuleConfig] # Optional override rules - default_target: # Optional default configuration for ShadowTLS v3 - handshake: { ... } # Handshake configuration (local or remote) - protocol: ServerProxyConfig - override_rules: string | [RuleConfig] -``` - -#### WebSocket -```yaml -protocol: - type: websocket | ws + type: websocket # Aliases: ws targets: - - matching_path: string? # Optional path to match - matching_headers: # Optional headers to match - header_name: string - protocol: ServerProxyConfig # Inner protocol configuration - ping_type: disabled | ping-frame | empty-frame # Default: ping-frame - override_rules: string | [RuleConfig] # Optional override rules + - matching_path: string? # Optional path filter (e.g., "/ws") + matching_headers: # Optional header filters + X-Custom-Header: "value" + protocol: ServerProxyConfig + ping_type: ping-frame # disabled | ping-frame | empty-frame + override_rules: [RuleConfig] ``` -#### Port Forward +### Port Forward ```yaml protocol: - type: forward | port_forward - targets: string | [string] # Target address(es) to forward to + type: forward # Aliases: port_forward, portforward + targets: string | [string] # Target address(es) ``` -#### Hysteria2 +### Hysteria2 ```yaml protocol: type: hysteria2 - password: string # Proxy password - udp_enabled: bool # Optional, defaults to true + password: string + udp_enabled: true # Default: true ``` -#### Tuic v5 +### TUIC v5 ```yaml protocol: - type: tuic | tuicv5 - uuid: string # UUID for identification + type: tuic # Aliases: tuicv5 + uuid: string # UUID password: string ``` -### Client Protocols +## Client Config -Client protocols (`ClientProxyConfig`) include all server protocols plus: +Used in rules to specify upstream proxies. -#### Direct Connection +```yaml +address: string # Proxy server address (e.g., "proxy.example.com:1080") +protocol: ClientProxyConfig +transport: tcp | quic # Default: tcp +bind_interface: string # Optional, Linux/Android/Fuchsia only + +tcp_settings: + no_delay: true + +quic_settings: + verify: true # Default: true + server_fingerprints: [string] + sni_hostname: string + alpn_protocols: [string] + cert: string # Client certificate for mTLS + key: string # Client key for mTLS +``` + +## Client Protocols + +### Direct ```yaml protocol: type: direct ``` -#### TLS Client +### HTTP +```yaml +protocol: + type: http + username: string? + password: string? +``` + +### SOCKS5 +```yaml +protocol: + type: socks + username: string? + password: string? +``` + +### Shadowsocks +```yaml +protocol: + type: shadowsocks + cipher: string + password: string +``` + +### Snell +```yaml +protocol: + type: snell + cipher: string + password: string +``` + +### VMess +```yaml +protocol: + type: vmess + cipher: string + user_id: string +``` + +### VLESS +```yaml +protocol: + type: vless + user_id: string +``` + +### Trojan +```yaml +protocol: + type: trojan + password: string + shadowsocks: # Optional + cipher: string + password: string +``` + +### TLS Client ```yaml protocol: type: tls - verify: bool # Default: true - server_fingerprints: [string] # Optional allowed server fingerprints - sni_hostname: string? # Optional SNI hostname - alpn_protocols: [string] # Optional ALPN protocols - key: string? # Optional client key (path, inline PEM, or named PEM reference) - cert: string? # Optional client cert (path, inline PEM, or named PEM reference) - protocol: ClientProxyConfig # Inner protocol configuration + verify: true # Default: true + server_fingerprints: [string] + sni_hostname: string + alpn_protocols: [string] + tls_buffer_size: int + cert: string # Client certificate for mTLS + key: string # Client key for mTLS + vision: false # Enable Vision (requires VLESS inner protocol) + protocol: ClientProxyConfig ``` -#### WebSocket Client +### Reality Client ```yaml protocol: - type: websocket | ws - matching_path: string? # Optional path to match - matching_headers: # Optional headers to match + type: reality + public_key: string # Server's X25519 public key (base64url) + short_id: string # Your client ID (hex, 0-16 chars) + sni_hostname: string # SNI to send (must match server's reality_targets key) + vision: false # Enable Vision (requires VLESS inner protocol) + protocol: ClientProxyConfig # Inner protocol (typically VLESS) +``` + +### ShadowTLS Client +```yaml +protocol: + type: shadowtls + password: string + sni_hostname: string? # Optional SNI override + protocol: ClientProxyConfig +``` + +### WebSocket Client +```yaml +protocol: + type: websocket + matching_path: string? + matching_headers: header_name: string - ping_type: disabled | ping-frame | empty-frame # Default: ping-frame - protocol: ClientProxyConfig # Inner protocol configuration + ping_type: ping-frame # disabled | ping-frame | empty-frame + protocol: ClientProxyConfig ``` -## Client Configuration - -A client configuration defines proxy client settings: - +### Port Forward (No-op) ```yaml -bind_interface: string? # Optional interface name (available on Linux, Android, or Fuchsia) -address: string # Target address; defaults to unspecified if omitted -protocol: ClientProxyConfig # Client protocol configuration -transport: tcp | quic | udp # Optional, defaults to tcp -tcp_settings: # Optional TCP settings - no_delay: bool # Default: true -quic_settings: # Optional QUIC settings (only applicable if transport is quic) - verify: bool # Default: true - server_fingerprints: [string] # Optional allowed server fingerprints (alias: server_fingerprint) - sni_hostname: string? # Optional SNI hostname - alpn_protocols: [string] # Optional ALPN protocols (alias: alpn_protocol) - key: string? # Optional client key (path, inline PEM, or named PEM reference; must be paired with cert) - cert: string? # Optional client cert (path, inline PEM, or named PEM reference; must be paired with key) +protocol: + type: portforward # Aliases: noop ``` -## Client Proxy Groups +Passes through the raw connection without protocol wrapping. Useful for testing or transparent proxying. -Client proxy groups allow defining reusable proxy configurations: +## Rules System +Rules determine how incoming connections are routed. + +### Rule Config ```yaml -client_group: string -client_proxies: ClientConfig | [ClientConfig] +rules: + - masks: string | [string] # IP/CIDR or hostname masks + action: allow | block + # For action: allow + override_address: string? # Optional address override + client_proxy: string | ClientConfig | [ClientConfig] ``` -## Rule Groups - -Rule groups define access control and routing rules: - +### Mask Syntax ```yaml -rule_group: string -rules: RuleConfig | [RuleConfig] +# IP/CIDR masks +masks: "0.0.0.0/0" # All IPv4 +masks: "::/0" # All IPv6 +masks: "192.168.0.0/16" # Subnet +masks: "10.0.0.1:80" # Specific IP and port + +# Hostname masks +masks: "*.google.com" # Wildcard subdomain +masks: "example.com" # Exact match + +# Multiple masks +masks: + - "192.168.0.0/16" + - "10.0.0.0/8" + - "*.internal.com" ``` -### Rule Configuration +### Built-in Rule Groups +- `allow-all-direct` - Allow all connections, direct routing +- `block-all` - Block all connections +### Example Rules ```yaml -masks: string | [string] # IP/CIDR masks to match -action: allow | block # Action to take -override_address: string? # Optional address override for allow action -client_proxies: string | ClientConfig | [string | ClientConfig] # Required for allow action (alias: client_proxy) +rules: + # Direct connection for local networks + - masks: ["192.168.0.0/16", "10.0.0.0/8", "172.16.0.0/12"] + action: allow + client_proxy: direct + + # Block specific domains + - masks: ["*.ads.example.com", "tracking.example.com"] + action: block + + # Route through upstream proxy + - masks: "0.0.0.0/0" + action: allow + client_proxy: + address: "proxy.example.com:1080" + protocol: + type: socks ``` -## Built-in Defaults - -The system includes these built-in defaults: - -### Client Groups -- `direct`: Direct connections without proxy - -### Rule Groups -- `allow-all-direct`: Allows all traffic directly -- `block-all`: Blocks all traffic - -## Security Considerations - -1. TLS/QUIC Security: - - Use strong certificates and private keys - - Consider enabling client certificate authentication - - Use secure cipher suites - - Verify certificate fingerprints when possible - -2. Authentication: - - Use strong passwords for all authentication methods - - Consider using client certificates where supported - - Rotate credentials regularly - -3. Network Security: - - Be cautious with 0.0.0.0 bind addresses - - Use firewalls to restrict access - - Consider binding to specific interfaces when possible - - Monitor logs for suspicious activity - -4. WebSocket Security: - - Use path and header matching to restrict access - - Consider using TLS for transport security - - Implement rate limiting if needed - -## Certificate and Key Specification - -Certificates and keys can be specified in three ways: - -1. **File Path**: Provide the path to a PEM-encoded file - ```yaml - cert: "/etc/certs/server.pem" - key: "/etc/certs/server.pem" # Can be the same file if it contains both - ``` - -2. **Inline PEM Data**: Embed the PEM data directly in the configuration - ```yaml - cert: | - -----BEGIN CERTIFICATE----- - MIIDXTCCAkWgAwIBAgIJAKl... - -----END CERTIFICATE----- - key: | - -----BEGIN PRIVATE KEY----- - MIIEvQIBADANBgkqhkiG9w0... - -----END PRIVATE KEY----- - ``` - -3. **Named PEM Reference**: Reference a previously defined named PEM - ```yaml - cert: "my-server-pem" # References a named PEM (certificate portion will be used) - key: "my-server-pem" # References a named PEM (private key portion will be used) - ``` - -**Important**: When a PEM file contains both certificate and private key, you can reference it in both `cert` and `key` fields. The system will automatically extract the appropriate portion. - -## Examples - -### Using Named PEM Files +## Named Groups +### Client Proxy Group ```yaml -# Define named PEM files -- pem: "server-pem" - path: "/etc/certs/server.pem" # Contains both cert and key +- client_group: my-upstream + client_proxy: # Or client_proxies for multiple + - address: "proxy1.example.com:1080" + protocol: + type: socks + - address: "proxy2.example.com:1080" + protocol: + type: socks -- pem: "ca-cert" - data: | - -----BEGIN CERTIFICATE----- - MIIDXTCCAkWgAwIBAgIJAKl... - -----END CERTIFICATE----- - -# Use named PEM files in server configurations -- bind_location: "0.0.0.0:443" - transport: quic - quic_settings: - cert: "server-pem" # Reference named PEM (cert portion) - key: "server-pem" # Reference named PEM (key portion) - client_ca_certs: - - "ca-cert" # Reference named CA certificate +# Reference in rules +- address: "0.0.0.0:8080" protocol: - type: socks + type: http + rules: + - masks: "0.0.0.0/0" + action: allow + client_proxy: my-upstream # Reference by name +``` -- bind_location: "0.0.0.0:8443" +### Rule Group +```yaml +- rule_group: standard-rules + rules: + - masks: ["192.168.0.0/16"] + action: allow + client_proxy: direct + - masks: "0.0.0.0/0" + action: allow + client_proxy: my-upstream + +# Reference in server config +- address: "0.0.0.0:8080" + protocol: + type: http + rules: standard-rules # Reference by name +``` + +## Named PEMs + +Define certificates once and reference throughout configuration. + +```yaml +# From file +- pem: my-cert + path: /path/to/certificate.pem + +# Inline data +- pem: my-key + data: | + -----BEGIN PRIVATE KEY----- + ... + -----END PRIVATE KEY----- + +# Reference in config +- address: "0.0.0.0:443" protocol: type: tls - sni_targets: + tls_targets: "example.com": - cert: "server-pem" # Reuse the same PEM file - key: "server-pem" # Reuse the same PEM file + cert: my-cert # Reference by name + key: my-key protocol: type: http ``` -### Basic HTTP Proxy Server -```yaml -bind_location: "127.0.0.1:8080" -protocol: - type: http - username: user - password: pass -``` +## Advanced Features + +### Vision (XTLS-Vision) + +Vision optimizes TLS-in-TLS scenarios by detecting inner TLS traffic and switching to direct mode for zero-copy performance. + +**Requirements:** +- Inner protocol MUST be VLESS +- Works with both TLS and Reality -### SOCKS5 with TLS ```yaml -bind_location: "0.0.0.0:1080" +# TLS + Vision protocol: type: tls - sni_targets: - "proxy.example.com": - cert: "cert.pem" - key: "key.pem" + tls_targets: + "example.com": + cert: cert.pem + key: key.pem + vision: true + alpn_protocols: ["http/1.1"] protocol: - type: socks - username: user - password: pass + type: vless + user_id: "uuid" + +# Reality + Vision +protocol: + type: tls + reality_targets: + "www.google.com": + private_key: "..." + short_ids: ["..."] + dest: "www.google.com:443" + vision: true + protocol: + type: vless + user_id: "uuid" ``` -### VMess over WebSocket +### XUDP Multiplexing + +Automatically enabled for VMess and VLESS when `udp_enabled: true`. Multiplexes UDP traffic over a single connection. + +### Proxy Chaining + +Chain multiple proxies by nesting client protocols: + ```yaml -bind_location: "0.0.0.0:443" -protocol: - type: websocket - targets: - - matching_path: "/vmess" +client_proxy: + address: "proxy1.example.com:1080" + protocol: + type: socks + protocol: + type: tls protocol: type: vmess - cipher: auto - user_id: "123e4567-e89b-12d3-a456-426614174000" + cipher: aes-128-gcm + user_id: "uuid" ``` -### Complex Routing Setup +### Hot Reloading + +Configuration changes are automatically detected and applied without restarting. Disable with `--no-reload` flag. + +### mTLS (Mutual TLS) + +Require client certificates for authentication: + ```yaml -# Define client proxies -client_group: "proxies" -client_proxies: - - protocol: - type: shadowsocks - cipher: chacha20-ietf-poly1305 - password: secret1 - - protocol: - type: vmess - cipher: auto - user_id: "123e4567-e89b-12d3-a456-426614174000" - -# Define routing rules -rule_group: "routing" -rules: - - masks: "192.168.0.0/16" - action: allow - client_proxy: direct - - masks: "0.0.0.0/0" - action: allow - client_proxy: proxies - -# Main server config -bind_location: "0.0.0.0:8080" +# Server side protocol: - type: http -rules: routing + type: tls + tls_targets: + "example.com": + cert: server.crt + key: server.key + client_ca_certs: [ca.crt] # Required CA + client_fingerprints: ["sha256:..."] # Optional specific certs + protocol: ... + +# Client side +client_proxy: + address: "example.com:443" + protocol: + type: tls + cert: client.crt + key: client.key + protocol: ... ``` -### Hysteria2 Proxy Server -```yaml -bind_location: "0.0.0.0:4443" -protocol: - type: hysteria2 - password: "hysteria_secret" - udp_enabled: true +## Command Line + +```bash +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 + +COMMANDS: + generate-reality-keypair Generate Reality X25519 keypair + generate-shadowsocks-password Generate Shadowsocks password ``` -### Tuic v5 Proxy Server -```yaml -bind_location: "0.0.0.0:5555" -protocol: - type: tuicv5 - uuid: "123e4567-e89b-12d3-a456-426614174000" - password: "tuic_secret" +## Tips + +### Generate Keys + +**Reality keypair:** +```bash +shoes generate-reality-keypair ``` -## Advanced Configuration Details +**Shadowsocks password:** +```bash +shoes generate-shadowsocks-password 2022-blake3-aes-256-gcm +``` -- Multiple configuration blocks can be combined in a single YAML file. -- Fields support aliases (e.g., "alpn_protocol" and "alpn_protocols", "client_proxy" for "client_proxies"). -- For QUIC transport, if "num_endpoints" is set to 0, it defaults to the number of available CPU threads. -- Unix domain socket binding is supported only when using TCP transport. -- Client configurations require both certificate and key to be specified together. -- UUIDs in Vless and Vmess protocols are validated for correct format. +**UUID:** +```bash +uuidgen +``` + +**TLS certificate fingerprint:** +```bash +openssl x509 -in cert.pem -noout -fingerprint -sha256 +``` + +### Security Best Practices + +- Use strong, random passwords +- Keep private keys secure +- Use `127.0.0.1` instead of `0.0.0.0` for local-only access +- Use firewall rules to restrict access +- Enable client certificate authentication for sensitive services +- Use Vision with Reality for maximum privacy + +### Performance Tips + +- Enable `vision: true` for TLS-in-TLS scenarios +- Use `tcp_settings.no_delay: true` for low latency +- Set `quic_settings.num_endpoints` to match worker threads +- Use QUIC transport for high-latency or lossy networks + +### Common Issues + +- **"Address already in use"**: Change port or stop conflicting service +- **"Permission denied"**: Ports < 1024 require root/admin +- **Reality connection fails**: Verify keys match, UUID matches, SNI matches server's reality_targets key +- **Vision not working**: Ensure inner protocol is VLESS +- **Config validation fails**: Run with `--dry-run` for detailed errors diff --git a/Cargo.lock b/Cargo.lock index 70feb889..dc52f83f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -22,6 +22,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + [[package]] name = "anstream" version = "0.6.21" @@ -54,22 +60,22 @@ dependencies = [ [[package]] name = "anstyle-query" -version = "1.1.4" +version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e231f6134f61b71076a3eab506c379d4f36122f2af15a9ff04415ea4c3339e2" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] name = "anstyle-wincon" -version = "3.0.10" +version = "3.0.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3e0633414522a32ffaac8ac6cc8f748e090c5717661fddeea04219e2344f5f2a" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -96,6 +102,45 @@ version = "0.7.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" +[[package]] +name = "asn1-rs" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56624a96882bb8c26d61312ae18cb45868e5a9992ea73c58e45c3101e56a1e60" +dependencies = [ + "asn1-rs-derive", + "asn1-rs-impl", + "displaydoc", + "nom", + "num-traits", + "rusticata-macros", + "thiserror 2.0.17", + "time", +] + +[[package]] +name = "asn1-rs-derive" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3109e49b1e4909e9db6515a30c633684d68cdeaa252f215214cb4fa1a5bfee2c" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "asn1-rs-impl" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b18050c2cd6fe86c3a76584ef5e0baf286d038cda203eb6223df2cc413565f7" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "async-trait" version = "0.1.89" @@ -108,10 +153,22 @@ dependencies = [ ] [[package]] -name = "aws-lc-rs" -version = "1.14.1" +name = "atomic-waker" +version = "1.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "879b6c89592deb404ba4dc0ae6b58ffd1795c78991cbb5b8bc441c48a070440d" +checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" + +[[package]] +name = "autocfg" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" + +[[package]] +name = "aws-lc-rs" +version = "1.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b5ce75405893cd713f9ab8e297d8e438f624dde7d706108285f7e17a25a180f" dependencies = [ "aws-lc-sys", "zeroize", @@ -119,11 +176,10 @@ dependencies = [ [[package]] name = "aws-lc-sys" -version = "0.32.3" +version = "0.34.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "107a4e9d9cab9963e04e84bb8dee0e25f2a987f9a8bad5ed054abd439caa8f8c" +checksum = "179c3777a8b5e70e90ea426114ffc565b2c1a9f82f6c4a0c5a34aa6ef5e781b6" dependencies = [ - "bindgen", "cc", "cmake", "dunce", @@ -142,26 +198,6 @@ version = "1.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "55248b47b0caf0546f7988906588779981c43bb1bc9d0c44087278f80cdb44ba" -[[package]] -name = "bindgen" -version = "0.72.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "993776b509cfb49c750f11b8f07a46fa23e0a1386ffc01fb1e7d343efc387895" -dependencies = [ - "bitflags 2.10.0", - "cexpr", - "clang-sys", - "itertools", - "log", - "prettyplease", - "proc-macro2", - "quote", - "regex", - "rustc-hash 2.1.1", - "shlex", - "syn", -] - [[package]] name = "bitflags" version = "1.3.2" @@ -213,15 +249,15 @@ checksum = "46c5e41b57b8bba42a04676d81cb89e9ee8e859a1a66f80a5a72e1cb76b34d43" [[package]] name = "bytes" -version = "1.10.1" +version = "1.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d71b6127be86fdcfddb610f7182ac57211d4b18a3e9c82eb2d17662f2227ad6a" +checksum = "b35204fbdc0b3f4446b89fc1ac2cf84a8a68971995d0bf2e925ec7cd960f9cb3" [[package]] name = "cc" -version = "1.2.44" +version = "1.2.47" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "37521ac7aabe3d13122dc382493e20c9416f299d2ccd5b3a5340a2570cdeb0f3" +checksum = "cd405d82c84ff7f35739f175f67d8b9fb7687a0e84ccdc78bd3568839827cf07" dependencies = [ "find-msvc-tools", "jobserver", @@ -235,24 +271,6 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6d43a04d8753f35258c91f8ec639f792891f748a1edbd759cf1dcea3382ad83c" -[[package]] -name = "cexpr" -version = "0.6.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766" -dependencies = [ - "nom", -] - -[[package]] -name = "cfb-mode" -version = "0.8.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "738b8d467867f80a71351933f70461f5b56f24d5c93e0cf216e59229c968d330" -dependencies = [ - "cipher", -] - [[package]] name = "cfg-if" version = "1.0.4" @@ -275,17 +293,6 @@ dependencies = [ "inout", ] -[[package]] -name = "clang-sys" -version = "1.8.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4" -dependencies = [ - "glob", - "libc", - "libloading", -] - [[package]] name = "cmake" version = "0.1.54" @@ -350,9 +357,9 @@ checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" [[package]] name = "crypto-common" -version = "0.1.6" +version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1bfb12502f3fc46cca1bb51ac28df9d618d813cdc3d2f25b9fe775a34af26bb3" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" dependencies = [ "generic-array", "typenum", @@ -372,6 +379,35 @@ dependencies = [ "parking_lot_core", ] +[[package]] +name = "data-encoding" +version = "2.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a2330da5de22e8a3cb63252ce2abb30116bf5265e89c0e01bc17015ce30a476" + +[[package]] +name = "der-parser" +version = "10.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07da5016415d5a3c4dd39b11ed26f915f52fc4e0dc197d87908bc916e51bc1a6" +dependencies = [ + "asn1-rs", + "displaydoc", + "nom", + "num-bigint", + "num-traits", + "rusticata-macros", +] + +[[package]] +name = "deranged" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ececcb659e7ba858fb4f10388c250a7252eb0a27373f1a72b8748afdd248e587" +dependencies = [ + "powerfmt", +] + [[package]] name = "digest" version = "0.10.7" @@ -383,18 +419,23 @@ dependencies = [ "subtle", ] +[[package]] +name = "displaydoc" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "dunce" version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" -[[package]] -name = "either" -version = "1.15.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" - [[package]] name = "env_filter" version = "0.1.4" @@ -442,15 +483,15 @@ checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be" [[package]] name = "find-msvc-tools" -version = "0.1.4" +version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52051878f80a721bb68ebfbc930e07b65ba72f2da88968ea5c06fd6ca3d3a127" +checksum = "3a3076410a55c90011c298b04d0cfa770b00fa04e1e3c97d3f6c9de105a03844" [[package]] -name = "fnv" -version = "1.0.7" +name = "foldhash" +version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" [[package]] name = "fs_extra" @@ -558,9 +599,9 @@ dependencies = [ [[package]] name = "generic-array" -version = "0.14.9" +version = "0.14.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4bb6743198531e02858aeaea5398fcc883e71851fcbcb5a2f773e2fb6cb1edf2" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" dependencies = [ "typenum", "version_check", @@ -593,12 +634,6 @@ dependencies = [ "wasm-bindgen", ] -[[package]] -name = "glob" -version = "0.3.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" - [[package]] name = "h3" version = "0.0.8" @@ -635,38 +670,104 @@ checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" [[package]] name = "hashbrown" -version = "0.16.0" +version = "0.16.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5419bdc4f6a9207fbeba6d11b604d481addf78ecd10c11ad51e76c2f6482748d" - -[[package]] -name = "hmac" -version = "0.12.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" dependencies = [ - "digest", + "allocator-api2", + "equivalent", + "foldhash", ] [[package]] name = "http" -version = "1.3.1" +version = "1.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f4a85d31aea989eead29a3aaf9e1115a180df8282431156e533de47660892565" +checksum = "e3ba2a386d7f85a81f119ad7498ebe444d2e22c2af0b86b069416ace48b3311a" dependencies = [ "bytes", - "fnv", "itoa", ] [[package]] -name = "indexmap" -version = "2.12.0" +name = "http-body" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6717a8d2a5a929a1a2eb43a12812498ed141a0bcfb7e8f7844fbdbe4303bba9f" +checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" +dependencies = [ + "bytes", + "http", +] + +[[package]] +name = "http-body-util" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "pin-project-lite", +] + +[[package]] +name = "httparse" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" + +[[package]] +name = "httpdate" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" + +[[package]] +name = "hyper" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ab2d4f250c3d7b1c9fcdff1cece94ea4e2dfbec68614f7b87cb205f24ca9d11" +dependencies = [ + "atomic-waker", + "bytes", + "futures-channel", + "futures-core", + "http", + "http-body", + "httparse", + "httpdate", + "itoa", + "pin-project-lite", + "pin-utils", + "smallvec", + "tokio", +] + +[[package]] +name = "hyper-util" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52e9a2a24dc5c6821e71a7030e1e14b7b632acac55c40e9d2e082c621261bb56" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "hyper", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "indexmap" +version = "2.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ad4bb2b565bca0645f4d68c5c9af97fba094e9791da685bf83cb5f3ce74acf2" dependencies = [ "equivalent", - "hashbrown 0.16.0", + "hashbrown 0.16.1", ] [[package]] @@ -704,15 +805,6 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" -[[package]] -name = "itertools" -version = "0.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" -dependencies = [ - "either", -] - [[package]] name = "itoa" version = "1.0.15" @@ -721,22 +813,22 @@ checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c" [[package]] name = "jiff" -version = "0.2.15" +version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "be1f93b8b1eb69c77f24bbb0afdf66f54b632ee39af40ca21c4365a1d7347e49" +checksum = "49cce2b81f2098e7e3efc35bc2e0a6b7abec9d34128283d7a26fa8f32a6dbb35" dependencies = [ "jiff-static", "log", "portable-atomic", "portable-atomic-util", - "serde", + "serde_core", ] [[package]] name = "jiff-static" -version = "0.2.15" +version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "03343451ff899767262ec32146f6d559dd759fdadf42ff0e227c7c48f72594b4" +checksum = "980af8b43c3ad5d8d349ace167ec8170839f753a42d233ba19e08afe1850fa69" dependencies = [ "proc-macro2", "quote", @@ -814,22 +906,18 @@ dependencies = [ "libc", ] +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + [[package]] name = "libc" version = "0.2.177" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2874a2af47a2325c2001a6e6fad9b16a53b802102b528163885171cf92b15976" -[[package]] -name = "libloading" -version = "0.8.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" -dependencies = [ - "cfg-if", - "windows-link", -] - [[package]] name = "linux-raw-sys" version = "0.11.0" @@ -851,6 +939,15 @@ version = "0.4.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "34080505efa8e45a4b816c349525ebe327ceaa8559756f0356cba97ef3bf7432" +[[package]] +name = "lru" +version = "0.16.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96051b46fc183dc9cd4a223960ef37b9af631b55191852a8274bfef064cda20f" +dependencies = [ + "hashbrown 0.16.1", +] + [[package]] name = "lru-slab" version = "0.1.2" @@ -925,6 +1022,49 @@ version = "2.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e0826a989adedc2a244799e823aece04662b66609d96af8dff7ac6df9a8925d" +[[package]] +name = "num-bigint" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" +dependencies = [ + "num-integer", + "num-traits", +] + +[[package]] +name = "num-conv" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51d515d32fb182ee37cda2ccdcb92950d6a3c2893aa280e540671c2cd0f3b1d9" + +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "oid-registry" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "12f40cff3dde1b6087cc5d5f5d4d65712f34016a03ed60e9c08dcc392736b5b7" +dependencies = [ + "asn1-rs", +] + [[package]] name = "once_cell" version = "1.21.3" @@ -977,6 +1117,16 @@ dependencies = [ "subtle", ] +[[package]] +name = "pem" +version = "3.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d30c53c26bc5b31a98cd02d20f25a7c8567146caf63ed593a9d87b2775291be" +dependencies = [ + "base64", + "serde_core", +] + [[package]] name = "percent-encoding" version = "2.3.2" @@ -1010,6 +1160,12 @@ dependencies = [ "portable-atomic", ] +[[package]] +name = "powerfmt" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" + [[package]] name = "ppv-lite86" version = "0.2.21" @@ -1019,16 +1175,6 @@ dependencies = [ "zerocopy", ] -[[package]] -name = "prettyplease" -version = "0.2.37" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" -dependencies = [ - "proc-macro2", - "syn", -] - [[package]] name = "proc-macro2" version = "1.0.103" @@ -1098,9 +1244,9 @@ dependencies = [ [[package]] name = "quote" -version = "1.0.41" +version = "1.0.42" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce25767e7b499d1b604768e7cde645d14cc8584231ea6b295e9c9eb22c02e1d1" +checksum = "a338cc41d27e6cc6dce6cefc13a0729dfbb81c262b1f519331575dd80ef3067f" dependencies = [ "proc-macro2", ] @@ -1146,6 +1292,20 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rcgen" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5fae430c6b28f1ad601274e78b7dffa0546de0b73b4cd32f46723c0c2a16f7a5" +dependencies = [ + "pem", + "ring", + "rustls-pki-types", + "time", + "x509-parser", + "yasna", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -1210,6 +1370,15 @@ version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d" +[[package]] +name = "rusticata-macros" +version = "4.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "faf0c4a6ece9950b9abdb62b1cfcf2a68b3b67a10ba445b3bb85be2a293d0632" +dependencies = [ + "nom", +] + [[package]] name = "rustix" version = "1.1.2" @@ -1225,9 +1394,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.34" +version = "0.23.35" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a9586e9ee2b4f8fab52a0048ca7334d7024eef48e2cb9407e3497bb7cab7fa7" +checksum = "533f54bc6a7d4f647e46ad909549eda97bf5afc1585190ef692b4286b198bd8f" dependencies = [ "aws-lc-rs", "log", @@ -1250,6 +1419,15 @@ dependencies = [ "security-framework", ] +[[package]] +name = "rustls-pemfile" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dce314e5fee3f39953d46bb63bb8a46d40c2f8fb7cc5a3b6cab2bde9721d6e50" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "rustls-pki-types" version = "1.13.0" @@ -1419,7 +1597,7 @@ checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" [[package]] name = "shoes" -version = "0.1.8" +version = "0.2.0" dependencies = [ "aes", "argon2", @@ -1428,16 +1606,18 @@ dependencies = [ "base64", "blake3", "bytes", - "cfb-mode", "dashmap", "digest", "env_logger", "futures", "h3", "h3-quinn", - "hmac", "http", + "http-body-util", + "hyper", + "hyper-util", "log", + "lru", "md-5", "memchr", "notify", @@ -1445,17 +1625,22 @@ dependencies = [ "percent-encoding", "quinn", "rand", + "rand_core 0.6.4", + "rcgen", "rustc-hash 1.1.0", "rustls", + "rustls-pemfile", "serde", "serde_yaml", "sha3", "socket2", + "subtle", "tempfile", "tikv-jemallocator", "tokio", "tokio-rustls", "webpki-roots", + "x509-parser", ] [[package]] @@ -1488,15 +1673,26 @@ checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" [[package]] name = "syn" -version = "2.0.108" +version = "2.0.111" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da58917d35242480a05c2897064da0a80589a2a0476c9a3f2fdc83b53502e917" +checksum = "390cc9a294ab71bdb1aa2e99d13be9c753cd2d7bd6560c77118597410c4d2e87" dependencies = [ "proc-macro2", "quote", "unicode-ident", ] +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "tempfile" version = "3.23.0" @@ -1570,6 +1766,37 @@ dependencies = [ "tikv-jemalloc-sys", ] +[[package]] +name = "time" +version = "0.3.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e7d9e3bb61134e77bde20dd4825b97c010155709965fedf0f49bb138e52a9d" +dependencies = [ + "deranged", + "itoa", + "num-conv", + "powerfmt", + "serde", + "time-core", + "time-macros", +] + +[[package]] +name = "time-core" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40868e7c1d2f0b8d73e4a8c7f0ff63af4f6d19be117e90bd73eb1d62cf831c6b" + +[[package]] +name = "time-macros" +version = "0.2.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30cfb0125f12d9c277f35663a0a33f8c30190f4e4574868a330595412d34ebf3" +dependencies = [ + "num-conv", + "time-core", +] + [[package]] name = "tinyvec" version = "1.10.0" @@ -2032,19 +2259,46 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f17a85883d4e6d00e8a97c586de764dabcc06133f7f1d55dce5cdc070ad7fe59" [[package]] -name = "zerocopy" -version = "0.8.27" +name = "x509-parser" +version = "0.18.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0894878a5fa3edfd6da3f88c4805f4c8558e2b996227a3d864f47fe11e38282c" +checksum = "eb3e137310115a65136898d2079f003ce33331a6c4b0d51f1531d1be082b6425" +dependencies = [ + "asn1-rs", + "data-encoding", + "der-parser", + "lazy_static", + "nom", + "oid-registry", + "ring", + "rusticata-macros", + "thiserror 2.0.17", + "time", +] + +[[package]] +name = "yasna" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e17bb3549cc1321ae1296b9cdc2698e2b6cb1992adfa19a8c72e5b7a738f44cd" +dependencies = [ + "time", +] + +[[package]] +name = "zerocopy" +version = "0.8.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ea879c944afe8a2b25fef16bb4ba234f47c694565e97383b36f3a878219065c" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.27" +version = "0.8.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "88d2b8d9c68ad2b9e4340d7832716a4d21a22a1154777ad56ea55c51a9cf3831" +checksum = "cf955aa904d6040f70dc8e9384444cb1030aed272ba3cb09bbc4ab9e7c1f34f5" dependencies = [ "proc-macro2", "quote", diff --git a/Cargo.toml b/Cargo.toml index faaad147..ef7d9a6a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "shoes" -version = "0.1.8" +version = "0.2.0" edition = "2021" license = "MIT" description = "A multi-protocol proxy server." @@ -18,16 +18,15 @@ aws-lc-rs = { version = "*", default-features = false } base64 = "*" blake3 = "*" bytes = "*" -cfb-mode = "*" dashmap = "*" digest = "*" env_logger = "*" futures = "*" h3 = "*" h3-quinn = "*" -hmac = "*" http = "*" -log = "*" +log = { version = "*", features = ["release_max_level_info"] } +lru = "*" md-5 = "*" memchr = "*" notify = "*" @@ -35,9 +34,12 @@ parking_lot = "*" percent-encoding = "*" quinn = { version = "*", default-features = false, features = ["log", "platform-verifier", "runtime-tokio", "rustls-aws-lc-rs"] } rand = "*" +rand_core = "*" +rcgen = { version = "*", features = ["x509-parser"] } rustc-hash = "1" rustls = "*" serde = { version = "*", features = ["derive", "std"] } +subtle = "*" serde_yaml = "*" sha3 = "*" socket2 = "*" @@ -50,11 +52,17 @@ tokio = { version = "*", features = [ "sync", "time", ] } -tokio-rustls = { version = "*", default-features = false, features = ["logging", "tls12", "aws-lc-rs"] } webpki-roots = { version = "*" } +x509-parser = "*" [dev-dependencies] tempfile = "*" +hyper = { version = "1", features = ["server", "http1"] } +hyper-util = { version = "0.1", features = ["tokio", "server"] } +http-body-util = "0.1" +rustls-pemfile = "*" +rcgen = "*" +tokio-rustls = "*" [target.'cfg(not(target_env = "msvc"))'.dependencies] tikv-jemallocator = "0.6" diff --git a/README.md b/README.md index f98b6979..b1edcd3f 100644 --- a/README.md +++ b/README.md @@ -1,123 +1,52 @@ # shoes -shoes is a multi-protocol proxy server written in Rust. +shoes is a high-performance multi-protocol proxy server written in Rust. -## Supported protocols +## Supported Protocols -- **HTTP/HTTPS** (TCP, QUIC) -- **SOCKS5** (TCP, QUIC) -- **Vmess** (TCP, QUIC, UDP-over-TCP) - - AEAD and Legacy modes - - Supported ciphers: - - aes-128-gcm - - chacha20-poly1305 -- **Vless** (TCP, QUIC, UDP-over-TCP) -- **Snell** v3 (TCP, QUIC, UDP-over-TCP) - - Supported ciphers: - - aes-128-gcm - - aes-256-gcm - - chacha20-ietf-poly1305 -- **Shadowsocks** (TCP, QUIC) - - Supported ciphers: - - aes-128-gcm - - aes-256-gcm - - chacha20-ietf-poly1305 - - 2022-blake3-aes-128-gcm - - 2022-blake3-aes-256-gcm - - 2022-blake3-chacha20-ietf-poly1305 -- **Trojan** (TCP, QUIC) - - Supported ciphers: - - aes-128-gcm - - aes-256-gcm - - chacha20-ietf-poly1305 -- **Hysteria2** (QUIC) -- **TUIC v5** (QUIC) +### Proxy Protocols +- **HTTP/HTTPS** +- **SOCKS5** +- **VMess AEAD** +- **VLESS** +- **Shadowsocks** +- **Trojan** +- **Snell v3** +- **Hysteria2** +- **TUIC v5** + +### Transport Protocols +All server protocols plus: +- **SagerNet UDP over TCP** (for Shadowsocks and SOCKS5) +- **ShadowTLS v3** +- **TLS** +- **WebSocket** (Shadowsocks SIP003) +- **XTLS Reality** +- **XTLS Vision** (for VLESS) + +### Supported Ciphers +- **VMess**: `aes-128-gcm`, `chacha20-poly1305`, `none` +- **Shadowsocks**: `aes-128-gcm`, `aes-256-gcm`, `chacha20-ietf-poly1305`, `2022-blake3-aes-128-gcm`, `2022-blake3-aes-256-gcm`, `2022-blake3-chacha20-ietf-poly1305` +- **Snell v3**: `aes-128-gcm`, `aes-256-gcm`, `chacha20-ietf-poly1305` ## Features -All supported protocols can be combined with the following features: +- **Multi-transport**: TCP or QUIC for all protocols +- **TLS with SNI routing**: Route by Server Name Indication +- **Upstream proxy chaining**: Route through multiple proxies +- **Rule-based routing**: Route by IP/CIDR or hostname masks +- **Named PEM certificates**: Define once, reference everywhere +- **TLS fingerprint authentication**: Certificate pinning for TLS/QUIC +- **Hot reloading**: Apply config changes without restart +- **Unix socket support**: Bind to Unix domain sockets -- **TLS support** with SNI based forwarding -- **Websocket obfs** (Shadowsocks SIP003) -- **ShadowTLS v3** -- **Upstream proxy support**: route connections through other proxy servers -- **Forwarding rules**: Redirect or block connections based on target IP or hostname -- **Hot reloading**: Updated configs are automatically reloaded -- **Netmask and proxy groups** - -For advanced access control of incoming connections (eg. IP allowlist/blocklists), check out [tobaru](https://github.com/cfal/tobaru). - -## Examples - -Here's an example of running a WSS vmess and shadowsocks server, with all requests routed through a SOCKS proxy: - -```yaml -# Listen on all IPv4 interfaces, port 443 (HTTPS) -- address: 0.0.0.0:443 - transport: tcp - # Use TLS as the first protocol layer - protocol: - type: tls - - # Set a default target, for any (or no) SNI - default_target: - cert: cert.pem - key: key.pem - # ..which goes to a websocket server - protocol: - type: ws - # .. where we have different supported proxy protocols, based on HTTP request path and headers. - targets: - - matching_path: /vmess - matching_headers: - X-Secret-Key: "secret" - protocol: - type: vmess - # allow any cipher, which means: none, aes-128-gcm, or chacha20-poly1305. - cipher: any - user_id: b0e80a62-8a51-47f0-91f1-f0f7faf8d9d4 - - matching_path: /shadowsocks - protocol: - type: shadowsocks - cipher: 2022-blake3-aes-256-gcm - password: Hax8btYlNao5qcaN/l/NUl9JgbwapfqG5QyAtH+aKPg= - - # Set a ShadowTLS v3 target by SNI - shadowtls_targets: - google.com: - # ShadowTLS password - password: 83a44859c0e7fbb589b - # Configure handshake server. - handshake: - address: google.com:443 - # Use the local SOCKS server to connect to the handshake server. - client_proxies: - - address: 127.0.0.1:1080 - protocol: - type: socks - username: socksuser - password: secretpass - - rules: - # Allow clients to connect to all IPs - - mask: 0.0.0.0/0 - action: allow - # Forward all requests through a local SOCKS server. - client_proxy: - address: 127.0.0.1:5000 - protocol: - type: socks - username: socksuser - password: secretpass -``` - -For other YAML config examples, see the [examples](./examples) directory. +For advanced access control (IP allowlist/blocklists), see [tobaru](https://github.com/cfal/tobaru). ## Installation Precompiled binaries for x86_64 and Apple aarch64 are available on [Github Releases](https://github.com/cfal/shoes/releases). -Else, if you have a fairly recent Rust and cargo installation on your system, shoes can be installed with `cargo`. +Or install with cargo: ```bash cargo install shoes @@ -126,34 +55,146 @@ cargo install shoes ## Usage ``` -shoes [OPTIONS] [YAML CONFIG PATH] [..] +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 + --no-reload Disable automatic config reloading on file changes - -t, --threads NUM - Set the number of worker threads. This usually defaults to the number of CPUs. +COMMANDS: + generate-reality-keypair Generate a new Reality X25519 keypair + generate-shadowsocks-password Generate a Shadowsocks password +``` - -d, --dry-run - Parse the config and exit. +### Examples +```bash +# Run with a single config file +shoes config.yaml + +# Run with multiple config files +shoes server1.yaml server2.yaml rules.yaml + +# Run with custom thread count +shoes --threads 8 config.yaml + +# Validate configuration without starting +shoes --dry-run config.yaml + +# Run without hot-reloading +shoes --no-reload config.yaml + +# Generate Reality keypair +shoes generate-reality-keypair + +# Generate Shadowsocks password +shoes generate-shadowsocks-password 2022-blake3-aes-256-gcm ``` ## Configuration -See [CONFIG.md](./CONFIG.md) for the YAML config format. You can also refer to the [examples](./examples), or open an issue if you need help. +See [CONFIG.md](./CONFIG.md) for the complete YAML configuration reference. -## Roadmap +## Examples -- Proxy client chaining -- SOCKS and Shadowsocks UDP support +See the [examples](./examples) directory for all examples. -## Similar and related projects +### Basic VMess Server +```yaml +- address: 0.0.0.0:16823 + protocol: + type: vmess + cipher: chacha20-poly1305 + user_id: b0e80a62-8a51-47f0-91f1-f0f7faf8d9d4 + udp_enabled: true +``` -- [shadowsocks/shadowsocks-rust](https://github.com/shadowsocks/shadowsocks-rust): A Rust port of shadowsocks +### VLESS with Vision over TLS +```yaml +- address: 0.0.0.0:443 + protocol: + type: tls + tls_targets: + "vless.example.com": + cert: cert.pem + key: key.pem + vision: true + alpn_protocols: ["http/1.1"] + protocol: + type: vless + user_id: b85798ef-e9dc-46a4-9a87-8da4499d36d0 + udp_enabled: true +``` -- [v2ray/v2ray-core](https://github.com/v2ray/v2ray-core): A full-featured proxy platform written in Go +### Reality Server +```yaml +- address: 0.0.0.0:443 + protocol: + type: tls + reality_targets: + "www.example.com": + private_key: "YOUR_BASE64URL_PRIVATE_KEY" + short_ids: ["0123456789abcdef", ""] + dest: "www.example.com:443" + protocol: + type: vless + user_id: b85798ef-e9dc-46a4-9a87-8da4499d36d0 + udp_enabled: true +``` -- [ihciah/shadow-tls](https://github.com/ihciah/shadow-tls): A proxy to expose real TLS handshake to the firewall +### Reality Client +```yaml +- address: 127.0.0.1:1080 + protocol: + type: socks + rules: + - masks: "0.0.0.0/0" + action: allow + client_proxy: + address: "server.example.com:443" + protocol: + type: reality + public_key: "SERVER_PUBLIC_KEY" + short_id: "0123456789abcdef" + sni_hostname: "www.example.com" + protocol: + type: vless + user_id: b85798ef-e9dc-46a4-9a87-8da4499d36d0 +``` -- [apernet/hysteria](https://github.com/apernet/hysteria): Hysteria is a powerful, lightning fast and censorship resistant proxy +### Hysteria2 Server +```yaml +- address: 0.0.0.0:443 + transport: quic + quic_settings: + cert: cert.pem + key: key.pem + alpn_protocols: ["h3"] + protocol: + type: hysteria2 + password: supersecret + udp_enabled: true +``` -- [tuic-protocol/tuic](https://github.com/tuic-protocol/tuic): Delicately-TUICed 0-RTT proxy protocol +### TUIC v5 Server +```yaml +- address: 0.0.0.0:443 + transport: quic + quic_settings: + cert: cert.pem + key: key.pem + protocol: + type: tuic + uuid: d685aef3-b3c4-4932-9a9d-d0c2f6727dfa + password: supersecret +``` + +## Similar Projects + +- [apernet/hysteria](https://github.com/apernet/hysteria) +- [ihciah/shadow-tls](https://github.com/ihciah/shadow-tls) +- [SagerNet/sing-box](https://github.com/SagerNet/sing-box) +- [shadowsocks/shadowsocks-rust](https://github.com/shadowsocks/shadowsocks-rust) +- [EAimTY/tuic](https://github.com/EAimTY/tuic) +- [v2fly/v2ray-core](https://github.com/v2fly/v2ray-core) +- [XTLS/Xray-core](https://github.com/XTLS/Xray-core) diff --git a/examples/named_certificates.yaml b/examples/named_certificates.yaml index 53e0b384..26002197 100644 --- a/examples/named_certificates.yaml +++ b/examples/named_certificates.yaml @@ -70,4 +70,43 @@ protocol: type: shadowsocks cipher: 2022-blake3-chacha20-ietf-poly1305 - password: "N5tp51JibWA9kIPEUoDHzyyjfgkDYToOERh+u9L47fk=" \ No newline at end of file + password: "N5tp51JibWA9kIPEUoDHzyyjfgkDYToOERh+u9L47fk=" + +# Reality server (note: Reality uses X25519 keys, not certificates) +# Reality doesn't use named PEMs for its keys since they're not certificates +- address: "0.0.0.0:10443" + protocol: + type: tls + reality_targets: + "www.cloudflare.com": + private_key: "YOUR_REALITY_PRIVATE_KEY_HERE" + short_ids: ["0123456789abcdef", ""] + dest: "www.cloudflare.com:443" + protocol: + type: vless + user_id: "b85798ef-e9dc-46a4-9a87-8da4499d36d0" + udp_enabled: true + +# Mixed configuration: Regular TLS and Reality on same server +- address: "0.0.0.0:11443" + protocol: + type: tls + # Regular TLS target using named certificates + tls_targets: + "tls.example.com": + cert: "main-server-cert" + key: "main-server-key" + protocol: + type: vless + user_id: "b85798ef-e9dc-46a4-9a87-8da4499d36d0" + # Reality target (uses its own X25519 keys) + reality_targets: + "www.google.com": + private_key: "YOUR_REALITY_PRIVATE_KEY_HERE" + short_ids: ["0123456789abcdef"] + dest: "www.google.com:443" + vision: true + protocol: + type: vless + user_id: "b85798ef-e9dc-46a4-9a87-8da4499d36d0" + udp_enabled: true \ No newline at end of file diff --git a/examples/reality_basic.yaml b/examples/reality_basic.yaml new file mode 100644 index 00000000..ea73a6e2 --- /dev/null +++ b/examples/reality_basic.yaml @@ -0,0 +1,39 @@ +# Basic Reality server with VLESS +# Provides enhanced privacy through TLS 1.3 mimicry +# +# Reality uses X25519 key exchange and HMAC-signed certificates to provide +# enhanced privacy. Non-Reality traffic is forwarded to the fallback +# destination, making the server look like a real website to observers. +# +# Generate a Reality keypair with: an X25519 keypair generation tool + +- address: "0.0.0.0:443" + protocol: + type: tls + reality_targets: + "www.cloudflare.com": + # X25519 private key (32 bytes, base64url encoded) + # Keep this secret! Clients need the corresponding public key. + private_key: "YOUR_PRIVATE_KEY_HERE" + + # List of valid short IDs (hex strings, 0-16 chars) + # Clients must use one of these IDs to connect + # Empty string "" allows all clients (less secure but convenient) + short_ids: + - "0123456789abcdef" # Specific client ID + - "" # Allow all clients + + # Fallback destination for non-Reality traffic + # This should be a real website - traffic that doesn't match Reality + # protocol will be forwarded here, making observers see a real site + dest: "www.cloudflare.com:443" + + # Optional: Maximum timestamp difference in milliseconds (default: 60000) + # Prevents replay attacks by rejecting connections with old timestamps + max_time_diff: 60000 # 1 minute + + # Inner protocol - the actual proxy protocol used after Reality handshake + protocol: + type: vless + user_id: "b85798ef-e9dc-46a4-9a87-8da4499d36d0" + udp_enabled: true # Enable XUDP for UDP multiplexing diff --git a/examples/reality_client.yaml b/examples/reality_client.yaml new file mode 100644 index 00000000..ec328e05 --- /dev/null +++ b/examples/reality_client.yaml @@ -0,0 +1,45 @@ +# Reality client configuration +# Connects to a Reality server with censorship resistance +# +# This example shows how to configure a client to connect to a Reality server +# The client needs the server's public key and a valid short ID + +# Local SOCKS proxy that uses Reality client +- address: "127.0.0.1:1080" + protocol: + type: socks + rules: + - masks: "0.0.0.0/0" + action: allow + client_proxy: + address: "your-server.example.com:443" + protocol: + type: reality + # Server's public key (from keypair generation) + public_key: "SERVER_PUBLIC_KEY_HERE" + # Your client ID (must match server's short_ids list) + short_id: "0123456789abcdef" + # Hostname to use (should match server's reality_targets key) + sni_hostname: "www.cloudflare.com" + # Inner protocol - must match what the server expects + protocol: + type: vless + user_id: "b85798ef-e9dc-46a4-9a87-8da4499d36d0" + +# Alternative: HTTP proxy using Reality +- address: "127.0.0.1:8080" + protocol: + type: http + rules: + - masks: "0.0.0.0/0" + action: allow + client_proxy: + address: "your-server.example.com:443" + protocol: + type: reality + public_key: "SERVER_PUBLIC_KEY_HERE" + short_id: "0123456789abcdef" + sni_hostname: "www.cloudflare.com" + protocol: + type: vless + user_id: "b85798ef-e9dc-46a4-9a87-8da4499d36d0" diff --git a/examples/reality_trojan.yaml b/examples/reality_trojan.yaml new file mode 100644 index 00000000..5a47b9c5 --- /dev/null +++ b/examples/reality_trojan.yaml @@ -0,0 +1,56 @@ +# Reality server with Trojan inner protocol +# Shows Reality works with protocols other than VLESS (when Vision is disabled) +# +# Reality provides censorship resistance regardless of the inner protocol +# However, Vision (TLS-in-TLS optimization) only works with VLESS +# This example uses Trojan without Vision + +- address: "0.0.0.0:443" + protocol: + type: tls + reality_targets: + "www.cloudflare.com": + private_key: "YOUR_PRIVATE_KEY_HERE" + short_ids: ["0123456789abcdef"] + dest: "www.cloudflare.com:443" + + # vision: false (default) + # Vision is not enabled because Trojan is the inner protocol + # Vision only works with VLESS + + protocol: + type: trojan + password: "your-trojan-password" + +# Multiple Reality targets with different inner protocols +- address: "0.0.0.0:8443" + protocol: + type: tls + reality_targets: + # VLESS target + "vless.example.com": + private_key: "YOUR_PRIVATE_KEY_1" + short_ids: ["0123456789abcdef"] + dest: "www.cloudflare.com:443" + protocol: + type: vless + user_id: "b85798ef-e9dc-46a4-9a87-8da4499d36d0" + + # Trojan target + "trojan.example.com": + private_key: "YOUR_PRIVATE_KEY_2" + short_ids: ["1234567890abcdef"] + dest: "www.google.com:443" + protocol: + type: trojan + password: "trojan-password" + + # Shadowsocks target + "ss.example.com": + private_key: "YOUR_PRIVATE_KEY_3" + short_ids: ["abcdef0123456789"] + dest: "www.microsoft.com:443" + protocol: + type: shadowsocks + cipher: 2022-blake3-chacha20-ietf-poly1305 + password: "N5tp51JibWA9kIPEUoDHzyyjfgkDYToOERh+u9L47fk=" diff --git a/examples/reality_vision.yaml b/examples/reality_vision.yaml new file mode 100644 index 00000000..67ee1d79 --- /dev/null +++ b/examples/reality_vision.yaml @@ -0,0 +1,31 @@ +# Reality + Vision server +# Combines censorship resistance with TLS-in-TLS optimization +# +# Reality provides censorship resistance through TLS 1.3 mimicry +# Vision provides zero-copy performance when TLS-in-TLS is detected +# Together they offer both benefits + +- address: "0.0.0.0:443" + protocol: + type: tls + reality_targets: + "www.google.com": + private_key: "YOUR_PRIVATE_KEY_HERE" + short_ids: ["0123456789abcdef"] + dest: "www.google.com:443" + + # Enable XTLS-Vision for TLS-in-TLS optimization + # When enabled, Vision detects TLS-in-TLS scenarios and switches + # to Direct mode for zero-copy performance + # IMPORTANT: Vision requires VLESS as the inner protocol + vision: true + + protocol: + type: vless # Vision REQUIRES VLESS + user_id: "b85798ef-e9dc-46a4-9a87-8da4499d36d0" + udp_enabled: true + +# Benefits of Reality + Vision: +# - Reality: Censorship resistance, looks like real TLS traffic +# - Vision: Zero-copy performance, better anti-detection +# - Combined: Maximum security and performance diff --git a/examples/reality_vision_client.yaml b/examples/reality_vision_client.yaml new file mode 100644 index 00000000..7f6c6094 --- /dev/null +++ b/examples/reality_vision_client.yaml @@ -0,0 +1,31 @@ +# Reality + Vision client configuration +# Maximum performance and censorship resistance +# +# Reality provides censorship resistance through TLS 1.3 mimicry +# Vision provides zero-copy performance for TLS-in-TLS scenarios +# Together they offer both benefits with VLESS as the inner protocol + +# Local SOCKS proxy using Reality + Vision +- address: "127.0.0.1:1080" + protocol: + type: socks + rules: + - masks: "0.0.0.0/0" + action: allow + client_proxy: + address: "your-server.example.com:443" + protocol: + type: reality + public_key: "SERVER_PUBLIC_KEY_HERE" + short_id: "0123456789abcdef" + sni_hostname: "www.google.com" + vision: true # Enable Vision (must use VLESS) + protocol: + type: vless # Vision REQUIRES VLESS + user_id: "b85798ef-e9dc-46a4-9a87-8da4499d36d0" + +# Benefits of Reality + Vision client: +# - Censorship resistance: Looks like normal TLS traffic to censors +# - Performance: Vision provides zero-copy for TLS-in-TLS +# - Anti-detection: Vision adds padding to mimic real TLS patterns +# - Compatibility: Works with CDNs and middleboxes diff --git a/examples/vless_variants.yaml b/examples/vless_variants.yaml index 81263ddc..6eae9c5a 100644 --- a/examples/vless_variants.yaml +++ b/examples/vless_variants.yaml @@ -50,3 +50,47 @@ address: "proxy.example.com:1080" protocol: type: socks + +# VLESS over Reality (censorship resistance) +- address: "0.0.0.0:8443" + protocol: + type: tls + reality_targets: + "www.cloudflare.com": + private_key: "YOUR_PRIVATE_KEY_HERE" + short_ids: ["0123456789abcdef", ""] + dest: "www.cloudflare.com:443" + protocol: + type: vless + user_id: "b85798ef-e9dc-46a4-9a87-8da4499d36d0" + udp_enabled: true + +# VLESS over Reality with Vision (censorship + performance) +- address: "0.0.0.0:9443" + protocol: + type: tls + reality_targets: + "www.google.com": + private_key: "YOUR_PRIVATE_KEY_HERE" + short_ids: ["0123456789abcdef"] + dest: "www.google.com:443" + vision: true # Enable Vision + protocol: + type: vless # Vision requires VLESS + user_id: "b85798ef-e9dc-46a4-9a87-8da4499d36d0" + udp_enabled: true + +# VLESS over TLS with Vision (no Reality) +- address: "0.0.0.0:10443" + protocol: + type: tls + tls_targets: + "vless.example.com": + cert: "cert.pem" + key: "key.pem" + vision: true # Enable Vision + alpn_protocols: ["http/1.1"] + protocol: + type: vless + user_id: "b85798ef-e9dc-46a4-9a87-8da4499d36d0" + udp_enabled: true diff --git a/src/address.rs b/src/address.rs index 831c957c..53e7ad4d 100644 --- a/src/address.rs +++ b/src/address.rs @@ -472,14 +472,69 @@ impl NetLocationMask { }; pub fn from(s: &str) -> std::io::Result { - let (address_mask_str, port) = match s.find(':') { - Some(i) => { - let port = s[i + 1..] + // Handle IPv6 with port: [::1/128]:80 + if s.starts_with('[') { + if let Some(bracket_end) = s.find(']') { + let address_mask_str = &s[1..bracket_end]; + let port = if s.len() > bracket_end + 1 && s.as_bytes()[bracket_end + 1] == b':' { + s[bracket_end + 2..] + .parse::() + .map_err(|e| std::io::Error::other(format!("Failed to parse port: {e}")))? + } else { + 0 + }; + return Ok(Self { + address_mask: AddressMask::from(address_mask_str)?, + port, + }); + } + } + + // For addresses without brackets, we need to distinguish IPv6 from port notation. + // IPv6 addresses contain multiple colons. If there's a `/` (CIDR), the port must come + // after it. If no `/`, use rfind(':') and check if what follows looks like a port. + let (address_mask_str, port) = if let Some(slash_pos) = s.rfind('/') { + // Has CIDR notation. Port (if any) must come after the slash. + // Format: addr/bits:port or addr/bits + let after_slash = &s[slash_pos + 1..]; + if let Some(colon_in_suffix) = after_slash.find(':') { + // There's a port after the CIDR + let bits_str = &after_slash[..colon_in_suffix]; + let port_str = &after_slash[colon_in_suffix + 1..]; + // Validate bits is a number + bits_str.parse::().map_err(|e| { + std::io::Error::other(format!("Failed to parse netmask bits: {e}")) + })?; + let port = port_str .parse::() .map_err(|e| std::io::Error::other(format!("Failed to parse port: {e}")))?; - (&s[0..i], port) + (&s[..slash_pos + 1 + colon_in_suffix], port) + } else { + // No port, just addr/bits + (s, 0) + } + } else { + // No CIDR notation. Could be: + // - hostname:port (example.com:80) + // - IPv4:port (1.2.3.4:80) + // - IPv6 (::1, fe80::1) + // - hostname (example.com) + // - IPv4 (1.2.3.4) + // + // IPv6 addresses have multiple colons, so if there's more than one colon, + // it's IPv6 and has no port. If there's exactly one colon, use it as port separator. + let colon_count = s.chars().filter(|&c| c == ':').count(); + if colon_count == 1 { + // Single colon - treat as port separator + let colon_pos = s.find(':').unwrap(); + let port = s[colon_pos + 1..] + .parse::() + .map_err(|e| std::io::Error::other(format!("Failed to parse port: {e}")))?; + (&s[..colon_pos], port) + } else { + // Zero or multiple colons - no port (IPv6 or plain hostname/IPv4) + (s, 0) } - None => (s, 0), }; Ok(Self { diff --git a/src/async_stream.rs b/src/async_stream.rs index 1edda55d..fe93182d 100644 --- a/src/async_stream.rs +++ b/src/async_stream.rs @@ -80,6 +80,28 @@ pub trait AsyncWriteSourcedMessage { ) -> Poll>; } +/// Session-based message reading trait. Used by protocols like XUDP that have session IDs. +/// Returns (session_id, data, source_addr) tuples. +pub trait AsyncReadSessionMessage { + fn poll_read_session_message( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll>; +} + +/// Session-based message writing trait. Used by protocols like XUDP that have session IDs. +/// Writes data for a specific session ID to a target address. +pub trait AsyncWriteSessionMessage { + fn poll_write_session_message( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + session_id: u16, + buf: &[u8], + target: &SocketAddr, + ) -> Poll>; +} + impl AsyncReadMessage for UdpSocket { fn poll_read_message( self: Pin<&mut Self>, @@ -158,6 +180,19 @@ pub trait AsyncSourcedMessageStream: { } +/// Session-based stream trait for protocols like XUDP that multiplex sessions over a single connection. +/// Reads return (session_id, data, source_addr) and writes target (session_id, data, target_addr). +pub trait AsyncSessionMessageStream: + AsyncReadSessionMessage + + AsyncWriteSessionMessage + + AsyncFlushMessage + + AsyncShutdownMessage + + AsyncPing + + Unpin + + Send +{ +} + impl AsyncPing for TcpStream { fn supports_ping(&self) -> bool { false @@ -196,38 +231,6 @@ impl AsyncPing for UdpSocket { impl AsyncMessageStream for UdpSocket {} -impl AsyncPing for tokio_rustls::client::TlsStream -where - AS: AsyncPing + Unpin, -{ - fn supports_ping(&self) -> bool { - self.get_ref().0.supports_ping() - } - - fn poll_write_ping(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let this = self.get_mut(); - Pin::new(this.get_mut().0).poll_write_ping(cx) - } -} - -impl AsyncStream for tokio_rustls::client::TlsStream where AS: AsyncStream {} - -impl AsyncPing for tokio_rustls::server::TlsStream -where - AS: AsyncPing + Unpin, -{ - fn supports_ping(&self) -> bool { - self.get_ref().0.supports_ping() - } - - fn poll_write_ping(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let this = self.get_mut(); - Pin::new(this.get_mut().0).poll_write_ping(cx) - } -} - -impl AsyncStream for tokio_rustls::server::TlsStream where AS: AsyncStream {} - // pattern copied from deref_async_read macro: https://docs.rs/tokio/latest/src/tokio/io/async_read.rs.html#60 impl AsyncPing for Box { fn supports_ping(&self) -> bool { @@ -426,3 +429,50 @@ impl AsyncTargetedMessageStream impl AsyncSourcedMessageStream for Box {} impl AsyncSourcedMessageStream for &mut T {} + +impl AsyncReadSessionMessage for Box { + fn poll_read_session_message( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + Pin::new(&mut **self).poll_read_session_message(cx, buf) + } +} + +impl AsyncReadSessionMessage for &mut T { + fn poll_read_session_message( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + Pin::new(&mut **self).poll_read_session_message(cx, buf) + } +} + +impl AsyncWriteSessionMessage for Box { + fn poll_write_session_message( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + session_id: u16, + buf: &[u8], + target: &SocketAddr, + ) -> Poll> { + Pin::new(&mut **self).poll_write_session_message(cx, session_id, buf, target) + } +} + +impl AsyncWriteSessionMessage for &mut T { + fn poll_write_session_message( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + session_id: u16, + buf: &[u8], + target: &SocketAddr, + ) -> Poll> { + Pin::new(&mut **self).poll_write_session_message(cx, session_id, buf, target) + } +} + +impl AsyncSessionMessageStream for Box {} +impl AsyncSessionMessageStream for &mut T {} diff --git a/src/client_proxy_selector.rs b/src/client_proxy_selector.rs index 63c2a5be..cdafd960 100644 --- a/src/client_proxy_selector.rs +++ b/src/client_proxy_selector.rs @@ -1,5 +1,9 @@ use log::{debug, error}; +use lru::LruCache; +use parking_lot::RwLock; +use std::hash::Hash; use std::net::{IpAddr, Ipv4Addr, Ipv6Addr}; +use std::num::NonZeroUsize; use std::sync::atomic::{AtomicU32, Ordering}; use std::sync::Arc; @@ -8,16 +12,104 @@ use crate::address::{AddressMask, NetLocationMask}; use crate::option_util::OneOrSome; use crate::resolver::{resolve_single_address, Resolver}; -// If a hostname is provided in a rule, it won't be resolved, -// and the netmask field will be ignored. -// This is useful when a huge blocklist or rule list is provided. -// However, this means that the user needs to make sure DNS -// resolutions are not done manually before hitting the server. -// TODO: make this configurable. -// TODO: when this is enabled, it would be much faster to use a -// hash map to do a single lookup rather than traversing the rules -// Vec. -const DONT_RESOLVE_RULE_HOSTNAMES: bool = true; +/// Cache key for routing decisions. +/// We cache based on the destination address and port. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +enum RoutingCacheKey { + /// Hostname destination (e.g., "example.com:443") + Hostname { hostname: String, port: u16 }, + /// IPv4 destination + Ipv4 { addr: Ipv4Addr, port: u16 }, + /// IPv6 destination + Ipv6 { addr: Ipv6Addr, port: u16 }, +} + +impl RoutingCacheKey { + fn from_location(location: &NetLocation) -> Self { + match location.address() { + Address::Hostname(hostname) => RoutingCacheKey::Hostname { + hostname: hostname.to_lowercase(), + port: location.port(), + }, + Address::Ipv4(addr) => RoutingCacheKey::Ipv4 { + addr: *addr, + port: location.port(), + }, + Address::Ipv6(addr) => RoutingCacheKey::Ipv6 { + addr: *addr, + port: location.port(), + }, + } + } +} + +/// Cached routing decision. +#[derive(Debug, Clone, Copy)] +pub(crate) enum CachedDecision { + /// Index into the rules Vec for an Allow decision + Allow(usize), + /// Block decision + Block, +} + +/// LRU cache for routing decisions. +/// Uses RwLock for concurrent access - reads use peek() which doesn't require mutable access. +pub struct RoutingCache { + inner: RwLock>, +} + +impl std::fmt::Debug for RoutingCache { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let cache = self.inner.read(); + f.debug_struct("RoutingCache") + .field("len", &cache.len()) + .field("cap", &cache.cap()) + .finish() + } +} + +impl RoutingCache { + /// Default cache capacity + pub const DEFAULT_CAPACITY: usize = 10_000; + + /// Create a new routing cache with the specified capacity. + pub fn new(capacity: usize) -> Self { + let cap = NonZeroUsize::new(capacity).unwrap_or(NonZeroUsize::new(1).unwrap()); + Self { + inner: RwLock::new(LruCache::new(cap)), + } + } + + /// Look up a cached decision. + /// Uses peek() which doesn't update LRU order (avoids write lock). + #[inline] + pub fn get(&self, location: &NetLocation) -> Option { + let key = RoutingCacheKey::from_location(location); + self.inner.read().peek(&key).copied() + } + + /// Insert a decision into the cache. + #[inline] + pub fn insert(&self, location: &NetLocation, decision: CachedDecision) { + let key = RoutingCacheKey::from_location(location); + self.inner.write().put(key, decision); + } + + #[cfg(test)] + fn clear(&self) { + self.inner.write().clear(); + } + + #[cfg(test)] + fn len(&self) -> usize { + self.inner.read().len() + } + + #[cfg(test)] + fn is_empty(&self) -> bool { + self.inner.read().is_empty() + } +} #[derive(Debug)] pub struct ConnectRule { @@ -87,13 +179,27 @@ impl ConnectAction { } } +/// Threshold for enabling cache based on rule count. +/// If rule count exceeds this value, caching is enabled even without DNS resolution. +const CACHE_RULE_THRESHOLD: usize = 16; + +// TODO: Replace linear rule matching with radix set/trie #[derive(Debug)] pub struct ClientProxySelector { rules: Vec>, default_rule_index: Option, + /// If false, hostname rules will not trigger DNS resolution to match against IP-based + /// destinations. This is useful when a huge blocklist or rule list is provided. + /// However, this means that the user needs to make sure DNS resolutions are not done + /// manually before hitting the server. Default is false (don't resolve). + resolve_rule_hostnames: bool, + /// LRU cache for routing decisions. Speeds up repeated lookups for the same destination. + /// None if caching is disabled (few rules and no DNS resolution). + cache: Option, } unsafe impl Send for ClientProxySelector {} +unsafe impl Sync for ClientProxySelector {} #[derive(Debug)] pub enum ConnectDecision<'a, T> { @@ -106,6 +212,45 @@ pub enum ConnectDecision<'a, T> { impl ClientProxySelector { pub fn new(rules: Vec>) -> Self { + Self::with_options(rules, false) + } + + /// Create a new ClientProxySelector with configurable hostname resolution behavior. + /// + /// # Arguments + /// * `rules` - The list of routing rules + /// * `resolve_rule_hostnames` - If true, hostname rules will be resolved via DNS to match + /// against IP-based destinations. If false (default), hostname rules only match hostname + /// destinations directly. Setting this to false is more performant for large rule sets. + pub fn with_options(rules: Vec>, resolve_rule_hostnames: bool) -> Self { + Self::with_options_and_cache_size( + rules, + resolve_rule_hostnames, + RoutingCache::DEFAULT_CAPACITY, + ) + } + + /// Create a new ClientProxySelector with configurable hostname resolution and cache size. + /// + /// # Arguments + /// * `rules` - The list of routing rules + /// * `resolve_rule_hostnames` - If true, hostname rules will be resolved via DNS to match + /// against IP-based destinations. If false (default), hostname rules only match hostname + /// destinations directly. + /// * `cache_capacity` - Maximum number of routing decisions to cache. Set to 0 to disable caching. + /// + /// # Caching behavior + /// Caching is automatically enabled when: + /// - `resolve_rule_hostnames` is true (DNS lookups are expensive), OR + /// - Rule count exceeds `CACHE_RULE_THRESHOLD` (16) + /// + /// For simple configurations with few rules and no DNS resolution, caching is disabled + /// as the overhead of cache key construction exceeds the cost of linear rule matching. + pub fn with_options_and_cache_size( + rules: Vec>, + resolve_rule_hostnames: bool, + cache_capacity: usize, + ) -> Self { let mut default_rule_index: Option = None; // find a default rule which we'll use for multidirectional forwarding.. // TODO: ideally, we'd check the rule for each target during multidirectional forwarding @@ -128,9 +273,21 @@ impl ClientProxySelector { break; } } + + // Enable caching if: + // 1. DNS resolution is enabled (expensive operation), OR + // 2. Many rules (linear scan becomes expensive) + let cache = if resolve_rule_hostnames || rules.len() > CACHE_RULE_THRESHOLD { + Some(RoutingCache::new(cache_capacity.max(1))) + } else { + None + }; + Self { rules, default_rule_index, + resolve_rule_hostnames, + cache, } } @@ -146,16 +303,101 @@ impl ClientProxySelector { } } + /// Judge a connection request, using the cache for faster repeated lookups. + /// + /// This is the primary method for routing decisions. It first checks the cache, + /// and only performs the full rule matching if there's a cache miss. + /// + /// Note: Caching is only enabled when `resolve_rule_hostnames` is true or there are + /// more than 16 rules. For simple configurations, direct rule matching is faster. pub async fn judge<'a>( &'a self, location: NetLocation, resolver: &Arc, ) -> std::io::Result> { - match match_rule(&self.rules, &location, resolver).await? { - Some(rule) => Ok(rule.action.to_decision(location)), + // If caching is disabled, go directly to rule matching + let cache = match &self.cache { + Some(c) => c, + None => return self.judge_uncached(location, resolver).await, + }; + + // Fast path: check cache first + if let Some(cached) = cache.get(&location) { + return Ok(self.cached_to_decision(cached, location)); + } + + // Slow path: full rule matching + match match_rule( + &self.rules, + &location, + resolver, + self.resolve_rule_hostnames, + ) + .await? + { + Some((rule_index, rule)) => { + // Cache the result + cache.insert(&location, CachedDecision::Allow(rule_index)); + Ok(rule.action.to_decision(location)) + } + None => { + // Cache the block decision + cache.insert(&location, CachedDecision::Block); + Ok(ConnectDecision::Block) + } + } + } + + /// Judge without using the cache. Useful for testing or when cache bypass is needed. + pub async fn judge_uncached<'a>( + &'a self, + location: NetLocation, + resolver: &Arc, + ) -> std::io::Result> { + match match_rule( + &self.rules, + &location, + resolver, + self.resolve_rule_hostnames, + ) + .await? + { + Some((_rule_index, rule)) => Ok(rule.action.to_decision(location)), None => Ok(ConnectDecision::Block), } } + + /// Convert a cached decision back to a ConnectDecision. + #[inline] + fn cached_to_decision( + &self, + cached: CachedDecision, + location: NetLocation, + ) -> ConnectDecision<'_, T> { + match cached { + CachedDecision::Allow(rule_index) => { + self.rules[rule_index].action.to_decision(location) + } + CachedDecision::Block => ConnectDecision::Block, + } + } + + #[cfg(test)] + fn cache_size(&self) -> usize { + self.cache.as_ref().map_or(0, |c| c.len()) + } + + #[cfg(test)] + fn clear_cache(&self) { + if let Some(cache) = &self.cache { + cache.clear(); + } + } + + #[cfg(test)] + fn is_cache_enabled(&self) -> bool { + self.cache.is_some() + } } #[inline] @@ -206,22 +448,32 @@ fn matches_domain(base_domain: &str, hostname: &str) -> bool { } } +/// Returns the matching rule and its index in the rules Vec. #[inline] async fn match_rule<'a, T>( rules: &'a [ConnectRule], location: &NetLocation, resolver: &Arc, -) -> std::io::Result>> { + resolve_rule_hostnames: bool, +) -> std::io::Result)>> { // We only resolve when necessary. let mut resolved_ip: Option = None; - for rule in rules.iter() { + for (rule_index, rule) in rules.iter().enumerate() { for mask in rule.masks.iter() { - match match_mask(mask, location, &mut resolved_ip, resolver).await { + match match_mask( + mask, + location, + &mut resolved_ip, + resolver, + resolve_rule_hostnames, + ) + .await + { Ok(is_match) => { if is_match { debug!("Found matching mask for {location} -> {mask:?}"); - return Ok(Some(rule)); + return Ok(Some((rule_index, rule))); } } Err(MatchMaskError::Fatal(e)) => { @@ -243,12 +495,19 @@ enum MatchMaskError { Fatal(std::io::Error), } +// Helper function for testing - exposed for unit tests +#[cfg(test)] +pub fn matches_domain_for_test(base_domain: &str, hostname: &str) -> bool { + matches_domain(base_domain, hostname) +} + #[inline] async fn match_mask( location_mask: &NetLocationMask, location: &NetLocation, resolved_ip: &mut Option, resolver: &Arc, + resolve_rule_hostnames: bool, ) -> std::result::Result { let NetLocationMask { address_mask: AddressMask { address, netmask }, @@ -273,8 +532,8 @@ async fn match_mask( // We don't care about netmasks when hostnames are provided, so we can do direct matching // without resolving when both the remote location and the provided rule address are hostnames, - // and simploy return if it doesn't match. - if DONT_RESOLVE_RULE_HOSTNAMES { + // and simply return if it doesn't match. + if !resolve_rule_hostnames { return Ok(false); } } @@ -323,3 +582,1713 @@ async fn match_mask( Ok(false) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::resolver::Resolver; + use std::future::Future; + use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr}; + use std::pin::Pin; + + // ========================================================================= + // Test Infrastructure + // ========================================================================= + + /// A mock resolver for testing that returns predefined results + #[derive(Debug)] + struct MockResolver { + mappings: std::collections::HashMap>, + } + + impl MockResolver { + fn new() -> Self { + Self { + mappings: std::collections::HashMap::new(), + } + } + + fn with_mapping(mut self, hostname: &str, port: u16, addrs: Vec) -> Self { + let key = format!("{}:{}", hostname, port); + let socket_addrs: Vec = addrs + .into_iter() + .map(|ip| SocketAddr::new(ip, port)) + .collect(); + self.mappings.insert(key, socket_addrs); + self + } + } + + impl Resolver for MockResolver { + fn resolve_location( + &self, + location: &NetLocation, + ) -> Pin>> + Send>> { + let key = format!("{}:{}", location.address(), location.port()); + let result = self.mappings.get(&key).cloned(); + Box::pin(async move { + result.ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::NotFound, + format!("No mock mapping for {}", key), + ) + }) + }) + } + } + + /// Helper to create a mock resolver as Arc + fn mock_resolver() -> Arc { + Arc::new(MockResolver::new()) + } + + /// Helper to create a mock resolver with mappings as Arc + #[allow(dead_code)] + fn mock_resolver_with_mappings(mappings: Vec<(&str, u16, Vec)>) -> Arc { + let mut resolver = MockResolver::new(); + for (hostname, port, addrs) in mappings { + resolver = resolver.with_mapping(hostname, port, addrs); + } + Arc::new(resolver) + } + + /// Simple proxy type for testing + #[derive(Debug, Clone, PartialEq)] + struct TestProxy { + name: String, + } + + impl TestProxy { + fn new(name: &str) -> Self { + Self { + name: name.to_string(), + } + } + } + + /// Helper to create an allow rule with a single proxy + fn allow_rule(masks: Vec<&str>, proxy_name: &str) -> ConnectRule { + let masks: Vec = masks + .into_iter() + .map(|s| NetLocationMask::from(s).unwrap()) + .collect(); + ConnectRule::new( + masks, + ConnectAction::new_allow(None, OneOrSome::One(TestProxy::new(proxy_name))), + ) + } + + /// Helper to create an allow rule with multiple proxies (for round-robin) + fn allow_rule_multi(masks: Vec<&str>, proxy_names: Vec<&str>) -> ConnectRule { + let masks: Vec = masks + .into_iter() + .map(|s| NetLocationMask::from(s).unwrap()) + .collect(); + let proxies: Vec = proxy_names.into_iter().map(TestProxy::new).collect(); + ConnectRule::new( + masks, + ConnectAction::new_allow(None, OneOrSome::Some(proxies)), + ) + } + + /// Helper to create a block rule + fn block_rule(masks: Vec<&str>) -> ConnectRule { + let masks: Vec = masks + .into_iter() + .map(|s| NetLocationMask::from(s).unwrap()) + .collect(); + ConnectRule::new(masks, ConnectAction::new_block()) + } + + /// Helper to create an allow rule with address override + fn allow_rule_with_override( + masks: Vec<&str>, + proxy_name: &str, + override_addr: &str, + ) -> ConnectRule { + let masks: Vec = masks + .into_iter() + .map(|s| NetLocationMask::from(s).unwrap()) + .collect(); + let override_location = NetLocation::from_str(override_addr, Some(0)).unwrap(); + ConnectRule::new( + masks, + ConnectAction::new_allow( + Some(override_location), + OneOrSome::One(TestProxy::new(proxy_name)), + ), + ) + } + + // ========================================================================= + // Domain Matching Tests (matches_domain function) + // ========================================================================= + + #[test] + fn test_matches_domain_exact_match() { + assert!(matches_domain_for_test("example.com", "example.com")); + assert!(matches_domain_for_test("foo.bar.com", "foo.bar.com")); + } + + #[test] + fn test_matches_domain_subdomain_match() { + assert!(matches_domain_for_test("example.com", "www.example.com")); + assert!(matches_domain_for_test( + "example.com", + "sub.www.example.com" + )); + assert!(matches_domain_for_test("bar.com", "foo.bar.com")); + } + + #[test] + fn test_matches_domain_no_match_different_domain() { + assert!(!matches_domain_for_test("example.com", "notexample.com")); + assert!(!matches_domain_for_test("example.com", "example.org")); + assert!(!matches_domain_for_test("foo.com", "bar.com")); + } + + #[test] + fn test_matches_domain_no_match_partial_suffix() { + // "malicious-example.com" ends with "example.com" as a string, + // but it's NOT a subdomain - this is a critical security test + assert!(!matches_domain_for_test( + "example.com", + "malicious-example.com" + )); + assert!(!matches_domain_for_test("example.com", "fakeexample.com")); + assert!(!matches_domain_for_test("bar.com", "foobar.com")); + } + + #[test] + fn test_matches_domain_no_match_base_longer() { + assert!(!matches_domain_for_test("www.example.com", "example.com")); + assert!(!matches_domain_for_test("sub.example.com", "example.com")); + } + + #[test] + fn test_matches_domain_empty_strings() { + assert!(!matches_domain_for_test("example.com", "")); + assert!(!matches_domain_for_test("", "example.com")); + assert!(matches_domain_for_test("", "")); // Both empty matches + } + + #[test] + fn test_matches_domain_single_label() { + assert!(matches_domain_for_test("localhost", "localhost")); + assert!(!matches_domain_for_test("localhost", "notlocalhost")); + assert!(!matches_domain_for_test("host", "localhost")); // partial suffix + } + + // ========================================================================= + // IPv4 Address Matching Tests + // ========================================================================= + + #[tokio::test] + async fn test_ipv4_exact_match() { + let rules = vec![ + allow_rule(vec!["192.168.1.1/32"], "proxy1"), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(192, 168, 1, 1)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "proxy1"); + } + ConnectDecision::Block => panic!("Expected Allow, got Block"), + } + } + + #[tokio::test] + async fn test_ipv4_exact_match_no_match() { + let rules = vec![ + allow_rule(vec!["192.168.1.1/32"], "proxy1"), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + // Different IP should fall through to default + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(192, 168, 1, 2)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "default"); + } + ConnectDecision::Block => panic!("Expected Allow, got Block"), + } + } + + #[tokio::test] + async fn test_ipv4_cidr_24() { + let rules = vec![ + allow_rule(vec!["192.168.1.0/24"], "lan"), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + // All addresses in 192.168.1.0/24 should match + for last_octet in [0u8, 1, 100, 254, 255] { + let location = + NetLocation::new(Address::Ipv4(Ipv4Addr::new(192, 168, 1, last_octet)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!( + client_proxy.name, "lan", + "Failed for 192.168.1.{}", + last_octet + ); + } + ConnectDecision::Block => panic!("Expected Allow for 192.168.1.{}", last_octet), + } + } + + // Address outside the range should not match + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(192, 168, 2, 1)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "default"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + } + + #[tokio::test] + async fn test_ipv4_cidr_16() { + let rules = vec![ + allow_rule(vec!["10.0.0.0/8"], "private_a"), + allow_rule(vec!["172.16.0.0/12"], "private_b"), + allow_rule(vec!["192.168.0.0/16"], "private_c"), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + // Test 10.0.0.0/8 range + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(10, 255, 255, 255)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "private_a"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + + // Test 172.16.0.0/12 range (172.16.0.0 - 172.31.255.255) + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(172, 20, 5, 1)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "private_b"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + + // Test 192.168.0.0/16 range + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(192, 168, 100, 50)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "private_c"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + + // Public IP should go to default + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(8, 8, 8, 8)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "default"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + } + + #[tokio::test] + async fn test_ipv4_localhost() { + let rules = vec![ + allow_rule(vec!["127.0.0.0/8"], "loopback"), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(127, 0, 0, 1)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "loopback"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + + // 127.x.x.x should all match loopback + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(127, 255, 255, 255)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "loopback"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + } + + // ========================================================================= + // IPv6 Address Matching Tests + // ========================================================================= + + #[tokio::test] + async fn test_ipv6_exact_match() { + let rules = vec![ + allow_rule(vec!["::1/128"], "loopback"), + allow_rule(vec!["::/0"], "default"), + ]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + let location = NetLocation::new(Address::Ipv6(Ipv6Addr::LOCALHOST), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "loopback"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + } + + #[tokio::test] + async fn test_ipv6_cidr() { + let rules = vec![ + allow_rule(vec!["fe80::/10"], "link_local"), + allow_rule(vec!["fc00::/7"], "unique_local"), + allow_rule(vec!["::/0"], "default"), + ]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + // Link-local address + let location = NetLocation::new(Address::Ipv6("fe80::1".parse().unwrap()), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "link_local"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + + // Unique local address (fc00::/7 covers fc00:: and fd00::) + let location = NetLocation::new(Address::Ipv6("fd00::1234".parse().unwrap()), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "unique_local"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + + // Global unicast should go to default + let location = NetLocation::new(Address::Ipv6("2001:db8::1".parse().unwrap()), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "default"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + } + + // ========================================================================= + // Hostname/Domain Matching Tests + // ========================================================================= + + #[tokio::test] + async fn test_hostname_exact_match() { + let rules = vec![ + allow_rule(vec!["example.com"], "example_proxy"), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + let location = NetLocation::new(Address::Hostname("example.com".to_string()), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "example_proxy"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + } + + #[tokio::test] + async fn test_hostname_subdomain_match() { + let rules = vec![ + allow_rule(vec!["example.com"], "example_proxy"), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + // Subdomain should match + let location = NetLocation::new(Address::Hostname("www.example.com".to_string()), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "example_proxy"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + + // Deep subdomain should match + let location = NetLocation::new(Address::Hostname("sub.www.example.com".to_string()), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "example_proxy"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + } + + #[tokio::test] + async fn test_hostname_no_false_subdomain_match() { + // SECURITY CRITICAL: malicious-example.com should NOT match example.com rule + let rules = vec![ + allow_rule(vec!["example.com"], "example_proxy"), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + let location = NetLocation::new(Address::Hostname("malicious-example.com".to_string()), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!( + client_proxy.name, "default", + "SECURITY: malicious-example.com should NOT match example.com rule" + ); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + + let location = NetLocation::new(Address::Hostname("fakeexample.com".to_string()), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "default"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + } + + #[tokio::test] + async fn test_hostname_different_tld() { + let rules = vec![ + allow_rule(vec!["example.com"], "example_proxy"), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + // Different TLD should not match + let location = NetLocation::new(Address::Hostname("example.org".to_string()), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "default"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + } + + // ========================================================================= + // Port Matching Tests + // ========================================================================= + + #[tokio::test] + async fn test_port_specific_rule() { + let rules = vec![ + allow_rule(vec!["0.0.0.0/0:443"], "https_proxy"), + allow_rule(vec!["0.0.0.0/0:80"], "http_proxy"), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + // Port 443 should match https_proxy + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(1, 2, 3, 4)), 443); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "https_proxy"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + + // Port 80 should match http_proxy + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(1, 2, 3, 4)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "http_proxy"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + + // Other ports should go to default + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(1, 2, 3, 4)), 8080); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "default"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + } + + #[tokio::test] + async fn test_port_with_ip_restriction() { + let rules = vec![ + allow_rule(vec!["192.168.1.0/24:22"], "ssh_lan"), + allow_rule(vec!["0.0.0.0/0:22"], "ssh_default"), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + // LAN SSH should match ssh_lan + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(192, 168, 1, 100)), 22); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "ssh_lan"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + + // External SSH should match ssh_default + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(8, 8, 8, 8)), 22); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "ssh_default"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + + // LAN HTTP should go to default (not ssh_lan because port doesn't match) + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(192, 168, 1, 100)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "default"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + } + + // ========================================================================= + // Rule Priority Tests (First Match Wins) + // ========================================================================= + + #[tokio::test] + async fn test_first_rule_wins() { + let rules = vec![ + allow_rule(vec!["192.168.1.0/24"], "specific"), + allow_rule(vec!["192.168.0.0/16"], "less_specific"), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + // 192.168.1.x should match the first (more specific) rule + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(192, 168, 1, 1)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "specific"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + + // 192.168.2.x should match the second rule + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(192, 168, 2, 1)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "less_specific"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + } + + #[tokio::test] + async fn test_rule_order_matters_security() { + // SECURITY: If a block rule comes after an allow rule for overlapping ranges, + // the allow rule takes precedence + let rules = vec![ + allow_rule(vec!["192.168.1.0/24"], "allowed"), + block_rule(vec!["192.168.0.0/16"]), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + // 192.168.1.x matches the allow rule first + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(192, 168, 1, 1)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "allowed"); + } + ConnectDecision::Block => panic!("Expected Allow for 192.168.1.1"), + } + + // 192.168.2.x should be blocked + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(192, 168, 2, 1)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { .. } => panic!("Expected Block for 192.168.2.1"), + ConnectDecision::Block => {} // Expected + } + } + + // ========================================================================= + // Block Action Tests + // ========================================================================= + + #[tokio::test] + async fn test_block_rule() { + let rules = vec![ + block_rule(vec!["192.168.1.0/24"]), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + // Blocked range + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(192, 168, 1, 1)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { .. } => panic!("Expected Block"), + ConnectDecision::Block => {} // Expected + } + + // Non-blocked range + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(8, 8, 8, 8)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "default"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + } + + #[tokio::test] + async fn test_block_hostname() { + let rules = vec![ + block_rule(vec!["blocked.com"]), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + // Blocked hostname + let location = NetLocation::new(Address::Hostname("blocked.com".to_string()), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { .. } => panic!("Expected Block for blocked.com"), + ConnectDecision::Block => {} // Expected + } + + // Subdomain of blocked hostname should also be blocked + let location = NetLocation::new(Address::Hostname("www.blocked.com".to_string()), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { .. } => panic!("Expected Block for www.blocked.com"), + ConnectDecision::Block => {} // Expected + } + } + + // ========================================================================= + // Default Rule Tests + // ========================================================================= + + #[tokio::test] + async fn test_no_default_rule_blocks() { + let rules = vec![allow_rule(vec!["192.168.1.0/24"], "lan")]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + // Non-matching address with no default rule should block + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(8, 8, 8, 8)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { .. } => panic!("Expected Block when no default rule"), + ConnectDecision::Block => {} // Expected + } + } + + #[tokio::test] + async fn test_default_decision() { + let rules = vec![ + allow_rule(vec!["192.168.1.0/24"], "lan"), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = ClientProxySelector::new(rules); + + // default_decision should return the catch-all rule + let decision = selector.default_decision(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "default"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + } + + #[tokio::test] + async fn test_default_decision_no_catchall() { + // Rules without a catch-all should return Block for default_decision + let rules = vec![allow_rule(vec!["192.168.1.0/24"], "lan")]; + let selector = ClientProxySelector::new(rules); + + let decision = selector.default_decision(); + match decision { + ConnectDecision::Allow { .. } => panic!("Expected Block when no catch-all rule"), + ConnectDecision::Block => {} // Expected + } + } + + #[tokio::test] + async fn test_default_decision_skips_override_rules() { + // Rules with override_address should not be used as default + let rules = vec![ + allow_rule_with_override(vec!["0.0.0.0/0"], "redirect", "10.0.0.1:8080"), + allow_rule(vec!["192.168.0.0/16"], "fallback"), + ]; + let selector = ClientProxySelector::new(rules); + + // default_decision should skip the override rule and use the next catch-all + let decision = selector.default_decision(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "fallback"); + } + ConnectDecision::Block => {} // Also acceptable if no valid default found + } + } + + // ========================================================================= + // Round-Robin Proxy Selection Tests + // ========================================================================= + + #[tokio::test] + async fn test_round_robin_proxy_selection() { + let rules = vec![allow_rule_multi( + vec!["0.0.0.0/0"], + vec!["proxy1", "proxy2", "proxy3"], + )]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + let mut seen = vec![]; + for _ in 0..6 { + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(1, 2, 3, 4)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + seen.push(client_proxy.name.clone()); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + } + + // Should cycle through proxies in order + assert_eq!(seen[0], "proxy1"); + assert_eq!(seen[1], "proxy2"); + assert_eq!(seen[2], "proxy3"); + assert_eq!(seen[3], "proxy1"); + assert_eq!(seen[4], "proxy2"); + assert_eq!(seen[5], "proxy3"); + } + + #[tokio::test] + async fn test_single_proxy_no_rotation() { + let rules = vec![allow_rule(vec!["0.0.0.0/0"], "only_proxy")]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + for _ in 0..5 { + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(1, 2, 3, 4)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "only_proxy"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + } + } + + // ========================================================================= + // Address Override Tests + // ========================================================================= + + #[tokio::test] + async fn test_address_override() { + let rules = vec![allow_rule_with_override( + vec!["0.0.0.0/0"], + "proxy", + "10.0.0.1:8080", + )]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(1, 2, 3, 4)), 443); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { + client_proxy, + remote_location, + } => { + assert_eq!(client_proxy.name, "proxy"); + assert_eq!(remote_location.address().to_string(), "10.0.0.1"); + assert_eq!(remote_location.port(), 8080); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + } + + #[tokio::test] + async fn test_address_override_preserve_port() { + // Override with port 0 should preserve the original port + let rules = vec![allow_rule_with_override( + vec!["0.0.0.0/0"], + "proxy", + "10.0.0.1:0", + )]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(1, 2, 3, 4)), 443); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { + remote_location, .. + } => { + assert_eq!(remote_location.address().to_string(), "10.0.0.1"); + assert_eq!(remote_location.port(), 443); // Original port preserved + } + ConnectDecision::Block => panic!("Expected Allow"), + } + } + + // ========================================================================= + // Multiple Masks Per Rule Tests + // ========================================================================= + + #[tokio::test] + async fn test_multiple_masks_in_rule() { + let rules = vec![ + allow_rule( + vec!["192.168.1.0/24", "192.168.2.0/24", "10.0.0.0/8"], + "internal", + ), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + // All three ranges should match "internal" + for ip in [ + Ipv4Addr::new(192, 168, 1, 1), + Ipv4Addr::new(192, 168, 2, 100), + Ipv4Addr::new(10, 50, 50, 50), + ] { + let location = NetLocation::new(Address::Ipv4(ip), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "internal", "Failed for {}", ip); + } + ConnectDecision::Block => panic!("Expected Allow for {}", ip), + } + } + + // Other IPs should go to default + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(8, 8, 8, 8)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "default"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + } + + #[tokio::test] + async fn test_multiple_hostnames_in_rule() { + let rules = vec![ + allow_rule( + vec!["google.com", "youtube.com", "gmail.com"], + "google_proxy", + ), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + for hostname in [ + "google.com", + "www.google.com", + "youtube.com", + "mail.gmail.com", + ] { + let location = NetLocation::new(Address::Hostname(hostname.to_string()), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "google_proxy", "Failed for {}", hostname); + } + ConnectDecision::Block => panic!("Expected Allow for {}", hostname), + } + } + } + + // ========================================================================= + // Edge Cases and Security Tests + // ========================================================================= + + #[tokio::test] + async fn test_empty_rules_blocks_all() { + let rules: Vec> = vec![]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(1, 2, 3, 4)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { .. } => panic!("Expected Block with empty rules"), + ConnectDecision::Block => {} // Expected + } + } + + #[tokio::test] + async fn test_ipv4_mapped_ipv6() { + // IPv4-mapped IPv6 addresses (::ffff:192.168.1.1) should be handled correctly + let rules = vec![ + allow_rule(vec!["192.168.1.0/24"], "lan"), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + // Regular IPv4 + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(192, 168, 1, 1)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "lan"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + } + + #[tokio::test] + async fn test_broadcast_and_special_addresses() { + let rules = vec![ + block_rule(vec!["255.255.255.255/32"]), // Broadcast + block_rule(vec!["0.0.0.0/32"]), // Unspecified + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + // Broadcast should be blocked + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::BROADCAST), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { .. } => panic!("Expected Block for broadcast"), + ConnectDecision::Block => {} // Expected + } + + // Unspecified should be blocked + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::UNSPECIFIED), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { .. } => panic!("Expected Block for unspecified"), + ConnectDecision::Block => {} // Expected + } + } + + #[tokio::test] + async fn test_case_sensitivity_hostname() { + let rules = vec![ + allow_rule(vec!["Example.COM"], "example_proxy"), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + // Note: DNS hostnames are case-insensitive by spec, but our implementation + // does case-sensitive matching. This test documents current behavior. + let location = NetLocation::new(Address::Hostname("example.com".to_string()), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + // Current behavior: case-sensitive, so this goes to default + // If we want case-insensitive, this should be "example_proxy" + assert_eq!(client_proxy.name, "default"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + } + + #[tokio::test] + async fn test_concurrent_access_round_robin() { + // Test that concurrent access to round-robin is safe (uses AtomicU32) + use std::sync::Arc as StdArc; + + let rules = vec![allow_rule_multi( + vec!["0.0.0.0/0"], + vec!["p1", "p2", "p3", "p4"], + )]; + let selector = StdArc::new(ClientProxySelector::new(rules)); + let resolver = mock_resolver(); + + let mut handles = vec![]; + for _ in 0..10 { + let selector = selector.clone(); + let resolver = resolver.clone(); + handles.push(tokio::spawn(async move { + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(1, 2, 3, 4)), 80); + let _ = selector.judge(location, &resolver).await; + })); + } + + for handle in handles { + handle.await.unwrap(); + } + // If we got here without panic, concurrent access is safe + } + + // ========================================================================= + // Regression Tests for Specific Scenarios + // ========================================================================= + + #[tokio::test] + async fn test_localhost_bypass() { + // Common pattern: bypass proxy for localhost + // Note: With resolve_rule_hostnames=false (default), hostname destinations + // won't match IP rules unless the hostname rule comes first. + let rules = vec![ + allow_rule(vec!["localhost"], "direct_host"), // Hostname rule first + allow_rule_with_override(vec!["127.0.0.0/8"], "direct", "127.0.0.1:0"), + allow_rule(vec!["0.0.0.0/0"], "proxy"), + ]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + // 127.0.0.1 should use direct (matches IP rule) + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(127, 0, 0, 1)), 8080); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "direct"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + + // localhost hostname should use direct_host (matches hostname rule) + let location = NetLocation::new(Address::Hostname("localhost".to_string()), 8080); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "direct_host"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + } + + #[tokio::test] + async fn test_mixed_ip_and_hostname_rules() { + // Note: With resolve_rule_hostnames=false (default), hostname destinations + // encountering IP rules will attempt to resolve the hostname. To avoid this, + // place hostname rules before IP rules, or use resolve_rule_hostnames=true + // with proper DNS mock mappings. + let rules = vec![ + block_rule(vec!["malware.com"]), + allow_rule(vec!["trusted.com"], "trusted"), // Hostname rules before IP rules + allow_rule(vec!["192.168.0.0/16"], "lan"), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = ClientProxySelector::new(rules); + let resolver = mock_resolver(); + + // Malware hostname blocked + let location = NetLocation::new(Address::Hostname("malware.com".to_string()), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + assert!(matches!(decision, ConnectDecision::Block)); + + // LAN IP allowed through lan proxy + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(192, 168, 1, 1)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "lan"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + + // Trusted hostname allowed through trusted proxy + let location = NetLocation::new(Address::Hostname("api.trusted.com".to_string()), 443); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "trusted"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + } + + // ========================================================================= + // Routing Cache Tests + // ========================================================================= + + #[test] + fn test_routing_cache_key_from_location_hostname() { + let location = NetLocation::new(Address::Hostname("Example.COM".to_string()), 443); + let key = RoutingCacheKey::from_location(&location); + match key { + RoutingCacheKey::Hostname { hostname, port } => { + assert_eq!(hostname, "example.com"); // Should be lowercased + assert_eq!(port, 443); + } + _ => panic!("Expected Hostname key"), + } + } + + #[test] + fn test_routing_cache_key_from_location_ipv4() { + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(192, 168, 1, 1)), 80); + let key = RoutingCacheKey::from_location(&location); + match key { + RoutingCacheKey::Ipv4 { addr, port } => { + assert_eq!(addr, Ipv4Addr::new(192, 168, 1, 1)); + assert_eq!(port, 80); + } + _ => panic!("Expected Ipv4 key"), + } + } + + #[test] + fn test_routing_cache_key_from_location_ipv6() { + let location = NetLocation::new(Address::Ipv6(Ipv6Addr::LOCALHOST), 8080); + let key = RoutingCacheKey::from_location(&location); + match key { + RoutingCacheKey::Ipv6 { addr, port } => { + assert_eq!(addr, Ipv6Addr::LOCALHOST); + assert_eq!(port, 8080); + } + _ => panic!("Expected Ipv6 key"), + } + } + + #[test] + fn test_routing_cache_basic_operations() { + let cache = RoutingCache::new(100); + assert!(cache.is_empty()); + assert_eq!(cache.len(), 0); + + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(1, 2, 3, 4)), 80); + + // Initially no entry + assert!(cache.get(&location).is_none()); + + // Insert and retrieve + cache.insert(&location, CachedDecision::Allow(5)); + assert_eq!(cache.len(), 1); + + let cached = cache.get(&location); + assert!(cached.is_some()); + match cached.unwrap() { + CachedDecision::Allow(idx) => assert_eq!(idx, 5), + CachedDecision::Block => panic!("Expected Allow"), + } + } + + #[test] + fn test_routing_cache_block_decision() { + let cache = RoutingCache::new(100); + let location = NetLocation::new(Address::Hostname("blocked.com".to_string()), 443); + + cache.insert(&location, CachedDecision::Block); + + let cached = cache.get(&location); + assert!(matches!(cached, Some(CachedDecision::Block))); + } + + #[test] + fn test_routing_cache_different_ports() { + let cache = RoutingCache::new(100); + let loc_80 = NetLocation::new(Address::Ipv4(Ipv4Addr::new(1, 2, 3, 4)), 80); + let loc_443 = NetLocation::new(Address::Ipv4(Ipv4Addr::new(1, 2, 3, 4)), 443); + + cache.insert(&loc_80, CachedDecision::Allow(1)); + cache.insert(&loc_443, CachedDecision::Allow(2)); + + // Same IP, different ports should be different cache entries + match cache.get(&loc_80).unwrap() { + CachedDecision::Allow(idx) => assert_eq!(idx, 1), + _ => panic!("Expected Allow(1)"), + } + match cache.get(&loc_443).unwrap() { + CachedDecision::Allow(idx) => assert_eq!(idx, 2), + _ => panic!("Expected Allow(2)"), + } + } + + #[test] + fn test_routing_cache_hostname_case_insensitive() { + let cache = RoutingCache::new(100); + + // Insert with uppercase + let loc_upper = NetLocation::new(Address::Hostname("EXAMPLE.COM".to_string()), 443); + cache.insert(&loc_upper, CachedDecision::Allow(1)); + + // Retrieve with lowercase - should find it + let loc_lower = NetLocation::new(Address::Hostname("example.com".to_string()), 443); + assert!(cache.get(&loc_lower).is_some()); + + // Retrieve with mixed case - should find it + let loc_mixed = NetLocation::new(Address::Hostname("Example.Com".to_string()), 443); + assert!(cache.get(&loc_mixed).is_some()); + } + + #[test] + fn test_routing_cache_clear() { + let cache = RoutingCache::new(100); + + // Add some entries + for i in 0..10 { + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(192, 168, 1, i)), 80); + cache.insert(&location, CachedDecision::Allow(i as usize)); + } + assert_eq!(cache.len(), 10); + + // Clear + cache.clear(); + assert!(cache.is_empty()); + assert_eq!(cache.len(), 0); + + // Verify entries are gone + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(192, 168, 1, 0)), 80); + assert!(cache.get(&location).is_none()); + } + + #[test] + fn test_routing_cache_lru_eviction() { + // Small cache to test eviction + let cache = RoutingCache::new(3); + + // Insert 3 entries + let loc1 = NetLocation::new(Address::Ipv4(Ipv4Addr::new(1, 0, 0, 1)), 80); + let loc2 = NetLocation::new(Address::Ipv4(Ipv4Addr::new(1, 0, 0, 2)), 80); + let loc3 = NetLocation::new(Address::Ipv4(Ipv4Addr::new(1, 0, 0, 3)), 80); + + cache.insert(&loc1, CachedDecision::Allow(1)); + cache.insert(&loc2, CachedDecision::Allow(2)); + cache.insert(&loc3, CachedDecision::Allow(3)); + assert_eq!(cache.len(), 3); + + // Insert 4th entry - should evict the oldest (loc1) + let loc4 = NetLocation::new(Address::Ipv4(Ipv4Addr::new(1, 0, 0, 4)), 80); + cache.insert(&loc4, CachedDecision::Allow(4)); + assert_eq!(cache.len(), 3); + + // loc1 should be evicted + assert!(cache.get(&loc1).is_none()); + // Others should still be present + assert!(cache.get(&loc2).is_some()); + assert!(cache.get(&loc3).is_some()); + assert!(cache.get(&loc4).is_some()); + } + + #[test] + fn test_routing_cache_overwrite() { + let cache = RoutingCache::new(100); + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(1, 2, 3, 4)), 80); + + // Insert Allow + cache.insert(&location, CachedDecision::Allow(1)); + match cache.get(&location).unwrap() { + CachedDecision::Allow(idx) => assert_eq!(idx, 1), + _ => panic!("Expected Allow(1)"), + } + + // Overwrite with Block + cache.insert(&location, CachedDecision::Block); + assert!(matches!(cache.get(&location), Some(CachedDecision::Block))); + + // Overwrite with different Allow + cache.insert(&location, CachedDecision::Allow(5)); + match cache.get(&location).unwrap() { + CachedDecision::Allow(idx) => assert_eq!(idx, 5), + _ => panic!("Expected Allow(5)"), + } + } + + /// Helper to create a selector with caching enabled (via resolve_rule_hostnames=true) + fn selector_with_cache(rules: Vec>) -> ClientProxySelector { + ClientProxySelector::with_options(rules, true) + } + + #[tokio::test] + async fn test_selector_cache_disabled_for_few_rules() { + // With few rules and resolve_rule_hostnames=false, cache should be disabled + let rules = vec![ + allow_rule(vec!["192.168.1.0/24"], "lan"), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = ClientProxySelector::new(rules); + assert!(!selector.is_cache_enabled()); + + // Routing should still work correctly + let resolver = mock_resolver(); + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(192, 168, 1, 100)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "lan"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + } + + #[tokio::test] + async fn test_selector_cache_enabled_with_dns_resolution() { + // With resolve_rule_hostnames=true, cache should be enabled even with few rules + let rules = vec![ + allow_rule(vec!["192.168.1.0/24"], "lan"), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = ClientProxySelector::with_options(rules, true); + assert!(selector.is_cache_enabled()); + } + + #[tokio::test] + async fn test_selector_cache_enabled_with_many_rules() { + // With many rules (>16), cache should be enabled even without DNS resolution + let mut rules: Vec> = (0..20) + .map(|i| allow_rule(vec![&format!("10.0.{}.0/24", i)], &format!("rule{}", i))) + .collect(); + rules.push(allow_rule(vec!["0.0.0.0/0"], "default")); + + let selector = ClientProxySelector::new(rules); + assert!(selector.is_cache_enabled()); + } + + #[tokio::test] + async fn test_selector_cache_hit() { + // Use selector_with_cache to ensure caching is enabled for this test + let rules = vec![ + allow_rule(vec!["192.168.1.0/24"], "lan"), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = selector_with_cache(rules); + let resolver = mock_resolver(); + + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(192, 168, 1, 100)), 80); + + // First call - cache miss + assert_eq!(selector.cache_size(), 0); + let decision = selector.judge(location.clone(), &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "lan"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + + // Should be cached now + assert_eq!(selector.cache_size(), 1); + + // Second call - cache hit (same result) + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "lan"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + + // Still only 1 entry (was a hit, not a new insert) + assert_eq!(selector.cache_size(), 1); + } + + #[tokio::test] + async fn test_selector_cache_block_decision() { + let rules = vec![ + block_rule(vec!["192.168.1.0/24"]), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = selector_with_cache(rules); + let resolver = mock_resolver(); + + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(192, 168, 1, 100)), 80); + + // First call - cache miss, should block + let decision = selector.judge(location.clone(), &resolver).await.unwrap(); + assert!(matches!(decision, ConnectDecision::Block)); + assert_eq!(selector.cache_size(), 1); + + // Second call - cache hit, still blocks + let decision = selector.judge(location, &resolver).await.unwrap(); + assert!(matches!(decision, ConnectDecision::Block)); + } + + #[tokio::test] + async fn test_selector_clear_cache() { + let rules = vec![allow_rule(vec!["0.0.0.0/0"], "default")]; + let selector = selector_with_cache(rules); + let resolver = mock_resolver(); + + // Make some cached entries + for i in 0..5 { + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(10, 0, 0, i)), 80); + let _ = selector.judge(location, &resolver).await.unwrap(); + } + assert_eq!(selector.cache_size(), 5); + + // Clear cache + selector.clear_cache(); + assert_eq!(selector.cache_size(), 0); + } + + #[tokio::test] + async fn test_selector_judge_uncached() { + let rules = vec![allow_rule(vec!["0.0.0.0/0"], "default")]; + let selector = selector_with_cache(rules); + let resolver = mock_resolver(); + + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(1, 2, 3, 4)), 80); + + // Use judge_uncached - should NOT populate cache + let decision = selector + .judge_uncached(location.clone(), &resolver) + .await + .unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "default"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + + // Cache should still be empty + assert_eq!(selector.cache_size(), 0); + } + + #[tokio::test] + async fn test_selector_cache_different_destinations() { + let rules = vec![ + allow_rule(vec!["192.168.0.0/16"], "lan"), + allow_rule(vec!["10.0.0.0/8"], "private"), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = selector_with_cache(rules); + let resolver = mock_resolver(); + + // Hit different rules + let loc_lan = NetLocation::new(Address::Ipv4(Ipv4Addr::new(192, 168, 1, 1)), 80); + let loc_private = NetLocation::new(Address::Ipv4(Ipv4Addr::new(10, 0, 0, 1)), 80); + let loc_public = NetLocation::new(Address::Ipv4(Ipv4Addr::new(8, 8, 8, 8)), 80); + + let d1 = selector.judge(loc_lan.clone(), &resolver).await.unwrap(); + let d2 = selector + .judge(loc_private.clone(), &resolver) + .await + .unwrap(); + let d3 = selector.judge(loc_public.clone(), &resolver).await.unwrap(); + + // Verify correct routing + match d1 { + ConnectDecision::Allow { client_proxy, .. } => assert_eq!(client_proxy.name, "lan"), + _ => panic!("Expected lan"), + } + match d2 { + ConnectDecision::Allow { client_proxy, .. } => assert_eq!(client_proxy.name, "private"), + _ => panic!("Expected private"), + } + match d3 { + ConnectDecision::Allow { client_proxy, .. } => assert_eq!(client_proxy.name, "default"), + _ => panic!("Expected default"), + } + + // All 3 should be cached + assert_eq!(selector.cache_size(), 3); + + // Verify cache hits return same results + let d1_cached = selector.judge(loc_lan, &resolver).await.unwrap(); + match d1_cached { + ConnectDecision::Allow { client_proxy, .. } => assert_eq!(client_proxy.name, "lan"), + _ => panic!("Expected lan from cache"), + } + } + + #[tokio::test] + async fn test_selector_cache_hostname_destinations() { + let rules = vec![ + allow_rule(vec!["google.com"], "google"), + allow_rule(vec!["0.0.0.0/0"], "default"), + ]; + let selector = selector_with_cache(rules); + let resolver = mock_resolver(); + + // Hostname match + let loc_google = NetLocation::new(Address::Hostname("www.google.com".to_string()), 443); + let decision = selector.judge(loc_google.clone(), &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "google"); + } + _ => panic!("Expected google"), + } + + // Should be cached + assert_eq!(selector.cache_size(), 1); + + // Case-insensitive cache hit + let loc_google_upper = + NetLocation::new(Address::Hostname("WWW.GOOGLE.COM".to_string()), 443); + let decision = selector.judge(loc_google_upper, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "google"); + } + _ => panic!("Expected google from cache"), + } + + // Should still be 1 entry (case-insensitive hit) + assert_eq!(selector.cache_size(), 1); + } + + #[tokio::test] + async fn test_selector_cache_with_custom_size() { + // Use resolve_rule_hostnames=true to enable caching with custom size + let rules = vec![allow_rule(vec!["0.0.0.0/0"], "default")]; + let selector = ClientProxySelector::with_options_and_cache_size(rules, true, 5); + let resolver = mock_resolver(); + + // Fill cache to capacity + for i in 0..5 { + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(10, 0, 0, i)), 80); + let _ = selector.judge(location, &resolver).await.unwrap(); + } + assert_eq!(selector.cache_size(), 5); + + // Add one more - should evict oldest + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(10, 0, 0, 100)), 80); + let _ = selector.judge(location, &resolver).await.unwrap(); + assert_eq!(selector.cache_size(), 5); // Still 5, one was evicted + + // First entry should be evicted - verify by checking the cache is full + // but we can't directly access the internal cache to verify eviction + // So we verify indirectly: if we add another entry, size stays at 5 + let location = NetLocation::new(Address::Ipv4(Ipv4Addr::new(10, 0, 0, 200)), 80); + let _ = selector.judge(location, &resolver).await.unwrap(); + assert_eq!(selector.cache_size(), 5); + } + + #[tokio::test] + async fn test_selector_concurrent_cache_access() { + use std::sync::Arc as StdArc; + + let rules = vec![allow_rule(vec!["0.0.0.0/0"], "default")]; + let selector = StdArc::new(selector_with_cache(rules)); + let resolver = mock_resolver(); + + // Spawn multiple tasks accessing the cache concurrently + let mut handles = vec![]; + for i in 0..10 { + let selector = selector.clone(); + let resolver = resolver.clone(); + handles.push(tokio::spawn(async move { + for j in 0..100 { + let location = + NetLocation::new(Address::Ipv4(Ipv4Addr::new(10, 0, i as u8, j as u8)), 80); + let decision = selector.judge(location, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "default"); + } + ConnectDecision::Block => panic!("Expected Allow"), + } + } + })); + } + + // Wait for all tasks + for handle in handles { + handle.await.unwrap(); + } + + // Should have cached entries (exact number depends on timing) + assert!(selector.cache_size() > 0); + } + + #[tokio::test] + async fn test_selector_cache_preserves_rule_priority() { + // Ensure caching doesn't break rule priority + let rules = vec![ + allow_rule(vec!["192.168.1.100/32"], "specific"), // Rule 0 + allow_rule(vec!["192.168.1.0/24"], "subnet"), // Rule 1 + allow_rule(vec!["192.168.0.0/16"], "network"), // Rule 2 + allow_rule(vec!["0.0.0.0/0"], "default"), // Rule 3 + ]; + let selector = selector_with_cache(rules); + let resolver = mock_resolver(); + + // Specific IP should match rule 0 (most specific) + let loc_specific = NetLocation::new(Address::Ipv4(Ipv4Addr::new(192, 168, 1, 100)), 80); + let decision = selector + .judge(loc_specific.clone(), &resolver) + .await + .unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "specific"); + } + _ => panic!("Expected specific"), + } + + // Verify from cache - should still be "specific" + let decision_cached = selector.judge(loc_specific, &resolver).await.unwrap(); + match decision_cached { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "specific"); + } + _ => panic!("Expected specific from cache"), + } + + // Different IP in same subnet should match rule 1 + let loc_subnet = NetLocation::new(Address::Ipv4(Ipv4Addr::new(192, 168, 1, 50)), 80); + let decision = selector.judge(loc_subnet, &resolver).await.unwrap(); + match decision { + ConnectDecision::Allow { client_proxy, .. } => { + assert_eq!(client_proxy.name, "subnet"); + } + _ => panic!("Expected subnet"), + } + } +} diff --git a/src/config/mod.rs b/src/config/mod.rs index 05771d41..291136fc 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -3,7 +3,6 @@ mod types; pub use types::*; use std::collections::HashMap; -use std::path::Path; use log::debug; @@ -100,7 +99,11 @@ pub async fn convert_cert_paths(all_configs: Vec) -> std::io::Result<(Ve // we do this as a separate step because we can't load the data in an async manner // right away due to recursion. for config in server_configs.iter_mut() { - gather_pem_file_paths_from_server_config(config, &path_pem_configs, &mut unknown_pem_paths); + gather_pem_file_paths_from_server_config( + config, + &path_pem_configs, + &mut unknown_pem_paths, + )?; } for config in client_group_configs.iter_mut() { @@ -227,7 +230,7 @@ fn gather_pem_file_paths_from_server_config( server_config: &mut ServerConfig, known_pem_paths: &HashMap, unknown_pem_paths: &mut HashMap, -) { +) -> std::io::Result<()> { // Check QUIC settings if let Some(ref mut quic_settings) = server_config.quic_settings { process_pem_path(&mut quic_settings.cert, known_pem_paths, unknown_pem_paths); @@ -242,12 +245,13 @@ fn gather_pem_file_paths_from_server_config( &mut server_config.protocol, known_pem_paths, unknown_pem_paths, - ); + )?; // Check rules for rule in server_config.rules.iter_mut() { gather_pem_file_paths_from_rule(rule, known_pem_paths, unknown_pem_paths); } + Ok(()) } fn gather_pem_file_paths_from_client_config_group( @@ -260,31 +264,126 @@ fn gather_pem_file_paths_from_client_config_group( } } +/// Validates that Vision is only enabled when the inner protocol is VLESS. +/// +/// Vision (XTLS-RPRX-Vision) is a TLS-in-TLS optimization protocol that requires +/// VLESS as the inner protocol to function correctly. +fn validate_vision_protocol( + vision_enabled: bool, + protocol: &ServerProxyConfig, + config_type: &str, +) -> std::io::Result<()> { + if !vision_enabled { + // Vision not enabled, no validation needed + return Ok(()); + } + + // Vision is enabled - verify inner protocol is VLESS + match protocol { + ServerProxyConfig::Vless { .. } => { + // Valid: Vision + VLESS + Ok(()) + } + other_protocol => { + // Invalid: Vision requires VLESS + Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!( + "{} config has vision=true but inner protocol is {} (not VLESS). \ + Vision (XTLS-RPRX-Vision) requires VLESS as the inner protocol. \ + Either set vision=false or change the inner protocol to VLESS.", + config_type, other_protocol + ), + )) + } + } +} + +/// Validates that Vision is only enabled when the inner protocol is VLESS (client-side) +fn validate_client_vision_protocol( + vision_enabled: bool, + protocol: &ClientProxyConfig, + config_type: &str, +) -> std::io::Result<()> { + if !vision_enabled { + return Ok(()); + } + + match protocol { + ClientProxyConfig::Vless { .. } => Ok(()), + other_protocol => Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!( + "{} client config has vision=true but inner protocol is {} (not VLESS). \ + Vision (XTLS-RPRX-Vision) requires VLESS as the inner protocol. \ + Either set vision=false or change the inner protocol to VLESS.", + config_type, + other_protocol.protocol_name() + ), + )), + } +} + +/// Recursive validation of client proxy config structure (Vision rules, etc.) +fn validate_client_proxy_structure(config: &ClientProxyConfig) -> std::io::Result<()> { + match config { + ClientProxyConfig::Tls(tls_config) => { + validate_client_vision_protocol(tls_config.vision, &tls_config.protocol, "TLS")?; + validate_client_proxy_structure(&tls_config.protocol) + } + + ClientProxyConfig::Reality { + vision, protocol, .. + } => { + validate_client_vision_protocol(*vision, protocol, "Reality")?; + validate_client_proxy_structure(protocol) + } + + ClientProxyConfig::ShadowTls { protocol, .. } => validate_client_proxy_structure(protocol), + + ClientProxyConfig::Websocket(ws_config) => { + validate_client_proxy_structure(&ws_config.protocol) + } + + // Leaf protocols + _ => Ok(()), + } +} + fn gather_pem_file_paths_from_server_proxy( server_proxy: &mut ServerProxyConfig, known_pem_paths: &HashMap, unknown_pem_paths: &mut HashMap, -) { +) -> std::io::Result<()> { match server_proxy { ServerProxyConfig::Tls { tls_targets, default_tls_target, shadowtls_targets, + reality_targets, .. } => { // Process TLS targets - for (_, tls_config) in tls_targets.iter_mut() { + for (sni, tls_config) in tls_targets.iter_mut() { process_pem_path(&mut tls_config.cert, known_pem_paths, unknown_pem_paths); process_pem_path(&mut tls_config.key, known_pem_paths, unknown_pem_paths); for cert in tls_config.client_ca_certs.iter_mut() { process_pem_path(cert, known_pem_paths, unknown_pem_paths); } + + // Validate Vision configuration + validate_vision_protocol( + tls_config.vision, + &tls_config.protocol, + &format!("TLS target '{}'", sni), + )?; + // Recurse into inner protocol gather_pem_file_paths_from_server_proxy( &mut tls_config.protocol, known_pem_paths, unknown_pem_paths, - ); + )?; // Check override rules for rule in tls_config.override_rules.iter_mut() { gather_pem_file_paths_from_rule(rule, known_pem_paths, unknown_pem_paths); @@ -298,12 +397,20 @@ fn gather_pem_file_paths_from_server_proxy( for cert in tls_config.client_ca_certs.iter_mut() { process_pem_path(cert, known_pem_paths, unknown_pem_paths); } + + // Validate Vision configuration + validate_vision_protocol( + tls_config.vision, + &tls_config.protocol, + "default TLS target", + )?; + // Recurse into inner protocol gather_pem_file_paths_from_server_proxy( &mut tls_config.protocol, known_pem_paths, unknown_pem_paths, - ); + )?; // Check override rules for rule in tls_config.override_rules.iter_mut() { gather_pem_file_paths_from_rule(rule, known_pem_paths, unknown_pem_paths); @@ -323,12 +430,33 @@ fn gather_pem_file_paths_from_server_proxy( &mut shadowtls_config.protocol, known_pem_paths, unknown_pem_paths, - ); + )?; // Check override rules for rule in shadowtls_config.override_rules.iter_mut() { gather_pem_file_paths_from_rule(rule, known_pem_paths, unknown_pem_paths); } } + + // Process Reality targets + for (sni, reality_config) in reality_targets.iter_mut() { + // Validate Vision configuration + validate_vision_protocol( + reality_config.vision, + &reality_config.protocol, + &format!("Reality target '{}'", sni), + )?; + + // Recurse into inner protocol + gather_pem_file_paths_from_server_proxy( + &mut reality_config.protocol, + known_pem_paths, + unknown_pem_paths, + )?; + // Check override rules + for rule in reality_config.override_rules.iter_mut() { + gather_pem_file_paths_from_rule(rule, known_pem_paths, unknown_pem_paths); + } + } } ServerProxyConfig::Websocket { targets } => { for websocket_config in targets.iter_mut() { @@ -337,15 +465,17 @@ fn gather_pem_file_paths_from_server_proxy( &mut websocket_config.protocol, known_pem_paths, unknown_pem_paths, - ); + )?; // Check override rules for rule in websocket_config.override_rules.iter_mut() { gather_pem_file_paths_from_rule(rule, known_pem_paths, unknown_pem_paths); } } } + ServerProxyConfig::Vless { .. } => {} _ => {} } + Ok(()) } fn gather_pem_file_paths_from_rule( @@ -499,7 +629,135 @@ fn validate_client_fingerprints( if client_fingerprints.iter().any(|fp| fp == "any") { let _ = std::mem::replace(client_fingerprints, NoneOrSome::Unspecified); } else { - let _ = crate::rustls_util::process_fingerprints(&client_fingerprints.clone().into_vec())?; + let _ = crate::rustls_config_util::process_fingerprints( + &client_fingerprints.clone().into_vec(), + )?; + } + + Ok(()) +} + +/// Validates Reality short_ids to ensure they are valid hexadecimal strings. +/// +/// Each short_id must be: +/// - Valid hexadecimal (0-9, a-f, A-F) +/// - Even length (to decode to whole bytes) +/// - Maximum 16 characters (8 bytes when decoded) +/// +/// Validates Reality private_key to ensure it's a valid base64url-encoded X25519 key. +/// +/// The private key must: +/// - Be valid base64url encoding +/// - Decode to exactly 32 bytes (X25519 key size) +fn validate_reality_private_key(private_key: &str, target_name: &str) -> std::io::Result<()> { + use crate::reality::decode_private_key; + + decode_private_key(private_key).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!( + "REALITY target '{}': invalid private_key: {}", + target_name, e + ), + ) + })?; + + Ok(()) +} + +/// Validates Reality client short_id to ensure it's a valid hexadecimal string +fn validate_reality_client_short_id(short_id: &str) -> std::io::Result<()> { + use crate::reality::decode_short_id; + + // Check length + if short_id.len() > 16 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!( + "Reality client short_id is too long: '{}' ({} chars, max 16)", + short_id, + short_id.len() + ), + )); + } + + // Check for valid hexadecimal characters + if !short_id.chars().all(|c| c.is_ascii_hexdigit()) { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!( + "Reality client short_id contains non-hexadecimal characters: '{}'. \ + Only 0-9, a-f, and A-F are allowed.", + short_id + ), + )); + } + + // Try to decode to validate format + decode_short_id(short_id).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("Reality client short_id decode failed: {}", e), + ) + })?; + + Ok(()) +} + +fn validate_reality_server_short_ids( + short_ids: &crate::option_util::OneOrSome, + target_name: &str, +) -> std::io::Result<()> { + // OneOrSome ensures at least one short_id is always present (default is all zeros) + + // Check if it's exactly one short_id with the default all-zeros value + let is_default = match short_ids { + crate::option_util::OneOrSome::One(id) => { + id == crate::config::types::DEFAULT_REALITY_SHORT_ID + } + crate::option_util::OneOrSome::Some(ids) => { + ids.len() == 1 && ids[0] == crate::config::types::DEFAULT_REALITY_SHORT_ID + } + }; + + if is_default { + log::warn!( + "Reality server '{}' using default short_ids (all zeros). \ + For better security in production, configure explicit short_ids.", + target_name + ); + } + + for (i, short_id) in short_ids.iter().enumerate() { + // Check length + if short_id.len() > 16 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!( + "REALITY target '{}': short_ids[{}] is too long: '{}' ({} chars, max 16)", + target_name, + i, + short_id, + short_id.len() + ), + )); + } + + // Check for valid hexadecimal characters + if !short_id.chars().all(|c| c.is_ascii_hexdigit()) { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!( + "REALITY target '{}': short_ids[{}] contains non-hexadecimal characters: '{}'. \ + Only 0-9, a-f, and A-F are allowed.", + target_name, i, short_id + ), + )); + } + + // Empty string is technically valid (will be left-padded to 16 zeros) + // but odd-length strings are also valid (will be left-padded) + // The decode_short_id function handles padding internally } Ok(()) @@ -566,7 +824,9 @@ fn validate_server_fingerprints( if server_fingerprints.iter().any(|fp| fp == "any") { let _ = std::mem::replace(server_fingerprints, NoneOrSome::Unspecified); } else { - let _ = crate::rustls_util::process_fingerprints(&server_fingerprints.clone().into_vec())?; + let _ = crate::rustls_config_util::process_fingerprints( + &server_fingerprints.clone().into_vec(), + )?; } Ok(()) @@ -576,21 +836,45 @@ fn validate_client_proxy_config( client_proxy_config: &mut ClientProxyConfig, named_pems: &HashMap, ) -> std::io::Result<()> { - if let ClientProxyConfig::Tls(ref mut tls_config) = client_proxy_config { - // Embed certificates before validation - embed_optional_pem_from_map(&mut tls_config.cert, named_pems); - embed_optional_pem_from_map(&mut tls_config.key, named_pems); + // First validate the structure (Vision rules, etc.) + validate_client_proxy_structure(client_proxy_config)?; - if tls_config.cert.is_none() != tls_config.key.is_none() { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidInput, - "Both client cert and key have to be specified, or both have to be omitted", - )); + match client_proxy_config { + ClientProxyConfig::Reality { + short_id, protocol, .. + } => { + // Validate short_id format + validate_reality_client_short_id(short_id)?; + + // Warn if using default all-zeros short_id + if short_id == crate::config::types::DEFAULT_REALITY_SHORT_ID { + log::warn!( + "Reality client using default short_id (all zeros). \ + For better security in production, configure an explicit short_id that matches your server." + ); + } + + // Recurse into inner protocol + validate_client_proxy_config(protocol, named_pems)?; } - validate_server_fingerprints(&mut tls_config.server_fingerprints)?; - // Recursively validate inner protocol - validate_client_proxy_config(&mut tls_config.protocol, named_pems)?; + ClientProxyConfig::Tls(ref mut tls_config) => { + // Embed certificates before validation + embed_optional_pem_from_map(&mut tls_config.cert, named_pems); + embed_optional_pem_from_map(&mut tls_config.key, named_pems); + + if tls_config.cert.is_none() != tls_config.key.is_none() { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "Both client cert and key have to be specified, or both have to be omitted", + )); + } + validate_server_fingerprints(&mut tls_config.server_fingerprints)?; + + // Recursively validate inner protocol + validate_client_proxy_config(&mut tls_config.protocol, named_pems)?; + } + _ => {} } Ok(()) } @@ -612,8 +896,19 @@ fn validate_server_proxy_config( tls_targets, default_tls_target, shadowtls_targets, + reality_targets, tls_buffer_size, } => { + if tls_targets.is_empty() + && default_tls_target.is_none() + && shadowtls_targets.is_empty() + && reality_targets.is_empty() + { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "TLS server has no entries", + )); + } for (_, tls_server_config) in tls_targets.iter_mut() { // Embed certificates embed_pem_from_map(&mut tls_server_config.cert, named_pems); @@ -704,6 +999,53 @@ fn validate_server_proxy_config( } } + for (sni_hostname, reality_config) in reality_targets.iter_mut() { + if tls_targets.contains_key(sni_hostname) { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!( + "duplicated SNI hostname between TLS and REALITY targets: {sni_hostname}" + ), + )); + } + if shadowtls_targets.contains_key(sni_hostname) { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!( + "duplicated SNI hostname between ShadowTLS and REALITY targets: {sni_hostname}" + ), + )); + } + + // Validate Reality private_key (must be valid base64url-encoded X25519 key) + validate_reality_private_key(&reality_config.private_key, sni_hostname)?; + + // Validate Reality short_ids (must be valid hex strings) + validate_reality_server_short_ids(&reality_config.short_ids, sni_hostname)?; + + // Validate inner protocol + validate_server_proxy_config( + &mut reality_config.protocol, + client_groups, + rule_groups, + named_pems, + )?; + + // Validate override_rules + ConfigSelection::replace_none_or_some_groups( + &mut reality_config.override_rules, + rule_groups, + )?; + + for rule_config_selection in reality_config.override_rules.iter_mut() { + validate_rule_config( + rule_config_selection.unwrap_config_mut(), + client_groups, + named_pems, + )?; + } + } + if let Some(size) = tls_buffer_size { if *size < MIN_TLS_BUFFER_SIZE { return Err(std::io::Error::new( @@ -821,19 +1163,6 @@ fn embed_optional_pem_from_map(pem: &mut Option, named_pems: &HashMap std::io::Result<()> { - let yaml_str = serde_yaml::to_string(configs).map_err(|e| { - std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!("Failed to serialize to YAML: {e}"), - ) - })?; - - tokio::fs::write(path, yaml_str).await?; - Ok(()) -} - #[cfg(test)] mod tests { use super::*; diff --git a/src/config/types.rs b/src/config/types.rs index 606973b1..efa7ddfe 100644 --- a/src/config/types.rs +++ b/src/config/types.rs @@ -6,6 +6,9 @@ use serde::{Deserialize, Serialize}; use crate::address::{NetLocation, NetLocationMask, NetLocationPortRange}; use crate::option_util::{NoneOrOne, NoneOrSome, OneOrSome}; +/// Default Reality short_id: all zeros (16 hex chars = 8 bytes of zeros) +pub const DEFAULT_REALITY_SHORT_ID: &str = "0000000000000000"; + fn default_true() -> bool { true } @@ -14,6 +17,86 @@ fn default_snell_udp_num_sockets() -> usize { 1 } +fn default_reality_client_short_id() -> String { + DEFAULT_REALITY_SHORT_ID.to_string() +} + +fn default_reality_server_short_ids() -> OneOrSome { + OneOrSome::One(DEFAULT_REALITY_SHORT_ID.to_string()) +} + +/// Custom deserializer for ServerProxyConfig::Vmess that validates legacy force_aead field +fn deserialize_vmess_server<'de, D>(deserializer: D) -> Result<(String, String, bool), D::Error> +where + D: serde::Deserializer<'de>, +{ + use serde::de::Error; + + #[derive(Deserialize)] + struct VmessServerTemp { + cipher: String, + user_id: String, + #[serde(default = "default_true")] + udp_enabled: bool, + #[serde(default)] + force_aead: Option, + } + + let temp = VmessServerTemp::deserialize(deserializer)?; + + // Check if force_aead was explicitly set + if let Some(force_aead_value) = temp.force_aead { + if !force_aead_value { + return Err(Error::custom( + "Non-AEAD VMess mode (force_aead=false) is no longer supported. \ + Please remove the force_aead field from your configuration, or set it to true.", + )); + } + // Warn about deprecated field + log::warn!( + "The 'force_aead' field in VMess server configuration is deprecated and will be removed in a future version. \ + AEAD mode is now always enabled. Please remove this field from your configuration." + ); + } + + Ok((temp.cipher, temp.user_id, temp.udp_enabled)) +} + +/// Custom deserializer for ClientProxyConfig::Vmess that validates legacy aead field +fn deserialize_vmess_client<'de, D>(deserializer: D) -> Result<(String, String), D::Error> +where + D: serde::Deserializer<'de>, +{ + use serde::de::Error; + + #[derive(Deserialize)] + struct VmessClientTemp { + cipher: String, + user_id: String, + #[serde(default)] + aead: Option, + } + + let temp = VmessClientTemp::deserialize(deserializer)?; + + // Check if aead was explicitly set + if let Some(aead_value) = temp.aead { + if !aead_value { + return Err(Error::custom( + "Non-AEAD VMess mode (aead=false) is no longer supported. \ + Please remove the aead field from your configuration, or set it to true.", + )); + } + // Warn about deprecated field + log::warn!( + "The 'aead' field in VMess client configuration is deprecated and will be removed in a future version. \ + AEAD mode is now always enabled. Please remove this field from your configuration." + ); + } + + Ok((temp.cipher, temp.user_id)) +} + #[derive(Debug, Clone, Deserialize, Serialize, PartialEq)] #[serde(rename_all = "lowercase")] pub enum BindLocation { @@ -30,20 +113,15 @@ impl std::fmt::Display for BindLocation { } } -#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, Default)] #[serde(rename_all = "lowercase")] pub enum Transport { + #[default] Tcp, Quic, Udp, } -impl Default for Transport { - fn default() -> Self { - Self::Tcp - } -} - #[derive(Debug, Clone, Deserialize, Serialize)] pub struct TcpConfig { #[serde(default = "default_true")] @@ -322,6 +400,51 @@ pub struct ShadowsocksConfig { pub password: String, } +// REALITY Protocol Configuration + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct RealityServerConfig { + /// X25519 private key (32 bytes, base64url encoded) + pub private_key: String, + + /// List of valid short IDs (hex strings, 0-16 chars each) + #[serde(alias = "short_id", default = "default_reality_server_short_ids")] + pub short_ids: OneOrSome, + + /// Fallback destination (e.g., "example.com:443") + pub dest: NetLocation, + + /// Maximum timestamp difference in milliseconds (optional) + #[serde(default = "default_reality_time_diff")] + pub max_time_diff: Option, + + /// Minimum client version [major, minor, patch] (optional) + #[serde(default)] + pub min_client_version: Option<[u8; 3]>, + + /// Maximum client version [major, minor, patch] (optional) + #[serde(default)] + pub max_client_version: Option<[u8; 3]>, + + /// Enable XTLS-Vision protocol for TLS-in-TLS optimization. + /// When enabled, the inner protocol MUST be VLESS. + /// Vision detects TLS-in-TLS scenarios and switches to Direct mode for zero-copy performance. + /// Reality provides censorship resistance while Vision provides performance optimization. + #[serde(default)] + pub vision: bool, + /// Inner protocol (VLESS, Trojan, etc.) + pub protocol: ServerProxyConfig, + + /// Override rules + #[serde(alias = "override_rule", default)] + pub override_rules: NoneOrSome>, +} + +fn default_reality_time_diff() -> Option { + // 1 minute + Some(1000 * 60) +} + #[derive(Debug, Clone, Deserialize, Serialize)] pub struct TlsServerConfig { pub cert: String, @@ -352,6 +475,12 @@ pub struct TlsServerConfig { #[serde(alias = "client_fingerprint", default)] pub client_fingerprints: NoneOrSome, + /// Enable XTLS-Vision protocol for TLS-in-TLS optimization. + /// When enabled, the inner protocol MUST be VLESS. + /// Vision detects TLS-in-TLS scenarios and switches to Direct mode for zero-copy performance. + /// Requires TLS 1.3. + #[serde(default)] + pub vision: bool, pub protocol: ServerProxyConfig, #[serde(alias = "override_rule", default)] @@ -556,16 +685,17 @@ pub enum ServerProxyConfig { default_tls_target: Option>, #[serde(default)] shadowtls_targets: HashMap, + #[serde(default)] + reality_targets: HashMap, #[serde(default)] tls_buffer_size: Option, }, + #[serde(deserialize_with = "deserialize_vmess_server")] Vmess { cipher: String, user_id: String, #[serde(default = "default_true")] - force_aead: bool, - #[serde(default = "default_true")] udp_enabled: bool, }, #[serde(alias = "ws")] @@ -589,39 +719,47 @@ pub enum ServerProxyConfig { impl std::fmt::Display for ServerProxyConfig { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!( - f, - "{}", - match self { - Self::Http { .. } => "HTTP", - Self::Socks { .. } => "SOCKS", - Self::Shadowsocks { .. } => "Shadowsocks", - Self::Snell { .. } => "Snell", - Self::Vless { .. } => "Vless", - Self::Trojan { .. } => "Trojan", - Self::Tls { - tls_targets, - default_tls_target, - shadowtls_targets, - .. - } => { - let has_tls = !tls_targets.is_empty() || !default_tls_target.is_none(); - let has_shadowtls = !shadowtls_targets.is_empty(); - if has_tls && has_shadowtls { - "TLS+ShadowTLSv3" - } else if has_shadowtls { - "ShadowTLSv3" - } else { - "TLS" - } + match self { + Self::Http { .. } => write!(f, "HTTP"), + Self::Socks { .. } => write!(f, "SOCKS"), + Self::Shadowsocks { .. } => write!(f, "Shadowsocks"), + Self::Snell { .. } => write!(f, "Snell"), + Self::Vless { .. } => write!(f, "Vless"), + Self::Trojan { .. } => write!(f, "Trojan"), + Self::Tls { + tls_targets, + default_tls_target, + shadowtls_targets, + reality_targets, + .. + } => { + let mut parts = vec![]; + + if !tls_targets.is_empty() { + parts.push("TLS"); } - Self::Vmess { .. } => "Vmess", - Self::Websocket { .. } => "Websocket", - Self::PortForward { .. } => "Portforward", - Self::Hysteria2 { .. } => "Hysteria2", - Self::TuicV5 { .. } => "TuicV5", + + if !reality_targets.is_empty() { + parts.push("REALITY"); + } + if !shadowtls_targets.is_empty() { + parts.push("ShadowTLSv3"); + } + if tls_targets.values().any(|cfg| cfg.vision) + || default_tls_target.as_ref().is_some_and(|cfg| cfg.vision) + || reality_targets.values().any(|cfg| cfg.vision) + { + parts.push("Vision"); + } + + write!(f, "{}", parts.join("+")) } - ) + Self::Vmess { .. } => write!(f, "Vmess"), + Self::Websocket { .. } => write!(f, "Websocket"), + Self::PortForward { .. } => write!(f, "Portforward"), + Self::Hysteria2 { .. } => write!(f, "Hysteria2"), + Self::TuicV5 { .. } => write!(f, "TuicV5"), + } } } @@ -860,24 +998,160 @@ pub enum ClientProxyConfig { #[serde(default)] shadowsocks: Option, }, + Reality { + public_key: String, + #[serde(default = "default_reality_client_short_id")] + short_id: String, + #[serde(default)] + sni_hostname: Option, + + /// Enable XTLS-Vision protocol for TLS-in-TLS optimization. + /// When enabled, the inner protocol MUST be VLESS. + #[serde(default)] + vision: bool, + + protocol: Box, + }, + #[serde(alias = "shadowtls")] + ShadowTls { + /// ShadowTLS password for authentication + password: String, + + /// Optional SNI hostname override + #[serde(default)] + sni_hostname: Option, + + /// Inner protocol (typically VLESS, Trojan, etc.) + protocol: Box, + }, + #[serde(deserialize_with = "deserialize_tls_variant")] Tls(TlsClientConfig), + #[serde(deserialize_with = "deserialize_vmess_client")] Vmess { cipher: String, user_id: String, - #[serde(default = "default_true")] - aead: bool, }, #[serde(alias = "ws")] Websocket(WebsocketClientConfig), + #[serde(alias = "noop")] + PortForward, } impl ClientProxyConfig { pub fn is_direct(&self) -> bool { matches!(self, ClientProxyConfig::Direct) } + + /// Returns the protocol name for display/error messages + pub fn protocol_name(&self) -> &str { + match self { + ClientProxyConfig::Direct => "Direct", + ClientProxyConfig::Http { .. } => "HTTP", + ClientProxyConfig::Socks { .. } => "SOCKS5", + ClientProxyConfig::Shadowsocks(..) => "Shadowsocks", + ClientProxyConfig::Snell(..) => "Snell", + ClientProxyConfig::Vless { .. } => "VLESS", + ClientProxyConfig::Trojan { .. } => "Trojan", + ClientProxyConfig::Reality { .. } => "Reality", + ClientProxyConfig::Tls(..) => "TLS", + ClientProxyConfig::ShadowTls { .. } => "ShadowTLS", + ClientProxyConfig::Vmess { .. } => "VMess", + ClientProxyConfig::Websocket(..) => "WebSocket", + ClientProxyConfig::PortForward => "PortForward", + } + } } -#[derive(Debug, Clone, Deserialize, Serialize)] +/// Custom deserializer for TlsClientConfig that handles deprecated shadowtls_password field +fn deserialize_tls_client_config<'de, D>(deserializer: D) -> Result +where + D: serde::Deserializer<'de>, +{ + use serde::de::Error; + + #[derive(Deserialize)] + struct TlsClientConfigTemp { + #[serde(default = "default_true")] + verify: bool, + #[serde(alias = "server_fingerprint", default)] + server_fingerprints: NoneOrSome, + #[serde(default)] + sni_hostname: NoneOrOne, + #[serde(alias = "alpn_protocol", default)] + alpn_protocols: NoneOrSome, + #[serde(default)] + tls_buffer_size: Option, + #[serde(default)] + key: Option, + #[serde(default)] + cert: Option, + #[serde(default)] + shadowtls_password: Option, + #[serde(default)] + vision: bool, + protocol: Box, + } + + let temp = TlsClientConfigTemp::deserialize(deserializer)?; + + // Check for mutually exclusive fields + if temp.vision && temp.shadowtls_password.is_some() { + return Err(Error::custom( + "TLS client config cannot have both vision=true and shadowtls_password set. \ + Vision and ShadowTLS are incompatible. \ + Use either 'vision: true' with regular TLS, or 'type: shadowtls' for ShadowTLS.", + )); + } + + // Check if deprecated shadowtls_password was used + if let Some(password) = temp.shadowtls_password { + log::warn!( + "The 'shadowtls_password' field in TLS client configuration is deprecated. \ + Please use 'type: shadowtls' with 'password' field instead. \ + This field will be removed in a future version." + ); + + // Transform to ShadowTLS variant internally by wrapping protocol + return Ok(TlsClientConfig { + verify: temp.verify, + server_fingerprints: temp.server_fingerprints, + sni_hostname: temp.sni_hostname.clone(), + alpn_protocols: temp.alpn_protocols, + tls_buffer_size: temp.tls_buffer_size, + key: temp.key, + cert: temp.cert, + vision: false, + protocol: Box::new(ClientProxyConfig::ShadowTls { + password, + sni_hostname: temp.sni_hostname.into_option(), + protocol: temp.protocol, + }), + }); + } + + // Normal case - no shadowtls_password + Ok(TlsClientConfig { + verify: temp.verify, + server_fingerprints: temp.server_fingerprints, + sni_hostname: temp.sni_hostname, + alpn_protocols: temp.alpn_protocols, + tls_buffer_size: temp.tls_buffer_size, + key: temp.key, + cert: temp.cert, + vision: temp.vision, + protocol: temp.protocol, + }) +} + +/// Variant deserializer for Tls in ClientProxyConfig enum +fn deserialize_tls_variant<'de, D>(deserializer: D) -> Result +where + D: serde::Deserializer<'de>, +{ + deserialize_tls_client_config(deserializer) +} + +#[derive(Debug, Clone, Serialize)] pub struct TlsClientConfig { #[serde(default = "default_true")] pub verify: bool, @@ -893,8 +1167,13 @@ pub struct TlsClientConfig { pub key: Option, #[serde(default)] pub cert: Option, + + /// Enable XTLS-Vision protocol for TLS-in-TLS optimization. + /// When enabled, the inner protocol MUST be VLESS. + /// Requires TLS 1.3. #[serde(default)] - pub shadowtls_password: Option, + pub vision: bool, + pub protocol: Box, } @@ -1037,7 +1316,7 @@ mod tests { transport: Transport::Tcp, tcp_settings: Some(TcpConfig { no_delay: true }), quic_settings: None, - rules: NoneOrSome::One(ConfigSelection::GroupName("test-rules".to_string())), + rules: NoneOrSome::None, } } @@ -1067,7 +1346,7 @@ mod tests { transport: Transport::Tcp, tcp_settings: None, quic_settings: None, - rules: NoneOrSome::One(ConfigSelection::GroupName("test-rules".to_string())), + rules: NoneOrSome::None, } } @@ -1124,6 +1403,7 @@ mod tests { alpn_protocols: NoneOrSome::Some(vec!["h2".to_string(), "http/1.1".to_string()]), client_ca_certs: NoneOrSome::One("ca.crt".to_string()), client_fingerprints: NoneOrSome::One("abc123".to_string()), + vision: false, protocol: ServerProxyConfig::Http { username: None, password: None, @@ -1144,6 +1424,7 @@ mod tests { alpn_protocols: NoneOrSome::None, client_ca_certs: NoneOrSome::None, client_fingerprints: NoneOrSome::None, + vision: false, protocol: ServerProxyConfig::Http { username: None, password: None, @@ -1151,6 +1432,7 @@ mod tests { override_rules: NoneOrSome::None, })), shadowtls_targets: HashMap::new(), + reality_targets: HashMap::new(), tls_buffer_size: Some(8192), }, transport: Transport::Tcp, @@ -1168,7 +1450,6 @@ mod tests { protocol: ServerProxyConfig::Vmess { cipher: "aes-128-gcm".to_string(), user_id: "b831381d-6324-4d53-ad4f-8cda48b30811".to_string(), - force_aead: true, udp_enabled: false, }, transport: Transport::Tcp, @@ -1628,6 +1909,7 @@ mod tests { client_ca_certs: NoneOrSome::None, client_fingerprints: NoneOrSome::Unspecified, alpn_protocols: NoneOrSome::Unspecified, + vision: false, protocol: ServerProxyConfig::Http { username: None, password: None, @@ -1639,13 +1921,14 @@ mod tests { .collect(), default_tls_target: None, shadowtls_targets: HashMap::new(), + reality_targets: HashMap::new(), tls_buffer_size: None, }, transport: Transport::Tcp, tcp_settings: None, quic_settings: None, rules: NoneOrSome::One(ConfigSelection::GroupName("allow-all-direct".to_string())), - }), + }), ]; let result = validate_configs_test(configs).await; @@ -2065,7 +2348,7 @@ mod tests { num_endpoints: 1, }), rules: NoneOrSome::One(ConfigSelection::GroupName("allow-all-direct".to_string())), - }), + }), ]; let result = validate_configs_test(configs).await; @@ -2125,13 +2408,14 @@ mod tests { )] .into_iter() .collect(), + reality_targets: HashMap::new(), tls_buffer_size: None, }, transport: Transport::Tcp, tcp_settings: None, quic_settings: None, rules: NoneOrSome::One(ConfigSelection::GroupName("allow-all-direct".to_string())), - }), + }), ]; let result = validate_configs_test(configs).await; @@ -2218,6 +2502,7 @@ mod tests { ]), client_fingerprints: NoneOrSome::Unspecified, alpn_protocols: NoneOrSome::Unspecified, + vision: false, protocol: ServerProxyConfig::Http { username: None, password: None, @@ -2229,6 +2514,7 @@ mod tests { .collect(), default_tls_target: None, shadowtls_targets: HashMap::new(), + reality_targets: HashMap::new(), tls_buffer_size: None, }, transport: Transport::Tcp, diff --git a/src/copy_session_messages.rs b/src/copy_session_messages.rs new file mode 100644 index 00000000..32313d31 --- /dev/null +++ b/src/copy_session_messages.rs @@ -0,0 +1,287 @@ +use futures::ready; +use tokio::io::ReadBuf; + +use std::future::Future; +use std::io; +use std::net::SocketAddr; +use std::pin::Pin; +use std::task::{Context, Poll}; +use std::time::Instant; + +use crate::async_stream::AsyncSessionMessageStream; + +// Informed by https://stackoverflow.com/questions/14856639/udp-hole-punching-timeout +pub const DEFAULT_ASSOCIATION_TIMEOUT_SECS: u32 = 200; + +/// Buffer for copying messages from one session stream to another +#[derive(Debug)] +struct CopySessionToSessionBuffer { + read_done: bool, + need_flush: bool, + need_write_ping: bool, + cache_length: usize, + buf: [u8; 65535], + session_id: u16, + source_addr: SocketAddr, + read_count: usize, + write_count: usize, +} + +impl CopySessionToSessionBuffer { + pub fn new(need_flush: bool) -> Self { + Self { + read_done: false, + need_flush, + need_write_ping: false, + cache_length: 0, + buf: [0u8; 65535], + session_id: 0, + source_addr: SocketAddr::from(([0, 0, 0, 0], 0)), + read_count: 0, + write_count: 0, + } + } + + pub fn poll_copy_session_to_session( + &mut self, + cx: &mut Context<'_>, + mut reader: Pin<&mut R>, + mut writer: Pin<&mut W>, + ) -> Poll> + where + R: AsyncSessionMessageStream + ?Sized, + W: AsyncSessionMessageStream + ?Sized, + { + loop { + let mut did_read = false; + let mut did_write = false; + let mut read_pending = false; + let mut write_pending = false; + + // Read from session stream + if !self.read_done && self.cache_length == 0 { + let me = &mut *self; + let mut buf = ReadBuf::new(&mut me.buf); + match reader.as_mut().poll_read_session_message(cx, &mut buf) { + Poll::Ready(val) => { + let (session_id, source_addr) = val?; + let n = buf.filled().len(); + if n == 0 { + self.read_done = true; + } else { + self.session_id = session_id; + self.source_addr = source_addr; + self.cache_length = n; + did_read = true; + self.read_count = self.read_count.wrapping_add(n); + } + } + Poll::Pending => { + read_pending = true; + } + } + } + + // Write to session stream + if self.cache_length > 0 { + let me = &mut *self; + match writer.as_mut().poll_write_session_message( + cx, + me.session_id, + &me.buf[0..me.cache_length], + &me.source_addr, + ) { + Poll::Ready(val) => { + val?; + self.write_count = self.write_count.wrapping_add(1); + self.cache_length = 0; + self.need_flush = true; + self.need_write_ping = false; + did_write = true; + } + Poll::Pending => { + write_pending = true; + } + } + } + + if !write_pending && self.need_write_ping { + match writer.as_mut().poll_write_ping(cx) { + Poll::Ready(val) => { + let written = val?; + self.need_write_ping = false; + if written { + self.need_flush = true; + } + } + Poll::Pending => { + write_pending = true; + } + } + } + + if did_read && did_write && !read_pending && !write_pending { + continue; + } + + if self.need_flush { + ready!(writer.as_mut().poll_flush_message(cx))?; + self.need_flush = false; + continue; + } + + if self.read_done && self.cache_length == 0 { + return Poll::Ready(Ok(())); + } + + if read_pending || write_pending { + return Poll::Pending; + } + } + } +} + +enum TransferState { + Running, + ShuttingDown, + Done, +} + +struct CopySessionMessages<'a, A: ?Sized, B: ?Sized> { + a: &'a mut A, + b: &'a mut B, + a_buf: CopySessionToSessionBuffer, + b_buf: CopySessionToSessionBuffer, + a_to_b: TransferState, + b_to_a: TransferState, + sleep_future: Pin>, + a_last_active: Instant, + b_last_active: Instant, +} + +fn transfer_session_to_session( + cx: &mut Context<'_>, + state: &mut TransferState, + buf: &mut CopySessionToSessionBuffer, + r: &mut A, + w: &mut B, +) -> Poll> +where + A: AsyncSessionMessageStream + ?Sized, + B: AsyncSessionMessageStream + ?Sized, +{ + let mut r = Pin::new(r); + let mut w = Pin::new(w); + + loop { + match state { + TransferState::Running => { + ready!(buf.poll_copy_session_to_session(cx, r.as_mut(), w.as_mut()))?; + *state = TransferState::ShuttingDown; + } + TransferState::ShuttingDown => { + ready!(w.as_mut().poll_shutdown_message(cx))?; + *state = TransferState::Done; + } + TransferState::Done => return Poll::Ready(Ok(())), + } + } +} + +impl Future for CopySessionMessages<'_, A, B> +where + A: AsyncSessionMessageStream + ?Sized, + B: AsyncSessionMessageStream + ?Sized, +{ + type Output = io::Result<()>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let CopySessionMessages { + a, + b, + a_buf, + b_buf, + a_to_b, + b_to_a, + sleep_future, + a_last_active, + b_last_active, + } = &mut *self; + + let ping_fired = sleep_future.as_mut().poll(cx).is_ready(); + if ping_fired { + a_buf.need_write_ping = b.supports_ping(); + b_buf.need_write_ping = a.supports_ping(); + sleep_future + .as_mut() + .reset(tokio::time::Instant::now() + std::time::Duration::from_secs(60)); + } + + let a_read_count = a_buf.read_count; + let a_write_count = a_buf.write_count; + let b_read_count = b_buf.read_count; + let b_write_count = b_buf.write_count; + + let a_to_b = transfer_session_to_session(cx, a_to_b, &mut *a_buf, &mut *a, &mut *b); + let b_to_a = transfer_session_to_session(cx, b_to_a, &mut *b_buf, &mut *b, &mut *a); + + if a_buf.read_count != a_read_count || a_buf.write_count != a_write_count { + *a_last_active = Instant::now(); + } else if a_last_active.elapsed().as_secs() >= DEFAULT_ASSOCIATION_TIMEOUT_SECS.into() { + return Poll::Ready(Ok(())); + } + + if b_buf.read_count != b_read_count || b_buf.write_count != b_write_count { + *b_last_active = Instant::now(); + } else if b_last_active.elapsed().as_secs() >= DEFAULT_ASSOCIATION_TIMEOUT_SECS.into() { + return Poll::Ready(Ok(())); + } + + if a_to_b.is_ready() { + return a_to_b; + } else if b_to_a.is_ready() { + return b_to_a; + } + + Poll::Pending + } +} + +/// Copies data in both directions between two session-based streams. +/// +/// This is used for session-based protocols like XUDP where: +/// - `a` is the protocol stream (e.g., VlessXudpMessageStream) that reads/writes XUDP frames with session IDs +/// - `b` is the UDP manager (e.g., SessionUdpManager) that manages dedicated UDP sockets per session +/// +/// The function returns a future that will: +/// - Read from `a` (XUDP frames with session IDs) and write to `b` (create/use UDP sockets for each session) +/// - Read from `b` (UDP socket responses with session IDs) and write to `a` (encode as XUDP frames) +/// +/// This happens in both directions concurrently with timeout handling. +pub async fn copy_session_messages( + a: &mut A, + b: &mut B, + a_initial_flush: bool, + b_initial_flush: bool, +) -> Result<(), std::io::Error> +where + A: AsyncSessionMessageStream + ?Sized, + B: AsyncSessionMessageStream + ?Sized, +{ + let sleep_future = Box::pin(tokio::time::sleep(std::time::Duration::from_secs(60))); + + CopySessionMessages { + a, + b, + // a_buf copies from a (XUDP reader) to b (UDP manager writer), so need_flush is for b + a_buf: CopySessionToSessionBuffer::new(b_initial_flush), + // b_buf copies from b (UDP manager reader) to a (XUDP writer), so need_flush is for a + b_buf: CopySessionToSessionBuffer::new(a_initial_flush), + a_to_b: TransferState::Running, + b_to_a: TransferState::Running, + sleep_future, + a_last_active: Instant::now(), + b_last_active: Instant::now(), + } + .await +} diff --git a/src/crypto/crypto_connection.rs b/src/crypto/crypto_connection.rs new file mode 100644 index 00000000..9df103bc --- /dev/null +++ b/src/crypto/crypto_connection.rs @@ -0,0 +1,296 @@ +// Unified connection enum for all cryptographic protocols +// +// This enum mirrors rustls::Connection, allowing different TLS-like +// protocols to be used interchangeably. + +use std::io::{self, Read, Write}; + +use crate::reality::{ + feed_reality_client_connection, feed_reality_server_connection, RealityClientConnection, + RealityServerConnection, +}; + +use crate::rustls_connection_util::{feed_rustls_client_connection, feed_rustls_server_connection}; + +use super::crypto_reader_writer::{CryptoReader, CryptoWriter}; + +/// Represents the I/O state after processing packets +#[derive(Debug, Clone, Copy)] +pub struct CryptoIoState { + /// Number of plaintext bytes available to read + plaintext_bytes_to_read: usize, +} + +impl CryptoIoState { + /// Create a new IoState + pub fn new(plaintext_bytes_to_read: usize) -> Self { + CryptoIoState { + plaintext_bytes_to_read, + } + } + + /// How many plaintext bytes could be obtained via Read without further I/O + pub fn plaintext_bytes_to_read(&self) -> usize { + self.plaintext_bytes_to_read + } +} + +/// Unified connection type supporting multiple crypto protocols +/// +/// This enum allows switching between rustls and REALITY implementations +/// while maintaining a consistent API that matches rustls. +pub enum CryptoConnection { + /// rustls server-side connection + RustlsServer(rustls::ServerConnection), + /// rustls client-side connection + RustlsClient(rustls::ClientConnection), + /// REALITY server-side connection + RealityServer(RealityServerConnection), + /// REALITY client-side connection + RealityClient(RealityClientConnection), +} + +impl CryptoConnection { + /// Create a new rustls server connection + pub fn new_rustls_server(conn: rustls::ServerConnection) -> Self { + CryptoConnection::RustlsServer(conn) + } + + /// Create a new rustls client connection + pub fn new_rustls_client(conn: rustls::ClientConnection) -> Self { + CryptoConnection::RustlsClient(conn) + } + + /// Create a new REALITY server connection + pub fn new_reality_server(conn: RealityServerConnection) -> Self { + CryptoConnection::RealityServer(conn) + } + + /// Create a new REALITY client connection + pub fn new_reality_client(conn: RealityClientConnection) -> Self { + CryptoConnection::RealityClient(conn) + } + + /// Check if this is a server-side connection + pub fn is_server(&self) -> bool { + matches!( + self, + CryptoConnection::RustlsServer(_) | CryptoConnection::RealityServer(_) + ) + } + + /// Check if this is a client-side connection + pub fn is_client(&self) -> bool { + matches!( + self, + CryptoConnection::RustlsClient(_) | CryptoConnection::RealityClient(_) + ) + } + + /// Read TLS messages from `rd` into internal buffers + /// + /// Returns the number of bytes read, or 0 if the connection is closed. + /// This does NOT decrypt data - call `process_new_packets()` for that. + pub fn read_tls(&mut self, rd: &mut dyn Read) -> io::Result { + match self { + CryptoConnection::RustlsServer(conn) => conn.read_tls(rd), + CryptoConnection::RustlsClient(conn) => conn.read_tls(rd), + CryptoConnection::RealityServer(conn) => conn.read_tls(rd), + CryptoConnection::RealityClient(conn) => conn.read_tls(rd), + } + } + + /// Process any buffered TLS messages and update internal state + /// + /// This decrypts data and advances the handshake state machine. + /// Returns the I/O state including how many plaintext bytes are available. + pub fn process_new_packets(&mut self) -> io::Result { + match self { + CryptoConnection::RustlsServer(conn) => { + let io_state = conn.process_new_packets().map_err(|e| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("rustls server error processing new packets: {:?}", e), + ) + })?; + + Ok(CryptoIoState::new(io_state.plaintext_bytes_to_read())) + } + CryptoConnection::RustlsClient(conn) => { + let io_state = conn.process_new_packets().map_err(|e| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("rustls client error processing new packets: {:?}", e), + ) + })?; + + Ok(CryptoIoState::new(io_state.plaintext_bytes_to_read())) + } + CryptoConnection::RealityServer(conn) => { + let io_state = conn.process_new_packets().map_err(|e| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("reality server error processing new packets: {:?}", e), + ) + })?; + + Ok(CryptoIoState::new(io_state.plaintext_bytes_to_read())) + } + CryptoConnection::RealityClient(conn) => { + let io_state = conn.process_new_packets().map_err(|e| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("reality client error processing new packets: {:?}", e), + ) + })?; + + Ok(CryptoIoState::new(io_state.plaintext_bytes_to_read())) + } + } + } + + /// Get a unified reader for reading decrypted plaintext + /// + /// Works for both Rustls and REALITY connections. + /// Returns a CryptoReader enum that abstracts over rustls::Reader and REALITY's reader. + pub fn reader(&mut self) -> CryptoReader<'_> { + match self { + CryptoConnection::RustlsServer(conn) => CryptoReader::Rustls(conn.reader()), + CryptoConnection::RustlsClient(conn) => CryptoReader::Rustls(conn.reader()), + CryptoConnection::RealityServer(conn) => CryptoReader::Reality(conn.reader()), + CryptoConnection::RealityClient(conn) => CryptoReader::Reality(conn.reader()), + } + } + + /// Get a unified writer for writing plaintext to be encrypted + /// + /// Works for both Rustls and REALITY connections. + /// Returns a CryptoWriter enum that abstracts over rustls::Writer and REALITY's writer. + pub fn writer(&mut self) -> CryptoWriter<'_> { + match self { + CryptoConnection::RustlsServer(conn) => CryptoWriter::Rustls(conn.writer()), + CryptoConnection::RustlsClient(conn) => CryptoWriter::Rustls(conn.writer()), + CryptoConnection::RealityServer(conn) => CryptoWriter::Reality(conn.writer()), + CryptoConnection::RealityClient(conn) => CryptoWriter::Reality(conn.writer()), + } + } + + /// Write any buffered TLS messages to `wr` + /// + /// This encrypts any pending plaintext and writes the ciphertext. + /// Returns the number of bytes written. + pub fn write_tls(&mut self, wr: &mut dyn Write) -> io::Result { + match self { + CryptoConnection::RustlsServer(conn) => conn.write_tls(wr), + CryptoConnection::RustlsClient(conn) => conn.write_tls(wr), + CryptoConnection::RealityServer(conn) => conn.write_tls(wr), + CryptoConnection::RealityClient(conn) => conn.write_tls(wr), + } + } + + /// Check if the connection wants to write data + /// + /// If true, the application should call `write_tls()` as soon as possible. + pub fn wants_write(&self) -> bool { + match self { + CryptoConnection::RustlsServer(conn) => conn.wants_write(), + CryptoConnection::RustlsClient(conn) => conn.wants_write(), + CryptoConnection::RealityServer(conn) => conn.wants_write(), + CryptoConnection::RealityClient(conn) => conn.wants_write(), + } + } + + /// Check whether the handshake is complete + /// + /// For both client and server connections, this returns true when + /// the handshake has completed and application data can be exchanged. + pub fn is_handshaking(&self) -> bool { + match self { + CryptoConnection::RustlsServer(conn) => conn.is_handshaking(), + CryptoConnection::RustlsClient(conn) => conn.is_handshaking(), + CryptoConnection::RealityServer(conn) => conn.is_handshaking(), + CryptoConnection::RealityClient(conn) => conn.is_handshaking(), + } + } + + /// Queue a close notification + /// + /// This queues a close notification to be sent to the peer. + /// Call `write_tls()` to actually send it. + pub fn send_close_notify(&mut self) { + match self { + CryptoConnection::RustlsServer(conn) => conn.send_close_notify(), + CryptoConnection::RustlsClient(conn) => conn.send_close_notify(), + CryptoConnection::RealityServer(conn) => conn.send_close_notify(), + CryptoConnection::RealityClient(conn) => conn.send_close_notify(), + } + } +} + +#[inline(always)] +pub fn feed_crypto_connection( + connection: &mut CryptoConnection, + data: &[u8], +) -> std::io::Result<()> { + match connection { + CryptoConnection::RustlsServer(conn) => feed_rustls_server_connection(conn, data), + CryptoConnection::RustlsClient(conn) => feed_rustls_client_connection(conn, data), + CryptoConnection::RealityServer(conn) => feed_reality_server_connection(conn, data), + CryptoConnection::RealityClient(conn) => feed_reality_client_connection(conn, data), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_connection_type_queries() { + use crate::address::{Address, NetLocation}; + use crate::reality::{ + RealityClientConfig, RealityClientConnection, RealityServerConfig, + RealityServerConnection, + }; + + // Test with REALITY server connection + let server_config = RealityServerConfig { + private_key: [0u8; 32], + short_ids: vec![[0u8; 8]], + dest: NetLocation::new(Address::UNSPECIFIED, 443), + max_time_diff: None, + min_client_version: None, + max_client_version: None, + }; + let reality_server = RealityServerConnection::new(server_config).unwrap(); + let server_conn = CryptoConnection::new_reality_server(reality_server); + + assert!(server_conn.is_server()); + assert!(!server_conn.is_client()); + assert!(matches!(server_conn, CryptoConnection::RealityServer(_))); + assert!(!matches!( + server_conn, + CryptoConnection::RustlsServer(_) | CryptoConnection::RustlsClient(_) + )); + + // Test with REALITY client connection + // Using a valid X25519 test public key (derived from private key [1u8; 32]) + let client_config = RealityClientConfig { + public_key: [ + 0x25, 0x31, 0x51, 0x68, 0x12, 0x8f, 0xf3, 0x7b, 0x46, 0x5e, 0x0c, 0x0c, 0xd8, 0x28, + 0x3d, 0xd5, 0x35, 0x86, 0x8d, 0x8d, 0x3e, 0x8c, 0x1c, 0x85, 0x1e, 0x0e, 0x86, 0x28, + 0xcf, 0x48, 0x0e, 0x66, + ], + short_id: [0u8; 8], + server_name: "test.example.com".to_string(), + }; + let reality_client = RealityClientConnection::new(client_config).unwrap(); + let client_conn = CryptoConnection::new_reality_client(reality_client); + assert!(!client_conn.is_server()); + assert!(client_conn.is_client()); + assert!(matches!(client_conn, CryptoConnection::RealityClient(_))); + assert!(!matches!( + client_conn, + CryptoConnection::RustlsServer(_) | CryptoConnection::RustlsClient(_) + )); + } +} diff --git a/src/crypto/crypto_reader_writer.rs b/src/crypto/crypto_reader_writer.rs new file mode 100644 index 00000000..c69c97f4 --- /dev/null +++ b/src/crypto/crypto_reader_writer.rs @@ -0,0 +1,61 @@ +// Reader and Writer types for crypto connections +// +// These types provide a consistent API for reading decrypted plaintext +// and writing plaintext to be encrypted, working with both rustls and REALITY. + +use std::io::{self, BufRead, Read, Write}; + +use crate::reality::{RealityReader, RealityWriter}; + +/// Unified reader that works with both Rustls and REALITY connections +pub enum CryptoReader<'a> { + Rustls(rustls::Reader<'a>), + Reality(RealityReader<'a>), +} + +impl<'a> Read for CryptoReader<'a> { + fn read(&mut self, buf: &mut [u8]) -> io::Result { + match self { + CryptoReader::Rustls(reader) => reader.read(buf), + CryptoReader::Reality(reader) => reader.read(buf), + } + } +} + +impl<'a> BufRead for CryptoReader<'a> { + fn fill_buf(&mut self) -> io::Result<&[u8]> { + match self { + CryptoReader::Rustls(reader) => reader.fill_buf(), + CryptoReader::Reality(reader) => reader.fill_buf(), + } + } + + fn consume(&mut self, amt: usize) { + match self { + CryptoReader::Rustls(reader) => reader.consume(amt), + CryptoReader::Reality(reader) => reader.consume(amt), + } + } +} + +/// Unified writer that works with both Rustls and REALITY connections +pub enum CryptoWriter<'a> { + Rustls(rustls::Writer<'a>), + Reality(RealityWriter<'a>), +} + +impl<'a> Write for CryptoWriter<'a> { + fn write(&mut self, buf: &[u8]) -> io::Result { + match self { + CryptoWriter::Rustls(writer) => writer.write(buf), + CryptoWriter::Reality(writer) => writer.write(buf), + } + } + + fn flush(&mut self) -> io::Result<()> { + match self { + CryptoWriter::Rustls(writer) => writer.flush(), + CryptoWriter::Reality(writer) => writer.flush(), + } + } +} diff --git a/src/crypto/crypto_tls_stream.rs b/src/crypto/crypto_tls_stream.rs new file mode 100644 index 00000000..c77f7753 --- /dev/null +++ b/src/crypto/crypto_tls_stream.rs @@ -0,0 +1,344 @@ +// CryptoTlsStream: Async wrapper around Connection enum +// +// This provides AsyncRead + AsyncWrite for the Connection enum, +// allowing it to work with tokio-based async code. +// +// Based on VisionStream's TLS mode logic (non-padded), which itself +// is adapted from tokio-rustls. + +use std::io::{self, BufRead, Write}; +use std::pin::Pin; +use std::task::{Context, Poll}; + +use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; + +use super::crypto_connection::CryptoConnection; +use crate::sync_adapter::{SyncReadAdapter, SyncWriteAdapter}; + +/// Async stream wrapper for Connection enum +/// +/// This combines a Connection (rustls or REALITY) with an async I/O stream, +/// providing AsyncRead + AsyncWrite interface. +pub struct CryptoTlsStream { + /// The underlying async I/O stream (e.g., TcpStream) + io: IO, + /// The crypto connection (rustls or REALITY) + session: CryptoConnection, + /// Whether we've hit EOF when reading + is_read_eof: bool, +} + +impl CryptoTlsStream +where + IO: AsyncRead + AsyncWrite + Unpin, +{ + /// Create a new CryptoTlsStream + pub fn new(io: IO, session: CryptoConnection) -> Self { + CryptoTlsStream { + io, + session, + is_read_eof: false, + } + } + + /// Extract the underlying I/O stream and crypto session + /// + /// This consumes the CryptoTlsStream and returns its components. + /// Useful for scenarios like Vision that need direct access to both. + pub fn into_inner(self) -> (IO, CryptoConnection) { + (self.io, self.session) + } + + /// Write TLS data directly to the underlying stream + /// + /// Returns Poll::Ready(Ok(n)) with bytes written, or Poll::Pending if would block. + fn write_tls_direct(&mut self, cx: &mut Context<'_>) -> Poll> { + let mut adapter = SyncWriteAdapter { + io: &mut self.io, + cx, + }; + match self.session.write_tls(&mut adapter) { + Ok(n) => Poll::Ready(Ok(n)), + Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => Poll::Pending, + Err(e) => Poll::Ready(Err(e)), + } + } + + /// Drain all pending TLS writes to the underlying stream + fn drain_all_writes(&mut self, cx: &mut Context<'_>) -> Poll> { + while self.session.wants_write() { + match self.write_tls_direct(cx) { + Poll::Ready(Ok(_)) => {} + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + Poll::Pending => return Poll::Pending, + } + } + Poll::Ready(Ok(())) + } + + /// Complete handshake if needed + fn complete_handshake_if_needed(&mut self, cx: &mut Context<'_>) -> Poll> { + while self.session.is_handshaking() { + // Write any pending TLS data + match self.drain_all_writes(cx) { + Poll::Ready(Ok(())) => {} + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + Poll::Pending => return Poll::Pending, + } + + // Read more data for handshake + let mut adapter = SyncReadAdapter { + io: &mut self.io, + cx, + }; + match self.session.read_tls(&mut adapter) { + Ok(0) => { + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "EOF during handshake", + ))); + } + Ok(_) => { + // Process the data + self.session.process_new_packets()?; + } + Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + return Poll::Pending; + } + Err(e) => return Poll::Ready(Err(e)), + } + } + + Poll::Ready(Ok(())) + } +} + +impl AsyncRead for CryptoTlsStream +where + IO: AsyncRead + AsyncWrite + Unpin, +{ + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + let this = self.get_mut(); + + // Check if session already has decrypted data available + { + let mut reader = this.session.reader(); + match reader.fill_buf() { + Ok(available) if !available.is_empty() => { + // Copy directly from session buffer to user buffer + let len = buf.remaining().min(available.len()); + buf.put_slice(&available[..len]); + reader.consume(len); + return Poll::Ready(Ok(())); + } + Ok(_) => { + // Empty buffer, need to read more + } + Err(e) if e.kind() == io::ErrorKind::WouldBlock => { + // No data available yet + } + Err(e) => return Poll::Ready(Err(e)), + } + } + + if this.is_read_eof { + return Poll::Ready(Ok(())); + } + + // Complete handshake first if needed + if this.session.is_handshaking() { + match this.complete_handshake_if_needed(cx) { + Poll::Ready(Ok(())) => {} + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + Poll::Pending => return Poll::Pending, + } + } + + // No data in session buffer, read from TCP and decrypt + // Split borrows manually to avoid borrow checker issues + + loop { + let bytes_read = { + let mut adapter = SyncReadAdapter { + io: &mut this.io, + cx, + }; + match this.session.read_tls(&mut adapter) { + Ok(n) => n, + Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + return Poll::Pending; + } + Err(e) => return Poll::Ready(Err(e)), + } + }; + + if bytes_read == 0 { + // EOF + this.is_read_eof = true; + return Poll::Ready(Ok(())); + } + + // Process encrypted data - try to send alerts on error + let io_state = match this.session.process_new_packets() { + Ok(state) => state, + Err(e) => { + // Try last-gasp write to send any pending TLS alerts (tokio-rustls pattern) + // Make best-effort attempt to drain all pending writes without blocking. + // This increases the chance that the peer receives an alert explaining + // the protocol error. We ignore write failures to ensure the primary + // TLS protocol error (e) is always returned. + while this.session.wants_write() { + let mut adapter = SyncWriteAdapter { + io: &mut this.io, + cx, + }; + match this.session.write_tls(&mut adapter) { + Ok(_) => {} // Continue draining + Err(ref write_err) if write_err.kind() == io::ErrorKind::WouldBlock => { + break + } + Err(_) => break, // Give up on write error + } + } + return Poll::Ready(Err(e)); + } + }; + + if io_state.plaintext_bytes_to_read() == 0 { + // No plaintext yet, need more data + continue; + } + + // Extract plaintext from session + let mut reader = this.session.reader(); + match reader.fill_buf() { + Ok(available) => { + if available.is_empty() { + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "Read zero bytes when plaintext is available", + ))); + } + let len = buf.remaining().min(available.len()); + buf.put_slice(&available[..len]); + reader.consume(len); + return Poll::Ready(Ok(())); + } + Err(e) => { + return Poll::Ready(Err(e)); + } + } + } + } +} + +impl AsyncWrite for CryptoTlsStream +where + IO: AsyncRead + AsyncWrite + Unpin, +{ + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + // Complete handshake first if needed + if self.session.is_handshaking() { + match self.complete_handshake_if_needed(cx) { + Poll::Ready(Ok(())) => {} + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + Poll::Pending => return Poll::Pending, + } + } + + let mut pos = 0; + + while pos < buf.len() { + let mut would_block = false; + + // Write plaintext to session + match self.session.writer().write(&buf[pos..]) { + Ok(n) => pos += n, + Err(e) => return Poll::Ready(Err(e)), + }; + + // Drain TLS output to TCP stream + while self.session.wants_write() { + match self.write_tls_direct(cx) { + Poll::Ready(Ok(0)) | Poll::Pending => { + would_block = true; + break; + } + Poll::Ready(Ok(_)) => (), + Poll::Ready(Err(err)) => return Poll::Ready(Err(err)), + } + } + + return match (pos, would_block) { + (0, true) => Poll::Pending, + (n, true) => Poll::Ready(Ok(n)), // Partial write + (_, false) => continue, // Keep writing + }; + } + + Poll::Ready(Ok(pos)) + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + // Drain all pending TLS writes + match self.drain_all_writes(cx) { + Poll::Ready(Ok(())) => {} + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + Poll::Pending => return Poll::Pending, + } + + // Flush underlying stream + Pin::new(&mut self.io).poll_flush(cx) + } + + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + // Send close notification + self.session.send_close_notify(); + + // Drain all pending writes + match self.drain_all_writes(cx) { + Poll::Ready(Ok(())) => {} + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + Poll::Pending => return Poll::Pending, + } + + // Shutdown underlying stream - ignore NotConnected errors (tokio-rustls pattern) + match Pin::new(&mut self.io).poll_shutdown(cx) { + Poll::Ready(Ok(())) => Poll::Ready(Ok(())), + Poll::Ready(Err(e)) if e.kind() == io::ErrorKind::NotConnected => { + // When trying to shutdown, not being connected is fine + Poll::Ready(Ok(())) + } + Poll::Ready(Err(e)) => Poll::Ready(Err(e)), + Poll::Pending => Poll::Pending, + } + } +} + +// Implement AsyncPing for CryptoTlsStream +impl crate::async_stream::AsyncPing for CryptoTlsStream +where + IO: AsyncRead + AsyncWrite + Unpin, +{ + fn supports_ping(&self) -> bool { + false // TLS doesn't have a built-in ping mechanism + } + + fn poll_write_ping(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(false)) + } +} + +// Implement AsyncStream blanket trait +impl crate::async_stream::AsyncStream for CryptoTlsStream where + IO: AsyncRead + AsyncWrite + Unpin + Send +{ +} diff --git a/src/crypto/mod.rs b/src/crypto/mod.rs new file mode 100644 index 00000000..0c2508ae --- /dev/null +++ b/src/crypto/mod.rs @@ -0,0 +1,13 @@ +// Unified cryptographic connection abstraction +// +// This module provides a common interface for different TLS-like protocols, +// allowing rustls and REALITY to be used interchangeably throughout the codebase. + +mod crypto_connection; +mod crypto_reader_writer; +mod crypto_tls_stream; + +// Re-export core types +pub use crypto_connection::{feed_crypto_connection, CryptoConnection}; +pub use crypto_tls_stream::CryptoTlsStream; +// CryptoReader, CryptoWriter, and IoState are used internally within this module diff --git a/src/main.rs b/src/main.rs index f3868c55..d8f70d02 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,6 +6,8 @@ mod config; mod copy_bidirectional; mod copy_bidirectional_message; mod copy_multidirectional_message; +mod copy_session_messages; +mod crypto; mod http_handler; mod hysteria2_server; mod noop_stream; @@ -13,28 +15,35 @@ mod option_util; mod port_forward_handler; mod quic_server; mod quic_stream; +mod reality; +mod reality_client_handler; mod resolver; -mod rustls_util; -mod salt_checker; +mod rustls_config_util; +mod rustls_connection_util; +mod rustls_handshake; mod shadow_tls; mod shadowsocks; +mod slide_buffer; mod snell; mod socket_util; mod socks_handler; mod stream_reader; +mod sync_adapter; mod tcp; mod thread_util; -mod timed_salt_checker; -mod tls_handler; +mod tls_client_handler; +mod tls_server_handler; mod trojan_handler; mod tuic_server; mod udp_message_stream; mod udp_multi_message_stream; +mod udp_session_message_stream; +mod uot; mod util; -mod vless_handler; -mod vless_message_stream; +mod vless; mod vmess; mod websocket; +mod xudp; #[cfg(not(target_env = "msvc"))] use tikv_jemallocator::Jemalloc; @@ -55,6 +64,8 @@ use tokio::task::JoinHandle; use crate::config::{ServerConfig, Transport}; use crate::quic_server::start_quic_servers; +use crate::reality::generate_keypair; +use crate::shadowsocks::ShadowsocksCipher; use crate::thread_util::set_num_threads; use tcp::*; @@ -191,6 +202,74 @@ fn main() { } } + if args.iter().any(|s| s == "generate-reality-keypair") { + let (private_key, public_key) = generate_keypair().unwrap(); + println!( + "--------------------------------------------------------------------------------" + ); + println!("REALITY private key: {}", private_key); + println!("REALITY public key: {}", public_key); + println!( + "--------------------------------------------------------------------------------" + ); + return; + } + + if let Some(pos) = args + .iter() + .position(|s| s == "generate-shadowsocks-password") + { + let cipher = args.get(pos + 1).map(|s| s.as_str()); + match cipher { + Some(c) => { + // Strip 2022-blake3- prefix if present for cipher lookup + let base_cipher = c.strip_prefix("2022-blake3-").unwrap_or(c); + match ShadowsocksCipher::try_from(base_cipher) { + Ok(cipher) => { + use aws_lc_rs::rand::{SecureRandom, SystemRandom}; + use base64::engine::{general_purpose::STANDARD, Engine as _}; + + let rng = SystemRandom::new(); + let mut key_bytes = vec![0u8; cipher.key_len()]; + rng.fill(&mut key_bytes).expect("RNG failed"); + let password = STANDARD.encode(&key_bytes); + println!( + "--------------------------------------------------------------------------------" + ); + println!("Cipher: {}", c); + println!("Password: {}", password); + println!( + "--------------------------------------------------------------------------------" + ); + } + Err(_) => { + eprintln!("Unknown cipher: {}", c); + eprintln!("Supported ciphers:"); + eprintln!(" aes-128-gcm"); + eprintln!(" aes-256-gcm"); + eprintln!(" chacha20-ietf-poly1305"); + eprintln!(" 2022-blake3-aes-128-gcm"); + eprintln!(" 2022-blake3-aes-256-gcm"); + eprintln!(" 2022-blake3-chacha20-poly1305"); + std::process::exit(1); + } + } + } + None => { + eprintln!("Usage: {} generate-shadowsocks-password ", arg0); + eprintln!("Supported ciphers:"); + eprintln!(" aes-128-gcm"); + eprintln!(" aes-256-gcm"); + eprintln!(" chacha20-ietf-poly1305"); + eprintln!(" 2022-blake3-aes-128-gcm"); + eprintln!(" 2022-blake3-aes-256-gcm"); + eprintln!(" 2022-blake3-chacha20-poly1305"); + std::process::exit(1); + } + } + return; + } + if args.is_empty() { println!("No config specified, assuming loading from file config.shoes.yaml"); args.push("config.shoes.yaml".to_string()) @@ -276,8 +355,6 @@ fn main() { return; } - println!("\nStarting {} server(s)..", configs.len()); - let mut join_handles = vec![]; let server_configs = match config::create_server_configs(configs).await { @@ -288,6 +365,9 @@ fn main() { return; } }; + + println!("\nStarting {} server(s)..", server_configs.len()); + for server_config in server_configs { join_handles.extend(start_servers(server_config).await.unwrap()); } diff --git a/src/port_forward_handler.rs b/src/port_forward_handler.rs index 324f9ef4..9e3bf00b 100644 --- a/src/port_forward_handler.rs +++ b/src/port_forward_handler.rs @@ -5,6 +5,7 @@ use async_trait::async_trait; use crate::address::NetLocation; use crate::async_stream::AsyncStream; use crate::option_util::NoneOrOne; +use crate::tcp_handler::{TcpClientHandler, TcpClientSetupResult}; use crate::tcp_handler::{TcpServerHandler, TcpServerSetupResult}; #[derive(Debug)] @@ -38,10 +39,25 @@ impl TcpServerHandler for PortForwardServerHandler { Ok(TcpServerSetupResult::TcpForward { remote_location: location.clone(), stream: server_stream, - need_initial_flush: false, + need_initial_flush: true, connection_success_response: None, initial_remote_data: None, override_proxy_provider: NoneOrOne::Unspecified, }) } } + +#[derive(Debug)] +pub struct PortForwardClientHandler; + +#[async_trait] +impl TcpClientHandler for PortForwardClientHandler { + async fn setup_client_stream( + &self, + _server_stream: &mut Box, + client_stream: Box, + _remote_location: NetLocation, + ) -> std::io::Result { + Ok(TcpClientSetupResult { client_stream }) + } +} diff --git a/src/quic_server.rs b/src/quic_server.rs index 299aaf95..1eab036e 100644 --- a/src/quic_server.rs +++ b/src/quic_server.rs @@ -16,9 +16,10 @@ use crate::config::{ use crate::copy_bidirectional::copy_bidirectional; use crate::copy_bidirectional_message::copy_bidirectional_message; use crate::copy_multidirectional_message::copy_multidirectional_message; +use crate::copy_session_messages::copy_session_messages; use crate::quic_stream::QuicStream; use crate::resolver::{resolve_single_address, NativeResolver, Resolver}; -use crate::rustls_util::create_server_config; +use crate::rustls_config_util::create_server_config; use crate::socket_util::new_socket2_udp_socket; use crate::tcp_client_connector::TcpClientConnector; use crate::tcp_handler::{TcpServerHandler, TcpServerSetupResult}; @@ -26,6 +27,7 @@ use crate::tcp_handler_util::{create_tcp_client_proxy_selector, create_tcp_serve use crate::tcp_server::setup_client_stream; use crate::udp_message_stream::UdpMessageStream; use crate::udp_multi_message_stream::UdpMultiMessageStream; +use crate::udp_session_message_stream::{UdpSessionMessageStream, UdpSocketConfig}; use crate::util::parse_uuid; async fn start_quic_server( @@ -320,6 +322,45 @@ async fn process_streams( } } } + TcpServerSetupResult::SessionBasedUdp { + stream: mut server_stream, + need_initial_flush: server_need_initial_flush, + override_proxy_provider, + } => { + let selected_proxy_provider = if override_proxy_provider.is_one() { + override_proxy_provider.unwrap() + } else { + client_proxy_selector + }; + let action = selected_proxy_provider.default_decision(); + match action { + ConnectDecision::Allow { + client_proxy, + remote_location: _, + } => { + // Create UdpSessionMessageStream with user's interface configuration + let config = UdpSocketConfig { + bind_interface: client_proxy.bind_interface().clone(), + }; + let mut session_manager = UdpSessionMessageStream::new(config); + + let copy_result = copy_session_messages( + &mut server_stream, + &mut session_manager, + server_need_initial_flush, + false, + ) + .await; + + copy_result?; + Ok(()) + } + ConnectDecision::Block => { + warn!("Blocked session-based udp forward, because the default action is to block."); + Ok(()) + } + } + } } } diff --git a/src/reality/common.rs b/src/reality/common.rs new file mode 100644 index 00000000..e38fc990 --- /dev/null +++ b/src/reality/common.rs @@ -0,0 +1,114 @@ +// Common constants and helpers shared between REALITY client and server implementations +// +// This module provides: +// - TLS constants (content types, alert codes, version bytes, handshake types) +// - Close notify alert construction + +use super::reality_aead::encrypt_tls13_record; +use std::io; + +// TLS ContentType values +pub const CONTENT_TYPE_CHANGE_CIPHER_SPEC: u8 = 0x14; +pub const CONTENT_TYPE_ALERT: u8 = 0x15; +pub const CONTENT_TYPE_HANDSHAKE: u8 = 0x16; +pub const CONTENT_TYPE_APPLICATION_DATA: u8 = 0x17; + +// TLS alert levels and descriptions +pub const ALERT_LEVEL_WARNING: u8 = 0x01; +pub const ALERT_DESC_CLOSE_NOTIFY: u8 = 0x00; + +// TLS version bytes (used on wire for compatibility) +// TLS 1.2 version bytes: 0x03, 0x03 +// Used in TLS 1.3 for compatibility (appears in record layer) +pub const VERSION_TLS_1_2_MAJOR: u8 = 0x03; +pub const VERSION_TLS_1_2_MINOR: u8 = 0x03; + +// TLS 1.3 handshake message types +pub const HANDSHAKE_TYPE_SERVER_HELLO: u8 = 2; +pub const HANDSHAKE_TYPE_ENCRYPTED_EXTENSIONS: u8 = 8; +pub const HANDSHAKE_TYPE_CERTIFICATE: u8 = 11; +pub const HANDSHAKE_TYPE_CERTIFICATE_VERIFY: u8 = 15; +pub const HANDSHAKE_TYPE_FINISHED: u8 = 20; + +// TLS 1.3 record size limits per RFC 8446 +// +// The TLS record header's `length` field specifies the size of the ENCRYPTED payload. +// Per RFC 8446, the TLS 1.3 limit is stricter than TLS 1.2: +// +// - TLS 1.3: Plaintext limit = 16,384 bytes (2^14) +// Encryption overhead allowance = 256 bytes +// Ciphertext limit = 16,384 + 256 = 16,640 bytes +// +// - TLS 1.2: Plaintext limit = 16,384 bytes (2^14) +// Encryption overhead allowance = 2,048 bytes +// Ciphertext limit = 16,384 + 2,048 = 18,432 bytes +// +// REALITY uses TLS 1.3, so we MUST use the TLS 1.3 limit. Using the larger +// TLS 1.2 limit causes "record overflow" errors in libraries like utls. + +/// Maximum TLS 1.3 ciphertext payload size (16,640 bytes) +pub const MAX_TLS_CIPHERTEXT_LEN: usize = 16384 + 256; + +/// Maximum plaintext payload size for a single TLS 1.3 record +/// +/// RFC 8446 Section 5.1: "The record layer fragments information blocks into +/// TLSPlaintext records carrying data in chunks of 2^14 bytes or less." +/// +/// This is the hard limit enforced by TLS implementations. +/// The 256-byte allowance in MAX_TLS_CIPHERTEXT_LEN is for: +/// - AEAD tag (16 bytes for AES-GCM) +/// - Content type byte (1 byte) +/// - Optional padding (up to 239 bytes) +/// +/// We MUST NOT exceed 16384 bytes of actual plaintext per record, or clients +/// will reject with "record overflow" error. +pub const MAX_TLS_PLAINTEXT_LEN: usize = 16384; + +/// TLS record header size (ContentType + ProtocolVersion + Length) +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; + +/// Buffer capacity for ciphertext read (2x TLS max record for safety) +pub const CIPHERTEXT_READ_BUF_CAPACITY: usize = TLS_MAX_RECORD_SIZE * 2; + +/// Buffer capacity for plaintext read +pub const PLAINTEXT_READ_BUF_CAPACITY: usize = TLS_MAX_RECORD_SIZE * 2; + +/// Build an encrypted close_notify alert for TLS 1.3 +/// +/// In TLS 1.3, alerts must be encrypted like application data. +pub fn build_close_notify_alert(key: &[u8], iv: &[u8], seq_num: u64) -> io::Result> { + // Build alert message: level(1) + description(0) + ContentType + let alert_with_type = vec![ + ALERT_LEVEL_WARNING, + ALERT_DESC_CLOSE_NOTIFY, + CONTENT_TYPE_ALERT, // ContentType byte for TLS 1.3 + ]; + + // Build TLS header with correct ciphertext length + let ciphertext_len = (alert_with_type.len() + 16) as u16; // plaintext + tag + let mut tls_header = [ + CONTENT_TYPE_APPLICATION_DATA, + VERSION_TLS_1_2_MAJOR, + VERSION_TLS_1_2_MINOR, + 0x00, + 0x00, // Length will be set + ]; + tls_header[3..5].copy_from_slice(&ciphertext_len.to_be_bytes()); + + // Encrypt the alert + let ciphertext = encrypt_tls13_record(key, iv, seq_num, &alert_with_type, &tls_header)?; + + // Build complete TLS record + let mut record = Vec::with_capacity(5 + ciphertext.len()); + record.push(CONTENT_TYPE_APPLICATION_DATA); + record.push(VERSION_TLS_1_2_MAJOR); + record.push(VERSION_TLS_1_2_MINOR); + record.push(((ciphertext.len() >> 8) & 0xff) as u8); + record.push((ciphertext.len() & 0xff) as u8); + record.extend_from_slice(&ciphertext); + + Ok(record) +} diff --git a/src/reality/mod.rs b/src/reality/mod.rs new file mode 100644 index 00000000..abe4a746 --- /dev/null +++ b/src/reality/mod.rs @@ -0,0 +1,33 @@ +// REALITY Protocol Implementation +// +// This module implements the REALITY obfuscation protocol for TLS connections. +// REALITY disguises proxy traffic as legitimate HTTPS connections using: +// - X25519 ECDH key exchange +// - HKDF-SHA256 key derivation +// - AES-256-GCM encryption +// - HMAC-SHA512 authentication + +mod common; +mod reality_aead; +mod reality_auth; +mod reality_certificate; +mod reality_client_connection; +mod reality_client_verify; +mod reality_io_state; +mod reality_reader_writer; +mod reality_records; +mod reality_server_connection; +mod reality_tls13_keys; +mod reality_tls13_messages; +mod reality_util; + +pub use reality_util::{decode_private_key, decode_public_key, decode_short_id, generate_keypair}; + +// Re-export connection types for crypto_connection module +pub use reality_client_connection::{ + feed_reality_client_connection, RealityClientConfig, RealityClientConnection, +}; +pub use reality_reader_writer::{RealityReader, RealityWriter}; +pub use reality_server_connection::{ + feed_reality_server_connection, RealityServerConfig, RealityServerConnection, +}; diff --git a/src/reality/reality_aead.rs b/src/reality/reality_aead.rs new file mode 100644 index 00000000..312ef354 --- /dev/null +++ b/src/reality/reality_aead.rs @@ -0,0 +1,365 @@ +// TLS 1.3 Encryption/Decryption Helpers +// +// AES-GCM encryption for TLS 1.3 records using aws-lc-rs + +use super::common::{ + CONTENT_TYPE_ALERT, CONTENT_TYPE_APPLICATION_DATA, VERSION_TLS_1_2_MAJOR, VERSION_TLS_1_2_MINOR, +}; +use aws_lc_rs::aead::{Aad, LessSafeKey, Nonce, UnboundKey, AES_128_GCM}; +use std::io::{Error, ErrorKind, Result}; + +/// Encrypt TLS 1.3 record using AES-128-GCM +/// +/// # Arguments +/// * `key` - AES key (16 bytes for AES-128) +/// * `iv` - Base IV (12 bytes) +/// * `sequence_number` - TLS record sequence number +/// * `plaintext` - Plaintext data (including ContentType trailer) +/// * `additional_data` - TLS record header for AEAD +/// +/// # Returns +/// Ciphertext with authentication tag appended +pub fn encrypt_tls13_record( + key: &[u8], + iv: &[u8], + sequence_number: u64, + plaintext: &[u8], + additional_data: &[u8], +) -> Result> { + if key.len() != 16 { + return Err(Error::new( + ErrorKind::InvalidInput, + format!("Invalid key length: {} (expected 16)", key.len()), + )); + } + if iv.len() != 12 { + return Err(Error::new( + ErrorKind::InvalidInput, + format!("Invalid IV length: {} (expected 12)", iv.len()), + )); + } + + // Construct nonce: IV XOR sequence_number + let mut nonce_bytes = [0u8; 12]; + nonce_bytes.copy_from_slice(iv); + + // XOR the last 8 bytes with the sequence number + let seq_bytes = sequence_number.to_be_bytes(); + for i in 0..8 { + nonce_bytes[4 + i] ^= seq_bytes[i]; + } + + // Create key and nonce for aws-lc-rs + let unbound_key = UnboundKey::new(&AES_128_GCM, key) + .map_err(|e| Error::new(ErrorKind::InvalidInput, format!("Invalid key: {:?}", e)))?; + let sealing_key = LessSafeKey::new(unbound_key); + + let nonce = Nonce::try_assume_unique_for_key(&nonce_bytes) + .map_err(|e| Error::new(ErrorKind::InvalidInput, format!("Invalid nonce: {:?}", e)))?; + + let aad = Aad::from(additional_data); + + // Note: The caller is responsible for adding the ContentType byte to plaintext + // For handshake: plaintext = handshake_msg || 0x16 + // For app data: plaintext = app_data || 0x17 + let mut in_out = plaintext.to_vec(); + + sealing_key + .seal_in_place_append_tag(nonce, aad, &mut in_out) + .map_err(|e| { + Error::new( + ErrorKind::InvalidData, + format!("Encryption failed: {:?}", e), + ) + })?; + + Ok(in_out) +} + +/// Decrypt TLS 1.3 record using AES-128-GCM +/// +/// # Arguments +/// * `key` - AES key (16 bytes for AES-128) +/// * `iv` - Base IV (12 bytes) +/// * `sequence_number` - TLS record sequence number +/// * `ciphertext` - Ciphertext with authentication tag +/// * `additional_data` - TLS record header for AEAD +/// +/// # Returns +/// Plaintext data (including ContentType trailer) +pub fn decrypt_tls13_record( + key: &[u8], + iv: &[u8], + sequence_number: u64, + ciphertext: &[u8], + additional_data: &[u8], +) -> Result> { + if key.len() != 16 { + return Err(Error::new( + ErrorKind::InvalidInput, + format!("Invalid key length: {} (expected 16)", key.len()), + )); + } + if iv.len() != 12 { + return Err(Error::new( + ErrorKind::InvalidInput, + format!("Invalid IV length: {} (expected 12)", iv.len()), + )); + } + + // Construct nonce: IV XOR sequence_number + let mut nonce_bytes = [0u8; 12]; + nonce_bytes.copy_from_slice(iv); + + // XOR the last 8 bytes with the sequence number + let seq_bytes = sequence_number.to_be_bytes(); + for i in 0..8 { + nonce_bytes[4 + i] ^= seq_bytes[i]; + } + + // Create key and nonce for aws-lc-rs + let unbound_key = UnboundKey::new(&AES_128_GCM, key) + .map_err(|e| Error::new(ErrorKind::InvalidInput, format!("Invalid key: {:?}", e)))?; + let opening_key = LessSafeKey::new(unbound_key); + + let nonce = Nonce::try_assume_unique_for_key(&nonce_bytes) + .map_err(|e| Error::new(ErrorKind::InvalidInput, format!("Invalid nonce: {:?}", e)))?; + + let aad = Aad::from(additional_data); + + // aws-lc-rs requires in-place decryption + let mut in_out = ciphertext.to_vec(); + let plaintext = opening_key + .open_in_place(nonce, aad, &mut in_out) + .map_err(|e| { + Error::new( + ErrorKind::InvalidData, + format!("Decryption failed: {:?}", e), + ) + })?; + + Ok(plaintext.to_vec()) +} + +/// Decrypt TLS 1.3 handshake message +/// +/// Decrypts and extracts handshake message, removing ContentType trailer +pub fn decrypt_handshake_message( + key: &[u8], + iv: &[u8], + sequence_number: u64, + ciphertext: &[u8], + record_length: u16, +) -> Result> { + // Additional data for decryption + let mut additional_data = Vec::new(); + additional_data.push(0x17); // ApplicationData + additional_data.extend_from_slice(&[VERSION_TLS_1_2_MAJOR, VERSION_TLS_1_2_MINOR]); // TLS 1.2 + additional_data.extend_from_slice(&record_length.to_be_bytes()); + + let mut plaintext = + decrypt_tls13_record(key, iv, sequence_number, ciphertext, &additional_data)?; + + // Remove ContentType trailer + if plaintext.is_empty() { + return Err(Error::new(ErrorKind::InvalidData, "Empty plaintext")); + } + + // TLS 1.3 has format: content | type_byte | padding (zeros) + // We need to find the type byte by removing trailing zeros first + + // Remove trailing zeros (padding) + while !plaintext.is_empty() && plaintext[plaintext.len() - 1] == 0 { + plaintext.pop(); + } + + if plaintext.is_empty() { + return Err(Error::new(ErrorKind::InvalidData, "Plaintext is all zeros")); + } + + // Now the last byte should be the content type + let content_type = plaintext.pop().unwrap(); + + if content_type != 0x16 + && content_type != CONTENT_TYPE_APPLICATION_DATA + && content_type != CONTENT_TYPE_ALERT + { + return Err(Error::new( + ErrorKind::InvalidData, + format!("Invalid content type: 0x{:02x}", content_type), + )); + } + + Ok(plaintext) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_encrypt_decrypt_record() { + let key = vec![0x42u8; 16]; + let iv = vec![0x99u8; 12]; + let plaintext = b"Hello, TLS 1.3!"; + let aad = b"additional data"; + + let ciphertext = encrypt_tls13_record(&key, &iv, 0, plaintext, aad).unwrap(); + + let decrypted = decrypt_tls13_record(&key, &iv, 0, &ciphertext, aad).unwrap(); + + assert_eq!(&decrypted[..], plaintext); + } + + #[test] + fn test_roundtrip_handshake() { + let key = vec![0x11u8; 16]; + let iv = vec![0x22u8; 12]; + let handshake_msg = vec![0x33u8; 50]; + + // Manually encrypt like encrypt_handshake_to_records does + let mut plaintext = handshake_msg.clone(); + plaintext.push(0x16); // ContentType: Handshake + + let ciphertext_length = (plaintext.len() + 16) as u16; + let aad: [u8; 5] = [ + 0x17, // ApplicationData + 0x03, + 0x03, // TLS 1.2 + (ciphertext_length >> 8) as u8, + (ciphertext_length & 0xff) as u8, + ]; + + let ciphertext = encrypt_tls13_record(&key, &iv, 0, &plaintext, &aad).unwrap(); + + let decrypted = + decrypt_handshake_message(&key, &iv, 0, &ciphertext, ciphertext_length).unwrap(); + + assert_eq!(decrypted, handshake_msg); + } + + #[test] + fn test_encrypt_with_sequence_number() { + let key = vec![0x42u8; 16]; + let iv = vec![0x99u8; 12]; + let plaintext = b"Test with sequence"; + let aad = b"aad"; + + // Test that different sequence numbers produce different ciphertexts + let cipher1 = encrypt_tls13_record(&key, &iv, 1, plaintext, aad).unwrap(); + let cipher2 = encrypt_tls13_record(&key, &iv, 2, plaintext, aad).unwrap(); + let cipher3 = encrypt_tls13_record(&key, &iv, 100, plaintext, aad).unwrap(); + + // Ciphertexts should all be different + assert_ne!(cipher1, cipher2); + assert_ne!(cipher2, cipher3); + assert_ne!(cipher1, cipher3); + + // But they should all decrypt correctly + let decrypt1 = decrypt_tls13_record(&key, &iv, 1, &cipher1, aad).unwrap(); + let decrypt2 = decrypt_tls13_record(&key, &iv, 2, &cipher2, aad).unwrap(); + let decrypt3 = decrypt_tls13_record(&key, &iv, 100, &cipher3, aad).unwrap(); + + assert_eq!(decrypt1, plaintext); + assert_eq!(decrypt2, plaintext); + assert_eq!(decrypt3, plaintext); + } + + #[test] + fn test_decrypt_with_wrong_sequence_number() { + let key = vec![0x42u8; 16]; + let iv = vec![0x99u8; 12]; + let plaintext = b"Test sequence"; + let aad = b"aad"; + + let ciphertext = encrypt_tls13_record(&key, &iv, 5, plaintext, aad).unwrap(); + + // Decrypting with wrong sequence number should fail + let result = decrypt_tls13_record(&key, &iv, 6, &ciphertext, aad); + assert!(result.is_err()); + } + + #[test] + fn test_decrypt_with_wrong_aad() { + let key = vec![0x42u8; 16]; + let iv = vec![0x99u8; 12]; + let plaintext = b"Test AAD"; + let aad = b"correct aad"; + let wrong_aad = b"wrong aad"; + + let ciphertext = encrypt_tls13_record(&key, &iv, 0, plaintext, aad).unwrap(); + + // Decrypting with wrong AAD should fail + let result = decrypt_tls13_record(&key, &iv, 0, &ciphertext, wrong_aad); + assert!(result.is_err()); + } + + #[test] + fn test_encrypt_invalid_key_length() { + let invalid_key = vec![0x42u8; 15]; // Wrong length (not 16) + let iv = vec![0x99u8; 12]; + let plaintext = b"Test"; + let aad = b"aad"; + + let result = encrypt_tls13_record(&invalid_key, &iv, 0, plaintext, aad); + assert!(result.is_err()); + assert_eq!(result.unwrap_err().kind(), ErrorKind::InvalidInput); + } + + #[test] + fn test_encrypt_invalid_iv_length() { + let key = vec![0x42u8; 16]; + let invalid_iv = vec![0x99u8; 11]; // Wrong length (not 12) + let plaintext = b"Test"; + let aad = b"aad"; + + let result = encrypt_tls13_record(&key, &invalid_iv, 0, plaintext, aad); + assert!(result.is_err()); + assert_eq!(result.unwrap_err().kind(), ErrorKind::InvalidInput); + } + + #[test] + fn test_decrypt_corrupted_ciphertext() { + let key = vec![0x42u8; 16]; + let iv = vec![0x99u8; 12]; + let plaintext = b"Test corruption"; + let aad = b"aad"; + + let mut ciphertext = encrypt_tls13_record(&key, &iv, 0, plaintext, aad).unwrap(); + + // Corrupt the ciphertext + ciphertext[5] ^= 0xFF; + + let result = decrypt_tls13_record(&key, &iv, 0, &ciphertext, aad); + assert!(result.is_err()); + } + + #[test] + fn test_encrypt_empty_plaintext() { + let key = vec![0x42u8; 16]; + let iv = vec![0x99u8; 12]; + let plaintext = b""; + let aad = b"aad"; + + let ciphertext = encrypt_tls13_record(&key, &iv, 0, plaintext, aad).unwrap(); + + // Should still produce a ciphertext with auth tag + assert!(ciphertext.len() >= 16); // At least the auth tag + + let decrypted = decrypt_tls13_record(&key, &iv, 0, &ciphertext, aad).unwrap(); + assert_eq!(decrypted, plaintext); + } + + #[test] + fn test_large_plaintext() { + let key = vec![0x42u8; 16]; + let iv = vec![0x99u8; 12]; + let plaintext = vec![0xAB; 16384]; // 16KB + let aad = b"aad"; + + let ciphertext = encrypt_tls13_record(&key, &iv, 42, &plaintext, aad).unwrap(); + let decrypted = decrypt_tls13_record(&key, &iv, 42, &ciphertext, aad).unwrap(); + + assert_eq!(decrypted, plaintext); + } +} diff --git a/src/reality/reality_auth.rs b/src/reality/reality_auth.rs new file mode 100644 index 00000000..4e51aab5 --- /dev/null +++ b/src/reality/reality_auth.rs @@ -0,0 +1,552 @@ +use aws_lc_rs::aead::{Aad, LessSafeKey, Nonce, UnboundKey, AES_256_GCM}; +use aws_lc_rs::agreement; +use aws_lc_rs::hkdf::{Salt, HKDF_SHA256}; + +#[cfg(test)] +use std::time::{SystemTime, UNIX_EPOCH}; + +/// Custom error type for REALITY cryptographic operations +#[derive(Debug)] +pub enum CryptoError { + InvalidKeyLength, + InvalidNonceLength, + InvalidCiphertextLength, + EncryptionFailed, + DecryptionFailed, + EcdhFailed, + HkdfFailed, + CertificateGenerationFailed(String), +} + +impl std::fmt::Display for CryptoError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + CryptoError::InvalidKeyLength => write!(f, "Invalid key length"), + CryptoError::InvalidNonceLength => write!(f, "Invalid nonce length"), + CryptoError::InvalidCiphertextLength => write!(f, "Invalid ciphertext length"), + CryptoError::EncryptionFailed => write!(f, "Encryption failed"), + CryptoError::DecryptionFailed => write!(f, "Decryption failed"), + CryptoError::EcdhFailed => write!(f, "ECDH key exchange failed"), + CryptoError::HkdfFailed => write!(f, "HKDF derivation failed"), + CryptoError::CertificateGenerationFailed(e) => { + write!(f, "Certificate generation failed: {}", e) + } + } + } +} + +impl std::error::Error for CryptoError {} + +impl From for std::io::Error { + fn from(err: CryptoError) -> Self { + std::io::Error::new(std::io::ErrorKind::InvalidData, err.to_string()) + } +} + +/// Performs X25519 ECDH key exchange +/// +/// # Arguments +/// * `private_key` - 32-byte X25519 private key +/// * `public_key` - 32-byte X25519 public key +/// +/// # Returns +/// 32-byte shared secret +pub fn perform_ecdh( + private_key: &[u8; 32], + public_key: &[u8; 32], +) -> Result<[u8; 32], CryptoError> { + // Create private key from raw bytes + let my_private_key = agreement::PrivateKey::from_private_key(&agreement::X25519, private_key) + .map_err(|_| CryptoError::EcdhFailed)?; + + // Compute public key from private key (not used, but verifies key is valid) + let _my_public_key = my_private_key + .compute_public_key() + .map_err(|_| CryptoError::EcdhFailed)?; + + // Create unparsed public key + let peer_public_key = + agreement::UnparsedPublicKey::new(&agreement::X25519, public_key.as_ref()); + + // Perform ECDH + let mut shared_secret = [0u8; 32]; + agreement::agree( + &my_private_key, + peer_public_key, + CryptoError::EcdhFailed, + |key_material| { + shared_secret.copy_from_slice(key_material); + Ok(()) + }, + )?; + + Ok(shared_secret) +} + +/// Derives authentication key using HKDF-SHA256 +/// +/// # Arguments +/// * `shared_secret` - 32-byte shared secret from ECDH +/// * `salt` - Salt bytes (must be exactly 20 bytes, from ClientHello.Random[0..20]) +/// * `info` - Context string (should be b"REALITY") +/// +/// # Returns +/// 32-byte derived authentication key +/// +/// # Panics +/// Panics if salt is not exactly 20 bytes. +pub fn derive_auth_key( + shared_secret: &[u8; 32], + salt: &[u8], + info: &[u8], +) -> Result<[u8; 32], CryptoError> { + debug_assert_eq!(salt.len(), 20, "salt must be exactly 20 bytes"); + let salt = Salt::new(HKDF_SHA256, salt); + let prk = salt.extract(shared_secret); + let info_pieces = [info]; + let okm = prk + .expand(&info_pieces, HKDF_SHA256) + .map_err(|_| CryptoError::HkdfFailed)?; + let mut auth_key = [0u8; 32]; + okm.fill(&mut auth_key) + .map_err(|_| CryptoError::HkdfFailed)?; + Ok(auth_key) +} + +/// Encrypts SessionId using AES-256-GCM +/// +/// # Arguments +/// * `plaintext` - 16-byte plaintext (first 16 bytes of SessionId) +/// * `auth_key` - 32-byte authentication key +/// * `nonce` - 12-byte nonce (ClientHello.Random[20..32]) +/// * `aad` - Additional authenticated data (entire ClientHello) +/// +/// # Returns +/// 32-byte result (16 bytes ciphertext + 16 bytes GCM tag) +/// +/// # Panics +/// Panics if nonce is not exactly 12 bytes. +pub fn encrypt_session_id( + plaintext: &[u8; 16], + auth_key: &[u8; 32], + nonce: &[u8], + aad: &[u8], +) -> Result<[u8; 32], CryptoError> { + debug_assert_eq!(nonce.len(), 12, "nonce must be exactly 12 bytes"); + let unbound_key = + UnboundKey::new(&AES_256_GCM, auth_key).map_err(|_| CryptoError::EncryptionFailed)?; + let sealing_key = LessSafeKey::new(unbound_key); + + let nonce_obj = + Nonce::try_assume_unique_for_key(nonce).map_err(|_| CryptoError::InvalidNonceLength)?; + + let aad_obj = Aad::from(aad); + + // aws-lc-rs requires in-place encryption + let mut in_out = plaintext.to_vec(); + sealing_key + .seal_in_place_append_tag(nonce_obj, aad_obj, &mut in_out) + .map_err(|_| CryptoError::EncryptionFailed)?; + + if in_out.len() != 32 { + return Err(CryptoError::EncryptionFailed); + } + + let mut result = [0u8; 32]; + result.copy_from_slice(&in_out); + Ok(result) +} + +/// Decrypts SessionId using AES-256-GCM +/// +/// # Arguments +/// * `ciphertext_and_tag` - 32-byte encrypted data (16 ciphertext + 16 tag) +/// * `auth_key` - 32-byte authentication key +/// * `nonce` - 12-byte nonce (ClientHello.Random[20..32]) +/// * `aad` - Additional authenticated data (entire ClientHello) +/// +/// # Returns +/// 16-byte decrypted plaintext +/// +/// # Panics +/// Panics if nonce is not exactly 12 bytes. +pub fn decrypt_session_id( + ciphertext_and_tag: &[u8; 32], + auth_key: &[u8; 32], + nonce: &[u8], + aad: &[u8], +) -> Result<[u8; 16], CryptoError> { + debug_assert_eq!(nonce.len(), 12, "nonce must be exactly 12 bytes"); + let unbound_key = + UnboundKey::new(&AES_256_GCM, auth_key).map_err(|_| CryptoError::DecryptionFailed)?; + let opening_key = LessSafeKey::new(unbound_key); + + let nonce_obj = + Nonce::try_assume_unique_for_key(nonce).map_err(|_| CryptoError::InvalidNonceLength)?; + + let aad_obj = Aad::from(aad); + + // aws-lc-rs requires in-place decryption + let mut in_out = ciphertext_and_tag.to_vec(); + let plaintext = opening_key + .open_in_place(nonce_obj, aad_obj, &mut in_out) + .map_err(|_| CryptoError::DecryptionFailed)?; + + if plaintext.len() != 16 { + return Err(CryptoError::DecryptionFailed); + } + + let mut result = [0u8; 16]; + result.copy_from_slice(plaintext); + Ok(result) +} + +#[cfg(test)] +/// Creates a REALITY SessionId (test helper) +fn create_session_id(version: [u8; 3], timestamp: u32, short_id: &[u8; 8]) -> [u8; 32] { + let mut session_id = [0u8; 32]; + session_id[0] = version[0]; // Major version + session_id[1] = version[1]; // Minor version + session_id[2] = version[2]; // Patch version + // session_id[3] = 0 (reserved) + session_id[4..8].copy_from_slice(×tamp.to_be_bytes()); + session_id[8..16].copy_from_slice(short_id); + // session_id[16..32] remain zeros + session_id +} + +#[cfg(test)] +/// Gets current Unix timestamp (test helper) +fn get_current_timestamp() -> u32 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("Time went backwards") + .as_secs() as u32 +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_ecdh() { + // Test vectors + let private_key_a = [1u8; 32]; + let private_key_b = [2u8; 32]; + + // Compute public keys using aws-lc-rs + let priv_a = + agreement::PrivateKey::from_private_key(&agreement::X25519, &private_key_a).unwrap(); + let pub_a = priv_a.compute_public_key().unwrap(); + + let priv_b = + agreement::PrivateKey::from_private_key(&agreement::X25519, &private_key_b).unwrap(); + let pub_b = priv_b.compute_public_key().unwrap(); + + // Perform ECDH both ways + let mut shared_a_bytes = [0u8; 32]; + agreement::agree( + &priv_a, + &agreement::UnparsedPublicKey::new(&agreement::X25519, pub_b.as_ref()), + (), + |key_material| { + shared_a_bytes.copy_from_slice(key_material); + Ok(()) + }, + ) + .unwrap(); + + let mut shared_b_bytes = [0u8; 32]; + agreement::agree( + &priv_b, + &agreement::UnparsedPublicKey::new(&agreement::X25519, pub_a.as_ref()), + (), + |key_material| { + shared_b_bytes.copy_from_slice(key_material); + Ok(()) + }, + ) + .unwrap(); + + assert_eq!(shared_a_bytes, shared_b_bytes); + } + + #[test] + fn test_session_id_structure() { + // Test that session ID has the correct structure + let version = [1, 8, 1]; // major, minor, patch + let timestamp = get_current_timestamp(); + let short_id = [0x12, 0x34, 0x56, 0x78, 0x9A, 0xBC, 0xDE, 0xF0]; + + let session_id = create_session_id(version, timestamp, &short_id); + + // Verify structure + assert_eq!(session_id[0], 1); // version major + assert_eq!(session_id[1], 8); // version minor + assert_eq!(session_id[2], 1); // version patch + assert_eq!(session_id[3], 0); // reserved + + // Verify timestamp + let extracted_timestamp = + u32::from_be_bytes([session_id[4], session_id[5], session_id[6], session_id[7]]); + assert_eq!(extracted_timestamp, timestamp); + + // Verify short_id + assert_eq!(&session_id[8..16], &short_id[..]); + + // Verify remaining bytes are zero + for i in 16..32 { + assert_eq!(session_id[i], 0); + } + } + + #[test] + fn test_version_comparison() { + // Test version comparison logic (same as used in reality_server_handler.rs) + let v1_8_1 = [1u8, 8, 1]; + let v1_8_0 = [1u8, 8, 0]; + let v1_9_0 = [1u8, 9, 0]; + let v2_0_0 = [2u8, 0, 0]; + + assert!(v1_8_0 < v1_8_1); + assert!(v1_8_1 < v1_9_0); + assert!(v1_9_0 < v2_0_0); + assert!(v1_8_1 > v1_8_0); + } + + #[test] + fn test_timestamp_validation_logic() { + // Test the timestamp validation logic used in reality_server_handler.rs + let now = get_current_timestamp(); + + // Test within bounds (60 seconds = 60000 milliseconds) + let max_diff_ms = 60000u64; + let max_diff_secs = max_diff_ms / 1000; + + // Current time - should pass + let diff = now.abs_diff(now); + assert!((diff as u64) <= max_diff_secs); + + // 30 seconds ago - should pass + let past_timestamp = now - 30; + let diff = now.abs_diff(past_timestamp); + assert!((diff as u64) <= max_diff_secs); + + // 30 seconds future - should pass + let future_timestamp = now + 30; + let diff = now.abs_diff(future_timestamp); + assert!((diff as u64) <= max_diff_secs); + + // 2 minutes ago - should fail + let old_timestamp = now.saturating_sub(120); + let diff = now.abs_diff(old_timestamp); + assert!((diff as u64) > max_diff_secs); + + // 2 minutes future - should fail + let future_timestamp = now + 120; + let diff = now.abs_diff(future_timestamp); + assert!((diff as u64) > max_diff_secs); + } + + #[test] + fn test_session_id_encryption_preserves_structure() { + // Create session ID with known values + let version = [1, 8, 1]; + let timestamp = 1234567890u32; + let short_id = [0xAB; 8]; + + let session_id = create_session_id(version, timestamp, &short_id); + + // Perform ECDH + let client_private = [0x01; 32]; + let server_public = [0x02; 32]; + let shared_secret = perform_ecdh(&client_private, &server_public).unwrap(); + + // Derive auth key + let salt = [0x03; 20]; + let auth_key = derive_auth_key(&shared_secret, &salt, b"REALITY").unwrap(); + + // Encrypt session ID (first 16 bytes) + let plaintext: [u8; 16] = session_id[0..16].try_into().unwrap(); + let nonce = [0x04; 12]; + let aad = b"test additional authenticated data"; + + let encrypted = encrypt_session_id(&plaintext, &auth_key, &nonce, aad).unwrap(); + + // Decrypt session ID + let decrypted = decrypt_session_id(&encrypted, &auth_key, &nonce, aad).unwrap(); + + // Verify we can recover the structure + assert_eq!(decrypted[0], version[0]); + assert_eq!(decrypted[1], version[1]); + assert_eq!(decrypted[2], version[2]); + assert_eq!(decrypted[3], 0); // reserved + + let recovered_timestamp = + u32::from_be_bytes([decrypted[4], decrypted[5], decrypted[6], decrypted[7]]); + assert_eq!(recovered_timestamp, timestamp); + + assert_eq!(&decrypted[8..16], &short_id[..]); + } + + #[test] + fn test_hkdf() { + let shared_secret = [0x42u8; 32]; + let salt = [0x43u8; 20]; + let info = b"REALITY"; + + let auth_key = derive_auth_key(&shared_secret, &salt, info).unwrap(); + + // Verify length + assert_eq!(auth_key.len(), 32); + + // Verify deterministic + let auth_key2 = derive_auth_key(&shared_secret, &salt, info).unwrap(); + assert_eq!(auth_key, auth_key2); + + // Verify different salt produces different key + let salt2 = [0x44u8; 20]; + let auth_key3 = derive_auth_key(&shared_secret, &salt2, info).unwrap(); + assert_ne!(auth_key, auth_key3); + } + + #[test] + fn test_aes_gcm_roundtrip() { + let plaintext = [0x55u8; 16]; + let auth_key = [0x66u8; 32]; + let nonce = [0x77u8; 12]; + let aad = b"additional authenticated data"; + + let encrypted = encrypt_session_id(&plaintext, &auth_key, &nonce, aad).unwrap(); + assert_eq!(encrypted.len(), 32); + + let decrypted = decrypt_session_id(&encrypted, &auth_key, &nonce, aad).unwrap(); + assert_eq!(plaintext, decrypted); + } + + #[test] + fn test_aes_gcm_wrong_key_fails() { + let plaintext = [0x55u8; 16]; + let auth_key = [0x66u8; 32]; + let wrong_key = [0x67u8; 32]; + let nonce = [0x77u8; 12]; + let aad = b"additional authenticated data"; + + let encrypted = encrypt_session_id(&plaintext, &auth_key, &nonce, aad).unwrap(); + + let result = decrypt_session_id(&encrypted, &wrong_key, &nonce, aad); + assert!(result.is_err()); + } + + #[test] + fn test_aes_gcm_wrong_aad_fails() { + let plaintext = [0x55u8; 16]; + let auth_key = [0x66u8; 32]; + let nonce = [0x77u8; 12]; + let aad = b"additional authenticated data"; + let wrong_aad = b"wrong additional authenticated data"; + + let encrypted = encrypt_session_id(&plaintext, &auth_key, &nonce, aad).unwrap(); + + let result = decrypt_session_id(&encrypted, &auth_key, &nonce, wrong_aad); + assert!(result.is_err()); + } + + #[test] + fn test_create_session_id() { + let version = [1, 8, 1]; + let timestamp = 1234567890u32; + let short_id = [0xAB; 8]; + + let session_id = create_session_id(version, timestamp, &short_id); + + assert_eq!(session_id[0], 1); + assert_eq!(session_id[1], 8); + assert_eq!(session_id[2], 1); + assert_eq!(session_id[3], 0); + + let ts = u32::from_be_bytes([session_id[4], session_id[5], session_id[6], session_id[7]]); + assert_eq!(ts, timestamp); + + assert_eq!(&session_id[8..16], &short_id[..]); + + // Verify remaining bytes are zeros + for i in 16..32 { + assert_eq!(session_id[i], 0); + } + } + + #[test] + fn test_validation_scenarios() { + println!("\n=== REALITY Validation Test Scenarios ===\n"); + + // Scenario 1: Version validation + println!("Scenario 1: Version Validation"); + println!(" Client version: [1, 8, 1]"); + println!(" Min version: [1, 8, 0]"); + println!(" Max version: [1, 9, 0]"); + println!(" Expected: PASS ✓"); + + let client_version = [1u8, 8, 1]; + let min_version = [1u8, 8, 0]; + let max_version = [1u8, 9, 0]; + assert!(client_version >= min_version); + assert!(client_version <= max_version); + + // Scenario 2: Version below minimum + println!("\nScenario 2: Version Below Minimum"); + println!(" Client version: [1, 7, 5]"); + println!(" Min version: [1, 8, 0]"); + println!(" Expected: FAIL ✗"); + + let old_client_version = [1u8, 7, 5]; + assert!(old_client_version < min_version); + + // Scenario 3: Version above maximum + println!("\nScenario 3: Version Above Maximum"); + println!(" Client version: [2, 0, 0]"); + println!(" Max version: [1, 9, 0]"); + println!(" Expected: FAIL ✗"); + + let new_client_version = [2u8, 0, 0]; + assert!(new_client_version > max_version); + + // Scenario 4: Timestamp validation (within bounds) + println!("\nScenario 4: Timestamp Within Bounds"); + let now = get_current_timestamp(); + let client_timestamp = now - 30; // 30 seconds ago + let max_diff_ms = 60000u64; // 60 seconds + + let diff = now.abs_diff(client_timestamp); + println!(" Server time: {}", now); + println!(" Client time: {} (30 seconds ago)", client_timestamp); + println!(" Max allowed: {} seconds", max_diff_ms / 1000); + println!(" Actual diff: {} seconds", diff); + println!(" Expected: PASS ✓"); + assert!((diff as u64) <= (max_diff_ms / 1000)); + + // Scenario 5: Timestamp validation (out of bounds) + println!("\nScenario 5: Timestamp Out of Bounds"); + let old_timestamp = now - 120; // 2 minutes ago + let diff = now.abs_diff(old_timestamp); + println!(" Server time: {}", now); + println!(" Client time: {} (2 minutes ago)", old_timestamp); + println!(" Max allowed: {} seconds", max_diff_ms / 1000); + println!(" Actual diff: {} seconds", diff); + println!(" Expected: FAIL ✗"); + assert!((diff as u64) > (max_diff_ms / 1000)); + + // Scenario 6: Future timestamp validation + println!("\nScenario 6: Future Timestamp Validation"); + let future_timestamp = now + 45; // 45 seconds in future + let diff = now.abs_diff(future_timestamp); + println!(" Server time: {}", now); + println!(" Client time: {} (45 seconds in future)", future_timestamp); + println!(" Max allowed: {} seconds", max_diff_ms / 1000); + println!(" Actual diff: {} seconds", diff); + println!(" Expected: PASS ✓"); + assert!((diff as u64) <= (max_diff_ms / 1000)); + + println!("\n=== All Validation Scenarios Passed ===\n"); + } +} diff --git a/src/reality/reality_certificate.rs b/src/reality/reality_certificate.rs new file mode 100644 index 00000000..4786eeff --- /dev/null +++ b/src/reality/reality_certificate.rs @@ -0,0 +1,178 @@ +use aws_lc_rs::hmac; +use aws_lc_rs::signature::Ed25519KeyPair; +use std::io::{Error, ErrorKind, Result}; + +/// Generate a template Ed25519 certificate for the given hostname +fn generate_template_cert(hostname: &str) -> (Vec, Ed25519KeyPair) { + // Generate Ed25519 keypair using rcgen + let key_pair = rcgen::KeyPair::generate_for(&rcgen::PKCS_ED25519) + .expect("Failed to generate Ed25519 key pair"); + + // Generate a minimal self-signed certificate using rcgen with Ed25519 + let params = rcgen::CertificateParams::new(vec![hostname.to_string()]) + .expect("Failed to create certificate params"); + + // Create self-signed certificate + let cert = params + .self_signed(&key_pair) + .expect("Failed to create self-signed certificate"); + + let cert_der = cert.der().to_vec(); + + // Create the Ed25519KeyPair from the PKCS#8 document + let signing_key = Ed25519KeyPair::from_pkcs8(key_pair.serialized_der()) + .expect("Failed to parse generated key"); + + use aws_lc_rs::signature::KeyPair; + let public_key_bytes = signing_key.public_key().as_ref(); + + log::debug!( + "REALITY DEBUG: Generated template certificate ({} bytes) for {}", + cert_der.len(), + hostname + ); + log::debug!( + "REALITY DEBUG: Using public key: {:?}", + &public_key_bytes[..16.min(public_key_bytes.len())] + ); + + // Debug: print last 70 bytes to see signature structure + if cert_der.len() > 70 { + log::debug!( + "REALITY DEBUG: Last 70 bytes: {:?}", + &cert_der[cert_der.len() - 70..] + ); + } + + (cert_der, signing_key) +} + +/// Generate HMAC-signed Ed25519 certificate +/// +/// This follows the uTLS approach: +/// 1. Generate a certificate for the destination hostname +/// 2. Replace the signature (last 64 bytes) with HMAC-SHA512(auth_key, ed25519_public_key) +/// +/// This is NOT a cryptographically valid certificate! +pub fn generate_hmac_certificate( + auth_key: &[u8; 32], + hostname: &str, +) -> Result<(Vec, Ed25519KeyPair)> { + // Generate a new certificate for the destination hostname + let (cert_der, signing_key) = generate_template_cert(hostname); + + log::debug!( + "REALITY DEBUG: Using template certificate ({} bytes)", + cert_der.len() + ); + + // Extract the Ed25519 public key from the signing key + use aws_lc_rs::signature::KeyPair; + let public_key_bytes = signing_key.public_key().as_ref(); + + // Replace signature with HMAC + let cert_with_hmac = replace_signature_with_hmac(cert_der, auth_key, public_key_bytes)?; + + log::debug!("REALITY DEBUG: Replaced signature with HMAC"); + + // Return the HMAC-signed certificate with the signing key + Ok((cert_with_hmac, signing_key)) +} + +/// Replace the certificate signature with HMAC-SHA512(auth_key, ed25519_public_key) +fn replace_signature_with_hmac( + mut cert_der: Vec, + auth_key: &[u8; 32], + public_key_bytes: &[u8], +) -> Result> { + // The signature is the last BIT STRING in the certificate DER + // For Ed25519, it's always 64 bytes + + // Find the signature BIT STRING tag + // It should be near the end of the certificate + // BIT STRING tag = 0x03 + // Length = 0x41 (65 bytes: 1 byte unused bits + 64 bytes signature) + + // Search backwards for BIT STRING tag + let mut sig_offset = None; + for i in (0..cert_der.len().saturating_sub(66)).rev() { + if cert_der[i] == 0x03 && cert_der[i + 1] == 0x41 && cert_der[i + 2] == 0x00 { + // Found it! + sig_offset = Some(i + 3); // Skip tag, length, unused bits + break; + } + } + + let sig_offset = sig_offset.ok_or_else(|| { + Error::new( + ErrorKind::InvalidData, + "Could not find signature in certificate DER", + ) + })?; + + if sig_offset + 64 > cert_der.len() { + return Err(Error::new( + ErrorKind::InvalidData, + "Invalid signature offset", + )); + } + + // Compute HMAC-SHA512(auth_key, ed25519_public_key) + let key = hmac::Key::new(hmac::HMAC_SHA512, auth_key); + let tag = hmac::sign(&key, public_key_bytes); + let hmac_bytes = tag.as_ref(); + + // Replace signature bytes with HMAC + cert_der[sig_offset..sig_offset + 64].copy_from_slice(hmac_bytes); + + log::debug!( + "REALITY DEBUG: HMAC signature (first 16 bytes): {:?}", + &hmac_bytes[..16] + ); + + Ok(cert_der) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_generate_hmac_certificate() { + let auth_key = [42u8; 32]; + let hostname = "test.example.com"; + let result = generate_hmac_certificate(&auth_key, hostname); + + assert!(result.is_ok()); + + let (cert_der, _signing_key) = result.unwrap(); + + // Certificate should be a reasonable size (few hundred bytes) + assert!(cert_der.len() > 100); + assert!(cert_der.len() < 1000); + + // Should start with SEQUENCE tag + assert_eq!(cert_der[0], 0x30); + } + + #[test] + fn test_hmac_signature_deterministic() { + let auth_key = [99u8; 32]; + let hostname = "test.example.com"; + + let (cert1, key1) = generate_hmac_certificate(&auth_key, hostname).unwrap(); + use aws_lc_rs::signature::KeyPair; + let public_key1 = key1.public_key().as_ref().to_vec(); + + // Generate another certificate with same auth_key + let (cert2, key2) = generate_hmac_certificate(&auth_key, hostname).unwrap(); + let public_key2 = key2.public_key().as_ref().to_vec(); + + // Different Ed25519 keys (random generation) + assert_ne!(public_key1, public_key2); + + // But certificates should have different signatures + // (because they're based on different public keys) + assert_ne!(cert1, cert2); + } +} diff --git a/src/reality/reality_client_connection.rs b/src/reality/reality_client_connection.rs new file mode 100644 index 00000000..d6c1ece7 --- /dev/null +++ b/src/reality/reality_client_connection.rs @@ -0,0 +1,865 @@ +// REALITY client-side connection +// +// This implements a rustls-compatible API for REALITY protocol client connections, +// allowing REALITY to be used as a drop-in replacement for rustls. + +use std::io::{self, Read, Write}; + +use crate::slide_buffer::SlideBuffer; + +use aws_lc_rs::{ + agreement, digest, + rand::{SecureRandom, SystemRandom}, +}; + +use super::common::{ + self, CIPHERTEXT_READ_BUF_CAPACITY, CONTENT_TYPE_ALERT, CONTENT_TYPE_APPLICATION_DATA, + CONTENT_TYPE_CHANGE_CIPHER_SPEC, CONTENT_TYPE_HANDSHAKE, HANDSHAKE_TYPE_CERTIFICATE, + HANDSHAKE_TYPE_CERTIFICATE_VERIFY, HANDSHAKE_TYPE_ENCRYPTED_EXTENSIONS, + HANDSHAKE_TYPE_FINISHED, PLAINTEXT_READ_BUF_CAPACITY, TLS_MAX_RECORD_SIZE, + TLS_RECORD_HEADER_SIZE, +}; +use super::reality_aead::{decrypt_handshake_message, decrypt_tls13_record}; +use super::reality_auth::{derive_auth_key, encrypt_session_id, perform_ecdh}; +use super::reality_client_verify::{extract_certificate_der, verify_certificate_hmac}; +use super::reality_io_state::RealityIoState; +use super::reality_reader_writer::{RealityReader, RealityWriter}; +use super::reality_records::{encrypt_handshake_to_records, encrypt_plaintext_to_records}; +use super::reality_tls13_keys::{ + compute_finished_verify_data, derive_application_secrets, derive_handshake_keys, + derive_traffic_keys, +}; +use super::reality_tls13_messages::*; +use super::reality_util::extract_server_public_key; + +/// Configuration for REALITY client connections +#[derive(Clone)] +pub struct RealityClientConfig { + /// Server's X25519 public key (32 bytes) + pub public_key: [u8; 32], + /// Short ID for authentication (8 bytes) + pub short_id: [u8; 8], + /// Server name for SNI + pub server_name: String, +} + +/// Handshake state machine for REALITY client +enum HandshakeState { + /// ClientHello sent, waiting for ServerHello + AwaitingServerHello { + client_hello_hash: [u8; 32], + client_hello_bytes: Vec, // Full ClientHello handshake message + client_private_key: [u8; 32], + auth_key: [u8; 32], // REALITY authentication key for HMAC verification + }, + /// ServerHello received, processing encrypted handshake messages + ProcessingHandshake { + client_handshake_traffic_secret: Vec, + server_handshake_traffic_secret: Vec, + master_secret: Vec, + cipher_suite: u16, + handshake_transcript_bytes: Vec, // Accumulated transcript for hash computation + auth_key: [u8; 32], // REALITY authentication key for HMAC verification + }, + /// Handshake complete, ready for application data + Complete, +} + +/// REALITY client-side connection implementing rustls-compatible API +pub struct RealityClientConnection { + // Configuration + config: RealityClientConfig, + + // Handshake state + handshake_state: HandshakeState, + + // TLS 1.3 application traffic encryption (post-handshake) + app_read_key: Option>, + app_read_iv: Option>, + app_write_key: Option>, + app_write_iv: Option>, + read_seq: u64, + write_seq: u64, + cipher_suite: u16, + + // Pre-allocated buffer for TLS read operations (reused across calls) + tls_read_buffer: Box<[u8; TLS_MAX_RECORD_SIZE]>, + + // Buffers for I/O - using SlideBuffer for efficient zero-alloc operations + ciphertext_read_buf: SlideBuffer, // Incoming encrypted TLS records + ciphertext_write_buf: Vec, // Outgoing encrypted TLS records + plaintext_read_buf: SlideBuffer, // Decrypted application data + plaintext_write_buf: Vec, // Application data to encrypt +} + +impl RealityClientConnection { + /// Create a new REALITY client connection and generate ClientHello + pub fn new(config: RealityClientConfig) -> io::Result { + let mut conn = RealityClientConnection { + config, + handshake_state: HandshakeState::AwaitingServerHello { + client_hello_hash: [0u8; 32], + client_hello_bytes: Vec::new(), + client_private_key: [0u8; 32], + auth_key: [0u8; 32], + }, + app_read_key: None, + app_read_iv: None, + app_write_key: None, + app_write_iv: None, + read_seq: 0, + write_seq: 0, + cipher_suite: 0, + tls_read_buffer: Box::new([0u8; TLS_MAX_RECORD_SIZE]), + ciphertext_read_buf: SlideBuffer::new(CIPHERTEXT_READ_BUF_CAPACITY), + ciphertext_write_buf: Vec::with_capacity(CIPHERTEXT_READ_BUF_CAPACITY), + plaintext_read_buf: SlideBuffer::new(PLAINTEXT_READ_BUF_CAPACITY), + plaintext_write_buf: Vec::with_capacity(TLS_MAX_RECORD_SIZE), + }; + + // Generate ClientHello immediately + conn.generate_client_hello()?; + + Ok(conn) + } + + /// Generate and buffer ClientHello + fn generate_client_hello(&mut self) -> io::Result<()> { + let rng = SystemRandom::new(); + + // Generate our X25519 keypair + let mut our_private_bytes = [0u8; 32]; + rng.fill(&mut our_private_bytes) + .map_err(|_| io::Error::other("RNG failed"))?; + + let our_private_key = + agreement::PrivateKey::from_private_key(&agreement::X25519, &our_private_bytes) + .map_err(|_| io::Error::other("Failed to create X25519 key"))?; + let our_public_key_bytes = our_private_key + .compute_public_key() + .map_err(|_| io::Error::other("Failed to compute public key"))?; + + // Generate client random + let mut client_random = [0u8; 32]; + rng.fill(&mut client_random) + .map_err(|_| io::Error::other("RNG failed"))?; + + // Perform ECDH with server's public key to derive auth key + let shared_secret = perform_ecdh(&our_private_bytes, &self.config.public_key) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()))?; + + // Use slice directly from client_random to avoid copying + let auth_key = derive_auth_key(&shared_secret, &client_random[0..20], b"REALITY") + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()))?; + + // Create session ID with REALITY metadata + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map_err(|_| io::Error::other("System time error"))? + .as_secs(); + + let mut session_id_plaintext = [0u8; 16]; + session_id_plaintext[0] = 1; // Protocol version major + session_id_plaintext[1] = 8; // Protocol version minor + session_id_plaintext[2] = 0; // Protocol version patch + session_id_plaintext[3] = 0; // Padding byte + // Timestamp (4 bytes as uint32, in seconds) + session_id_plaintext[4..8].copy_from_slice(&(timestamp as u32).to_be_bytes()); + // Short ID (8 bytes) + session_id_plaintext[8..16].copy_from_slice(&self.config.short_id); + + // Create a 32-byte SessionId (16 bytes plaintext + 16 bytes zeros for padding) + let mut session_id_for_hello = [0u8; 32]; + session_id_for_hello[0..16].copy_from_slice(&session_id_plaintext); + + // Build ClientHello with plaintext SessionId first + let mut client_hello = construct_client_hello( + &client_random, + &session_id_for_hello, + our_public_key_bytes.as_ref(), + &self.config.server_name, + )?; + + // Now encrypt the SessionId using the ClientHello with zeroed SessionId as AAD + // Use slice directly from client_random to avoid copying + let nonce = &client_random[20..32]; + + // Zero out the SessionId in ClientHello to create AAD (matches what server will use) + // SessionId is at offset 39 in ClientHello handshake + client_hello[39..71].fill(0); + + log::debug!("REALITY CLIENT: Encrypting SessionId"); + log::debug!(" auth_key={:02x?}", &auth_key); + log::debug!(" nonce={:02x?}", nonce); + log::debug!(" plaintext={:02x?}", &session_id_plaintext); + log::debug!( + " aad_len={} (ClientHello with zero SessionId)", + client_hello.len() + ); + log::debug!(" aad[0..4]={:02x?}", &client_hello[0..4]); + + let encrypted_session_id = + encrypt_session_id(&session_id_plaintext, &auth_key, nonce, &client_hello) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()))?; + + log::debug!( + "REALITY CLIENT: Encrypted SessionId={:02x?}", + &encrypted_session_id + ); + + // Replace the zeros with the encrypted SessionId + client_hello[39..71].copy_from_slice(&encrypted_session_id); + + // Compute ClientHello hash + let mut ch_transcript = digest::Context::new(&digest::SHA256); + ch_transcript.update(&client_hello); + let client_hello_hash = ch_transcript.finish(); + let mut client_hello_hash_arr = [0u8; 32]; + client_hello_hash_arr.copy_from_slice(client_hello_hash.as_ref()); + + // Wrap in TLS record + let mut record = write_record_header(CONTENT_TYPE_HANDSHAKE, client_hello.len() as u16); + record.extend_from_slice(&client_hello); + + // Buffer for sending + self.ciphertext_write_buf.extend_from_slice(&record); + + // Update state + self.handshake_state = HandshakeState::AwaitingServerHello { + client_hello_hash: client_hello_hash_arr, + client_hello_bytes: client_hello.clone(), // Save the actual ClientHello bytes + client_private_key: our_private_bytes, + auth_key, // Save auth_key for HMAC certificate verification + }; + + log::debug!( + "REALITY: ClientHello generated and buffered ({} bytes)", + record.len() + ); + + Ok(()) + } + + /// Read TLS messages from the provided reader into internal buffer + /// + /// Uses pre-allocated buffer to avoid allocation on every call. + pub fn read_tls(&mut self, rd: &mut dyn Read) -> io::Result { + // Compact if remaining capacity is insufficient for a full TLS record + if self.ciphertext_read_buf.remaining_capacity() < TLS_MAX_RECORD_SIZE { + self.ciphertext_read_buf.compact(); + } + + // Read into pre-allocated buffer + let n = rd.read(&mut self.tls_read_buffer[..])?; + if n > 0 { + self.ciphertext_read_buf + .extend_from_slice(&self.tls_read_buffer[..n]); + } + Ok(n) + } + + /// Process buffered packets and advance state machine + pub fn process_new_packets(&mut self) -> io::Result { + match &self.handshake_state { + HandshakeState::AwaitingServerHello { .. } => { + self.process_server_hello()?; + } + HandshakeState::ProcessingHandshake { .. } => { + self.process_encrypted_handshake()?; + } + HandshakeState::Complete => { + self.process_application_data()?; + } + } + + Ok(RealityIoState::new(self.plaintext_read_buf.len())) + } + + /// Process ServerHello + fn process_server_hello(&mut self) -> io::Result<()> { + // Extract state + let (client_hello_hash, client_private_key, auth_key) = match &self.handshake_state { + HandshakeState::AwaitingServerHello { + client_hello_hash, + client_hello_bytes: _, + client_private_key, + auth_key, + } => (*client_hello_hash, *client_private_key, *auth_key), + _ => return Ok(()), // Wrong state + }; + + // Check if we have enough data for a TLS record header + if self.ciphertext_read_buf.len() < TLS_RECORD_HEADER_SIZE { + return Ok(()); // Need more data + } + + // Parse TLS record length + let record_len = self + .ciphertext_read_buf + .get_u16_be(3) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "Buffer too short"))? + as usize; + + // Check if we have the complete record + let total_record_len = TLS_RECORD_HEADER_SIZE + record_len; + if self.ciphertext_read_buf.len() < total_record_len { + return Ok(()); // Need more data + } + + // Copy ServerHello record to Vec for processing + let record: Vec = self.ciphertext_read_buf[..total_record_len].to_vec(); + self.ciphertext_read_buf.consume(total_record_len); + let server_hello = &record[TLS_RECORD_HEADER_SIZE..]; // Skip TLS record header (includes handshake header) + + log::debug!( + "REALITY CLIENT: ServerHello for transcript: len={}, bytes={:02x?}", + server_hello.len(), + server_hello + ); + + // Extract server public key from ServerHello + let server_public_key = extract_server_public_key(&record)?; + + // Compute transcript hash including ClientHello and ServerHello + // NOTE: We need to update with the actual ClientHello bytes, not its hash! + // We need to preserve the ClientHello bytes to include them in the transcript + // For now, we'll need to get them from the state + + // Get the actual ClientHello bytes from our saved state + let client_hello_bytes = match &self.handshake_state { + HandshakeState::AwaitingServerHello { + client_hello_bytes, .. + } => client_hello_bytes.clone(), + _ => { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "Invalid handshake state", + )) + } + }; + + let mut full_transcript = digest::Context::new(&digest::SHA256); + log::debug!( + "REALITY CLIENT: Transcript includes ClientHello ({} bytes), first bytes: {:02x?}", + client_hello_bytes.len(), + &client_hello_bytes[..client_hello_bytes.len().min(20)] + ); + log::debug!( + "REALITY CLIENT: Transcript includes ServerHello ({} bytes), first bytes: {:02x?}", + server_hello.len(), + &server_hello[..server_hello.len().min(20)] + ); + full_transcript.update(&client_hello_bytes); // Use actual ClientHello bytes, not hash! + full_transcript.update(server_hello); // ServerHello already includes handshake header + let server_hello_hash = full_transcript.finish(); + let mut server_hello_hash_arr = [0u8; 32]; + server_hello_hash_arr.copy_from_slice(server_hello_hash.as_ref()); + + // Perform ECDH for TLS 1.3 key derivation + let peer_public_key = + agreement::UnparsedPublicKey::new(&agreement::X25519, &server_public_key); + let my_private_key = + agreement::PrivateKey::from_private_key(&agreement::X25519, &client_private_key) + .map_err(|_| io::Error::other("Failed to create private key"))?; + + let mut tls_shared_secret = [0u8; 32]; + agreement::agree( + &my_private_key, + peer_public_key, + io::Error::other("ECDH failed"), + |key_material| { + tls_shared_secret.copy_from_slice(key_material); + Ok(()) + }, + )?; + + // Derive handshake keys + // Removed verbose secret logging for security + let hs_keys = derive_handshake_keys( + &tls_shared_secret, + &client_hello_hash, + &server_hello_hash_arr, + )?; + + // Use standard cipher suite + const CIPHER_SUITE: u16 = 0x1301; // TLS_AES_128_GCM_SHA256 + + log::debug!("REALITY: ServerHello processed, handshake keys derived"); + + // Initialize transcript with actual ClientHello and ServerHello bytes (not hashes!) + let mut transcript_bytes = Vec::new(); + transcript_bytes.extend_from_slice(&client_hello_bytes); // Actual ClientHello bytes + transcript_bytes.extend_from_slice(server_hello); // Actual ServerHello bytes + + // Update state + self.handshake_state = HandshakeState::ProcessingHandshake { + client_handshake_traffic_secret: hs_keys.client_handshake_traffic_secret.clone(), + server_handshake_traffic_secret: hs_keys.server_handshake_traffic_secret.clone(), + master_secret: hs_keys.master_secret.clone(), + cipher_suite: CIPHER_SUITE, + handshake_transcript_bytes: transcript_bytes, + auth_key, // Pass auth_key for certificate HMAC verification + }; + + Ok(()) + } + + /// Process encrypted handshake messages (EncryptedExtensions, Certificate, CertificateVerify, Finished) + fn process_encrypted_handshake(&mut self) -> io::Result<()> { + // Extract state - we need to preserve it across multiple calls + let ( + client_hs_secret, + server_hs_secret, + master_secret, + cipher_suite, + transcript_bytes, + auth_key, + ) = match &self.handshake_state { + HandshakeState::ProcessingHandshake { + client_handshake_traffic_secret, + server_handshake_traffic_secret, + master_secret, + cipher_suite, + handshake_transcript_bytes, + auth_key, + } => ( + client_handshake_traffic_secret.clone(), + server_handshake_traffic_secret.clone(), + master_secret.clone(), + *cipher_suite, + handshake_transcript_bytes.clone(), + *auth_key, + ), + _ => return Ok(()), + }; + + // Derive server handshake traffic keys for decryption + let (server_hs_key, server_hs_iv) = derive_traffic_keys(&server_hs_secret, cipher_suite)?; + + log::debug!( + "REALITY CLIENT: Server HS key={:02x?}, iv={:02x?}", + &server_hs_key[..16], + &server_hs_iv + ); + + // Check if we have enough data for a TLS record header + if self.ciphertext_read_buf.len() < TLS_RECORD_HEADER_SIZE { + return Ok(()); // Need more data + } + + // Check record type + let record_type = self.ciphertext_read_buf[0]; + let tls_version = self + .ciphertext_read_buf + .get_u16_be(1) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "Buffer too short"))?; + let record_len = self + .ciphertext_read_buf + .get_u16_be(3) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "Buffer too short"))? + as usize; + + log::debug!( + "REALITY CLIENT: TLS record header: type=0x{:02x}, version=0x{:04x}, len={}", + record_type, + tls_version, + record_len + ); + + // Check if we have the complete record + let total_record_len = TLS_RECORD_HEADER_SIZE + record_len; + if self.ciphertext_read_buf.len() < total_record_len { + return Ok(()); // Need more data + } + + // Skip ChangeCipherSpec records - these are dummy in TLS 1.3 + if record_type == CONTENT_TYPE_CHANGE_CIPHER_SPEC { + log::debug!( + "REALITY CLIENT: Skipping ChangeCipherSpec record ({} bytes)", + record_len + ); + self.ciphertext_read_buf.consume(total_record_len); + // Process next record recursively + return self.process_encrypted_handshake(); + } + + // If it's not Application Data, we have a problem + if record_type != CONTENT_TYPE_APPLICATION_DATA { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "Expected Application Data record, got 0x{:02x}", + record_type + ), + )); + } + + // Copy and extract the combined handshake record + let ciphertext: Vec = + self.ciphertext_read_buf[TLS_RECORD_HEADER_SIZE..total_record_len].to_vec(); + self.ciphertext_read_buf.consume(total_record_len); + + log::debug!("REALITY CLIENT: Decrypting combined handshake record - record_type=0x{:02x}, record_len={}, seq=0", + record_type, record_len); + + // Try different sequence numbers - sing-box might count differently + let mut combined_plaintext = None; + for seq in 0..3 { + match decrypt_handshake_message( + &server_hs_key, + &server_hs_iv, + seq, + &ciphertext, + record_len as u16, + ) { + Ok(plaintext) => { + combined_plaintext = Some(plaintext); + break; + } + Err(_) if seq < 2 => { + log::debug!( + "REALITY CLIENT: Decryption failed with seq={}, trying next", + seq + ); + continue; + } + Err(e) => { + log::error!("REALITY CLIENT: Failed to decrypt combined handshake with all sequence numbers: {}", e); + return Err(e); + } + } + } + + let combined_plaintext = combined_plaintext.unwrap(); + + log::info!( + "REALITY CLIENT: Successfully decrypted combined handshake record ({} bytes plaintext)", + combined_plaintext.len() + ); + + // Now parse the individual handshake messages from the combined plaintext + let mut offset = 0; + let mut messages_found = 0; + let mut certificate_verified = false; + + while offset < combined_plaintext.len() && messages_found < 4 { + // Each handshake message has: type (1 byte) + length (3 bytes) + data + if offset + 4 > combined_plaintext.len() { + break; + } + + let msg_type = combined_plaintext[offset]; + let msg_len = u32::from_be_bytes([ + 0, + combined_plaintext[offset + 1], + combined_plaintext[offset + 2], + combined_plaintext[offset + 3], + ]) as usize; + + if offset + 4 + msg_len > combined_plaintext.len() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "Invalid handshake message length: {} at offset {}", + msg_len, offset + ), + )); + } + + let msg_name = match msg_type { + HANDSHAKE_TYPE_ENCRYPTED_EXTENSIONS => "EncryptedExtensions", + HANDSHAKE_TYPE_CERTIFICATE => "Certificate", + HANDSHAKE_TYPE_CERTIFICATE_VERIFY => "CertificateVerify", + HANDSHAKE_TYPE_FINISHED => "Finished", + _ => "Unknown", + }; + + log::info!( + "REALITY CLIENT: Found {} message (type={}, len={})", + msg_name, + msg_type, + msg_len + ); + + // Verify HMAC signature when we encounter the Certificate message + if msg_type == HANDSHAKE_TYPE_CERTIFICATE { + let cert_der = + extract_certificate_der(&combined_plaintext[offset..offset + 4 + msg_len])?; + verify_certificate_hmac(cert_der, &auth_key)?; + certificate_verified = true; + } + + messages_found += 1; + offset += 4 + msg_len; + } + + // Check if we got all 4 messages AND the certificate was verified + if messages_found != 4 { + // Not all messages received yet - this shouldn't happen with combined record + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Expected 4 handshake messages, found {}", messages_found), + )); + } + + // Ensure the Certificate message was present and HMAC verified + if !certificate_verified { + return Err(io::Error::new( + io::ErrorKind::PermissionDenied, + "REALITY handshake failed: Certificate message not received or not verified", + )); + } + + // Build handshake transcript + // We need the actual handshake message bytes, not their hashes! + // The transcript_bytes already contains the actual ClientHello and ServerHello bytes + let mut handshake_transcript = digest::Context::new(&digest::SHA256); + handshake_transcript.update(&transcript_bytes); // Contains actual ClientHello + ServerHello bytes + handshake_transcript.update(&combined_plaintext); // EncryptedExtensions + Certificate + CertificateVerify + Finished + + let handshake_hash = handshake_transcript.finish(); + let mut handshake_hash_arr = [0u8; 32]; + handshake_hash_arr.copy_from_slice(handshake_hash.as_ref()); + + log::info!( + "REALITY CLIENT: Handshake hash for client Finished: {:02x?}", + handshake_hash_arr + ); + log::info!( + "REALITY CLIENT: Transcript bytes len={}, combined_plaintext len={}", + transcript_bytes.len(), + combined_plaintext.len() + ); + + // Generate client Finished message + let client_verify_data = + compute_finished_verify_data(&client_hs_secret, &handshake_hash_arr)?; + log::info!( + "REALITY CLIENT: Client verify data: {:02x?}", + client_verify_data + ); + let client_finished = construct_finished(&client_verify_data)?; + + // Derive client handshake traffic keys for encryption + let (client_hs_key, client_hs_iv) = derive_traffic_keys(&client_hs_secret, cipher_suite)?; + + // Encrypt Finished message + let mut client_hs_seq = 0u64; + let buf_len_before = self.ciphertext_write_buf.len(); + encrypt_handshake_to_records( + &client_finished, + &client_hs_key, + &client_hs_iv, + &mut client_hs_seq, + &mut self.ciphertext_write_buf, + )?; + + log::info!( + "REALITY CLIENT: Client Finished message generated and buffered ({} bytes)", + self.ciphertext_write_buf.len() - buf_len_before + ); + + // Derive application secrets + let (client_app_secret, server_app_secret) = + derive_application_secrets(&master_secret, &handshake_hash_arr)?; + + // Derive application traffic keys + let (client_app_key, client_app_iv) = + derive_traffic_keys(&client_app_secret, cipher_suite)?; + let (server_app_key, server_app_iv) = + derive_traffic_keys(&server_app_secret, cipher_suite)?; + + // Store application keys + self.app_read_key = Some(server_app_key); + self.app_read_iv = Some(server_app_iv); + self.app_write_key = Some(client_app_key); + self.app_write_iv = Some(client_app_iv); + self.read_seq = 0; + self.write_seq = 0; + self.cipher_suite = cipher_suite; + + // Mark handshake complete + self.handshake_state = HandshakeState::Complete; + log::info!("REALITY CLIENT: Handshake complete, application keys derived"); + + Ok(()) + } + + /// Decrypt application data using TLS 1.3 keys + fn process_application_data(&mut self) -> io::Result<()> { + let (app_read_key, app_read_iv) = match (&self.app_read_key, &self.app_read_iv) { + (Some(key), Some(iv)) => (key, iv), + _ => return Ok(()), + }; + + while self.ciphertext_read_buf.len() >= TLS_RECORD_HEADER_SIZE { + let record_len = self + .ciphertext_read_buf + .get_u16_be(3) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "Buffer too short"))? + as usize; + + let total_record_len = TLS_RECORD_HEADER_SIZE + record_len; + if self.ciphertext_read_buf.len() < total_record_len { + break; + } + + // Copy record header for decryption AAD + let tls_header: [u8; TLS_RECORD_HEADER_SIZE] = [ + self.ciphertext_read_buf[0], + self.ciphertext_read_buf[1], + self.ciphertext_read_buf[2], + self.ciphertext_read_buf[3], + self.ciphertext_read_buf[4], + ]; + let ciphertext: Vec = + self.ciphertext_read_buf[TLS_RECORD_HEADER_SIZE..total_record_len].to_vec(); + self.ciphertext_read_buf.consume(total_record_len); + + let mut plaintext = decrypt_tls13_record( + app_read_key, + app_read_iv, + self.read_seq, + &ciphertext, + &tls_header, + )?; + + self.read_seq += 1; + + // TLS 1.3: Remove ContentType trailer byte + // The last byte of the plaintext is the actual ContentType + if !plaintext.is_empty() { + let content_type = plaintext.pop().unwrap(); + + // For now, we only handle ApplicationData and Alert + // Alerts should be handled separately, but for now just pass through + if content_type != CONTENT_TYPE_APPLICATION_DATA + && content_type != CONTENT_TYPE_ALERT + { + log::warn!( + "REALITY CLIENT: Unexpected ContentType: 0x{:02x}", + content_type + ); + } + } + + // Compact plaintext buffer if needed before extending + self.plaintext_read_buf.maybe_compact(4096); + self.plaintext_read_buf.extend_from_slice(&plaintext); + } + + Ok(()) + } + + /// Get a reader for accessing decrypted plaintext + pub fn reader(&mut self) -> RealityReader<'_> { + // SlideBuffer handles compaction internally via maybe_compact() + // Compact before returning reader if we've consumed significant data + self.plaintext_read_buf.maybe_compact(4096); + RealityReader::new(&mut self.plaintext_read_buf) + } + + /// Get a writer for buffering plaintext to be encrypted + pub fn writer(&mut self) -> RealityWriter<'_> { + RealityWriter::new(&mut self.plaintext_write_buf) + } + + /// Write buffered TLS messages to the provided writer + /// + /// Large plaintext is automatically fragmented into multiple TLS records + /// to comply with the TLS 1.3 record size limit. + pub fn write_tls(&mut self, wr: &mut dyn Write) -> io::Result { + // If handshake not complete, just write buffered handshake data + if !matches!(self.handshake_state, HandshakeState::Complete) { + let n = wr.write(&self.ciphertext_write_buf)?; + self.ciphertext_write_buf.drain(..n); + return Ok(n); + } + + // Encrypt any pending plaintext (with automatic fragmentation for large data) + if !self.plaintext_write_buf.is_empty() { + let (app_write_key, app_write_iv) = match (&self.app_write_key, &self.app_write_iv) { + (Some(key), Some(iv)) => (key, iv), + _ => { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "Application keys not available", + )) + } + }; + + encrypt_plaintext_to_records( + &mut self.plaintext_write_buf, + app_write_key, + app_write_iv, + &mut self.write_seq, + &mut self.ciphertext_write_buf, + )?; + } + + let n = wr.write(&self.ciphertext_write_buf)?; + self.ciphertext_write_buf.drain(..n); + Ok(n) + } + + /// Check if the connection wants to write data + pub fn wants_write(&self) -> bool { + !self.ciphertext_write_buf.is_empty() || !self.plaintext_write_buf.is_empty() + } + + /// Check if handshake is still in progress + pub fn is_handshaking(&self) -> bool { + !matches!(self.handshake_state, HandshakeState::Complete) + } + + /// Queue a close notification alert + pub fn send_close_notify(&mut self) { + // In TLS 1.3, alerts must be encrypted like application data + if !matches!(self.handshake_state, HandshakeState::Complete) { + log::debug!("REALITY CLIENT: Cannot send close_notify - handshake not complete"); + return; + } + + // Get application keys + let (app_write_key, app_write_iv) = match (&self.app_write_key, &self.app_write_iv) { + (Some(key), Some(iv)) => (key, iv), + _ => { + log::debug!( + "REALITY CLIENT: Cannot send close_notify - application keys not available" + ); + return; + } + }; + + // Use common helper to build encrypted close_notify alert + match common::build_close_notify_alert(app_write_key, app_write_iv, self.write_seq) { + Ok(record) => { + self.write_seq += 1; + self.ciphertext_write_buf.extend_from_slice(&record); + log::debug!("REALITY CLIENT: Encrypted close_notify alert queued"); + } + Err(e) => { + log::error!("REALITY CLIENT: Failed to encrypt close_notify: {}", e); + } + } + } +} + +#[inline(always)] +pub fn feed_reality_client_connection( + client_connection: &mut RealityClientConnection, + data: &[u8], +) -> std::io::Result<()> { + let mut cursor = std::io::Cursor::new(data); + let mut i = 0; + while i < data.len() { + let n = client_connection.read_tls(&mut cursor).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("failed to feed TLS connection: {e}"), + ) + })?; + i += n; + } + Ok(()) +} diff --git a/src/reality/reality_client_verify.rs b/src/reality/reality_client_verify.rs new file mode 100644 index 00000000..f3bce4a5 --- /dev/null +++ b/src/reality/reality_client_verify.rs @@ -0,0 +1,314 @@ +// REALITY client certificate verification +// +// This module handles HMAC verification of REALITY server certificates. +// In REALITY protocol, the server embeds HMAC-SHA512(auth_key, ed25519_public_key) +// in the signature field of the certificate. + +use std::io; + +use aws_lc_rs::hmac; +use subtle::ConstantTimeEq; +use x509_parser::certificate::X509Certificate; +use x509_parser::prelude::FromDer; + +/// Extract the DER-encoded certificate from a TLS 1.3 Certificate message +/// +/// Certificate message structure: +/// - certificate_request_context (1 byte length + data) +/// - certificate_list length (3 bytes) +/// - For each certificate entry: +/// - cert_data length (3 bytes) +/// - cert_data (DER-encoded X.509 certificate) +/// - extensions length (2 bytes) +/// - extensions data +/// +#[inline] +pub fn extract_certificate_der(certificate_message: &[u8]) -> io::Result<&[u8]> { + // Skip handshake header (type + 3-byte length) + if certificate_message.len() < 4 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "Certificate message too short", + )); + } + + let mut pos = 4; // Skip handshake header + + // certificate_request_context length (1 byte) + if pos >= certificate_message.len() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "Certificate message truncated at context length", + )); + } + let context_len = certificate_message[pos] as usize; + pos += 1 + context_len; + + // certificate_list length (3 bytes) + if pos + 3 > certificate_message.len() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "Certificate message truncated at list length", + )); + } + let _list_len = u32::from_be_bytes([ + 0, + certificate_message[pos], + certificate_message[pos + 1], + certificate_message[pos + 2], + ]) as usize; + pos += 3; + + // First certificate entry: cert_data length (3 bytes) + if pos + 3 > certificate_message.len() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "Certificate message truncated at cert length", + )); + } + let cert_len = u32::from_be_bytes([ + 0, + certificate_message[pos], + certificate_message[pos + 1], + certificate_message[pos + 2], + ]) as usize; + pos += 3; + + // Extract the DER-encoded certificate + if pos + cert_len > certificate_message.len() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "Certificate message truncated at cert data", + )); + } + + Ok(&certificate_message[pos..pos + cert_len]) +} + +/// Verify the HMAC signature embedded in the REALITY certificate +/// +/// In REALITY protocol, the server embeds HMAC-SHA512(auth_key, ed25519_public_key) +/// in the signature field of the certificate. We compare the first 32 bytes. +/// +/// Uses proper X.509 parsing via x509-parser crate for robust extraction. +#[inline] +pub fn verify_certificate_hmac(cert_der: &[u8], auth_key: &[u8; 32]) -> io::Result<()> { + // Parse the X.509 certificate properly + let (_, cert) = X509Certificate::from_der(cert_der).map_err(|e| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("Failed to parse X.509 certificate: {}", e), + ) + })?; + + // Extract the public key from SubjectPublicKeyInfo + let spki = cert.public_key(); + let pubkey_data: &[u8] = &spki.subject_public_key.data; + let signature: &[u8] = &cert.signature_value.data; + + log::debug!( + "REALITY CLIENT: Parsed certificate - pubkey len={}, sig len={}", + pubkey_data.len(), + signature.len() + ); + + // Verify this is an Ed25519 public key (32 bytes) + if pubkey_data.len() != 32 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "Expected Ed25519 public key (32 bytes), got {} bytes", + pubkey_data.len() + ), + )); + } + + // Verify signature is long enough (should be 64 bytes for Ed25519, containing HMAC-SHA512) + if signature.len() < 32 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "Signature too short for HMAC verification: {} bytes", + signature.len() + ), + )); + } + + // Compute expected HMAC: HMAC-SHA512(auth_key, ed25519_public_key) + let hmac_key = hmac::Key::new(hmac::HMAC_SHA512, auth_key); + let hmac_tag = hmac::sign(&hmac_key, pubkey_data); + let expected_signature = &hmac_tag.as_ref()[..32]; // First 32 bytes + + log::debug!( + "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[..32] + ); + + // Compare first 32 bytes of signature with expected HMAC using constant-time comparison + // to prevent timing attacks that could leak information about the expected signature + if expected_signature.ct_eq(&signature[..32]).unwrap_u8() == 0 { + return Err(io::Error::new( + io::ErrorKind::PermissionDenied, + "Certificate HMAC verification failed - signature mismatch", + )); + } + + log::info!("REALITY CLIENT: Certificate HMAC verified successfully"); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use aws_lc_rs::signature::KeyPair; + + /// Test certificate extraction with a properly formatted TLS 1.3 Certificate message + #[test] + fn test_extract_certificate_der_valid() { + // Build a minimal TLS 1.3 Certificate message: + // Handshake header: type (0x0b) + 3-byte length + // certificate_request_context: 1 byte length (0) + no data + // certificate_list: 3 byte length + entries + // First entry: 3 byte cert length + cert data + 2 byte extensions length (0) + + let cert_data = b"fake_certificate_der_data"; + let cert_len = cert_data.len(); + + // Build certificate entry: cert_len (3 bytes) + cert + extensions_len (2 bytes) + let entry_len = 3 + cert_len + 2; + + // Build certificate_list: list_len (3 bytes) + entry + let list_len = entry_len; + + // Build message body: context_len (1) + list_len (3) + entry + let body_len = 1 + 3 + entry_len; + + let mut message = Vec::new(); + // Handshake type (Certificate = 0x0b) + message.push(0x0b); + // Handshake length (3 bytes, big-endian) + message.push(((body_len >> 16) & 0xff) as u8); + message.push(((body_len >> 8) & 0xff) as u8); + message.push((body_len & 0xff) as u8); + // certificate_request_context length = 0 + message.push(0x00); + // certificate_list length (3 bytes) + message.push(((list_len >> 16) & 0xff) as u8); + message.push(((list_len >> 8) & 0xff) as u8); + message.push((list_len & 0xff) as u8); + // First certificate: cert_data length (3 bytes) + message.push(((cert_len >> 16) & 0xff) as u8); + message.push(((cert_len >> 8) & 0xff) as u8); + message.push((cert_len & 0xff) as u8); + // cert_data + message.extend_from_slice(cert_data); + // extensions length = 0 + message.push(0x00); + message.push(0x00); + + let result = extract_certificate_der(&message); + assert!(result.is_ok()); + assert_eq!(result.unwrap(), cert_data); + } + + /// Test certificate extraction with message too short + #[test] + fn test_extract_certificate_der_too_short() { + let message = vec![0x0b, 0x00, 0x00]; // Only 3 bytes, need at least 4 + let result = extract_certificate_der(&message); + assert!(result.is_err()); + assert_eq!(result.unwrap_err().kind(), io::ErrorKind::InvalidData); + } + + /// Test certificate extraction with truncated context + #[test] + fn test_extract_certificate_der_truncated_context() { + // Handshake header but no body + let message = vec![0x0b, 0x00, 0x00, 0x10]; // Claims 16 bytes but none present + let result = extract_certificate_der(&message); + assert!(result.is_err()); + } + + /// Test HMAC verification with a real Ed25519 certificate + #[test] + fn test_verify_certificate_hmac_with_real_cert() { + // Generate a real Ed25519 certificate using rcgen + let key_pair = rcgen::KeyPair::generate_for(&rcgen::PKCS_ED25519) + .expect("Failed to generate Ed25519 key pair"); + + let params = rcgen::CertificateParams::new(vec!["test.example.com".to_string()]) + .expect("Failed to create certificate params"); + + let cert = params + .self_signed(&key_pair) + .expect("Failed to create self-signed certificate"); + + let mut cert_der = cert.der().to_vec(); + + // Get the public key + let signing_key = + aws_lc_rs::signature::Ed25519KeyPair::from_pkcs8(key_pair.serialized_der()) + .expect("Failed to parse key"); + let public_key_bytes = signing_key.public_key().as_ref(); + + // Parse the certificate to find where the signature actually is + let (_, parsed_cert) = + X509Certificate::from_der(&cert_der).expect("Failed to parse certificate"); + let sig_offset = + parsed_cert.signature_value.data.as_ptr() as usize - cert_der.as_ptr() as usize; + + // Create auth_key and compute HMAC + let auth_key = [0x42u8; 32]; + let hmac_key = hmac::Key::new(hmac::HMAC_SHA512, &auth_key); + let hmac_tag = hmac::sign(&hmac_key, public_key_bytes); + + // Replace the signature bytes at the correct offset + cert_der[sig_offset..sig_offset + 64].copy_from_slice(hmac_tag.as_ref()); + + // Now verify + let result = verify_certificate_hmac(&cert_der, &auth_key); + assert!(result.is_ok(), "HMAC verification should succeed"); + } + + /// Test HMAC verification with invalid signature + #[test] + fn test_verify_certificate_hmac_invalid_signature() { + // Generate a real Ed25519 certificate but don't modify the signature + let key_pair = rcgen::KeyPair::generate_for(&rcgen::PKCS_ED25519) + .expect("Failed to generate Ed25519 key pair"); + + let params = rcgen::CertificateParams::new(vec!["test.example.com".to_string()]) + .expect("Failed to create certificate params"); + + let cert = params + .self_signed(&key_pair) + .expect("Failed to create self-signed certificate"); + + let cert_der = cert.der().to_vec(); + let auth_key = [0x42u8; 32]; + + // This should fail because the signature is a real Ed25519 signature, not HMAC + let result = verify_certificate_hmac(&cert_der, &auth_key); + assert!(result.is_err()); + assert_eq!(result.unwrap_err().kind(), io::ErrorKind::PermissionDenied); + } + + /// Test HMAC verification with invalid certificate DER + #[test] + fn test_verify_certificate_hmac_invalid_der() { + let auth_key = [0x11u8; 32]; + let invalid_der = vec![0x30, 0x00]; // Empty SEQUENCE, not a valid certificate + + let result = verify_certificate_hmac(&invalid_der, &auth_key); + assert!(result.is_err()); + assert_eq!(result.unwrap_err().kind(), io::ErrorKind::InvalidData); + } +} diff --git a/src/reality/reality_io_state.rs b/src/reality/reality_io_state.rs new file mode 100644 index 00000000..35ba5140 --- /dev/null +++ b/src/reality/reality_io_state.rs @@ -0,0 +1,20 @@ +/// Represents the I/O state after processing packets +#[derive(Debug, Clone, Copy)] +pub struct RealityIoState { + /// Number of plaintext bytes available to read + plaintext_bytes_to_read: usize, +} + +impl RealityIoState { + /// Create a new RealityIoState + pub fn new(plaintext_bytes_to_read: usize) -> Self { + Self { + plaintext_bytes_to_read, + } + } + + /// How many plaintext bytes could be obtained via Read without further I/O + pub fn plaintext_bytes_to_read(&self) -> usize { + self.plaintext_bytes_to_read + } +} diff --git a/src/reality/reality_reader_writer.rs b/src/reality/reality_reader_writer.rs new file mode 100644 index 00000000..4e61ee2a --- /dev/null +++ b/src/reality/reality_reader_writer.rs @@ -0,0 +1,117 @@ +use std::io::{BufRead, Read, Write}; + +use crate::slide_buffer::SlideBuffer; + +/// Reader for accessing decrypted plaintext from REALITY connections +/// +/// This reader provides a view over a SlideBuffer and consumes data from it +/// as it is read. The SlideBuffer handles efficient memory management. +pub struct RealityReader<'a> { + buffer: &'a mut SlideBuffer, +} + +impl<'a> RealityReader<'a> { + pub fn new(buffer: &'a mut SlideBuffer) -> Self { + RealityReader { buffer } + } +} + +impl<'a> Read for RealityReader<'a> { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + let available = self.buffer.len(); + let to_read = buf.len().min(available); + if to_read > 0 { + buf[..to_read].copy_from_slice(&self.buffer[..to_read]); + self.buffer.consume(to_read); + } + Ok(to_read) + } +} + +impl<'a> BufRead for RealityReader<'a> { + fn fill_buf(&mut self) -> std::io::Result<&[u8]> { + Ok(self.buffer.as_slice()) + } + + fn consume(&mut self, amt: usize) { + let actual = amt.min(self.buffer.len()); + self.buffer.consume(actual); + } +} + +/// Writer for buffering plaintext to be encrypted in REALITY connections +pub struct RealityWriter<'a> { + buffer: &'a mut Vec, +} + +impl<'a> RealityWriter<'a> { + pub fn new(buffer: &'a mut Vec) -> Self { + RealityWriter { buffer } + } +} + +impl<'a> Write for RealityWriter<'a> { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.buffer.extend_from_slice(buf); + Ok(buf.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_reality_crypto_reader() { + let mut buffer = SlideBuffer::new(1024); + buffer.extend_from_slice(b"hello world"); + + { + let mut reader = RealityReader::new(&mut buffer); + let mut buf = [0u8; 5]; + assert_eq!(reader.read(&mut buf).unwrap(), 5); + assert_eq!(&buf, b"hello"); + } + + // Buffer should have consumed 5 bytes + assert_eq!(buffer.len(), 6); // " world" remaining + + { + let mut reader = RealityReader::new(&mut buffer); + let mut buf = [0u8; 5]; + assert_eq!(reader.read(&mut buf).unwrap(), 5); + assert_eq!(&buf[..5], b" worl"); + } + + assert_eq!(buffer.len(), 1); // "d" remaining + } + + #[test] + fn test_reality_crypto_reader_bufread() { + let mut buffer = SlideBuffer::new(1024); + buffer.extend_from_slice(b"test data"); + + let mut reader = RealityReader::new(&mut buffer); + + let buf = reader.fill_buf().unwrap(); + assert_eq!(buf, b"test data"); + + reader.consume(5); + let buf = reader.fill_buf().unwrap(); + assert_eq!(buf, b"data"); + } + + #[test] + fn test_reality_crypto_writer() { + let mut buffer = Vec::new(); + let mut writer = RealityWriter::new(&mut buffer); + + assert_eq!(writer.write(b"hello").unwrap(), 5); + assert_eq!(writer.write(b" world").unwrap(), 6); + assert_eq!(buffer.as_slice(), b"hello world"); + } +} diff --git a/src/reality/reality_records.rs b/src/reality/reality_records.rs new file mode 100644 index 00000000..f116440d --- /dev/null +++ b/src/reality/reality_records.rs @@ -0,0 +1,1594 @@ +// TLS 1.3 record encryption with automatic fragmentation +// +// Handles encrypting plaintext into TLS records, automatically splitting +// large data into multiple records to stay within TLS 1.3 size limits. + +use std::io; + +use super::common::{ + CONTENT_TYPE_APPLICATION_DATA, CONTENT_TYPE_HANDSHAKE, MAX_TLS_CIPHERTEXT_LEN, + MAX_TLS_PLAINTEXT_LEN, TLS_RECORD_HEADER_SIZE, +}; +use super::reality_aead::encrypt_tls13_record; + +/// Encrypt plaintext into TLS 1.3 application data records, fragmenting if necessary. +/// +/// This function handles the TLS 1.3 record size limit by splitting large plaintext +/// into multiple records of at most `MAX_TLS_PLAINTEXT_LEN` bytes each. +/// +/// # Arguments +/// * `plaintext` - The plaintext data to encrypt (will be cleared after encryption) +/// * `app_write_key` - The application traffic key (16 bytes for AES-128-GCM) +/// * `app_write_iv` - The application traffic IV (12 bytes) +/// * `write_seq` - Mutable reference to the write sequence number (incremented per record) +/// * `ciphertext_buf` - Buffer to append encrypted TLS records to +/// +/// # Returns +/// * `Ok(())` on success +/// * `Err` if encryption fails +#[inline] +pub fn encrypt_plaintext_to_records( + plaintext: &mut Vec, + app_write_key: &[u8], + app_write_iv: &[u8], + write_seq: &mut u64, + ciphertext_buf: &mut Vec, +) -> io::Result<()> { + if plaintext.is_empty() { + return Ok(()); + } + + // Fast path: single record (most common case, ~16KB or less) + if plaintext.len() <= MAX_TLS_PLAINTEXT_LEN { + encrypt_single_record( + plaintext, + app_write_key, + app_write_iv, + write_seq, + ciphertext_buf, + )?; + plaintext.clear(); + return Ok(()); + } + + // Slow path: fragment into multiple records + let total_len = plaintext.len(); + let num_records = total_len.div_ceil(MAX_TLS_PLAINTEXT_LEN); + log::debug!( + "REALITY: Fragmenting {} bytes into {} TLS records (max {} bytes/record)", + total_len, + num_records, + MAX_TLS_PLAINTEXT_LEN + ); + + let mut offset = 0; + while offset < plaintext.len() { + let chunk_end = (offset + MAX_TLS_PLAINTEXT_LEN).min(plaintext.len()); + let chunk = &plaintext[offset..chunk_end]; + + encrypt_chunk_with_type( + chunk, + app_write_key, + app_write_iv, + write_seq, + ciphertext_buf, + CONTENT_TYPE_APPLICATION_DATA, + )?; + offset = chunk_end; + } + + plaintext.clear(); + Ok(()) +} + +/// Encrypt a single record where plaintext fits within MAX_TLS_PLAINTEXT_LEN. +/// This is the fast path - we can use mem::take to avoid allocation for plaintext_with_type. +#[inline] +fn encrypt_single_record( + plaintext: &mut Vec, + app_write_key: &[u8], + app_write_iv: &[u8], + write_seq: &mut u64, + ciphertext_buf: &mut Vec, +) -> io::Result<()> { + // TLS 1.3: plaintext_with_type = plaintext || ContentType + // Take ownership to avoid allocation, we'll restore and clear it + let mut plaintext_with_type = std::mem::take(plaintext); + plaintext_with_type.push(CONTENT_TYPE_APPLICATION_DATA); + + // Ciphertext length = plaintext_with_type + GCM tag (16 bytes) + let ciphertext_len = (plaintext_with_type.len() + 16) as u16; + + // Build TLS record header (also serves as AAD for AEAD) + let tls_header: [u8; TLS_RECORD_HEADER_SIZE] = [ + CONTENT_TYPE_APPLICATION_DATA, + 0x03, + 0x03, // TLS 1.2 version for compatibility + (ciphertext_len >> 8) as u8, + (ciphertext_len & 0xff) as u8, + ]; + + let ciphertext = encrypt_tls13_record( + app_write_key, + app_write_iv, + *write_seq, + &plaintext_with_type, + &tls_header, + )?; + + *write_seq += 1; + + // Append TLS record to output buffer + ciphertext_buf.reserve(TLS_RECORD_HEADER_SIZE + ciphertext.len()); + ciphertext_buf.extend_from_slice(&tls_header); + ciphertext_buf.extend_from_slice(&ciphertext); + + // Restore the vec for reuse by caller (will be cleared by caller) + *plaintext = plaintext_with_type; + + Ok(()) +} + +/// Encrypt a chunk of data into a single TLS 1.3 record. +/// +/// This is the core encryption function used by both application data and handshake +/// fragmentation. The `inner_content_type` specifies the content type byte that gets +/// encrypted inside the record (0x17 for app data, 0x16 for handshake). +/// +/// The outer record type is always 0x17 (ApplicationData) for TLS 1.3 encrypted records. +#[inline] +fn encrypt_chunk_with_type( + chunk: &[u8], + key: &[u8], + iv: &[u8], + write_seq: &mut u64, + ciphertext_buf: &mut Vec, + inner_content_type: u8, +) -> io::Result<()> { + // TLS 1.3: plaintext_with_type = chunk || ContentType + let mut plaintext_with_type = Vec::with_capacity(chunk.len() + 1); + plaintext_with_type.extend_from_slice(chunk); + plaintext_with_type.push(inner_content_type); + + // Ciphertext length = plaintext_with_type + GCM tag (16 bytes) + let ciphertext_len = (plaintext_with_type.len() + 16) as u16; + + // Verify we don't exceed TLS 1.3 limit (this should never happen if callers fragment correctly) + debug_assert!( + (ciphertext_len as usize) <= MAX_TLS_CIPHERTEXT_LEN, + "BUG: ciphertext_len {} exceeds MAX_TLS_CIPHERTEXT_LEN {}", + ciphertext_len, + MAX_TLS_CIPHERTEXT_LEN + ); + + // Build TLS record header (outer type is always ApplicationData for encrypted records) + let tls_header: [u8; TLS_RECORD_HEADER_SIZE] = [ + CONTENT_TYPE_APPLICATION_DATA, + 0x03, + 0x03, // TLS 1.2 version for compatibility + (ciphertext_len >> 8) as u8, + (ciphertext_len & 0xff) as u8, + ]; + + let ciphertext = encrypt_tls13_record(key, iv, *write_seq, &plaintext_with_type, &tls_header)?; + + *write_seq += 1; + + // Append TLS record to output buffer + ciphertext_buf.reserve(TLS_RECORD_HEADER_SIZE + ciphertext.len()); + ciphertext_buf.extend_from_slice(&tls_header); + ciphertext_buf.extend_from_slice(&ciphertext); + + Ok(()) +} + +/// Encrypt handshake data into TLS 1.3 records, fragmenting if necessary. +/// +/// This function handles the TLS 1.3 record size limit by splitting large handshake +/// data into multiple records of at most `MAX_TLS_PLAINTEXT_LEN` bytes each. +/// +/// Unlike application data encryption, handshake records use content type 0x16 (handshake) +/// inside the encrypted payload, though the outer record type is still 0x17 (application_data) +/// as per TLS 1.3 encrypted record format. +/// +/// # Arguments +/// * `handshake_data` - The combined handshake messages to encrypt +/// * `key` - The handshake traffic key (16 bytes for AES-128-GCM) +/// * `iv` - The handshake traffic IV (12 bytes) +/// * `write_seq` - Mutable reference to the write sequence number (incremented per record) +/// * `ciphertext_buf` - Buffer to append encrypted TLS records to +/// +/// # Returns +/// * `Ok(())` on success +/// * `Err` if encryption fails +#[inline] +pub fn encrypt_handshake_to_records( + handshake_data: &[u8], + key: &[u8], + iv: &[u8], + write_seq: &mut u64, + ciphertext_buf: &mut Vec, +) -> io::Result<()> { + if handshake_data.is_empty() { + return Ok(()); + } + + // Check if fragmentation is needed + if handshake_data.len() <= MAX_TLS_PLAINTEXT_LEN { + // Fast path: single record (most handshakes fit in one record) + encrypt_chunk_with_type( + handshake_data, + key, + iv, + write_seq, + ciphertext_buf, + CONTENT_TYPE_HANDSHAKE, + )?; + } else { + // Slow path: fragment into multiple records (large certificates) + let total_len = handshake_data.len(); + let num_records = total_len.div_ceil(MAX_TLS_PLAINTEXT_LEN); + log::debug!( + "REALITY: Fragmenting {} bytes of handshake data into {} TLS records (max {} bytes/record)", + total_len, + num_records, + MAX_TLS_PLAINTEXT_LEN + ); + + let mut offset = 0; + while offset < handshake_data.len() { + let chunk_end = (offset + MAX_TLS_PLAINTEXT_LEN).min(handshake_data.len()); + let chunk = &handshake_data[offset..chunk_end]; + + encrypt_chunk_with_type( + chunk, + key, + iv, + write_seq, + ciphertext_buf, + CONTENT_TYPE_HANDSHAKE, + )?; + offset = chunk_end; + } + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::reality::common::{CIPHERTEXT_READ_BUF_CAPACITY, TLS_MAX_RECORD_SIZE}; + + #[test] + fn test_constants() { + // Verify constants match expected TLS 1.3 values + assert_eq!(MAX_TLS_CIPHERTEXT_LEN, 16640); // TLS 1.3 limit: 16384 + 256 + assert_eq!(MAX_TLS_PLAINTEXT_LEN, 16384); // RFC 8446: 2^14 bytes max plaintext + assert_eq!(TLS_RECORD_HEADER_SIZE, 5); + assert_eq!(TLS_MAX_RECORD_SIZE, 16645); // 16640 + 5 + assert_eq!(CIPHERTEXT_READ_BUF_CAPACITY, 33290); // 2 * TLS_MAX_RECORD_SIZE + } + + #[test] + fn test_encrypt_empty_plaintext() { + let mut plaintext = Vec::new(); + let key = [0u8; 16]; + let iv = [0u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + let result = + encrypt_plaintext_to_records(&mut plaintext, &key, &iv, &mut seq, &mut ciphertext_buf); + assert!(result.is_ok()); + assert!(ciphertext_buf.is_empty()); + assert_eq!(seq, 0); // No records encrypted + } + + #[test] + fn test_encrypt_small_plaintext_single_record() { + let mut plaintext = vec![0x41u8; 100]; // 100 bytes of 'A' + let key = [0x01u8; 16]; + let iv = [0x02u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + let result = + encrypt_plaintext_to_records(&mut plaintext, &key, &iv, &mut seq, &mut ciphertext_buf); + assert!(result.is_ok()); + assert!(plaintext.is_empty()); // Should be cleared + + // Should have one record: header(5) + plaintext(100) + content_type(1) + tag(16) + let expected_len = TLS_RECORD_HEADER_SIZE + 100 + 1 + 16; + assert_eq!(ciphertext_buf.len(), expected_len); + assert_eq!(seq, 1); // One record encrypted + + // Verify TLS header + assert_eq!(ciphertext_buf[0], CONTENT_TYPE_APPLICATION_DATA); + assert_eq!(ciphertext_buf[1], 0x03); + assert_eq!(ciphertext_buf[2], 0x03); + let record_len = u16::from_be_bytes([ciphertext_buf[3], ciphertext_buf[4]]) as usize; + assert_eq!(record_len, 100 + 1 + 16); + } + + #[test] + fn test_encrypt_max_single_record() { + // Test exactly at the boundary - should still be single record + let mut plaintext = vec![0x42u8; MAX_TLS_PLAINTEXT_LEN]; + let key = [0x03u8; 16]; + let iv = [0x04u8; 12]; + let mut seq = 5u64; + let mut ciphertext_buf = Vec::new(); + + let result = + encrypt_plaintext_to_records(&mut plaintext, &key, &iv, &mut seq, &mut ciphertext_buf); + assert!(result.is_ok()); + assert!(plaintext.is_empty()); + + // Should have one record + let expected_len = TLS_RECORD_HEADER_SIZE + MAX_TLS_PLAINTEXT_LEN + 1 + 16; + assert_eq!(ciphertext_buf.len(), expected_len); + assert_eq!(seq, 6); // One record encrypted + } + + #[test] + fn test_encrypt_fragmentation_two_records() { + // One byte over the limit - should produce two records + let mut plaintext = vec![0x43u8; MAX_TLS_PLAINTEXT_LEN + 1]; + let key = [0x05u8; 16]; + let iv = [0x06u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + let result = + encrypt_plaintext_to_records(&mut plaintext, &key, &iv, &mut seq, &mut ciphertext_buf); + assert!(result.is_ok()); + assert!(plaintext.is_empty()); + assert_eq!(seq, 2); // Two records encrypted + + // First record: full MAX_TLS_PLAINTEXT_LEN + let first_record_len = TLS_RECORD_HEADER_SIZE + MAX_TLS_PLAINTEXT_LEN + 1 + 16; + // Second record: 1 byte + let second_record_len = TLS_RECORD_HEADER_SIZE + 1 + 1 + 16; + assert_eq!(ciphertext_buf.len(), first_record_len + second_record_len); + + // Verify first record header + assert_eq!(ciphertext_buf[0], CONTENT_TYPE_APPLICATION_DATA); + let first_payload_len = u16::from_be_bytes([ciphertext_buf[3], ciphertext_buf[4]]) as usize; + assert_eq!(first_payload_len, MAX_TLS_PLAINTEXT_LEN + 1 + 16); + + // Verify second record header + let second_header_offset = first_record_len; + assert_eq!( + ciphertext_buf[second_header_offset], + CONTENT_TYPE_APPLICATION_DATA + ); + let second_payload_len = u16::from_be_bytes([ + ciphertext_buf[second_header_offset + 3], + ciphertext_buf[second_header_offset + 4], + ]) as usize; + assert_eq!(second_payload_len, 1 + 1 + 16); + } + + #[test] + fn test_encrypt_fragmentation_multiple_records() { + // Test 3x the max size - should produce 3 records + let size = MAX_TLS_PLAINTEXT_LEN * 3; + let mut plaintext = vec![0x44u8; size]; + let key = [0x07u8; 16]; + let iv = [0x08u8; 12]; + let mut seq = 10u64; + let mut ciphertext_buf = Vec::new(); + + let result = + encrypt_plaintext_to_records(&mut plaintext, &key, &iv, &mut seq, &mut ciphertext_buf); + assert!(result.is_ok()); + assert!(plaintext.is_empty()); + assert_eq!(seq, 13); // Three records encrypted + + // Each record: header(5) + MAX_TLS_PLAINTEXT_LEN + content_type(1) + tag(16) + let record_len = TLS_RECORD_HEADER_SIZE + MAX_TLS_PLAINTEXT_LEN + 1 + 16; + assert_eq!(ciphertext_buf.len(), record_len * 3); + } + + #[test] + fn test_encrypt_fragmentation_uneven_split() { + // Test size that doesn't divide evenly + let size = MAX_TLS_PLAINTEXT_LEN * 2 + 1000; + let mut plaintext = vec![0x45u8; size]; + let key = [0x09u8; 16]; + let iv = [0x0Au8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + let result = + encrypt_plaintext_to_records(&mut plaintext, &key, &iv, &mut seq, &mut ciphertext_buf); + assert!(result.is_ok()); + assert!(plaintext.is_empty()); + assert_eq!(seq, 3); // Three records + + // Two full records + one partial + let full_record_len = TLS_RECORD_HEADER_SIZE + MAX_TLS_PLAINTEXT_LEN + 1 + 16; + let partial_record_len = TLS_RECORD_HEADER_SIZE + 1000 + 1 + 16; + assert_eq!( + ciphertext_buf.len(), + full_record_len * 2 + partial_record_len + ); + } + + #[test] + fn test_sequence_number_increments() { + let key = [0x0Bu8; 16]; + let iv = [0x0Cu8; 12]; + let mut seq = 100u64; + let mut ciphertext_buf = Vec::new(); + + // First call with small data + let mut plaintext1 = vec![0x46u8; 50]; + encrypt_plaintext_to_records(&mut plaintext1, &key, &iv, &mut seq, &mut ciphertext_buf) + .unwrap(); + assert_eq!(seq, 101); + + // Second call with data requiring fragmentation + let mut plaintext2 = vec![0x47u8; MAX_TLS_PLAINTEXT_LEN + 500]; + encrypt_plaintext_to_records(&mut plaintext2, &key, &iv, &mut seq, &mut ciphertext_buf) + .unwrap(); + assert_eq!(seq, 103); // Two more records + } + + #[test] + fn test_ciphertext_buf_appends() { + let key = [0x0Du8; 16]; + let iv = [0x0Eu8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + // Pre-populate buffer + ciphertext_buf.extend_from_slice(b"existing data"); + let initial_len = ciphertext_buf.len(); + + let mut plaintext = vec![0x48u8; 100]; + encrypt_plaintext_to_records(&mut plaintext, &key, &iv, &mut seq, &mut ciphertext_buf) + .unwrap(); + + // Should append, not overwrite + assert!(ciphertext_buf.len() > initial_len); + assert_eq!(&ciphertext_buf[..initial_len], b"existing data"); + } + + #[test] + fn test_encrypt_single_byte() { + // Minimum non-empty plaintext + let mut plaintext = vec![0x42u8; 1]; + let key = [0x0Fu8; 16]; + let iv = [0x10u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + let result = + encrypt_plaintext_to_records(&mut plaintext, &key, &iv, &mut seq, &mut ciphertext_buf); + assert!(result.is_ok()); + assert!(plaintext.is_empty()); + assert_eq!(seq, 1); + + // header(5) + plaintext(1) + content_type(1) + tag(16) = 23 + assert_eq!(ciphertext_buf.len(), 23); + } + + #[test] + fn test_encrypt_boundary_minus_one() { + // One byte below the limit - should still be single record (fast path) + let mut plaintext = vec![0x49u8; MAX_TLS_PLAINTEXT_LEN - 1]; + let key = [0x11u8; 16]; + let iv = [0x12u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + let result = + encrypt_plaintext_to_records(&mut plaintext, &key, &iv, &mut seq, &mut ciphertext_buf); + assert!(result.is_ok()); + assert!(plaintext.is_empty()); + assert_eq!(seq, 1); // Still single record + + let expected_len = TLS_RECORD_HEADER_SIZE + (MAX_TLS_PLAINTEXT_LEN - 1) + 1 + 16; + assert_eq!(ciphertext_buf.len(), expected_len); + } + + #[test] + fn test_encrypt_large_500kb() { + // 500KB = 512000 bytes, requires 32 records (512000 / 16384 = 31.25) + let size = 512000; + let expected_records = (size + MAX_TLS_PLAINTEXT_LEN - 1) / MAX_TLS_PLAINTEXT_LEN; + assert_eq!(expected_records, 32); // Verify our math + + let mut plaintext = vec![0x4Au8; size]; + let key = [0x13u8; 16]; + let iv = [0x14u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + let result = + encrypt_plaintext_to_records(&mut plaintext, &key, &iv, &mut seq, &mut ciphertext_buf); + assert!(result.is_ok()); + assert!(plaintext.is_empty()); + assert_eq!(seq, 32); // 32 records + + // 31 full records + 1 partial + let full_records = 31; + let last_record_plaintext = size - (full_records * MAX_TLS_PLAINTEXT_LEN); + let full_record_len = TLS_RECORD_HEADER_SIZE + MAX_TLS_PLAINTEXT_LEN + 1 + 16; + let last_record_len = TLS_RECORD_HEADER_SIZE + last_record_plaintext + 1 + 16; + let expected_total = full_records * full_record_len + last_record_len; + assert_eq!(ciphertext_buf.len(), expected_total); + } + + #[test] + fn test_sequence_number_high_values() { + // Test with sequence number near u64::MAX to ensure no overflow issues + let key = [0x15u8; 16]; + let iv = [0x16u8; 12]; + let mut seq = u64::MAX - 5; + let mut ciphertext_buf = Vec::new(); + + // Small data - single record + let mut plaintext = vec![0x4Bu8; 100]; + let result = + encrypt_plaintext_to_records(&mut plaintext, &key, &iv, &mut seq, &mut ciphertext_buf); + assert!(result.is_ok()); + assert_eq!(seq, u64::MAX - 4); + + // Another small record + let mut plaintext2 = vec![0x4Cu8; 100]; + let result2 = + encrypt_plaintext_to_records(&mut plaintext2, &key, &iv, &mut seq, &mut ciphertext_buf); + assert!(result2.is_ok()); + assert_eq!(seq, u64::MAX - 3); + } + + #[test] + fn test_record_headers_valid() { + // Verify all record headers in a fragmented output are valid TLS 1.3 headers + let size = MAX_TLS_PLAINTEXT_LEN * 3 + 100; // 3 full + 1 partial = 4 records + let mut plaintext = vec![0x4Du8; size]; + let key = [0x17u8; 16]; + let iv = [0x18u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + encrypt_plaintext_to_records(&mut plaintext, &key, &iv, &mut seq, &mut ciphertext_buf) + .unwrap(); + assert_eq!(seq, 4); + + // Parse and verify each record header + let mut offset = 0; + for i in 0..4 { + assert!( + offset + 5 <= ciphertext_buf.len(), + "Record {} header incomplete", + i + ); + + // Content type should be application_data (0x17) + assert_eq!( + ciphertext_buf[offset], 0x17, + "Record {} wrong content type", + i + ); + + // Version should be TLS 1.2 (0x0303) for compatibility + assert_eq!( + ciphertext_buf[offset + 1], + 0x03, + "Record {} wrong version major", + i + ); + assert_eq!( + ciphertext_buf[offset + 2], + 0x03, + "Record {} wrong version minor", + i + ); + + // Length field + let record_len = + u16::from_be_bytes([ciphertext_buf[offset + 3], ciphertext_buf[offset + 4]]) + as usize; + + // Record length must not exceed TLS 1.3 limit + assert!( + record_len <= MAX_TLS_CIPHERTEXT_LEN, + "Record {} length {} exceeds TLS 1.3 limit {}", + i, + record_len, + MAX_TLS_CIPHERTEXT_LEN + ); + + offset += 5 + record_len; + } + + // Should have consumed exactly all data + assert_eq!(offset, ciphertext_buf.len(), "Unexpected trailing data"); + } + + // ==================== Additional app data edge case tests ==================== + + #[test] + fn test_app_data_different_keys_produce_different_output() { + let key1 = [0x50u8; 16]; + let key2 = [0x51u8; 16]; + let iv = [0x52u8; 12]; + + let mut data1 = vec![0xAAu8; 100]; + let mut data2 = vec![0xAAu8; 100]; + let mut seq1 = 0u64; + let mut seq2 = 0u64; + let mut buf1 = Vec::new(); + let mut buf2 = Vec::new(); + + encrypt_plaintext_to_records(&mut data1, &key1, &iv, &mut seq1, &mut buf1).unwrap(); + encrypt_plaintext_to_records(&mut data2, &key2, &iv, &mut seq2, &mut buf2).unwrap(); + + assert_eq!(buf1.len(), buf2.len()); + assert_ne!(buf1, buf2); // Different keys = different ciphertext + } + + #[test] + fn test_app_data_different_ivs_produce_different_output() { + let key = [0x53u8; 16]; + let iv1 = [0x54u8; 12]; + let iv2 = [0x55u8; 12]; + + let mut data1 = vec![0xBBu8; 100]; + let mut data2 = vec![0xBBu8; 100]; + let mut seq1 = 0u64; + let mut seq2 = 0u64; + let mut buf1 = Vec::new(); + let mut buf2 = Vec::new(); + + encrypt_plaintext_to_records(&mut data1, &key, &iv1, &mut seq1, &mut buf1).unwrap(); + encrypt_plaintext_to_records(&mut data2, &key, &iv2, &mut seq2, &mut buf2).unwrap(); + + assert_eq!(buf1.len(), buf2.len()); + assert_ne!(buf1, buf2); // Different IVs = different ciphertext + } + + #[test] + fn test_app_data_different_seq_produce_different_output() { + let key = [0x56u8; 16]; + let iv = [0x57u8; 12]; + + let mut data1 = vec![0xCCu8; 100]; + let mut data2 = vec![0xCCu8; 100]; + let mut seq1 = 0u64; + let mut seq2 = 50u64; + let mut buf1 = Vec::new(); + let mut buf2 = Vec::new(); + + encrypt_plaintext_to_records(&mut data1, &key, &iv, &mut seq1, &mut buf1).unwrap(); + encrypt_plaintext_to_records(&mut data2, &key, &iv, &mut seq2, &mut buf2).unwrap(); + + assert_eq!(buf1.len(), buf2.len()); + assert_ne!(buf1, buf2); // Different seq = different nonce = different ciphertext + } + + #[test] + fn test_app_data_exactly_double_limit() { + let mut plaintext = vec![0xDDu8; MAX_TLS_PLAINTEXT_LEN * 2]; + let key = [0x58u8; 16]; + let iv = [0x59u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + encrypt_plaintext_to_records(&mut plaintext, &key, &iv, &mut seq, &mut ciphertext_buf) + .unwrap(); + assert!(plaintext.is_empty()); + assert_eq!(seq, 2); + + // Both records: header(5) + plaintext + content_type(1) + tag(16) + let expected_len = 2 * (TLS_RECORD_HEADER_SIZE + MAX_TLS_PLAINTEXT_LEN + 1 + 16); + assert_eq!(ciphertext_buf.len(), expected_len); + } + + #[test] + fn test_app_data_exactly_triple_limit() { + let mut plaintext = vec![0xEEu8; MAX_TLS_PLAINTEXT_LEN * 3]; + let key = [0x5Au8; 16]; + let iv = [0x5Bu8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + encrypt_plaintext_to_records(&mut plaintext, &key, &iv, &mut seq, &mut ciphertext_buf) + .unwrap(); + assert!(plaintext.is_empty()); + assert_eq!(seq, 3); + + // Three records: header(5) + plaintext + content_type(1) + tag(16) + let expected_len = 3 * (TLS_RECORD_HEADER_SIZE + MAX_TLS_PLAINTEXT_LEN + 1 + 16); + assert_eq!(ciphertext_buf.len(), expected_len); + } + + #[test] + fn test_app_data_256kb_stress() { + let size = 256 * 1024; + let expected_records = (size + MAX_TLS_PLAINTEXT_LEN - 1) / MAX_TLS_PLAINTEXT_LEN; + assert_eq!(expected_records, 16); + + let mut plaintext = vec![0xFFu8; size]; + let key = [0x5Cu8; 16]; + let iv = [0x5Du8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + encrypt_plaintext_to_records(&mut plaintext, &key, &iv, &mut seq, &mut ciphertext_buf) + .unwrap(); + assert!(plaintext.is_empty()); + assert_eq!(seq, 16); + + // Verify all records valid + let mut offset = 0; + let mut count = 0; + while offset < ciphertext_buf.len() { + let record_len = + u16::from_be_bytes([ciphertext_buf[offset + 3], ciphertext_buf[offset + 4]]) + as usize; + assert!(record_len <= MAX_TLS_CIPHERTEXT_LEN); + offset += 5 + record_len; + count += 1; + } + assert_eq!(count, 16); + } + + #[test] + fn test_app_data_preserves_structure() { + // Verify total plaintext reconstructable from records + let mut plaintext = vec![0x11u8; 20000]; + let key = [0x5Eu8; 16]; + let iv = [0x5Fu8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + encrypt_plaintext_to_records(&mut plaintext, &key, &iv, &mut seq, &mut ciphertext_buf) + .unwrap(); + + let mut offset = 0; + let mut total_plaintext = 0; + while offset < ciphertext_buf.len() { + let record_len = + u16::from_be_bytes([ciphertext_buf[offset + 3], ciphertext_buf[offset + 4]]) + as usize; + total_plaintext += record_len - 16 - 1; // minus tag and content type + offset += 5 + record_len; + } + assert_eq!(total_plaintext, 20000); + } + + #[test] + fn test_app_data_small_prime_sizes() { + for &size in &[17usize, 31, 127, 251, 509, 1021, 2039, 4093, 8191] { + let mut plaintext = vec![0x22u8; size]; + let key = [0x60u8; 16]; + let iv = [0x61u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + encrypt_plaintext_to_records(&mut plaintext, &key, &iv, &mut seq, &mut ciphertext_buf) + .unwrap(); + assert!(plaintext.is_empty()); + assert_eq!(seq, 1, "Size {} should be single record", size); + + let record_len = u16::from_be_bytes([ciphertext_buf[3], ciphertext_buf[4]]) as usize; + assert_eq!(record_len, size + 1 + 16); + } + } + + #[test] + fn test_app_data_near_boundary_primes() { + let primes = [16619, 16631, 16633, 16649]; + for &size in &primes { + let mut plaintext = vec![0x33u8; size]; + let key = [0x62u8; 16]; + let iv = [0x63u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + encrypt_plaintext_to_records(&mut plaintext, &key, &iv, &mut seq, &mut ciphertext_buf) + .unwrap(); + assert!(plaintext.is_empty()); + + let expected_records = if size <= MAX_TLS_PLAINTEXT_LEN { 1 } else { 2 }; + assert_eq!( + seq, expected_records, + "Wrong record count for size {}", + size + ); + } + } + + #[test] + fn test_app_data_multiple_calls_accumulate() { + let key = [0x64u8; 16]; + let iv = [0x65u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + let mut data1 = vec![0x44u8; 1000]; + encrypt_plaintext_to_records(&mut data1, &key, &iv, &mut seq, &mut ciphertext_buf).unwrap(); + let len1 = ciphertext_buf.len(); + assert_eq!(seq, 1); + + let mut data2 = vec![0x55u8; MAX_TLS_PLAINTEXT_LEN + 500]; + encrypt_plaintext_to_records(&mut data2, &key, &iv, &mut seq, &mut ciphertext_buf).unwrap(); + assert!(ciphertext_buf.len() > len1); + assert_eq!(seq, 3); // 1 + 2 more + } + + #[test] + fn test_app_data_high_sequence_fragmentation() { + let key = [0x66u8; 16]; + let iv = [0x67u8; 12]; + let mut seq = u64::MAX - 5; + let mut ciphertext_buf = Vec::new(); + + // Fragmentation at high seq numbers + let mut plaintext = vec![0x66u8; MAX_TLS_PLAINTEXT_LEN + 100]; + encrypt_plaintext_to_records(&mut plaintext, &key, &iv, &mut seq, &mut ciphertext_buf) + .unwrap(); + assert_eq!(seq, u64::MAX - 3); // 2 records + } + + #[test] + fn test_app_data_boundary_exactly_at_limit() { + let mut plaintext = vec![0x77u8; MAX_TLS_PLAINTEXT_LEN]; + let key = [0x68u8; 16]; + let iv = [0x69u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + encrypt_plaintext_to_records(&mut plaintext, &key, &iv, &mut seq, &mut ciphertext_buf) + .unwrap(); + assert_eq!(seq, 1); + + // Record length = plaintext + content_type(1) + tag(16) + let record_len = u16::from_be_bytes([ciphertext_buf[3], ciphertext_buf[4]]) as usize; + assert_eq!(record_len, MAX_TLS_PLAINTEXT_LEN + 1 + 16); + } + + #[test] + fn test_app_data_boundary_one_over() { + let mut plaintext = vec![0x88u8; MAX_TLS_PLAINTEXT_LEN + 1]; + let key = [0x6Au8; 16]; + let iv = [0x6Bu8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + encrypt_plaintext_to_records(&mut plaintext, &key, &iv, &mut seq, &mut ciphertext_buf) + .unwrap(); + assert_eq!(seq, 2); + + // First at max: plaintext + content_type(1) + tag(16) + let first_len = u16::from_be_bytes([ciphertext_buf[3], ciphertext_buf[4]]) as usize; + assert_eq!(first_len, MAX_TLS_PLAINTEXT_LEN + 1 + 16); + + // Second has 1 byte + let second_offset = 5 + first_len; + let second_len = u16::from_be_bytes([ + ciphertext_buf[second_offset + 3], + ciphertext_buf[second_offset + 4], + ]) as usize; + assert_eq!(second_len, 1 + 1 + 16); + } + + // ==================== Handshake encryption tests ==================== + + #[test] + fn test_handshake_encrypt_empty() { + let handshake_data: &[u8] = &[]; + let key = [0u8; 16]; + let iv = [0u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + let result = + encrypt_handshake_to_records(handshake_data, &key, &iv, &mut seq, &mut ciphertext_buf); + assert!(result.is_ok()); + assert!(ciphertext_buf.is_empty()); + assert_eq!(seq, 0); // No records encrypted + } + + #[test] + fn test_handshake_encrypt_small_single_record() { + // Typical small handshake (EncryptedExtensions + small cert + CV + Finished) + let handshake_data = vec![0x16u8; 500]; + let key = [0x01u8; 16]; + let iv = [0x02u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + let result = + encrypt_handshake_to_records(&handshake_data, &key, &iv, &mut seq, &mut ciphertext_buf); + assert!(result.is_ok()); + assert_eq!(seq, 1); // One record encrypted + + // Should have one record: header(5) + handshake(500) + content_type(1) + tag(16) + let expected_len = TLS_RECORD_HEADER_SIZE + 500 + 1 + 16; + assert_eq!(ciphertext_buf.len(), expected_len); + + // Verify TLS header - outer type is ApplicationData (0x17) + assert_eq!(ciphertext_buf[0], CONTENT_TYPE_APPLICATION_DATA); + assert_eq!(ciphertext_buf[1], 0x03); + assert_eq!(ciphertext_buf[2], 0x03); + let record_len = u16::from_be_bytes([ciphertext_buf[3], ciphertext_buf[4]]) as usize; + assert_eq!(record_len, 500 + 1 + 16); + } + + #[test] + fn test_handshake_encrypt_max_single_record() { + // Exactly at the boundary - should still be single record + let handshake_data = vec![0x16u8; MAX_TLS_PLAINTEXT_LEN]; + let key = [0x03u8; 16]; + let iv = [0x04u8; 12]; + let mut seq = 5u64; + let mut ciphertext_buf = Vec::new(); + + let result = + encrypt_handshake_to_records(&handshake_data, &key, &iv, &mut seq, &mut ciphertext_buf); + assert!(result.is_ok()); + assert_eq!(seq, 6); // One record encrypted + + let expected_len = TLS_RECORD_HEADER_SIZE + MAX_TLS_PLAINTEXT_LEN + 1 + 16; + assert_eq!(ciphertext_buf.len(), expected_len); + } + + #[test] + fn test_handshake_encrypt_fragmentation_two_records() { + // One byte over the limit - should produce two records + let handshake_data = vec![0x16u8; MAX_TLS_PLAINTEXT_LEN + 1]; + let key = [0x05u8; 16]; + let iv = [0x06u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + let result = + encrypt_handshake_to_records(&handshake_data, &key, &iv, &mut seq, &mut ciphertext_buf); + assert!(result.is_ok()); + assert_eq!(seq, 2); // Two records encrypted + + // First record: full MAX_TLS_PLAINTEXT_LEN + let first_record_len = TLS_RECORD_HEADER_SIZE + MAX_TLS_PLAINTEXT_LEN + 1 + 16; + // Second record: 1 byte + let second_record_len = TLS_RECORD_HEADER_SIZE + 1 + 1 + 16; + assert_eq!(ciphertext_buf.len(), first_record_len + second_record_len); + } + + #[test] + fn test_handshake_encrypt_large_certificate() { + // Simulate a large certificate chain (e.g., 20KB) + // This is a realistic scenario that would cause record overflow without fragmentation + let handshake_data = vec![0x16u8; 20000]; + let key = [0x07u8; 16]; + let iv = [0x08u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + let result = + encrypt_handshake_to_records(&handshake_data, &key, &iv, &mut seq, &mut ciphertext_buf); + assert!(result.is_ok()); + assert_eq!(seq, 2); // Should be 2 records (20000 / 16384 = 1.2) + + // Verify all records are valid + let mut offset = 0; + for i in 0..2 { + assert!( + offset + 5 <= ciphertext_buf.len(), + "Record {} header incomplete", + i + ); + + // Outer type is ApplicationData (0x17) + assert_eq!(ciphertext_buf[offset], CONTENT_TYPE_APPLICATION_DATA); + assert_eq!(ciphertext_buf[offset + 1], 0x03); + assert_eq!(ciphertext_buf[offset + 2], 0x03); + + let record_len = + u16::from_be_bytes([ciphertext_buf[offset + 3], ciphertext_buf[offset + 4]]) + as usize; + + // Record length must not exceed TLS 1.3 limit + assert!( + record_len <= MAX_TLS_CIPHERTEXT_LEN, + "Record {} length {} exceeds TLS 1.3 limit {}", + i, + record_len, + MAX_TLS_CIPHERTEXT_LEN + ); + + offset += 5 + record_len; + } + assert_eq!(offset, ciphertext_buf.len()); + } + + #[test] + fn test_handshake_encrypt_very_large_certificate_chain() { + // Simulate a very large certificate chain (e.g., 50KB - 3+ records) + let handshake_data = vec![0x16u8; 50000]; + let key = [0x09u8; 16]; + let iv = [0x0Au8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + let result = + encrypt_handshake_to_records(&handshake_data, &key, &iv, &mut seq, &mut ciphertext_buf); + assert!(result.is_ok()); + + // 50000 / 16384 = 3.01 -> 4 records + let expected_records = (50000 + MAX_TLS_PLAINTEXT_LEN - 1) / MAX_TLS_PLAINTEXT_LEN; + assert_eq!(expected_records, 4); + assert_eq!(seq, 4); + + // Verify all records + let mut offset = 0; + let mut record_count = 0; + while offset < ciphertext_buf.len() { + assert!(offset + 5 <= ciphertext_buf.len()); + + let record_len = + u16::from_be_bytes([ciphertext_buf[offset + 3], ciphertext_buf[offset + 4]]) + as usize; + + assert!(record_len <= MAX_TLS_CIPHERTEXT_LEN); + offset += 5 + record_len; + record_count += 1; + } + assert_eq!(record_count, 4); + } + + #[test] + fn test_handshake_sequence_number_increments() { + let key = [0x0Bu8; 16]; + let iv = [0x0Cu8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + // First handshake - single record + let handshake1 = vec![0x16u8; 1000]; + encrypt_handshake_to_records(&handshake1, &key, &iv, &mut seq, &mut ciphertext_buf) + .unwrap(); + assert_eq!(seq, 1); + + // Second handshake - requires fragmentation + let handshake2 = vec![0x16u8; MAX_TLS_PLAINTEXT_LEN + 5000]; + encrypt_handshake_to_records(&handshake2, &key, &iv, &mut seq, &mut ciphertext_buf) + .unwrap(); + assert_eq!(seq, 3); // 1 + 2 more records + } + + #[test] + fn test_handshake_realistic_sizes() { + // Test with realistic handshake message sizes + // EncryptedExtensions: ~50 bytes + // Certificate (typical): ~2-3KB + // CertificateVerify: ~200 bytes + // Finished: ~36 bytes + // Total: ~3KB - should fit in one record + let typical_handshake = vec![0x16u8; 3000]; + let key = [0x0Du8; 16]; + let iv = [0x0Eu8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + let result = encrypt_handshake_to_records( + &typical_handshake, + &key, + &iv, + &mut seq, + &mut ciphertext_buf, + ); + assert!(result.is_ok()); + assert_eq!(seq, 1); // Single record for typical handshake + } + + #[test] + fn test_handshake_realistic_large_cert_chain() { + // Large certificate chain scenario: + // EncryptedExtensions: ~50 bytes + // Certificate (with chain): ~15KB + // CertificateVerify: ~200 bytes + // Finished: ~36 bytes + // Total: ~15.3KB - should fit in one record + let large_cert_handshake = vec![0x16u8; 15300]; + let key = [0x0Fu8; 16]; + let iv = [0x10u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + let result = encrypt_handshake_to_records( + &large_cert_handshake, + &key, + &iv, + &mut seq, + &mut ciphertext_buf, + ); + assert!(result.is_ok()); + assert_eq!(seq, 1); // Still fits in one record (15300 < 16384) + } + + #[test] + fn test_handshake_realistic_very_large_cert_chain() { + // Very large certificate chain scenario: + // EncryptedExtensions: ~50 bytes + // Certificate (with long chain + extensions): ~18KB + // CertificateVerify: ~200 bytes + // Finished: ~36 bytes + // Total: ~18.3KB - requires 2 records + let very_large_handshake = vec![0x16u8; 18300]; + let key = [0x11u8; 16]; + let iv = [0x12u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + let result = encrypt_handshake_to_records( + &very_large_handshake, + &key, + &iv, + &mut seq, + &mut ciphertext_buf, + ); + assert!(result.is_ok()); + assert_eq!(seq, 2); // Requires 2 records (18300 > 16384) + } + + #[test] + fn test_handshake_boundary_exactly_at_limit() { + // Test exactly at MAX_TLS_PLAINTEXT_LEN (16384 bytes) + let handshake_data = vec![0x16u8; MAX_TLS_PLAINTEXT_LEN]; + let key = [0x13u8; 16]; + let iv = [0x14u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + let result = + encrypt_handshake_to_records(&handshake_data, &key, &iv, &mut seq, &mut ciphertext_buf); + assert!(result.is_ok()); + assert_eq!(seq, 1); // Exactly one record + + // Verify the ciphertext length: plaintext + content_type(1) + tag(16) + let record_len = u16::from_be_bytes([ciphertext_buf[3], ciphertext_buf[4]]) as usize; + assert_eq!(record_len, MAX_TLS_PLAINTEXT_LEN + 1 + 16); + } + + #[test] + fn test_handshake_boundary_one_over_limit() { + // Test one byte over MAX_TLS_PLAINTEXT_LEN (16385 bytes) + let handshake_data = vec![0x16u8; MAX_TLS_PLAINTEXT_LEN + 1]; + let key = [0x15u8; 16]; + let iv = [0x16u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + let result = + encrypt_handshake_to_records(&handshake_data, &key, &iv, &mut seq, &mut ciphertext_buf); + assert!(result.is_ok()); + assert_eq!(seq, 2); // Requires two records + + // First record at max: plaintext + content_type(1) + tag(16) + let first_record_len = u16::from_be_bytes([ciphertext_buf[3], ciphertext_buf[4]]) as usize; + assert_eq!(first_record_len, MAX_TLS_PLAINTEXT_LEN + 1 + 16); + + // Second record should have just 1 byte of payload + let second_offset = 5 + first_record_len; + let second_record_len = u16::from_be_bytes([ + ciphertext_buf[second_offset + 3], + ciphertext_buf[second_offset + 4], + ]) as usize; + assert_eq!(second_record_len, 1 + 1 + 16); // 1 byte + content_type + tag + } + + #[test] + fn test_handshake_single_byte() { + let handshake_data = vec![0x16u8; 1]; + let key = [0x17u8; 16]; + let iv = [0x18u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + let result = + encrypt_handshake_to_records(&handshake_data, &key, &iv, &mut seq, &mut ciphertext_buf); + assert!(result.is_ok()); + assert_eq!(seq, 1); + + // header(5) + handshake(1) + content_type(1) + tag(16) = 23 + assert_eq!(ciphertext_buf.len(), 23); + } + + #[test] + fn test_handshake_record_headers_all_valid() { + // Test that all fragmented records have valid headers + let size = MAX_TLS_PLAINTEXT_LEN * 3 + 100; + let handshake_data = vec![0x16u8; size]; + let key = [0x19u8; 16]; + let iv = [0x1Au8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + encrypt_handshake_to_records(&handshake_data, &key, &iv, &mut seq, &mut ciphertext_buf) + .unwrap(); + assert_eq!(seq, 4); // 3 full + 1 partial + + let mut offset = 0; + for i in 0..4 { + // Outer content type is ApplicationData (0x17) + assert_eq!( + ciphertext_buf[offset], CONTENT_TYPE_APPLICATION_DATA, + "Record {} has wrong outer content type", + i + ); + + // Version is TLS 1.2 (0x0303) + assert_eq!(ciphertext_buf[offset + 1], 0x03); + assert_eq!(ciphertext_buf[offset + 2], 0x03); + + let record_len = + u16::from_be_bytes([ciphertext_buf[offset + 3], ciphertext_buf[offset + 4]]) + as usize; + + // All records must be within TLS 1.3 limit + assert!( + record_len <= MAX_TLS_CIPHERTEXT_LEN, + "Record {} length {} exceeds limit", + i, + record_len + ); + + offset += 5 + record_len; + } + assert_eq!(offset, ciphertext_buf.len()); + } + + // ==================== Additional edge case tests ==================== + + #[test] + fn test_handshake_high_sequence_numbers() { + // Test with sequence numbers near u64::MAX + let key = [0x1Bu8; 16]; + let iv = [0x1Cu8; 12]; + let mut seq = u64::MAX - 10; + let mut ciphertext_buf = Vec::new(); + + // Single record + let handshake1 = vec![0x16u8; 1000]; + encrypt_handshake_to_records(&handshake1, &key, &iv, &mut seq, &mut ciphertext_buf) + .unwrap(); + assert_eq!(seq, u64::MAX - 9); + + // Two records (fragmentation at high seq) + let handshake2 = vec![0x16u8; MAX_TLS_PLAINTEXT_LEN + 100]; + encrypt_handshake_to_records(&handshake2, &key, &iv, &mut seq, &mut ciphertext_buf) + .unwrap(); + assert_eq!(seq, u64::MAX - 7); // 2 more records + } + + #[test] + fn test_handshake_appends_to_existing_buffer() { + let key = [0x1Du8; 16]; + let iv = [0x1Eu8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + // Pre-populate buffer (simulating ServerHello already written) + ciphertext_buf.extend_from_slice(b"existing server hello data"); + let initial_len = ciphertext_buf.len(); + + let handshake = vec![0x16u8; 500]; + encrypt_handshake_to_records(&handshake, &key, &iv, &mut seq, &mut ciphertext_buf).unwrap(); + + // Should append, not overwrite + assert!(ciphertext_buf.len() > initial_len); + assert_eq!( + &ciphertext_buf[..initial_len], + b"existing server hello data" + ); + } + + #[test] + fn test_handshake_multiple_calls_accumulate() { + // Simulating how the server might send multiple handshake batches + let key = [0x1Fu8; 16]; + let iv = [0x20u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + // First batch (e.g., just EE + Cert) + let batch1 = vec![0x16u8; 5000]; + encrypt_handshake_to_records(&batch1, &key, &iv, &mut seq, &mut ciphertext_buf).unwrap(); + let len_after_batch1 = ciphertext_buf.len(); + assert_eq!(seq, 1); + + // Second batch (e.g., CV + Finished) - note seq continues + let batch2 = vec![0x16u8; 300]; + encrypt_handshake_to_records(&batch2, &key, &iv, &mut seq, &mut ciphertext_buf).unwrap(); + assert!(ciphertext_buf.len() > len_after_batch1); + assert_eq!(seq, 2); + } + + #[test] + fn test_handshake_boundary_minus_one() { + // Test one byte below the limit (16622 bytes) + let handshake_data = vec![0x16u8; MAX_TLS_PLAINTEXT_LEN - 1]; + let key = [0x21u8; 16]; + let iv = [0x22u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + let result = + encrypt_handshake_to_records(&handshake_data, &key, &iv, &mut seq, &mut ciphertext_buf); + assert!(result.is_ok()); + assert_eq!(seq, 1); // Still single record (fast path) + + let expected_len = TLS_RECORD_HEADER_SIZE + (MAX_TLS_PLAINTEXT_LEN - 1) + 1 + 16; + assert_eq!(ciphertext_buf.len(), expected_len); + } + + #[test] + fn test_handshake_exactly_double_limit() { + // Test exactly 2x the limit (32768 bytes) - should produce exactly 2 full records + let handshake_data = vec![0x16u8; MAX_TLS_PLAINTEXT_LEN * 2]; + let key = [0x23u8; 16]; + let iv = [0x24u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + let result = + encrypt_handshake_to_records(&handshake_data, &key, &iv, &mut seq, &mut ciphertext_buf); + assert!(result.is_ok()); + assert_eq!(seq, 2); + + // Both records: header(5) + plaintext + content_type(1) + tag(16) + let expected_len = 2 * (TLS_RECORD_HEADER_SIZE + MAX_TLS_PLAINTEXT_LEN + 1 + 16); + assert_eq!(ciphertext_buf.len(), expected_len); + } + + #[test] + fn test_handshake_exactly_triple_limit() { + // Test exactly 3x the limit (49152 bytes) + let handshake_data = vec![0x16u8; MAX_TLS_PLAINTEXT_LEN * 3]; + let key = [0x25u8; 16]; + let iv = [0x26u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + let result = + encrypt_handshake_to_records(&handshake_data, &key, &iv, &mut seq, &mut ciphertext_buf); + assert!(result.is_ok()); + assert_eq!(seq, 3); + + // Three records: header(5) + plaintext + content_type(1) + tag(16) + let expected_len = 3 * (TLS_RECORD_HEADER_SIZE + MAX_TLS_PLAINTEXT_LEN + 1 + 16); + assert_eq!(ciphertext_buf.len(), expected_len); + } + + #[test] + fn test_handshake_different_keys_produce_different_output() { + let handshake = vec![0x16u8; 100]; + let iv = [0x27u8; 12]; + let mut seq1 = 0u64; + let mut seq2 = 0u64; + + let key1 = [0x28u8; 16]; + let key2 = [0x29u8; 16]; + + let mut buf1 = Vec::new(); + let mut buf2 = Vec::new(); + + encrypt_handshake_to_records(&handshake, &key1, &iv, &mut seq1, &mut buf1).unwrap(); + encrypt_handshake_to_records(&handshake, &key2, &iv, &mut seq2, &mut buf2).unwrap(); + + // Same structure but different ciphertext + assert_eq!(buf1.len(), buf2.len()); + assert_ne!(buf1, buf2); // Different keys = different ciphertext + } + + #[test] + fn test_handshake_different_ivs_produce_different_output() { + let handshake = vec![0x16u8; 100]; + let key = [0x2Au8; 16]; + let mut seq1 = 0u64; + let mut seq2 = 0u64; + + let iv1 = [0x2Bu8; 12]; + let iv2 = [0x2Cu8; 12]; + + let mut buf1 = Vec::new(); + let mut buf2 = Vec::new(); + + encrypt_handshake_to_records(&handshake, &key, &iv1, &mut seq1, &mut buf1).unwrap(); + encrypt_handshake_to_records(&handshake, &key, &iv2, &mut seq2, &mut buf2).unwrap(); + + assert_eq!(buf1.len(), buf2.len()); + assert_ne!(buf1, buf2); // Different IVs = different ciphertext + } + + #[test] + fn test_handshake_different_sequence_numbers_produce_different_output() { + let handshake = vec![0x16u8; 100]; + let key = [0x2Du8; 16]; + let iv = [0x2Eu8; 12]; + + let mut seq1 = 0u64; + let mut seq2 = 100u64; + + let mut buf1 = Vec::new(); + let mut buf2 = Vec::new(); + + encrypt_handshake_to_records(&handshake, &key, &iv, &mut seq1, &mut buf1).unwrap(); + encrypt_handshake_to_records(&handshake, &key, &iv, &mut seq2, &mut buf2).unwrap(); + + assert_eq!(buf1.len(), buf2.len()); + assert_ne!(buf1, buf2); // Different seq = different nonce = different ciphertext + } + + #[test] + fn test_plaintext_and_handshake_produce_different_inner_content_type() { + // While we can't directly inspect the encrypted inner content type, + // we can verify that the same plaintext encrypted as app data vs handshake + // produces different ciphertext (due to different content type byte) + let data = vec![0xAAu8; 100]; + let key = [0x2Fu8; 16]; + let iv = [0x30u8; 12]; + + let mut seq1 = 0u64; + let mut seq2 = 0u64; + + let mut app_buf = Vec::new(); + let mut hs_buf = Vec::new(); + + let mut app_data = data.clone(); + encrypt_plaintext_to_records(&mut app_data, &key, &iv, &mut seq1, &mut app_buf).unwrap(); + encrypt_handshake_to_records(&data, &key, &iv, &mut seq2, &mut hs_buf).unwrap(); + + // Same size but different content (inner content type differs) + assert_eq!(app_buf.len(), hs_buf.len()); + assert_ne!(app_buf, hs_buf); + } + + #[test] + fn test_handshake_256kb_stress_test() { + // 256KB handshake data (extreme case) + let size = 256 * 1024; + let expected_records = (size + MAX_TLS_PLAINTEXT_LEN - 1) / MAX_TLS_PLAINTEXT_LEN; + assert_eq!(expected_records, 16); // 262144 / 16384 = 15.77 -> 16 + + let handshake_data = vec![0x16u8; size]; + let key = [0x31u8; 16]; + let iv = [0x32u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + let result = + encrypt_handshake_to_records(&handshake_data, &key, &iv, &mut seq, &mut ciphertext_buf); + assert!(result.is_ok()); + assert_eq!(seq, 16); + + // Verify all records + let mut offset = 0; + let mut count = 0; + while offset < ciphertext_buf.len() { + let record_len = + u16::from_be_bytes([ciphertext_buf[offset + 3], ciphertext_buf[offset + 4]]) + as usize; + assert!(record_len <= MAX_TLS_CIPHERTEXT_LEN); + offset += 5 + record_len; + count += 1; + } + assert_eq!(count, 16); + assert_eq!(offset, ciphertext_buf.len()); + } + + #[test] + fn test_app_data_and_handshake_same_fragmentation_boundaries() { + // Both app data and handshake should fragment at the same boundaries + let size = MAX_TLS_PLAINTEXT_LEN + 1; + let key = [0x33u8; 16]; + let iv = [0x34u8; 12]; + + let mut app_seq = 0u64; + let mut hs_seq = 0u64; + let mut app_buf = Vec::new(); + let mut hs_buf = Vec::new(); + + let mut app_data = vec![0xBBu8; size]; + let hs_data = vec![0xCCu8; size]; + + encrypt_plaintext_to_records(&mut app_data, &key, &iv, &mut app_seq, &mut app_buf).unwrap(); + encrypt_handshake_to_records(&hs_data, &key, &iv, &mut hs_seq, &mut hs_buf).unwrap(); + + // Both should produce exactly 2 records + assert_eq!(app_seq, 2); + assert_eq!(hs_seq, 2); + + // Same total size + assert_eq!(app_buf.len(), hs_buf.len()); + } + + #[test] + fn test_handshake_preserves_data_integrity_structure() { + // Verify that the record structure allows proper parsing + let handshake_data = vec![0x16u8; 20000]; // 2 records + let key = [0x35u8; 16]; + let iv = [0x36u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + encrypt_handshake_to_records(&handshake_data, &key, &iv, &mut seq, &mut ciphertext_buf) + .unwrap(); + + // Parse records and verify total plaintext size matches + let mut offset = 0; + let mut total_plaintext = 0; + while offset < ciphertext_buf.len() { + let record_len = + u16::from_be_bytes([ciphertext_buf[offset + 3], ciphertext_buf[offset + 4]]) + as usize; + + // Plaintext size = record_len - 16 (tag) - 1 (content type) + total_plaintext += record_len - 16 - 1; + offset += 5 + record_len; + } + + assert_eq!(total_plaintext, 20000); + } + + #[test] + fn test_small_prime_sized_handshake() { + // Test with prime-number sizes to catch off-by-one errors + for &size in &[17usize, 31, 127, 251, 509, 1021, 2039, 4093, 8191] { + let handshake_data = vec![0x16u8; size]; + let key = [0x37u8; 16]; + let iv = [0x38u8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + let result = encrypt_handshake_to_records( + &handshake_data, + &key, + &iv, + &mut seq, + &mut ciphertext_buf, + ); + assert!(result.is_ok(), "Failed for size {}", size); + assert_eq!(seq, 1, "Size {} should be single record", size); + + let record_len = u16::from_be_bytes([ciphertext_buf[3], ciphertext_buf[4]]) as usize; + assert_eq!(record_len, size + 1 + 16); + } + } + + #[test] + fn test_handshake_near_boundary_prime_sizes() { + // Prime numbers near the boundary + let near_boundary_primes = [16619, 16631, 16633, 16649]; + + for &size in &near_boundary_primes { + let handshake_data = vec![0x16u8; size]; + let key = [0x39u8; 16]; + let iv = [0x3Au8; 12]; + let mut seq = 0u64; + let mut ciphertext_buf = Vec::new(); + + let result = encrypt_handshake_to_records( + &handshake_data, + &key, + &iv, + &mut seq, + &mut ciphertext_buf, + ); + assert!(result.is_ok(), "Failed for size {}", size); + + let expected_records = if size <= MAX_TLS_PLAINTEXT_LEN { 1 } else { 2 }; + assert_eq!( + seq, expected_records, + "Wrong record count for size {}", + size + ); + } + } +} diff --git a/src/reality/reality_server_connection.rs b/src/reality/reality_server_connection.rs new file mode 100644 index 00000000..8a2a988e --- /dev/null +++ b/src/reality/reality_server_connection.rs @@ -0,0 +1,901 @@ +// REALITY server-side connection +// +// This implements a rustls-compatible API for REALITY protocol server connections, +// allowing REALITY to be used as a drop-in replacement for rustls. + +use std::io::{self, Read, Write}; + +use crate::address::{Address, NetLocation}; +use crate::slide_buffer::SlideBuffer; + +use super::common::{ + self, CIPHERTEXT_READ_BUF_CAPACITY, CONTENT_TYPE_ALERT, CONTENT_TYPE_APPLICATION_DATA, + CONTENT_TYPE_CHANGE_CIPHER_SPEC, CONTENT_TYPE_HANDSHAKE, HANDSHAKE_TYPE_FINISHED, + PLAINTEXT_READ_BUF_CAPACITY, TLS_MAX_RECORD_SIZE, TLS_RECORD_HEADER_SIZE, +}; +use super::reality_aead::{decrypt_handshake_message, decrypt_tls13_record}; +use super::reality_auth::{decrypt_session_id, derive_auth_key, perform_ecdh}; +use super::reality_certificate::generate_hmac_certificate; +use super::reality_io_state::RealityIoState; +use super::reality_reader_writer::{RealityReader, RealityWriter}; +use super::reality_records::{encrypt_handshake_to_records, encrypt_plaintext_to_records}; +use super::reality_tls13_keys::{ + compute_finished_verify_data, derive_application_secrets, derive_handshake_keys, + derive_traffic_keys, +}; +use super::reality_tls13_messages::*; +use super::reality_util::{ + extract_client_public_key, extract_client_random, extract_session_id_slice, +}; +use aws_lc_rs::{ + agreement, digest, + rand::{SecureRandom, SystemRandom}, +}; +use subtle::ConstantTimeEq; + +/// Configuration for REALITY server connections +#[derive(Clone)] +pub struct RealityServerConfig { + /// Server's X25519 private key (32 bytes) + pub private_key: [u8; 32], + /// List of valid short IDs for authentication (8 bytes each) + pub short_ids: Vec<[u8; 8]>, + /// Destination server for certificate generation + pub dest: NetLocation, + /// Maximum allowed time difference in milliseconds (None = no check) + pub max_time_diff: Option, + /// Minimum accepted client version (3 bytes: major.minor.patch) + pub min_client_version: Option<[u8; 3]>, + /// Maximum accepted client version (3 bytes: major.minor.patch) + pub max_client_version: Option<[u8; 3]>, +} + +/// Handshake state machine for REALITY server +enum HandshakeState { + /// Initial state, waiting for ClientHello + Initial, + /// ServerHello and encrypted handshake messages sent, waiting for client Finished + ServerHelloSent { + handshake_hash: [u8; 32], // Hash before server Finished + handshake_hash_with_server_finished: [u8; 32], // Hash including server Finished (for verifying client Finished) + client_handshake_traffic_secret: Vec, + master_secret: Vec, + cipher_suite: u16, + }, + /// Handshake complete, ready for application data + Complete, +} + +/// REALITY server-side connection implementing rustls-compatible API +pub struct RealityServerConnection { + // Configuration + config: RealityServerConfig, + + // Handshake state + handshake_state: HandshakeState, + + // TLS 1.3 application traffic encryption (post-handshake) + app_read_key: Option>, + app_read_iv: Option>, + app_write_key: Option>, + app_write_iv: Option>, + read_seq: u64, + write_seq: u64, + cipher_suite: u16, + + // Pre-allocated buffer for TLS read operations (reused across calls) + tls_read_buffer: Box<[u8; TLS_MAX_RECORD_SIZE]>, + + // Buffers for I/O - using SlideBuffer for efficient zero-alloc operations + ciphertext_read_buf: SlideBuffer, // Incoming encrypted TLS records + ciphertext_write_buf: Vec, // Outgoing encrypted TLS records + plaintext_read_buf: SlideBuffer, // Decrypted application data + plaintext_write_buf: Vec, // Application data to encrypt +} + +impl RealityServerConnection { + /// Create a new REALITY server connection + pub fn new(config: RealityServerConfig) -> io::Result { + Ok(RealityServerConnection { + config, + handshake_state: HandshakeState::Initial, + app_read_key: None, + app_read_iv: None, + app_write_key: None, + app_write_iv: None, + read_seq: 0, + write_seq: 0, + cipher_suite: 0, + tls_read_buffer: Box::new([0u8; TLS_MAX_RECORD_SIZE]), + ciphertext_read_buf: SlideBuffer::new(CIPHERTEXT_READ_BUF_CAPACITY), + ciphertext_write_buf: Vec::with_capacity(CIPHERTEXT_READ_BUF_CAPACITY), + plaintext_read_buf: SlideBuffer::new(PLAINTEXT_READ_BUF_CAPACITY), + plaintext_write_buf: Vec::with_capacity(TLS_MAX_RECORD_SIZE), + }) + } + + /// Read TLS messages from the provided reader into internal buffer + /// + /// This does NOT decrypt - call process_new_packets() for that. + /// Uses pre-allocated buffer to avoid allocation on every call. + pub fn read_tls(&mut self, rd: &mut dyn Read) -> io::Result { + // Compact if remaining capacity is insufficient for a full TLS record + if self.ciphertext_read_buf.remaining_capacity() < TLS_MAX_RECORD_SIZE { + self.ciphertext_read_buf.compact(); + } + + // Read into pre-allocated buffer + let n = rd.read(&mut self.tls_read_buffer[..])?; + if n > 0 { + self.ciphertext_read_buf + .extend_from_slice(&self.tls_read_buffer[..n]); + } + Ok(n) + } + + /// Process buffered TLS messages and advance handshake/decrypt data + /// + /// Returns I/O state with available plaintext bytes and write status. + pub fn process_new_packets(&mut self) -> io::Result { + match &self.handshake_state { + HandshakeState::Initial => { + self.process_client_hello()?; + } + HandshakeState::ServerHelloSent { .. } => { + self.process_client_finished()?; + } + HandshakeState::Complete => { + self.process_application_data()?; + } + } + + Ok(RealityIoState::new(self.plaintext_read_buf.len())) + } + + /// Process ClientHello message and send ServerHello + fn process_client_hello(&mut self) -> io::Result<()> { + // Need at least TLS record header (5 bytes) + if self.ciphertext_read_buf.len() < TLS_RECORD_HEADER_SIZE { + return Ok(()); // Need more data + } + + // Parse TLS record length + let record_len = self + .ciphertext_read_buf + .get_u16_be(3) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "Buffer too short"))? + as usize; + + // Check if we have the complete record + let total_record_len = TLS_RECORD_HEADER_SIZE + record_len; + if self.ciphertext_read_buf.len() < total_record_len { + return Ok(()); // Need more data + } + + // Copy the ClientHello record to a Vec for processing + // (We need to keep it around for transcript hashing and AAD modification) + let client_hello: Vec = self.ciphertext_read_buf[..total_record_len].to_vec(); + self.ciphertext_read_buf.consume(total_record_len); + + // Step 1: Extract fields from ClientHello (using slice for session_id to avoid allocation) + let client_random = extract_client_random(&client_hello)?; + let session_id = extract_session_id_slice(&client_hello)?; + let client_public_key = extract_client_public_key(&client_hello)?; + + log::debug!( + "REALITY: ClientHello received, client_random: {:?}", + &client_random[..8] + ); + + // Step 2: Perform ECDH to derive auth key + let shared_secret = perform_ecdh(&self.config.private_key, &client_public_key) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()))?; + + // Use slices directly from client_random to avoid copying + let salt = &client_random[0..20]; + let auth_key = derive_auth_key(&shared_secret, salt, b"REALITY") + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()))?; + + // Step 3: Validate session ID (contains encrypted metadata) + if session_id.len() != 32 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "Invalid session ID length", + )); + } + + let nonce = &client_random[20..32]; + let mut encrypted_session_id_arr = [0u8; 32]; + encrypted_session_id_arr.copy_from_slice(session_id); + + // CRITICAL: Reconstruct AAD with zeros at SessionId location + // The AAD during encryption had zeros where the session ID would go + // SessionId is at offset 39 in ClientHello handshake (after type(1) + length(3) + version(2) + random(32) + sessionid_length(1)) + let client_hello_handshake = &client_hello[TLS_RECORD_HEADER_SIZE..]; + let mut aad_for_decryption = client_hello_handshake.to_vec(); + if aad_for_decryption.len() >= 39 + 32 { + // Replace encrypted SessionId with zeros + aad_for_decryption[39..39 + 32].fill(0); + } + + let decrypted_session_id = decrypt_session_id( + &encrypted_session_id_arr, + &auth_key, + nonce, + &aad_for_decryption, + ) + .map_err(|e| { + io::Error::new( + io::ErrorKind::PermissionDenied, + format!("Session ID decrypt failed: {:?}", e), + ) + })?; + + // Validate session ID contents + // Bytes 0-2: Version + // Byte 3: Padding + // Bytes 4-7: Timestamp (uint32) + // Bytes 8-15: ShortId + let client_version = &decrypted_session_id[0..3]; + let client_timestamp = u32::from_be_bytes([ + decrypted_session_id[4], + decrypted_session_id[5], + decrypted_session_id[6], + decrypted_session_id[7], + ]) as u64; + let client_short_id = &decrypted_session_id[8..16]; + + log::debug!("REALITY: Client version: {:?}", client_version); + log::debug!("REALITY: Client timestamp: {}", client_timestamp); + log::debug!("REALITY: Client short_id: {:02x?}", client_short_id); + + // Validate short ID - check if client's short_id is in the configured list + let mut client_short_id_arr = [0u8; 8]; + client_short_id_arr.copy_from_slice(client_short_id); + let short_id_ok = self.config.short_ids.contains(&client_short_id_arr); + + if !short_id_ok { + log::warn!( + "REALITY: Client short_id {:02x?} not in configured list", + client_short_id + ); + return Err(io::Error::new( + io::ErrorKind::PermissionDenied, + format!("Invalid short_id: {:02x?}", client_short_id), + )); + } + + // Validate timestamp if max_time_diff is configured + if let Some(max_diff_ms) = self.config.max_time_diff { + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map_err(|_| io::Error::other("System time error"))? + .as_secs(); + + let time_diff_secs = now.abs_diff(client_timestamp); + let max_diff_secs = max_diff_ms / 1000; + + if time_diff_secs > max_diff_secs { + log::warn!( + "REALITY: Client timestamp {} differs from server {} by {} seconds (max: {} seconds)", + client_timestamp, now, time_diff_secs, max_diff_secs + ); + return Err(io::Error::new( + io::ErrorKind::PermissionDenied, + format!( + "Timestamp difference {} seconds exceeds maximum {} seconds", + time_diff_secs, max_diff_secs + ), + )); + } + } + + // Validate client version (min) + if let Some(min_ver) = &self.config.min_client_version { + if client_version < &min_ver[..] { + log::warn!( + "REALITY: Client version {:?} is below minimum {:?}", + client_version, + min_ver + ); + return Err(io::Error::new( + io::ErrorKind::PermissionDenied, + format!( + "Client version {:?} is below minimum {:?}", + client_version, min_ver + ), + )); + } + } + + // Validate client version (max) + if let Some(max_ver) = &self.config.max_client_version { + if client_version > &max_ver[..] { + log::warn!( + "REALITY: Client version {:?} is above maximum {:?}", + client_version, + max_ver + ); + return Err(io::Error::new( + io::ErrorKind::PermissionDenied, + format!( + "Client version {:?} is above maximum {:?}", + client_version, max_ver + ), + )); + } + } + + log::info!("REALITY: Client authentication successful - short_id: {:02x?}, version: {:?}, timestamp: {}", + client_short_id, client_version, client_timestamp); + + // Step 4: Generate our server X25519 keypair + let rng = SystemRandom::new(); + let mut our_private_bytes = [0u8; 32]; + rng.fill(&mut our_private_bytes) + .map_err(|_| io::Error::other("RNG failed"))?; + + let our_private_key = + agreement::PrivateKey::from_private_key(&agreement::X25519, &our_private_bytes) + .map_err(|_| io::Error::other("Failed to create X25519 key"))?; + let our_public_key_bytes = our_private_key + .compute_public_key() + .map_err(|_| io::Error::other("Failed to compute public key"))?; + + // Step 5: Generate server random + let mut server_random = [0u8; 32]; + rng.fill(&mut server_random) + .map_err(|_| io::Error::other("RNG failed"))?; + + // Step 6: Use standard cipher suite (TLS_AES_128_GCM_SHA256) + const CIPHER_SUITE: u16 = 0x1301; + + // Step 7: Build ServerHello + let server_hello = construct_server_hello( + &server_random, + session_id, + CIPHER_SUITE, + our_public_key_bytes.as_ref(), + )?; + + // Step 8: Compute transcript hashes + let client_hello_handshake = &client_hello[TLS_RECORD_HEADER_SIZE..]; // Skip TLS record header + + let mut ch_transcript = digest::Context::new(&digest::SHA256); + ch_transcript.update(client_hello_handshake); + let client_hello_hash = ch_transcript.finish(); + + let mut ch_sh_transcript = digest::Context::new(&digest::SHA256); + ch_sh_transcript.update(client_hello_handshake); + ch_sh_transcript.update(&server_hello); + + // Clone before finalizing + let mut handshake_transcript = ch_sh_transcript.clone(); + let server_hello_hash = ch_sh_transcript.finish(); + + // Step 9: Perform ECDH for TLS 1.3 key derivation + let peer_public_key = + agreement::UnparsedPublicKey::new(&agreement::X25519, &client_public_key); + let mut tls_shared_secret = [0u8; 32]; + agreement::agree( + &our_private_key, + peer_public_key, + io::Error::other("ECDH failed"), + |key_material| { + tls_shared_secret.copy_from_slice(key_material); + Ok(()) + }, + )?; + + // Step 10: Derive TLS 1.3 keys + let mut client_hello_hash_arr = [0u8; 32]; + client_hello_hash_arr.copy_from_slice(client_hello_hash.as_ref()); + let mut server_hello_hash_arr = [0u8; 32]; + server_hello_hash_arr.copy_from_slice(server_hello_hash.as_ref()); + + let hs_keys = derive_handshake_keys( + &tls_shared_secret, + &client_hello_hash_arr, + &server_hello_hash_arr, + )?; + + // Get destination hostname for certificate + let dest_hostname = match self.config.dest.address() { + Address::Hostname(h) => h.as_str(), + _ => { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "REALITY requires a hostname destination", + )) + } + }; + + // Step 11: Generate HMAC-signed certificate + let (cert_der, signing_key) = generate_hmac_certificate(&auth_key, dest_hostname)?; + + // Step 12: Build encrypted handshake messages + let encrypted_extensions = construct_encrypted_extensions()?; + handshake_transcript.update(&encrypted_extensions); + + let certificate = construct_certificate(&cert_der)?; + handshake_transcript.update(&certificate); + + let cert_verify_hash = handshake_transcript.clone().finish(); + let mut cert_verify_hash_arr = [0u8; 32]; + cert_verify_hash_arr.copy_from_slice(cert_verify_hash.as_ref()); + let certificate_verify = construct_certificate_verify(&signing_key, &cert_verify_hash_arr)?; + handshake_transcript.update(&certificate_verify); + + let handshake_hash_before_finished = handshake_transcript.clone().finish(); + let mut handshake_hash_arr = [0u8; 32]; + handshake_hash_arr.copy_from_slice(handshake_hash_before_finished.as_ref()); + + // Step 13: Derive server handshake traffic keys for encryption + let (server_hs_key, server_hs_iv) = + derive_traffic_keys(&hs_keys.server_handshake_traffic_secret, CIPHER_SUITE)?; + + // Step 14: Build server Finished message first (before encryption) + let server_verify_data = compute_finished_verify_data( + &hs_keys.server_handshake_traffic_secret, + &handshake_hash_arr, + )?; + let server_finished = construct_finished(&server_verify_data)?; + + // Step 15: Combine all 4 handshake messages into one plaintext buffer + // This is required by REALITY protocol - all 4 messages in one encrypted record + let mut combined_plaintext = Vec::new(); + combined_plaintext.extend_from_slice(&encrypted_extensions); + combined_plaintext.extend_from_slice(&certificate); + combined_plaintext.extend_from_slice(&certificate_verify); + combined_plaintext.extend_from_slice(&server_finished); + + log::debug!("REALITY SERVER: Combining 4 handshake messages: EE={}, Cert={}, CV={}, Fin={}, Total={}", + encrypted_extensions.len(), certificate.len(), certificate_verify.len(), + server_finished.len(), combined_plaintext.len()); + + // Step 16: Encrypt the combined handshake messages, fragmenting into multiple records if needed + // TLS 1.3 limits records to 16,640 bytes ciphertext (16,623 bytes plaintext). + // Large certificates can exceed this, so we fragment like uTLS does. + let mut handshake_ciphertext = Vec::new(); + let mut handshake_seq = 0u64; + encrypt_handshake_to_records( + &combined_plaintext, + &server_hs_key, + &server_hs_iv, + &mut handshake_seq, + &mut handshake_ciphertext, + )?; + + // Update transcript with server Finished (needed for client Finished verification) + handshake_transcript.update(&server_finished); + let handshake_hash_with_server_finished = handshake_transcript.finish(); + let mut handshake_hash_with_finished_arr = [0u8; 32]; + handshake_hash_with_finished_arr + .copy_from_slice(handshake_hash_with_server_finished.as_ref()); + + // Step 16: Buffer all handshake messages to write buffer + // ServerHello (plaintext) + self.ciphertext_write_buf + .extend_from_slice(&write_record_header( + CONTENT_TYPE_HANDSHAKE, + server_hello.len() as u16, + )); + self.ciphertext_write_buf.extend_from_slice(&server_hello); + + // ChangeCipherSpec (for compatibility) + self.ciphertext_write_buf + .extend_from_slice(&write_record_header(CONTENT_TYPE_CHANGE_CIPHER_SPEC, 1)); + self.ciphertext_write_buf.push(0x01); + + // Encrypted handshake record(s) - may be fragmented into multiple records + self.ciphertext_write_buf + .extend_from_slice(&handshake_ciphertext); + + log::info!( + "REALITY: ServerHello and encrypted handshake messages buffered ({} bytes)", + self.ciphertext_write_buf.len() + ); + + // Step 17: Update handshake state + self.handshake_state = HandshakeState::ServerHelloSent { + handshake_hash: handshake_hash_arr, + handshake_hash_with_server_finished: handshake_hash_with_finished_arr, + client_handshake_traffic_secret: hs_keys.client_handshake_traffic_secret.clone(), + master_secret: hs_keys.master_secret, + cipher_suite: CIPHER_SUITE, + }; + + Ok(()) + } + + /// Process client's Finished message and complete handshake + fn process_client_finished(&mut self) -> io::Result<()> { + // Check if we have enough data for a TLS record header BEFORE extracting state + if self.ciphertext_read_buf.len() < TLS_RECORD_HEADER_SIZE { + return Ok(()); // Need more data + } + + // Check for ChangeCipherSpec (TLS 1.3 compatibility message) + if self.ciphertext_read_buf[0] == CONTENT_TYPE_CHANGE_CIPHER_SPEC { + // ChangeCipherSpec record + let ccs_len = self + .ciphertext_read_buf + .get_u16_be(3) + .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 { + return Ok(()); // Need more data + } + + // Skip ChangeCipherSpec (compatibility message) + log::debug!("REALITY: Skipping ChangeCipherSpec (compatibility message)"); + self.ciphertext_read_buf + .consume(TLS_RECORD_HEADER_SIZE + ccs_len); + + // Check if we have the next record header + if self.ciphertext_read_buf.len() < TLS_RECORD_HEADER_SIZE { + return Ok(()); // Need more data + } + } + + // Parse TLS record length + let record_len = self + .ciphertext_read_buf + .get_u16_be(3) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "Buffer too short"))? + as usize; + + // Check if we have the complete record + let total_record_len = TLS_RECORD_HEADER_SIZE + record_len; + if self.ciphertext_read_buf.len() < total_record_len { + return Ok(()); // Need more data + } + + // Verify it's ApplicationData (encrypted Finished) + if self.ciphertext_read_buf[0] != CONTENT_TYPE_APPLICATION_DATA { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "Expected ApplicationData (0x17), got 0x{:02x}", + self.ciphertext_read_buf[0] + ), + )); + } + + // NOW we're committed to processing - take ownership of handshake state + // This avoids cloning Vec fields + let old_state = std::mem::replace(&mut self.handshake_state, HandshakeState::Complete); + let ( + client_handshake_traffic_secret, + master_secret, + cipher_suite, + _handshake_hash, + handshake_hash_with_server_finished, + ) = match old_state { + HandshakeState::ServerHelloSent { + client_handshake_traffic_secret, + master_secret, + cipher_suite, + handshake_hash, + handshake_hash_with_server_finished, + } => ( + client_handshake_traffic_secret, // moved, not cloned + master_secret, // moved, not cloned + cipher_suite, + handshake_hash, + handshake_hash_with_server_finished, + ), + _ => { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "Invalid state for process_client_finished", + )) + } + }; + + // Extract the encrypted Finished record (copy to Vec for decryption) + let record: Vec = self.ciphertext_read_buf[..total_record_len].to_vec(); + self.ciphertext_read_buf.consume(total_record_len); + let ciphertext = &record[TLS_RECORD_HEADER_SIZE..]; // Skip TLS record header + + // Derive client handshake traffic keys for decryption + let (client_hs_key, client_hs_iv) = + derive_traffic_keys(&client_handshake_traffic_secret, cipher_suite)?; + + // Decrypt the Finished message (sequence number = 0 for client's first encrypted record) + let plaintext = decrypt_handshake_message( + &client_hs_key, + &client_hs_iv, + 0, // Client's first encrypted record + ciphertext, + record_len as u16, + )?; + + // Verify it's a Finished message (type 0x14) + if plaintext.is_empty() || plaintext[0] != HANDSHAKE_TYPE_FINISHED { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "Expected Finished message", + )); + } + + // Extract verify_data (skip type(1) + length(3) = 4 bytes) + if plaintext.len() < 4 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "Finished message too short", + )); + } + let client_verify_data = &plaintext[4..]; + + // Compute expected client Finished verify_data + // IMPORTANT: Use hash that includes server Finished (per TLS 1.3 RFC 8446) + let expected_verify_data = compute_finished_verify_data( + &client_handshake_traffic_secret, + &handshake_hash_with_server_finished, + )?; + + // Verify it matches using constant-time comparison to prevent timing attacks + if client_verify_data + .ct_eq(expected_verify_data.as_slice()) + .unwrap_u8() + == 0 + { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "Client Finished verify_data mismatch", + )); + } + + log::debug!("REALITY: Client Finished verified successfully"); + + // Derive application secrets + // IMPORTANT: Use hash that includes server Finished (per TLS 1.3 RFC 8446) + let (client_app_secret, server_app_secret) = + derive_application_secrets(&master_secret, &handshake_hash_with_server_finished)?; + + // Derive application traffic keys + let (client_app_key, client_app_iv) = + derive_traffic_keys(&client_app_secret, cipher_suite)?; + let (server_app_key, server_app_iv) = + derive_traffic_keys(&server_app_secret, cipher_suite)?; + + // Store application traffic keys + self.app_read_key = Some(client_app_key); + self.app_read_iv = Some(client_app_iv); + self.app_write_key = Some(server_app_key); + self.app_write_iv = Some(server_app_iv); + self.read_seq = 0; + self.write_seq = 0; + self.cipher_suite = cipher_suite; + + // Handshake state already set to Complete above + + log::debug!("REALITY: Handshake complete, application keys derived"); + + Ok(()) + } + + /// Decrypt application data using TLS 1.3 keys + fn process_application_data(&mut self) -> io::Result<()> { + // Check if we have application keys + let (app_read_key, app_read_iv) = match (&self.app_read_key, &self.app_read_iv) { + (Some(key), Some(iv)) => (key, iv), + _ => return Ok(()), // Keys not ready yet + }; + + // Process all complete TLS records in the buffer + while self.ciphertext_read_buf.len() >= 5 { + // Parse TLS record header + let record_len = self + .ciphertext_read_buf + .get_u16_be(3) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "Buffer too short"))? + as usize; + + // Check if we have the complete record + let total_record_len = TLS_RECORD_HEADER_SIZE + record_len; + if self.ciphertext_read_buf.len() < total_record_len { + break; // Need more data + } + + // Copy record header for decryption AAD + let tls_header: [u8; TLS_RECORD_HEADER_SIZE] = [ + self.ciphertext_read_buf[0], + self.ciphertext_read_buf[1], + self.ciphertext_read_buf[2], + self.ciphertext_read_buf[3], + self.ciphertext_read_buf[4], + ]; + let ciphertext: Vec = + self.ciphertext_read_buf[TLS_RECORD_HEADER_SIZE..total_record_len].to_vec(); + self.ciphertext_read_buf.consume(total_record_len); + + // Decrypt the application data + let mut plaintext = decrypt_tls13_record( + app_read_key, + app_read_iv, + self.read_seq, + &ciphertext, + &tls_header, + )?; + + self.read_seq += 1; + + // TLS 1.3: Remove ContentType trailer byte + // The last byte of the plaintext is the actual ContentType + if !plaintext.is_empty() { + let content_type = plaintext.pop().unwrap(); + + // For now, we only handle ApplicationData (0x17) and Alert (0x15) + // Alerts should be handled separately, but for now just pass through + if content_type != CONTENT_TYPE_APPLICATION_DATA + && content_type != CONTENT_TYPE_ALERT + { + log::warn!("REALITY: Unexpected ContentType: 0x{:02x}", content_type); + } + } + + // Compact plaintext buffer if needed before extending + self.plaintext_read_buf.maybe_compact(4096); + + // Append to plaintext buffer (without ContentType) + self.plaintext_read_buf.extend_from_slice(&plaintext); + } + + Ok(()) + } + + /// Get a reader for accessing decrypted plaintext + pub fn reader(&mut self) -> RealityReader<'_> { + // SlideBuffer handles compaction internally via maybe_compact() + // Compact before returning reader if we've consumed significant data + self.plaintext_read_buf.maybe_compact(4096); + RealityReader::new(&mut self.plaintext_read_buf) + } + + /// Get a writer for buffering plaintext to be encrypted + pub fn writer(&mut self) -> RealityWriter<'_> { + RealityWriter::new(&mut self.plaintext_write_buf) + } + + /// Write buffered TLS messages to the provided writer + /// + /// This encrypts any pending plaintext and writes ciphertext. + /// Large plaintext is automatically fragmented into multiple TLS records + /// to comply with the TLS 1.3 record size limit. + pub fn write_tls(&mut self, wr: &mut dyn Write) -> io::Result { + // If handshake not complete, just write buffered handshake data + if !matches!(self.handshake_state, HandshakeState::Complete) { + let n = wr.write(&self.ciphertext_write_buf)?; + self.ciphertext_write_buf.drain(..n); + return Ok(n); + } + + // Encrypt any pending plaintext (with automatic fragmentation for large data) + if !self.plaintext_write_buf.is_empty() { + let (app_write_key, app_write_iv) = match (&self.app_write_key, &self.app_write_iv) { + (Some(key), Some(iv)) => (key, iv), + _ => { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "Application keys not available", + )) + } + }; + + encrypt_plaintext_to_records( + &mut self.plaintext_write_buf, + app_write_key, + app_write_iv, + &mut self.write_seq, + &mut self.ciphertext_write_buf, + )?; + } + + // Write buffered ciphertext + let n = wr.write(&self.ciphertext_write_buf)?; + self.ciphertext_write_buf.drain(..n); + Ok(n) + } + + /// Check if the connection wants to write data + pub fn wants_write(&self) -> bool { + !self.ciphertext_write_buf.is_empty() || !self.plaintext_write_buf.is_empty() + } + + /// Check if handshake is still in progress + pub fn is_handshaking(&self) -> bool { + !matches!(self.handshake_state, HandshakeState::Complete) + } + + /// Queue a close notification alert + pub fn send_close_notify(&mut self) { + // In TLS 1.3, alerts must be encrypted like application data + if !matches!(self.handshake_state, HandshakeState::Complete) { + log::debug!("REALITY: Cannot send close_notify - handshake not complete"); + return; + } + + // Get application keys + let (app_write_key, app_write_iv) = match (&self.app_write_key, &self.app_write_iv) { + (Some(key), Some(iv)) => (key, iv), + _ => { + log::debug!("REALITY: Cannot send close_notify - application keys not available"); + return; + } + }; + + // Use common helper to build encrypted close_notify alert + match common::build_close_notify_alert(app_write_key, app_write_iv, self.write_seq) { + Ok(record) => { + self.write_seq += 1; + self.ciphertext_write_buf.extend_from_slice(&record); + log::debug!("REALITY: Encrypted close_notify alert queued"); + } + Err(e) => { + log::error!("REALITY: Failed to encrypt close_notify: {}", e); + } + } + } +} + +#[inline(always)] +pub fn feed_reality_server_connection( + server_connection: &mut RealityServerConnection, + data: &[u8], +) -> std::io::Result<()> { + let mut cursor = std::io::Cursor::new(data); + let mut i = 0; + while i < data.len() { + let n = server_connection.read_tls(&mut cursor).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("failed to feed TLS connection: {e}"), + ) + })?; + i += n; + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_reality_server_connection_creation() { + let config = RealityServerConfig { + private_key: [0u8; 32], + short_ids: vec![[0u8; 8]], + dest: NetLocation::new(Address::UNSPECIFIED, 443), + max_time_diff: Some(60000), + min_client_version: None, + max_client_version: None, + }; + + let conn = RealityServerConnection::new(config).unwrap(); + assert!(conn.is_handshaking()); + assert!(!conn.wants_write()); + } + + #[test] + fn test_io_state() { + let config = RealityServerConfig { + private_key: [0u8; 32], + short_ids: vec![[0u8; 8]], + dest: NetLocation::new(Address::UNSPECIFIED, 443), + max_time_diff: None, + min_client_version: None, + max_client_version: None, + }; + + let mut conn = RealityServerConnection::new(config).unwrap(); + let state = conn.process_new_packets().unwrap(); + + assert_eq!(state.plaintext_bytes_to_read(), 0); + assert!(!conn.wants_write()); + } +} diff --git a/src/reality/reality_tls13_keys.rs b/src/reality/reality_tls13_keys.rs new file mode 100644 index 00000000..8ae47e48 --- /dev/null +++ b/src/reality/reality_tls13_keys.rs @@ -0,0 +1,429 @@ +// TLS 1.3 Key Schedule Implementation +// +// Implements RFC 8446 key derivation for TLS 1.3 +// Using HKDF-SHA256 for all key derivation operations + +use aws_lc_rs::digest; +use std::io::{Error, ErrorKind, Result}; + +/// Intermediate TLS 1.3 keys (handshake secrets + master secret) +/// Used for two-phase key derivation where application secrets +/// must be derived after server Finished message +#[derive(Debug, Clone)] +pub struct Tls13HandshakeKeys { + /// Client handshake traffic secret + pub client_handshake_traffic_secret: Vec, + /// Server handshake traffic secret + pub server_handshake_traffic_secret: Vec, + /// Master secret (for deriving application secrets later) + pub master_secret: Vec, +} + +/// HKDF-Expand implementation using HMAC-SHA256 directly +/// This follows RFC 5869 Section 2.3 +pub fn hkdf_expand_sha256(prk: &[u8], info: &[u8], length: usize) -> Result> { + use aws_lc_rs::hmac; + + const HASH_LEN: usize = 32; // SHA256 output length + let n = length.div_ceil(HASH_LEN); // Number of iterations + + if n > 255 { + return Err(Error::new(ErrorKind::InvalidData, "HKDF output too long")); + } + + let mut output = Vec::new(); + let mut prev = Vec::new(); + + for i in 1..=n { + let key = hmac::Key::new(hmac::HMAC_SHA256, prk); + let mut ctx = hmac::Context::with_key(&key); + + log::debug!( + "HKDF iteration {}: prev_len={}, info_len={}", + i, + prev.len(), + info.len() + ); + + ctx.update(&prev); + ctx.update(info); + ctx.update(&[i as u8]); + let tag = ctx.sign(); + + log::debug!( + "HKDF iteration {}: output={:02x?}", + i, + &tag.as_ref()[..tag.as_ref().len().min(16)] + ); + + prev = tag.as_ref().to_vec(); + output.extend_from_slice(tag.as_ref()); + } + + output.truncate(length); + Ok(output) +} + +/// HKDF-Expand-Label as defined in RFC 8446 Section 7.1 +fn hkdf_expand_label( + secret: &[u8], + label: &[u8], + context: &[u8], + length: usize, +) -> Result> { + log::debug!( + "DEBUG hkdf_expand_label: secret len={}, label={:?}, context len={}, length={}", + secret.len(), + std::str::from_utf8(label).unwrap_or(""), + context.len(), + length + ); + // HkdfLabel structure: + // struct { + // uint16 length = Length; + // opaque label<7..255> = "tls13 " + Label; + // opaque context<0..255> = Context; + // } HkdfLabel; + + let mut hkdf_label = Vec::new(); + + // Length (2 bytes, big-endian) + hkdf_label.extend_from_slice(&(length as u16).to_be_bytes()); + + // Label length and content + let full_label = format!("tls13 {}", std::str::from_utf8(label).unwrap()); + hkdf_label.push(full_label.len() as u8); + hkdf_label.extend_from_slice(full_label.as_bytes()); + + // Context length and content + hkdf_label.push(context.len() as u8); + hkdf_label.extend_from_slice(context); + + log::debug!("HKDF_LABEL_BYTES: {:02x?}", hkdf_label); + + // Use our helper function + hkdf_expand_sha256(secret, &hkdf_label, length) +} + +/// Derive-Secret as defined in RFC 8446 Section 7.1 +fn derive_secret(secret: &[u8], label: &[u8], messages_hash: &[u8]) -> Result> { + hkdf_expand_label(secret, label, messages_hash, 32) +} + +/// HKDF-Extract operation +fn hkdf_extract(salt: &[u8], ikm: &[u8]) -> Vec { + // Use hmac directly for extract operation since aws-lc-rs doesn't expose PRK bytes + use aws_lc_rs::hmac; + let key = hmac::Key::new(hmac::HMAC_SHA256, salt); + let tag = hmac::sign(&key, ikm); + tag.as_ref().to_vec() +} + +/// Derive TLS 1.3 handshake keys and master secret (Phase 1) +/// +/// This function derives handshake traffic secrets and the master secret, +/// but NOT the application traffic secrets. Application secrets must be +/// derived separately after the server Finished message is sent (Phase 2). +/// +/// # Arguments +/// * `shared_secret` - ECDH shared secret (32 bytes for X25519) +/// * `client_hello_hash` - SHA256 hash of ClientHello (32 bytes) +/// * `server_hello_hash` - SHA256 hash of ClientHello...ServerHello (32 bytes) +/// +/// # Returns +/// Handshake traffic secrets and master secret +pub fn derive_handshake_keys( + shared_secret: &[u8], + client_hello_hash: &[u8], + server_hello_hash: &[u8], +) -> Result { + // Validate input lengths + if shared_secret.len() != 32 { + return Err(Error::new( + ErrorKind::InvalidInput, + format!( + "Invalid shared_secret length: {} (expected 32)", + shared_secret.len() + ), + )); + } + if client_hello_hash.len() != 32 || server_hello_hash.len() != 32 { + return Err(Error::new( + ErrorKind::InvalidInput, + "All hashes must be 32 bytes (SHA256)", + )); + } + + log::debug!("TLS13 DEBUG: Deriving handshake keys (Phase 1)..."); + + // 1. Early Secret = HKDF-Extract(salt=0, IKM=0) + let zero_salt = vec![0u8; 32]; + let early_secret = hkdf_extract(&zero_salt, &zero_salt); + + // 2. Derive-Secret(., "derived", "") + let mut empty_ctx = digest::Context::new(&digest::SHA256); + empty_ctx.update(b""); + let empty_hash = empty_ctx.finish(); + let derived_secret = derive_secret(&early_secret, b"derived", empty_hash.as_ref())?; + + // 3. Handshake Secret = HKDF-Extract(salt=derived_secret, IKM=shared_secret) + let handshake_secret = hkdf_extract(&derived_secret, shared_secret); + + // 4. Client Handshake Traffic Secret + let client_handshake_traffic_secret = + derive_secret(&handshake_secret, b"c hs traffic", server_hello_hash)?; + + // 5. Server Handshake Traffic Secret + let server_handshake_traffic_secret = + derive_secret(&handshake_secret, b"s hs traffic", server_hello_hash)?; + + // 6. Derive-Secret(., "derived", "") + let mut empty_ctx_2 = digest::Context::new(&digest::SHA256); + empty_ctx_2.update(b""); + let empty_hash_2 = empty_ctx_2.finish(); + let derived_secret_2 = derive_secret(&handshake_secret, b"derived", empty_hash_2.as_ref())?; + + // 7. Master Secret = HKDF-Extract(salt=derived_secret, IKM=0) + let master_secret = hkdf_extract(&derived_secret_2, &zero_salt); + + log::debug!(" master_secret: {:?}", &master_secret[..8]); + + Ok(Tls13HandshakeKeys { + client_handshake_traffic_secret, + server_handshake_traffic_secret, + master_secret, + }) +} + +/// Derive TLS 1.3 application traffic secrets (Phase 2) +/// +/// This function must be called AFTER the server Finished message is sent, +/// with a transcript hash that includes the Finished message. +/// +/// # Arguments +/// * `master_secret` - Master secret from Phase 1 +/// * `handshake_hash` - SHA256 hash including server Finished (32 bytes) +/// +/// # Returns +/// (client_application_traffic_secret, server_application_traffic_secret) +pub fn derive_application_secrets( + master_secret: &[u8], + handshake_hash: &[u8], +) -> Result<(Vec, Vec)> { + if master_secret.len() != 32 || handshake_hash.len() != 32 { + return Err(Error::new( + ErrorKind::InvalidInput, + "Master secret and handshake hash must be 32 bytes", + )); + } + + log::debug!("TLS13 DEBUG: Deriving application secrets (Phase 2)..."); + log::info!( + " handshake_hash (with Finished): {:?}", + &handshake_hash[..8] + ); + + // Client Application Traffic Secret + let client_application_traffic_secret = + derive_secret(master_secret, b"c ap traffic", handshake_hash)?; + + log::debug!( + " client_app_traffic: {:?}", + &client_application_traffic_secret[..8] + ); + log::info!( + "DERIVE_APP_SECRETS: ClientAppSecret(full)={:02x?}", + client_application_traffic_secret + ); + + // Server Application Traffic Secret + let server_application_traffic_secret = + derive_secret(master_secret, b"s ap traffic", handshake_hash)?; + + log::debug!( + " server_app_traffic: {:?}", + &server_application_traffic_secret[..8] + ); + log::info!( + "DERIVE_APP_SECRETS: ServerAppSecret(full)={:02x?}", + server_application_traffic_secret + ); + + Ok(( + client_application_traffic_secret, + server_application_traffic_secret, + )) +} + +/// Derive traffic keys and IV from traffic secret +/// +/// # Arguments +/// * `traffic_secret` - Traffic secret (32 bytes) +/// * `cipher_suite` - TLS cipher suite (e.g., 0x1301 for TLS_AES_128_GCM_SHA256) +/// +/// # Returns +/// (key, iv) tuple for AES-GCM +pub fn derive_traffic_keys(traffic_secret: &[u8], cipher_suite: u16) -> Result<(Vec, Vec)> { + // For TLS_AES_128_GCM_SHA256 (0x1301): + // - key_length = 16 bytes + // - iv_length = 12 bytes + + let (key_length, iv_length) = match cipher_suite { + 0x1301 => (16, 12), // TLS_AES_128_GCM_SHA256 + 0x1302 => (32, 12), // TLS_AES_256_GCM_SHA384 + 0x1303 => (32, 12), // TLS_CHACHA20_POLY1305_SHA256 + _ => { + return Err(Error::new( + ErrorKind::InvalidInput, + format!("Unsupported cipher suite: 0x{:04x}", cipher_suite), + )) + } + }; + + log::debug!( + "TRAFFIC_KEY_DERIVE: cipher_suite=0x{:04x}, key_len={}, iv_len={}", + cipher_suite, + key_length, + iv_length + ); + log::debug!("TRAFFIC_KEY_DERIVE: traffic_secret={:02x?}", traffic_secret); + + // key = HKDF-Expand-Label(Secret, "key", "", key_length) + let key = hkdf_expand_label(traffic_secret, b"key", b"", key_length)?; + + // iv = HKDF-Expand-Label(Secret, "iv", "", iv_length) + let iv = hkdf_expand_label(traffic_secret, b"iv", b"", iv_length)?; + + log::info!("TRAFFIC_KEY_DERIVE: key={:02x?}", key); + log::info!("TRAFFIC_KEY_DERIVE: iv={:02x?}", iv); + + Ok((key, iv)) +} + +/// Compute "Finished" verify data +/// +/// finished_key = HKDF-Expand-Label(BaseKey, "finished", "", Hash.length) +/// verify_data = HMAC(finished_key, Transcript-Hash(Handshake Context)) +pub fn compute_finished_verify_data(base_key: &[u8], handshake_hash: &[u8]) -> Result> { + use aws_lc_rs::hmac; + + // finished_key = HKDF-Expand-Label(BaseKey, "finished", "", 32) + let finished_key = hkdf_expand_label(base_key, b"finished", b"", 32)?; + + // verify_data = HMAC(finished_key, handshake_hash) + let key = hmac::Key::new(hmac::HMAC_SHA256, &finished_key); + let tag = hmac::sign(&key, handshake_hash); + let verify_data = tag.as_ref().to_vec(); + + Ok(verify_data) +} + +#[cfg(test)] +mod tests { + use super::*; + + // Test vectors from RFC 5869 Appendix A + #[test] + fn test_hkdf_expand_sha256_rfc_vector() { + // Test Case 1 from RFC 5869 (simplified - using extracted PRK directly) + let prk = [ + 0x07, 0x77, 0x09, 0x36, 0x2c, 0x2e, 0x32, 0xdf, 0x0d, 0xdc, 0x3f, 0x0d, 0xc4, 0x7b, + 0xba, 0x63, 0x90, 0xb6, 0xc7, 0x3b, 0xb5, 0x0f, 0x9c, 0x31, 0x22, 0xec, 0x84, 0x4a, + 0xd7, 0xc2, 0xb3, 0xe5, + ]; + let info = [0xf0, 0xf1, 0xf2, 0xf3, 0xf4, 0xf5, 0xf6, 0xf7, 0xf8, 0xf9]; + + let result = hkdf_expand_sha256(&prk, &info, 42).unwrap(); + assert_eq!(result.len(), 42); + + // Check first few bytes match expected pattern + assert_eq!(result[0], 0x3c); + assert_eq!(result[1], 0xb2); + assert_eq!(result[2], 0x5f); + } + + #[test] + fn test_hkdf_expand_sha256_empty_info() { + let prk = vec![0x42u8; 32]; + let result = hkdf_expand_sha256(&prk, &[], 16); + assert!(result.is_ok()); + assert_eq!(result.unwrap().len(), 16); + } + + #[test] + fn test_hkdf_expand_sha256_max_length() { + let prk = vec![0x42u8; 32]; + let info = b"test info"; + + // Maximum output length is 255 * hash_len (32 for SHA256) = 8160 bytes + let result = hkdf_expand_sha256(&prk, info, 8160); + assert!(result.is_ok()); + assert_eq!(result.unwrap().len(), 8160); + + // Should fail for length > 8160 + let result = hkdf_expand_sha256(&prk, info, 8161); + assert!(result.is_err()); + } + + #[test] + fn test_hkdf_expand_label() { + let secret = vec![0x42u8; 32]; + let result = hkdf_expand_label(&secret, b"test", b"", 16); + assert!(result.is_ok()); + assert_eq!(result.unwrap().len(), 16); + } + + #[test] + fn test_hkdf_expand_label_with_context() { + let secret = vec![0x42u8; 32]; + let context = vec![0x11u8; 32]; + let result = hkdf_expand_label(&secret, b"finished", &context, 32); + assert!(result.is_ok()); + let output = result.unwrap(); + assert_eq!(output.len(), 32); + + // Result should be deterministic + let result2 = hkdf_expand_label(&secret, b"finished", &context, 32).unwrap(); + assert_eq!(output, result2); + } + + #[test] + fn test_hkdf_extract() { + let salt = vec![0x11u8; 32]; + let ikm = vec![0x22u8; 32]; + + let result1 = hkdf_extract(&salt, &ikm); + assert_eq!(result1.len(), 32); // SHA256 output length + + // Should be deterministic + let result2 = hkdf_extract(&salt, &ikm); + assert_eq!(result1, result2); + + // Different input should give different output + let ikm2 = vec![0x33u8; 32]; + let result3 = hkdf_extract(&salt, &ikm2); + assert_ne!(result1, result3); + } + + #[test] + fn test_derive_traffic_keys() { + let traffic_secret = vec![0x99u8; 32]; + + // Test TLS_AES_128_GCM_SHA256 + let result = derive_traffic_keys(&traffic_secret, 0x1301); + assert!(result.is_ok()); + let (key, iv) = result.unwrap(); + assert_eq!(key.len(), 16); + assert_eq!(iv.len(), 12); + } + + #[test] + fn test_compute_finished_verify_data() { + let base_key = vec![0xAAu8; 32]; + let handshake_hash = vec![0xBBu8; 32]; + + let result = compute_finished_verify_data(&base_key, &handshake_hash); + assert!(result.is_ok()); + assert_eq!(result.unwrap().len(), 32); + } +} diff --git a/src/reality/reality_tls13_messages.rs b/src/reality/reality_tls13_messages.rs new file mode 100644 index 00000000..98c5a3be --- /dev/null +++ b/src/reality/reality_tls13_messages.rs @@ -0,0 +1,431 @@ +// TLS 1.3 Message Construction +// +// Construct TLS 1.3 handshake messages for REALITY protocol + +use super::common::{ + HANDSHAKE_TYPE_CERTIFICATE, HANDSHAKE_TYPE_CERTIFICATE_VERIFY, + HANDSHAKE_TYPE_ENCRYPTED_EXTENSIONS, HANDSHAKE_TYPE_FINISHED, HANDSHAKE_TYPE_SERVER_HELLO, + VERSION_TLS_1_2_MAJOR, VERSION_TLS_1_2_MINOR, +}; +use aws_lc_rs::signature::Ed25519KeyPair; +use std::io::Result; + +/// Construct ServerHello message +/// +/// # Arguments +/// * `server_random` - 32 bytes of server random +/// * `session_id` - Session ID from ClientHello (for compatibility) +/// * `cipher_suite` - Selected cipher suite (e.g., 0x1301) +/// * `key_share_data` - Server's X25519 public key (32 bytes) +pub fn construct_server_hello( + server_random: &[u8; 32], + session_id: &[u8], + cipher_suite: u16, + key_share_data: &[u8], +) -> Result> { + let mut server_hello = Vec::new(); + + // ServerHello structure: + // - handshake_type (1 byte) = 2 + // - length (3 bytes) + // - version (2 bytes) = 0x0303 (TLS 1.2 for compatibility) + // - random (32 bytes) + // - session_id_length (1 byte) + // - session_id (variable) + // - cipher_suite (2 bytes) + // - compression_method (1 byte) = 0 + // - extensions_length (2 bytes) + // - extensions (variable) + + let mut payload = Vec::new(); + + // Version: 0x0303 (TLS 1.2 for compatibility) + payload.extend_from_slice(&[VERSION_TLS_1_2_MAJOR, VERSION_TLS_1_2_MINOR]); + + // Random (32 bytes) + payload.extend_from_slice(server_random); + + // Session ID + payload.push(session_id.len() as u8); + payload.extend_from_slice(session_id); + + // Cipher suite + payload.extend_from_slice(&cipher_suite.to_be_bytes()); + + // Compression method = 0 + payload.push(0x00); + + // Extensions + let mut extensions = Vec::new(); + + // supported_versions extension (type=43) + extensions.extend_from_slice(&[0x00, 0x2b]); // type = 43 + extensions.extend_from_slice(&[0x00, 0x02]); // length = 2 + extensions.extend_from_slice(&[0x03, 0x04]); // TLS 1.3 + + // key_share extension (type=51) + let key_share_length = 2 + 2 + key_share_data.len(); // group + length + data + extensions.extend_from_slice(&[0x00, 0x33]); // type = 51 + extensions.extend_from_slice(&(key_share_length as u16).to_be_bytes()); + extensions.extend_from_slice(&[0x00, 0x1d]); // group = X25519 (0x001d) + extensions.extend_from_slice(&(key_share_data.len() as u16).to_be_bytes()); + extensions.extend_from_slice(key_share_data); + + // Extensions length + payload.extend_from_slice(&(extensions.len() as u16).to_be_bytes()); + payload.extend_from_slice(&extensions); + + // Handshake header + server_hello.push(HANDSHAKE_TYPE_SERVER_HELLO); + + // Payload length (3 bytes, big-endian) + let length_bytes = [ + ((payload.len() >> 16) & 0xff) as u8, + ((payload.len() >> 8) & 0xff) as u8, + (payload.len() & 0xff) as u8, + ]; + server_hello.extend_from_slice(&length_bytes); + server_hello.extend_from_slice(&payload); + + Ok(server_hello) +} + +/// Construct EncryptedExtensions message +pub fn construct_encrypted_extensions() -> Result> { + let mut encrypted_extensions = Vec::new(); + + // EncryptedExtensions structure: + // - handshake_type (1 byte) = 8 + // - length (3 bytes) + // - extensions_length (2 bytes) + // - extensions (variable, usually empty for minimal setup) + + encrypted_extensions.push(HANDSHAKE_TYPE_ENCRYPTED_EXTENSIONS); + + // Empty extensions for minimal setup + let extensions_length: u16 = 0; + let payload_length = 2; // Just the extensions_length field + + // Payload length (3 bytes) + encrypted_extensions.extend_from_slice(&[0x00, 0x00, payload_length as u8]); + + // Extensions length (2 bytes) + encrypted_extensions.extend_from_slice(&extensions_length.to_be_bytes()); + + Ok(encrypted_extensions) +} + +/// Construct Certificate message with HMAC-signed Ed25519 certificate +/// +/// # Arguments +/// * `cert_der` - DER-encoded certificate (with HMAC signature) +pub fn construct_certificate(cert_der: &[u8]) -> Result> { + let mut certificate = Vec::new(); + + // Certificate structure: + // - handshake_type (1 byte) = 11 + // - length (3 bytes) + // - certificate_request_context (1 byte length + data, usually empty) + // - certificate_list (3 bytes length + entries) + // - certificate_entry: + // - cert_data (3 bytes length + DER) + // - extensions (2 bytes length, usually empty) + + let mut payload = Vec::new(); + + // Certificate request context (empty for server certificates) + payload.push(0x00); + + // Certificate list + let mut cert_list = Vec::new(); + + // Certificate entry + // Cert data length (3 bytes) + cert_list.extend_from_slice(&[ + ((cert_der.len() >> 16) & 0xff) as u8, + ((cert_der.len() >> 8) & 0xff) as u8, + (cert_der.len() & 0xff) as u8, + ]); + cert_list.extend_from_slice(cert_der); + + // Extensions (empty) + cert_list.extend_from_slice(&[0x00, 0x00]); + + // Certificate list length (3 bytes) + payload.extend_from_slice(&[ + ((cert_list.len() >> 16) & 0xff) as u8, + ((cert_list.len() >> 8) & 0xff) as u8, + (cert_list.len() & 0xff) as u8, + ]); + payload.extend_from_slice(&cert_list); + + // Handshake header + certificate.push(HANDSHAKE_TYPE_CERTIFICATE); + + // Payload length (3 bytes) + certificate.extend_from_slice(&[ + ((payload.len() >> 16) & 0xff) as u8, + ((payload.len() >> 8) & 0xff) as u8, + (payload.len() & 0xff) as u8, + ]); + certificate.extend_from_slice(&payload); + + Ok(certificate) +} + +/// Construct CertificateVerify message +/// +/// # Arguments +/// * `signing_key` - Ed25519 signing key +/// * `handshake_hash` - SHA256 hash of all handshake messages up to this point +pub fn construct_certificate_verify( + signing_key: &Ed25519KeyPair, + handshake_hash: &[u8], +) -> Result> { + let mut certificate_verify = Vec::new(); + + // CertificateVerify structure: + // - handshake_type (1 byte) = 15 + // - length (3 bytes) + // - signature_algorithm (2 bytes) = 0x0807 for Ed25519 + // - signature (2 bytes length + data) + + // Construct the signed content + // TLS 1.3 uses a specific prefix for CertificateVerify: + // " " * 64 + "TLS 1.3, server CertificateVerify" + 0x00 + handshake_hash + let mut signed_content = Vec::new(); + signed_content.extend_from_slice(&[0x20u8; 64]); // 64 spaces + signed_content.extend_from_slice(b"TLS 1.3, server CertificateVerify"); + signed_content.push(0x00); + signed_content.extend_from_slice(handshake_hash); + + // Sign the content + let signature = signing_key.sign(&signed_content); + let signature_bytes = signature.as_ref(); + + let mut payload = Vec::new(); + + // Signature algorithm: Ed25519 (0x0807) + payload.extend_from_slice(&[0x08, 0x07]); + + // Signature length and data + payload.extend_from_slice(&(signature_bytes.len() as u16).to_be_bytes()); + payload.extend_from_slice(signature_bytes); + + // Handshake header + certificate_verify.push(HANDSHAKE_TYPE_CERTIFICATE_VERIFY); + + // Payload length (3 bytes) + certificate_verify.extend_from_slice(&[ + ((payload.len() >> 16) & 0xff) as u8, + ((payload.len() >> 8) & 0xff) as u8, + (payload.len() & 0xff) as u8, + ]); + certificate_verify.extend_from_slice(&payload); + + Ok(certificate_verify) +} + +/// Construct Finished message +/// +/// # Arguments +/// * `verify_data` - HMAC of handshake transcript (32 bytes for SHA256) +pub fn construct_finished(verify_data: &[u8]) -> Result> { + let mut finished = Vec::new(); + + // Finished structure: + // - handshake_type (1 byte) = 20 + // - length (3 bytes) + // - verify_data (variable, 32 bytes for SHA256) + + finished.push(HANDSHAKE_TYPE_FINISHED); + + // Payload length (3 bytes) + finished.extend_from_slice(&[ + ((verify_data.len() >> 16) & 0xff) as u8, + ((verify_data.len() >> 8) & 0xff) as u8, + (verify_data.len() & 0xff) as u8, + ]); + + finished.extend_from_slice(verify_data); + + Ok(finished) +} + +/// Write TLS record header +/// +/// Construct TLS 1.3 ClientHello message +/// +/// Returns handshake message bytes (without record header) +pub fn construct_client_hello( + client_random: &[u8; 32], + session_id: &[u8; 32], + client_public_key: &[u8], + server_name: &str, +) -> Result> { + let mut hello = Vec::with_capacity(512); + + // Handshake message type: ClientHello (0x01) + hello.push(0x01); + + // Placeholder for handshake message length (3 bytes) + let length_offset = hello.len(); + hello.extend_from_slice(&[0u8; 3]); + + // TLS version: 3.3 (TLS 1.2 for compatibility) + hello.extend_from_slice(&[VERSION_TLS_1_2_MAJOR, VERSION_TLS_1_2_MINOR]); + + // Client random (32 bytes) + hello.extend_from_slice(client_random); + + // Session ID length (1 byte) + Session ID (32 bytes) + hello.push(32); + hello.extend_from_slice(session_id); + + // Cipher suites + // Support only TLS_AES_128_GCM_SHA256 (0x1301) + hello.extend_from_slice(&[0x00, 0x02]); // Cipher suites length: 2 bytes + hello.extend_from_slice(&[0x13, 0x01]); // TLS_AES_128_GCM_SHA256 + + // Compression methods (1 method: null) + hello.extend_from_slice(&[0x01, 0x00]); + + // Extensions + let extensions_offset = hello.len(); + hello.extend_from_slice(&[0u8; 2]); // Placeholder for extensions length + + let mut extensions = Vec::new(); + + // 1. server_name extension (type 0) + { + let server_name_bytes = server_name.as_bytes(); + let server_name_len = server_name_bytes.len(); + + extensions.extend_from_slice(&[0x00, 0x00]); // Extension type: server_name + let ext_len = 5 + server_name_len; + extensions.extend_from_slice(&(ext_len as u16).to_be_bytes()); // Extension length + extensions.extend_from_slice(&((server_name_len + 3) as u16).to_be_bytes()); // Server name list length + extensions.push(0x00); // Name type: host_name + extensions.extend_from_slice(&(server_name_len as u16).to_be_bytes()); // Name length + extensions.extend_from_slice(server_name_bytes); // Server name + } + + // 2. supported_versions extension (type 43) + { + extensions.extend_from_slice(&[0x00, 0x2b]); // Extension type: supported_versions + extensions.extend_from_slice(&[0x00, 0x03]); // Extension length: 3 + extensions.push(0x02); // Supported versions length: 2 + extensions.extend_from_slice(&[0x03, 0x04]); // TLS 1.3 + } + + // 3. supported_groups extension (type 10) + { + extensions.extend_from_slice(&[0x00, 0x0a]); // Extension type: supported_groups + extensions.extend_from_slice(&[0x00, 0x04]); // Extension length: 4 + extensions.extend_from_slice(&[0x00, 0x02]); // Supported groups length: 2 + extensions.extend_from_slice(&[0x00, 0x1d]); // x25519 + } + + // 4. key_share extension (type 51) + { + extensions.extend_from_slice(&[0x00, 0x33]); // Extension type: key_share + let key_share_len = 2 + 4 + client_public_key.len(); + extensions.extend_from_slice(&(key_share_len as u16).to_be_bytes()); // Extension length + let key_share_list_len = 4 + client_public_key.len(); + extensions.extend_from_slice(&(key_share_list_len as u16).to_be_bytes()); // Key share list length + extensions.extend_from_slice(&[0x00, 0x1d]); // Group: x25519 + extensions.extend_from_slice(&(client_public_key.len() as u16).to_be_bytes()); // Key length + extensions.extend_from_slice(client_public_key); // Public key + } + + // 5. signature_algorithms extension (type 13) + { + extensions.extend_from_slice(&[0x00, 0x0d]); // Extension type: signature_algorithms + extensions.extend_from_slice(&[0x00, 0x04]); // Extension length: 4 + extensions.extend_from_slice(&[0x00, 0x02]); // Signature algorithms length: 2 + extensions.extend_from_slice(&[0x08, 0x07]); // ed25519 + } + + // Write extensions length + let extensions_length = extensions.len(); + hello[extensions_offset..extensions_offset + 2] + .copy_from_slice(&(extensions_length as u16).to_be_bytes()); + + // Append extensions + hello.extend_from_slice(&extensions); + + // Write handshake message length + let message_length = hello.len() - 4; // Exclude type (1) and length (3) + hello[length_offset..length_offset + 3] + .copy_from_slice(&(message_length as u32).to_be_bytes()[1..]); + + Ok(hello) +} + +/// # Arguments +/// * `record_type` - TLS record type (0x16 for Handshake, 0x17 for ApplicationData) +/// * `length` - Length of record payload +pub fn write_record_header(record_type: u8, length: u16) -> Vec { + let mut header = Vec::new(); + header.push(record_type); + header.extend_from_slice(&[VERSION_TLS_1_2_MAJOR, VERSION_TLS_1_2_MINOR]); // Version: TLS 1.2 + header.extend_from_slice(&length.to_be_bytes()); + header +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::reality::common::CONTENT_TYPE_HANDSHAKE; + + #[test] + fn test_construct_server_hello() { + let server_random = [0x42u8; 32]; + let session_id = vec![0x99u8; 32]; + let cipher_suite = 0x1301; // TLS_AES_128_GCM_SHA256 + let key_share = vec![0xAAu8; 32]; + + let result = construct_server_hello(&server_random, &session_id, cipher_suite, &key_share); + + assert!(result.is_ok()); + let msg = result.unwrap(); + assert_eq!(msg[0], HANDSHAKE_TYPE_SERVER_HELLO); + } + + #[test] + fn test_construct_encrypted_extensions() { + let result = construct_encrypted_extensions(); + assert!(result.is_ok()); + let msg = result.unwrap(); + assert_eq!(msg[0], HANDSHAKE_TYPE_ENCRYPTED_EXTENSIONS); + } + + #[test] + fn test_construct_certificate() { + let cert_der = vec![0xBBu8; 100]; + let result = construct_certificate(&cert_der); + assert!(result.is_ok()); + let msg = result.unwrap(); + assert_eq!(msg[0], HANDSHAKE_TYPE_CERTIFICATE); + } + + #[test] + fn test_construct_finished() { + let verify_data = vec![0xCCu8; 32]; + let result = construct_finished(&verify_data); + assert!(result.is_ok()); + let msg = result.unwrap(); + assert_eq!(msg[0], HANDSHAKE_TYPE_FINISHED); + assert_eq!(msg.len(), 1 + 3 + 32); // type + length + verify_data + } + + #[test] + fn test_write_record_header() { + let header = write_record_header(CONTENT_TYPE_HANDSHAKE, 100); + assert_eq!(header.len(), 5); + assert_eq!(header[0], 0x16); // Handshake + assert_eq!(header[1], 0x03); // TLS 1.2 + assert_eq!(header[2], 0x03); + assert_eq!(u16::from_be_bytes([header[3], header[4]]), 100); + } +} diff --git a/src/reality/reality_util.rs b/src/reality/reality_util.rs new file mode 100644 index 00000000..cfb3a31b --- /dev/null +++ b/src/reality/reality_util.rs @@ -0,0 +1,458 @@ +use aws_lc_rs::{ + agreement, + rand::{SecureRandom, SystemRandom}, +}; +use base64::engine::{general_purpose::URL_SAFE_NO_PAD, Engine as _}; + +use crate::buf_reader::BufReader; + +/// 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| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("Invalid base64: {}", e), + ) + })?; + + if decoded.len() != 32 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("Invalid public key length: {} (expected 32)", decoded.len()), + )); + } + + let mut key = [0u8; 32]; + key.copy_from_slice(&decoded); + Ok(key) +} + +/// Decodes a base64url-encoded private key +pub fn decode_private_key(encoded: &str) -> Result<[u8; 32], std::io::Error> { + let decoded = URL_SAFE_NO_PAD.decode(encoded).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("Invalid base64: {}", e), + ) + })?; + + if decoded.len() != 32 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!( + "Invalid private key length: {} (expected 32)", + decoded.len() + ), + )); + } + + let mut key = [0u8; 32]; + key.copy_from_slice(&decoded); + Ok(key) +} + +/// Decodes a hex-encoded short ID with zero-padding +/// +/// Short IDs can be 0-16 hex characters (0-8 bytes). +/// If shorter than 16 characters, they are left-padded with zeros. +pub fn decode_short_id(hex: &str) -> Result<[u8; 8], std::io::Error> { + if hex.len() > 16 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("Short ID too long: {} (max 16 hex chars)", hex.len()), + )); + } + + // Left-pad with zeros to make 16 chars + let padded = format!("{:0>16}", hex); + + let mut short_id = [0u8; 8]; + decode_hex_to_slice(&padded, &mut short_id).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("Invalid hex: {}", e), + ) + })?; + + Ok(short_id) +} + +/// Decode hex string to byte slice +fn decode_hex_to_slice(hex: &str, output: &mut [u8]) -> Result<(), &'static str> { + if hex.len() != output.len() * 2 { + return Err("Invalid hex length"); + } + + for (i, chunk) in hex.as_bytes().chunks(2).enumerate() { + let high = hex_char_to_value(chunk[0])?; + let low = hex_char_to_value(chunk[1])?; + output[i] = (high << 4) | low; + } + + Ok(()) +} + +/// Convert hex character to its numeric value +fn hex_char_to_value(c: u8) -> Result { + match c { + b'0'..=b'9' => Ok(c - b'0'), + b'a'..=b'f' => Ok(c - b'a' + 10), + b'A'..=b'F' => Ok(c - b'A' + 10), + _ => Err("Invalid hex character"), + } +} + +/// Extracts ClientRandom from ClientHello +/// +/// ClientHello structure (simplified): +/// - TLS Header: 5 bytes +/// - Handshake Type: 1 byte +/// - Length: 3 bytes +/// - Protocol Version: 2 bytes +/// - ClientRandom: 32 bytes (starts at offset 11) +pub fn extract_client_random(client_hello: &[u8]) -> Result<[u8; 32], std::io::Error> { + const TLS_HEADER_LEN: usize = 5; + const HANDSHAKE_HEADER_LEN: usize = 4; // type(1) + length(3) + const PROTOCOL_VERSION_LEN: usize = 2; + const RANDOM_OFFSET: usize = TLS_HEADER_LEN + HANDSHAKE_HEADER_LEN + PROTOCOL_VERSION_LEN; + + if client_hello.len() < RANDOM_OFFSET + 32 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "ClientHello too short to extract random", + )); + } + + let mut random = [0u8; 32]; + random.copy_from_slice(&client_hello[RANDOM_OFFSET..RANDOM_OFFSET + 32]); + Ok(random) +} + +/// Extracts SessionId slice from ClientHello without allocation +/// +/// SessionId comes after ClientRandom and has a 1-byte length prefix. +/// Returns a slice pointing into the original buffer. +pub fn extract_session_id_slice(client_hello: &[u8]) -> Result<&[u8], std::io::Error> { + const TLS_HEADER_LEN: usize = 5; + const HANDSHAKE_HEADER_LEN: usize = 4; + const PROTOCOL_VERSION_LEN: usize = 2; + const RANDOM_LEN: usize = 32; + const SESSION_ID_LEN_OFFSET: usize = + TLS_HEADER_LEN + HANDSHAKE_HEADER_LEN + PROTOCOL_VERSION_LEN + RANDOM_LEN; + + if client_hello.len() < SESSION_ID_LEN_OFFSET + 1 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "ClientHello too short to extract session ID", + )); + } + + let session_id_len = client_hello[SESSION_ID_LEN_OFFSET] as usize; + + if client_hello.len() < SESSION_ID_LEN_OFFSET + 1 + session_id_len { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "ClientHello session ID extends past end", + )); + } + + Ok(&client_hello[SESSION_ID_LEN_OFFSET + 1..SESSION_ID_LEN_OFFSET + 1 + session_id_len]) +} + +/// Extracts X25519 public key from ClientHello KeyShare extension +/// +/// This parses the TLS 1.3 extensions to find the KeyShare extension +/// and extracts the X25519 public key (group 0x001d). +pub fn extract_client_public_key(client_hello: &[u8]) -> Result<[u8; 32], std::io::Error> { + const TLS_HEADER_LEN: usize = 5; + + if client_hello.len() < TLS_HEADER_LEN { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "ClientHello too short", + )); + } + + let mut reader = BufReader::new(&client_hello[TLS_HEADER_LEN..]); + + // Parse handshake header + let _handshake_type = reader.read_u8()?; + let _handshake_len = reader.read_u24_be()?; + let _protocol_version = reader.read_u16_be()?; + + // Skip random + reader.skip(32)?; + + // Read session ID + let session_id_len = reader.read_u8()? as usize; + reader.skip(session_id_len)?; + + // Read cipher suites + let cipher_suites_len = reader.read_u16_be()? as usize; + reader.skip(cipher_suites_len)?; + + // Read compression methods + let compression_len = reader.read_u8()? as usize; + reader.skip(compression_len)?; + + // Parse extensions + let extensions_len = reader.read_u16_be()? as usize; + let extensions_start = reader.position(); + let extensions_end = extensions_start + extensions_len; + + while reader.position() < extensions_end { + let ext_type = reader.read_u16_be()?; + let ext_len = reader.read_u16_be()? as usize; + + if ext_type == 51 { + // KeyShare extension (0x0033) + return parse_keyshare_extension(reader.read_slice(ext_len)?); + } else { + reader.skip(ext_len)?; + } + } + + Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "KeyShare extension not found in ClientHello", + )) +} + +/// Parses the KeyShare extension to extract X25519 public key +fn parse_keyshare_extension(data: &[u8]) -> Result<[u8; 32], std::io::Error> { + let mut reader = BufReader::new(data); + + // KeyShare extension format (client): + // - Client Key Share Length: 2 bytes + // - Key Share Entries (repeated): + // - Group: 2 bytes + // - Key Exchange Length: 2 bytes + // - Key Exchange Data: variable + + let _client_shares_len = reader.read_u16_be()?; + + // Find X25519 key share (group 0x001d = 29) + loop { + // Check if we have at least 4 bytes for group and length + if reader.position() + 4 > data.len() { + break; + } + + let group = reader.read_u16_be()?; + let key_len = reader.read_u16_be()? as usize; + + if reader.position() + key_len > data.len() { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "KeyShare entry extends past end", + )); + } + + if group == 0x001d { + // X25519 + if key_len != 32 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("Invalid X25519 key length: {}", key_len), + )); + } + + let key_bytes = reader.read_slice(32)?; + let mut key = [0u8; 32]; + key.copy_from_slice(key_bytes); + return Ok(key); + } else { + reader.skip(key_len)?; + } + } + + Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "X25519 key share not found in KeyShare extension", + )) +} + +/// Extract server's X25519 public key from ServerHello message +/// +/// This parses the TLS 1.3 ServerHello to find the KeyShare extension +/// and extracts the X25519 public key (group 0x001d). +pub fn extract_server_public_key(server_hello: &[u8]) -> Result<[u8; 32], std::io::Error> { + const TLS_HEADER_LEN: usize = 5; + + if server_hello.len() < TLS_HEADER_LEN { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "ServerHello too short", + )); + } + + let mut reader = BufReader::new(&server_hello[TLS_HEADER_LEN..]); + + // Parse handshake header + let _handshake_type = reader.read_u8()?; + let _handshake_len = reader.read_u24_be()?; + let _protocol_version = reader.read_u16_be()?; + + // Skip random + reader.skip(32)?; + + // Read session ID (ServerHello can echo it back) + let session_id_len = reader.read_u8()? as usize; + reader.skip(session_id_len)?; + + // Read cipher suite (single 2-byte value in ServerHello) + reader.skip(2)?; + + // Read compression method (single byte in ServerHello) + reader.skip(1)?; + + // Parse extensions + let extensions_len = reader.read_u16_be()? as usize; + let extensions_start = reader.position(); + let extensions_end = extensions_start + extensions_len; + + while reader.position() < extensions_end { + let ext_type = reader.read_u16_be()?; + let ext_len = reader.read_u16_be()? as usize; + + if ext_type == 51 { + // KeyShare extension (0x0033) + return parse_server_keyshare_extension(reader.read_slice(ext_len)?); + } else { + reader.skip(ext_len)?; + } + } + + Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "KeyShare extension not found in ServerHello", + )) +} + +/// Parses the ServerHello KeyShare extension to extract X25519 public key +fn parse_server_keyshare_extension(data: &[u8]) -> Result<[u8; 32], std::io::Error> { + let mut reader = BufReader::new(data); + + // ServerHello KeyShare extension format: + // - Group: 2 bytes + // - Key Exchange Length: 2 bytes + // - Key Exchange Data: variable + + let group = reader.read_u16_be()?; + let key_len = reader.read_u16_be()? as usize; + + if group == 29 { + // X25519 + if key_len != 32 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("Invalid X25519 key length: {}", key_len), + )); + } + + let key_bytes = reader.read_slice(32)?; + let mut key = [0u8; 32]; + key.copy_from_slice(key_bytes); + return Ok(key); + } + + Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "X25519 key share not found in ServerHello KeyShare extension", + )) +} + +pub fn generate_keypair() -> std::io::Result<(String, String)> { + // Step 1: Generate 32 random bytes for private key + let rng = SystemRandom::new(); + let mut private_key_bytes = [0u8; 32]; + rng.fill(&mut private_key_bytes) + .map_err(|_| std::io::Error::other("RNG failed"))?; + + // Step 2: Create X25519 private key from the random bytes + let private_key = + agreement::PrivateKey::from_private_key(&agreement::X25519, &private_key_bytes) + .map_err(|_| std::io::Error::other("Failed to create X25519 key"))?; + + // Step 3: Derive public key from private key + let public_key_bytes = private_key + .compute_public_key() + .map_err(|_| std::io::Error::other("Failed to compute public key"))?; + + // Step 4: Encode both keys as base64url (no padding) + let private_key_b64 = URL_SAFE_NO_PAD.encode(private_key_bytes); + let public_key_b64 = URL_SAFE_NO_PAD.encode(public_key_bytes.as_ref()); + + Ok((private_key_b64, public_key_b64)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_decode_short_id() { + // Full 16-char hex + let short_id = decode_short_id("0123456789abcdef").unwrap(); + assert_eq!(short_id, [0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef]); + + // Partial hex (should be zero-padded on left) + let short_id2 = decode_short_id("abcdef").unwrap(); + assert_eq!(short_id2, [0x00, 0x00, 0x00, 0x00, 0x00, 0xab, 0xcd, 0xef]); + + // Empty (all zeros) + let short_id3 = decode_short_id("").unwrap(); + assert_eq!(short_id3, [0; 8]); + + // Too long should error + let result = decode_short_id("0123456789abcdef0"); + assert!(result.is_err()); + } + + #[test] + fn test_extract_client_random() { + // Create a minimal ClientHello with random + let mut client_hello = vec![0u8; 100]; + + // TLS Header + client_hello[0] = 0x16; // Handshake + client_hello[1] = 0x03; // Version major + client_hello[2] = 0x03; // Version minor (TLS 1.2) + + // Handshake header + client_hello[5] = 0x01; // ClientHello type + + // Protocol version in handshake + client_hello[9] = 0x03; // Major + client_hello[10] = 0x03; // Minor + + // ClientRandom starts at offset 11 + for i in 0..32 { + client_hello[11 + i] = (i + 1) as u8; + } + + let random = extract_client_random(&client_hello).unwrap(); + for i in 0..32 { + assert_eq!(random[i], (i + 1) as u8); + } + } + + #[test] + fn test_decode_public_key() { + use base64::engine::{general_purpose::URL_SAFE_NO_PAD, Engine as _}; + + // Valid 32-byte key + let key_bytes = [0x42u8; 32]; + let encoded = URL_SAFE_NO_PAD.encode(&key_bytes); + let decoded = decode_public_key(&encoded).unwrap(); + assert_eq!(decoded, key_bytes); + + // Invalid length + let short_key = [0x42u8; 16]; + let encoded_short = URL_SAFE_NO_PAD.encode(&short_key); + assert!(decode_public_key(&encoded_short).is_err()); + + // Invalid base64 + assert!(decode_public_key("not-valid-base64!!!").is_err()); + } +} diff --git a/src/reality_client_handler.rs b/src/reality_client_handler.rs new file mode 100644 index 00000000..5fc5942b --- /dev/null +++ b/src/reality_client_handler.rs @@ -0,0 +1,209 @@ +// REALITY client handler +// +// This file lives at src/ (not in reality/) for two reasons: +// 1. Symmetry with tls_server_handler.rs which handles both TLS and REALITY server-side +// 2. Avoids a dependency cycle: crypto → reality (core types), but this handler → crypto + +use async_trait::async_trait; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; + +use crate::address::NetLocation; +use crate::async_stream::AsyncStream; +use crate::crypto::{CryptoConnection, CryptoTlsStream}; +use crate::reality::{RealityClientConfig, RealityClientConnection}; +use crate::tcp_handler::{TcpClientHandler, TcpClientSetupResult}; + +/// REALITY client handler using buffered Connection API +/// +/// This handler establishes REALITY-obfuscated TLS connections using +/// the buffered sans-I/O pattern with RealityClientConnection. +#[derive(Debug)] +pub struct RealityClientHandler { + public_key: [u8; 32], + short_id: [u8; 8], + server_name: rustls::pki_types::ServerName<'static>, + handler: RealityInnerClientHandler, +} + +#[derive(Debug)] +pub enum RealityInnerClientHandler { + Default(Box), + VisionVless { uuid: Box<[u8]> }, +} + +impl RealityClientHandler { + pub fn new( + public_key: [u8; 32], + short_id: [u8; 8], + server_name: rustls::pki_types::ServerName<'static>, + handler: Box, + ) -> Self { + Self { + public_key, + short_id, + server_name, + handler: RealityInnerClientHandler::Default(handler), + } + } + pub fn new_vision_vless( + public_key: [u8; 32], + short_id: [u8; 8], + server_name: rustls::pki_types::ServerName<'static>, + user_id: Box<[u8]>, + ) -> Self { + Self { + public_key, + short_id, + server_name, + handler: RealityInnerClientHandler::VisionVless { uuid: user_id }, + } + } +} + +#[async_trait] +impl TcpClientHandler for RealityClientHandler { + async fn setup_client_stream( + &self, + server_stream: &mut Box, + mut client_stream: Box, + remote_location: NetLocation, + ) -> std::io::Result { + // Extract server name as string + let server_name_str = match &self.server_name { + rustls::pki_types::ServerName::DnsName(name) => name.as_ref(), + rustls::pki_types::ServerName::IpAddress(ip) => { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("REALITY requires DNS name, got IP address: {:?}", ip), + )) + } + _ => { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "REALITY requires DNS name", + )) + } + }; + + // Step 1: Create buffered REALITY client connection + log::debug!("REALITY CLIENT: Creating buffered RealityClientConnection"); + let reality_config = RealityClientConfig { + public_key: self.public_key, + short_id: self.short_id, + server_name: server_name_str.to_string(), + }; + + let mut reality_conn = RealityClientConnection::new(reality_config)?; + + // Step 2: Write ClientHello to server + log::debug!("REALITY CLIENT: Writing ClientHello"); + { + let mut write_buf = Vec::new(); + while reality_conn.wants_write() { + reality_conn.write_tls(&mut write_buf)?; + } + if !write_buf.is_empty() { + client_stream.write_all(&write_buf).await?; + client_stream.flush().await?; + } + } + + // Step 3: Read server's handshake messages + log::debug!("REALITY CLIENT: Reading server handshake"); + { + let mut buf = vec![0u8; 16384]; // Large enough for server handshake + let n = client_stream.read(&mut buf).await?; + if n == 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + "EOF while waiting for server handshake", + )); + } + let mut cursor = std::io::Cursor::new(&buf[..n]); + reality_conn.read_tls(&mut cursor)?; + } + + // Step 4: Process server handshake + log::debug!("REALITY CLIENT: Processing server handshake"); + reality_conn.process_new_packets()?; + + // Continue processing packets until handshake is complete + while reality_conn.is_handshaking() { + log::debug!("REALITY CLIENT: Handshake still in progress, checking for more work"); + + // Check if we need to write (e.g., client Finished) + if reality_conn.wants_write() { + log::debug!("REALITY CLIENT: Writing buffered handshake data"); + let mut write_buf = Vec::new(); + while reality_conn.wants_write() { + reality_conn.write_tls(&mut write_buf)?; + } + if !write_buf.is_empty() { + client_stream.write_all(&write_buf).await?; + client_stream.flush().await?; + log::debug!("REALITY CLIENT: Sent {} bytes to server", write_buf.len()); + } + } + + // Try to process more packets + let prev_state = reality_conn.is_handshaking(); + reality_conn.process_new_packets()?; + + // If we're still handshaking and state didn't change, read more data + if reality_conn.is_handshaking() && prev_state == reality_conn.is_handshaking() { + log::debug!("REALITY CLIENT: Reading more server handshake data"); + let mut buf = vec![0u8; 16384]; + let n = client_stream.read(&mut buf).await?; + if n == 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + "EOF while reading server handshake", + )); + } + log::debug!("REALITY CLIENT: Read {} bytes from server", n); + let mut cursor = std::io::Cursor::new(&buf[..n]); + reality_conn.read_tls(&mut cursor)?; + + // Process the new data + reality_conn.process_new_packets()?; + } + } + + // Step 5: Final check for any remaining writes + if reality_conn.wants_write() { + log::debug!("REALITY CLIENT: Writing final handshake data"); + let mut write_buf = Vec::new(); + while reality_conn.wants_write() { + reality_conn.write_tls(&mut write_buf)?; + } + if !write_buf.is_empty() { + client_stream.write_all(&write_buf).await?; + client_stream.flush().await?; + } + } + + // Step 6: Wrap in Connection enum + log::debug!("REALITY CLIENT: Creating Connection"); + let connection = CryptoConnection::new_reality_client(reality_conn); + + log::debug!("REALITY CLIENT: Handshake completed successfully"); + + let tls_stream = CryptoTlsStream::new(client_stream, connection); + + match self.handler { + RealityInnerClientHandler::Default(ref handler) => { + handler + .setup_client_stream(server_stream, Box::new(tls_stream), remote_location) + .await + } + RealityInnerClientHandler::VisionVless { ref uuid } => { + crate::vless::vless_client_handler::setup_custom_tls_vision_vless_client_stream( + tls_stream, + uuid, + &remote_location, + ) + .await + } + } + } +} diff --git a/src/rustls_util.rs b/src/rustls_config_util.rs similarity index 100% rename from src/rustls_util.rs rename to src/rustls_config_util.rs diff --git a/src/rustls_connection_util.rs b/src/rustls_connection_util.rs new file mode 100644 index 00000000..60ecc4f4 --- /dev/null +++ b/src/rustls_connection_util.rs @@ -0,0 +1,77 @@ +use std::io::Cursor; + +use rustls::Connection; + +/// Feed all data to a rustls connection, looping until all bytes are consumed. +/// +/// rustls's read_tls() may not consume all bytes in one call, so we must loop +/// until all data has been fed. This is critical to avoid losing bytes. +#[inline(always)] +pub fn feed_rustls_connection(connection: &mut Connection, data: &[u8]) -> std::io::Result<()> { + match connection { + Connection::Client(client_connection) => { + feed_rustls_client_connection(client_connection, data) + } + Connection::Server(server_connection) => { + feed_rustls_server_connection(server_connection, data) + } + } +} + +#[inline(always)] +pub fn feed_rustls_server_connection( + connection: &mut rustls::ServerConnection, + data: &[u8], +) -> std::io::Result<()> { + let mut cursor = Cursor::new(data); + let mut i = 0; + while i < data.len() { + let n = connection.read_tls(&mut cursor).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("failed to feed rustls server connection: {e}"), + ) + })?; + if n == 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!( + "rustls server connection did not consume all bytes: fed {}/{} bytes", + i, + data.len() + ), + )); + } + i += n; + } + Ok(()) +} + +#[inline(always)] +pub fn feed_rustls_client_connection( + connection: &mut rustls::ClientConnection, + data: &[u8], +) -> std::io::Result<()> { + let mut cursor = Cursor::new(data); + let mut i = 0; + while i < data.len() { + let n = connection.read_tls(&mut cursor).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("failed to feed rustls client connection: {e}"), + ) + })?; + if n == 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!( + "rustls client connection did not consume all bytes: fed {}/{} bytes", + i, + data.len() + ), + )); + } + i += n; + } + Ok(()) +} diff --git a/src/rustls_handshake.rs b/src/rustls_handshake.rs new file mode 100644 index 00000000..abebc4fb --- /dev/null +++ b/src/rustls_handshake.rs @@ -0,0 +1,263 @@ +//! Generic rustls handshake helper that works for both client and server connections. +//! +//! This module provides a unified handshake implementation that works correctly +//! with both TLS 1.2 and TLS 1.3, for both client and server connections. +//! +//! The key insight from tokio-rustls is to use rustls's `wants_read()` and +//! `wants_write()` methods to let the TLS state machine guide the handshake, +//! rather than assuming packet boundaries or TLS version-specific message patterns. + +use tokio::io::{AsyncReadExt, AsyncWriteExt}; + +use crate::async_stream::AsyncStream; +use crate::rustls_connection_util::feed_rustls_connection; + +/// Perform a complete TLS handshake using the provided rustls connection and stream. +/// +/// This function works for both client and server connections (via `rustls::Connection`), +/// and handles both TLS 1.2 and TLS 1.3 correctly by following rustls's state machine +/// guidance via `wants_read()` and `wants_write()`. +/// +/// # Arguments +/// +/// * `connection` - The rustls::Connection (either Client or Server variant) +/// * `stream` - The underlying async stream to read/write TLS data +/// * `buffer_size` - Size of the read buffer (typically 16384) +/// +/// # Returns +/// +/// Returns `Ok(())` if the handshake completes successfully. +/// +/// # Errors +/// +/// Returns an error if: +/// - EOF is encountered during handshake +/// - TLS processing fails (certificate validation, protocol error, etc.) +/// - I/O error occurs +/// - Handshake stalls (neither wants_read nor wants_write but still handshaking) +/// +/// # Example +/// +/// ```ignore +/// use rustls::Connection; +/// +/// let client_conn = rustls::ClientConnection::new(config, server_name)?; +/// let mut connection = Connection::Client(client_conn); +/// perform_handshake(&mut connection, &mut tcp_stream, 16384).await?; +/// // Handshake complete, connection ready for application data +/// ``` +pub async fn perform_handshake( + connection: &mut rustls::Connection, + stream: &mut Box, + buffer_size: usize, +) -> std::io::Result<()> { + log::debug!("TLS handshake starting"); + + let mut iteration = 0; + let mut eof = false; + + // Infinite loop - we explicitly break when done + // This matches tokio-rustls's structure: loop { ... return match ... } + loop { + iteration += 1; + let mut write_would_block = false; + let mut read_would_block = false; + + log::debug!( + "TLS handshake iteration {}: wants_read={}, wants_write={}", + iteration, + connection.wants_read(), + connection.wants_write() + ); + + // Phase 1: Drain all pending writes + while connection.wants_write() { + match write_pending_data(connection, stream).await { + Ok(()) => {} + Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => { + write_would_block = true; + break; + } + Err(e) => return Err(e), + } + } + + // Phase 2: Flush the stream to ensure data is sent + match stream.flush().await { + Ok(()) => {} + Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => { + write_would_block = true; + } + Err(e) => return Err(e), + } + + // Phase 3: Read ONCE if rustls wants more data + // CRITICAL: Unlike tokio-rustls which uses non-blocking I/O with Poll::Pending, + // we use blocking async I/O. Therefore we can NOT loop on wants_read() because + // each read() blocks. We only read ONCE per main loop iteration. + if !eof && connection.wants_read() { + match read_and_process_data(connection, stream, buffer_size).await { + Ok(0) => { + log::debug!("TLS handshake: EOF from peer"); + eof = true; + } + Ok(n) => { + log::debug!("TLS handshake: read and processed {} bytes", n); + } + Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => { + read_would_block = true; + } + Err(e) => return Err(e), + } + } + + // Phase 4: Check handshake state and decide what to do next + // This EXACTLY matches tokio-rustls's logic at lines 173-187 + log::debug!( + "TLS handshake iteration {} complete: eof={}, is_handshaking={}, write_would_block={}, read_would_block={}", + iteration, eof, connection.is_handshaking(), write_would_block, read_would_block + ); + + match (eof, connection.is_handshaking()) { + // EOF + still handshaking = error + (true, true) => { + return Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + "EOF during TLS handshake", + )); + } + // Handshake complete! (regardless of EOF state) + (_, false) => { + log::debug!("TLS handshake complete after {} iterations", iteration); + break; // Exit the loop + } + // Still handshaking + would block on I/O + (_, true) if write_would_block || read_would_block => { + // In blocking mode, this shouldn't happen, but if it does we're stuck + log::error!("TLS handshake would block at iteration {}", iteration); + return Err(std::io::Error::new( + std::io::ErrorKind::WouldBlock, + "TLS handshake would block (should not happen in blocking mode)", + )); + } + // Still handshaking, no I/O blocking - loop again + (..) => { + // Safety check to prevent infinite loops + if iteration > 100 { + log::error!("TLS handshake exceeded 100 iterations"); + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "TLS handshake exceeded maximum iterations", + )); + } + log::debug!("TLS handshake continuing to iteration {}", iteration + 1); + continue; + } + } + } + + // Final write if there's any pending data after handshake completes + while connection.wants_write() { + log::debug!("TLS handshake: final write"); + write_pending_data(connection, stream).await?; + } + + Ok(()) +} + +/// Write all pending TLS data from the connection to the stream. +/// +/// This drains rustls's write buffer and sends all data to the peer, +/// ensuring it's flushed to the network. +async fn write_pending_data( + connection: &mut rustls::Connection, + stream: &mut Box, +) -> std::io::Result<()> { + let mut write_buf = Vec::new(); + + // Drain all pending writes into our buffer + while connection.wants_write() { + connection.write_tls(&mut write_buf)?; + } + + // Send to peer if we have data + if !write_buf.is_empty() { + stream.write_all(&write_buf).await?; + stream.flush().await?; + } + + Ok(()) +} + +/// Read data from the stream and process it through the TLS connection. +/// +/// This reads available data from the network, feeds ALL of it to rustls via the +/// `feed_rustls_connection` helper (which loops until all bytes are consumed), +/// and processes it via `process_new_packets()` to advance the handshake state machine. +/// +/// Returns the number of bytes read (0 indicates EOF). +async fn read_and_process_data( + connection: &mut rustls::Connection, + stream: &mut Box, + buffer_size: usize, +) -> std::io::Result { + let mut buf = vec![0u8; buffer_size]; + let n = stream.read(&mut buf).await?; + + log::debug!("TLS: read {} bytes from peer", n); + + if n == 0 { + // EOF - return 0 to let caller handle it + return Ok(0); + } + + // Feed ALL the data to rustls using the helper that loops until all bytes are consumed + // This is CRITICAL - rustls may not consume all bytes in one read_tls() call + feed_rustls_connection(connection, &buf[..n])?; + + // Process the new data to advance the state machine + // This is called after every feed, just like tokio-rustls does + connection.process_new_packets().map_err(|e| { + log::error!("TLS error processing packets: {}", e); + std::io::Error::new(std::io::ErrorKind::InvalidData, format!("TLS error: {}", e)) + })?; + + log::debug!("TLS: processed packets successfully"); + + Ok(n) +} + +#[cfg(test)] +mod tests { + #[test] + fn test_rustls_connection_enum_exists() { + // Compile-time test that rustls::Connection exists + // This will fail to compile if rustls::Connection doesn't exist + fn _assert_connection_type(_: &rustls::Connection) {} + } + + #[test] + fn test_connection_has_required_methods() { + // This test ensures rustls::Connection has the methods we need + // It won't compile if the methods don't exist + use std::sync::Arc; + + let config = Arc::new( + rustls::ClientConfig::builder() + .with_root_certificates(rustls::RootCertStore::empty()) + .with_no_client_auth(), + ); + + let server_name = rustls::pki_types::ServerName::try_from("example.com") + .unwrap() + .to_owned(); + + let client_conn = rustls::ClientConnection::new(config, server_name).unwrap(); + let connection = rustls::Connection::Client(client_conn); + + // These method calls will fail to compile if the methods don't exist + let _handshaking = connection.is_handshaking(); + let _wants_read = connection.wants_read(); + let _wants_write = connection.wants_write(); + } +} diff --git a/src/shadow_tls/mod.rs b/src/shadow_tls/mod.rs index 3e15d7dd..52c667a7 100644 --- a/src/shadow_tls/mod.rs +++ b/src/shadow_tls/mod.rs @@ -4,8 +4,8 @@ mod shadow_tls_server_handler; mod shadow_tls_stream; pub use shadow_tls_server_handler::{ - feed_server_connection, read_client_hello, setup_shadowtls_server_stream, ParsedClientHello, - ShadowTlsServerTarget, ShadowTlsServerTargetHandshake, + read_client_hello, setup_shadowtls_server_stream, ParsedClientHello, ShadowTlsServerTarget, + ShadowTlsServerTargetHandshake, }; pub use shadow_tls_client_handler::ShadowTlsClientHandler; diff --git a/src/shadow_tls/shadow_tls_client_handler.rs b/src/shadow_tls/shadow_tls_client_handler.rs index a3916cbe..c22b6f03 100644 --- a/src/shadow_tls/shadow_tls_client_handler.rs +++ b/src/shadow_tls/shadow_tls_client_handler.rs @@ -8,6 +8,7 @@ use tokio::io::AsyncWriteExt; // For write_all // For random bytes use crate::address::NetLocation; use crate::async_stream::AsyncStream; use crate::buf_reader::BufReader; +use crate::rustls_connection_util::feed_rustls_client_connection; use crate::shadow_tls::shadow_tls_hmac::ShadowTlsHmac; use crate::shadow_tls::shadow_tls_stream::ShadowTlsStream; use crate::stream_reader::StreamReader; @@ -94,7 +95,7 @@ impl TcpClientHandler for ShadowTlsClientHandler { let parsed_server_hello = parse_server_hello(server_hello_frame)?; - feed_client_connection(&mut client_conn, server_hello_frame)?; + feed_rustls_client_connection(&mut client_conn, server_hello_frame)?; client_conn.process_new_packets().map_err(|e| { std::io::Error::new( std::io::ErrorKind::InvalidData, @@ -169,7 +170,7 @@ impl TcpClientHandler for ShadowTlsClientHandler { )); } - feed_client_connection(&mut client_conn, server_frame)?; + feed_rustls_client_connection(&mut client_conn, server_frame)?; client_conn.process_new_packets().map_err(|e| { std::io::Error::new( std::io::ErrorKind::InvalidData, @@ -345,22 +346,3 @@ async fn read_tls_frame<'a>( Ok(&read_buf[..TLS_HEADER_LEN + payload_len]) } - -#[inline] -fn feed_client_connection( - client_connection: &mut rustls::ClientConnection, - data: &[u8], -) -> std::io::Result<()> { - let mut cursor = Cursor::new(data); - let mut i = 0; - while i < data.len() { - let n = client_connection.read_tls(&mut cursor).map_err(|e| { - std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!("failed to feed client connection: {e}"), - ) - })?; - i += n; - } - Ok(()) -} diff --git a/src/shadow_tls/shadow_tls_server_handler.rs b/src/shadow_tls/shadow_tls_server_handler.rs index 76ac7e3f..a58e0ca5 100644 --- a/src/shadow_tls/shadow_tls_server_handler.rs +++ b/src/shadow_tls/shadow_tls_server_handler.rs @@ -20,6 +20,7 @@ use crate::client_proxy_selector::ClientProxySelector; use crate::noop_stream::NoopStream; use crate::option_util::NoneOrOne; use crate::resolver::Resolver; +use crate::rustls_connection_util::feed_rustls_server_connection; use crate::stream_reader::StreamReader; use crate::tcp_client_connector::TcpClientConnector; use crate::tcp_handler::{TcpServerHandler, TcpServerSetupResult}; @@ -841,7 +842,7 @@ async fn setup_local_handshake( ) })?; - feed_server_connection(&mut server_connection, &client_hello_frame)?; + feed_rustls_server_connection(&mut server_connection, &client_hello_frame)?; server_connection.process_new_packets().map_err(|e| { std::io::Error::new( @@ -1041,7 +1042,7 @@ async fn setup_local_handshake( hmac_client_data.update(initial_client_data); hmac_client_data.update(&hmac_client_data.digest()); - let shadow_tls_stream = ShadowTlsStream::new( + let mut shadow_tls_stream = ShadowTlsStream::new( server_stream, initial_client_data, hmac_client_data, @@ -1049,12 +1050,18 @@ async fn setup_local_handshake( None, )?; + // Feed 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_server_connection(&mut server_connection, &client_header_bytes)?; - feed_server_connection(&mut server_connection, client_payload_bytes)?; + feed_rustls_server_connection(&mut server_connection, &client_header_bytes)?; + feed_rustls_server_connection(&mut server_connection, client_payload_bytes)?; server_connection.process_new_packets().map_err(|e| { std::io::Error::new( @@ -1065,25 +1072,6 @@ async fn setup_local_handshake( } } -#[inline] -pub fn feed_server_connection( - server_connection: &mut rustls::ServerConnection, - data: &[u8], -) -> std::io::Result<()> { - let mut cursor = Cursor::new(data); - let mut i = 0; - while i < data.len() { - let n = server_connection.read_tls(&mut cursor).map_err(|e| { - std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!("failed to feed server connection: {e}"), - ) - })?; - i += n; - } - Ok(()) -} - #[inline] fn read_server_connection( server_connection: &mut rustls::ServerConnection, diff --git a/src/shadowsocks/mod.rs b/src/shadowsocks/mod.rs index 91579ceb..9406e275 100644 --- a/src/shadowsocks/mod.rs +++ b/src/shadowsocks/mod.rs @@ -1,11 +1,13 @@ mod aead_util; mod blake3_key; mod default_key; +mod salt_checker; mod shadowsocks_cipher; mod shadowsocks_key; mod shadowsocks_stream; mod shadowsocks_stream_type; mod shadowsocks_tcp_handler; +mod timed_salt_checker; pub use default_key::DefaultKey; pub use shadowsocks_cipher::ShadowsocksCipher; diff --git a/src/salt_checker.rs b/src/shadowsocks/salt_checker.rs similarity index 100% rename from src/salt_checker.rs rename to src/shadowsocks/salt_checker.rs diff --git a/src/shadowsocks/shadowsocks_cipher.rs b/src/shadowsocks/shadowsocks_cipher.rs index 890cc4c5..7b9b4888 100644 --- a/src/shadowsocks/shadowsocks_cipher.rs +++ b/src/shadowsocks/shadowsocks_cipher.rs @@ -39,19 +39,23 @@ impl ShadowsocksCipher { pub fn salt_len(&self) -> usize { self.salt_len } + + pub fn key_len(&self) -> usize { + self.algorithm.key_len() + } } -impl From<&str> for ShadowsocksCipher { - fn from(name: &str) -> Self { +impl TryFrom<&str> for ShadowsocksCipher { + type Error = std::io::Error; + + fn try_from(name: &str) -> Result { match name { "chacha20-ietf-poly1305" | "chacha20-poly1305" => { - ShadowsocksCipher::chacha20_ietf_poly1305() - } - "aes-256-gcm" => ShadowsocksCipher::aes_256_gcm(), - "aes-128-gcm" => ShadowsocksCipher::aes_128_gcm(), - _ => { - panic!("Unknown cipher: {name}"); + Ok(ShadowsocksCipher::chacha20_ietf_poly1305()) } + "aes-256-gcm" => Ok(ShadowsocksCipher::aes_256_gcm()), + "aes-128-gcm" => Ok(ShadowsocksCipher::aes_128_gcm()), + _ => Err(std::io::Error::other(format!("Unknown cipher: {name}"))), } } } diff --git a/src/shadowsocks/shadowsocks_stream.rs b/src/shadowsocks/shadowsocks_stream.rs index 345f4121..a45e0374 100644 --- a/src/shadowsocks/shadowsocks_stream.rs +++ b/src/shadowsocks/shadowsocks_stream.rs @@ -10,16 +10,17 @@ use aws_lc_rs::error::Unspecified; use futures::ready; use parking_lot::Mutex; use rand::RngCore; +use subtle::ConstantTimeEq; use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; use super::aead_util::TAG_LEN; +use super::salt_checker::SaltChecker; use super::shadowsocks_key::ShadowsocksKey; use super::shadowsocks_stream_type::ShadowsocksStreamType; use crate::async_stream::{ AsyncFlushMessage, AsyncMessageStream, AsyncPing, AsyncReadMessage, AsyncShutdownMessage, AsyncStream, AsyncWriteMessage, }; -use crate::salt_checker::SaltChecker; use crate::util::allocate_vec; fn generate_iv(buf: &mut [u8]) { @@ -534,13 +535,12 @@ impl ShadowsocksStream { let request_salt = &self.unprocessed_buf[self.salt_len + 9..self.salt_len + 9 + self.salt_len]; - for (a, b) in request_salt.iter().zip(self.encrypt_iv.iter()) { - if a != b { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - "server returned request salt does not match", - )); - } + // Use constant-time comparison to prevent timing attacks + if request_salt.ct_eq(&self.encrypt_iv[..]).unwrap_u8() == 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "server returned request salt does not match", + )); } let first_chunk_len = diff --git a/src/shadowsocks/shadowsocks_tcp_handler.rs b/src/shadowsocks/shadowsocks_tcp_handler.rs index 5a91a666..66c9b807 100644 --- a/src/shadowsocks/shadowsocks_tcp_handler.rs +++ b/src/shadowsocks/shadowsocks_tcp_handler.rs @@ -5,16 +5,17 @@ use parking_lot::Mutex; use rand::{Rng, RngCore}; use tokio::io::AsyncWriteExt; -use crate::address::NetLocation; +use super::salt_checker::SaltChecker; +use super::timed_salt_checker::TimedSaltChecker; +use crate::address::{Address, NetLocation}; use crate::async_stream::AsyncStream; use crate::option_util::NoneOrOne; -use crate::salt_checker::SaltChecker; use crate::socks_handler::{read_location, write_location_to_vec}; use crate::stream_reader::StreamReader; use crate::tcp_handler::{ TcpClientHandler, TcpClientSetupResult, TcpServerHandler, TcpServerSetupResult, }; -use crate::timed_salt_checker::TimedSaltChecker; +use crate::uot::{UotV1Stream, UotV2Stream, UOT_V1_MAGIC_ADDRESS, UOT_V2_MAGIC_ADDRESS}; use crate::util::write_all; use super::blake3_key::Blake3Key; @@ -34,7 +35,7 @@ pub struct ShadowsocksTcpHandler { impl ShadowsocksTcpHandler { pub fn new(cipher_name: &str, password: &str) -> Self { - let cipher: ShadowsocksCipher = cipher_name.into(); + let cipher: ShadowsocksCipher = cipher_name.try_into().unwrap(); let key: Arc> = Arc::new(Box::new(DefaultKey::new( password, cipher.algorithm().key_len(), @@ -48,7 +49,7 @@ impl ShadowsocksTcpHandler { } pub fn new_aead2022(cipher_name: &str, key_bytes: &[u8]) -> Self { - let cipher: ShadowsocksCipher = cipher_name.into(); + let cipher: ShadowsocksCipher = cipher_name.try_into().unwrap(); let key: Arc> = Arc::new(Box::new(Blake3Key::new( key_bytes.to_vec().into_boxed_slice(), cipher.algorithm().key_len(), @@ -105,6 +106,80 @@ impl TcpServerHandler for ShadowsocksTcpHandler { } } + // Check for UDP-over-TCP (UoT) magic addresses + if let Address::Hostname(ref host) = remote_location.address() { + if host == UOT_V1_MAGIC_ADDRESS { + // UoT V1: Multi-destination UDP + // Each packet has: ATYP + address + port + length + data + let mut uot_stream = UotV1Stream::new(Box::new(server_stream)); + + // Feed any unparsed data (first UoT packet might be in same TCP segment) + let unparsed_data = stream_reader.unparsed_data(); + if !unparsed_data.is_empty() { + log::debug!( + "Shadowsocks UoT V1: feeding {} bytes of initial data", + unparsed_data.len() + ); + uot_stream.feed_initial_data(unparsed_data); + } + + return Ok(TcpServerSetupResult::MultiDirectionalUdp { + stream: Box::new(uot_stream), + need_initial_flush: false, + override_proxy_provider: NoneOrOne::Unspecified, + num_sockets: 4, // TODO: make configurable + }); + } else if host == UOT_V2_MAGIC_ADDRESS { + // UoT V2: Read request header first + // Request: isConnect(u8) + ATYP + address + port + // Note: V2 uses SOCKS address format (0x01=IPv4, 0x03=Domain, 0x04=IPv6), + // NOT UoT address format! + let is_connect = stream_reader.read_u8(&mut server_stream).await?; + log::debug!("Shadowsocks UoT V2: is_connect = {}", is_connect); + + // Read destination address using SOCKS address format + let destination = read_location(&mut server_stream, &mut stream_reader).await?; + log::debug!("Shadowsocks UoT V2: destination = {:?}", destination); + + if is_connect == 1 { + // V2 Connect mode: Single destination, length-prefixed packets only + // Reuse UotV2Stream which has identical format: length(u16be) + data + let unparsed_data = stream_reader.unparsed_data(); + let mut uot_v2_stream = UotV2Stream::new(Box::new(server_stream)); + if !unparsed_data.is_empty() { + uot_v2_stream.feed_initial_read_data(unparsed_data)?; + } + + return Ok(TcpServerSetupResult::BidirectionalUdp { + remote_location: destination, + stream: Box::new(uot_v2_stream), + need_initial_flush: false, + override_proxy_provider: NoneOrOne::Unspecified, + }); + } else { + // V2 Non-connect mode: Same as V1 (multi-destination) + let mut uot_stream = UotV1Stream::new(Box::new(server_stream)); + + // Feed any unparsed data + let unparsed_data = stream_reader.unparsed_data(); + if !unparsed_data.is_empty() { + log::debug!( + "Shadowsocks UoT V2 non-connect: feeding {} bytes of initial data", + unparsed_data.len() + ); + uot_stream.feed_initial_data(unparsed_data); + } + + return Ok(TcpServerSetupResult::MultiDirectionalUdp { + stream: Box::new(uot_stream), + need_initial_flush: false, + override_proxy_provider: NoneOrOne::Unspecified, + num_sockets: 4, + }); + } + } + } + Ok(TcpServerSetupResult::TcpForward { remote_location, stream: Box::new(server_stream), diff --git a/src/timed_salt_checker.rs b/src/shadowsocks/timed_salt_checker.rs similarity index 97% rename from src/timed_salt_checker.rs rename to src/shadowsocks/timed_salt_checker.rs index 898d9e95..0497d8af 100644 --- a/src/timed_salt_checker.rs +++ b/src/shadowsocks/timed_salt_checker.rs @@ -1,7 +1,7 @@ use std::collections::{HashSet, VecDeque}; use std::time::Instant; -use crate::salt_checker::SaltChecker; +use super::salt_checker::SaltChecker; #[derive(Debug)] struct TimeEntry { diff --git a/src/slide_buffer.rs b/src/slide_buffer.rs new file mode 100644 index 00000000..b33dde33 --- /dev/null +++ b/src/slide_buffer.rs @@ -0,0 +1,417 @@ +//! Zero-allocation sliding buffer for streaming data +//! +//! This module provides a fixed-capacity sliding buffer optimized for +//! streaming protocols where data is written at one end and read from another. +//! Unlike a true ring buffer, this uses a linear layout with lazy compaction +//! via `copy_within()`, which is optimal for use cases requiring contiguous +//! slices (like TLS record processing). +//! +//! # Design Rationale +//! +//! A true ring buffer wraps around, but that creates non-contiguous data +//! regions which require either two-slice APIs or copies for consumers +//! expecting `&[u8]`. Since TLS encryption/decryption functions need +//! contiguous slices, this linear design with lazy compaction is more +//! efficient for our use case. + +use std::io::{BufRead, Read}; + +/// A fixed-capacity sliding buffer with zero-allocation read/write operations. +/// +/// This buffer maintains start and end offsets into a pre-allocated storage, +/// allowing efficient append and consume operations without allocating. +/// When the write end approaches capacity, `compact()` or `maybe_compact()` +/// can be used to move remaining data to the front using `copy_within()`. +/// +/// # Example +/// ```ignore +/// let mut buf = SlideBuffer::new(1024); +/// buf.extend_from_slice(b"hello"); +/// assert_eq!(buf.as_slice(), b"hello"); +/// buf.consume(2); +/// assert_eq!(buf.as_slice(), b"llo"); +/// ``` +pub struct SlideBuffer { + /// Pre-allocated buffer storage + data: Box<[u8]>, + /// Start offset of valid data (inclusive) + start: usize, + /// End offset of valid data (exclusive) + end: usize, +} + +impl SlideBuffer { + /// Create a new slide buffer with the specified capacity. + /// + /// The buffer is allocated once and reused for all operations. + #[inline] + pub fn new(capacity: usize) -> Self { + Self { + data: vec![0u8; capacity].into_boxed_slice(), + start: 0, + end: 0, + } + } + + /// Returns the number of bytes currently stored in the buffer. + #[inline] + pub fn len(&self) -> usize { + self.end - self.start + } + + /// Returns the remaining space available for writing. + /// + /// This is the space at the end of the buffer. To reclaim space + /// consumed from the front, call `compact()`. + #[inline] + pub fn remaining_capacity(&self) -> usize { + self.data.len() - self.end + } + + /// Get a slice of the readable data. + #[inline] + pub fn as_slice(&self) -> &[u8] { + &self.data[self.start..self.end] + } + + /// Get a mutable slice for writing new data at the end. + /// + /// Returns the writable portion at the end of the buffer. + /// After writing, call `advance_write(n)` to mark bytes as written. + /// + /// # Example + /// ```ignore + /// let mut buf = SlideBuffer::new(1024); + /// let write_buf = buf.write_slice(); + /// write_buf[..5].copy_from_slice(b"hello"); + /// buf.advance_write(5); + /// ``` + #[inline] + pub fn write_slice(&mut self) -> &mut [u8] { + &mut self.data[self.end..] + } + + /// Extend the buffer with data from a slice. + /// + /// # Panics + /// Panics in debug mode if there isn't enough capacity. + #[inline] + pub fn extend_from_slice(&mut self, data: &[u8]) { + debug_assert!( + self.remaining_capacity() >= data.len(), + "SlideBuffer overflow: need {} bytes, have {}", + data.len(), + self.remaining_capacity() + ); + let end = self.end; + self.data[end..end + data.len()].copy_from_slice(data); + self.end += data.len(); + } + + /// Mark n bytes as written (after writing to `write_slice()`). + #[inline] + pub fn advance_write(&mut self, n: usize) { + debug_assert!( + self.end + n <= self.data.len(), + "SlideBuffer advance_write overflow: end={}, n={}, capacity={}", + self.end, + n, + self.data.len() + ); + self.end += n; + } + + /// Consume n bytes from the front of the buffer. + /// + /// # Panics + /// Panics in debug mode if n exceeds the available data. + #[inline] + pub fn consume(&mut self, n: usize) { + debug_assert!( + n <= self.len(), + "SlideBuffer consume underflow: n={}, len={}", + n, + self.len() + ); + self.start += n; + + // Reset offsets if buffer is now empty + if self.start >= self.end { + self.start = 0; + self.end = 0; + } + } + + /// Compact the buffer by moving data to the front. + /// + /// This reclaims the space that was consumed from the front, + /// making it available for writing at the end. + /// Uses `copy_within()` which is optimized by the compiler. + #[inline] + pub fn compact(&mut self) { + if self.start > 0 && self.start < self.end { + self.data.copy_within(self.start..self.end, 0); + self.end -= self.start; + self.start = 0; + } else if self.start >= self.end { + self.start = 0; + self.end = 0; + } + } + + /// Compact only if we've consumed more than the threshold. + /// + /// This amortizes the cost of compaction over many operations, + /// avoiding unnecessary copies when little space would be reclaimed. + #[inline] + pub fn maybe_compact(&mut self, threshold: usize) { + if self.start > threshold { + self.compact(); + } + } + + /// Returns a two-byte value at the given offset as big-endian u16. + #[inline] + pub fn get_u16_be(&self, offset: usize) -> Option { + if offset + 2 <= self.len() { + let idx = self.start + offset; + Some(u16::from_be_bytes([self.data[idx], self.data[idx + 1]])) + } else { + None + } + } +} + +impl Read for SlideBuffer { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + let available = self.len(); + if available == 0 { + return Ok(0); + } + let to_read = buf.len().min(available); + buf[..to_read].copy_from_slice(&self.data[self.start..self.start + to_read]); + self.consume(to_read); + Ok(to_read) + } +} + +impl BufRead for SlideBuffer { + fn fill_buf(&mut self) -> std::io::Result<&[u8]> { + Ok(self.as_slice()) + } + + fn consume(&mut self, amt: usize) { + SlideBuffer::consume(self, amt); + } +} + +// Implement indexing for convenient access +impl std::ops::Index for SlideBuffer { + type Output = u8; + + #[inline] + fn index(&self, index: usize) -> &Self::Output { + &self.data[self.start + index] + } +} + +impl std::ops::Index> for SlideBuffer { + type Output = [u8]; + + #[inline] + fn index(&self, range: std::ops::Range) -> &Self::Output { + &self.data[self.start + range.start..self.start + range.end] + } +} + +impl std::ops::Index> for SlideBuffer { + type Output = [u8]; + + #[inline] + fn index(&self, range: std::ops::RangeFrom) -> &Self::Output { + &self.data[self.start + range.start..self.end] + } +} + +impl std::ops::Index> for SlideBuffer { + type Output = [u8]; + + #[inline] + fn index(&self, range: std::ops::RangeTo) -> &Self::Output { + &self.data[self.start..self.start + range.end] + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_new_buffer() { + let buf = SlideBuffer::new(1024); + assert_eq!(buf.len(), 0); + assert_eq!(buf.remaining_capacity(), 1024); + } + + #[test] + fn test_extend_from_slice() { + let mut buf = SlideBuffer::new(1024); + buf.extend_from_slice(b"hello"); + assert_eq!(buf.len(), 5); + assert_eq!(buf.as_slice(), b"hello"); + assert_eq!(buf.remaining_capacity(), 1024 - 5); + } + + #[test] + fn test_consume() { + let mut buf = SlideBuffer::new(1024); + buf.extend_from_slice(b"hello world"); + buf.consume(6); + assert_eq!(buf.as_slice(), b"world"); + assert_eq!(buf.len(), 5); + } + + #[test] + fn test_consume_all_resets() { + let mut buf = SlideBuffer::new(1024); + buf.extend_from_slice(b"hello"); + buf.consume(5); + assert_eq!(buf.len(), 0); + assert_eq!(buf.remaining_capacity(), 1024); + } + + #[test] + fn test_compact() { + let mut buf = SlideBuffer::new(1024); + buf.extend_from_slice(b"hello world"); + buf.consume(6); + assert_eq!(buf.remaining_capacity(), 1024 - 11); + + buf.compact(); + assert_eq!(buf.as_slice(), b"world"); + assert_eq!(buf.remaining_capacity(), 1024 - 5); + } + + #[test] + fn test_maybe_compact() { + let mut buf = SlideBuffer::new(1024); + buf.extend_from_slice(b"0123456789"); // 10 bytes + buf.consume(5); + + // Threshold not met + buf.maybe_compact(10); + assert_eq!(buf.remaining_capacity(), 1024 - 10); + + // Threshold met + buf.maybe_compact(4); + assert_eq!(buf.remaining_capacity(), 1024 - 5); + } + + #[test] + fn test_write_slice() { + let mut buf = SlideBuffer::new(1024); + let write_buf = buf.write_slice(); + write_buf[..5].copy_from_slice(b"hello"); + buf.advance_write(5); + + assert_eq!(buf.as_slice(), b"hello"); + assert_eq!(buf.len(), 5); + } + + #[test] + fn test_read_trait() { + let mut buf = SlideBuffer::new(1024); + buf.extend_from_slice(b"hello world"); + + let mut output = [0u8; 5]; + let n = buf.read(&mut output).unwrap(); + assert_eq!(n, 5); + assert_eq!(&output, b"hello"); + + let n = buf.read(&mut output).unwrap(); + assert_eq!(n, 5); + assert_eq!(&output, b" worl"); + + let n = buf.read(&mut output).unwrap(); + assert_eq!(n, 1); + assert_eq!(&output[..1], b"d"); + } + + #[test] + fn test_bufread_trait() { + let mut buf = SlideBuffer::new(1024); + buf.extend_from_slice(b"hello world"); + + { + let slice = buf.fill_buf().unwrap(); + assert_eq!(slice, b"hello world"); + } + + buf.consume(6); + + { + let slice = buf.fill_buf().unwrap(); + assert_eq!(slice, b"world"); + } + } + + #[test] + fn test_indexing() { + let mut buf = SlideBuffer::new(1024); + buf.extend_from_slice(b"hello world"); + + assert_eq!(buf[0], b'h'); + assert_eq!(buf[6], b'w'); + assert_eq!(&buf[0..5], b"hello"); + assert_eq!(&buf[6..], b"world"); + assert_eq!(&buf[..5], b"hello"); + } + + #[test] + fn test_indexing_after_consume() { + let mut buf = SlideBuffer::new(1024); + buf.extend_from_slice(b"hello world"); + buf.consume(6); + + assert_eq!(buf[0], b'w'); + assert_eq!(&buf[0..5], b"world"); + } + + #[test] + fn test_get_u16_be() { + let mut buf = SlideBuffer::new(1024); + buf.extend_from_slice(&[0x12, 0x34, 0x56, 0x78]); + + assert_eq!(buf.get_u16_be(0), Some(0x1234)); + assert_eq!(buf.get_u16_be(2), Some(0x5678)); + assert_eq!(buf.get_u16_be(3), None); + } + + #[test] + fn test_multiple_extend_consume_cycles() { + let mut buf = SlideBuffer::new(100); + + for i in 0..10 { + buf.extend_from_slice(b"0123456789"); + assert_eq!(buf.len(), 10); + buf.consume(10); + assert_eq!(buf.len(), 0); + assert_eq!(buf.remaining_capacity(), 100, "iteration {} failed", i); + } + } + + #[test] + fn test_partial_consume_and_extend() { + let mut buf = SlideBuffer::new(100); + + buf.extend_from_slice(b"hello world"); // 11 bytes + buf.consume(6); // consume "hello " + + buf.extend_from_slice(b"!!!"); // add "!!!" + assert_eq!(buf.as_slice(), b"world!!!"); + + buf.compact(); + assert_eq!(buf.as_slice(), b"world!!!"); + assert_eq!(buf.remaining_capacity(), 100 - 8); + } +} diff --git a/src/snell/snell_handler.rs b/src/snell/snell_handler.rs index 55eb3b0e..362ec354 100644 --- a/src/snell/snell_handler.rs +++ b/src/snell/snell_handler.rs @@ -71,7 +71,7 @@ impl SnellServerHandler { udp_enabled: bool, udp_num_sockets: usize, ) -> Self { - let cipher: ShadowsocksCipher = cipher_name.into(); + let cipher: ShadowsocksCipher = cipher_name.try_into().unwrap(); let key: Arc> = Arc::new(Box::new(SnellKey::new( password, cipher.algorithm().key_len(), @@ -206,7 +206,7 @@ pub struct SnellClientHandler { impl SnellClientHandler { pub fn new(cipher_name: &str, password: &str) -> Self { - let cipher: ShadowsocksCipher = cipher_name.into(); + let cipher: ShadowsocksCipher = cipher_name.try_into().unwrap(); let key: Arc> = Arc::new(Box::new(SnellKey::new( password, cipher.algorithm().key_len(), diff --git a/src/socks_handler.rs b/src/socks_handler.rs index 881fd6dc..2939a647 100644 --- a/src/socks_handler.rs +++ b/src/socks_handler.rs @@ -11,6 +11,7 @@ use crate::stream_reader::StreamReader; use crate::tcp_handler::{ TcpClientHandler, TcpClientSetupResult, TcpServerHandler, TcpServerSetupResult, }; +use crate::uot::{UotV1Stream, UotV2Stream, UOT_V1_MAGIC_ADDRESS, UOT_V2_MAGIC_ADDRESS}; use crate::util::write_all; pub const VER_SOCKS5: u8 = 0x05; @@ -199,6 +200,84 @@ impl TcpServerHandler for SocksTcpServerHandler { let location = read_location(&mut server_stream, &mut stream_reader).await?; + // Check for UDP-over-TCP (UoT) magic addresses + if let Address::Hostname(ref host) = location.address() { + if host == UOT_V1_MAGIC_ADDRESS { + // UoT V1: Multi-destination UDP + // Each packet has: ATYP + address + port + length + data + // Send success response first + write_all(&mut server_stream, connection_success_response).await?; + server_stream.flush().await?; + + let mut uot_stream = UotV1Stream::new(Box::new(server_stream)); + + // Feed any unparsed data (first UoT packet might be in same TCP segment) + let unparsed_data = stream_reader.unparsed_data(); + if !unparsed_data.is_empty() { + log::debug!( + "SOCKS UoT V1: feeding {} bytes of initial data", + unparsed_data.len() + ); + uot_stream.feed_initial_data(unparsed_data); + } + + return Ok(TcpServerSetupResult::MultiDirectionalUdp { + stream: Box::new(uot_stream), + need_initial_flush: false, + override_proxy_provider: NoneOrOne::Unspecified, + num_sockets: 4, + }); + } else if host == UOT_V2_MAGIC_ADDRESS { + // UoT V2: Read request header first + // Request: isConnect(u8) + ATYP + address + port + // Send success response first + write_all(&mut server_stream, connection_success_response).await?; + server_stream.flush().await?; + + let is_connect = stream_reader.read_u8(&mut server_stream).await?; + log::debug!("SOCKS UoT V2: is_connect = {}", is_connect); + + // Read destination address using SOCKS address format + let destination = read_location(&mut server_stream, &mut stream_reader).await?; + log::debug!("SOCKS UoT V2: destination = {:?}", destination); + + if is_connect == 1 { + // V2 Connect mode: Single destination, length-prefixed packets only + let unparsed_data = stream_reader.unparsed_data(); + let mut uot_v2_stream = UotV2Stream::new(Box::new(server_stream)); + if !unparsed_data.is_empty() { + uot_v2_stream.feed_initial_read_data(unparsed_data)?; + } + + return Ok(TcpServerSetupResult::BidirectionalUdp { + remote_location: destination, + stream: Box::new(uot_v2_stream), + need_initial_flush: false, + override_proxy_provider: NoneOrOne::Unspecified, + }); + } else { + // V2 Non-connect mode: Same as V1 (multi-destination) + let mut uot_stream = UotV1Stream::new(Box::new(server_stream)); + + let unparsed_data = stream_reader.unparsed_data(); + if !unparsed_data.is_empty() { + log::debug!( + "SOCKS UoT V2 non-connect: feeding {} bytes of initial data", + unparsed_data.len() + ); + uot_stream.feed_initial_data(unparsed_data); + } + + return Ok(TcpServerSetupResult::MultiDirectionalUdp { + stream: Box::new(uot_stream), + need_initial_flush: false, + override_proxy_provider: NoneOrOne::Unspecified, + num_sockets: 4, + }); + } + } + } + Ok(TcpServerSetupResult::TcpForward { remote_location: location, stream: server_stream, diff --git a/src/sync_adapter.rs b/src/sync_adapter.rs new file mode 100644 index 00000000..84b60560 --- /dev/null +++ b/src/sync_adapter.rs @@ -0,0 +1,56 @@ +/// Sync adapters for bridging async I/O with rustls's synchronous API +/// +/// These adapters allow rustls's synchronous `read_tls()` and `write_tls()` methods +/// to work with Tokio's async I/O primitives. +/// +/// Adapted from tokio-rustls: +/// https://github.com/rustls/tokio-rustls/blob/ba767aeb51611107e7cb6aa756f10a2f49e70926/src/common/mod.rs#L403 +use std::io::{self, Write}; +use std::pin::Pin; +use std::task::{Context, Poll}; +use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; + +/// Adapter to bridge async read to synchronous read for rustls +/// +/// This allows rustls's `read_tls()` method to read from async TCP sockets. +/// When the async socket would block (Poll::Pending), this returns WouldBlock error. +pub struct SyncReadAdapter<'a, 'b, T> { + pub io: &'a mut T, + pub cx: &'a mut Context<'b>, +} + +impl std::io::Read for SyncReadAdapter<'_, '_, T> { + fn read(&mut self, buf: &mut [u8]) -> io::Result { + let mut read_buf = ReadBuf::new(buf); + match Pin::new(&mut self.io).poll_read(self.cx, &mut read_buf) { + Poll::Ready(Ok(())) => Ok(read_buf.filled().len()), + Poll::Ready(Err(e)) => Err(e), + Poll::Pending => Err(io::ErrorKind::WouldBlock.into()), + } + } +} + +/// Adapter to bridge async write to synchronous write for rustls +/// +/// This allows rustls's `write_tls()` method to write to async TCP sockets. +/// When the async socket would block (Poll::Pending), this returns WouldBlock error. +pub struct SyncWriteAdapter<'a, 'b, T> { + pub io: &'a mut T, + pub cx: &'a mut Context<'b>, +} + +impl Write for SyncWriteAdapter<'_, '_, T> { + fn write(&mut self, buf: &[u8]) -> io::Result { + match Pin::new(&mut self.io).poll_write(self.cx, buf) { + Poll::Ready(result) => result, + Poll::Pending => Err(io::ErrorKind::WouldBlock.into()), + } + } + + fn flush(&mut self) -> io::Result<()> { + match Pin::new(&mut self.io).poll_flush(self.cx) { + Poll::Ready(result) => result, + Poll::Pending => Err(io::ErrorKind::WouldBlock.into()), + } + } +} diff --git a/src/tcp/tcp_client_connector.rs b/src/tcp/tcp_client_connector.rs index 950025b8..a54e9743 100644 --- a/src/tcp/tcp_client_connector.rs +++ b/src/tcp/tcp_client_connector.rs @@ -8,7 +8,7 @@ use crate::async_stream::AsyncStream; use crate::config::{ClientConfig, ClientQuicConfig, TcpConfig, Transport}; use crate::quic_stream::QuicStream; use crate::resolver::{resolve_single_address, Resolver}; -use crate::rustls_util::create_client_config; +use crate::rustls_config_util::create_client_config; use crate::socket_util::{ new_reuse_udp_sockets, new_tcp_socket, new_udp_socket, set_tcp_keepalive, }; @@ -177,6 +177,10 @@ impl TcpClientConnector { }) } + pub fn bind_interface(&self) -> &Option { + &self.bind_interface + } + pub fn configure_udp_socket(&self, is_ipv6: bool) -> std::io::Result { let udp_socket = new_udp_socket(is_ipv6, self.bind_interface.clone())?; Ok(udp_socket) diff --git a/src/tcp/tcp_handler.rs b/src/tcp/tcp_handler.rs index 063ce59c..13936e4a 100644 --- a/src/tcp/tcp_handler.rs +++ b/src/tcp/tcp_handler.rs @@ -34,6 +34,11 @@ pub enum TcpServerSetupResult { override_proxy_provider: NoneOrOne>>, num_sockets: usize, }, + SessionBasedUdp { + need_initial_flush: bool, + stream: Box, + override_proxy_provider: NoneOrOne>>, + }, } impl TcpServerSetupResult { @@ -50,6 +55,10 @@ impl TcpServerSetupResult { | TcpServerSetupResult::MultiDirectionalUdp { need_initial_flush: ref mut flush, .. + } + | TcpServerSetupResult::SessionBasedUdp { + need_initial_flush: ref mut flush, + .. } => { *flush = need_initial_flush; } @@ -68,6 +77,10 @@ impl TcpServerSetupResult { | TcpServerSetupResult::MultiDirectionalUdp { override_proxy_provider, .. + } + | TcpServerSetupResult::SessionBasedUdp { + override_proxy_provider, + .. } => override_proxy_provider.is_unspecified(), } } @@ -88,6 +101,10 @@ impl TcpServerSetupResult { | TcpServerSetupResult::MultiDirectionalUdp { override_proxy_provider: ref mut provider, .. + } + | TcpServerSetupResult::SessionBasedUdp { + override_proxy_provider: ref mut provider, + .. } => { *provider = override_proxy_provider; } diff --git a/src/tcp/tcp_handler_util.rs b/src/tcp/tcp_handler_util.rs index 886273ae..4918e44d 100644 --- a/src/tcp/tcp_handler_util.rs +++ b/src/tcp/tcp_handler_util.rs @@ -6,14 +6,15 @@ use rustc_hash::FxHashMap; use crate::client_proxy_selector::{ClientProxySelector, ConnectAction, ConnectRule}; use crate::config::{ - ClientConfig, ClientProxyConfig, ConfigSelection, RuleActionConfig, RuleConfig, - ServerProxyConfig, ShadowTlsServerConfig, ShadowTlsServerHandshakeConfig, ShadowsocksConfig, - TlsClientConfig, TlsServerConfig, WebsocketClientConfig, WebsocketServerConfig, + ClientConfig, ClientProxyConfig, ConfigSelection, RealityServerConfig, RuleActionConfig, + RuleConfig, ServerProxyConfig, ShadowTlsServerConfig, ShadowTlsServerHandshakeConfig, + ShadowsocksConfig, TlsClientConfig, TlsServerConfig, WebsocketClientConfig, + WebsocketServerConfig, }; use crate::http_handler::{HttpTcpClientHandler, HttpTcpServerHandler}; use crate::option_util::NoneOrOne; -use crate::port_forward_handler::PortForwardServerHandler; -use crate::rustls_util::{create_client_config, create_server_config}; +use crate::port_forward_handler::{PortForwardClientHandler, PortForwardServerHandler}; +use crate::rustls_config_util::{create_client_config, create_server_config}; use crate::shadow_tls::{ ShadowTlsClientHandler, ShadowTlsServerTarget, ShadowTlsServerTargetHandshake, }; @@ -22,9 +23,13 @@ use crate::snell::snell_handler::{SnellClientHandler, SnellServerHandler}; use crate::socks_handler::{SocksTcpClientHandler, SocksTcpServerHandler}; use crate::tcp_client_connector::TcpClientConnector; use crate::tcp_handler::{TcpClientHandler, TcpServerHandler}; -use crate::tls_handler::{TlsClientHandler, TlsServerHandler, TlsServerTarget}; +use crate::tls_client_handler::TlsClientHandler; +use crate::tls_server_handler::{ + RealityServerTarget, TlsServerHandler, TlsServerTarget, VisionConfig, +}; use crate::trojan_handler::TrojanTcpHandler; -use crate::vless_handler::{VlessTcpClientHandler, VlessTcpServerHandler}; +use crate::vless::vless_client_handler::VlessTcpClientHandler; +use crate::vless::vless_server_handler::VlessTcpServerHandler; use crate::vmess::{VmessTcpClientHandler, VmessTcpServerHandler}; use crate::websocket::{ WebsocketServerTarget, WebsocketTcpClientHandler, WebsocketTcpServerHandler, @@ -84,6 +89,7 @@ pub fn create_tcp_server_handler( tls_targets, default_tls_target, shadowtls_targets, + reality_targets, tls_buffer_size, } => { let mut all_targets = tls_targets @@ -97,6 +103,11 @@ pub fn create_tcp_server_handler( .map(|(sni, config)| (sni, create_shadow_tls_server_target(config, rules_stack))) .collect::>(); all_targets.extend(shadowtls_targets); + let reality_server_targets = reality_targets + .into_iter() + .map(|(sni, config)| (sni, create_reality_server_target(config, rules_stack))) + .collect::>(); + all_targets.extend(reality_server_targets); Box::new(TlsServerHandler::new( all_targets, default_tls_target, @@ -106,14 +117,8 @@ pub fn create_tcp_server_handler( ServerProxyConfig::Vmess { cipher, user_id, - force_aead, udp_enabled, - } => Box::new(VmessTcpServerHandler::new( - &cipher, - &user_id, - force_aead, - udp_enabled, - )), + } => Box::new(VmessTcpServerHandler::new(&cipher, &user_id, udp_enabled)), ServerProxyConfig::Websocket { targets } => { let server_targets: Vec = targets .into_vec() @@ -142,6 +147,7 @@ fn create_tls_server_target( alpn_protocols, client_ca_certs, client_fingerprints, + vision, protocol, override_rules, } = tls_server_config; @@ -173,6 +179,28 @@ fn create_tls_server_target( ); } + // Create vision_config if vision is enabled + let vision_config = if vision { + // Vision requires VLESS protocol (validated in config/mod.rs) + if let ServerProxyConfig::Vless { + user_id, + udp_enabled, + } = &protocol + { + let user_id_bytes = crate::util::parse_uuid(user_id) + .expect("Invalid user_id UUID") + .into_boxed_slice(); + Some(VisionConfig { + user_id: user_id_bytes, + udp_enabled: *udp_enabled, + }) + } else { + unreachable!("Vision requires VLESS (should be validated during config load)") + } + } else { + None + }; + let handler = create_tcp_server_handler(protocol, rules_stack); let override_proxy_provider = if override_rules.is_empty() { @@ -190,6 +218,7 @@ fn create_tls_server_target( server_config, handler, override_proxy_provider, + vision_config, } } @@ -272,6 +301,90 @@ fn create_shadow_tls_server_target( )) } +fn create_reality_server_target( + reality_server_config: RealityServerConfig, + rules_stack: &mut Vec>, +) -> TlsServerTarget { + let RealityServerConfig { + private_key, + short_ids, + dest, + max_time_diff, + min_client_version, + max_client_version, + vision, + protocol, + override_rules, + } = reality_server_config; + + // Decode private key from base64url (validated during config load) + let private_key_bytes = crate::reality::decode_private_key(&private_key) + .expect("Invalid REALITY private key (should be validated during config load)"); + + // Decode short IDs from hex strings (validated during config load) + // OneOrSome ensures at least one short_id is always present (default is all zeros) + let short_id_bytes: Vec<[u8; 8]> = short_ids + .into_vec() + .into_iter() + .map(|s| { + crate::reality::decode_short_id(&s) + .expect("Invalid REALITY short_id (should be validated during config load)") + }) + .collect(); + + // Create vision_config if vision is enabled + let vision_config = if vision { + // Vision requires VLESS protocol (validated in config/mod.rs) + if let ServerProxyConfig::Vless { + user_id, + udp_enabled, + } = &protocol + { + let user_id_bytes = crate::util::parse_uuid(user_id) + .expect("Invalid user_id UUID") + .into_boxed_slice(); + Some(VisionConfig { + user_id: user_id_bytes, + udp_enabled: *udp_enabled, + }) + } else { + unreachable!("Vision requires VLESS (should be validated during config load)") + } + } else { + None + }; + + // Create inner handler + let handler = create_tcp_server_handler(protocol, rules_stack); + + let override_proxy_provider = if override_rules.is_empty() { + NoneOrOne::None + } else { + rules_stack.push( + override_rules + .clone() + .map(ConfigSelection::unwrap_config) + .into_vec(), + ); + let rules = rules_stack.last().unwrap().clone(); + let provider = NoneOrOne::One(Arc::new(create_tcp_client_proxy_selector(rules))); + rules_stack.pop().unwrap(); + provider + }; + + TlsServerTarget::Reality(RealityServerTarget { + private_key: private_key_bytes, + short_ids: short_id_bytes, + dest, + max_time_diff, + min_client_version, + max_client_version, + handler, + override_proxy_provider, + vision_config, + }) +} + fn create_websocket_server_target( websocket_server_config: WebsocketServerConfig, rules_stack: &mut Vec>, @@ -352,7 +465,10 @@ pub fn create_tcp_client_handler( ClientProxyConfig::Snell(ShadowsocksConfig { cipher, password }) => { Box::new(SnellClientHandler::new(&cipher, &password)) } - ClientProxyConfig::Vless { user_id } => Box::new(VlessTcpClientHandler::new(&user_id)), + ClientProxyConfig::Vless { user_id } => { + // Plain VLESS without TLS + Box::new(VlessTcpClientHandler::new(&user_id)) + } ClientProxyConfig::Trojan { password, shadowsocks, @@ -367,7 +483,7 @@ pub fn create_tcp_client_handler( protocol, key, cert, - shadowtls_password, + vision, } = tls_client_config; let sni_hostname = if sni_hostname.is_unspecified() { @@ -405,33 +521,122 @@ pub fn create_tcp_client_handler( None => "example.com".try_into().unwrap(), }; - let handler = create_tcp_client_handler(*protocol, None); + if vision { + let ClientProxyConfig::Vless { user_id } = protocol.as_ref() else { + // Validated when loading config + unreachable!(); + }; + let user_id_bytes = crate::util::parse_uuid(user_id) + .expect("Invalid user_id UUID") + .into_boxed_slice(); + Box::new(TlsClientHandler::new_vision_vless( + client_config, + tls_buffer_size, + server_name, + user_id_bytes, + )) + } else { + let handler = create_tcp_client_handler(*protocol, None); - match shadowtls_password { - None => Box::new(TlsClientHandler::new( + Box::new(TlsClientHandler::new( client_config, tls_buffer_size, server_name, handler, - )), - Some(password) => { - // TODO: Ensure client_config is suitable for TLS 1.3. - // Rustls ClientConfig by default supports TLS 1.3 if server does. - // The server handler enforces TLS 1.3 from client and for negotiation. - Box::new(ShadowTlsClientHandler::new( - password, - client_config, - server_name, - handler, - )) - } + )) } } - ClientProxyConfig::Vmess { - cipher, - user_id, - aead, - } => Box::new(VmessTcpClientHandler::new(&cipher, &user_id, aead)), + ClientProxyConfig::Reality { + public_key, + short_id, + sni_hostname, + vision, + protocol, + } => { + eprintln!("========== CREATING REALITY CLIENT HANDLER =========="); + + // 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 + 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"); + } + }; + + if vision { + let ClientProxyConfig::Vless { user_id } = protocol.as_ref() else { + unreachable!("Vision requires VLESS (should be validated during config load)") + }; + let user_id_bytes = crate::util::parse_uuid(user_id) + .expect("Invalid user_id UUID") + .into_boxed_slice(); + Box::new( + crate::reality_client_handler::RealityClientHandler::new_vision_vless( + public_key_bytes, + short_id_bytes, + server_name, + user_id_bytes, + ), + ) + } else { + let inner_handler = create_tcp_client_handler(*protocol, None); + Box::new(crate::reality_client_handler::RealityClientHandler::new( + public_key_bytes, + short_id_bytes, + server_name, + inner_handler, + )) + } + } + ClientProxyConfig::ShadowTls { + password, + sni_hostname, + protocol, + } => { + // ShadowTLS client handler + let sni_hostname = sni_hostname.or(default_sni_hostname); + let enable_sni = sni_hostname.is_some(); + + let server_name = match sni_hostname { + Some(s) => rustls::pki_types::ServerName::try_from(s).unwrap(), + None => "example.com".try_into().unwrap(), // Fallback + }; + + // Create TLS config for ShadowTLS + // TODO: Ensure client_config is suitable for TLS 1.3. + // Rustls ClientConfig by default supports TLS 1.3 if server does. + // The server handler enforces TLS 1.3 from client and for negotiation. + let client_config = Arc::new(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 + )); + + let handler = create_tcp_client_handler(*protocol, None); + + Box::new(ShadowTlsClientHandler::new( + password, + client_config, + server_name, + handler, + )) + } + ClientProxyConfig::Vmess { cipher, user_id } => { + Box::new(VmessTcpClientHandler::new(&cipher, &user_id)) + } ClientProxyConfig::Websocket(websocket_client_config) => { let WebsocketClientConfig { matching_path, @@ -449,6 +654,7 @@ pub fn create_tcp_client_handler( handler, )) } + ClientProxyConfig::PortForward => Box::new(PortForwardClientHandler), } } diff --git a/src/tcp/tcp_server.rs b/src/tcp/tcp_server.rs index 370f6be7..8d91068c 100644 --- a/src/tcp/tcp_server.rs +++ b/src/tcp/tcp_server.rs @@ -15,6 +15,7 @@ use crate::config::{BindLocation, ConfigSelection, ServerConfig, TcpConfig}; use crate::copy_bidirectional::copy_bidirectional; use crate::copy_bidirectional_message::copy_bidirectional_message; use crate::copy_multidirectional_message::copy_multidirectional_message; +use crate::copy_session_messages::copy_session_messages; use crate::resolver::{resolve_single_address, NativeResolver, Resolver}; use crate::socket_util::{new_tcp_listener, set_tcp_keepalive}; use crate::tcp_client_connector::TcpClientConnector; @@ -22,6 +23,7 @@ use crate::tcp_handler::{TcpServerHandler, TcpServerSetupResult}; use crate::tcp_handler_util::{create_tcp_client_proxy_selector, create_tcp_server_handler}; use crate::udp_message_stream::UdpMessageStream; use crate::udp_multi_message_stream::UdpMultiMessageStream; +use crate::udp_session_message_stream::{UdpSessionMessageStream, UdpSocketConfig}; use crate::util::write_all; async fn run_tcp_server( @@ -332,6 +334,56 @@ where } } } + TcpServerSetupResult::SessionBasedUdp { + stream: mut server_stream, + need_initial_flush: server_need_initial_flush, + override_proxy_provider, + } => { + let selected_proxy_provider = if override_proxy_provider.is_one() { + override_proxy_provider.unwrap() + } else { + client_proxy_selector + }; + let action = selected_proxy_provider.default_decision(); + match action { + ConnectDecision::Allow { + client_proxy, + remote_location: _, + } => { + // Create UdpSessionMessageStream with user's interface configuration + let config = UdpSocketConfig { + bind_interface: client_proxy.bind_interface().clone(), + }; + let mut session_manager = UdpSessionMessageStream::new(config); + + // IMPORTANT: We do NOT create any sessions upfront + // Sessions will be created on-demand as XUDP frames arrive + // This is handled by the copy_session_messages function + // which will call poll_write_session_message, which in turn + // will create sessions via SessionUdpManager::create_session + + let copy_result = copy_session_messages( + &mut server_stream, + &mut session_manager, + server_need_initial_flush, + false, + ) + .await; + + // TODO: add async trait ext and make this work + //let (_, _) = futures::join!(server_stream.shutdown_message(), client_stream.shutdown_message()); + + copy_result?; + Ok(()) + } + ConnectDecision::Block => { + warn!("Blocked session-based udp forward, because the default action is to block."); + // TODO: add async trait ext and make this work + // let _ = server_stream.shutdown_message().await; + Ok(()) + } + } + } } } diff --git a/src/tls_client_handler.rs b/src/tls_client_handler.rs new file mode 100644 index 00000000..56df6d3b --- /dev/null +++ b/src/tls_client_handler.rs @@ -0,0 +1,112 @@ +use std::sync::Arc; + +use async_trait::async_trait; + +use crate::address::NetLocation; +use crate::async_stream::AsyncStream; +use crate::tcp_handler::{TcpClientHandler, TcpClientSetupResult}; + +#[derive(Debug)] +pub struct TlsClientHandler { + pub client_config: Arc, + pub tls_buffer_size: Option, + pub server_name: rustls::pki_types::ServerName<'static>, + pub handler: TlsInnerClientHandler, +} + +#[derive(Debug)] +pub enum TlsInnerClientHandler { + Default(Box), + VisionVless { uuid: Box<[u8]> }, +} + +impl TlsClientHandler { + pub fn new( + client_config: Arc, + tls_buffer_size: Option, + server_name: rustls::pki_types::ServerName<'static>, + handler: Box, + ) -> Self { + Self { + client_config, + tls_buffer_size, + server_name, + handler: TlsInnerClientHandler::Default(handler), + } + } + + pub fn new_vision_vless( + client_config: Arc, + tls_buffer_size: Option, + server_name: rustls::pki_types::ServerName<'static>, + uuid: Box<[u8]>, + ) -> Self { + Self { + client_config, + tls_buffer_size, + server_name, + handler: TlsInnerClientHandler::VisionVless { uuid }, + } + } +} + +#[async_trait] +impl TcpClientHandler for TlsClientHandler { + async fn setup_client_stream( + &self, + server_stream: &mut Box, + mut client_stream: Box, + remote_location: NetLocation, + ) -> std::io::Result { + use crate::crypto::{CryptoConnection, CryptoTlsStream}; + use crate::rustls_handshake::perform_handshake; + + // Create rustls ClientConnection + let client_conn = + rustls::ClientConnection::new(self.client_config.clone(), self.server_name.clone()) + .map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("Failed to create client connection: {e}"), + ) + })?; + + // Wrap in rustls::Connection enum for handshake + let mut rustls_connection = rustls::Connection::Client(client_conn); + + // Set buffer limits if configured + if let Some(size) = self.tls_buffer_size { + rustls_connection.set_buffer_limit(Some(size)); + } + + // Perform the TLS handshake using the generic helper + // This works for both TLS 1.2 and TLS 1.3 + perform_handshake(&mut rustls_connection, &mut client_stream, 16384).await?; + + // Extract the ClientConnection back from the enum + let client_conn = match rustls_connection { + rustls::Connection::Client(conn) => conn, + _ => unreachable!("We created a Client variant"), + }; + + // Wrap in CryptoTlsStream + let connection = CryptoConnection::new_rustls_client(client_conn); + let tls_stream = CryptoTlsStream::new(client_stream, connection); + + match self.handler { + TlsInnerClientHandler::Default(ref handler) => { + handler + .setup_client_stream(server_stream, Box::new(tls_stream), remote_location) + .await + } + TlsInnerClientHandler::VisionVless { ref uuid } => { + crate::vless::vless_client_handler::setup_custom_tls_vision_vless_client_stream( + tls_stream, + uuid, + &remote_location, + ) + .await + } + } + } +} diff --git a/src/tls_handler.rs b/src/tls_handler.rs deleted file mode 100644 index 2d87af84..00000000 --- a/src/tls_handler.rs +++ /dev/null @@ -1,210 +0,0 @@ -use std::sync::Arc; - -use async_trait::async_trait; -use rustc_hash::FxHashMap; -use tokio_rustls::TlsAcceptor; - -use crate::address::NetLocation; -use crate::async_stream::AsyncStream; -use crate::client_proxy_selector::ClientProxySelector; -use crate::option_util::NoneOrOne; -use crate::resolver::{NativeResolver, Resolver}; -use crate::shadow_tls::{ - feed_server_connection, read_client_hello, setup_shadowtls_server_stream, ParsedClientHello, - ShadowTlsServerTarget, -}; -use crate::tcp_client_connector::TcpClientConnector; -use crate::tcp_handler::{ - TcpClientHandler, TcpClientSetupResult, TcpServerHandler, TcpServerSetupResult, -}; - -#[derive(Debug)] -pub struct TlsServerHandler { - sni_targets: FxHashMap, - default_target: Option, - // used to resolve handshake server hostnames - shadowtls_resolver: Arc, - tls_buffer_size: Option, -} - -#[derive(Debug)] -#[allow(clippy::large_enum_variant)] -pub enum TlsServerTarget { - Tls { - server_config: Arc, - handler: Box, - override_proxy_provider: NoneOrOne>>, - }, - ShadowTls(ShadowTlsServerTarget), -} - -impl TlsServerHandler { - pub fn new( - sni_targets: FxHashMap, - default_target: Option, - tls_buffer_size: Option, - ) -> Self { - Self { - sni_targets, - default_target, - shadowtls_resolver: Arc::new(NativeResolver::new()), - tls_buffer_size, - } - } -} - -#[async_trait] -impl TcpServerHandler for TlsServerHandler { - async fn setup_server_stream( - &self, - mut server_stream: Box, - ) -> std::io::Result { - let parsed_client_hello = read_client_hello(&mut server_stream).await?; - - let target = match parsed_client_hello.requested_server_name.as_ref() { - None => match self.default_target { - Some(ref t) => t, - None => { - return Err(std::io::Error::other( - "No default target for unspecified SNI", - )); - } - }, - Some(hostname) => match self.sni_targets.get(hostname) { - Some(t) => t, - None => match self.default_target { - Some(ref t) => t, - None => { - return Err(std::io::Error::other(format!( - "No default target for unknown SNI: {hostname}" - ))); - } - }, - }, - }; - - match target { - TlsServerTarget::Tls { - ref server_config, - ref handler, - ref override_proxy_provider, - } => { - let ParsedClientHello { - client_hello_frame, - client_reader, - .. - } = parsed_client_hello; - let tls_acceptor = TlsAcceptor::from(server_config.clone()); - - let accept_future = { - let mut accept_error: Option = None; - - let accept_future = tls_acceptor.accept_with(server_stream, |server_conn| { - if let Some(size) = self.tls_buffer_size { - server_conn.set_buffer_limit(Some(size)); - } - - if let Err(e) = feed_server_connection(server_conn, &client_hello_frame) { - let _ = accept_error.insert(std::io::Error::other(format!( - "Failed to feed initial frame to server connection: {e}" - ))); - return; - } - let unparsed_data = client_reader.unparsed_data(); - if !unparsed_data.is_empty() { - if let Err(e) = feed_server_connection(server_conn, unparsed_data) { - let _ = accept_error.insert(std::io::Error::other(format!( - "Failed to feed unparsed data to server connection: {e}" - ))); - return; - } - } - if let Err(e) = server_conn.process_new_packets() { - let _ = accept_error.insert(std::io::Error::other(format!( - "Failed to process new packets: {e}" - ))); - } - }); - - if let Some(e) = accept_error { - return Err(e); - } - - accept_future - }; - - let tls_stream: Box = Box::new(accept_future.await?); - - let mut target_setup_result = handler.setup_server_stream(tls_stream).await; - if let Ok(ref mut setup_result) = target_setup_result { - setup_result.set_need_initial_flush(true); - if setup_result.override_proxy_provider_unspecified() - && !override_proxy_provider.is_unspecified() - { - setup_result.set_override_proxy_provider(override_proxy_provider.clone()); - } - } - - target_setup_result - } - TlsServerTarget::ShadowTls(ref target) => { - setup_shadowtls_server_stream( - server_stream, - target, - parsed_client_hello, - &self.shadowtls_resolver, - ) - .await - } - } - } -} - -#[derive(Debug)] -pub struct TlsClientHandler { - pub client_config: Arc, - pub tls_buffer_size: Option, - pub server_name: rustls::pki_types::ServerName<'static>, - pub handler: Box, -} - -impl TlsClientHandler { - pub fn new( - client_config: Arc, - tls_buffer_size: Option, - server_name: rustls::pki_types::ServerName<'static>, - handler: Box, - ) -> Self { - Self { - client_config, - tls_buffer_size, - server_name, - handler, - } - } -} - -#[async_trait] -impl TcpClientHandler for TlsClientHandler { - async fn setup_client_stream( - &self, - server_stream: &mut Box, - client_stream: Box, - remote_location: NetLocation, - ) -> std::io::Result { - let connector: tokio_rustls::TlsConnector = self.client_config.clone().into(); - let tls_stream = Box::new( - connector - .connect_with(self.server_name.clone(), client_stream, |client_conn| { - if let Some(size) = self.tls_buffer_size { - client_conn.set_buffer_limit(Some(size)); - } - }) - .await?, - ); - - self.handler - .setup_client_stream(server_stream, tls_stream, remote_location) - .await - } -} diff --git a/src/tls_server_handler.rs b/src/tls_server_handler.rs new file mode 100644 index 00000000..507ca099 --- /dev/null +++ b/src/tls_server_handler.rs @@ -0,0 +1,397 @@ +use std::sync::{Arc, LazyLock}; + +use async_trait::async_trait; +use rustc_hash::FxHashMap; + +use crate::address::NetLocation; +use crate::async_stream::AsyncStream; +use crate::client_proxy_selector::ClientProxySelector; +use crate::option_util::NoneOrOne; +use crate::resolver::{NativeResolver, Resolver}; +use crate::shadow_tls::{ + read_client_hello, setup_shadowtls_server_stream, ParsedClientHello, ShadowTlsServerTarget, +}; +use crate::tcp_client_connector::TcpClientConnector; +use crate::tcp_handler::{TcpServerHandler, TcpServerSetupResult}; + +#[derive(Debug, Clone)] +pub struct VisionConfig { + pub user_id: Box<[u8]>, + pub udp_enabled: bool, +} + +#[derive(Debug)] +pub struct RealityServerTarget { + pub private_key: [u8; 32], + pub short_ids: Vec<[u8; 8]>, + pub dest: NetLocation, + pub max_time_diff: Option, // in milliseconds + pub min_client_version: Option<[u8; 3]>, + pub max_client_version: Option<[u8; 3]>, + pub handler: Box, + pub override_proxy_provider: NoneOrOne>>, + pub vision_config: Option, +} + +#[derive(Debug)] +#[allow(clippy::large_enum_variant)] +pub enum TlsServerTarget { + Tls { + server_config: Arc, + handler: Box, + override_proxy_provider: NoneOrOne>>, + vision_config: Option, + }, + ShadowTls(ShadowTlsServerTarget), + Reality(RealityServerTarget), +} + +#[derive(Debug)] +pub struct TlsServerHandler { + sni_targets: FxHashMap, + default_target: Option, + // used to resolve handshake server hostnames + shadowtls_resolver: Arc, + tls_buffer_size: Option, +} + +impl TlsServerHandler { + pub fn new( + sni_targets: FxHashMap, + default_target: Option, + tls_buffer_size: Option, + ) -> Self { + Self { + sni_targets, + default_target, + shadowtls_resolver: Arc::new(NativeResolver::new()), + tls_buffer_size, + } + } +} + +#[async_trait] +impl TcpServerHandler for TlsServerHandler { + async fn setup_server_stream( + &self, + mut server_stream: Box, + ) -> std::io::Result { + let parsed_client_hello = read_client_hello(&mut server_stream).await?; + + let target = match parsed_client_hello.requested_server_name.as_ref() { + None => match self.default_target { + Some(ref t) => t, + None => { + return Err(std::io::Error::other( + "No default target for unspecified SNI", + )); + } + }, + Some(hostname) => match self.sni_targets.get(hostname) { + Some(t) => t, + None => match self.default_target { + Some(ref t) => t, + None => { + return Err(std::io::Error::other(format!( + "No default target for unknown SNI: {hostname}" + ))); + } + }, + }, + }; + + match target { + TlsServerTarget::Tls { + ref server_config, + ref handler, + ref override_proxy_provider, + ref vision_config, + } => { + use crate::crypto::{CryptoConnection, CryptoTlsStream}; + use crate::rustls_handshake::perform_handshake; + + let ParsedClientHello { + client_hello_frame, + client_reader, + .. + } = parsed_client_hello; + + // Create rustls ServerConnection + let server_conn = + rustls::ServerConnection::new(server_config.clone()).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("Failed to create server connection: {e}"), + ) + })?; + + // Wrap in rustls::Connection enum for handshake + let mut rustls_connection = rustls::Connection::Server(server_conn); + + // Set buffer limits if configured + if let Some(size) = self.tls_buffer_size { + rustls_connection.set_buffer_limit(Some(size)); + } + + // Feed the ClientHello we already parsed + { + use std::io::Cursor; + + let mut cursor = Cursor::new(&client_hello_frame); + rustls_connection.read_tls(&mut cursor)?; + + let unparsed_data = client_reader.unparsed_data(); + if !unparsed_data.is_empty() { + let mut cursor = Cursor::new(unparsed_data); + rustls_connection.read_tls(&mut cursor)?; + } + + rustls_connection.process_new_packets().map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("Failed to process ClientHello: {e}"), + ) + })?; + } + + // Perform the TLS handshake using the generic helper + // This works for both TLS 1.2 and TLS 1.3 + perform_handshake(&mut rustls_connection, &mut server_stream, 16384).await?; + + // Extract the ServerConnection back from the enum + let server_conn = match rustls_connection { + rustls::Connection::Server(conn) => conn, + _ => unreachable!("We created a Server variant"), + }; + + // Wrap in CryptoTlsStream + let connection = CryptoConnection::new_rustls_server(server_conn); + let tls_stream = CryptoTlsStream::new(server_stream, connection); + + let mut target_setup_result = if let Some(vision_cfg) = vision_config { + // Vision is enabled - call setup_custom_tls_vision_vless_server_stream + crate::vless::vless_server_handler::setup_custom_tls_vision_vless_server_stream( + tls_stream, + &vision_cfg.user_id, + vision_cfg.udp_enabled, + ) + .await + } else { + // Regular TLS - call inner handler + handler.setup_server_stream(Box::new(tls_stream)).await + }; + + if let Ok(ref mut setup_result) = target_setup_result { + setup_result.set_need_initial_flush(true); + if setup_result.override_proxy_provider_unspecified() + && !override_proxy_provider.is_unspecified() + { + setup_result.set_override_proxy_provider(override_proxy_provider.clone()); + } + } + + target_setup_result + } + TlsServerTarget::ShadowTls(ref target) => { + setup_shadowtls_server_stream( + server_stream, + target, + parsed_client_hello, + &self.shadowtls_resolver, + ) + .await + } + TlsServerTarget::Reality(ref target) => { + use crate::crypto::{CryptoConnection, CryptoTlsStream}; + use crate::reality::{RealityServerConfig, RealityServerConnection}; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + // Note: Vision over Reality would require updating setup_vision_vless_server_stream + // to accept CryptoTlsStream instead of tokio_rustls::server::TlsStream. + // For now, vision_config should always be None for Reality targets. + // This will be enforced in tcp_handler_util.rs + + log::debug!("REALITY DEBUG: setup_reality_server_stream called"); + let client_hello_frame = &parsed_client_hello.client_hello_frame; + log::debug!( + "REALITY DEBUG: ClientHello frame length: {}", + client_hello_frame.len() + ); + + // Create buffered REALITY connection + log::debug!("REALITY DEBUG: Creating buffered RealityServerConnection"); + let reality_config = RealityServerConfig { + private_key: target.private_key, + short_ids: target.short_ids.clone(), + dest: target.dest.clone(), + max_time_diff: target.max_time_diff, + min_client_version: target.min_client_version, + max_client_version: target.max_client_version, + }; + + let mut reality_conn = RealityServerConnection::new(reality_config)?; + + // Feed the ClientHello to the connection + log::debug!("REALITY DEBUG: Feeding ClientHello to connection via read_tls"); + { + let mut cursor = std::io::Cursor::new(client_hello_frame); + reality_conn.read_tls(&mut cursor)?; + } + + // Process the ClientHello to advance handshake (this validates everything) + log::debug!("REALITY DEBUG: Processing ClientHello via process_new_packets"); + if let Err(e) = reality_conn.process_new_packets() { + // Check if this is an authentication failure + if e.kind() == std::io::ErrorKind::PermissionDenied { + log::warn!( + "REALITY: Authentication failed, falling back to dest: {} - reason: {}", + target.dest, + e + ); + // Implement fallback mechanism + return reality_fallback_to_dest( + server_stream, + client_hello_frame, + &target.dest, + ) + .await; + } else { + // Other errors should propagate normally + return Err(e); + } + } + + // Write the server's handshake response to the stream + log::debug!("REALITY DEBUG: Writing server handshake response"); + { + let mut write_buf = Vec::new(); + while reality_conn.wants_write() { + reality_conn.write_tls(&mut write_buf)?; + } + if !write_buf.is_empty() { + server_stream.write_all(&write_buf).await?; + server_stream.flush().await?; + } + } + + // Read client's Finished message + log::debug!("REALITY DEBUG: Reading client Finished"); + { + let mut buf = vec![0u8; 4096]; + let n = server_stream.read(&mut buf).await?; + if n == 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + "EOF while waiting for client Finished", + )); + } + let mut cursor = std::io::Cursor::new(&buf[..n]); + reality_conn.read_tls(&mut cursor)?; + } + + // Process client Finished to complete handshake + log::debug!("REALITY DEBUG: Processing client Finished"); + let io_state = reality_conn.process_new_packets()?; + log::debug!( + "REALITY DEBUG: After processing, plaintext_bytes_to_read={}", + io_state.plaintext_bytes_to_read() + ); + + // Wrap in Connection enum and CryptoTlsStream + log::debug!("REALITY DEBUG: Wrapping in CryptoTlsStream"); + let connection = CryptoConnection::new_reality_server(reality_conn); + let tls_stream = CryptoTlsStream::new(server_stream, connection); + + log::debug!("REALITY DEBUG: TLS 1.3 handshake completed successfully"); + + // Check if Vision is enabled + let mut target_setup_result = if let Some(vision_cfg) = &target.vision_config { + log::info!("REALITY+Vision: Setting up VLESS+VISION stream"); + // Vision is enabled - call setup_vision_vless_server_stream_from_custom_tls + crate::vless::vless_server_handler::setup_custom_tls_vision_vless_server_stream( + tls_stream, + &vision_cfg.user_id, + vision_cfg.udp_enabled, + ) + .await + } else { + // Regular REALITY (no Vision) - call inner handler + target + .handler + .setup_server_stream(Box::new(tls_stream)) + .await + }; + if let Ok(ref mut setup_result) = target_setup_result { + setup_result.set_need_initial_flush(true); + if setup_result.override_proxy_provider_unspecified() + && !target.override_proxy_provider.is_unspecified() + { + setup_result + .set_override_proxy_provider(target.override_proxy_provider.clone()); + } + } + + target_setup_result + } + } + } +} + +/// Fallback mechanism for Reality authentication failures +/// +/// When a client fails Reality authentication (invalid short_id, timestamp, or version), +/// instead of dropping the connection, we transparently forward it to the configured +/// destination server. This makes the server indistinguishable from a legitimate +/// reverse proxy or CDN, defeating active probing attacks. +async fn reality_fallback_to_dest( + mut client_stream: Box, + client_hello_bytes: &[u8], + dest: &NetLocation, +) -> std::io::Result { + use tokio::io::AsyncWriteExt; + use tokio::net::TcpStream; + + log::info!("REALITY FALLBACK: Connecting to dest server: {}", dest); + + /// Shared resolver instance for Reality fallback connections + /// Using LazyLock to avoid creating a new resolver for each fallback connection + static REALITY_FALLBACK_RESOLVER: LazyLock> = + LazyLock::new(|| Arc::new(NativeResolver::new())); + + // Resolve and connect to the fallback destination using shared resolver + let dest_addr = + crate::resolver::resolve_single_address(&REALITY_FALLBACK_RESOLVER, dest).await?; + + log::debug!("REALITY FALLBACK: Resolved {} to {}", dest, dest_addr); + + let mut dest_stream: Box = Box::new(TcpStream::connect(dest_addr).await?); + + log::info!( + "REALITY FALLBACK: Connected to dest, forwarding ClientHello ({} bytes)", + client_hello_bytes.len() + ); + + // Forward the raw ClientHello that we received from the client + dest_stream.write_all(client_hello_bytes).await?; + dest_stream.flush().await?; + + log::info!("REALITY FALLBACK: ClientHello forwarded, entering bidirectional mode"); + + // Enter bidirectional forwarding mode + // This will copy all traffic between client and dest transparently + crate::copy_bidirectional::copy_bidirectional( + &mut *client_stream, + &mut *dest_stream, + false, // client doesn't need initial flush + false, // dest doesn't need initial flush + ) + .await?; + + log::info!("REALITY FALLBACK: Bidirectional forwarding completed"); + + // Connection handled, no further setup needed + // Return an error to signal that this connection has been fully handled + Err(std::io::Error::other( + "Connection handled by Reality fallback mechanism", + )) +} diff --git a/src/trojan_handler.rs b/src/trojan_handler.rs index 7edfc713..0fcb10d0 100644 --- a/src/trojan_handler.rs +++ b/src/trojan_handler.rs @@ -2,6 +2,7 @@ use std::sync::Arc; use async_trait::async_trait; use aws_lc_rs::digest::SHA224; +use subtle::ConstantTimeEq; use tokio::io::AsyncWriteExt; use crate::address::NetLocation; @@ -38,7 +39,7 @@ impl TrojanTcpHandler { cipher, password: shadowsocks_password, } = config; - let cipher: ShadowsocksCipher = cipher.as_str().into(); + let cipher: ShadowsocksCipher = cipher.as_str().try_into().unwrap(); let key: Arc> = Arc::new(Box::new(DefaultKey::new( shadowsocks_password, cipher.algorithm().key_len(), @@ -88,10 +89,9 @@ impl TcpServerHandler for TrojanTcpHandler { ))); } - for (b1, b2) in self.password_hash.iter().zip(received_hash.iter()) { - if b1 != b2 { - return Err(std::io::Error::other("Invalid password hash")); - } + // Use constant-time comparison to prevent timing attacks + if self.password_hash.ct_eq(received_hash).unwrap_u8() == 0 { + return Err(std::io::Error::other("Invalid password hash")); } let command_type = stream_reader.read_u8(&mut server_stream).await?; diff --git a/src/udp_session_message_stream.rs b/src/udp_session_message_stream.rs new file mode 100644 index 00000000..66debb72 --- /dev/null +++ b/src/udp_session_message_stream.rs @@ -0,0 +1,266 @@ +use std::collections::HashMap; +use std::net::SocketAddr; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use tokio::io::ReadBuf; +use tokio::sync::mpsc; + +use crate::async_stream::{ + AsyncFlushMessage, AsyncPing, AsyncReadSessionMessage, AsyncSessionMessageStream, + AsyncShutdownMessage, AsyncWriteSessionMessage, +}; +use crate::socket_util::new_udp_socket; + +/// Configuration for creating UDP sockets for each session. +/// Preserves user's interface binding configuration. +#[derive(Clone, Debug)] +pub struct UdpSocketConfig { + /// Optional network interface to bind sockets to (e.g., "tun0" for VPN) + pub bind_interface: Option, +} + +/// Message from a session socket reader task containing data received from a UDP peer. +struct SessionMessage { + session_id: u16, + data: Vec, + source_addr: SocketAddr, +} + +/// Session-based UDP message stream that manages dedicated UDP sockets for each session. +/// Used for protocols like XUDP that require socket-per-session architecture. +/// This eliminates the need for destination-based session mapping by using socket identity +/// as session identity. +pub struct UdpSessionMessageStream { + /// Receives messages from all session socket readers + receiver: mpsc::UnboundedReceiver, + /// Used to send to session socket readers (cloned for each new session) + sender: mpsc::UnboundedSender, + /// Maps session ID to the sender for that session's writer task + session_writers: HashMap, SocketAddr)>>, + /// Configuration for creating new UDP sockets + config: UdpSocketConfig, + /// Counter for session IDs (not used yet, reserved for future session creation API) + _next_session_id: u16, +} + +impl UdpSessionMessageStream { + /// Creates a new UdpSessionMessageStream with the specified configuration. + pub fn new(config: UdpSocketConfig) -> Self { + // Use unbounded channel for maximum performance (no backpressure) + let (sender, receiver) = mpsc::unbounded_channel(); + + Self { + receiver, + sender, + session_writers: HashMap::new(), + config, + _next_session_id: 0, + } + } + + /// Creates a new session with a dedicated UDP socket. + /// Returns the session ID that was created. + /// + /// This spawns two tasks: + /// - Reader task: reads from UDP socket and sends to manager's channel + /// - Writer task: receives write requests and sends to UDP socket + pub fn create_session(&mut self, session_id: u16, is_ipv6: bool) -> std::io::Result<()> { + // Create dedicated UDP socket for this session + let socket = new_udp_socket(is_ipv6, self.config.bind_interface.clone())?; + let socket = Arc::new(socket); + + // Create channel for write requests to this session + let (write_tx, mut write_rx) = mpsc::unbounded_channel::<(Vec, SocketAddr)>(); + + // Spawn reader task + let reader_socket = socket.clone(); + let reader_sender = self.sender.clone(); + tokio::spawn(async move { + let mut buf = vec![0u8; 65536]; + loop { + match reader_socket.recv_from(&mut buf).await { + Ok((len, source_addr)) => { + let data = buf[..len].to_vec(); + let msg = SessionMessage { + session_id, + data, + source_addr, + }; + + // For unbounded channels, send() doesn't need to be awaited + if reader_sender.send(msg).is_err() { + // Manager dropped, exit reader task + log::debug!( + "[Session {}] Manager dropped, stopping reader", + session_id + ); + break; + } + } + Err(e) => { + log::error!("[Session {}] UDP recv error: {}", session_id, e); + break; + } + } + } + }); + + // Spawn writer task + let writer_socket = socket.clone(); + tokio::spawn(async move { + while let Some((data, target_addr)) = write_rx.recv().await { + match writer_socket.send_to(&data, target_addr).await { + Ok(sent) => { + log::debug!( + "[Session {}] Sent {} bytes to {}", + session_id, + sent, + target_addr + ); + } + Err(e) => { + log::error!("[Session {}] UDP send error: {}", session_id, e); + break; + } + } + } + log::debug!("[Session {}] Writer task exiting", session_id); + }); + + // Store the writer channel + self.session_writers.insert(session_id, write_tx); + + log::info!( + "[Session {}] Created dedicated UDP socket (IPv6: {})", + session_id, + is_ipv6 + ); + Ok(()) + } + + /// Removes a session and its associated socket. + /// The socket will be closed when the writer task exits. + /// TODO: why isn't this used? add session cleanup + #[allow(dead_code)] + pub fn remove_session(&mut self, session_id: u16) { + if self.session_writers.remove(&session_id).is_some() { + log::info!("[Session {}] Removed session", session_id); + } + } +} + +impl AsyncReadSessionMessage for UdpSessionMessageStream { + fn poll_read_session_message( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + // Poll the receiver for messages from any session socket + match self.receiver.poll_recv(cx) { + Poll::Ready(Some(msg)) => { + // Write the data to the buffer + let len = msg.data.len(); + if buf.remaining() < len { + return Poll::Ready(Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!( + "Buffer too small: need {} bytes, have {}", + len, + buf.remaining() + ), + ))); + } + + buf.put_slice(&msg.data); + Poll::Ready(Ok((msg.session_id, msg.source_addr))) + } + Poll::Ready(None) => { + // All session sockets dropped + Poll::Ready(Err(std::io::Error::new( + std::io::ErrorKind::ConnectionAborted, + "All session sockets closed", + ))) + } + Poll::Pending => Poll::Pending, + } + } +} + +impl AsyncWriteSessionMessage for UdpSessionMessageStream { + fn poll_write_session_message( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + session_id: u16, + buf: &[u8], + target: &SocketAddr, + ) -> Poll> { + // Auto-create session if it doesn't exist yet + if !self.session_writers.contains_key(&session_id) { + let is_ipv6 = target.is_ipv6(); + if let Err(e) = self.create_session(session_id, is_ipv6) { + return Poll::Ready(Err(e)); + } + } + + // Get the writer channel for this session + match self.session_writers.get(&session_id) { + Some(writer) => { + let data = buf.to_vec(); + let target_addr = *target; + + // Send write request to the writer task + // For unbounded channels, send() doesn't need to be awaited + if writer.send((data, target_addr)).is_err() { + return Poll::Ready(Err(std::io::Error::new( + std::io::ErrorKind::BrokenPipe, + format!("Session {} writer task closed", session_id), + ))); + } + + Poll::Ready(Ok(())) + } + None => { + // Should never happen because we just created it above + Poll::Ready(Err(std::io::Error::new( + std::io::ErrorKind::NotFound, + format!("Session {} not found", session_id), + ))) + } + } + } +} + +impl AsyncFlushMessage for UdpSessionMessageStream { + fn poll_flush_message( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll> { + // UDP is connectionless, no flushing needed + Poll::Ready(Ok(())) + } +} + +impl AsyncShutdownMessage for UdpSessionMessageStream { + fn poll_shutdown_message( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll> { + // Close all sessions + self.session_writers.clear(); + Poll::Ready(Ok(())) + } +} + +impl AsyncPing for UdpSessionMessageStream { + fn supports_ping(&self) -> bool { + false + } + + fn poll_write_ping(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(false)) + } +} + +impl AsyncSessionMessageStream for UdpSessionMessageStream {} diff --git a/src/uot/mod.rs b/src/uot/mod.rs new file mode 100644 index 00000000..2d359500 --- /dev/null +++ b/src/uot/mod.rs @@ -0,0 +1,41 @@ +//! SagerNet UDP-over-TCP (UoT) protocol implementation +//! +//! This module implements the sing-box UoT protocol for tunneling UDP over TCP. +//! It supports both V1 and V2 formats. +//! +//! ## Magic Addresses +//! - V1: `sp.udp-over-tcp.arpa` - Multi-destination mode, each packet has full address +//! - V2: `sp.v2.udp-over-tcp.arpa` - Optional connect mode for single destination +//! +//! ## V1 Packet Format +//! ```text +//! | ATYP | address | port | length | data | +//! | u8 | variable | u16be | u16be | variable | +//! ``` +//! +//! ## V2 Request Format +//! ```text +//! | isConnect | ATYP | address | port | +//! | u8 | u8 | variable | u16be | +//! ``` +//! +//! If isConnect=1, subsequent packets are length-prefixed only (V2 connect mode). +//! If isConnect=0, subsequent packets use V1 format (multi-destination). +//! +//! ## ATYP Values (different from SOCKS5!) +//! - 0x00: IPv4 Address (4 bytes) +//! - 0x01: IPv6 Address (16 bytes) +//! - 0x02: Domain Name (1 byte length + domain) + +mod uot_v1_stream; + +pub use uot_v1_stream::UotV1Stream; + +/// UoT V2 connect mode stream - identical format to VlessMessageStream (length-prefixed u16be + data) +pub type UotV2Stream = crate::vless::VlessMessageStream; + +/// Magic address used to signal UoT V1 mode (multi-destination) +pub const UOT_V1_MAGIC_ADDRESS: &str = "sp.udp-over-tcp.arpa"; + +/// Magic address used to signal UoT V2 mode (optional connect mode) +pub const UOT_V2_MAGIC_ADDRESS: &str = "sp.v2.udp-over-tcp.arpa"; diff --git a/src/uot/uot_v1_stream.rs b/src/uot/uot_v1_stream.rs new file mode 100644 index 00000000..3eeb5469 --- /dev/null +++ b/src/uot/uot_v1_stream.rs @@ -0,0 +1,400 @@ +//! UDP-over-TCP V1 stream implementation +//! +//! V1 Packet format (each packet has full address): +//! ```text +//! | ATYP | address | port | length | data | +//! | u8 | variable | u16be | u16be | variable | +//! ``` + +use futures::ready; +use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr}; +use std::pin::Pin; +use std::task::{Context, Poll}; +use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; + +use crate::address::{Address, NetLocation}; +use crate::async_stream::{ + AsyncFlushMessage, AsyncPing, AsyncReadTargetedMessage, AsyncShutdownMessage, AsyncStream, + AsyncTargetedMessageStream, AsyncWriteSourcedMessage, +}; +use crate::slide_buffer::SlideBuffer; + +/// ATYP values for UoT (different from SOCKS5!) +const ATYP_IPV4: u8 = 0x00; +const ATYP_IPV6: u8 = 0x01; +const ATYP_DOMAIN: u8 = 0x02; + +/// Buffer size for reading/writing UoT packets +const BUFFER_SIZE: usize = 65535; + +/// UoT V1 stream for multi-destination UDP packets +/// +/// Each packet includes the full destination address, making it suitable +/// for applications that send UDP packets to multiple destinations. +/// +/// This stream wraps an `AsyncStream` (raw byte stream) and handles +/// UoT packet framing internally. +pub struct UotV1Stream { + stream: Box, + + /// Buffer for reading - accumulates bytes until we have a complete packet + read_buf: SlideBuffer, + + /// Buffer for writing - holds a complete packet before sending + write_buf: [u8; BUFFER_SIZE], + write_buf_len: usize, + write_buf_sent: usize, + + is_eof: bool, +} + +impl UotV1Stream { + pub fn new(stream: Box) -> Self { + Self { + stream, + read_buf: SlideBuffer::new(BUFFER_SIZE), + write_buf: [0u8; BUFFER_SIZE], + write_buf_len: 0, + write_buf_sent: 0, + is_eof: false, + } + } + + /// Feed initial data that was read before the stream was created. + /// This is needed when the first UoT packet arrives in the same TCP segment + /// as the magic address header. + pub fn feed_initial_data(&mut self, data: &[u8]) { + if !data.is_empty() { + let len = data.len().min(self.read_buf.remaining_capacity()); + self.read_buf.extend_from_slice(&data[..len]); + } + } + + /// Try to parse a complete UoT packet from the read buffer. + /// Returns Some((location, payload_start, payload_len)) if a complete packet is available. + fn try_parse_packet(&self) -> Option<(NetLocation, usize, usize)> { + let data = self.read_buf.as_slice(); + if data.is_empty() { + return None; + } + + // Parse address to determine header length + let (location, addr_len) = parse_uot_address(data).ok()?; + + // Need at least addr_len + 2 (length prefix) bytes + if data.len() < addr_len + 2 { + return None; + } + + // Read length prefix + let payload_len = u16::from_be_bytes([data[addr_len], data[addr_len + 1]]) as usize; + let payload_start = addr_len + 2; + let total_len = payload_start + payload_len; + + // Check if we have the complete packet + if data.len() < total_len { + return None; + } + + Some((location, payload_start, payload_len)) + } +} + +/// Parse UoT address format (ATYP + address + port) +/// Returns (NetLocation, bytes consumed) +fn parse_uot_address(data: &[u8]) -> std::io::Result<(NetLocation, usize)> { + if data.is_empty() { + return Err(std::io::Error::other("empty UoT address")); + } + + let atyp = data[0]; + match atyp { + ATYP_IPV4 => { + // ATYP(1) + IPv4(4) + Port(2) = 7 bytes + if data.len() < 7 { + return Err(std::io::Error::other("UoT IPv4 address truncated")); + } + let ip = Ipv4Addr::new(data[1], data[2], data[3], data[4]); + let port = u16::from_be_bytes([data[5], data[6]]); + Ok((NetLocation::new(Address::Ipv4(ip), port), 7)) + } + ATYP_IPV6 => { + // ATYP(1) + IPv6(16) + Port(2) = 19 bytes + if data.len() < 19 { + return Err(std::io::Error::other("UoT IPv6 address truncated")); + } + let ip_bytes: [u8; 16] = data[1..17].try_into().unwrap(); + let ip = Ipv6Addr::from(ip_bytes); + let port = u16::from_be_bytes([data[17], data[18]]); + Ok((NetLocation::new(Address::Ipv6(ip), port), 19)) + } + ATYP_DOMAIN => { + // ATYP(1) + DomainLen(1) + Domain(variable) + Port(2) + if data.len() < 2 { + return Err(std::io::Error::other("UoT domain length truncated")); + } + let domain_len = data[1] as usize; + let total_len = 1 + 1 + domain_len + 2; // ATYP + len + domain + port + if data.len() < total_len { + return Err(std::io::Error::other("UoT domain address truncated")); + } + let domain = std::str::from_utf8(&data[2..2 + domain_len]) + .map_err(|e| std::io::Error::other(format!("invalid domain: {e}")))?; + let port = u16::from_be_bytes([data[2 + domain_len], data[3 + domain_len]]); + Ok(( + NetLocation::new(Address::Hostname(domain.to_string()), port), + total_len, + )) + } + _ => Err(std::io::Error::other(format!("unknown UoT ATYP: {atyp}"))), + } +} + +/// Write UoT address format (ATYP + address + port) +/// Returns number of bytes written +fn write_uot_address(buf: &mut [u8], addr: &SocketAddr) -> usize { + match addr { + SocketAddr::V4(v4) => { + buf[0] = ATYP_IPV4; + buf[1..5].copy_from_slice(&v4.ip().octets()); + buf[5..7].copy_from_slice(&v4.port().to_be_bytes()); + 7 + } + SocketAddr::V6(v6) => { + buf[0] = ATYP_IPV6; + buf[1..17].copy_from_slice(&v6.ip().octets()); + buf[17..19].copy_from_slice(&v6.port().to_be_bytes()); + 19 + } + } +} + +impl AsyncReadTargetedMessage for UotV1Stream { + fn poll_read_targeted_message( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + let this = self.get_mut(); + + if this.is_eof { + return Poll::Ready(Ok(NetLocation::UNSPECIFIED)); + } + + loop { + // Try to parse a complete packet from buffered data + if let Some((location, payload_start, payload_len)) = this.try_parse_packet() { + let data = this.read_buf.as_slice(); + buf.put_slice(&data[payload_start..payload_start + payload_len]); + + // Consume the entire packet from the buffer + let total_consumed = payload_start + payload_len; + this.read_buf.consume(total_consumed); + + return Poll::Ready(Ok(location)); + } + + // Need more data - compact buffer if needed + this.read_buf.maybe_compact(4096); + + if this.read_buf.remaining_capacity() == 0 { + return Poll::Ready(Err(std::io::Error::other( + "UoT read buffer full but no complete packet", + ))); + } + + // Read more data from the underlying stream + let write_slice = this.read_buf.write_slice(); + let mut read_buf = ReadBuf::new(write_slice); + + match Pin::new(&mut this.stream).poll_read(cx, &mut read_buf) { + Poll::Ready(Ok(())) => { + let bytes_read = read_buf.filled().len(); + if bytes_read == 0 { + this.is_eof = true; + return Poll::Ready(Ok(NetLocation::UNSPECIFIED)); + } + this.read_buf.advance_write(bytes_read); + // Loop to try parsing again + } + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + Poll::Pending => return Poll::Pending, + } + } + } +} + +impl AsyncWriteSourcedMessage for UotV1Stream { + fn poll_write_sourced_message( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + source: &SocketAddr, + ) -> Poll> { + let this = self.get_mut(); + + // If there's pending data to send, flush it first + while this.write_buf_sent < this.write_buf_len { + let remaining = &this.write_buf[this.write_buf_sent..this.write_buf_len]; + match Pin::new(&mut this.stream).poll_write(cx, remaining) { + Poll::Ready(Ok(n)) => { + this.write_buf_sent += n; + } + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + Poll::Pending => return Poll::Pending, + } + } + + // Reset write buffer + this.write_buf_len = 0; + this.write_buf_sent = 0; + + // Calculate required space: address + length(2) + payload + let addr_len = match source { + SocketAddr::V4(_) => 7, // ATYP + IPv4(4) + Port(2) + SocketAddr::V6(_) => 19, // ATYP + IPv6(16) + Port(2) + }; + let total_len = addr_len + 2 + buf.len(); + + if total_len > this.write_buf.len() { + return Poll::Ready(Err(std::io::Error::other(format!( + "UoT packet too large: {total_len} > {}", + this.write_buf.len() + )))); + } + + // Write UoT address format + let offset = write_uot_address(&mut this.write_buf, source); + + // Write length prefix (u16be) + let len_bytes = (buf.len() as u16).to_be_bytes(); + this.write_buf[offset..offset + 2].copy_from_slice(&len_bytes); + let data_start = offset + 2; + + // Write payload + this.write_buf[data_start..data_start + buf.len()].copy_from_slice(buf); + this.write_buf_len = data_start + buf.len(); + + Poll::Ready(Ok(())) + } +} + +impl AsyncFlushMessage for UotV1Stream { + fn poll_flush_message(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + + // Flush any pending write data + while this.write_buf_sent < this.write_buf_len { + let remaining = &this.write_buf[this.write_buf_sent..this.write_buf_len]; + match Pin::new(&mut this.stream).poll_write(cx, remaining) { + Poll::Ready(Ok(n)) => { + this.write_buf_sent += n; + } + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + Poll::Pending => return Poll::Pending, + } + } + + this.write_buf_len = 0; + this.write_buf_sent = 0; + + // Flush the underlying stream + Pin::new(&mut this.stream).poll_flush(cx) + } +} + +impl AsyncShutdownMessage for UotV1Stream { + fn poll_shutdown_message( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + let mut this = self.get_mut(); + ready!(Pin::new(&mut this).poll_flush_message(cx))?; + Pin::new(&mut this.stream).poll_shutdown(cx) + } +} + +impl AsyncPing for UotV1Stream { + fn supports_ping(&self) -> bool { + false + } + + fn poll_write_ping(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(false)) + } +} + +impl AsyncTargetedMessageStream for UotV1Stream {} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_uot_ipv4_address() { + // ATYP=0x00, IP=192.168.1.1, Port=8080 + let data = [0x00, 192, 168, 1, 1, 0x1F, 0x90]; + let (location, len) = parse_uot_address(&data).unwrap(); + assert_eq!(len, 7); + assert_eq!(location.port(), 8080); + match location.address() { + Address::Ipv4(ip) => assert_eq!(*ip, Ipv4Addr::new(192, 168, 1, 1)), + _ => panic!("expected IPv4"), + } + } + + #[test] + fn test_parse_uot_ipv6_address() { + // ATYP=0x01, IP=::1, Port=443 + let mut data = vec![ATYP_IPV6]; + data.extend_from_slice(&Ipv6Addr::LOCALHOST.octets()); + data.extend_from_slice(&443u16.to_be_bytes()); + + let (location, len) = parse_uot_address(&data).unwrap(); + assert_eq!(len, 19); + assert_eq!(location.port(), 443); + match location.address() { + Address::Ipv6(ip) => assert_eq!(*ip, Ipv6Addr::LOCALHOST), + _ => panic!("expected IPv6"), + } + } + + #[test] + fn test_parse_uot_domain_address() { + // ATYP=0x02, Domain="example.com", Port=53 + let domain = b"example.com"; + let mut data = vec![ATYP_DOMAIN, domain.len() as u8]; + data.extend_from_slice(domain); + data.extend_from_slice(&53u16.to_be_bytes()); + + let (location, len) = parse_uot_address(&data).unwrap(); + assert_eq!(len, 1 + 1 + domain.len() + 2); + assert_eq!(location.port(), 53); + match location.address() { + Address::Hostname(h) => assert_eq!(h, "example.com"), + _ => panic!("expected hostname"), + } + } + + #[test] + fn test_write_uot_ipv4_address() { + let addr: SocketAddr = "192.168.1.1:8080".parse().unwrap(); + let mut buf = [0u8; 32]; + let len = write_uot_address(&mut buf, &addr); + assert_eq!(len, 7); + assert_eq!(buf[0], ATYP_IPV4); + assert_eq!(&buf[1..5], &[192, 168, 1, 1]); + assert_eq!(&buf[5..7], &0x1F90u16.to_be_bytes()); + } + + #[test] + fn test_write_uot_ipv6_address() { + let addr: SocketAddr = "[::1]:443".parse().unwrap(); + let mut buf = [0u8; 32]; + let len = write_uot_address(&mut buf, &addr); + assert_eq!(len, 19); + assert_eq!(buf[0], ATYP_IPV6); + assert_eq!(&buf[1..17], &Ipv6Addr::LOCALHOST.octets()); + assert_eq!(&buf[17..19], &443u16.to_be_bytes()); + } +} diff --git a/src/vless/mod.rs b/src/vless/mod.rs new file mode 100644 index 00000000..92233ffb --- /dev/null +++ b/src/vless/mod.rs @@ -0,0 +1,20 @@ +// VLESS protocol implementation with VISION support + +// Public API +pub mod vless_client_handler; +pub mod vless_server_handler; + +// Internal implementation details +pub mod tls_deframer; +mod tls_fuzzy_deframer; +mod tls_handshake_util; +mod vision_filter; +mod vision_pad; +mod vision_stream; +mod vision_unpad; +mod vless_message_stream; +mod vless_response_stream; +mod vless_util; + +// Re-export VlessMessageStream for use by other protocols (e.g., Shadowsocks UoT V2) +pub use vless_message_stream::VlessMessageStream; diff --git a/src/vless/tls_deframer.rs b/src/vless/tls_deframer.rs new file mode 100644 index 00000000..10a26fb1 --- /dev/null +++ b/src/vless/tls_deframer.rs @@ -0,0 +1,599 @@ +use bytes::{Bytes, BytesMut}; +use std::io; + +/// TLS record header size (ContentType + ProtocolVersion + Length) +const TLS_RECORD_HEADER_SIZE: usize = 5; + +/// Maximum TLS ciphertext payload size per RFC specifications +/// +/// The TLS record header's `length` field specifies the size of the ENCRYPTED payload (ciphertext). +/// Per RFC 8446 (TLS 1.3) and RFC 5246 (TLS 1.2), the limits are: +/// +/// - TLS 1.3: Plaintext limit = 16,384 bytes (2^14) +/// Encryption overhead allowance = 256 bytes +/// Ciphertext limit = 16,384 + 256 = 16,640 bytes +/// +/// - TLS 1.2: Plaintext limit = 16,384 bytes (2^14) +/// Encryption overhead allowance = 2,048 bytes +/// Ciphertext limit = 16,384 + 2,048 = 18,432 bytes +/// +/// We use the TLS 1.2 limit for maximum compatibility with both protocol versions. +/// See /shoes/VISION_TLS_RECORD_SIZE.md for detailed analysis and RFC references. +const MAX_TLS_CIPHERTEXT_LEN: usize = 16384 + 2048; // 18,432 bytes (TLS 1.2 limit) + +/// Maximum TLS record size (ciphertext + header) +pub const TLS_MAX_RECORD_SIZE: usize = MAX_TLS_CIPHERTEXT_LEN + TLS_RECORD_HEADER_SIZE; + +/// TLS protocol versions we expect (0x0303 = TLS 1.2 for compatibility) +const TLS_PROTOCOL_VERSION_MAJOR: u8 = 0x03; +const TLS_PROTOCOL_VERSION_MINOR_MIN: u8 = 0x01; +const TLS_PROTOCOL_VERSION_MINOR_MAX: u8 = 0x03; + +/// Deframer that reassembles TLS records from partial reads/writes +/// +/// ```ignore +/// let mut deframer = TlsDeframer::new(); +/// +/// // Feed data as it arrives (may be partial) +/// deframer.feed(&chunk1); +/// deframer.feed(&chunk2); +/// +/// // Extract complete records +/// while let Some(record) = deframer.next_record()? { +/// process_tls_record(record); +/// } +/// ``` +#[derive(Debug, PartialEq)] +pub struct TlsDeframer { + /// Buffer holding partial TLS record data + buffer: BytesMut, + + /// Current parsing state + state: DeframerState, +} + +#[derive(Debug, Clone, Copy, PartialEq)] +enum DeframerState { + /// Waiting for or reading the 5-byte TLS record header + ReadingHeader, + + /// Reading payload of known length + /// Stores the expected payload length + ReadingPayload { payload_len: usize }, +} + +impl TlsDeframer { + pub fn new() -> Self { + Self { + buffer: BytesMut::with_capacity(TLS_MAX_RECORD_SIZE), + state: DeframerState::ReadingHeader, + } + } + + /// Feed data into the deframer + /// + /// This data may contain partial TLS records, complete records, or multiple records. + /// Call `next_record()` after feeding to extract complete records. + pub fn feed(&mut self, data: &[u8]) { + self.buffer.extend_from_slice(data); + } + + /// Extract the next complete TLS record if available + /// + /// This method automatically advances the internal buffer after extracting the record. + /// + /// Returns: + /// - `Ok(Some(Bytes))` if a complete record is available (zero-copy slice) + /// - `Ok(None)` if more data is needed + /// - `Err(...)` if the data is malformed + pub fn next_record(&mut self) -> io::Result> { + loop { + match self.state { + DeframerState::ReadingHeader => { + // Need at least 5 bytes for header + if self.buffer.len() < TLS_RECORD_HEADER_SIZE { + return Ok(None); + } + + // Parse TLS record header + let content_type = self.buffer[0]; + let version_major = self.buffer[1]; + let version_minor = self.buffer[2]; + let payload_len = u16::from_be_bytes([self.buffer[3], self.buffer[4]]) as usize; + + // Validate protocol version (should be 0x03XX for TLS 1.x) + // Accept TLS 1.0 (0x0301), TLS 1.1 (0x0302), TLS 1.2 (0x0303), and TLS 1.3 (0x0303) + // TLS 1.3 uses 0x0303 in the record layer for backwards compatibility + if version_major != TLS_PROTOCOL_VERSION_MAJOR + || !(TLS_PROTOCOL_VERSION_MINOR_MIN..=TLS_PROTOCOL_VERSION_MINOR_MAX) + .contains(&version_minor) + { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "Invalid TLS protocol version: 0x{:02x}{:02x} (expected 0x0301 to 0x0303)", + version_major, version_minor + ), + )); + } + + // Validate ciphertext payload length + // The `payload_len` from the TLS record header specifies encrypted data size. + // TLS 1.2 allows up to 18,432 bytes (16,384 plaintext + 2,048 overhead). + if payload_len > MAX_TLS_CIPHERTEXT_LEN { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "TLS record ciphertext length {} exceeds maximum {} (TLS 1.2 limit)", + payload_len, MAX_TLS_CIPHERTEXT_LEN + ), + )); + } + + // Validate content type (basic sanity check) + // Valid types: 0x14-0x18 (ChangeCipherSpec, Alert, Handshake, Application Data, Heartbeat) + if !(0x14..=0x18).contains(&content_type) { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Invalid TLS content type: 0x{:02x}", content_type), + )); + } + + // Transition to reading payload + self.state = DeframerState::ReadingPayload { payload_len }; + } + + DeframerState::ReadingPayload { payload_len } => { + let total_len = TLS_RECORD_HEADER_SIZE + payload_len; + + // Check if we have the complete record + if self.buffer.len() < total_len { + return Ok(None); + } + + // We have a complete record! Extract it using split_to for zero-copy. + // BytesMut::split_to() returns a Bytes object that shares the underlying + // buffer without copying, and advances the BytesMut past the split point. + let record = self.buffer.split_to(total_len).freeze(); + + // Reset state for next record + self.state = DeframerState::ReadingHeader; + + return Ok(Some(record)); + } + } + } + } + + /// Get the number of bytes currently buffered + pub fn pending_bytes(&self) -> usize { + self.buffer.len() + } + + /// Extract all available complete TLS records + /// + /// This is equivalent to calling `next_record()` repeatedly until it returns `Ok(None)`, + /// but more convenient when you want to process all available records at once. + /// + /// Returns a vector of all complete records that can be extracted from the current buffer. + /// If no complete records are available, returns an empty vector. + /// + /// Returns: + /// - `Ok(Vec)` containing all available complete records (may be empty) + /// - `Err(...)` if any record is malformed + pub fn next_records(&mut self) -> io::Result> { + let mut records = Vec::new(); + + while let Some(record) = self.next_record()? { + records.push(record); + } + + Ok(records) + } + + /// Consume the deframer and return all remaining buffered data + pub fn into_remaining_data(self) -> Bytes { + self.buffer.freeze() + } + + pub fn remaining_data(&self) -> &[u8] { + &self.buffer + } + + /// Clear all buffered data + pub fn clear(&mut self) { + self.buffer.clear(); + self.state = DeframerState::ReadingHeader; + } + + pub fn deallocate(&mut self) { + drop(std::mem::take(&mut self.buffer)); + } +} + +impl Default for TlsDeframer { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Helper to create a TLS record with given content type and payload + fn make_tls_record(content_type: u8, payload: &[u8]) -> Vec { + let mut record = Vec::new(); + record.push(content_type); + record.push(0x03); // Version major + record.push(0x03); // Version minor + let len = payload.len() as u16; + record.extend_from_slice(&len.to_be_bytes()); + record.extend_from_slice(payload); + record + } + + #[test] + fn test_single_complete_packet() { + let mut deframer = TlsDeframer::new(); + + let payload = b"Hello, TLS!"; + let record = make_tls_record(0x17, payload); // Application Data + + deframer.feed(&record); + + let read_record = deframer.next_record().unwrap(); + assert!(read_record.is_some()); + let read_record = read_record.unwrap(); + assert_eq!(read_record, &record[..]); + + // Verify no more packets + assert!(deframer.next_record().unwrap().is_none()); + } + + #[test] + fn test_partial_packet_split_header() { + let mut deframer = TlsDeframer::new(); + + let payload = b"Partial header test"; + let record = make_tls_record(0x16, payload); // Handshake + + // Feed first 3 bytes (partial header) + deframer.feed(&record[..3]); + assert!(deframer.next_record().unwrap().is_none()); + + // Feed rest of header + deframer.feed(&record[3..5]); + assert!(deframer.next_record().unwrap().is_none()); + + // Feed payload + deframer.feed(&record[5..]); + let extracted_record = deframer.next_record().unwrap(); + assert!(extracted_record.is_some()); + assert_eq!(extracted_record.unwrap(), &record[..]); + } + + #[test] + fn test_partial_packet_split_payload() { + let mut deframer = TlsDeframer::new(); + + let payload = vec![0x42; 100]; // 100 bytes of 0x42 + let record = make_tls_record(0x17, &payload); + + // Feed header + half payload + let split_point = 5 + 50; + deframer.feed(&record[..split_point]); + assert!(deframer.next_record().unwrap().is_none()); + + // Feed rest of payload + deframer.feed(&record[split_point..]); + let extracted_record = deframer.next_record().unwrap(); + assert!(extracted_record.is_some()); + assert_eq!(extracted_record.unwrap(), &record[..]); + } + + #[test] + fn test_multiple_packets_at_once() { + let mut deframer = TlsDeframer::new(); + + let record1 = make_tls_record(0x16, b"First"); + let record2 = make_tls_record(0x17, b"Second"); + let record3 = make_tls_record(0x16, b"Third"); + + // Feed all three at once + let mut combined = Vec::new(); + combined.extend_from_slice(&record1); + combined.extend_from_slice(&record2); + combined.extend_from_slice(&record3); + + deframer.feed(&combined); + + // Extract all three (automatically consumed) + let rec1 = deframer.next_record().unwrap().unwrap(); + assert_eq!(rec1, &record1[..]); + + let rec2 = deframer.next_record().unwrap().unwrap(); + assert_eq!(rec2, &record2[..]); + + let rec3 = deframer.next_record().unwrap().unwrap(); + assert_eq!(rec3, &record3[..]); + + assert!(deframer.next_record().unwrap().is_none()); + } + + #[test] + fn test_multiple_packets_partial_boundary() { + let mut deframer = TlsDeframer::new(); + + let record1 = make_tls_record(0x16, b"Record One"); + let record2 = make_tls_record(0x17, b"Record Two"); + + let mut combined = Vec::new(); + combined.extend_from_slice(&record1); + combined.extend_from_slice(&record2); + + // Feed up to middle of second record + let split = record1.len() + 8; + deframer.feed(&combined[..split]); + + // Should get first record (automatically consumed) + let rec1 = deframer.next_record().unwrap().unwrap(); + assert_eq!(rec1, &record1[..]); + + // Second record not complete yet + assert!(deframer.next_record().unwrap().is_none()); + + // Feed rest + deframer.feed(&combined[split..]); + let rec2 = deframer.next_record().unwrap().unwrap(); + assert_eq!(rec2, &record2[..]); + } + + #[test] + fn test_tls10_accepted() { + let mut deframer = TlsDeframer::new(); + + let mut record = vec![0x16, 0x03, 0x01, 0x00, 0x05]; // TLS 1.0 + record.extend_from_slice(b"Hello"); + + deframer.feed(&record); + + let result = deframer.next_record(); + assert!(result.is_ok()); + assert_eq!(result.unwrap().unwrap(), &record[..]); + } + + #[test] + fn test_invalid_protocol_version() { + let mut deframer = TlsDeframer::new(); + + let mut record = vec![0x16, 0x02, 0x00, 0x00, 0x05]; // SSL 2.0 (invalid) + record.extend_from_slice(b"Hello"); + + deframer.feed(&record); + + let result = deframer.next_record(); + assert!(result.is_err()); + assert!(result + .unwrap_err() + .to_string() + .contains("Invalid TLS protocol version")); + } + + #[test] + fn test_invalid_content_type() { + let mut deframer = TlsDeframer::new(); + + let mut record = vec![0xFF, 0x03, 0x03, 0x00, 0x05]; // Invalid content type + record.extend_from_slice(b"Hello"); + + deframer.feed(&record); + + let result = deframer.next_record(); + assert!(result.is_err()); + assert!(result + .unwrap_err() + .to_string() + .contains("Invalid TLS content type")); + } + + #[test] + fn test_excessive_payload_length() { + let mut deframer = TlsDeframer::new(); + + // Payload length > 18432 (TLS 1.2 max ciphertext) + let record = vec![0x17, 0x03, 0x03, 0xFF, 0xFF]; // 65535 bytes + + deframer.feed(&record); + + let result = deframer.next_record(); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("exceeds maximum")); + } + + #[test] + fn test_byte_by_byte_feed() { + let mut deframer = TlsDeframer::new(); + + let record = make_tls_record(0x16, b"Byte by byte"); + + // Feed one byte at a time + for byte in &record[..record.len() - 1] { + deframer.feed(&[*byte]); + assert!(deframer.next_record().unwrap().is_none()); + } + + // Feed last byte + deframer.feed(&[record[record.len() - 1]]); + let extracted_record = deframer.next_record().unwrap().unwrap(); + assert_eq!(extracted_record, &record[..]); + } + + #[test] + fn test_clear() { + let mut deframer = TlsDeframer::new(); + + let record = make_tls_record(0x16, b"Test"); + deframer.feed(&record[..3]); // Partial + + assert_eq!(deframer.pending_bytes(), 3); + + deframer.clear(); + + assert_eq!(deframer.pending_bytes(), 0); + assert!(deframer.next_record().unwrap().is_none()); + } + + #[test] + fn test_empty_payload() { + let mut deframer = TlsDeframer::new(); + + let record = make_tls_record(0x16, b""); // Empty payload + + deframer.feed(&record); + + let extracted_record = deframer.next_record().unwrap().unwrap(); + assert_eq!(extracted_record.len(), 5); // Just the header + assert_eq!(extracted_record, &record[..]); + } + + #[test] + fn test_max_size_payload() { + let mut deframer = TlsDeframer::new(); + + // Test with maximum TLS 1.2 ciphertext size (18,432 bytes) + let payload = vec![0xAA; 18432]; + let record = make_tls_record(0x17, &payload); + + deframer.feed(&record); + + let extracted_record = deframer.next_record().unwrap().unwrap(); + assert_eq!(extracted_record, &record[..]); + } + + #[test] + fn test_reuse_after_consume() { + let mut deframer = TlsDeframer::new(); + + let record1 = make_tls_record(0x16, b"First"); + let record2 = make_tls_record(0x17, b"Second"); + + // First packet (automatically consumed) + deframer.feed(&record1); + let rec1 = deframer.next_record().unwrap().unwrap(); + assert_eq!(rec1, &record1[..]); + + // Second packet (automatically consumed) + deframer.feed(&record2); + let rec2 = deframer.next_record().unwrap().unwrap(); + assert_eq!(rec2, &record2[..]); + + assert_eq!(deframer.pending_bytes(), 0); + } + + #[test] + fn test_next_records_empty() { + let mut deframer = TlsDeframer::new(); + + let records = deframer.next_records().unwrap(); + assert!(records.is_empty()); + } + + #[test] + fn test_next_records_single() { + let mut deframer = TlsDeframer::new(); + + let record = make_tls_record(0x16, b"Single"); + deframer.feed(&record); + + let records = deframer.next_records().unwrap(); + assert_eq!(records.len(), 1); + assert_eq!(records[0], &record[..]); + } + + #[test] + fn test_next_records_multiple() { + let mut deframer = TlsDeframer::new(); + + let record1 = make_tls_record(0x16, b"First"); + let record2 = make_tls_record(0x17, b"Second"); + let record3 = make_tls_record(0x16, b"Third"); + + let mut combined = Vec::new(); + combined.extend_from_slice(&record1); + combined.extend_from_slice(&record2); + combined.extend_from_slice(&record3); + + deframer.feed(&combined); + + let records = deframer.next_records().unwrap(); + assert_eq!(records.len(), 3); + assert_eq!(records[0], &record1[..]); + assert_eq!(records[1], &record2[..]); + assert_eq!(records[2], &record3[..]); + + assert_eq!(deframer.pending_bytes(), 0); + } + + #[test] + fn test_next_records_partial_remaining() { + let mut deframer = TlsDeframer::new(); + + let record1 = make_tls_record(0x16, b"Complete"); + let record2 = make_tls_record(0x17, b"Incomplete"); + + let mut combined = Vec::new(); + combined.extend_from_slice(&record1); + combined.extend_from_slice(&record2[..10]); // Only partial second record + + deframer.feed(&combined); + + let records = deframer.next_records().unwrap(); + assert_eq!(records.len(), 1); + assert_eq!(records[0], &record1[..]); + + // Should still have partial data + assert_eq!(deframer.pending_bytes(), 10); + } + + #[test] + fn test_next_records_error_propagation() { + let mut deframer = TlsDeframer::new(); + + // Invalid record with bad content type + let mut bad_record = vec![0xFF, 0x03, 0x03, 0x00, 0x05]; + bad_record.extend_from_slice(b"Hello"); + + deframer.feed(&bad_record); + + let result = deframer.next_records(); + assert!(result.is_err()); + assert!(result + .unwrap_err() + .to_string() + .contains("Invalid TLS content type")); + } + + #[test] + fn test_next_records_mixed_complete_and_partial() { + let mut deframer = TlsDeframer::new(); + + let record1 = make_tls_record(0x16, b"First"); + let record2 = make_tls_record(0x17, b"Second"); + + // Feed first complete record and partial second + deframer.feed(&record1); + deframer.feed(&record2[..8]); + + let records = deframer.next_records().unwrap(); + assert_eq!(records.len(), 1); + assert_eq!(records[0], &record1[..]); + + // Complete the second record + deframer.feed(&record2[8..]); + + let records = deframer.next_records().unwrap(); + assert_eq!(records.len(), 1); + assert_eq!(records[0], &record2[..]); + } +} diff --git a/src/vless/tls_fuzzy_deframer.rs b/src/vless/tls_fuzzy_deframer.rs new file mode 100644 index 00000000..79c706a4 --- /dev/null +++ b/src/vless/tls_fuzzy_deframer.rs @@ -0,0 +1,531 @@ +/// Fuzzy TLS deframer that can handle unknown prefixes before TLS data +/// +/// This deframer wraps TlsDeframer and adds the ability to: +/// 1. Search for TLS handshake patterns in buffers that may contain protocol headers +/// 2. Proactively return unknown prefix bytes that cannot be part of TLS patterns +/// 3. Minimize buffering by discarding non-TLS prefixes as soon as possible +/// +/// The deframer operates in two states: +/// - **Searching**: Looking for the TLS handshake pattern [0x16, 0x03, 0x01-0x03, ?, ?, 0x01|0x02] +/// - Returns UnknownPrefix for bytes that cannot be part of the pattern +/// - Keeps minimal suffix that might complete into a pattern (max 5 bytes) +/// - **Deframing**: Found TLS pattern, now extracting complete TLS records +/// - Delegates to inner TlsDeframer for standard TLS record parsing +/// +/// This is used in VLESS Vision protocol when chaining proxies, where the inner TLS data +/// may be prefixed with VLESS protocol headers from intermediate proxies. +use bytes::{Bytes, BytesMut}; +use std::io; + +use super::tls_deframer::TlsDeframer; +use super::tls_handshake_util::{ + find_tls_handshake_start, get_partial_tls_pattern_suffix_len, MIN_TLS_HANDSHAKE_PATTERN_LEN, +}; + +/// Result of attempting to extract data from the fuzzy deframer +#[derive(Debug, PartialEq)] +pub enum DeframeResult { + /// Complete TLS record extracted + TlsRecord(Bytes), + + /// Unknown prefix bytes that were skipped + /// Returned proactively as soon as we know bytes aren't part of TLS pattern + UnknownPrefix(Bytes), + + /// Need more data (waiting for complete record or more pattern bytes) + NeedData, +} + +/// Fuzzy TLS deframer that can skip unknown prefixes before TLS data +pub struct FuzzyTlsDeframer { + /// Underlying TLS deframer (only used in Deframing state) + inner: TlsDeframer, + + /// Search buffer (only used in Searching state) + search_buffer: BytesMut, + + /// Current state + state: FuzzyState, + + /// Total prefix bytes discarded (for logging) + total_prefix_bytes: usize, +} + +#[derive(Debug, PartialEq)] +enum FuzzyState { + /// Searching for TLS handshake pattern + Searching, + + /// Found TLS pattern, now in normal deframing mode + Deframing, +} + +impl FuzzyTlsDeframer { + pub fn new() -> Self { + Self { + inner: TlsDeframer::new(), + search_buffer: BytesMut::new(), + state: FuzzyState::Searching, + total_prefix_bytes: 0, + } + } + + /// Feed data into the deframer + /// + /// Data may contain protocol headers, partial TLS patterns, or complete TLS records. + /// Call `next_record()` after feeding to extract available data. + pub fn feed(&mut self, data: &[u8]) { + match self.state { + FuzzyState::Searching => { + // Just accumulate in search buffer + // Don't optimize here - let next_record() handle it for proper accounting + self.search_buffer.extend_from_slice(data); + } + FuzzyState::Deframing => { + // Normal mode, feed directly to inner deframer + self.inner.feed(data); + } + } + } + + /// Extract the next available data + /// + /// Returns: + /// - `TlsRecord`: A complete TLS record (in Deframing state) + /// - `UnknownPrefix`: Bytes that cannot be part of TLS pattern (in Searching state) + /// - `IncompleteRecord`: Need more data + /// + /// This method should be called in a loop until it returns `IncompleteRecord`. + pub fn next_record(&mut self) -> io::Result { + loop { + match self.state { + FuzzyState::Searching => { + if self.search_buffer.is_empty() { + return Ok(DeframeResult::NeedData); + } + + // Try to find complete TLS pattern (need at least 6 bytes) + if self.search_buffer.len() >= MIN_TLS_HANDSHAKE_PATTERN_LEN { + if let Some(offset) = find_tls_handshake_start(&self.search_buffer) { + // Found complete pattern! + + if offset > 0 { + // There's a prefix before the TLS pattern + // Extract it and return as UnknownPrefix + let prefix = self.search_buffer.split_to(offset).freeze(); + self.total_prefix_bytes += prefix.len(); + + // Feed remaining data (TLS) to inner deframer + self.inner.feed(&self.search_buffer); + self.search_buffer.clear(); + + // Transition to Deframing state + self.state = FuzzyState::Deframing; + + log::debug!( + "FuzzyTlsDeframer: Found TLS pattern after {} byte prefix (total {} bytes discarded), transitioned to Deframing", + prefix.len(), + self.total_prefix_bytes + ); + + return Ok(DeframeResult::UnknownPrefix(prefix)); + } else { + // offset == 0, pattern at start, no prefix + self.inner.feed(&self.search_buffer); + self.search_buffer.clear(); + self.state = FuzzyState::Deframing; + + log::debug!( + "FuzzyTlsDeframer: Found TLS pattern at offset 0 (total {} bytes discarded), transitioned to Deframing", + self.total_prefix_bytes + ); + + // Continue loop to try extracting a record from inner + continue; + } + } else { + // No complete pattern found, but we have enough bytes to check + // Check if we can discard some prefix bytes + let suffix_len = + get_partial_tls_pattern_suffix_len(&self.search_buffer); + + // Only discard if suffix is shorter than buffer + // This means we have definite non-pattern bytes at the start + if suffix_len < self.search_buffer.len() { + let prefix_len = self.search_buffer.len() - suffix_len; + let prefix = self.search_buffer.split_to(prefix_len).freeze(); + self.total_prefix_bytes += prefix.len(); + + log::debug!( + "FuzzyTlsDeframer: Returning {} byte UnknownPrefix (total {} bytes discarded, kept {} suffix bytes)", + prefix.len(), + self.total_prefix_bytes, + suffix_len + ); + + return Ok(DeframeResult::UnknownPrefix(prefix)); + } + } + } + + // Either: + // - Buffer too small (< 6 bytes) - need more data to check for complete pattern + // - Buffer >= 6, no complete pattern, entire buffer might be partial pattern + // Need more data in all cases + return Ok(DeframeResult::NeedData); + } + FuzzyState::Deframing => { + // Already in deframing mode, delegate to inner + match self.inner.next_record()? { + Some(record) => return Ok(DeframeResult::TlsRecord(record)), + None => return Ok(DeframeResult::NeedData), + } + } + } + } + } + + /// Get the number of bytes currently buffered + pub fn pending_bytes(&self) -> usize { + match self.state { + FuzzyState::Searching => self.search_buffer.len(), + FuzzyState::Deframing => self.inner.pending_bytes(), + } + } + + /// Get a reference to remaining buffered data + pub fn remaining_data(&self) -> &[u8] { + match self.state { + FuzzyState::Searching => &self.search_buffer, + FuzzyState::Deframing => self.inner.remaining_data(), + } + } + + /// Clear all buffered data and reset to searching state + pub fn clear(&mut self) { + self.search_buffer.clear(); + self.inner.clear(); + self.state = FuzzyState::Searching; + self.total_prefix_bytes = 0; + } + + /// Deallocate buffers + pub fn deallocate(&mut self) { + self.search_buffer = BytesMut::new(); + self.inner.deallocate(); + } +} + +impl Default for FuzzyTlsDeframer { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Helper to create a TLS record with given content type and payload + /// For Handshake records (0x16), prepends handshake type to payload + fn make_tls_record(content_type: u8, payload: &[u8]) -> Vec { + let mut record = Vec::new(); + record.push(content_type); + record.push(0x03); // Version major + record.push(0x03); // Version minor + + // For Handshake records, prepend handshake type (0x01 = ClientHello) + if content_type == 0x16 { + let total_len = 1 + payload.len(); // handshake type byte + payload + record.extend_from_slice(&(total_len as u16).to_be_bytes()); + record.push(0x01); // ClientHello handshake type + record.extend_from_slice(payload); + } else { + let len = payload.len() as u16; + record.extend_from_slice(&len.to_be_bytes()); + record.extend_from_slice(payload); + } + + record + } + + #[test] + fn test_no_prefix_complete_record() { + let mut deframer = FuzzyTlsDeframer::new(); + + // TLS record with no prefix + let record = make_tls_record(0x16, b"Hello"); + deframer.feed(&record); + + // Should immediately transition to Deframing and return the record + match deframer.next_record().unwrap() { + DeframeResult::TlsRecord(r) => assert_eq!(r, &record[..]), + other => panic!("Expected TlsRecord, got {:?}", other), + } + } + + #[test] + fn test_prefix_then_complete_record() { + let mut deframer = FuzzyTlsDeframer::new(); + + // VLESS-like prefix (26 bytes of zeros) + TLS record + let mut data = vec![0x00; 26]; + let record = make_tls_record(0x16, b"Hello"); + data.extend_from_slice(&record); + + deframer.feed(&data); + + // First call should return the prefix + match deframer.next_record().unwrap() { + DeframeResult::UnknownPrefix(prefix) => assert_eq!(prefix.len(), 26), + other => panic!("Expected UnknownPrefix, got {:?}", other), + } + + // Second call should return the TLS record + match deframer.next_record().unwrap() { + DeframeResult::TlsRecord(r) => assert_eq!(r, &record[..]), + other => panic!("Expected TlsRecord, got {:?}", other), + } + + // Third call should return incomplete + match deframer.next_record().unwrap() { + DeframeResult::NeedData => {} + other => panic!("Expected IncompleteRecord, got {:?}", other), + } + } + + #[test] + fn test_small_prefix_not_returned_immediately() { + let mut deframer = FuzzyTlsDeframer::new(); + + // Small prefix (3 bytes) that doesn't match any partial pattern + deframer.feed(&[0x00, 0x01, 0x02]); + + // Should wait for more data (buffer <= MAX_PARTIAL_SUFFIX_LEN) + match deframer.next_record().unwrap() { + DeframeResult::NeedData => {} + other => panic!("Expected IncompleteRecord, got {:?}", other), + } + + assert_eq!(deframer.pending_bytes(), 3); + } + + #[test] + fn test_large_prefix_returned_proactively() { + let mut deframer = FuzzyTlsDeframer::new(); + + // Large prefix (10 bytes) with no partial pattern at end + deframer.feed(&[0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09]); + + // Should return prefix, keeping nothing (suffix_len = 0) + match deframer.next_record().unwrap() { + DeframeResult::UnknownPrefix(prefix) => assert_eq!(prefix.len(), 10), + other => panic!("Expected UnknownPrefix, got {:?}", other), + } + + assert_eq!(deframer.pending_bytes(), 0); + } + + #[test] + fn test_prefix_with_partial_pattern_suffix() { + let mut deframer = FuzzyTlsDeframer::new(); + + // 10 bytes of junk + partial TLS pattern [0x16, 0x03] + deframer.feed(&[ + 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x16, 0x03, + ]); + + // Should return 10 bytes as prefix, keep [0x16, 0x03] as potential pattern + match deframer.next_record().unwrap() { + DeframeResult::UnknownPrefix(prefix) => assert_eq!(prefix.len(), 10), + other => panic!("Expected UnknownPrefix, got {:?}", other), + } + + assert_eq!(deframer.pending_bytes(), 2); + + // Now feed the rest of the pattern + deframer.feed(&[0x01, 0x00, 0x05, 0x01, b'H', b'e', b'l', b'l', b'o']); + + // Should get the complete TLS record + match deframer.next_record().unwrap() { + DeframeResult::TlsRecord(r) => { + assert_eq!(r[0], 0x16); // Handshake + assert_eq!(r.len(), 5 + 5); // header + payload + } + other => panic!("Expected TlsRecord, got {:?}", other), + } + } + + #[test] + fn test_partial_pattern_alone_not_discarded() { + let mut deframer = FuzzyTlsDeframer::new(); + + // Just a partial pattern [0x16, 0x03] + deframer.feed(&[0x16, 0x03]); + + // Should wait for more data, not discard + match deframer.next_record().unwrap() { + DeframeResult::NeedData => {} + other => panic!("Expected IncompleteRecord, got {:?}", other), + } + + assert_eq!(deframer.pending_bytes(), 2); + } + + #[test] + fn test_multiple_records_after_prefix() { + let mut deframer = FuzzyTlsDeframer::new(); + + // Prefix + two TLS records + let mut data = vec![0x00; 26]; + let record1 = make_tls_record(0x16, b"First"); + let record2 = make_tls_record(0x17, b"Second"); + data.extend_from_slice(&record1); + data.extend_from_slice(&record2); + + deframer.feed(&data); + + // Get prefix + match deframer.next_record().unwrap() { + DeframeResult::UnknownPrefix(prefix) => assert_eq!(prefix.len(), 26), + other => panic!("Expected UnknownPrefix, got {:?}", other), + } + + // Get first record + match deframer.next_record().unwrap() { + DeframeResult::TlsRecord(r) => assert_eq!(r, &record1[..]), + other => panic!("Expected TlsRecord, got {:?}", other), + } + + // Get second record + match deframer.next_record().unwrap() { + DeframeResult::TlsRecord(r) => assert_eq!(r, &record2[..]), + other => panic!("Expected TlsRecord, got {:?}", other), + } + + // No more data + match deframer.next_record().unwrap() { + DeframeResult::NeedData => {} + other => panic!("Expected IncompleteRecord, got {:?}", other), + } + } + + #[test] + fn test_split_feed_across_pattern() { + let mut deframer = FuzzyTlsDeframer::new(); + + // Feed prefix + partial pattern in first call + deframer.feed(&[0x00, 0x01, 0x02, 0x16, 0x03]); + + // Buffer too small, wait for more + match deframer.next_record().unwrap() { + DeframeResult::NeedData => {} + other => panic!("Expected IncompleteRecord, got {:?}", other), + } + + // Feed rest of pattern + record + deframer.feed(&[0x01, 0x00, 0x05, 0x01, b'H', b'e', b'l', b'l', b'o']); + + // Now should get prefix + match deframer.next_record().unwrap() { + DeframeResult::UnknownPrefix(prefix) => assert_eq!(prefix.len(), 3), + other => panic!("Expected UnknownPrefix, got {:?}", other), + } + + // Then the record + match deframer.next_record().unwrap() { + DeframeResult::TlsRecord(_) => {} + other => panic!("Expected TlsRecord, got {:?}", other), + } + } + + #[test] + fn test_pending_bytes_accounting() { + let mut deframer = FuzzyTlsDeframer::new(); + + assert_eq!(deframer.pending_bytes(), 0); + + // Feed some data + deframer.feed(&[0x00, 0x01, 0x02]); + assert_eq!(deframer.pending_bytes(), 3); + + // Feed more to exceed threshold + deframer.feed(&[0x03, 0x04, 0x05, 0x06, 0x07, 0x08]); + assert_eq!(deframer.pending_bytes(), 9); + + // Extract prefix + match deframer.next_record().unwrap() { + DeframeResult::UnknownPrefix(prefix) => assert_eq!(prefix.len(), 9), + _ => panic!("Expected UnknownPrefix"), + } + + assert_eq!(deframer.pending_bytes(), 0); + } + + #[test] + fn test_pending_bytes_during_state_transition() { + let mut deframer = FuzzyTlsDeframer::new(); + + // Feed prefix + complete TLS record + let mut data = vec![0x00; 26]; // VLESS prefix + let record = make_tls_record(0x16, b"Test"); + data.extend_from_slice(&record); + + deframer.feed(&data); + + // Before extracting anything, pending_bytes should be full buffer + assert_eq!(deframer.pending_bytes(), 26 + record.len()); + + // Extract prefix (state transitions to Deframing, TLS data fed to inner) + match deframer.next_record().unwrap() { + DeframeResult::UnknownPrefix(prefix) => { + assert_eq!(prefix.len(), 26); + } + other => panic!("Expected UnknownPrefix, got {:?}", other), + } + + // After extracting prefix, pending_bytes should show TLS record length + // (TLS data is now in inner deframer) + assert_eq!(deframer.pending_bytes(), record.len()); + + // Extract the TLS record + match deframer.next_record().unwrap() { + DeframeResult::TlsRecord(r) => assert_eq!(r.len(), record.len()), + other => panic!("Expected TlsRecord, got {:?}", other), + } + + // After extracting record, buffer should be empty + assert_eq!(deframer.pending_bytes(), 0); + } + + #[test] + fn test_clear() { + let mut deframer = FuzzyTlsDeframer::new(); + + deframer.feed(&[0x00, 0x01, 0x02]); + assert_eq!(deframer.pending_bytes(), 3); + + deframer.clear(); + assert_eq!(deframer.pending_bytes(), 0); + assert_eq!(deframer.state, FuzzyState::Searching); + } + + #[test] + fn test_false_tls_pattern_continues_search() { + let mut deframer = FuzzyTlsDeframer::new(); + + // Data with [0x16, 0x03] but not a valid TLS pattern (wrong minor version) + // Pattern needs: [0x16, 0x03, 0x01-0x03, ?, ?, 0x01|0x02] + let data = vec![0x00, 0x01, 0x16, 0x03, 0x04, 0x00, 0x00, 0x01]; // 0x04 is invalid minor version + + deframer.feed(&data); + + // Should return as prefix since pattern search fails + match deframer.next_record().unwrap() { + DeframeResult::UnknownPrefix(prefix) => assert_eq!(prefix.len(), 8), + other => panic!("Expected UnknownPrefix, got {:?}", other), + } + + assert_eq!(deframer.state, FuzzyState::Searching); + } +} diff --git a/src/vless/tls_handshake_util.rs b/src/vless/tls_handshake_util.rs new file mode 100644 index 00000000..f1225cb6 --- /dev/null +++ b/src/vless/tls_handshake_util.rs @@ -0,0 +1,241 @@ +/// Utilities for detecting TLS handshake patterns in Vision protocol +/// +/// These functions search for TLS handshake patterns to identify where the actual +/// TLS data begins in buffers that may contain protocol headers (e.g., VLESS headers). +/// Minimum bytes needed for a complete TLS handshake pattern +/// Pattern: [0x16, 0x03, 0x01-0x03, ?, ?, 0x01|0x02] = 6 bytes +pub const MIN_TLS_HANDSHAKE_PATTERN_LEN: usize = 6; + +/// Search for TLS handshake pattern using SIMD-optimized memchr +/// +/// Returns Some(offset) if pattern found, None otherwise +/// +/// Pattern: [0x16, 0x03, 0x01-0x03, ?, ?, 0x01|0x02] +/// - Byte 0: 0x16 (Handshake record type) +/// - Byte 1: 0x03 (TLS version major) +/// - Byte 2: 0x01-0x03 (TLS version minor: 1.0, 1.1, 1.2, or 1.3 compat) +/// - Bytes 3-4: [any] (Record length) +/// - Byte 5: 0x01 or 0x02 (ClientHello or ServerHello) +use memchr::memmem; + +pub fn find_tls_handshake_start(data: &[u8]) -> Option { + if data.len() < MIN_TLS_HANDSHAKE_PATTERN_LEN { + return None; + } + + // Use SIMD-optimized search for the 2-byte prefix [0x16, 0x03] + // This is much faster than naive byte-by-byte scanning + let finder = memmem::Finder::new(&[0x16, 0x03]); + + let mut pos = 0; + while let Some(offset) = finder.find(&data[pos..]) { + let absolute_offset = pos + offset; + + // Check if we have enough bytes for full pattern + if absolute_offset + MIN_TLS_HANDSHAKE_PATTERN_LEN > data.len() { + return None; + } + + // Verify remaining bytes manually + let candidate = &data[absolute_offset..]; + if (candidate[2] >= 0x01 && candidate[2] <= 0x03) // TLS minor version + && (candidate[5] == 0x01 || candidate[5] == 0x02) + // ClientHello or ServerHello + { + return Some(absolute_offset); + } + + // Not a match, continue searching after this position + pos = absolute_offset + 1; + } + + None +} + +/// Optimized suffix matching using decision tree based on last byte +/// +/// Returns the length of suffix to keep (0-5 bytes) +/// +/// Uses a decision tree: checks the last byte first to narrow down possibilities, +/// then validates from longest to shortest match for that last byte. +/// +/// This is used to minimize buffering when searching for TLS patterns - if the data +/// doesn't end with a potential pattern prefix, we can flush it immediately. +/// +/// If the data is shorter than the maximum pattern length (5 bytes), returns the +/// full buffer length since we can't make assumptions about partial data. +pub fn get_partial_tls_pattern_suffix_len(data: &[u8]) -> usize { + if data.is_empty() { + return 0; + } + + let len = data.len(); + + // If buffer is smaller than max pattern length, return full length + // We can't make assumptions about what might grow into a pattern + if len <= 5 { + return len; + } + let last = data[len - 1]; + + match last { + // Last byte is 0x16 - can only be position 0 of pattern + 0x16 => { + // Pattern: [0x16, ...] + // Only need to check 1-byte suffix + 1 + } + + // Last byte is 0x03 - could be position 1 (major version) or 2 (minor version) + 0x03 => { + // Check longest first (3-byte), fall back to 2-byte + if len >= 3 && data[len - 3] == 0x16 && data[len - 2] == 0x03 { + 3 // [0x16, 0x03, 0x03] + } else if len >= 2 && data[len - 2] == 0x16 { + 2 // [0x16, 0x03] + } else { + 0 + } + } + + // Last byte is 0x01 or 0x02 - could be position 2 (minor version if 0x01) + // or position 3/4 (length bytes - wildcard) or position 5 (handshake type) + 0x01 | 0x02 => { + // Check from longest to shortest: 5-byte → 4-byte → 3-byte + if len >= 5 + && data[len - 5] == 0x16 + && data[len - 4] == 0x03 + && (data[len - 3] >= 0x01 && data[len - 3] <= 0x03) + { + 5 // [0x16, 0x03, 0x01-0x03, ?, 0x01/0x02] + } else if len >= 4 + && data[len - 4] == 0x16 + && data[len - 3] == 0x03 + && (data[len - 2] >= 0x01 && data[len - 2] <= 0x03) + { + 4 // [0x16, 0x03, 0x01-0x03, 0x01/0x02] + } else if len >= 3 && data[len - 3] == 0x16 && data[len - 2] == 0x03 && last == 0x01 { + // Only valid if last byte is 0x01 (valid minor version) + // 0x02 is ServerHello handshake type, not minor version + 3 // [0x16, 0x03, 0x01] + } else { + 0 + } + } + + // Any other byte - can only be in wildcard length field (position 3 or 4) + _ => { + // Check 5-byte, then 4-byte + if len >= 5 + && data[len - 5] == 0x16 + && data[len - 4] == 0x03 + && (data[len - 3] >= 0x01 && data[len - 3] <= 0x03) + { + 5 // [0x16, 0x03, 0x01-0x03, ?, last] + } else if len >= 4 + && data[len - 4] == 0x16 + && data[len - 3] == 0x03 + && (data[len - 2] >= 0x01 && data[len - 2] <= 0x03) + { + 4 // [0x16, 0x03, 0x01-0x03, last] + } else { + 0 + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_find_tls_handshake_clienthello() { + // TLS 1.3 ClientHello: [0x16, 0x03, 0x03, length_bytes, 0x01] + let data = vec![0x16, 0x03, 0x03, 0x00, 0x10, 0x01]; + assert_eq!(find_tls_handshake_start(&data), Some(0)); + } + + #[test] + fn test_find_tls_handshake_serverhello() { + // TLS 1.3 ServerHello: [0x16, 0x03, 0x03, length_bytes, 0x02] + let data = vec![0x16, 0x03, 0x03, 0x00, 0x10, 0x02]; + assert_eq!(find_tls_handshake_start(&data), Some(0)); + } + + #[test] + fn test_find_tls_handshake_with_prefix() { + // VLESS header (26 bytes) + TLS ClientHello + let mut data = vec![0x00; 26]; // VLESS header + data.extend_from_slice(&[0x16, 0x03, 0x01, 0x02, 0x00, 0x01]); // TLS 1.0 ClientHello + assert_eq!(find_tls_handshake_start(&data), Some(26)); + } + + #[test] + fn test_find_tls_handshake_not_found() { + let data = vec![0x00, 0x01, 0x02, 0x03, 0x04, 0x05]; + assert_eq!(find_tls_handshake_start(&data), None); + } + + #[test] + fn test_find_tls_handshake_too_short() { + let data = vec![0x16, 0x03, 0x03]; + assert_eq!(find_tls_handshake_start(&data), None); + } + + #[test] + fn test_suffix_len_no_pattern() { + // Small buffer (≤ 5 bytes) returns full length + let data = vec![0x00, 0x01, 0x02, 0x03]; + assert_eq!(get_partial_tls_pattern_suffix_len(&data), 4); + } + + #[test] + fn test_suffix_len_single_byte() { + // Small buffer (≤ 5 bytes) returns full length + let data = vec![0x00, 0x01, 0x16]; + assert_eq!(get_partial_tls_pattern_suffix_len(&data), 3); + } + + #[test] + fn test_suffix_len_two_bytes() { + // Small buffer (≤ 5 bytes) returns full length + let data = vec![0x00, 0x01, 0x16, 0x03]; + assert_eq!(get_partial_tls_pattern_suffix_len(&data), 4); + } + + #[test] + fn test_suffix_len_three_bytes() { + // Small buffer (≤ 5 bytes) returns full length + let data = vec![0x00, 0x16, 0x03, 0x01]; + assert_eq!(get_partial_tls_pattern_suffix_len(&data), 4); + } + + #[test] + fn test_suffix_len_five_bytes() { + // Exactly 5 bytes, returns full length + let data = vec![0x16, 0x03, 0x03, 0xAB, 0x01]; + assert_eq!(get_partial_tls_pattern_suffix_len(&data), 5); + } + + #[test] + fn test_suffix_len_false_positive() { + // Exactly 5 bytes, returns full length (even if not valid pattern) + let data = vec![0x16, 0x03, 0x04, 0x00, 0x01]; // 0x04 is not valid TLS minor version + assert_eq!(get_partial_tls_pattern_suffix_len(&data), 5); + } + + #[test] + fn test_suffix_len_large_buffer_no_pattern() { + // Large buffer (> 5 bytes) with no pattern at end + let data = vec![0x00, 0x01, 0x02, 0x03, 0x04, 0x05]; + assert_eq!(get_partial_tls_pattern_suffix_len(&data), 0); + } + + #[test] + fn test_suffix_len_large_buffer_with_pattern() { + // Large buffer (> 5 bytes) ending with partial pattern + let data = vec![0x00, 0x01, 0x02, 0x03, 0x16, 0x03]; + assert_eq!(get_partial_tls_pattern_suffix_len(&data), 2); + } +} diff --git a/src/vless/vision_filter.rs b/src/vless/vision_filter.rs new file mode 100644 index 00000000..e9e0f160 --- /dev/null +++ b/src/vless/vision_filter.rs @@ -0,0 +1,512 @@ +use crate::buf_reader::BufReader; + +// TLS record types +const TLS_CONTENT_TYPE_HANDSHAKE: u8 = 0x16; + +// TLS handshake types +const TLS_HANDSHAKE_TYPE_CLIENT_HELLO: u8 = 0x01; +const TLS_HANDSHAKE_TYPE_SERVER_HELLO: u8 = 0x02; + +// TLS ciphers +const TLS13_CIPHER_AES_128_CCM_8_SHA256: u16 = 0x1305; + +/// TLS extension type for supported_versions +const TLS_EXT_SUPPORTED_VERSIONS: u16 = 0x002b; + +/// VISION protocol filter state +/// +/// This structure maintains only the state needed for TLS record filtering +/// and detection, tracking TLS patterns and remaining records to analyze. +#[derive(Debug)] +pub struct VisionFilter { + /// Number of records remaining to analyze for TLS patterns + record_filter_count: usize, + + /// Any TLS version detected (ClientHello seen) + is_tls: bool, + + /// TLS 1.2 or above detected (ServerHello seen) + is_tls12_or_above: bool, + + /// Whether XTLS can be enabled (TLS 1.3 + supported cipher) + supports_xtls: bool, +} + +impl VisionFilter { + /// Create a new VisionFilter with default record count + pub fn new() -> Self { + Self::new_with_record_count(8) + } + + fn new_with_record_count(record_count: usize) -> Self { + Self { + record_filter_count: record_count, + supports_xtls: false, + is_tls: false, + is_tls12_or_above: false, + } + } + + /// Check if filtering is complete + pub fn is_filtering(&self) -> bool { + self.record_filter_count > 0 + } + + pub fn decrement_filter_count(&mut self) { + self.record_filter_count = self.record_filter_count.saturating_sub(1); + } + + pub fn stop_filtering(&mut self, reason: String) { + log::debug!("VISION: Stopping filtering - {}", reason); + self.record_filter_count = 0; + } + + /// Returns true if any TLS version has been detected (ClientHello seen) + pub fn is_tls(&self) -> bool { + self.is_tls + } + + /// Returns true if TLS 1.2 or above has been detected (ServerHello seen) + pub fn is_tls12_or_above(&self) -> bool { + self.is_tls12_or_above + } + + /// Returns true if XTLS can be enabled (TLS 1.3 + supported cipher) + pub fn supports_xtls(&self) -> bool { + self.supports_xtls + } + + /// Analyze a complete TLS record for patterns and update state accordingly + /// + /// Caller must extract complete TLS records using TlsDeframer before calling this. + /// This method detects ClientHello, ServerHello, TLS version, cipher suite, and + /// ApplicationData to determine when to switch to direct copy mode. + /// + /// This is a public method now - VisionStream will call it after recordizing. + pub fn filter_record(&mut self, data: &[u8]) { + if self.record_filter_count == 0 { + return; + } + + self.record_filter_count = self.record_filter_count.saturating_sub(1); + + if data.len() < 5 { + // invalid record length + self.stop_filtering(format!("invalid record length: {}", data.len())); + return; + } + + // Detect ClientHello - minimal check for TLS handshake record + if !self.is_tls + && data.len() >= 6 + && data[0] == TLS_CONTENT_TYPE_HANDSHAKE + && data[1] == 0x03 // TLS version major + && data[5] == TLS_HANDSHAKE_TYPE_CLIENT_HELLO + { + self.is_tls = true; + log::debug!("VISION: Detected TLS ClientHello"); + } + + // Detect and parse ServerHello + if !self.is_tls12_or_above + && data.len() >= 6 + && data[0] == TLS_CONTENT_TYPE_HANDSHAKE + && data[1] == 0x03 // TLS version major + && data[2] == 0x03 // TLS version minor (1.2) + && data[5] == TLS_HANDSHAKE_TYPE_SERVER_HELLO + { + self.is_tls12_or_above = true; + self.is_tls = true; + + // Parse ServerHello properly - delegate to vision_tls_util + match parse_server_hello(data) { + Ok(parsed) => { + log::debug!( + "VISION: Detected TLS cipher suite: 0x{:04x}", + parsed.cipher_suite + ); + + if parsed.has_tls13_version { + log::debug!("VISION: Detected TLS 1.3 via supported_versions extension"); + + // Check if cipher is supported for XTLS + // Following Xray-core: all TLS 1.3 ciphers EXCEPT 0x1305 support XTLS + if parsed.cipher_suite == TLS13_CIPHER_AES_128_CCM_8_SHA256 { + log::warn!( + "VISION: TLS 1.3 detected but cipher 0x{:04x} (TLS_AES_128_CCM_8_SHA256) not supported for XTLS", + parsed.cipher_suite + ); + } else { + // Assume all other TLS 1.3 ciphers support XTLS + self.supports_xtls = true; + log::debug!( + "VISION: TLS 1.3 with supported cipher 0x{:04x} - XTLS enabled", + parsed.cipher_suite + ); + } + + // Stop filtering early + self.stop_filtering("TLS 1.3 handshake detected".to_string()); + } + } + Err(e) => { + self.stop_filtering(format!("invalid ServerHello: {}", e)); + } + } + } + } +} + +/// Result of parsing a ServerHello message +struct ParsedServerHello { + pub cipher_suite: u16, + pub has_tls13_version: bool, +} + +/// Parse ServerHello structure following RFC 8446 format +/// +/// ServerHello format: +/// - TLS Record Header (5 bytes): +/// - ContentType (1 byte): 0x16 (Handshake) +/// - ProtocolVersion (2 bytes): 0x0303 (TLS 1.2 for compatibility) +/// - Length (2 bytes): payload length +/// - Handshake Header (4 bytes): +/// - HandshakeType (1 byte): 0x02 (ServerHello) +/// - Length (3 bytes): message length +/// - ServerHello Content: +/// - LegacyVersion (2 bytes): 0x0303 (TLS 1.2) +/// - Random (32 bytes): server random +/// - LegacySessionIdEcho (1 + N bytes): session ID length + session ID +/// - CipherSuite (2 bytes): selected cipher +/// - LegacyCompressionMethod (1 byte): 0x00 +/// - Extensions (2 + N bytes): extensions length + extensions +/// +/// TODO: consolidate with ShadowTLS version +fn parse_server_hello(data: &[u8]) -> std::io::Result { + // Minimum size: 5 (record header) + 4 (handshake header) + 2 (version) + 32 (random) + 1 (session id len) + 2 (cipher) + 1 (compression) = 47 + if data.len() < 47 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "ServerHello too short", + )); + } + + // Skip TLS record header (5 bytes) and handshake type (1 byte) + let mut reader = BufReader::new(&data[6..]); + + // Read handshake message length (3 bytes) + let _message_len = reader.read_u24_be()?; + + // Read legacy version (2 bytes) - should be 0x0303 for TLS 1.2 + let legacy_version_major = reader.read_u8()?; + let legacy_version_minor = reader.read_u8()?; + if legacy_version_major != 3 || legacy_version_minor != 3 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!( + "expected TLS 1.2 (major/minor 3.3), got major/minor {}.{}", + legacy_version_major, legacy_version_minor + ), + )); + } + + // Skip server random (32 bytes) + reader.skip(32)?; + + // Read session ID length and skip session ID + let session_id_len = reader.read_u8()?; + if session_id_len > 0 { + reader.skip(session_id_len as usize)?; + } + + // Read cipher suite (2 bytes) + let cipher_suite = reader.read_u16_be()?; + + // Skip compression method (1 byte) + reader.skip(1)?; + + // Check if we have extensions + if reader.is_consumed() { + // No extensions, assume not TLS 1.3 + return Ok(ParsedServerHello { + cipher_suite, + has_tls13_version: false, + }); + } + + // Read extensions length + let extensions_len = match reader.read_u16_be() { + Ok(len) => len as usize, + Err(_) => { + // No extensions + return Ok(ParsedServerHello { + cipher_suite, + has_tls13_version: false, + }); + } + }; + + // Parse extensions to find supported_versions + let mut has_tls13_version = false; + let extensions_data = match reader.read_slice(extensions_len) { + Ok(data) => data, + Err(_) => { + // Invalid extensions length + return Ok(ParsedServerHello { + cipher_suite, + has_tls13_version: false, + }); + } + }; + + let mut ext_reader = BufReader::new(extensions_data); + + while !ext_reader.is_consumed() { + let extension_type = match ext_reader.read_u16_be() { + Ok(t) => t, + Err(_) => break, + }; + let extension_len = match ext_reader.read_u16_be() { + Ok(l) => l as usize, + Err(_) => break, + }; + + if extension_type == TLS_EXT_SUPPORTED_VERSIONS { + // supported_versions extension found + // In ServerHello, this is 2 bytes: the selected version + if extension_len == 2 { + match ext_reader.read_slice(2) { + Ok(version_bytes) => { + // TLS 1.3 is 0x0304 + if version_bytes[0] == 0x03 && version_bytes[1] == 0x04 { + has_tls13_version = true; + break; + } + } + Err(_) => break, + } + } else if ext_reader.skip(extension_len).is_err() { + break; + } + } else if ext_reader.skip(extension_len).is_err() { + // Skip unknown extension + break; + } + } + + Ok(ParsedServerHello { + cipher_suite, + has_tls13_version, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_vision_filter_initialization() { + let filter = VisionFilter::new(); + + assert!(filter.is_filtering()); + assert!(!filter.supports_xtls()); + assert!(!filter.is_tls()); + assert!(!filter.is_tls12_or_above()); + } + + #[test] + fn test_vision_filter_with_custom_record_count() { + let filter = VisionFilter::new_with_record_count(5); + + assert!(filter.is_filtering()); + assert!(!filter.supports_xtls()); + assert!(!filter.is_tls()); + } + + #[test] + fn test_stop_filtering() { + let mut filter = VisionFilter::new(); + + assert!(filter.is_filtering()); + + filter.stop_filtering("test".to_string()); + + assert!(!filter.is_filtering()); + } + + #[test] + fn test_detect_client_hello() { + let mut filter = VisionFilter::new(); + + // Mock ClientHello record + let client_hello = vec![ + 0x16, 0x03, 0x01, // TLS Handshake, version 3.1 + 0x00, 0x05, // Length: 5 + 0x01, // ClientHello type + 0x00, 0x00, 0x00, // Rest of handshake + ]; + + filter.filter_record(&client_hello); + + assert!(filter.is_tls()); + assert!(!filter.supports_xtls()); // Not enabled until TLS 1.3 ServerHello + } + + #[test] + fn test_detect_server_hello_tls12() { + let mut filter = VisionFilter::new(); + + // Mock ServerHello record (TLS 1.2) + let server_hello = vec![ + 0x16, 0x03, 0x03, // TLS Handshake, version 3.3 (TLS 1.2) + 0x00, 0x05, // Length: 5 + 0x02, // ServerHello type + 0x00, 0x00, 0x00, // Rest of handshake + ]; + + filter.filter_record(&server_hello); + + assert!(filter.is_tls()); + assert!(filter.is_tls12_or_above()); + } + + #[test] + fn test_detect_application_data() { + let mut filter = VisionFilter::new(); + + // First detect TLS + let client_hello = vec![ + 0x16, 0x03, 0x01, // TLS Handshake + 0x00, 0x05, // Length: 5 + 0x01, // ClientHello type + 0x00, 0x00, 0x00, + ]; + filter.filter_record(&client_hello); + + // Then ApplicationData + let app_data = vec![ + 0x17, 0x03, 0x03, // ApplicationData, TLS 1.2 + 0x00, 0x05, // Length: 5 + 0x00, 0x00, 0x00, 0x00, 0x00, // Encrypted data + ]; + + filter.filter_record(&app_data); + + assert!(filter.is_tls()); + assert!(filter.is_filtering()); + } + + #[test] + fn test_record_count_decrements() { + let mut filter = VisionFilter::new_with_record_count(3); + + let dummy_record = vec![0x16, 0x03, 0x01, 0x00, 0x05, 0x00]; + + filter.filter_record(&dummy_record); + assert!(filter.is_filtering()); + + filter.filter_record(&dummy_record); + assert!(filter.is_filtering()); + + filter.filter_record(&dummy_record); + assert!(!filter.is_filtering()); + + // Further calls should not decrement + filter.filter_record(&dummy_record); + assert!(!filter.is_filtering()); + } + + #[test] + fn test_filter_record_no_processing_when_complete() { + let mut filter = VisionFilter::new(); + filter.stop_filtering("test".to_string()); + + let data = vec![0x16, 0x03, 0x01, 0x00, 0x05, 0x01, 0x00, 0x00, 0x00]; + + // Should return immediately without processing + filter.filter_record(&data); + + assert!(!filter.is_tls()); // State shouldn't change + assert!(!filter.is_filtering()); + } + + #[test] + fn test_short_record_ignored() { + let mut filter = VisionFilter::new(); + + let short_record = vec![0x16, 0x03, 0x01]; // Less than 6 bytes + + filter.filter_record(&short_record); + + assert!(!filter.is_tls()); // Should not detect anything + assert!(!filter.is_filtering()); // Filtering stopped due to invalid record + } + + #[test] + fn test_non_tls_record_ignored() { + let mut filter = VisionFilter::new(); + + let non_tls_record = vec![0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF]; + + filter.filter_record(&non_tls_record); + + assert!(!filter.is_tls()); + assert!(!filter.supports_xtls()); + assert!(filter.is_filtering()); + } + + #[test] + fn test_filter_record_client_hello() { + let mut filter = VisionFilter::new(); + + // Simulate ClientHello + let client_hello = vec![0x16, 0x03, 0x01, 0x00, 0x05, 0x01, 0x00, 0x00, 0x00, 0x00]; + + filter.filter_record(&client_hello); + + assert!(filter.is_tls()); + assert!(!filter.supports_xtls()); // Not optimistically enabled anymore + } + + #[test] + fn test_filter_record_server_hello() { + let mut filter = VisionFilter::new(); + + // Simulate ServerHello + let server_hello = vec![0x16, 0x03, 0x03, 0x00, 0x05, 0x02, 0x00, 0x00, 0x00, 0x00]; + + filter.filter_record(&server_hello); + + assert!(filter.is_tls()); + assert!(filter.is_tls12_or_above()); + } + + #[test] + fn test_client_hello_only_detected_once() { + let mut filter = VisionFilter::new(); + + let client_hello = vec![0x16, 0x03, 0x01, 0x00, 0x05, 0x01, 0x00, 0x00, 0x00, 0x00]; + + filter.filter_record(&client_hello); + assert!(filter.is_tls()); + assert!(filter.is_filtering()); + + // Send again - should not reset state + filter.filter_record(&client_hello); + assert!(filter.is_tls()); + assert!(filter.is_filtering()); + } + + #[test] + fn test_server_hello_only_detected_once() { + let mut filter = VisionFilter::new(); + + let server_hello = vec![0x16, 0x03, 0x03, 0x00, 0x05, 0x02, 0x00, 0x00, 0x00, 0x00]; + + filter.filter_record(&server_hello); + assert!(filter.is_tls12_or_above()); + // Note: This mock ServerHello is too short and will cause parse_server_hello to fail, + // which stops filtering. This is expected behavior - invalid records stop filtering. + assert!(!filter.is_filtering()); + } +} diff --git a/src/vless/vision_pad.rs b/src/vless/vision_pad.rs new file mode 100644 index 00000000..07f06629 --- /dev/null +++ b/src/vless/vision_pad.rs @@ -0,0 +1,67 @@ +use bytes::{BufMut, Bytes, BytesMut}; +use rand::Rng; + +// Padding parameters +const LONG_PADDING_MIN: usize = 900; +const LONG_PADDING_RANDOM_MAX: usize = 500; +const SHORT_PADDING_RANDOM_MAX: usize = 256; +const MAX_PADDING_SIZE: usize = 8171; // buf.Size - 21 + +pub fn pad_with_uuid_and_command(data: &[u8], uuid: &[u8; 16], command: u8, is_tls: bool) -> Bytes { + pad(data, Some(uuid), command, is_tls) +} +// +pub fn pad_with_command(data: &[u8], command: u8, is_tls: bool) -> Bytes { + pad(data, None, command, is_tls) +} + +fn pad(data: &[u8], uuid: Option<&[u8; 16]>, command: u8, is_tls: bool) -> Bytes { + let content_len = data.len() as u16; + let padding_len = calculate_padding_length(content_len as usize, is_tls); + + let uuid_len = if uuid.is_some() { 16 } else { 0 }; + let total_size = uuid_len + 1 + 2 + 2 + data.len() + padding_len; + + let mut output = BytesMut::with_capacity(total_size); + + if let Some(uuid) = uuid { + output.put_slice(uuid); + } + + output.put_u8(command); + output.put_u16(content_len); + output.put_u16(padding_len as u16); + output.put_slice(data); + + // Write random padding + if padding_len > 0 { + let padding_start = output.len(); + output.resize(padding_start + padding_len, 0); + rand::rng().fill(&mut output[padding_start..]); + } + + output.freeze() +} + +/// Calculate padding length based on content size and TLS detection +fn calculate_padding_length(content_len: usize, is_tls: bool) -> usize { + let mut rng = rand::rng(); + + // Calculate maximum allowable padding (avoid overflow) + // Matches Xray-core logic: buf.Size - 21 - contentLen + // MAX_PADDING_SIZE is already (buf.Size - 21), so we only subtract content_len + let max_allowable = MAX_PADDING_SIZE.saturating_sub(content_len); + + if is_tls && content_len < LONG_PADDING_MIN { + // Long padding for TLS handshake phase + let random_part = rng.random_range(0..LONG_PADDING_RANDOM_MAX); + let padding = LONG_PADDING_MIN + .saturating_sub(content_len) + .saturating_add(random_part); + std::cmp::min(padding, max_allowable) + } else { + // Short random padding + let padding = rng.random_range(0..SHORT_PADDING_RANDOM_MAX); + std::cmp::min(padding, max_allowable) + } +} diff --git a/src/vless/vision_stream.rs b/src/vless/vision_stream.rs new file mode 100644 index 00000000..f2fd95e2 --- /dev/null +++ b/src/vless/vision_stream.rs @@ -0,0 +1,1480 @@ +use crate::crypto::CryptoConnection; +/// VISION stream implementation +/// +/// VisionStream wraps an IO stream and TLS session, allowing it to: +/// 1. Use TLS with VISION padding to hide protocol fingerprints +/// 2. Detect TLS-in-TLS scenarios by analyzing traffic patterns +/// 3. Switch to direct I/O mode, bypassing TLS for zero-copy performance +use bytes::{Buf, BytesMut}; +use futures::ready; +use std::io::{self, BufRead, Write}; +use std::pin::Pin; +use std::task::{Context, Poll}; +use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; + +use crate::async_stream::AsyncStream; +use crate::crypto::feed_crypto_connection; +use crate::sync_adapter::{SyncReadAdapter, SyncWriteAdapter}; +use crate::util::allocate_vec; + +use super::tls_deframer::TlsDeframer; +use super::tls_fuzzy_deframer::{DeframeResult, FuzzyTlsDeframer}; +use super::vision_filter::VisionFilter; +use super::vision_unpad::{UnpadCommand, UnpadResult, VisionUnpadder}; + +#[inline] +fn feed_and_process_crypto_connection( + session: &mut CryptoConnection, + data: &[u8], +) -> std::io::Result { + feed_crypto_connection(session, data)?; + Ok(session + .process_new_packets() + .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidData, err))? + .plaintext_bytes_to_read()) +} + +/// Current operating mode of the VISION stream +#[derive(Debug, PartialEq)] +enum VisionMode { + /// Using TLS with VISION padding/unpadding + PaddingTls, + /// Regular TLS (padding ended, but still using TLS encryption) + Tls, + /// Direct I/O, TLS bypassed completely + Direct, +} + +impl std::fmt::Display for VisionMode { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + VisionMode::PaddingTls => write!(f, "PaddingTls"), + VisionMode::Tls => write!(f, "Tls"), + VisionMode::Direct => write!(f, "Direct"), + } + } +} + +pub struct VisionStream { + /// The underlying transport stream (e.g., TcpStream) + tcp: IO, + + /// rustls session for TLS encryption and decryption + /// Used in PaddingTls and Tls mode + session: CryptoConnection, + + /// Current READ operating mode (independent from write mode) + read_mode: VisionMode, + + /// Current WRITE operating mode (independent from read mode) + write_mode: VisionMode, + + /// Outer TLS deframer for read path + /// Deframes TCP stream into complete TLS records before feeding to rustls + /// Used in PaddingTls mode + outer_read_deframer: Option, + + /// Unpadding state machine for read path + read_unpadder: VisionUnpadder, + + /// Inner TLS deframer for filtering read packets + /// Used in PaddingTls mode + inner_read_deframer: FuzzyTlsDeframer, + + /// Inner TLS deframer for filtering write packets + /// Used in PaddingTls mode + inner_write_deframer: FuzzyTlsDeframer, + + /// Whether this is the first write (includes UUID in padding) + /// Used in PaddingTls mode + write_first_packet: bool, + + /// User UUID for padding validation + /// Used in PaddingTls mode + user_uuid: [u8; 16], + + /// TLS pattern filter for inner TLS + /// Used in PaddingTls mode + filter: VisionFilter, + + /// Buffer for leftover data when switching modes or when output buffer is too small + /// Used in PaddingTls mode + pending_read: BytesMut, + + /// Whether we need to read VLESS response header (client-side only) + /// Similar to shadowsocks's `is_initial_read` pattern + /// Used in PaddingTls mode + vless_response_pending: bool, + + /// Partial VLESS response data accumulated across multiple TLS records + /// Used to handle cases where the response header is split across TLS records + /// Used in PaddingTls mode + partial_vless_response: BytesMut, + + /// Whether we need to send VLESS response header (server-side only) + /// The response will be prepended to the first write, like sing-box does + /// Used in PaddingTls mode + vless_response_to_send: bool, + + /// Flag indicating we should switch to direct mode on the NEXT write call + /// Used in PaddingTls mode + pending_direct_mode_switch: bool, + + /// Flag indicating we should switch to Tls mode on the NEXT write call + /// Used in PaddingTls mode + pending_tls_mode_switch: bool, + + /// Reusable buffer for TLS read operations (TCP → read_tls → decrypt) + /// Used in PaddingTls mode + tls_read_buffer: Vec, + + /// Buffer for plaintext data waiting to be written to the TLS session buffer + /// When the rustls Writer buffer fills (write length 0), we store the remainder here + /// Used in PaddingTls mode - drained immediately before switching modes + pending_plain_writes: BytesMut, + + /// Caches whether we've hit EOF when reading + /// Used in PaddingTls and Tls mode + is_read_eof: bool, +} + +impl VisionStream +where + IO: AsyncRead + AsyncWrite + Unpin, +{ + /// Create a new VisionStream for server-side (inbound) connections with VLESS response writing + pub fn new_server( + tcp: IO, + session: CryptoConnection, + user_uuid: [u8; 16], + initial_read_data: &[u8], + ) -> std::io::Result { + // Validate that this is a server connection + if !session.is_server() { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "VisionStream::new_server requires a server-side connection", + )); + } + + let mut stream = Self::new_common( + tcp, session, user_uuid, false, // vless_response_pending + true, // vless_response_to_send + ); + + // feed in the initial read data that came from reading the VLESS request + stream.feed_initial_read_data(initial_read_data)?; + // read all the remaining plaintext data from the session since we start directly + // with TCP reads in poll_read_padding_tls + stream.feed_initial_session_data()?; + + Ok(stream) + } + + /// Create a new VisionStream for client-side connections with VLESS response handling + pub fn new_client(tcp: IO, session: CryptoConnection, user_uuid: [u8; 16]) -> Self { + // Validate that this is a client connection + if !session.is_client() { + panic!("VisionStream::new_client requires a client-side connection"); + } + + Self::new_common( + tcp, session, user_uuid, true, // vless_response_pending + false, // vless_response_to_send + ) + } + + fn new_common( + tcp: IO, + session: CryptoConnection, + user_uuid: [u8; 16], + vless_response_pending: bool, + vless_response_to_send: bool, + ) -> Self { + Self { + tcp, + session, + // this should be somewhat smaller than TLS_MAX_RECORD_SIZE when used with session.read_tls, + // else we will hit "message buffer full" (see other comment) - because if we set it at + // TLS_MAX_RECORD_SIZE, and a small record is read first, there can be a large partial record + // left in the buffer and then prevoius_record_size + TLS_MAX_RECORD_SIZE > TLS_MAX_RECORD_SIZE + // will trigger the error. + tls_read_buffer: allocate_vec(8192), + read_mode: VisionMode::PaddingTls, + write_mode: VisionMode::PaddingTls, + outer_read_deframer: Some(TlsDeframer::new()), + inner_read_deframer: FuzzyTlsDeframer::new(), + inner_write_deframer: FuzzyTlsDeframer::new(), + read_unpadder: VisionUnpadder::new(user_uuid), + write_first_packet: true, + user_uuid, + filter: VisionFilter::new(), + pending_read: BytesMut::new(), + vless_response_pending, + partial_vless_response: BytesMut::new(), + vless_response_to_send, + pending_direct_mode_switch: false, + pending_tls_mode_switch: false, + pending_plain_writes: BytesMut::new(), + is_read_eof: false, + } + } + + /// Feed initial read data that was already decrypted by TLS but needs VISION unpadding + /// This is used when the VLESS header parser read extra bytes from the TLS stream. + /// This can be called multiple times in order to successfully drain the session. + /// Those bytes are already TLS-decrypted but still VISION-padded. + fn feed_initial_read_data(&mut self, data: &[u8]) -> std::io::Result<()> { + if data.is_empty() { + return Ok(()); + } + + log::debug!( + "VISION: Feeding {} initial bytes (already TLS-decrypted, needs unpadding)", + data.len() + ); + + self.handle_padded_bytes(data)?; + + Ok(()) + } + + // Read all the available plaintext data from the session. + // This should only be necessary for server streams since a read occurred in order to read the + // VLESS request before the session is passed to VisionStream. + fn feed_initial_session_data(&mut self) -> std::io::Result<()> { + let plaintext_len = self + .session + .process_new_packets() + .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidData, err))? + .plaintext_bytes_to_read(); + let mut decrypted = Vec::with_capacity(plaintext_len); + + let mut reader = self.session.reader(); + let mut i = 0; + while i < plaintext_len { + match reader.fill_buf() { + Ok(buf) => { + let len = buf.len(); + if len == 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "TLS session filled buffer with zero bytes", + )); + } + decrypted.extend_from_slice(buf); + reader.consume(len); + log::debug!("VISION: Filled more unparsed data of length {}", len); + i += len; + } + Err(e) => { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("failed to read TLS session: {e}"), + )); + } + } + } + + self.handle_padded_bytes(&decrypted) + } + + fn switch_read_to_direct_mode(&mut self) -> io::Result<()> { + if self.read_mode != VisionMode::PaddingTls { + return Err(std::io::Error::other(format!( + "switch_read_to_direct_mode called from mode {}", + self.read_mode + ))); + } + + log::debug!("VISION READ: Switching to direct copy mode (asymmetric - write side may still use padding)"); + + // Set read mode to Direct FIRST to avoid inconsistent state + self.read_mode = VisionMode::Direct; + + log::debug!("VISION READ: Switched to direct mode"); + + self.post_padding_cleanup(); + + Ok(()) + } + + fn switch_read_to_tls_mode(&mut self) -> io::Result<()> { + if self.read_mode != VisionMode::PaddingTls { + return Err(std::io::Error::other(format!( + "switch_read_to_tls_mode called from mode {}", + self.read_mode + ))); + } + + log::debug!("VISION READ: Switching to Tls mode (no XTLS support)"); + + // Set read mode to Tls + self.read_mode = VisionMode::Tls; + + self.post_padding_cleanup(); + + Ok(()) + } + + /// Switch WRITE side from TLS mode to direct I/O mode + /// This doesn't need to preserve any data, just updates the mode + fn switch_write_to_direct_mode(&mut self) -> io::Result<()> { + if self.write_mode != VisionMode::PaddingTls { + return Err(std::io::Error::other(format!( + "switch_write_to_direct_mode called from mode {}", + self.write_mode + ))); + } + + log::debug!("VISION WRITE: Switching to direct copy mode (asymmetric - read side may still use padding)"); + + // Set write mode to Direct + self.write_mode = VisionMode::Direct; + + self.post_padding_cleanup(); + + Ok(()) + } + + /// Switch WRITE side from PaddingTls to Tls mode + /// This happens when XTLS is not supported or TLS 1.2 is detected + fn switch_write_to_tls_mode(&mut self) -> io::Result<()> { + if self.write_mode != VisionMode::PaddingTls { + return Err(std::io::Error::other(format!( + "switch_write_to_tls_mode called from mode {}", + self.write_mode + ))); + } + + log::debug!("VISION WRITE: Switching to Tls mode (no XTLS support)"); + + // Set write mode to Tls + self.write_mode = VisionMode::Tls; + + self.post_padding_cleanup(); + + Ok(()) + } + + fn post_padding_cleanup(&mut self) { + if self.read_mode == VisionMode::PaddingTls || self.write_mode == VisionMode::PaddingTls { + return; + } + + log::debug!("VISION: Cleaning up after read and write padding mode switch"); + + // TODO: consider using an Option for all PaddingTls fields instead + self.inner_read_deframer.deallocate(); + self.inner_write_deframer.deallocate(); + self.tls_read_buffer = Vec::new(); + self.pending_plain_writes = BytesMut::new(); + } + + /// Drain pending plaintext writes to TLS session buffer + /// Returns Poll::Ready(Ok(())) when fully drained + /// Returns Poll::Pending if session buffer is full (needs TLS/TCP drain first) + fn drain_pending_plain_writes(&mut self, _cx: &mut Context<'_>) -> Poll> { + use std::io::Write; + + if self.pending_plain_writes.is_empty() { + return Poll::Ready(Ok(())); + } + + log::debug!( + "VISION WRITE: Draining {} pending plaintext bytes", + self.pending_plain_writes.len() + ); + + // Try to write pending plaintext to TLS session + while !self.pending_plain_writes.is_empty() { + match self.session.writer().write(&self.pending_plain_writes) { + Ok(n) => { + if n == 0 { + // Session buffer full - return Pending so drain_all_writes can drain TLS/TCP first + log::debug!( + "VISION WRITE: Session buffer full ({} plaintext bytes remaining), returning Pending", + self.pending_plain_writes.len() + ); + return Poll::Pending; + } + + self.pending_plain_writes.advance(n); + } + Err(e) => return Poll::Ready(Err(e)), + } + } + + Poll::Ready(Ok(())) + } + + /// Drain all pending writes in PaddingTls mode: plain → TLS → TCP + /// + /// This drains the write pipeline in priority order: + /// 1. TLS session → TCP socket (highest priority, creates space in session) + /// 2. Plain buffer → TLS session (writes to session) + /// + /// Returns Poll::Ready(Ok(())) when all buffers are fully drained + /// Returns Poll::Pending if TCP blocks (backpressure - nothing else can make progress) + fn drain_all_writes_padding(&mut self, cx: &mut Context<'_>) -> Poll> { + loop { + if self.session.wants_write() { + match self.write_tls_direct(cx) { + Poll::Ready(Ok(0)) => return Poll::Ready(Err(io::ErrorKind::WriteZero.into())), + Poll::Ready(Ok(_)) => { + // Wrote some data. This might have freed up space in session buffer. + // Loop back to try draining plaintext buffer again. + continue; + } + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + Poll::Pending => { + // TCP is blocked. We can't drain session, so we can't drain plain writes either. + return Poll::Pending; + } + } + } + + if !self.pending_plain_writes.is_empty() { + match self.drain_pending_plain_writes(cx) { + Poll::Ready(Ok(())) => { + // Drained plain writes, session now likely wants to write. + continue; + } + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + Poll::Pending => { + // Session buffer full. But we didn't drain session above (wants_write was false). + // So we are stuck. + // TODO: this is probably an error + return Poll::Pending; + } + } + } + + // If we are here, everything is empty. + return Poll::Ready(Ok(())); + } + } + + /// Drain all pending writes in Tls mode + /// + /// This is simpler than PaddingTls mode because `pending_plain_writes` is no longer used. + /// + /// Returns Poll::Ready(Ok(())) when fully drained + /// Returns Poll::Pending if TCP blocks + fn drain_all_writes_tls(&mut self, cx: &mut Context<'_>) -> Poll> { + // Drain TLS session output directly to TCP + while self.session.wants_write() { + if ready!(self.write_tls_direct(cx))? == 0 { + return Poll::Ready(Err(io::ErrorKind::WriteZero.into())); + } + } + + Poll::Ready(Ok(())) + } + + /// Write data to TLS session, handling buffer full by saving to pending buffer + /// Used in PaddingTls mode + fn write_to_session(&mut self, data: &[u8]) -> io::Result<()> { + use std::io::Write; + + if !self.pending_plain_writes.is_empty() { + // If there's already pending data, just append to it + // (don't try to write, as that would require cx for polling) + self.pending_plain_writes.extend_from_slice(data); + return Ok(()); + } + + let mut written = 0; + while written < data.len() { + match self.session.writer().write(&data[written..]) { + Ok(n) => { + if n == 0 { + // Session buffer full, save remainder + log::debug!( + "VISION WRITE: Session buffer full, saving {} bytes to pending", + data.len() - written + ); + self.pending_plain_writes + .extend_from_slice(&data[written..]); + return Ok(()); + } + written += n; + } + Err(e) => return Err(e), + } + } + + Ok(()) + } + + /// Write TLS session output directly to TCP stream + /// Returns the number of bytes written from TLS session to TCP + fn write_tls_direct(&mut self, cx: &mut Context<'_>) -> Poll> { + let mut writer = SyncWriteAdapter { + io: &mut self.tcp, + cx, + }; + + // Write TLS output DIRECTLY to TCP (zero intermediate copies) + match self.session.write_tls(&mut writer) { + Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => Poll::Pending, + result => Poll::Ready(result), + } + } + + /// Write data in Tls mode + fn poll_write_tls(&mut self, cx: &mut Context<'_>, buf: &[u8]) -> Poll> { + let mut pos = 0; + + while pos < buf.len() { + let mut would_block = false; + + // Write plaintext to TLS session + match self.session.writer().write(&buf[pos..]) { + Ok(n) => pos += n, + Err(e) => return Poll::Ready(Err(e)), + }; + + // Drain TLS output to TCP stream + while self.session.wants_write() { + match self.write_tls_direct(cx) { + Poll::Ready(Ok(0)) | Poll::Pending => { + would_block = true; + break; + } + Poll::Ready(Ok(_)) => (), + Poll::Ready(Err(err)) => return Poll::Ready(Err(err)), + } + } + + return match (pos, would_block) { + (0, true) => Poll::Pending, + (n, true) => Poll::Ready(Ok(n)), // Partial write + (_, false) => continue, // Keep writing + }; + } + + Poll::Ready(Ok(pos)) + } + + /// Read the VLESS response header from the TLS session + /// + /// The VLESS response consists of: + /// - 1 byte: version (must be 0) + /// - 1 byte: addon length + /// - N bytes: addon data (if addon_length > 0) + /// + /// This is called lazily on the first read when vless_response_pending is true. + /// Uses packetization to feed TLS records one at a time to avoid bad MAC errors. + /// + /// Returns any unused data that came after the response + fn poll_read_vless_response(&mut self, cx: &mut Context<'_>) -> Poll> { + log::debug!( + "VLESS: poll_read_vless_response called, partial_vless_response has {} bytes", + self.partial_vless_response.len() + ); + + loop { + // Always read from TCP to feed the deframer + let mut read_buf = ReadBuf::new(&mut self.tls_read_buffer); + match ready!(Pin::new(&mut self.tcp).poll_read(cx, &mut read_buf)) { + Ok(()) => {} + Err(e) => return Poll::Ready(Err(e)), + }; + + let tcp_bytes = read_buf.filled(); + + if tcp_bytes.is_empty() { + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "Connection closed while reading VLESS response", + ))); + } + + log::debug!( + "VLESS: Read {} bytes from TCP for response header", + tcp_bytes.len() + ); + + // Feed to outer read deframer + self.outer_read_deframer + .as_mut() + .expect("outer_read_deframer must exist in PaddingTls mode") + .feed(tcp_bytes); + + // We have to handle all packets and drain the outer_read_deframer because + // poll_read_padded_tls() will always read from the TCP and assume there's outstanding + // data. + // We don't have to worry about consuming some data post-Direct command because the TLS + // flow is one-directional atm and it's not possible to complete a handshake. + // If the peer sends a Direct command without handshake completing, this will error out + // when we feed the TLS record to the session. + let tls_records = self + .outer_read_deframer + .as_mut() + .expect("outer_read_deframer must exist in PaddingTls mode") + .next_records()?; + + for tls_record in tls_records.into_iter() { + // Feed single TLS record to session + let plaintext_len = + feed_and_process_crypto_connection(&mut self.session, tls_record.as_ref())?; + + // Access the decrypted buffer and accumulate into partial_vless_response + let mut reader = self.session.reader(); + + let mut i = 0; + while i < plaintext_len { + let decrypted_buf = match reader.fill_buf() { + Ok(buf) => buf, + Err(e) => return Poll::Ready(Err(e)), + }; + + if decrypted_buf.is_empty() { + return Poll::Ready(Err(io::Error::other( + "no plaintext data when some is available", + ))); + } + + // Accumulate decrypted data into partial buffer + self.partial_vless_response.extend_from_slice(decrypted_buf); + + let consumed_len = decrypted_buf.len(); + reader.consume(consumed_len); + i += consumed_len; + } + } + + log::debug!( + "VLESS: Accumulated vless response with {} bytes", + self.partial_vless_response.len() + ); + + // Try to parse the VLESS response from accumulated data + if self.partial_vless_response.len() < 2 { + // Need at least 2 bytes for header, try to read from the TCP stream again + continue; + } + + // Read the 2-byte header + let version = self.partial_vless_response[0]; + let addon_length = self.partial_vless_response[1]; + + // Validate version + if version != 0 { + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Invalid VLESS response version: {}", version), + ))); + } + + // Check if we have enough data for the full response + let total_response_len = 2 + addon_length as usize; + if self.partial_vless_response.len() < total_response_len { + // Need more data - read more TCP + continue; + } + + // We have the complete VLESS response! Consume it from the buffer + self.partial_vless_response.advance(total_response_len); + + log::debug!( + "VLESS: Successfully parsed {} byte response header (version={}, addon_length={}), {} bytes remaining in buffer", + total_response_len, + version, + addon_length, + self.partial_vless_response.len() + ); + + // If there's any remaining data in partial buffer, it belongs to the actual payload + let remaining_data = std::mem::take(&mut self.partial_vless_response); + return Poll::Ready(Ok(remaining_data)); + } + } + + fn poll_read_padding_tls( + &mut self, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + if self.is_read_eof { + return Poll::Ready(Ok(())); + } + + loop { + // Read from TCP to feed the deframer. + // At this point, there is no data ready in neither the session or the outer read + // deframer. + let mut read_buf = ReadBuf::new(&mut self.tls_read_buffer); + match ready!(Pin::new(&mut self.tcp).poll_read(cx, &mut read_buf)) { + Ok(()) => {} + Err(e) => return Poll::Ready(Err(e)), + }; + + let tcp_bytes = read_buf.filled(); + + if tcp_bytes.is_empty() { + log::debug!("VISION READ: TCP connection EOF in PaddingTls mode"); + self.is_read_eof = true; + return Poll::Ready(Ok(())); + } + + log::debug!("VISION READ: Read {} bytes from TCP", tcp_bytes.len()); + + // Feed to outer TLS deframer + self.outer_read_deframer + .as_mut() + .expect("outer_read_deframer must exist in PaddingTls mode") + .feed(tcp_bytes); + + // Process TLS records one at a time + loop { + let tls_record = match self + .outer_read_deframer + .as_mut() + .expect("outer_read_deframer must exist in PaddingTls mode") + .next_record()? + { + Some(record) => record, + None => { + log::debug!("VISION READ: No more complete TLS records in deframer, need more TCP data"); + break; // Need more TCP data + } + }; + + log::debug!( + "VISION READ: Processing {} byte TLS record", + tls_record.len() + ); + + // Feed single TLS record to session + let plaintext_len = + feed_and_process_crypto_connection(&mut self.session, tls_record.as_ref())?; + + if plaintext_len == 0 { + continue; // Process next TLS record + } + + let mut reader = self.session.reader(); + + let mut decrypted = Vec::with_capacity(plaintext_len); + let mut i = 0; + while i < plaintext_len { + let decrypted_part = match reader.fill_buf() { + Ok(buf) => buf, + Err(e) => { + // TODO: better error eg expected available plaintext but got error + return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, e))); + } + }; + + let part_len = decrypted_part.len(); + decrypted.extend_from_slice(decrypted_part); + reader.consume(part_len); + i += part_len; + } + + log::debug!("VISION READ: Got {} decrypted bytes", decrypted.len()); + self.handle_padded_bytes(&decrypted)?; + + // If the mode changed, the outer read deframer has been consumed and + // we should return any pending read data and use the new mode's read function + match self.read_mode { + VisionMode::PaddingTls => { + // Need more TCP data, loop back to read + } + VisionMode::Tls => { + if !self.pending_read.is_empty() { + let len = buf.remaining().min(self.pending_read.len()); + buf.put_slice(&self.pending_read[..len]); + self.pending_read.advance(len); + return Poll::Ready(Ok(())); + } + return self.poll_read_tls(cx, buf); + } + VisionMode::Direct => { + if !self.pending_read.is_empty() { + let len = buf.remaining().min(self.pending_read.len()); + buf.put_slice(&self.pending_read[..len]); + self.pending_read.advance(len); + return Poll::Ready(Ok(())); + } + return Pin::new(&mut self.tcp).poll_read(cx, buf); + } + } + } + + // Finished iterating through all TLS records, return data if we have any, else loop + // back to read more TCP data + if !self.pending_read.is_empty() { + let len = buf.remaining().min(self.pending_read.len()); + buf.put_slice(&self.pending_read[..len]); + self.pending_read.advance(len); + return Poll::Ready(Ok(())); + } + } + } + + fn handle_padded_bytes(&mut self, decrypted: &[u8]) -> std::io::Result<()> { + // Feed newly decrypted bytes to the unpadder. + // All data should be consumed from the rustls session, as any actual content + // will eventually be returned by the unpadder. + // The unpadder maintains internal state for partial parsing + let UnpadResult { + content: mut unpadded, + command: maybe_command, + } = self.read_unpadder.unpad(decrypted)?; + + log::debug!("VISION READ: Unpadded to {} bytes", unpadded.len()); + + if !unpadded.is_empty() && self.filter.is_filtering() { + // Feed to deframer + self.inner_read_deframer.feed(&unpadded); + + // Process all complete TLS packets + loop { + match self.inner_read_deframer.next_record() { + Ok(DeframeResult::TlsRecord(record)) => { + self.filter.filter_record(&record); + if !self.filter.is_filtering() { + break; + } + } + Ok(DeframeResult::UnknownPrefix(prefix)) => { + // Skip unknown prefix bytes (e.g., VLESS headers from proxy chain) + if prefix.is_empty() { + return Err(io::Error::other("FuzzyTlsDeframer returned empty prefix")); + } + if prefix.len() > 512 { + log::warn!( + "VISION READ: Unusually large prefix discarded: {} bytes", + prefix.len() + ); + } + log::debug!("VISION READ: Skipped {} byte prefix", prefix.len()); + + // Decrement once for this chunk + self.filter.decrement_filter_count(); + + // Continue processing + } + Ok(DeframeResult::NeedData) => break, // Need more data + Err(e) => { + // Invalid TLS packet - stop filtering + // This only occurs if we've already seen valid TLS records, and + // then encountered invalid ones. + log::error!("VISION READ: Read invalid TLS data after valid records - stopping filtering: {}", e); + self.filter + .stop_filtering("read invalid TLS data".to_string()); + break; + } + } + } + } + + match maybe_command { + Some(UnpadCommand::Direct) => { + log::debug!( + "VISION READ: Received DIRECT command with {} bytes", + unpadded.len() + ); + + // The unpadded data already contains: + // 1. The Direct command content + // 2. Any remaining bytes after the padding (already appended by unpadder) + + // Extract remaining data from outer read deframer (raw TCP data in Direct mode) + let remaining = self + .outer_read_deframer + .take() + .expect("outer_read_deframer must exist") + .into_remaining_data(); + + if !remaining.is_empty() { + log::debug!( + "VISION READ: Extracted {} bytes from deframer as raw data", + remaining.len() + ); + unpadded.extend_from_slice(&remaining); + } + + self.pending_read.extend_from_slice(&unpadded); + + self.switch_read_to_direct_mode()?; + } + Some(UnpadCommand::End) => { + log::debug!("VISION READ: Received END command, switching to Tls mode"); + + // The `unpadded` data is already decrypted, store it in pending_read to be + // returned + if !unpadded.is_empty() { + log::debug!( + "VISION READ: Storing {} decrypted bytes in pending_read (End command - already decrypted)", + unpadded.len() + ); + self.pending_read.extend_from_slice(&unpadded); + } + + // Process the remaining data from the outer read deframer. + // Previously we called into_remaining_data and fed it all at once to + // feed_and_process_crypto_connection, but if there was too much outstanding + // data, we hit "message buffer full" at + // https://github.com/rustls/rustls/blob/58fbe9e7ad91951a7df148c2854f57db728a717c/rustls/src/msgs/deframer/buffers.rs#L235 + + let mut outer_read_deframer = self + .outer_read_deframer + .take() + .expect("outer_read_deframer must exist"); + + while let Some(record) = outer_read_deframer.next_record()? { + // Feed the encrypted TLS records to the session + let plaintext_len = + feed_and_process_crypto_connection(&mut self.session, &record)?; + + if plaintext_len > 0 { + // Extract the decrypted data + let mut reader = self.session.reader(); + + // Read all available decrypted data, possibly in multiple chunks + let mut total_read = 0; + while total_read < plaintext_len { + let chunk = match reader.fill_buf() { + Ok(buf) => buf, + Err(e) => { + return Err(io::Error::new(io::ErrorKind::InvalidData, e)); + } + }; + + let chunk_len = chunk.len(); + self.pending_read.extend_from_slice(chunk); + reader.consume(chunk_len); + total_read += chunk_len; + } + + log::debug!( + "VISION READ: Decrypted {} bytes from deframer remaining data (End command)", + total_read + ); + } + } + + let remaining = outer_read_deframer.into_remaining_data(); + + // The `remaining` data from the deframer is ENCRYPTED TLS record bytes + // We need to feed it to the TLS session for decryption, not append to unpadded + if !remaining.is_empty() { + // This should only be a partial TLS record so there should be no new plaintext + // to read + let plaintext_len = + feed_and_process_crypto_connection(&mut self.session, &remaining)?; + assert!(plaintext_len == 0); + } + + // Switch to Tls mode - End command means: stop padding, continue normal relay through outer TLS + self.switch_read_to_tls_mode()?; + } + Some(UnpadCommand::Continue) => { + self.pending_read.extend_from_slice(&unpadded); + } + None => { + self.pending_read.extend_from_slice(&unpadded); + } + } + Ok(()) + } + + fn poll_read_tls( + &mut self, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + // Check if rustls buffer has decrypted data available first + { + let mut reader = self.session.reader(); + match reader.fill_buf() { + Ok(available) if !available.is_empty() => { + // Copy directly from rustls buffer to user buffer (single copy) + let len = buf.remaining().min(available.len()); + buf.put_slice(&available[..len]); + reader.consume(len); + // If we didn't consume all data, it stays in rustls buffer for next read + return Poll::Ready(Ok(())); + } + Ok(_) => { + // Empty buffer, fall through to slow path + } + Err(e) if e.kind() == io::ErrorKind::WouldBlock => { + // No data available, fall through to slow path + } + Err(e) => return Poll::Ready(Err(e)), + } + } + + if self.is_read_eof { + return Poll::Ready(Ok(())); + } + + // No data in rustls buffer, need to read from TCP + // To prevent "message buffer full", we use a SyncReadAdapter to allow + // rustls to decide how much to read from TCP, versus reading to a buffer + // and then feeding like in PaddingTls mode. + // + // In read_tls, rustls will internally: + // 1. Check if it has buffer space via prepare_read() + // 2. Read only up to 4KB at a time + // 3. Prevent "message buffer full" by reading less + let mut reader = SyncReadAdapter { + io: &mut self.tcp, + cx, + }; + + loop { + match self.session.read_tls(&mut reader) { + Ok(n) => { + if n == 0 { + log::debug!("VISION READ: TCP connection EOF in Tls mode"); + self.is_read_eof = true; + return Poll::Ready(Ok(())); + } + + log::debug!("VISION READ: Read {} bytes from TCP via rustls", n); + + // Process the encrypted data + let plaintext_len = self + .session + .process_new_packets() + .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))? + .plaintext_bytes_to_read(); + + if plaintext_len == 0 { + // No plaintext yet, need more TCP data + continue; + } + + // Extract plaintext from rustls and return it + let mut reader = self.session.reader(); + match reader.fill_buf() { + Ok(available) => { + // TODO: this should probably be a debug_assert! + if available.is_empty() { + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "Read zero bytes when plaintext is available", + ))); + } + let len = buf.remaining().min(available.len()); + buf.put_slice(&available[..len]); + reader.consume(len); + return Poll::Ready(Ok(())); + } + Err(e) => { + return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, e))); + } + } + } + Err(e) if e.kind() == io::ErrorKind::WouldBlock => { + return Poll::Pending; + } + Err(e) => { + return Poll::Ready(Err(e)); + } + } + } + } + + fn poll_write_padding_tls( + &mut self, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + if self.vless_response_to_send { + log::debug!("VISION WRITE: Sending VLESS response to TLS session"); + + // VLESS response header: [version=0, addon_length=0] + const VLESS_RESPONSE: [u8; 2] = [0, 0]; + + // Write VLESS response to TLS session (will be encrypted) + self.write_to_session(&VLESS_RESPONSE)?; + + // Clear flag so we don't send it again + self.vless_response_to_send = false; + } + + // Drain all pending writes - this must be done before mode switch + // as the session could contain the DIRECT/END packet that triggered the mode switch flag, + // and because we no longer use `pending_plain_writes` afterwards. + if !self.pending_plain_writes.is_empty() || self.session.wants_write() { + ready!(self.drain_all_writes_padding(cx))?; + } + + if self.pending_direct_mode_switch { + log::debug!("VISION WRITE: Switching to direct mode (flag set by previous write)"); + self.pending_direct_mode_switch = false; + self.switch_write_to_direct_mode()?; + return Pin::new(&mut self.tcp).poll_write(cx, buf); + } + + if self.pending_tls_mode_switch { + log::debug!("VISION WRITE: Switching to Tls mode (flag set by previous write)"); + self.pending_tls_mode_switch = false; + self.switch_write_to_tls_mode()?; + return self.poll_write_tls(cx, buf); + } + + // Feed write buffer to inner deframer and check for ApplicationData + let existing_inner_len = self.inner_write_deframer.pending_bytes(); + self.inner_write_deframer.feed(buf); + + // Process all complete TLS records in the buffer + let mut processed_len = 0; + loop { + match self.inner_write_deframer.next_record() { + Ok(DeframeResult::TlsRecord(record)) => { + processed_len += record.len(); + + // Feed to filter for TLS pattern detection + self.filter.filter_record(&record); + + // Check if this packet is ApplicationData to switch to Direct mode + let is_app_data = self.filter.is_tls() && record.len() >= 3 + && record[0] == 0x17 // ApplicationData + && record[1] == 0x03; + + // Check if filtering ended and we are not TLS 1.2 or above + let non_tls_filtering_ended = !is_app_data + && !self.filter.is_filtering() + && !self.filter.is_tls12_or_above(); + + if is_app_data || non_tls_filtering_ended { + if is_app_data { + log::debug!( + "VISION WRITE: Detected ApplicationData in {} byte packet", + record.len() + ); + } else { + log::debug!("VISION WRITE: Filtering ended, not TLS 1.2 or above"); + } + + // DIRECT or END command + // Toggle mode switch flag, if `non_tls_filtering_ended`, it never supports XTLS + let command = if self.filter.supports_xtls() { + self.pending_direct_mode_switch = true; + 0x02 + } else { + self.pending_tls_mode_switch = true; + 0x01 + }; + + let final_padded_packet = if self.write_first_packet { + self.write_first_packet = false; + super::vision_pad::pad_with_uuid_and_command( + &record, + &self.user_uuid, + command, + true, // is_tls + ) + } else { + super::vision_pad::pad_with_command( + &record, command, true, // is_tls + ) + }; + + self.write_to_session(&final_padded_packet)?; + + // this must be true because else we would have a successful next_record call on previous iteration + assert!(processed_len > existing_inner_len); + + // Drain and handle result + match self.drain_all_writes_padding(cx) { + Poll::Ready(Ok(())) => { + // Fully drained, continue + } + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + Poll::Pending => { + // Still draining + } + } + + // Clear deframer since caller will re-feed remaining data + self.inner_write_deframer.clear(); + + // tell caller to resend the next packets after the mode switch + return Poll::Ready(Ok(processed_len - existing_inner_len)); + } + + // if we got here, it's: + // - a non-TLS packet of data and we're still filtering + // - a TLS record that is not app data, and we have not yet seen app data, even + // though we might be done filtering + // .. so we need to continue padding + // ref: https://github.com/XTLS/Xray-core/blob/9f5dcb15910aadc7ef450514747576827a389853/proxy/proxy.go#L371 + + let padded_packet = if self.write_first_packet { + self.write_first_packet = false; + super::vision_pad::pad_with_uuid_and_command( + &record, + &self.user_uuid, + 0x00, // CONTINUE + self.filter.is_tls(), + ) + } else { + super::vision_pad::pad_with_command( + &record, + 0x00, // CONTINUE + self.filter.is_tls(), + ) + }; + + self.write_to_session(&padded_packet)?; + match self.drain_all_writes_padding(cx) { + Poll::Pending => { + let unprocessed_buf_len = + buf.len() - (processed_len - existing_inner_len); + // sanity check + assert!( + self.inner_write_deframer.pending_bytes() == unprocessed_buf_len + ); + // clear since the user will re-feed + self.inner_write_deframer.clear(); + return Poll::Ready(Ok(processed_len - existing_inner_len)); + } + Poll::Ready(Ok(())) => { + // continue since it drained + } + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + } + } + Ok(DeframeResult::NeedData) => { + // need more data + // TODO: investigate this case. it's possible that this is not TLS data and the deframer is waiting + // for more writes from the caller to complete a full TLS record, but no more writes are coming. + // do we need to change our deframer approach? + break; + } + Ok(DeframeResult::UnknownPrefix(prefix)) => { + // TODO: check lengths and see if we need to split into multiple packets, see + // https://github.com/XTLS/Xray-core/blob/9f5dcb15910aadc7ef450514747576827a389853/proxy/proxy.go#L390 + + processed_len += prefix.len(); + + self.filter.decrement_filter_count(); + + if self.filter.is_filtering() { + // We don't assume the deframer won't return any more packets, + // so we continue here, and stop if we're no longer filtering. + let padded_packet = if self.write_first_packet { + self.write_first_packet = false; + super::vision_pad::pad_with_uuid_and_command( + &prefix, + &self.user_uuid, + 0x0, + self.filter.is_tls(), + ) + } else { + super::vision_pad::pad_with_command(&prefix, 0x0, self.filter.is_tls()) + }; + + self.write_to_session(&padded_packet)?; + match self.drain_all_writes_padding(cx) { + Poll::Pending => {} + Poll::Ready(Ok(())) => {} + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + } + + // Continue processing + } else { + // Sending end command to switch to TLS, there are no more padded packets. + self.pending_tls_mode_switch = true; + + let padded_packet = if self.write_first_packet { + self.write_first_packet = false; + super::vision_pad::pad_with_uuid_and_command( + &prefix, + &self.user_uuid, + 0x1, + self.filter.is_tls(), + ) + } else { + super::vision_pad::pad_with_command(&prefix, 0x1, self.filter.is_tls()) + }; + + // Clear deframe, not really necessary since deallocate will occur and this + // will never be used again. + self.inner_write_deframer.clear(); + + self.write_to_session(&padded_packet)?; + // Drain and handle result + match self.drain_all_writes_padding(cx) { + Poll::Ready(Ok(())) => { + // Fully drained + } + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + Poll::Pending => { + // Still draining + } + } + + // tell caller to resend the next packets after the mode switch + return Poll::Ready(Ok(processed_len - existing_inner_len)); + } + } + Err(e) => { + // Deframing failed, this error means we've already seen a valid TLS record on + // the write side and then invalid TLS data is encountered. + // Switch immediately to TLS mode and send the invalid data as a single packet. + // TODO: check lengths and see if we need to split into multiple packets, see + // https://github.com/XTLS/Xray-core/blob/9f5dcb15910aadc7ef450514747576827a389853/proxy/proxy.go#L390 + // + log::error!( + "VISION WRITE: Deframing failed, invalid data after valid records - stopping filtering: {}", + e + ); + + self.filter + .stop_filtering("write invalid TLS data".to_string()); + + let remaining_data = self.inner_write_deframer.remaining_data(); + + let padded_packet = if self.write_first_packet { + self.write_first_packet = false; + super::vision_pad::pad_with_uuid_and_command( + remaining_data, + &self.user_uuid, + 0x1, + self.filter.is_tls(), + ) + } else { + super::vision_pad::pad_with_command( + remaining_data, + 0x1, + self.filter.is_tls(), + ) + }; + + self.write_to_session(&padded_packet)?; + + match self.drain_all_writes_padding(cx) { + Poll::Pending => {} + Poll::Ready(Ok(())) => {} + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + } + + return Poll::Ready(Ok(buf.len())); + } + } + } + + // we processed all records if we got here + Poll::Ready(Ok(buf.len())) + } +} + +// AsyncRead implementation +impl AsyncRead for VisionStream +where + IO: AsyncRead + AsyncWrite + Unpin, +{ + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + let this = self.get_mut(); + + // If we have leftover buffered data, return it first + if !this.pending_read.is_empty() { + let len = buf.remaining().min(this.pending_read.len()); + buf.put_slice(&this.pending_read[..len]); + this.pending_read.advance(len); + return Poll::Ready(Ok(())); + } + + match this.read_mode { + VisionMode::PaddingTls => { + if this.vless_response_pending { + let decrypted_data = ready!(this.poll_read_vless_response(cx))?; + this.vless_response_pending = false; + if !decrypted_data.is_empty() { + this.handle_padded_bytes(&decrypted_data)?; + if !this.pending_read.is_empty() { + let len = buf.remaining().min(this.pending_read.len()); + buf.put_slice(&this.pending_read[..len]); + this.pending_read.advance(len); + return Poll::Ready(Ok(())); + } + match this.read_mode { + VisionMode::PaddingTls => { + // Fall-through + } + VisionMode::Tls => { + return this.poll_read_tls(cx, buf); + } + VisionMode::Direct => { + return Pin::new(&mut this.tcp).poll_read(cx, buf); + } + } + } + } + this.poll_read_padding_tls(cx, buf) + } + VisionMode::Tls => this.poll_read_tls(cx, buf), + VisionMode::Direct => Pin::new(&mut this.tcp).poll_read(cx, buf), + } + } +} + +// AsyncWrite implementation +impl AsyncWrite for VisionStream +where + IO: AsyncRead + AsyncWrite + Unpin, +{ + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let this = self.get_mut(); + + match this.write_mode { + VisionMode::PaddingTls => this.poll_write_padding_tls(cx, buf), + VisionMode::Tls => this.poll_write_tls(cx, buf), + VisionMode::Direct => Pin::new(&mut this.tcp).poll_write(cx, buf), + } + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + + match this.write_mode { + VisionMode::PaddingTls => { + ready!(this.drain_all_writes_padding(cx))?; + Pin::new(&mut this.tcp).poll_flush(cx) + } + VisionMode::Tls => { + ready!(this.drain_all_writes_tls(cx))?; + Pin::new(&mut this.tcp).poll_flush(cx) + } + VisionMode::Direct => Pin::new(&mut this.tcp).poll_flush(cx), + } + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.get_mut().tcp).poll_shutdown(cx) + } +} + +// Implement AsyncPing trait +impl crate::async_stream::AsyncPing for VisionStream +where + IO: AsyncStream, +{ + fn supports_ping(&self) -> bool { + self.tcp.supports_ping() + } + + fn poll_write_ping(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.get_mut().tcp).poll_write_ping(cx) + } +} + +impl AsyncStream for VisionStream where IO: AsyncStream {} diff --git a/src/vless/vision_unpad.rs b/src/vless/vision_unpad.rs new file mode 100644 index 00000000..f63b9853 --- /dev/null +++ b/src/vless/vision_unpad.rs @@ -0,0 +1,547 @@ +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum UnpadCommand { + Continue = 0, + End = 1, + Direct = 2, +} + +impl TryFrom for UnpadCommand { + type Error = std::io::Error; + + fn try_from(value: u8) -> Result { + match value { + 0 => Ok(UnpadCommand::Continue), + 1 => Ok(UnpadCommand::End), + 2 => Ok(UnpadCommand::Direct), + _ => Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("Invalid padding command: {}", value), + )), + } + } +} + +/// Result of a completed padding block +#[derive(Debug, Default)] +pub struct UnpadResult { + /// Content data extracted from the block + pub content: Vec, + /// The command from the block, or None if this is not XTLS data + pub command: Option, +} + +#[derive(Debug, Clone)] +enum UnpadState { + /// Initial state: expecting UUID (16 bytes) before first command + Initial { expected_uuid: [u8; 16] }, + /// Reading the command byte (first of 5 header bytes) + ReadingCommand, + /// Reading content length (2 bytes) + ReadingContentLength { + command: UnpadCommand, + first_byte: Option, + }, + /// Reading padding length (2 bytes) + ReadingPaddingLength { + command: UnpadCommand, + content_len: u16, + first_byte: Option, + }, + /// Reading content data + ReadingContent { + command: UnpadCommand, + partial_content: Vec, + remaining_content_len: u16, + padding_len: u16, + }, + /// Reading (skipping) padding data + ReadingPadding { + command: UnpadCommand, + content: Vec, + remaining_padding_len: u16, + }, + /// Finished processing (End or Direct command completed) + Done, +} + +#[derive(Debug, Clone)] +pub struct VisionUnpadder { + state: UnpadState, + first_block: bool, + /// Reusable buffer for accumulated content to avoid allocation on every unpad() call + accumulated_buffer: Vec, +} + +impl VisionUnpadder { + pub fn new(expected_uuid: [u8; 16]) -> Self { + Self { + state: UnpadState::Initial { expected_uuid }, + first_block: true, + accumulated_buffer: Vec::new(), + } + } + + /// Process input data and extract padding blocks + /// Returns UnpadResult when a complete block is parsed, empty result if more data is needed + /// Accumulates content across Continue commands and returns on End/Direct or when needing more data + pub fn unpad(&mut self, mut data: &[u8]) -> std::io::Result { + // Reuse the accumulated buffer - clear but keep capacity + self.accumulated_buffer.clear(); + log::debug!( + "UNPADDER: unpad() called with {} bytes, state before: {:?}, first_block: {}", + data.len(), + self.state, + self.first_block + ); + + loop { + match &mut self.state { + UnpadState::Initial { expected_uuid } => { + if data.len() < 16 { + return Ok(UnpadResult::default()); // Need more data + } + + // Check if this is XTLS padding data + if &data[..16] != expected_uuid { + // Not XTLS padding data + return Ok(UnpadResult { + content: data.to_vec(), + command: None, + }); + } + + // Consume UUID + data = &data[16..]; + self.state = UnpadState::ReadingCommand; + } + + UnpadState::ReadingCommand => { + if data.is_empty() { + // Out of data + if !self.first_block { + // We've completed at least one block, return what we have + return Ok(UnpadResult { + content: std::mem::take(&mut self.accumulated_buffer), + command: Some(UnpadCommand::Continue), + }); + } + debug_assert!(self.accumulated_buffer.is_empty(), "self.accumulated_buffer should be empty when returning default in ReadingCommand"); + return Ok(UnpadResult::default()); + } + let command_byte = data[0]; + data = &data[1..]; + + let command = UnpadCommand::try_from(command_byte)?; + + self.state = UnpadState::ReadingContentLength { + command, + first_byte: None, + }; + } + + UnpadState::ReadingContentLength { + command, + first_byte, + } => { + if data.is_empty() { + // Out of data - return accumulated content if any + if !self.first_block { + return Ok(UnpadResult { + content: std::mem::take(&mut self.accumulated_buffer), + command: Some(UnpadCommand::Continue), + }); + } + debug_assert!(self.accumulated_buffer.is_empty(), "self.accumulated_buffer should be empty when returning default in ReadingContentLength"); + return Ok(UnpadResult::default()); + } + match first_byte { + None => { + *first_byte = Some(data[0]); + data = &data[1..]; + } + Some(high_byte) => { + let low_byte = data[0]; + data = &data[1..]; + let content_len = ((*high_byte as u16) << 8) | (low_byte as u16); + + self.state = UnpadState::ReadingPaddingLength { + command: *command, + content_len, + first_byte: None, + }; + } + } + } + + UnpadState::ReadingPaddingLength { + command, + content_len, + first_byte, + } => { + if data.is_empty() { + // Out of data - return accumulated content if any + if !self.first_block { + return Ok(UnpadResult { + content: std::mem::take(&mut self.accumulated_buffer), + command: Some(UnpadCommand::Continue), + }); + } + debug_assert!(self.accumulated_buffer.is_empty(), "self.accumulated_buffer should be empty when returning default in ReadingPaddingLength"); + return Ok(UnpadResult::default()); + } + match first_byte { + None => { + *first_byte = Some(data[0]); + data = &data[1..]; + } + Some(high_byte) => { + let low_byte = data[0]; + data = &data[1..]; + let padding_len = ((*high_byte as u16) << 8) | (low_byte as u16); + log::debug!("UNPADDER: Parsed header - command={:?}, content_len={}, padding_len={}", *command, *content_len, padding_len); + // Pre-allocate based on min of content_len and available data + // (we eagerly return when data runs out, so we may not need full content_len) + let prealloc_size = (*content_len as usize).min(data.len()); + self.state = UnpadState::ReadingContent { + command: *command, + partial_content: Vec::with_capacity(prealloc_size), + remaining_content_len: *content_len, + padding_len, + }; + } + } + } + + UnpadState::ReadingContent { + command, + partial_content, + remaining_content_len, + padding_len, + } => { + if *remaining_content_len > 0 { + if data.is_empty() { + // Out of data - return accumulated content if any. + // We need to return everything we have including the partial content + // else some clients will not return - see below. + self.accumulated_buffer.append(partial_content); + if !self.first_block { + return Ok(UnpadResult { + content: std::mem::take(&mut self.accumulated_buffer), + command: Some(UnpadCommand::Continue), + }); + } else { + return Ok(UnpadResult { + content: std::mem::take(&mut self.accumulated_buffer), + command: None, + }); + } + } + let to_read = (*remaining_content_len as usize).min(data.len()); + partial_content.extend_from_slice(&data[..to_read]); + data = &data[to_read..]; + *remaining_content_len -= to_read as u16; + } + + if *remaining_content_len == 0 { + log::debug!("UNPADDER: Content read complete ({} bytes), accumulated at {}, transitioning to padding", + partial_content.len(), self.accumulated_buffer.len()); + let content = std::mem::take(partial_content); + self.state = UnpadState::ReadingPadding { + command: *command, + content, + remaining_padding_len: *padding_len, + }; + } + } + + UnpadState::ReadingPadding { + command, + content, + remaining_padding_len, + } => { + if *remaining_padding_len > 0 { + if data.is_empty() { + // We need to return _all_ the content including from this incomplete + // block because some clients will not write the remaining padding until we've + // generated a response to the current content. Waiting for the remaining + // padding before returning the current block's content caused a + // deadlock. This is also the behavior on the official client: + // https://github.com/XTLS/Xray-core/blob/9f5dcb15910aadc7ef450514747576827a389853/proxy/proxy.go#L518 + // + // It's important that we do not return the command until the rest of + // the padding has been consumed so that unpad() consumes all of the + // padded data before we switch modes due to a Direct or End command. + // + // Note that append moves all of `content` and leaves it empty. + self.accumulated_buffer.append(content); + + log::debug!("UNPADDER: Out of data in ReadingPadding with {} padding bytes remaining, returning {} bytes (first_block={})", + *remaining_padding_len, self.accumulated_buffer.len(), self.first_block); + + if !self.first_block { + return Ok(UnpadResult { + content: std::mem::take(&mut self.accumulated_buffer), + command: Some(UnpadCommand::Continue), + }); + } else { + return Ok(UnpadResult { + content: std::mem::take(&mut self.accumulated_buffer), + command: None, + }); + } + } + let to_skip = (*remaining_padding_len as usize).min(data.len()); + data = &data[to_skip..]; + *remaining_padding_len -= to_skip as u16; + } + + if *remaining_padding_len == 0 { + self.first_block = false; + + // Block complete (all padding skipped) + self.accumulated_buffer.append(content); + + log::debug!("UNPADDER: Padding complete for command {:?}, {} accumulated bytes total", *command, self.accumulated_buffer.len()); + match *command { + UnpadCommand::Continue => { + // Continue to next block + log::debug!("UNPADDER: Continue command - transitioning to ReadingCommand for next block"); + self.state = UnpadState::ReadingCommand; + } + end_or_direct_command => { + // Return all accumulated content with final command + // Any remaining data in the input slice is unconsumed (comes after the final padding) + // Append it to content so caller gets everything in one place + let remaining_len = data.len(); + self.accumulated_buffer.extend_from_slice(data); + log::debug!("UNPADDER: {:?} command - transitioning to Done state, returning {} bytes content ({} from padding, {} remaining after)", + end_or_direct_command, self.accumulated_buffer.len(), self.accumulated_buffer.len() - remaining_len, remaining_len); + self.state = UnpadState::Done; + return Ok(UnpadResult { + content: std::mem::take(&mut self.accumulated_buffer), + command: Some(end_or_direct_command), + }); + } + } + } + } + + UnpadState::Done => { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "Already in UnpadState::Done", + )); + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_padding_state_basic() { + let user_uuid = [0u8; 16]; + let mut state = VisionUnpadder::new(user_uuid); + + let data = vec![ + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, // UUID + 1, // command (End) + 0, 5, // content length = 5 + 0, 3, // padding length = 3 + 1, 2, 3, 4, 5, // content + 0, 0, 0, // padding + ]; + + let result = state.unpad(&data[..]).unwrap(); + assert_eq!(result.content, vec![1, 2, 3, 4, 5]); + assert_eq!(result.command, Some(UnpadCommand::End)); + assert!(matches!(state.state, UnpadState::Done)); + } + + #[test] + fn test_padding_command_continue() { + let user_uuid = [1u8; 16]; + let mut state = VisionUnpadder::new(user_uuid); + + let mut data = vec![1u8; 16]; // UUID + data.extend_from_slice(&[ + 0, // command (Continue) + 0, 3, // content length = 3 + 0, 2, // padding length = 2 + 10, 11, 12, // content + 0, 0, // padding + 0, // command (Continue) + 0, 2, // content length = 2 + 0, 1, // padding length = 1 + 20, 21, // content + 0, // padding + ]); + + // Now it processes both Continue blocks and accumulates content + let result = state.unpad(&data[..]).unwrap(); + assert_eq!(result.content, vec![10, 11, 12, 20, 21]); + assert_eq!(result.command, Some(UnpadCommand::Continue)); + assert!(!matches!(state.state, UnpadState::Done)); + } + + #[test] + fn test_non_xtls_data() { + let user_uuid = [0u8; 16]; + let mut state = VisionUnpadder::new(user_uuid); + + // Data that's less than 16 bytes (too short to check UUID) + let short_data = vec![1, 2, 3, 4, 5]; + let result = state.unpad(&short_data[..]).unwrap(); + assert!(result.content.is_empty()); // Need more data + assert!(result.command.is_none()); + + // Data that's exactly 16 bytes but doesn't match UUID + let non_xtls_data = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]; + let result = state.unpad(&non_xtls_data[..]).unwrap(); + assert_eq!(result.content, non_xtls_data); + assert_eq!(result.command, None); + // State remains Initial (doesn't transition to Done for non-XTLS data) + assert!(!matches!(state.state, UnpadState::Done)); + } + + #[test] + fn test_direct_command() { + let user_uuid = [2u8; 16]; + let mut state = VisionUnpadder::new(user_uuid); + + let mut data = vec![2u8; 16]; // UUID + data.extend_from_slice(&[ + 2, // command (Direct) + 0, 3, // content length = 3 + 0, 2, // padding length = 2 + 10, 11, 12, // content + 0, 0, // padding + 99, 98, 97, // remaining data after Direct + ]); + + // Process the padded block + let result = state.unpad(&data[..]).unwrap(); + // Direct command appends remaining data (99, 98, 97) to content + assert_eq!(result.content, vec![10, 11, 12, 99, 98, 97]); + assert_eq!(result.command, Some(UnpadCommand::Direct)); + assert!(matches!(state.state, UnpadState::Done)) + } + + #[test] + fn test_incremental_processing() { + let user_uuid = [3u8; 16]; + let mut state = VisionUnpadder::new(user_uuid); + + // First chunk: UUID + partial command + let chunk1 = { + let mut d = vec![3u8; 16]; + d.extend_from_slice(&[1, 0, 3]); // command (End), content_len high byte, partial + d + }; + + let result1 = state.unpad(&chunk1[..]).unwrap(); + assert!(result1.content.is_empty()); // Not enough data yet + assert!(result1.command.is_none()); + assert!(!matches!(state.state, UnpadState::Done)); + + // Second chunk: rest of command + content + padding + let chunk2 = vec![ + 0, 2, // rest of content length (=3), padding length = 2 + 10, 11, 12, // content + 0, 0, // padding + ]; + + let result2 = state.unpad(&chunk2[..]).unwrap(); + assert_eq!(result2.content, vec![10, 11, 12]); + assert_eq!(result2.command, Some(UnpadCommand::End)); + assert!(matches!(state.state, UnpadState::Done)); + } + + #[test] + fn test_empty_content_block() { + let user_uuid = [4u8; 16]; + let mut state = VisionUnpadder::new(user_uuid); + + let mut data = vec![4u8; 16]; // UUID + data.extend_from_slice(&[ + 0, // command (Continue) + 0, 0, // content length = 0 + 0, 5, // padding length = 5 + 0, 0, 0, 0, 0, // padding + ]); + + let result = state.unpad(&data[..]).unwrap(); + assert_eq!(result.content, Vec::::new()); + assert_eq!(result.command, Some(UnpadCommand::Continue)); + assert!(!matches!(state.state, UnpadState::Done)); + } + + #[test] + fn test_multiple_continue_with_end() { + let user_uuid = [5u8; 16]; + let mut state = VisionUnpadder::new(user_uuid); + + let mut data = vec![5u8; 16]; // UUID + data.extend_from_slice(&[ + 0, // command (Continue) + 0, 2, // content length = 2 + 0, 1, // padding length = 1 + 10, 11, // content + 0, // padding + 0, // command (Continue) + 0, 2, // content length = 2 + 0, 1, // padding length = 1 + 20, 21, // content + 0, // padding + 1, // command (End) + 0, 2, // content length = 2 + 0, 1, // padding length = 1 + 30, 31, // content + 0, // padding + ]); + + // Process all blocks - should accumulate and return with End + let result = state.unpad(&data[..]).unwrap(); + assert_eq!(result.content, vec![10, 11, 20, 21, 30, 31]); + assert_eq!(result.command, Some(UnpadCommand::End)); + assert!(matches!(state.state, UnpadState::Done)); + } + + #[test] + fn test_incomplete_padding_returns_content_to_prevent_deadlock() { + let user_uuid = [6u8; 16]; + let mut state = VisionUnpadder::new(user_uuid); + + // First chunk: UUID + command + header + content + partial padding + let mut chunk1 = vec![6u8; 16]; // UUID + chunk1.extend_from_slice(&[ + 1, // command (End) + 0, 3, // content length = 3 + 0, 5, // padding length = 5 + 10, 11, 12, // content + 0, 0, // only 2 bytes of padding (need 5 total) + ]); + + // Should return content immediately to prevent deadlock + // (some clients don't send remaining padding until they get a response) + // However, command should be None because padding is incomplete + let result1 = state.unpad(&chunk1[..]).unwrap(); + assert_eq!(result1.content, vec![10, 11, 12]); + assert_eq!(result1.command, None); // Command not returned until padding complete + assert!(matches!(state.state, UnpadState::ReadingPadding { .. })); + + // Second chunk: remaining padding + let chunk2 = vec![ + 0, 0, 0, // remaining 3 bytes of padding + ]; + + // Now padding is complete, should return End command (no more content since it was already returned) + let result2 = state.unpad(&chunk2[..]).unwrap(); + assert!(result2.content.is_empty()); // Content already returned in result1 + assert_eq!(result2.command, Some(UnpadCommand::End)); + assert!(matches!(state.state, UnpadState::Done)); + } +} diff --git a/src/vless/vless_client_handler.rs b/src/vless/vless_client_handler.rs new file mode 100644 index 00000000..43d7ec5b --- /dev/null +++ b/src/vless/vless_client_handler.rs @@ -0,0 +1,151 @@ +use async_trait::async_trait; +use tokio::io::AsyncWriteExt; + +use crate::address::{Address, NetLocation}; +use crate::async_stream::AsyncStream; +use crate::crypto::CryptoTlsStream; +use crate::tcp_handler::{TcpClientHandler, TcpClientSetupResult}; +use crate::util::{allocate_vec, parse_uuid, write_all}; + +use super::vision_stream::VisionStream; +use super::vless_response_stream::VlessResponseStream; +use super::vless_util::{vision_flow_addon_data, COMMAND_TCP}; + +pub struct VlessTcpClientHandler { + user_id: Box<[u8]>, +} + +impl std::fmt::Debug for VlessTcpClientHandler { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("VlessTcpClientHandler") + .field("user_id", &self.user_id) + .finish() + } +} + +impl VlessTcpClientHandler { + pub fn new(user_id: &str) -> Self { + Self { + user_id: parse_uuid(user_id).unwrap().into_boxed_slice(), + } + } +} + +#[async_trait] +impl TcpClientHandler for VlessTcpClientHandler { + async fn setup_client_stream( + &self, + _server_stream: &mut Box, + mut client_stream: Box, + remote_location: NetLocation, + ) -> std::io::Result { + write_vless_header(&mut client_stream, &self.user_id, &[], &remote_location).await?; + client_stream.flush().await?; + + // Wrap stream to read VLESS response on first read + let client_stream = Box::new(VlessResponseStream::new(client_stream)); + + Ok(TcpClientSetupResult { client_stream }) + } +} + +// REMOVED: Old tokio-rustls vision function - now using setup_custom_tls_vision_vless_client_stream + +pub async fn setup_custom_tls_vision_vless_client_stream( + mut tls_stream: CryptoTlsStream, + user_id: &[u8], + remote_location: &NetLocation, +) -> std::io::Result +where + IO: crate::async_stream::AsyncStream + 'static, +{ + // Write VLESS request header through TLS + write_vless_header( + &mut tls_stream, + user_id, + vision_flow_addon_data(), + remote_location, + ) + .await?; + tls_stream.flush().await?; + + // Create stream with VLESS response reading support + // CryptoTlsStream already wraps a Connection, so we don't need to wrap again + let (io, connection) = tls_stream.into_inner(); + + // VisionStream will read VLESS response header on first read + let mut user_uuid = [0u8; 16]; + user_uuid.copy_from_slice(user_id); + let vision_stream = VisionStream::new_client(io, connection, user_uuid); + + Ok(TcpClientSetupResult { + client_stream: Box::new(vision_stream), + }) +} + +async fn write_vless_header( + stream: &mut S, + user_id: &[u8], + addon_data: &[u8], + remote_location: &NetLocation, +) -> std::io::Result<()> { + // VLESS header format: + // version (1 byte) + user_id (16 bytes) + addon_length (1 byte) + addon_data + command (1 byte) + port (2 bytes) + address_type (1 byte) + address + + // Calculate base header size: version + user_id + addon_length + addon_data + command + port + address_type + let base_header_size = 1 + 16 + 1 + addon_data.len() + 1 + 2 + 1; + let mut header_bytes = allocate_vec(base_header_size); + + // version 0, we need to write since it's uninitialized + header_bytes[0] = 0; + // Copy user_id + header_bytes[1..17].copy_from_slice(user_id); + + // addon length + header_bytes[17] = addon_data.len() as u8; + + // Copy addon data if present + if !addon_data.is_empty() { + header_bytes[18..18 + addon_data.len()].copy_from_slice(addon_data); + } + + let addon_end = 18 + addon_data.len(); + + // command (1 = tcp) + header_bytes[addon_end] = COMMAND_TCP; + + // port (2 bytes, big-endian) + let remote_port = remote_location.port(); + header_bytes[addon_end + 1] = (remote_port >> 8) as u8; + header_bytes[addon_end + 2] = (remote_port & 0xff) as u8; + + // address_type + let address_type_offset = addon_end + 3; + + match remote_location.address() { + Address::Ipv4(v4addr) => { + header_bytes[address_type_offset] = 1; + header_bytes.extend_from_slice(&v4addr.octets()); + } + Address::Ipv6(v6addr) => { + header_bytes[address_type_offset] = 3; + header_bytes.extend_from_slice(&v6addr.octets()); + } + Address::Hostname(hostname) => { + if hostname.len() > 255 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("Hostname is too long: {hostname}"), + )); + } + + header_bytes[address_type_offset] = 2; + header_bytes.push(hostname.len() as u8); + header_bytes.extend_from_slice(hostname.as_bytes()); + } + } + + write_all(stream, &header_bytes).await?; + + Ok(()) +} diff --git a/src/vless_message_stream.rs b/src/vless/vless_message_stream.rs similarity index 100% rename from src/vless_message_stream.rs rename to src/vless/vless_message_stream.rs diff --git a/src/vless/vless_response_stream.rs b/src/vless/vless_response_stream.rs new file mode 100644 index 00000000..f806b484 --- /dev/null +++ b/src/vless/vless_response_stream.rs @@ -0,0 +1,137 @@ +use std::pin::Pin; +use std::task::{Context, Poll}; + +use log::info; +use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; + +use crate::async_stream::AsyncStream; + +/// A wrapper stream that reads VLESS response header on first read, similar +/// to vmess and shadowsocks. +pub struct VlessResponseStream { + inner: Box, + response_pending: bool, + response_buffer: Vec, +} + +impl VlessResponseStream { + pub fn new(inner: Box) -> Self { + Self { + inner, + response_pending: true, + response_buffer: Vec::new(), + } + } + + fn poll_read_response(&mut self, cx: &mut Context<'_>) -> Poll> { + use futures::ready; + + // Need to read at least 2 bytes for the response header + while self.response_buffer.len() < 2 { + let mut buf = [0u8; 2]; + let mut read_buf = ReadBuf::new(&mut buf); + ready!(Pin::new(&mut self.inner).poll_read(cx, &mut read_buf))?; + + let filled = read_buf.filled(); + if filled.is_empty() { + return Poll::Ready(Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + "Connection closed while reading VLESS response", + ))); + } + self.response_buffer.extend_from_slice(filled); + } + + // Validate version + let version = self.response_buffer[0]; + if version != 0 { + return Poll::Ready(Err(std::io::Error::other(format!( + "invalid server protocol version, expected 0, got {version}" + )))); + } + + // Check if we need to read addons + let addon_length = self.response_buffer[1] as usize; + let total_response_len = 2 + addon_length; + + // Read addon data if needed + while self.response_buffer.len() < total_response_len { + let remaining = total_response_len - self.response_buffer.len(); + let mut buf = vec![0u8; remaining]; + let mut read_buf = ReadBuf::new(&mut buf); + ready!(Pin::new(&mut self.inner).poll_read(cx, &mut read_buf))?; + + let filled = read_buf.filled(); + if filled.is_empty() { + return Poll::Ready(Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + "Connection closed while reading VLESS response addons", + ))); + } + self.response_buffer.extend_from_slice(filled); + } + + info!( + "VLESS: Successfully read and consumed {} byte response header (version={}, addon_length={})", + total_response_len, version, addon_length + ); + + Poll::Ready(Ok(())) + } +} + +impl AsyncRead for VlessResponseStream { + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + // Read VLESS response header on first read + if self.response_pending { + match self.poll_read_response(cx) { + Poll::Ready(Ok(())) => { + self.response_pending = false; + // Response is consumed, now read actual data + } + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + Poll::Pending => return Poll::Pending, + } + } + + // Pass through to inner stream + Pin::new(&mut self.inner).poll_read(cx, buf) + } +} + +impl AsyncWrite for VlessResponseStream { + 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> { + 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 crate::async_stream::AsyncPing for VlessResponseStream { + fn supports_ping(&self) -> bool { + self.inner.supports_ping() + } + + fn poll_write_ping( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + Pin::new(&mut self.inner).poll_write_ping(cx) + } +} + +impl AsyncStream for VlessResponseStream {} diff --git a/src/vless/vless_server_handler.rs b/src/vless/vless_server_handler.rs new file mode 100644 index 00000000..5f70c3e2 --- /dev/null +++ b/src/vless/vless_server_handler.rs @@ -0,0 +1,327 @@ +use async_trait::async_trait; +use log::info; +use subtle::ConstantTimeEq; + +use crate::async_stream::AsyncStream; +use crate::crypto::CryptoTlsStream; +use crate::option_util::NoneOrOne; +use crate::stream_reader::StreamReader; +use crate::tcp_handler::{TcpServerHandler, TcpServerSetupResult}; +use crate::util::{parse_uuid, write_all}; +use crate::xudp::XudpMessageStream; + +use super::vision_stream::VisionStream; +use super::vless_message_stream::VlessMessageStream; +use super::vless_util::{ + parse_addons_from_reader, parse_remote_location_from_reader, COMMAND_MUX, COMMAND_TCP, + COMMAND_UDP, XTLS_VISION_FLOW, +}; + +pub struct VlessTcpServerHandler { + user_id: Box<[u8]>, + udp_enabled: bool, +} + +impl std::fmt::Debug for VlessTcpServerHandler { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("VlessTcpServerHandler") + .field("user_id", &self.user_id) + .field("udp_enabled", &self.udp_enabled) + .finish() + } +} + +impl VlessTcpServerHandler { + pub fn new(user_id: &str, udp_enabled: bool) -> Self { + Self { + user_id: parse_uuid(user_id).unwrap().into_boxed_slice(), + udp_enabled, + } + } +} + +const SERVER_RESPONSE_HEADER: &[u8] = &[ + 0u8, // version + 0u8, // addons length +]; + +#[async_trait] +impl TcpServerHandler for VlessTcpServerHandler { + async fn setup_server_stream( + &self, + mut server_stream: Box, + ) -> std::io::Result { + // Parse VLESS header without TLS + let mut stream_reader = StreamReader::new_with_buffer_size(800); + + let client_version = stream_reader.read_u8(&mut server_stream).await?; + if client_version != 0 { + return Err(std::io::Error::other(format!( + "invalid client protocol version, expected 0, got {client_version}" + ))); + } + + let target_id = stream_reader.read_slice(&mut server_stream, 16).await?; + if self.user_id.ct_eq(target_id).unwrap_u8() == 0 { + return Err(std::io::Error::other("Unknown user id")); + } + + let addon_length = stream_reader.read_u8(&mut server_stream).await?; + if addon_length > 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "VLESS addons not supported in current configuration, use TLS protocol for VISION support", + )); + } + + let instruction = stream_reader.read_u8(&mut server_stream).await?; + + match instruction { + COMMAND_TCP => { + let remote_location = + parse_remote_location_from_reader(&mut stream_reader, &mut server_stream) + .await?; + let unparsed_data = stream_reader.unparsed_data(); + + Ok(TcpServerSetupResult::TcpForward { + remote_location, + stream: server_stream, + need_initial_flush: false, + connection_success_response: Some( + SERVER_RESPONSE_HEADER.to_vec().into_boxed_slice(), + ), + initial_remote_data: if unparsed_data.is_empty() { + None + } else { + Some(unparsed_data.to_vec().into_boxed_slice()) + }, + override_proxy_provider: NoneOrOne::Unspecified, + }) + } + COMMAND_UDP => { + if !self.udp_enabled { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "UDP not enabled", + )); + } + + let remote_location = + parse_remote_location_from_reader(&mut stream_reader, &mut server_stream) + .await?; + let unparsed_data = stream_reader.unparsed_data(); + + write_all(&mut server_stream, SERVER_RESPONSE_HEADER).await?; + let mut vless_stream = VlessMessageStream::new(server_stream); + if !unparsed_data.is_empty() { + vless_stream.feed_initial_read_data(unparsed_data)?; + } + + Ok(TcpServerSetupResult::BidirectionalUdp { + remote_location, + stream: Box::new(vless_stream), + need_initial_flush: false, + override_proxy_provider: NoneOrOne::Unspecified, + }) + } + COMMAND_MUX => { + if !self.udp_enabled { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "MUX/XUDP requires UDP to be enabled", + )); + } + + // MUX/XUDP: Destination is NOT in the VLESS header - it comes in XUDP frames + // Don't read destination from wire - protocol spec says it's not present for command 3 + info!( + "MUX/XUDP: No destination in VLESS header (destinations come in XUDP frames)" + ); + let unparsed_data = stream_reader.unparsed_data(); + + // Send VLESS response header immediately + write_all(&mut server_stream, SERVER_RESPONSE_HEADER).await?; + + // Wrap raw stream in XUDP stream + let mut xudp_stream = XudpMessageStream::new(server_stream); + + // Feed any unparsed data to XUDP stream + if !unparsed_data.is_empty() { + xudp_stream.feed_initial_read_data(unparsed_data)?; + } + + Ok(TcpServerSetupResult::SessionBasedUdp { + stream: Box::new(xudp_stream), + need_initial_flush: false, + override_proxy_provider: NoneOrOne::Unspecified, + }) + } + unknown_protocol_type => { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("Unknown requested protocol: {unknown_protocol_type}"), + )); + } + } + } +} + +// REMOVED: Old tokio-rustls vision function - now using setup_custom_tls_vision_vless_server_stream + +/// Setup a VISION+VLESS stream from a CryptoTlsStream (for REALITY+Vision support) +pub async fn setup_custom_tls_vision_vless_server_stream( + mut tls_stream: CryptoTlsStream, + user_id: &[u8], + udp_enabled: bool, +) -> std::io::Result +where + IO: AsyncStream + 'static, +{ + // Parse VLESS header from TLS stream + let mut stream_reader = StreamReader::new_with_buffer_size(800); + + let client_version = stream_reader.read_u8(&mut tls_stream).await?; + if client_version != 0 { + return Err(std::io::Error::other(format!( + "invalid client protocol version, expected 0, got {client_version}" + ))); + } + + let target_id = stream_reader.read_slice(&mut tls_stream, 16).await?; + let mut user_uuid = [0u8; 16]; + user_uuid.copy_from_slice(target_id); + + // Verify user ID using constant-time comparison to prevent timing attacks + if user_id.ct_eq(target_id).unwrap_u8() == 0 { + return Err(std::io::Error::other("Unknown user id")); + } + + let addon_length = stream_reader.read_u8(&mut tls_stream).await?; + let flow = if addon_length > 0 { + parse_addons_from_reader(&mut stream_reader, &mut tls_stream, addon_length).await? + } else { + String::new() + }; + + let instruction = stream_reader.read_u8(&mut tls_stream).await?; + + match instruction { + COMMAND_TCP => { + if flow != XTLS_VISION_FLOW { + return Err(std::io::Error::other("expected vision flow for TCP")); + } + + info!("Parsing remote location..."); + let remote_location = + parse_remote_location_from_reader(&mut stream_reader, &mut tls_stream).await?; + info!("Remote location parsed: {}", remote_location); + let unparsed_data = stream_reader.unparsed_data(); + + let flow_stream: Box = if flow == XTLS_VISION_FLOW { + info!("Creating VISION stream (Custom TLS) for flow: {}", flow); + let (io, session) = tls_stream.into_inner(); + + Box::new(VisionStream::new_server( + io, + session, + user_uuid, + unparsed_data, + )?) + } else { + Box::new(tls_stream) + }; + + Ok(TcpServerSetupResult::TcpForward { + remote_location, + stream: flow_stream, + need_initial_flush: false, + connection_success_response: None, // VisionStream will send VLESS response with first write + initial_remote_data: None, // Data fed to VisionStream instead + override_proxy_provider: NoneOrOne::Unspecified, + }) + } + COMMAND_UDP => { + if !udp_enabled { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "UDP not enabled", + )); + } + + info!("Parsing remote location..."); + let remote_location = + parse_remote_location_from_reader(&mut stream_reader, &mut tls_stream).await?; + info!("Remote location parsed: {}", remote_location); + let unparsed_data = stream_reader.unparsed_data(); + + write_all(&mut tls_stream, SERVER_RESPONSE_HEADER).await?; + let mut vless_stream = VlessMessageStream::new(Box::new(tls_stream)); + if !unparsed_data.is_empty() { + vless_stream.feed_initial_read_data(unparsed_data)?; + } + + Ok(TcpServerSetupResult::BidirectionalUdp { + remote_location, + stream: Box::new(vless_stream), + need_initial_flush: false, + override_proxy_provider: NoneOrOne::Unspecified, + }) + } + COMMAND_MUX => { + if !udp_enabled { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "MUX/XUDP requires UDP to be enabled", + )); + } + // MUX/XUDP: Destination is NOT in the VLESS header - it comes in XUDP frames + info!("MUX/XUDP: No destination in VLESS header (destinations come in XUDP frames)"); + let unparsed_data = stream_reader.unparsed_data(); + + if flow == XTLS_VISION_FLOW { + info!("Creating VISION+XUDP stream (Custom TLS) with session-based UDP sockets"); + + // Extract components from CryptoTlsStream + let (io, session) = tls_stream.into_inner(); + + // Create VISION stream (will send VLESS response automatically on first write) + let vision_stream = + VisionStream::new_server(io, session, user_uuid, unparsed_data)?; + + // Wrap VISION stream in XUDP stream + let xudp_stream = XudpMessageStream::new(Box::new(vision_stream)); + + Ok(TcpServerSetupResult::SessionBasedUdp { + stream: Box::new(xudp_stream), + need_initial_flush: false, // VisionStream sends VLESS response on first write + override_proxy_provider: NoneOrOne::Unspecified, + }) + } else { + info!( + "Creating XUDP stream (Custom TLS, no VISION) with session-based UDP sockets" + ); + + // Send VLESS response header immediately + write_all(&mut tls_stream, SERVER_RESPONSE_HEADER).await?; + + // Wrap TLS stream in XUDP stream + let mut xudp_stream = XudpMessageStream::new(Box::new(tls_stream)); + + // Feed any unparsed data to XUDP stream + if !unparsed_data.is_empty() { + xudp_stream.feed_initial_read_data(unparsed_data)?; + } + + Ok(TcpServerSetupResult::SessionBasedUdp { + stream: Box::new(xudp_stream), + need_initial_flush: false, // Response already sent above + override_proxy_provider: NoneOrOne::Unspecified, + }) + } + } + unknown_protocol_type => Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("Unknown requested protocol: {unknown_protocol_type}"), + )), + } +} diff --git a/src/vless/vless_util.rs b/src/vless/vless_util.rs new file mode 100644 index 00000000..990295d1 --- /dev/null +++ b/src/vless/vless_util.rs @@ -0,0 +1,214 @@ +use std::net::{Ipv4Addr, Ipv6Addr}; +use std::sync::LazyLock; + +use tokio::io::AsyncReadExt; + +use crate::address::{Address, NetLocation}; +use crate::stream_reader::StreamReader; + +// VLESS protocol command types +pub const COMMAND_TCP: u8 = 1; +pub const COMMAND_UDP: u8 = 2; +pub const COMMAND_MUX: u8 = 3; // Also known as XUDP - multiplexes UDP over single TCP connection + +pub const XTLS_VISION_FLOW: &str = "xtls-rprx-vision"; + +pub async fn parse_addons_from_reader( + stream_reader: &mut StreamReader, + stream: &mut S, + addon_length: u8, +) -> std::io::Result { + let addon_bytes = stream_reader + .read_slice(stream, addon_length as usize) + .await?; + + log::debug!( + "Parsing addons: length={}, bytes={:?}", + addon_length, + addon_bytes + ); + + // Parse protobuf-encoded addons + // Format: field_tag length data [field_tag length data ...] + // Field 1 = flow (string) + // Field 2 = seed (bytes) + + let mut addon_cursor = 0; + let mut flow_string = String::new(); + + while addon_cursor < addon_bytes.len() { + // Read field tag + let field_tag = addon_bytes[addon_cursor]; + addon_cursor += 1; + + let field_number = field_tag >> 3; + let wire_type = field_tag & 0x07; + + if wire_type != 2 { + return Err(std::io::Error::other(format!( + "Unexpected wire type {} for field {}", + wire_type, field_number + ))); + } + + // Read length + if addon_cursor >= addon_bytes.len() { + return Err(std::io::Error::other( + "Unexpected end of addon data reading length", + )); + } + let (field_length, bytes_used) = read_varint(&addon_bytes[addon_cursor..])?; + addon_cursor += bytes_used; + + // Validate field_length is within bounds + if addon_cursor + field_length as usize > addon_bytes.len() { + return Err(std::io::Error::other(format!( + "Field {} length {} exceeds remaining addon bytes (cursor: {}, total: {})", + field_number, + field_length, + addon_cursor, + addon_bytes.len() + ))); + } + + // Read field data + let field_data = &addon_bytes[addon_cursor..addon_cursor + field_length as usize]; + addon_cursor += field_length as usize; + + match field_number { + 1 => { + // Flow field + flow_string = std::str::from_utf8(field_data) + .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))? + .to_string(); + log::debug!("Parsed flow: {}", flow_string); + } + 2 => { + // Seed field (ignored for now) + log::debug!("Parsed seed: {} bytes", field_data.len()); + } + _ => { + log::debug!( + "Unknown field {}, skipping {} bytes", + field_number, + field_data.len() + ); + } + } + } + + Ok(flow_string) +} + +/// Parse VLESS remote location from stream +pub async fn parse_remote_location_from_reader( + stream_reader: &mut StreamReader, + stream: &mut S, +) -> std::io::Result { + let port = stream_reader.read_u16_be(stream).await?; + let address_type = stream_reader.read_u8(stream).await?; + + match address_type { + 1 => { + // 4 byte ipv4 address + let address_bytes = stream_reader.read_slice(stream, 4).await?; + let v4addr = Ipv4Addr::new( + address_bytes[0], + address_bytes[1], + address_bytes[2], + address_bytes[3], + ); + Ok(NetLocation::new(Address::Ipv4(v4addr), port)) + } + 2 => { + // domain name + let domain_name_len = stream_reader.read_u8(stream).await?; + let domain_name_bytes = stream_reader + .read_slice(stream, domain_name_len as usize) + .await?; + + let address_str = std::str::from_utf8(domain_name_bytes).map_err(|e| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("Failed to decode address: {e}"), + ) + })?; + + Ok(NetLocation::new(Address::from(address_str)?, port)) + } + 3 => { + // 16 byte ipv6 address + let address_bytes = stream_reader.read_slice(stream, 16).await?; + let v6addr = Ipv6Addr::new( + ((address_bytes[0] as u16) << 8) | (address_bytes[1] as u16), + ((address_bytes[2] as u16) << 8) | (address_bytes[3] as u16), + ((address_bytes[4] as u16) << 8) | (address_bytes[5] as u16), + ((address_bytes[6] as u16) << 8) | (address_bytes[7] as u16), + ((address_bytes[8] as u16) << 8) | (address_bytes[9] as u16), + ((address_bytes[10] as u16) << 8) | (address_bytes[11] as u16), + ((address_bytes[12] as u16) << 8) | (address_bytes[13] as u16), + ((address_bytes[14] as u16) << 8) | (address_bytes[15] as u16), + ); + + Ok(NetLocation::new(Address::Ipv6(v6addr), port)) + } + invalid_type => Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("Invalid address type: {invalid_type}"), + )), + } +} + +pub fn vision_flow_addon_data() -> &'static [u8] { + static INSTANCE: LazyLock> = LazyLock::new(|| { + encode_flow_addon(XTLS_VISION_FLOW) + .expect("Failed to encode vision flow addon at initialization") + }); + &INSTANCE +} + +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)); + } + if cursor == 7 || cursor == data.len() { + return Err(std::io::Error::other("Varint is too long")); + } + cursor += 1; + } +} + +/// Encode a flow string as protobuf addon data +/// Format: field_tag(0x0a) + length + data +/// Field 1 = flow (string), wire type 2 (length-delimited) +fn encode_flow_addon(flow: &str) -> std::io::Result> { + let flow_bytes = flow.as_bytes(); + let flow_len = flow_bytes.len(); + + if flow_len > 127 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "Flow string too long for simple varint encoding", + )); + } + + let mut result = Vec::new(); + + // Field 1, wire type 2 (0x0a = (1 << 3) | 2) + result.push(0x0a); + + // Length as varint (simple case: < 128) + result.push(flow_len as u8); + + // Flow string data + result.extend_from_slice(flow_bytes); + + Ok(result) +} diff --git a/src/vless_handler.rs b/src/vless_handler.rs deleted file mode 100644 index f20be16c..00000000 --- a/src/vless_handler.rs +++ /dev/null @@ -1,324 +0,0 @@ -use std::net::{Ipv4Addr, Ipv6Addr}; - -use async_trait::async_trait; -use log::info; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; - -use crate::address::{Address, NetLocation}; -use crate::async_stream::AsyncStream; -use crate::option_util::NoneOrOne; -use crate::stream_reader::StreamReader; -use crate::tcp_handler::{ - TcpClientHandler, TcpClientSetupResult, TcpServerHandler, TcpServerSetupResult, -}; -use crate::util::{allocate_vec, parse_uuid, write_all}; -use crate::vless_message_stream::VlessMessageStream; - -#[derive(Debug)] -pub struct VlessTcpServerHandler { - user_id: Box<[u8]>, - udp_enabled: bool, -} - -impl VlessTcpServerHandler { - pub fn new(user_id: &str, udp_enabled: bool) -> Self { - Self { - user_id: parse_uuid(user_id).unwrap().into_boxed_slice(), - udp_enabled, - } - } -} - -const SERVER_RESPONSE_HEADER: &[u8] = &[ - 0u8, // version - 0u8, // addons length -]; - -#[async_trait] -impl TcpServerHandler for VlessTcpServerHandler { - async fn setup_server_stream( - &self, - mut server_stream: Box, - ) -> std::io::Result { - // this needs to be less than `read_buf` size in VlessMessageStream so that - // feed_initial_data can handle the unparsed data. - let mut stream_reader = StreamReader::new_with_buffer_size(800); - - let client_version = stream_reader.read_u8(&mut server_stream).await?; - if client_version != 0 { - return Err(std::io::Error::other(format!( - "invalid client protocol version, expected 0, got {client_version}" - ))); - } - - let target_id = stream_reader.read_slice(&mut server_stream, 16).await?; - for (b1, b2) in self.user_id.iter().zip(target_id.iter()) { - if b1 != b2 { - return Err(std::io::Error::other("Unknown user id")); - } - } - - let addon_length = stream_reader.read_u8(&mut server_stream).await?; - if addon_length > 0 { - stream_reader - .read_slice(&mut server_stream, addon_length as usize) - .await?; - } - - let instruction = stream_reader.read_u8(&mut server_stream).await?; - let is_udp = match instruction { - 1 => { - // tcp - false - } - 2 => { - if !self.udp_enabled { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - "UDP not enabled", - )); - } - true - } - unknown_protocol_type => { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!("Unknown requested protocol: {unknown_protocol_type}"), - )); - } - }; - - let port = stream_reader.read_u16_be(&mut server_stream).await?; - - let address_type = stream_reader.read_u8(&mut server_stream).await?; - let remote_location = match address_type { - 1 => { - // 4 byte ipv4 address - let address_bytes = stream_reader.read_slice(&mut server_stream, 4).await?; - let v4addr = Ipv4Addr::new( - address_bytes[0], - address_bytes[1], - address_bytes[2], - address_bytes[3], - ); - NetLocation::new(Address::Ipv4(v4addr), port) - } - 2 => { - // domain name - let domain_name_len = stream_reader.read_u8(&mut server_stream).await?; - let domain_name_bytes = stream_reader - .read_slice(&mut server_stream, domain_name_len as usize) - .await?; - - let address_str = match std::str::from_utf8(domain_name_bytes) { - Ok(s) => s, - Err(e) => { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!("Failed to decode address: {e}"), - )); - } - }; - - // Although this is supposed to be a hostname, some clients will pass - // ipv4 and ipv6 addresses as well, so parse it rather than directly - // using Address:Hostname enum. - NetLocation::new(Address::from(address_str)?, port) - } - 3 => { - // 16 byte ipv6 address - let address_bytes = stream_reader.read_slice(&mut server_stream, 16).await?; - let v6addr = Ipv6Addr::new( - ((address_bytes[0] as u16) << 8) | (address_bytes[1] as u16), - ((address_bytes[2] as u16) << 8) | (address_bytes[3] as u16), - ((address_bytes[4] as u16) << 8) | (address_bytes[5] as u16), - ((address_bytes[6] as u16) << 8) | (address_bytes[7] as u16), - ((address_bytes[8] as u16) << 8) | (address_bytes[9] as u16), - ((address_bytes[10] as u16) << 8) | (address_bytes[11] as u16), - ((address_bytes[12] as u16) << 8) | (address_bytes[13] as u16), - ((address_bytes[14] as u16) << 8) | (address_bytes[15] as u16), - ); - - NetLocation::new(Address::Ipv6(v6addr), port) - } - invalid_type => { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!("Invalid address type: {invalid_type}"), - )); - } - }; - - let unparsed_data = stream_reader.unparsed_data(); - - if !is_udp { - Ok(TcpServerSetupResult::TcpForward { - remote_location, - stream: server_stream, - need_initial_flush: false, - connection_success_response: Some( - SERVER_RESPONSE_HEADER.to_vec().into_boxed_slice(), - ), - initial_remote_data: if unparsed_data.is_empty() { - None - } else { - Some(unparsed_data.to_vec().into_boxed_slice()) - }, - override_proxy_provider: NoneOrOne::Unspecified, - }) - } else { - write_all(&mut server_stream, SERVER_RESPONSE_HEADER).await?; - - let mut vless_stream = VlessMessageStream::new(server_stream); - if !unparsed_data.is_empty() { - vless_stream.feed_initial_read_data(unparsed_data)?; - } - - Ok(TcpServerSetupResult::BidirectionalUdp { - remote_location, - stream: Box::new(vless_stream), - // don't flush and write the server response with the initial data. - need_initial_flush: false, - override_proxy_provider: NoneOrOne::Unspecified, - }) - } - } -} - -#[derive(Debug)] -pub struct VlessTcpClientHandler { - user_id: Box<[u8]>, -} - -impl VlessTcpClientHandler { - pub fn new(user_id: &str) -> Self { - Self { - user_id: parse_uuid(user_id).unwrap().into_boxed_slice(), - } - } -} - -#[async_trait] -impl TcpClientHandler for VlessTcpClientHandler { - async fn setup_client_stream( - &self, - _server_stream: &mut Box, - mut client_stream: Box, - remote_location: NetLocation, - ) -> std::io::Result { - // version + user id + header addon length + command + port + address type - let mut header_bytes = [0u8; 1 + 16 + 1 + 1 + 2 + 1]; - - // version 0 is fine, no need to write. - header_bytes[1..17].copy_from_slice(&self.user_id); - - // header addon length of 0 is fine, no need to write. - // tcp - header_bytes[18] = 1; - - let (remote_address, remote_port) = remote_location.unwrap_components(); - header_bytes[19] = (remote_port >> 8) as u8; - header_bytes[20] = (remote_port & 0xff) as u8; - - match remote_address { - Address::Ipv4(v4addr) => { - header_bytes[21] = 1; - write_all(&mut client_stream, &header_bytes).await?; - - let address_bytes = v4addr.octets(); - write_all(&mut client_stream, &address_bytes).await?; - } - Address::Ipv6(v6addr) => { - header_bytes[21] = 3; - write_all(&mut client_stream, &header_bytes).await?; - - let address_bytes = v6addr.octets(); - write_all(&mut client_stream, &address_bytes).await?; - } - Address::Hostname(hostname) => { - if hostname.len() > 255 { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!("Hostname is too long: {hostname}"), - )); - } - - header_bytes[21] = 2; - write_all(&mut client_stream, &header_bytes).await?; - - let hostname_len_byte: [u8; 1] = [hostname.len() as u8]; - write_all(&mut client_stream, &hostname_len_byte).await?; - - let hostname_bytes = hostname.into_bytes(); - write_all(&mut client_stream, &hostname_bytes).await?; - } - } - client_stream.flush().await?; - - // TODO: remove read_exact, used buffered writes above. - let mut response_header = [0u8; 2]; - client_stream.read_exact(&mut response_header).await?; - - if response_header[0] != 0 { - return Err(std::io::Error::other(format!( - "invalid server protocol version, expected 0, got {}", - response_header[0] - ))); - } - - let addon_length = response_header[1]; - if addon_length > 0 { - read_addons(&mut client_stream, addon_length).await?; - } - - Ok(TcpClientSetupResult { client_stream }) - } -} - -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)); - } - if cursor == 7 || cursor == data.len() { - return Err(std::io::Error::other("Varint is too long")); - } - cursor += 1; - } -} - -async fn read_addons(stream: &mut Box, addon_length: u8) -> std::io::Result<()> { - let mut addon_bytes = allocate_vec(addon_length as usize).into_boxed_slice(); - stream.read_exact(&mut addon_bytes).await?; - - let mut addon_cursor = 0; - let (flow_length, bytes_used) = read_varint(&addon_bytes)?; - addon_cursor += bytes_used; - - let flow_bytes = &addon_bytes[addon_cursor..addon_cursor + flow_length as usize]; - addon_cursor += flow_length as usize; - - let (seed_length, bytes_used) = read_varint(&addon_bytes[addon_cursor..])?; - addon_cursor += bytes_used; - let seed_bytes = &addon_bytes[addon_cursor..addon_cursor + seed_length as usize]; - addon_cursor += seed_length as usize; - - if addon_cursor as u8 != addon_length { - return Err(std::io::Error::other(format!( - "Did not consume all addon bytes, cursor is at {addon_cursor}, length is {addon_length}" - ))); - } - - info!( - "Read addon bytes: flow: {:?}, seed: {:?}", - &flow_bytes, &seed_bytes - ); - - Ok(()) -} diff --git a/src/vmess/md5.rs b/src/vmess/md5.rs index 227dbc61..378a19a5 100644 --- a/src/vmess/md5.rs +++ b/src/vmess/md5.rs @@ -1,4 +1,3 @@ -use hmac::{Hmac, Mac}; use md5::{Digest, Md5}; #[inline] @@ -8,15 +7,6 @@ pub fn compute_md5(data: &[u8]) -> [u8; 16] { context.finalize().into() } -#[inline] -pub fn compute_md5_repeating(data: &[u8], data_times: usize) -> [u8; 16] { - let mut context = Md5::new(); - for _ in 0..data_times { - md5::Digest::update(&mut context, data); - } - context.finalize().into() -} - #[inline] pub fn create_chacha_key(data: &[u8]) -> [u8; 32] { let mut ret = [0u8; 32]; @@ -27,12 +17,3 @@ pub fn create_chacha_key(data: &[u8]) -> [u8; 32] { context.finalize_into((&mut ret[16..]).into()); ret } - -type HmacMd5 = Hmac; - -#[inline] -pub fn compute_hmac_md5(key: &[u8], data: &[u8]) -> [u8; 16] { - let mut mac = HmacMd5::new_from_slice(key).unwrap(); - Mac::update(&mut mac, data); - mac.finalize().into_bytes().into() -} diff --git a/src/vmess/typed.rs b/src/vmess/typed.rs index 46a9fbf1..96b6fdb5 100644 --- a/src/vmess/typed.rs +++ b/src/vmess/typed.rs @@ -1,3 +1 @@ -pub type Aes128CfbEnc = cfb_mode::Encryptor; -pub type Aes128CfbDec = cfb_mode::Decryptor; pub type VmessReader = digest::core_api::XofReaderCoreWrapper; diff --git a/src/vmess/vmess_handler.rs b/src/vmess/vmess_handler.rs index 0f885b25..649f0579 100644 --- a/src/vmess/vmess_handler.rs +++ b/src/vmess/vmess_handler.rs @@ -1,25 +1,21 @@ -use std::collections::HashMap; use std::net::{Ipv4Addr, Ipv6Addr}; use std::time::SystemTime; -use aes::cipher::{BlockDecrypt, BlockEncrypt, KeyIvInit}; +use aes::cipher::{BlockDecrypt, BlockEncrypt}; use aes::Aes128; use async_trait::async_trait; use aws_lc_rs::aead::{ Aad, BoundKey, OpeningKey, SealingKey, UnboundKey, AES_128_GCM, CHACHA20_POLY1305, }; -use cfb_mode::cipher::AsyncStreamCipher; use digest::KeyInit; -use parking_lot::Mutex; use rand::{Rng, RngCore}; use sha3::digest::{ExtendableOutput, Update}; use sha3::Shake128; use tokio::io::AsyncWriteExt; use super::fnv1a::Fnv1aHasher; -use super::md5::{compute_hmac_md5, compute_md5, compute_md5_repeating, create_chacha_key}; +use super::md5::{compute_md5, create_chacha_key}; use super::nonce::{SingleUseNonce, VmessNonceSequence}; -use super::typed::{Aes128CfbDec, Aes128CfbEnc}; use super::vmess_stream::{ReadHeaderInfo, VmessStream}; use crate::address::{Address, NetLocation}; use crate::async_stream::AsyncStream; @@ -29,9 +25,15 @@ use crate::tcp_handler::{ TcpClientHandler, TcpClientSetupResult, TcpServerHandler, TcpServerSetupResult, }; use crate::util::{allocate_vec, parse_uuid, write_all}; +use crate::xudp::XudpMessageStream; const TAG_LEN: usize = 16; +// VMess protocol command types +const COMMAND_TCP: u8 = 1; +const COMMAND_UDP: u8 = 2; +const COMMAND_MUX: u8 = 3; // MUX/XUDP mode + #[derive(Debug, Clone, PartialEq, Eq)] enum DataCipher { Any, @@ -54,80 +56,17 @@ impl From<&str> for DataCipher { } } -type UserHash = [u8; 16]; - -#[derive(Debug)] -struct CertHashProvider { - user_key: [u8; 16], - hashes: HashMap, - last_hash_time_secs: u64, -} - -impl CertHashProvider { - pub fn new(user_id_bytes: &[u8]) -> Self { - if user_id_bytes.len() != 16 { - panic!("invalid user id bytes length ({})", user_id_bytes.len()); - } - let mut user_key = [0u8; 16]; - user_key.copy_from_slice(user_id_bytes); - Self { - user_key, - hashes: HashMap::with_capacity(64), - last_hash_time_secs: 0, - } - } - - pub fn check(&mut self, hash: &UserHash) -> Option { - if let Some(time_secs) = self.hashes.get(hash) { - return Some(*time_secs); - } - self.update_hashes(); - self.hashes.get(hash).copied() - } - - fn update_hashes(&mut self) { - let current_time_secs = SystemTime::UNIX_EPOCH.elapsed().unwrap().as_secs(); - - let to_time_secs = current_time_secs + 32; - if self.last_hash_time_secs >= to_time_secs { - return; - } - - let from_time_secs = current_time_secs - 32; - - self.hashes - .retain(|_, hash_time_secs| *hash_time_secs >= from_time_secs); - - let mut create_time_secs = std::cmp::max(from_time_secs, self.last_hash_time_secs); - while create_time_secs <= to_time_secs { - let hash_bytes: [u8; 16] = - compute_hmac_md5(&self.user_key, &create_time_secs.to_be_bytes()); - self.hashes.insert(hash_bytes, create_time_secs); - create_time_secs += 1; - } - - self.last_hash_time_secs = to_time_secs; - } -} - #[derive(Debug)] pub struct VmessTcpServerHandler { data_cipher: DataCipher, instruction_key: [u8; 16], aead_cipher: Aes128, - cert_hash_provider: Option>, udp_enabled: bool, } impl VmessTcpServerHandler { - pub fn new(cipher_name: &str, user_id: &str, force_aead: bool, udp_enabled: bool) -> Self { + pub fn new(cipher_name: &str, user_id: &str, udp_enabled: bool) -> Self { let mut user_id_bytes = parse_uuid(user_id).unwrap(); - let cert_hash_provider = if force_aead { - None - } else { - Some(Mutex::new(CertHashProvider::new(&user_id_bytes))) - }; - user_id_bytes.extend(b"c48619fe-8f02-49e0-b9e9-edf763e17e21"); let instruction_key: [u8; 16] = compute_md5(&user_id_bytes); @@ -138,7 +77,6 @@ impl VmessTcpServerHandler { data_cipher: cipher_name.into(), aead_cipher, instruction_key, - cert_hash_provider, udp_enabled, } } @@ -165,232 +103,221 @@ impl TcpServerHandler for VmessTcpServerHandler { self.aead_cipher.decrypt_block((&mut aead_bytes).into()); let checksum = super::crc32::crc32c(&aead_bytes[0..12]); let expected_checksum = u32::from_be_bytes(aead_bytes[12..16].try_into().unwrap()); - let is_aead_request = checksum == expected_checksum; - let mut header_reader = if is_aead_request { - 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)"), - )); - } + if checksum != expected_checksum { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "AEAD authentication failed: checksum mismatch", + )); + } - let mut encrypted_payload_length = [0u8; 18]; - stream_reader - .read_slice_into(&mut server_stream, &mut encrypted_payload_length) - .await?; + 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)"), + )); + } - let mut nonce = [0u8; 8]; - stream_reader - .read_slice_into(&mut server_stream, &mut nonce) - .await?; + let mut encrypted_payload_length = [0u8; 18]; + stream_reader + .read_slice_into(&mut server_stream, &mut encrypted_payload_length) + .await?; - let header_length_aead_key = super::sha2::kdf( - &self.instruction_key, - &[b"VMess Header AEAD Key_Length", &cert_hash, &nonce], - ); + let mut nonce = [0u8; 8]; + stream_reader + .read_slice_into(&mut server_stream, &mut nonce) + .await?; - let header_length_nonce = super::sha2::kdf( - &self.instruction_key, - &[b"VMess Header AEAD Nonce_Length", &cert_hash, &nonce], - ); + let header_length_aead_key = super::sha2::kdf( + &self.instruction_key, + &[b"VMess Header AEAD Key_Length", &cert_hash, &nonce], + ); - // TODO: don't unwrap - let unbound_key = - UnboundKey::new(&AES_128_GCM, &header_length_aead_key[0..16]).unwrap(); + let header_length_nonce = super::sha2::kdf( + &self.instruction_key, + &[b"VMess Header AEAD Nonce_Length", &cert_hash, &nonce], + ); - let mut opening_key = OpeningKey::new( - unbound_key, - SingleUseNonce::new(&header_length_nonce[0..12]), - ); + // TODO: don't unwrap + let unbound_key = UnboundKey::new(&AES_128_GCM, &header_length_aead_key[0..16]).unwrap(); - if opening_key - .open_in_place(Aad::from(&cert_hash), &mut encrypted_payload_length) - .is_err() - { - return Err(std::io::Error::other( - "failed to open encrypted header length", - )); - } + let mut opening_key = OpeningKey::new( + unbound_key, + SingleUseNonce::new(&header_length_nonce[0..12]), + ); - let payload_length = - u16::from_be_bytes(encrypted_payload_length[0..2].try_into().unwrap()); + if opening_key + .open_in_place(Aad::from(&cert_hash), &mut encrypted_payload_length) + .is_err() + { + return Err(std::io::Error::other( + "failed to open encrypted header length", + )); + } - let header_aead_key = super::sha2::kdf( - &self.instruction_key, - &[b"VMess Header AEAD Key", &cert_hash, &nonce], - ); + let payload_length = u16::from_be_bytes(encrypted_payload_length[0..2].try_into().unwrap()); - let header_nonce = super::sha2::kdf( - &self.instruction_key, - &[b"VMess Header AEAD Nonce", &cert_hash, &nonce], - ); + let header_aead_key = super::sha2::kdf( + &self.instruction_key, + &[b"VMess Header AEAD Key", &cert_hash, &nonce], + ); - let mut encrypted_header = - allocate_vec(payload_length as usize + TAG_LEN).into_boxed_slice(); + let header_nonce = super::sha2::kdf( + &self.instruction_key, + &[b"VMess Header AEAD Nonce", &cert_hash, &nonce], + ); - stream_reader - .read_slice_into(&mut server_stream, &mut encrypted_header) - .await?; + let mut encrypted_header = + allocate_vec(payload_length as usize + TAG_LEN).into_boxed_slice(); - // TODO: don't unwrap - let unbound_key = UnboundKey::new(&AES_128_GCM, &header_aead_key[0..16]).unwrap(); + stream_reader + .read_slice_into(&mut server_stream, &mut encrypted_header) + .await?; - let mut opening_key = - OpeningKey::new(unbound_key, SingleUseNonce::new(&header_nonce[0..12])); + // TODO: don't unwrap + let unbound_key = UnboundKey::new(&AES_128_GCM, &header_aead_key[0..16]).unwrap(); - if opening_key - .open_in_place(Aad::from(&cert_hash), &mut encrypted_header) - .is_err() - { - return Err(std::io::Error::other("failed to open encrypted header")); - } + let mut opening_key = + OpeningKey::new(unbound_key, SingleUseNonce::new(&header_nonce[0..12])); - HeaderReader::Aead(AeadHeaderReader { - server_stream, - decrypted_header: encrypted_header, - cursor: 0, - }) - } else { - let hash_time_secs = match self.cert_hash_provider { - Some(ref provider) => match provider.lock().check(&cert_hash) { - Some(t) => t, - None => { - return Err(std::io::Error::other("unauthorized request, unknown hash")); - } - }, - None => { - return Err(std::io::Error::other( - "unauthorized request, unknown aead hash", - )); - } - }; + if opening_key + .open_in_place(Aad::from(&cert_hash), &mut encrypted_header) + .is_err() + { + return Err(std::io::Error::other("failed to open encrypted header")); + } - let instruction_iv: [u8; 16] = { - let time_bytes = hash_time_secs.to_be_bytes(); - compute_md5_repeating(&time_bytes, 4) - }; - - let request_cipher = - Aes128CfbDec::new(&self.instruction_key.into(), &instruction_iv.into()); - - HeaderReader::AesCfb(AesCfbHeaderReader { - server_stream, - request_cipher, - }) + let mut header_reader = AeadHeaderReader { + server_stream, + decrypted_header: encrypted_header, + cursor: 0, }; let mut fnv_hasher = Fnv1aHasher::new(); - let mut instructions_to_addr_type = [0u8; 41]; - header_reader - .read_slice_into(&mut stream_reader, &mut instructions_to_addr_type) - .await?; - fnv_hasher.write(&instructions_to_addr_type); + // Read fixed 38-byte header first + let mut fixed_header = [0u8; 38]; + header_reader.read_slice_into(&mut fixed_header)?; + fnv_hasher.write(&fixed_header); - if instructions_to_addr_type[0] != 1 { + if fixed_header[0] != 1 { return Err(std::io::Error::new( std::io::ErrorKind::InvalidData, - format!("Invalid version {}", instructions_to_addr_type[0]), + format!("Invalid version {}", fixed_header[0]), )); } - let port = u16::from_be_bytes(instructions_to_addr_type[38..40].try_into().unwrap()); + let command = fixed_header[37]; - let remote_location = match instructions_to_addr_type[40] { - 1 => { - // 4 byte ipv4 address - let mut address_bytes = [0u8; 4]; - header_reader - .read_slice_into(&mut stream_reader, &mut address_bytes) - .await?; - fnv_hasher.write(&address_bytes); + log::info!("VMess command: {}", command); - let v4addr = Ipv4Addr::new( - address_bytes[0], - address_bytes[1], - address_bytes[2], - address_bytes[3], - ); - NetLocation::new(Address::Ipv4(v4addr), port) - } - 2 => { - // domain name - let mut domain_name_len = [0u8; 1]; - header_reader - .read_slice_into(&mut stream_reader, &mut domain_name_len) - .await?; - fnv_hasher.write(&domain_name_len); + // For MUX/XUDP command (0x03), there is NO destination in the VMess header + // Destinations come in XUDP frames themselves + let remote_location = if command == COMMAND_MUX { + // Use a placeholder address for XUDP - actual destinations come from XUDP frames + log::info!( + "VMess MUX/XUDP: No destination in VMess header (destinations come in XUDP frames)" + ); + NetLocation::new(Address::Ipv4(Ipv4Addr::new(0, 0, 0, 0)), 0) + } else { + // For TCP/UDP commands, read port (2 bytes) and address + let mut port_and_addr_type = [0u8; 3]; + header_reader.read_slice_into(&mut port_and_addr_type)?; + fnv_hasher.write(&port_and_addr_type); - let mut domain_name_bytes = allocate_vec(domain_name_len[0] as usize); - header_reader - .read_slice_into(&mut stream_reader, &mut domain_name_bytes) - .await?; - fnv_hasher.write(&domain_name_bytes); + let port = u16::from_be_bytes(port_and_addr_type[0..2].try_into().unwrap()); - let address_str = match std::str::from_utf8(&domain_name_bytes) { - Ok(s) => s, - Err(e) => { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!("Failed to decode address: {e}"), - )); - } - }; + match port_and_addr_type[2] { + 1 => { + // 4 byte ipv4 address + let mut address_bytes = [0u8; 4]; + header_reader.read_slice_into(&mut address_bytes)?; + fnv_hasher.write(&address_bytes); - // Although this is supposed to be a hostname, some clients will pass - // ipv4 and ipv6 addresses as well, so parse it rather than directly - // using Address:Hostname enum. - NetLocation::new(Address::from(address_str)?, port) - } - 3 => { - // 16 byte ipv6 address - let mut address_bytes = [0u8; 16]; - header_reader - .read_slice_into(&mut stream_reader, &mut address_bytes) - .await?; - fnv_hasher.write(&address_bytes); + let v4addr = Ipv4Addr::new( + address_bytes[0], + address_bytes[1], + address_bytes[2], + address_bytes[3], + ); + NetLocation::new(Address::Ipv4(v4addr), port) + } + 2 => { + // domain name + let mut domain_name_len = [0u8; 1]; + header_reader.read_slice_into(&mut domain_name_len)?; + fnv_hasher.write(&domain_name_len); - let v6addr = Ipv6Addr::new( - u16::from_be_bytes(address_bytes[0..2].try_into().unwrap()), - u16::from_be_bytes(address_bytes[2..4].try_into().unwrap()), - u16::from_be_bytes(address_bytes[4..6].try_into().unwrap()), - u16::from_be_bytes(address_bytes[6..8].try_into().unwrap()), - u16::from_be_bytes(address_bytes[8..10].try_into().unwrap()), - u16::from_be_bytes(address_bytes[10..12].try_into().unwrap()), - u16::from_be_bytes(address_bytes[12..14].try_into().unwrap()), - u16::from_be_bytes(address_bytes[14..16].try_into().unwrap()), - ); + let mut domain_name_bytes = allocate_vec(domain_name_len[0] as usize); + header_reader.read_slice_into(&mut domain_name_bytes)?; + fnv_hasher.write(&domain_name_bytes); - NetLocation::new(Address::Ipv6(v6addr), port) - } - invalid_type => { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!("Invalid address type: {invalid_type}"), - )); + let address_str = match std::str::from_utf8(&domain_name_bytes) { + Ok(s) => s, + Err(e) => { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("Failed to decode address: {e}"), + )); + } + }; + + // Although this is supposed to be a hostname, some clients will pass + // ipv4 and ipv6 addresses as well, so parse it rather than directly + // using Address:Hostname enum. + NetLocation::new(Address::from(address_str)?, port) + } + 3 => { + // 16 byte ipv6 address + let mut address_bytes = [0u8; 16]; + header_reader.read_slice_into(&mut address_bytes)?; + fnv_hasher.write(&address_bytes); + + let v6addr = Ipv6Addr::new( + u16::from_be_bytes(address_bytes[0..2].try_into().unwrap()), + u16::from_be_bytes(address_bytes[2..4].try_into().unwrap()), + u16::from_be_bytes(address_bytes[4..6].try_into().unwrap()), + u16::from_be_bytes(address_bytes[6..8].try_into().unwrap()), + u16::from_be_bytes(address_bytes[8..10].try_into().unwrap()), + u16::from_be_bytes(address_bytes[10..12].try_into().unwrap()), + u16::from_be_bytes(address_bytes[12..14].try_into().unwrap()), + u16::from_be_bytes(address_bytes[14..16].try_into().unwrap()), + ); + + NetLocation::new(Address::Ipv6(v6addr), port) + } + invalid_type => { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("Invalid address type: {invalid_type}"), + )); + } } }; - let margin_len: u8 = instructions_to_addr_type[35] >> 4; + let margin_len: u8 = fixed_header[35] >> 4; + log::info!("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 stream_reader, &mut margin_bytes) - .await?; + 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 stream_reader, &mut check_bytes) - .await?; + 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, @@ -402,10 +329,10 @@ impl TcpServerHandler for VmessTcpServerHandler { let server_stream = header_reader.into_stream(); - let data_encryption_iv: &[u8] = &instructions_to_addr_type[1..17]; - let data_encryption_key: &[u8] = &instructions_to_addr_type[17..33]; - let response_authentication_v = instructions_to_addr_type[33]; - let option = instructions_to_addr_type[34]; + let data_encryption_iv: &[u8] = &fixed_header[1..17]; + let data_encryption_key: &[u8] = &fixed_header[17..33]; + let response_authentication_v = fixed_header[33]; + let option = fixed_header[34]; if option & 0x01 != 0x01 { return Err(std::io::Error::new( @@ -433,7 +360,7 @@ impl TcpServerHandler for VmessTcpServerHandler { // the developer docs have incorrect values for the data type, // see headers.pb.go in v2ray-core for the correct values. - let requested_data_cipher = match instructions_to_addr_type[35] & 0b1111 { + let requested_data_cipher = match fixed_header[35] & 0b1111 { 1 => { return Err(std::io::Error::new( std::io::ErrorKind::InvalidData, @@ -461,45 +388,20 @@ impl TcpServerHandler for VmessTcpServerHandler { )); } - let is_udp = match instructions_to_addr_type[37] { - 1 => false, - 2 => { - if !self.udp_enabled { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - "UDP not enabled", - )); - } - true - } - unknown_protocol_type => { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - format!("Unknown requested protocol: {unknown_protocol_type}"), - )); - } - }; - - let mut response_header: [u8; 4] = [ + let response_header: [u8; 4] = [ response_authentication_v, 0, // option 0, // command 0, // command length ]; - let (response_header_iv, response_header_key): ([u8; 16], [u8; 16]) = if is_aead_request { - let mut truncated_iv = [0u8; 16]; - let mut truncated_key = [0u8; 16]; - truncated_iv.copy_from_slice(&super::sha2::compute_sha256(data_encryption_iv)[0..16]); - truncated_key.copy_from_slice(&super::sha2::compute_sha256(data_encryption_key)[0..16]); - - (truncated_iv, truncated_key) - } else { - ( - compute_md5(data_encryption_iv), - compute_md5(data_encryption_key), - ) - }; + // AEAD mode only - use SHA256 for response header keys + let mut truncated_iv = [0u8; 16]; + let mut truncated_key = [0u8; 16]; + truncated_iv.copy_from_slice(&super::sha2::compute_sha256(data_encryption_iv)[0..16]); + truncated_key.copy_from_slice(&super::sha2::compute_sha256(data_encryption_key)[0..16]); + let response_header_iv = truncated_iv; + let response_header_key = truncated_key; let unbound_keys = match requested_data_cipher { // TODO: stop unwrapping @@ -554,146 +456,157 @@ impl TcpServerHandler for VmessTcpServerHandler { // store the response header as prefix bytes to read when we are streaming. // writing the response header immediately without reading causes Surge to fail with // "Got short header" error. - let prefix_bytes: Box<[u8]> = if is_aead_request { - let response_header_length_aead_key = - super::sha2::kdf(&response_header_key, &[b"AEAD Resp Header Len Key"]); - let response_header_length_nonce = - super::sha2::kdf(&response_header_iv, &[b"AEAD Resp Header Len IV"]); + // AEAD mode only + let response_header_length_aead_key = + super::sha2::kdf(&response_header_key, &[b"AEAD Resp Header Len Key"]); + let response_header_length_nonce = + super::sha2::kdf(&response_header_iv, &[b"AEAD Resp Header Len IV"]); - let mut encrypted_response_header = [0u8; 2 + TAG_LEN + 4 + TAG_LEN]; + let mut encrypted_response_header = [0u8; 2 + TAG_LEN + 4 + TAG_LEN]; - // we know the size of response_header already. - encrypted_response_header[1] = 4; + // we know the size of response_header already. + encrypted_response_header[1] = 4; - // TODO: don't unwrap - let unbound_key = - UnboundKey::new(&AES_128_GCM, &response_header_length_aead_key[0..16]).unwrap(); - let mut sealing_key = SealingKey::new( - unbound_key, - SingleUseNonce::new(&response_header_length_nonce[0..12]), - ); - let tag = sealing_key - .seal_in_place_separate_tag(Aad::empty(), &mut encrypted_response_header[0..2]) - .unwrap(); - encrypted_response_header[2..2 + TAG_LEN].copy_from_slice(tag.as_ref()); + // TODO: don't unwrap + let unbound_key = + UnboundKey::new(&AES_128_GCM, &response_header_length_aead_key[0..16]).unwrap(); + let mut sealing_key = SealingKey::new( + unbound_key, + SingleUseNonce::new(&response_header_length_nonce[0..12]), + ); + let tag = sealing_key + .seal_in_place_separate_tag(Aad::empty(), &mut encrypted_response_header[0..2]) + .unwrap(); + encrypted_response_header[2..2 + TAG_LEN].copy_from_slice(tag.as_ref()); - let response_header_aead_key = - super::sha2::kdf(&response_header_key, &[b"AEAD Resp Header Key"]); - let response_header_nonce = - super::sha2::kdf(&response_header_iv, &[b"AEAD Resp Header IV"]); - let unbound_key = - UnboundKey::new(&AES_128_GCM, &response_header_aead_key[0..16]).unwrap(); - let mut sealing_key = SealingKey::new( - unbound_key, - SingleUseNonce::new(&response_header_nonce[0..12]), - ); - - encrypted_response_header[2 + TAG_LEN..2 + TAG_LEN + 4] - .copy_from_slice(&response_header); - - let tag = sealing_key - .seal_in_place_separate_tag( - Aad::empty(), - &mut encrypted_response_header[2 + TAG_LEN..2 + TAG_LEN + 4], - ) - .unwrap(); - encrypted_response_header[2 + TAG_LEN + 4..].copy_from_slice(tag.as_ref()); - - Box::new(encrypted_response_header) - } else { - let response_cipher = - Aes128CfbEnc::new(&response_header_key.into(), &response_header_iv.into()); - response_cipher.encrypt(&mut response_header); - Box::new(response_header) - }; - - let mut vmess_stream = VmessStream::new( - server_stream, - is_udp, - data_keys, - read_length_shake_reader, - write_length_shake_reader, - enable_global_padding, - Some(prefix_bytes), - None, + let response_header_aead_key = + super::sha2::kdf(&response_header_key, &[b"AEAD Resp Header Key"]); + let response_header_nonce = + super::sha2::kdf(&response_header_iv, &[b"AEAD Resp Header IV"]); + let unbound_key = UnboundKey::new(&AES_128_GCM, &response_header_aead_key[0..16]).unwrap(); + let mut sealing_key = SealingKey::new( + unbound_key, + SingleUseNonce::new(&response_header_nonce[0..12]), ); - let unparsed_data = stream_reader.unparsed_data(); - if !unparsed_data.is_empty() { - vmess_stream.feed_initial_read_data(unparsed_data)?; - } + encrypted_response_header[2 + TAG_LEN..2 + TAG_LEN + 4].copy_from_slice(&response_header); - let server_stream = Box::new(vmess_stream); + let tag = sealing_key + .seal_in_place_separate_tag( + Aad::empty(), + &mut encrypted_response_header[2 + TAG_LEN..2 + TAG_LEN + 4], + ) + .unwrap(); + encrypted_response_header[2 + TAG_LEN + 4..].copy_from_slice(tag.as_ref()); - match is_udp { - false => Ok(TcpServerSetupResult::TcpForward { - remote_location, - stream: server_stream, - // Wait until there is data to send the response header. - need_initial_flush: false, - connection_success_response: None, - initial_remote_data: None, - override_proxy_provider: NoneOrOne::Unspecified, - }), - true => Ok(TcpServerSetupResult::BidirectionalUdp { - remote_location, - stream: server_stream, - need_initial_flush: false, - override_proxy_provider: NoneOrOne::Unspecified, - }), - } - } -} + let prefix_bytes: Box<[u8]> = Box::new(encrypted_response_header); -#[allow(clippy::large_enum_variant)] -enum HeaderReader { - AesCfb(AesCfbHeaderReader), - Aead(AeadHeaderReader), -} + match command { + COMMAND_TCP => { + let mut vmess_stream = VmessStream::new( + server_stream, + false, // is_udp = false + data_keys, + read_length_shake_reader, + write_length_shake_reader, + enable_global_padding, + Some(prefix_bytes), + None, + ); -impl HeaderReader { - async fn read_slice_into( - &mut self, - stream_reader: &mut StreamReader, - data: &mut [u8], - ) -> std::io::Result<()> { - match self { - HeaderReader::AesCfb(ref mut reader) => { - reader.read_slice_into(stream_reader, data).await + let unparsed_data = stream_reader.unparsed_data(); + if !unparsed_data.is_empty() { + vmess_stream.feed_initial_read_data(unparsed_data)?; + } + + let server_stream = Box::new(vmess_stream); + + Ok(TcpServerSetupResult::TcpForward { + remote_location, + stream: server_stream, + // Wait until there is data to send the response header. + need_initial_flush: false, + connection_success_response: None, + initial_remote_data: None, + override_proxy_provider: NoneOrOne::Unspecified, + }) } - HeaderReader::Aead(ref mut reader) => reader.read_slice_into(data), + COMMAND_UDP => { + if !self.udp_enabled { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "UDP not enabled", + )); + } + + let mut vmess_stream = VmessStream::new( + server_stream, + true, // is_udp = true + data_keys, + read_length_shake_reader, + write_length_shake_reader, + enable_global_padding, + Some(prefix_bytes), + None, + ); + + let unparsed_data = stream_reader.unparsed_data(); + if !unparsed_data.is_empty() { + vmess_stream.feed_initial_read_data(unparsed_data)?; + } + + let server_stream = Box::new(vmess_stream); + + Ok(TcpServerSetupResult::BidirectionalUdp { + remote_location, + stream: server_stream, + need_initial_flush: false, + override_proxy_provider: NoneOrOne::Unspecified, + }) + } + COMMAND_MUX => { + if !self.udp_enabled { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "MUX/XUDP requires UDP to be enabled", + )); + } + + // For XUDP mode, use is_udp=false since XUDP wraps the stream + let mut vmess_stream = VmessStream::new( + server_stream, + false, // XUDP handles UDP multiplexing, VmessStream sees it as TCP-like + data_keys, + read_length_shake_reader, + write_length_shake_reader, + enable_global_padding, + Some(prefix_bytes), + None, + ); + + let unparsed_data = stream_reader.unparsed_data(); + if !unparsed_data.is_empty() { + vmess_stream.feed_initial_read_data(unparsed_data)?; + } + + // Wrap VmessStream with XudpMessageStream for session multiplexing + let xudp_stream = XudpMessageStream::new(Box::new(vmess_stream)); + + // No unparsed data to feed since VmessStream already consumed it + // (XUDP framing starts after VMess header) + + Ok(TcpServerSetupResult::SessionBasedUdp { + stream: Box::new(xudp_stream), + need_initial_flush: false, + override_proxy_provider: NoneOrOne::Unspecified, + }) + } + unknown_protocol_type => Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("Unknown requested protocol: {unknown_protocol_type}"), + )), } } - - fn into_stream(self) -> Box { - match self { - HeaderReader::AesCfb(reader) => reader.into_stream(), - HeaderReader::Aead(reader) => reader.into_stream(), - } - } -} - -struct AesCfbHeaderReader { - server_stream: Box, - request_cipher: Aes128CfbDec, -} - -impl AesCfbHeaderReader { - async fn read_slice_into( - &mut self, - stream_reader: &mut StreamReader, - data: &mut [u8], - ) -> std::io::Result<()> { - stream_reader - .read_slice_into(&mut self.server_stream, data) - .await?; - self.request_cipher.clone().decrypt(data); - Ok(()) - } - - fn into_stream(self) -> Box { - self.server_stream - } } struct AeadHeaderReader { @@ -718,18 +631,13 @@ impl AeadHeaderReader { #[derive(Debug)] pub struct VmessTcpClientHandler { data_cipher: DataCipher, - user_key: [u8; 16], instruction_key: [u8; 16], aead_cipher: Aes128, - is_aead: bool, } impl VmessTcpClientHandler { - pub fn new(cipher_name: &str, user_id: &str, is_aead: bool) -> Self { + pub fn new(cipher_name: &str, user_id: &str) -> Self { let mut user_id_bytes = parse_uuid(user_id).unwrap(); - let mut user_key = [0u8; 16]; - user_key.copy_from_slice(&user_id_bytes); - user_id_bytes.extend(b"c48619fe-8f02-49e0-b9e9-edf763e17e21"); let instruction_key: [u8; 16] = compute_md5(&user_id_bytes); @@ -739,9 +647,7 @@ impl VmessTcpClientHandler { Self { data_cipher: cipher_name.into(), aead_cipher, - user_key, instruction_key, - is_aead, } } } @@ -754,35 +660,25 @@ impl TcpClientHandler for VmessTcpClientHandler { mut client_stream: Box, remote_location: NetLocation, ) -> std::io::Result { - let (cert_hash, time_bytes): ([u8; 16], [u8; 8]) = if self.is_aead { - // 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; + // 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 mut aead_bytes = [0u8; 16]; - let time_bytes = time_secs.to_be_bytes(); - aead_bytes[0..8].copy_from_slice(&time_bytes); + let mut aead_bytes = [0u8; 16]; + let time_bytes = time_secs.to_be_bytes(); + aead_bytes[0..8].copy_from_slice(&time_bytes); - rand::rng().fill_bytes(&mut aead_bytes[8..12]); + rand::rng().fill_bytes(&mut aead_bytes[8..12]); - let checksum = super::crc32::crc32c(&aead_bytes[0..12]).to_be_bytes(); - aead_bytes[12..16].copy_from_slice(&checksum); + let checksum_value = super::crc32::crc32c(&aead_bytes[0..12]); + let checksum = checksum_value.to_be_bytes(); + aead_bytes[12..16].copy_from_slice(&checksum); - self.aead_cipher.encrypt_block((&mut aead_bytes).into()); - (aead_bytes, time_bytes) - } else { - // non-AEAD only allows 30 second delta. - let random_delta: u64 = rand::rng().random_range(0..61); - let time_secs: u64 = - SystemTime::UNIX_EPOCH.elapsed().unwrap().as_secs() - 30u64 + random_delta; - let time_bytes = time_secs.to_be_bytes(); - let hash_bytes = compute_hmac_md5(&self.user_key, &time_bytes); - (hash_bytes, time_bytes) - }; + self.aead_cipher.encrypt_block((&mut aead_bytes).into()); - write_all(&mut client_stream, &cert_hash).await?; + let cert_hash = aead_bytes; // max length of encrypted header: // 41 (instructions up to addr type) + 256 (max domain name length 255 + 1 length byte) + @@ -803,19 +699,13 @@ impl TcpClientHandler for VmessTcpClientHandler { // construct everything where we need data_encryption_iv and data_encryption_key now, // because instructions_to_addr_type will be encrypted once it's filled. - let (response_header_iv, response_header_key): ([u8; 16], [u8; 16]) = if self.is_aead { - let mut truncated_iv = [0u8; 16]; - let mut truncated_key = [0u8; 16]; - truncated_iv.copy_from_slice(&super::sha2::compute_sha256(data_encryption_iv)[0..16]); - truncated_key.copy_from_slice(&super::sha2::compute_sha256(data_encryption_key)[0..16]); - - (truncated_iv, truncated_key) - } else { - ( - compute_md5(data_encryption_iv), - compute_md5(data_encryption_key), - ) - }; + // AEAD mode only - use SHA256 for response header keys + let mut truncated_iv = [0u8; 16]; + let mut truncated_key = [0u8; 16]; + truncated_iv.copy_from_slice(&super::sha2::compute_sha256(data_encryption_iv)[0..16]); + truncated_key.copy_from_slice(&super::sha2::compute_sha256(data_encryption_key)[0..16]); + let response_header_iv = truncated_iv; + let response_header_key = truncated_key; let (read_length_shake_reader, write_length_shake_reader) = { let mut request_hasher = Shake128::default(); @@ -929,79 +819,75 @@ impl TcpClientHandler for VmessTcpClientHandler { header_bytes[cursor..cursor + 4].copy_from_slice(&check_bytes); cursor += 4; - if self.is_aead { - let mut encrypted_payload_length = [0u8; 18]; - let mut nonce = [0u8; 8]; - rand::rng().fill_bytes(&mut nonce); + // AEAD mode only + let mut encrypted_payload_length = [0u8; 18]; + let mut nonce = [0u8; 8]; + rand::rng().fill_bytes(&mut nonce); - let header_length_aead_key = super::sha2::kdf( - &self.instruction_key, - &[b"VMess Header AEAD Key_Length", &cert_hash, &nonce], - ); + let header_length_aead_key = super::sha2::kdf( + &self.instruction_key, + &[b"VMess Header AEAD Key_Length", &cert_hash, &nonce], + ); - let header_length_nonce = super::sha2::kdf( - &self.instruction_key, - &[b"VMess Header AEAD Nonce_Length", &cert_hash, &nonce], - ); + let header_length_nonce = super::sha2::kdf( + &self.instruction_key, + &[b"VMess Header AEAD Nonce_Length", &cert_hash, &nonce], + ); - let unbound_key = - UnboundKey::new(&AES_128_GCM, &header_length_aead_key[0..16]).unwrap(); + let unbound_key = UnboundKey::new(&AES_128_GCM, &header_length_aead_key[0..16]).unwrap(); - let mut sealing_key = SealingKey::new( - unbound_key, - SingleUseNonce::new(&header_length_nonce[0..12]), - ); + let mut sealing_key = SealingKey::new( + unbound_key, + SingleUseNonce::new(&header_length_nonce[0..12]), + ); - encrypted_payload_length[0] = (cursor >> 8) as u8; - encrypted_payload_length[1] = (cursor & 0xff) as u8; + encrypted_payload_length[0] = (cursor >> 8) as u8; + encrypted_payload_length[1] = (cursor & 0xff) as u8; - let tag = sealing_key - .seal_in_place_separate_tag( - Aad::from(&cert_hash), - &mut encrypted_payload_length[0..2], - ) - .unwrap(); + let tag = sealing_key + .seal_in_place_separate_tag(Aad::from(&cert_hash), &mut encrypted_payload_length[0..2]) + .unwrap(); - encrypted_payload_length[2..].copy_from_slice(tag.as_ref()); + encrypted_payload_length[2..].copy_from_slice(tag.as_ref()); - write_all(&mut client_stream, &encrypted_payload_length).await?; - write_all(&mut client_stream, &nonce).await?; + let header_aead_key = super::sha2::kdf( + &self.instruction_key, + &[b"VMess Header AEAD Key", &cert_hash, &nonce], + ); - let header_aead_key = super::sha2::kdf( - &self.instruction_key, - &[b"VMess Header AEAD Key", &cert_hash, &nonce], - ); + let header_nonce = super::sha2::kdf( + &self.instruction_key, + &[b"VMess Header AEAD Nonce", &cert_hash, &nonce], + ); - let header_nonce = super::sha2::kdf( - &self.instruction_key, - &[b"VMess Header AEAD Nonce", &cert_hash, &nonce], - ); + // TODO: don't unwrap + let unbound_key = UnboundKey::new(&AES_128_GCM, &header_aead_key[0..16]).unwrap(); + let mut sealing_key = + SealingKey::new(unbound_key, SingleUseNonce::new(&header_nonce[0..12])); + let tag = sealing_key + .seal_in_place_separate_tag(Aad::from(&cert_hash), &mut header_bytes[0..cursor]) + .unwrap(); - // TODO: don't unwrap - let unbound_key = UnboundKey::new(&AES_128_GCM, &header_aead_key[0..16]).unwrap(); - let mut sealing_key = - SealingKey::new(unbound_key, SingleUseNonce::new(&header_nonce[0..12])); - let tag = sealing_key - .seal_in_place_separate_tag(Aad::from(&cert_hash), &mut header_bytes[0..cursor]) - .unwrap(); + header_bytes[cursor..cursor + TAG_LEN].copy_from_slice(tag.as_ref()); + cursor += TAG_LEN; - header_bytes[cursor..cursor + TAG_LEN].copy_from_slice(tag.as_ref()); - cursor += TAG_LEN; - write_all(&mut client_stream, &header_bytes[0..cursor]).await?; - } else { - let instruction_iv: [u8; 16] = compute_md5_repeating(&time_bytes, 4); - let cipher = Aes128CfbEnc::new(&self.instruction_key.into(), &instruction_iv.into()); - let sized_header_bytes = &mut header_bytes[0..cursor]; - cipher.encrypt(sized_header_bytes); - write_all(&mut client_stream, sized_header_bytes).await?; - } + // Build complete AEAD request in a single buffer to avoid multiple writes as some + // servers expect to read the entire header in one shot + let total_len = 16 + 18 + 8 + cursor; + let mut complete_request = Vec::with_capacity(total_len); + complete_request.extend_from_slice(&cert_hash); + complete_request.extend_from_slice(&encrypted_payload_length); + complete_request.extend_from_slice(&nonce); + complete_request.extend_from_slice(&header_bytes[0..cursor]); + + write_all(&mut client_stream, &complete_request).await?; // Flush the entire request. client_stream.flush().await?; // Info for reading the server response, which arrives along with the initial data. + // Always AEAD mode let read_header_info = ReadHeaderInfo { - is_aead: self.is_aead, response_header_key, response_header_iv, response_authentication_v, diff --git a/src/vmess/vmess_stream.rs b/src/vmess/vmess_stream.rs index e01c9f68..3518dc8b 100644 --- a/src/vmess/vmess_stream.rs +++ b/src/vmess/vmess_stream.rs @@ -9,14 +9,13 @@ use rand::RngCore; use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; use super::nonce::{SingleUseNonce, VmessNonceSequence}; -use super::typed::{Aes128CfbDec, VmessReader}; +use super::typed::VmessReader; use crate::async_stream::{ AsyncFlushMessage, AsyncMessageStream, AsyncPing, AsyncReadMessage, AsyncShutdownMessage, AsyncStream, AsyncWriteMessage, }; use crate::util::allocate_vec; -use aes::cipher::{AsyncStreamCipher, KeyIvInit}; // this should be the same as vmess_handler.rs TAG_LEN. const HEADER_TAG_LEN: usize = 16; const ENCRYPTION_TAG_LEN: usize = 16; @@ -111,7 +110,6 @@ enum DecryptState { // Expected VMess header response, when we're a client. pub struct ReadHeaderInfo { - pub is_aead: bool, pub response_header_key: [u8; 16], pub response_header_iv: [u8; 16], pub response_authentication_v: u8, @@ -121,7 +119,6 @@ pub struct ReadHeaderInfo { enum ReadHeaderState { ReadAeadLength, ReadAeadContent(usize), - ReadLegacyContent, Done, } @@ -206,13 +203,7 @@ impl VmessStream { }; let read_header_state = match read_header_info { - Some(ref info) => { - if info.is_aead { - ReadHeaderState::ReadAeadLength - } else { - ReadHeaderState::ReadLegacyContent - } - } + Some(_) => ReadHeaderState::ReadAeadLength, None => ReadHeaderState::Done, }; @@ -290,7 +281,6 @@ impl VmessStream { ReadHeaderState::ReadAeadContent(content_len) => { self.process_read_header_aead_content(content_len) } - ReadHeaderState::ReadLegacyContent => self.process_read_header_legacy_content(), ReadHeaderState::Done => { panic!("process_read_header called with Done state"); } @@ -400,47 +390,6 @@ impl VmessStream { check_header_response(encrypted_response_header, response_authentication_v) } - fn process_read_header_legacy_content(&mut self) -> std::io::Result<()> { - if self.unprocessed_end_offset - self.unprocessed_start_offset < 4 { - return Ok(()); - } - - let ReadHeaderInfo { - response_header_key, - response_header_iv, - .. - } = self.read_header_info.as_ref().unwrap(); - - let response_header_bytes = &mut self.unprocessed_buf - [self.unprocessed_start_offset..self.unprocessed_start_offset + 4]; - let response_cipher = Aes128CfbDec::new( - (&response_header_key[..]).into(), - (&response_header_iv[..]).into(), - ); - response_cipher.decrypt(response_header_bytes); - - // do this here, because we would already have read/decrypted it in the aead clause. - let command_len = response_header_bytes[3]; - if command_len > 0 { - // if this becomes an issue, we should read the command bytes and ignore them. - return Err(std::io::Error::other("extra command bytes")); - } - - self.read_header_state = ReadHeaderState::Done; - self.unprocessed_start_offset += 4; - if self.unprocessed_start_offset == self.unprocessed_end_offset { - self.unprocessed_start_offset = 0; - self.unprocessed_end_offset = 0; - } - - let ReadHeaderInfo { - response_authentication_v, - .. - } = self.read_header_info.take().unwrap(); - - check_header_response(response_header_bytes, response_authentication_v) - } - fn try_decrypt(&mut self) -> std::io::Result { // returns true if a full packet was decrypted, false if not (ie. more data required) let available_len = self.unprocessed_end_offset - self.unprocessed_start_offset; diff --git a/src/xudp/frame.rs b/src/xudp/frame.rs new file mode 100644 index 00000000..f502cdc7 --- /dev/null +++ b/src/xudp/frame.rs @@ -0,0 +1,846 @@ +// XUDP frame protocol implementation +// Protocol-agnostic UDP multiplexing used by VLESS and VMess + +use bytes::{Buf, BufMut, BytesMut}; +use std::net::{Ipv4Addr, Ipv6Addr}; + +use crate::address::{Address, NetLocation}; + +/// XUDP session status +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(u8)] +pub enum SessionStatus { + New = 0x01, // Create new session + Keep = 0x02, // Continue session + End = 0x03, // Close session + KeepAlive = 0x04, // Heartbeat +} + +impl TryFrom for SessionStatus { + type Error = std::io::Error; + + fn try_from(value: u8) -> Result { + match value { + 0x01 => Ok(SessionStatus::New), + 0x02 => Ok(SessionStatus::Keep), + 0x03 => Ok(SessionStatus::End), + 0x04 => Ok(SessionStatus::KeepAlive), + other => Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("invalid session status: {}", other), + )), + } + } +} + +/// XUDP frame options (bitmask) +#[derive(Debug, Clone, Copy)] +pub struct FrameOption(u8); + +impl FrameOption { + pub const DATA: u8 = 0x01; + pub const ERROR: u8 = 0x02; + + pub fn new() -> Self { + Self(0) + } + + pub fn with_data(mut self) -> Self { + self.0 |= Self::DATA; + self + } + + pub fn has_data(&self) -> bool { + (self.0 & Self::DATA) != 0 + } + + pub fn has_error(&self) -> bool { + (self.0 & Self::ERROR) != 0 + } + + pub fn raw(&self) -> u8 { + self.0 + } +} + +impl Default for FrameOption { + fn default() -> Self { + Self::new() + } +} + +impl From for FrameOption { + fn from(value: u8) -> Self { + Self(value) + } +} + +/// Target network type +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(u8)] +pub enum TargetNetwork { + Tcp = 0x01, + Udp = 0x02, +} + +impl TryFrom for TargetNetwork { + type Error = std::io::Error; + + fn try_from(value: u8) -> Result { + match value { + 0x01 => Ok(TargetNetwork::Tcp), + 0x02 => Ok(TargetNetwork::Udp), + other => Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("invalid network type: {}", other), + )), + } + } +} + +/// XUDP frame metadata +#[derive(Debug)] +pub struct FrameMetadata { + pub session_id: u16, + pub status: SessionStatus, + pub option: FrameOption, + pub target: Option, + pub network: Option, +} + +impl FrameMetadata { + /// Encode frame metadata to bytes + pub fn encode(&self, buf: &mut BytesMut) -> std::io::Result<()> { + // Reserve space for length (will fill in at end) + let length_pos = buf.len(); + buf.put_u16(0); + + let metadata_start = buf.len(); + + // Session ID + buf.put_u16(self.session_id); + + // Status and Option + buf.put_u8(self.status as u8); + buf.put_u8(self.option.raw()); + + // For New status or Keep+UDP: write destination + if self.status == SessionStatus::New { + let network = self.network.ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "network required for SessionStatusNew", + ) + })?; + buf.put_u8(network as u8); + + let target = self.target.as_ref().ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "target required for SessionStatusNew", + ) + })?; + + // Port + buf.put_u16(target.port()); + + // Address + encode_address(buf, target.address())?; + } else if matches!(self.status, SessionStatus::Keep) + && matches!(self.network, Some(TargetNetwork::Udp)) + { + // For Keep frames with UDP, include destination + buf.put_u8(TargetNetwork::Udp as u8); + + let target = self.target.as_ref().ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "target required for Keep+UDP", + ) + })?; + + buf.put_u16(target.port()); + encode_address(buf, target.address())?; + } + + // Calculate and write length + let metadata_len = buf.len() - metadata_start; + let length_bytes = &mut buf[length_pos..length_pos + 2]; + length_bytes.copy_from_slice(&(metadata_len as u16).to_be_bytes()); + + Ok(()) + } + + /// Decode frame metadata from bytes + pub fn decode(buf: &mut BytesMut) -> std::io::Result> { + // Need at least length field + if buf.len() < 2 { + log::debug!( + "[XUDP DECODE] Buffer too short for length field: {} bytes", + buf.len() + ); + return Ok(None); + } + + let metadata_len = u16::from_be_bytes([buf[0], buf[1]]) as usize; + log::debug!( + "[XUDP DECODE] metadata_len={}, buf.len()={}, need={}", + metadata_len, + buf.len(), + 2 + metadata_len + ); + + // Need complete metadata + if buf.len() < 2 + metadata_len { + log::debug!("[XUDP DECODE] Incomplete metadata, need more data"); + return Ok(None); + } + + // Skip length field + buf.advance(2); + + if metadata_len < 4 { + log::error!( + "[XUDP DECODE] Metadata too short: {} (buffer was {} bytes, first 8 bytes: {:?})", + metadata_len, + buf.len() + 2, + &buf[..std::cmp::min(8, buf.len())] + ); + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("metadata too short: {}", metadata_len), + )); + } + + // Track how many metadata bytes we consume + let metadata_start = buf.len(); + + let session_id = buf.get_u16(); + let status = SessionStatus::try_from(buf.get_u8())?; + let option = FrameOption::from(buf.get_u8()); + + let mut network = None; + let mut target = None; + + // Calculate remaining metadata bytes (after session_id, status, option = 4 bytes) + let remaining_metadata = metadata_len.saturating_sub(4); + + // Parse destination for New or Keep+UDP + // Check remaining_metadata > 0 to know if there's address data + if remaining_metadata > 0 + && (status == SessionStatus::New + || (status == SessionStatus::Keep && buf.remaining() > 0 && buf[0] == 0x02)) + { + let net_byte = buf.get_u8(); + network = Some(TargetNetwork::try_from(net_byte)?); + + let port = buf.get_u16(); + let address = decode_address(buf)?; + target = Some(NetLocation::new(address, port)); + } + + // Consume any remaining metadata bytes we didn't parse + let consumed = metadata_start - buf.len(); + let unconsumed = metadata_len.saturating_sub(consumed); + if unconsumed > 0 { + log::debug!( + "[XUDP DECODE] Skipping {} unconsumed metadata bytes (GlobalID or padding)", + unconsumed + ); + buf.advance(unconsumed); + } + + Ok(Some(FrameMetadata { + session_id, + status, + option, + target, + network, + })) + } +} + +/// Encode address to buffer (VLESS address format) +fn encode_address(buf: &mut BytesMut, address: &Address) -> std::io::Result<()> { + match address { + Address::Ipv4(v4) => { + buf.put_u8(0x01); + buf.put_slice(&v4.octets()); + } + Address::Hostname(hostname) => { + if hostname.len() > 255 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "hostname too long", + )); + } + buf.put_u8(0x02); + buf.put_u8(hostname.len() as u8); + buf.put_slice(hostname.as_bytes()); + } + Address::Ipv6(v6) => { + buf.put_u8(0x03); + buf.put_slice(&v6.octets()); + } + } + Ok(()) +} + +/// Decode address from buffer (VLESS address format) +fn decode_address(buf: &mut BytesMut) -> std::io::Result
{ + if buf.remaining() < 1 { + return Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + "incomplete address type", + )); + } + + let addr_type = buf.get_u8(); + match addr_type { + 0x01 => { + // IPv4 + if buf.remaining() < 4 { + return Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + "incomplete IPv4 address", + )); + } + let mut octets = [0u8; 4]; + buf.copy_to_slice(&mut octets); + Ok(Address::Ipv4(Ipv4Addr::from(octets))) + } + 0x02 => { + // Hostname + if buf.remaining() < 1 { + return Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + "incomplete hostname length", + )); + } + let len = buf.get_u8() as usize; + if buf.remaining() < len { + return Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + "incomplete hostname", + )); + } + let hostname_bytes = buf.copy_to_bytes(len); + let hostname = std::str::from_utf8(&hostname_bytes) + .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; + Ok(Address::from(hostname)?) + } + 0x03 => { + // IPv6 + if buf.remaining() < 16 { + return Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + "incomplete IPv6 address", + )); + } + let mut octets = [0u8; 16]; + buf.copy_to_slice(&mut octets); + Ok(Address::Ipv6(Ipv6Addr::from(octets))) + } + other => Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("invalid address type: {}", other), + )), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + // ========================================================================== + // Basic encode/decode tests + // ========================================================================== + + #[test] + fn test_encode_decode_new_frame_udp() { + let metadata = FrameMetadata { + session_id: 42, + status: SessionStatus::New, + option: FrameOption::new().with_data(), + target: Some(NetLocation::new( + Address::Ipv4("1.1.1.1".parse().unwrap()), + 53, + )), + network: Some(TargetNetwork::Udp), + }; + + let mut buf = BytesMut::new(); + metadata.encode(&mut buf).unwrap(); + + let decoded = FrameMetadata::decode(&mut buf).unwrap().unwrap(); + assert_eq!(decoded.session_id, 42); + assert_eq!(decoded.status, SessionStatus::New); + assert!(decoded.option.has_data()); + assert_eq!(decoded.network, Some(TargetNetwork::Udp)); + assert!(decoded.target.is_some()); + let target = decoded.target.unwrap(); + assert_eq!(target.port(), 53); + } + + #[test] + fn test_encode_decode_keep_frame_udp() { + let metadata = FrameMetadata { + session_id: 100, + status: SessionStatus::Keep, + option: FrameOption::new().with_data(), + target: Some(NetLocation::new( + Address::Ipv4("8.8.8.8".parse().unwrap()), + 53, + )), + network: Some(TargetNetwork::Udp), + }; + + let mut buf = BytesMut::new(); + metadata.encode(&mut buf).unwrap(); + + let decoded = FrameMetadata::decode(&mut buf).unwrap().unwrap(); + assert_eq!(decoded.session_id, 100); + assert_eq!(decoded.status, SessionStatus::Keep); + assert!(decoded.option.has_data()); + assert_eq!(decoded.network, Some(TargetNetwork::Udp)); + } + + #[test] + fn test_encode_decode_end_frame() { + let metadata = FrameMetadata { + session_id: 50, + status: SessionStatus::End, + option: FrameOption::new(), + target: None, + network: None, + }; + + let mut buf = BytesMut::new(); + metadata.encode(&mut buf).unwrap(); + + let decoded = FrameMetadata::decode(&mut buf).unwrap().unwrap(); + assert_eq!(decoded.session_id, 50); + assert_eq!(decoded.status, SessionStatus::End); + assert!(!decoded.option.has_data()); + } + + #[test] + fn test_encode_decode_hostname() { + let metadata = FrameMetadata { + session_id: 1, + status: SessionStatus::New, + option: FrameOption::new(), + target: Some(NetLocation::new( + Address::Hostname("example.com".to_string()), + 443, + )), + network: Some(TargetNetwork::Tcp), + }; + + let mut buf = BytesMut::new(); + metadata.encode(&mut buf).unwrap(); + + let decoded = FrameMetadata::decode(&mut buf).unwrap().unwrap(); + assert_eq!(decoded.session_id, 1); + let target = decoded.target.unwrap(); + assert_eq!(target.port(), 443); + match target.address() { + Address::Hostname(h) => assert_eq!(h, "example.com"), + _ => panic!("expected hostname"), + } + } + + /// Test decoding a frame with extra metadata bytes (GlobalID simulation) + /// Xray-core appends 8-byte GlobalID for UDP session continuity. + /// We skip these bytes like sing-box + #[test] + fn test_decode_frame_with_globalid_extra_bytes() { + let mut buf = BytesMut::new(); + + // Manually construct a frame with GlobalID (8 extra bytes) + // Format: length(2) + session_id(2) + status(1) + option(1) + network(1) + port(2) + addr_type(1) + ipv4(4) + GlobalID(8) + let metadata_len: u16 = 4 + 1 + 2 + 1 + 4 + 8; // 20 bytes + buf.put_u16(metadata_len); + buf.put_u16(123); // session_id + buf.put_u8(0x01); // SessionStatus::New + buf.put_u8(0x01); // FrameOption::DATA + buf.put_u8(0x02); // TargetNetwork::Udp + buf.put_u16(53); // port + buf.put_u8(0x01); // address type: IPv4 + buf.put_slice(&[8, 8, 8, 8]); // 8.8.8.8 + buf.put_slice(&[0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08]); // GlobalID (8 bytes) + + let decoded = FrameMetadata::decode(&mut buf).unwrap().unwrap(); + assert_eq!(decoded.session_id, 123); + assert_eq!(decoded.status, SessionStatus::New); + assert!(decoded.option.has_data()); + assert_eq!(decoded.network, Some(TargetNetwork::Udp)); + let target = decoded.target.unwrap(); + assert_eq!(target.port(), 53); + assert_eq!(target.address(), &Address::Ipv4("8.8.8.8".parse().unwrap())); + + // Buffer should be fully consumed + assert_eq!(buf.len(), 0, "Buffer should be fully consumed after decode"); + } + + /// Test Keep frame without destination (TCP keep) + /// This is a minimal Keep frame with just session_id + status + option + #[test] + fn test_decode_keep_frame_tcp_no_destination() { + let mut buf = BytesMut::new(); + + // Minimal Keep frame: length(2) + session_id(2) + status(1) + option(1) = 4 bytes metadata + buf.put_u16(4); // metadata length + buf.put_u16(42); // session_id + buf.put_u8(0x02); // SessionStatus::Keep + buf.put_u8(0x01); // FrameOption::DATA + + let decoded = FrameMetadata::decode(&mut buf).unwrap().unwrap(); + assert_eq!(decoded.session_id, 42); + assert_eq!(decoded.status, SessionStatus::Keep); + assert!(decoded.option.has_data()); + assert!(decoded.network.is_none(), "TCP Keep should have no network"); + assert!( + decoded.target.is_none(), + "TCP Keep should have no destination" + ); + assert_eq!(buf.len(), 0); + } + + /// Test Keep frame with UDP destination change + /// Per Xray-core: Keep frames CAN have destination if network byte is 0x02 (UDP) + #[test] + fn test_decode_keep_frame_udp_with_destination() { + let mut buf = BytesMut::new(); + + // Keep+UDP frame with destination + // Format: length(2) + session_id(2) + status(1) + option(1) + network(1) + port(2) + addr_type(1) + ipv4(4) + let metadata_len: u16 = 4 + 1 + 2 + 1 + 4; // 12 bytes + buf.put_u16(metadata_len); + buf.put_u16(99); // session_id + buf.put_u8(0x02); // SessionStatus::Keep + buf.put_u8(0x01); // FrameOption::DATA + buf.put_u8(0x02); // TargetNetwork::Udp - this triggers destination parsing + buf.put_u16(1234); // port + buf.put_u8(0x01); // address type: IPv4 + buf.put_slice(&[192, 168, 1, 1]); // 192.168.1.1 + + let decoded = FrameMetadata::decode(&mut buf).unwrap().unwrap(); + assert_eq!(decoded.session_id, 99); + assert_eq!(decoded.status, SessionStatus::Keep); + assert_eq!(decoded.network, Some(TargetNetwork::Udp)); + let target = decoded.target.unwrap(); + assert_eq!(target.port(), 1234); + assert_eq!(buf.len(), 0); + } + + /// Test Keep frame where first byte after header is NOT 0x02 (not UDP) + /// Should NOT parse destination even if there's extra data + #[test] + fn test_decode_keep_frame_non_udp_extra_bytes() { + let mut buf = BytesMut::new(); + + // Keep frame where network byte is 0x01 (TCP) - should NOT parse destination + let metadata_len: u16 = 4 + 1 + 2 + 1 + 4; // 12 bytes total + buf.put_u16(metadata_len); + buf.put_u16(50); // session_id + buf.put_u8(0x02); // SessionStatus::Keep + buf.put_u8(0x01); // FrameOption::DATA + buf.put_u8(0x01); // TargetNetwork::Tcp (NOT UDP!) + buf.put_u16(443); // port (should be skipped) + buf.put_u8(0x01); // address type + buf.put_slice(&[1, 2, 3, 4]); // address (should be skipped) + + let decoded = FrameMetadata::decode(&mut buf).unwrap().unwrap(); + assert_eq!(decoded.session_id, 50); + assert_eq!(decoded.status, SessionStatus::Keep); + // Network and target should be None because first byte wasn't 0x02 + assert!( + decoded.network.is_none(), + "Keep+TCP should not parse network" + ); + assert!( + decoded.target.is_none(), + "Keep+TCP should not parse destination" + ); + // Extra bytes should be consumed/skipped + assert_eq!(buf.len(), 0, "Extra bytes should be consumed"); + } + + /// Test KeepAlive frame (status 0x04) + #[test] + fn test_decode_keepalive_frame() { + let mut buf = BytesMut::new(); + + buf.put_u16(4); // metadata length + buf.put_u16(1); // session_id + buf.put_u8(0x04); // SessionStatus::KeepAlive + buf.put_u8(0x00); // no options + + let decoded = FrameMetadata::decode(&mut buf).unwrap().unwrap(); + assert_eq!(decoded.session_id, 1); + assert_eq!(decoded.status, SessionStatus::KeepAlive); + assert!(!decoded.option.has_data()); + assert!(!decoded.option.has_error()); + assert_eq!(buf.len(), 0); + } + + /// Test End frame with ERROR option + #[test] + fn test_decode_end_frame_with_error() { + let mut buf = BytesMut::new(); + + buf.put_u16(4); // metadata length + buf.put_u16(77); // session_id + buf.put_u8(0x03); // SessionStatus::End + buf.put_u8(0x02); // FrameOption::ERROR + + let decoded = FrameMetadata::decode(&mut buf).unwrap().unwrap(); + assert_eq!(decoded.session_id, 77); + assert_eq!(decoded.status, SessionStatus::End); + assert!(!decoded.option.has_data()); + assert!(decoded.option.has_error()); + assert_eq!(buf.len(), 0); + } + + /// Test IPv6 destination + #[test] + fn test_decode_new_frame_ipv6() { + let mut buf = BytesMut::new(); + + // New frame with IPv6 destination + let metadata_len: u16 = 4 + 1 + 2 + 1 + 16; // 24 bytes + buf.put_u16(metadata_len); + buf.put_u16(200); // session_id + buf.put_u8(0x01); // SessionStatus::New + buf.put_u8(0x01); // FrameOption::DATA + buf.put_u8(0x02); // TargetNetwork::Udp + buf.put_u16(53); // port + buf.put_u8(0x03); // address type: IPv6 + buf.put_slice(&[ + 0x20, 0x01, 0x48, 0x60, 0x48, 0x60, 0, 0, 0, 0, 0, 0, 0, 0, 0x88, 0x88, + ]); // 2001:4860:4860::8888 + + let decoded = FrameMetadata::decode(&mut buf).unwrap().unwrap(); + assert_eq!(decoded.session_id, 200); + let target = decoded.target.unwrap(); + assert_eq!(target.port(), 53); + match target.address() { + Address::Ipv6(v6) => { + assert_eq!(v6.to_string(), "2001:4860:4860::8888"); + } + _ => panic!("expected IPv6"), + } + assert_eq!(buf.len(), 0); + } + + // ========================================================================== + // Error handling tests + // ========================================================================== + + /// Test incomplete buffer (not enough for length field) + #[test] + fn test_decode_incomplete_length() { + let mut buf = BytesMut::new(); + buf.put_u8(0x00); // Only 1 byte, need 2 for length + + let result = FrameMetadata::decode(&mut buf).unwrap(); + assert!(result.is_none(), "Should return None for incomplete length"); + assert_eq!(buf.len(), 1, "Buffer should not be consumed"); + } + + /// Test incomplete metadata (length says more than available) + #[test] + fn test_decode_incomplete_metadata() { + let mut buf = BytesMut::new(); + buf.put_u16(10); // metadata length = 10 + buf.put_u16(1); // session_id + // Only 2 bytes of metadata, but length says 10 + + let result = FrameMetadata::decode(&mut buf).unwrap(); + assert!( + result.is_none(), + "Should return None for incomplete metadata" + ); + } + + /// Test metadata too short (less than 4 bytes) + #[test] + fn test_decode_metadata_too_short() { + let mut buf = BytesMut::new(); + buf.put_u16(2); // metadata length = 2 (too short, need at least 4) + buf.put_u16(1); // only 2 bytes of metadata + + let result = FrameMetadata::decode(&mut buf); + assert!(result.is_err(), "Should error on metadata too short"); + } + + /// Test invalid session status + #[test] + fn test_decode_invalid_session_status() { + let mut buf = BytesMut::new(); + buf.put_u16(4); // metadata length + buf.put_u16(1); // session_id + buf.put_u8(0x99); // Invalid status! + buf.put_u8(0x00); // option + + let result = FrameMetadata::decode(&mut buf); + assert!(result.is_err(), "Should error on invalid session status"); + } + + /// Test invalid network type in New frame + #[test] + fn test_decode_invalid_network_type() { + let mut buf = BytesMut::new(); + let metadata_len: u16 = 4 + 1 + 2 + 1 + 4; + buf.put_u16(metadata_len); + buf.put_u16(1); // session_id + buf.put_u8(0x01); // SessionStatus::New + buf.put_u8(0x01); // FrameOption::DATA + buf.put_u8(0x99); // Invalid network type! + buf.put_u16(53); // port + buf.put_u8(0x01); // address type + buf.put_slice(&[1, 2, 3, 4]); + + let result = FrameMetadata::decode(&mut buf); + assert!(result.is_err(), "Should error on invalid network type"); + } + + /// Test invalid address type + #[test] + fn test_decode_invalid_address_type() { + let mut buf = BytesMut::new(); + let metadata_len: u16 = 4 + 1 + 2 + 1 + 4; + buf.put_u16(metadata_len); + buf.put_u16(1); // session_id + buf.put_u8(0x01); // SessionStatus::New + buf.put_u8(0x01); // FrameOption::DATA + buf.put_u8(0x02); // TargetNetwork::Udp + buf.put_u16(53); // port + buf.put_u8(0x99); // Invalid address type! + buf.put_slice(&[1, 2, 3, 4]); + + let result = FrameMetadata::decode(&mut buf); + assert!(result.is_err(), "Should error on invalid address type"); + } + + #[test] + fn test_frame_format_compatibility() { + let mut buf = BytesMut::new(); + + // Xray-core format for New+UDP: + // 2 bytes - length (of metadata after this field) + // 2 bytes - session id + // 1 byte - status + // 1 byte - option + // 1 byte - network + // 2 bytes - port + // n bytes - address (1 byte type + address bytes) + + let metadata_len: u16 = 2 + 1 + 1 + 1 + 2 + 1 + 4; // = 12 + buf.put_u16(metadata_len); + buf.put_u16(0x0001); // session_id = 1 + buf.put_u8(0x01); // SessionStatusNew + buf.put_u8(0x01); // OptionData + buf.put_u8(0x02); // TargetNetworkUDP + buf.put_u16(443); // port (big-endian) + buf.put_u8(0x01); // AddressTypeIPv4 + buf.put_slice(&[93, 184, 216, 34]); // 93.184.216.34 (example.com) + + let decoded = FrameMetadata::decode(&mut buf).unwrap().unwrap(); + + assert_eq!(decoded.session_id, 1); + assert_eq!(decoded.status, SessionStatus::New); + assert!(decoded.option.has_data()); + assert!(!decoded.option.has_error()); + assert_eq!(decoded.network, Some(TargetNetwork::Udp)); + + let target = decoded.target.unwrap(); + assert_eq!(target.port(), 443); + assert_eq!( + target.address(), + &Address::Ipv4("93.184.216.34".parse().unwrap()) + ); + } + + /// Test domain name encoding + #[test] + fn test_domain_format() { + let mut buf = BytesMut::new(); + + let domain = "dns.google"; + let metadata_len: u16 = 4 + 1 + 2 + 1 + 1 + domain.len() as u16; // 19 + buf.put_u16(metadata_len); + buf.put_u16(5); // session_id + buf.put_u8(0x01); // SessionStatusNew + buf.put_u8(0x01); // OptionData + buf.put_u8(0x02); // TargetNetworkUDP + buf.put_u16(853); // DNS over TLS port + buf.put_u8(0x02); // AddressTypeDomain + buf.put_u8(domain.len() as u8); + buf.put_slice(domain.as_bytes()); + + let decoded = FrameMetadata::decode(&mut buf).unwrap().unwrap(); + let target = decoded.target.unwrap(); + assert_eq!(target.port(), 853); + match target.address() { + Address::Hostname(h) => assert_eq!(h, "dns.google"), + _ => panic!("expected hostname"), + } + } + + /// Test that our encoder produces Xray-compatible output + #[test] + fn test_encode_xray_compatible() { + let metadata = FrameMetadata { + session_id: 100, + status: SessionStatus::New, + option: FrameOption::new().with_data(), + target: Some(NetLocation::new( + Address::Ipv4("1.1.1.1".parse().unwrap()), + 53, + )), + network: Some(TargetNetwork::Udp), + }; + + let mut buf = BytesMut::new(); + metadata.encode(&mut buf).unwrap(); + + // Verify the wire format + assert_eq!(buf[0..2], [0x00, 0x0c]); // length = 12 (big-endian) + assert_eq!(buf[2..4], [0x00, 0x64]); // session_id = 100 (big-endian) + assert_eq!(buf[4], 0x01); // SessionStatusNew + assert_eq!(buf[5], 0x01); // OptionData + assert_eq!(buf[6], 0x02); // TargetNetworkUDP + assert_eq!(buf[7..9], [0x00, 0x35]); // port = 53 (big-endian) + assert_eq!(buf[9], 0x01); // AddressTypeIPv4 + assert_eq!(buf[10..14], [1, 1, 1, 1]); // 1.1.1.1 + } + + /// Test maximum session ID (u16::MAX) + #[test] + fn test_max_session_id() { + let mut buf = BytesMut::new(); + buf.put_u16(4); + buf.put_u16(0xFFFF); // max session_id + buf.put_u8(0x02); // Keep + buf.put_u8(0x00); + + let decoded = FrameMetadata::decode(&mut buf).unwrap().unwrap(); + assert_eq!(decoded.session_id, 65535); + } + + /// Test both DATA and ERROR options set + #[test] + fn test_data_and_error_options() { + let mut buf = BytesMut::new(); + buf.put_u16(4); + buf.put_u16(1); + buf.put_u8(0x03); // End + buf.put_u8(0x03); // DATA | ERROR + + let decoded = FrameMetadata::decode(&mut buf).unwrap().unwrap(); + assert!(decoded.option.has_data()); + assert!(decoded.option.has_error()); + } +} diff --git a/src/xudp/message_stream.rs b/src/xudp/message_stream.rs new file mode 100644 index 00000000..a46bf3de --- /dev/null +++ b/src/xudp/message_stream.rs @@ -0,0 +1,621 @@ +// XUDP message stream - protocol-agnostic UDP session multiplexing +// Wraps any AsyncStream and provides XUDP frame encoding/decoding with session management +// Used by both VLESS and VMess protocols + +use bytes::{Buf, BufMut, BytesMut}; +use futures::ready; +use std::collections::HashMap; +use std::net::SocketAddr; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; +use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; + +use crate::address::{Address, NetLocation}; +use crate::async_stream::{ + AsyncFlushMessage, AsyncPing, AsyncReadSessionMessage, AsyncSessionMessageStream, + AsyncShutdownMessage, AsyncStream, AsyncWriteSessionMessage, +}; +use crate::resolver::{NativeResolver, ResolverCache}; + +use super::frame::{FrameMetadata, FrameOption, SessionStatus, TargetNetwork}; + +pub struct XudpMessageStream { + /// Underlying byte stream (VLESS VisionStream, VMess stream, or any TLS stream) that reads/writes raw XUDP frame bytes + inner_stream: Box, + + /// Read buffer for incoming XUDP frames + read_buffer: BytesMut, + + /// Write buffer for outgoing XUDP frames + write_buffer: BytesMut, + + /// Session ID counter (starts at 1, wraps at u16::MAX) + next_session_id: u16, + + /// Active sessions: destination (RESOLVED IP) -> session_id + /// IMPORTANT: Keys are always resolved IP addresses, never hostnames + /// This ensures responses from UDP sockets (which give us IPs) can find the correct session + destination_to_session: HashMap, + + /// Reverse mapping: session_id -> destination (RESOLVED IP) + /// Used to remember which session_id maps to which destination + session_to_destination: HashMap, + + /// Maps session_id -> ORIGINAL destination (before resolution) + /// This preserves hostnames for encoding in response frames + /// CRITICAL FOR WEBRTC: Client expects responses to reference original hostname, not resolved IP + session_to_original_destination: HashMap, + + /// Resolver cache for hostname resolution + /// Resolves hostnames to IPs before storing in session maps + resolver_cache: ResolverCache, + + /// Buffered incoming message (if we read a complete frame) + /// Stores: (data, original_destination, optional_resolved_destination) + /// If resolved_destination is None, we still need to resolve it + incoming_message: Option<(Vec, NetLocation, Option)>, + + /// EOF flag + is_eof: bool, +} + +impl XudpMessageStream { + pub fn new(inner_stream: Box) -> Self { + // Create resolver implicitly like UdpMessageStream does + let resolver = Arc::new(NativeResolver::new()); + Self { + inner_stream, + read_buffer: BytesMut::with_capacity(65536), + write_buffer: BytesMut::with_capacity(65536), + next_session_id: 1, + destination_to_session: HashMap::new(), + session_to_destination: HashMap::new(), + session_to_original_destination: HashMap::new(), + resolver_cache: ResolverCache::new(resolver), + incoming_message: None, + is_eof: false, + } + } + + /// Feed initial unparsed data into the read buffer + /// Used when protocol header parsing (VLESS/VMess) consumed data that belongs to XUDP frames + pub fn feed_initial_read_data(&mut self, data: &[u8]) -> std::io::Result<()> { + if data.is_empty() { + return Ok(()); + } + + log::debug!( + "[XUDP] Feeding {} bytes of initial data to read buffer", + data.len() + ); + self.read_buffer.extend_from_slice(data); + Ok(()) + } + + /// Allocate new session ID + fn allocate_session_id(&mut self) -> u16 { + loop { + let id = self.next_session_id; + self.next_session_id = self.next_session_id.wrapping_add(1); + if self.next_session_id == 0 { + self.next_session_id = 1; // Skip 0 + } + + // Check if ID is already in use (unlikely with u16 space) + if !self.session_to_destination.contains_key(&id) { + return id; + } + } + } + + /// Get or create session ID for destination, preserving original address + /// resolved_destination: IP address (used for reverse lookup from UDP responses) + /// original_destination: Original address from XUDP frame (may be hostname, used in response frames) + fn get_or_create_session( + &mut self, + resolved_destination: &NetLocation, + original_destination: &NetLocation, + ) -> (u16, bool) { + if let Some(&session_id) = self.destination_to_session.get(resolved_destination) { + log::debug!( + "[XUDP] Found existing session {} for destination {}", + session_id, + resolved_destination + ); + (session_id, false) // Existing session + } else { + let session_id = self.allocate_session_id(); + log::debug!( + "[XUDP] Creating NEW session {} for resolved dest {} (original: {})", + session_id, + resolved_destination, + original_destination + ); + self.destination_to_session + .insert(resolved_destination.clone(), session_id); + self.session_to_destination + .insert(session_id, resolved_destination.clone()); + // CRITICAL: Store original destination for use in response frames + self.session_to_original_destination + .insert(session_id, original_destination.clone()); + log::debug!( + "[XUDP] Session maps updated. Total sessions: {}", + self.destination_to_session.len() + ); + (session_id, true) // New session + } + } + + /// Try to decode one complete XUDP frame from the read buffer. + /// Returns: + /// Ok(Some((data, destination))) - Successfully decoded a complete frame + /// Ok(None) - Buffer doesn't contain a complete frame yet (need more data) + /// Err(e) - Error during decoding + /// + /// IMPORTANT: This function must NOT consume any bytes from the buffer unless + /// it successfully decodes a complete frame. Otherwise, partial frames would + /// be lost when the function is called again with more data. + fn try_decode_one_frame(&mut self) -> std::io::Result, NetLocation)>> { + log::debug!( + "[XUDP READ] Attempting to decode frame, buffer len: {}", + self.read_buffer.len() + ); + + // First, peek at the buffer to determine total frame size WITHOUT consuming anything. + // We need to check: metadata_len (2 bytes) + metadata + data_len (2 bytes) + data + + // Need at least 2 bytes for metadata length + if self.read_buffer.len() < 2 { + log::debug!("[XUDP READ] Buffer too short for metadata length field"); + return Ok(None); + } + + let metadata_len = u16::from_be_bytes([self.read_buffer[0], self.read_buffer[1]]) as usize; + + // Check if we have complete metadata + if self.read_buffer.len() < 2 + metadata_len { + log::debug!("[XUDP READ] Buffer too short for complete metadata"); + return Ok(None); + } + + // Peek at metadata to check if frame has data + if metadata_len < 4 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("metadata too short: {}", metadata_len), + )); + } + + // Peek at status and option bytes (at offset 2 + 2 + 1 = 5 for status, 6 for option) + let status_byte = self.read_buffer[2 + 2]; // Skip metadata_len(2) + session_id(2) + let option_byte = self.read_buffer[2 + 3]; // Skip metadata_len(2) + session_id(2) + status(1) + + let has_data = (option_byte & 0x01) != 0; // FrameOption::DATA = 0x01 + let is_end = status_byte == 0x03; // SessionStatus::End + let is_keepalive = status_byte == 0x04; // SessionStatus::KeepAlive + + // For End or KeepAlive frames, or frames without data, we only need the metadata + if is_end || is_keepalive || !has_data { + // We have enough data to decode this frame - proceed with actual decode + let metadata = FrameMetadata::decode(&mut self.read_buffer)? + .expect("metadata decode should succeed after length check"); + + log::debug!("[XUDP READ] Decoded frame: session_id={}, status={:?}, has_data={}, target={:?}, network={:?}", + metadata.session_id, metadata.status, metadata.option.has_data(), metadata.target, metadata.network); + + // Handle session mappings + if let Some(ref target) = metadata.target { + log::debug!( + "[XUDP READ] Updating session {} mapping to target {}", + metadata.session_id, + target + ); + self.session_to_destination + .insert(metadata.session_id, target.clone()); + self.session_to_original_destination + .insert(metadata.session_id, target.clone()); + } + + if metadata.status == SessionStatus::End { + if let Some(destination) = self.session_to_destination.remove(&metadata.session_id) + { + self.destination_to_session.remove(&destination); + } + self.session_to_original_destination + .remove(&metadata.session_id); + return self.try_decode_one_frame(); + } + + // No data, try to decode next frame + return self.try_decode_one_frame(); + } + + // Frame has data - check if we have the data length and data + let data_len_offset = 2 + metadata_len; + if self.read_buffer.len() < data_len_offset + 2 { + log::debug!("[XUDP READ] Buffer too short for data length field"); + return Ok(None); + } + + let data_len = u16::from_be_bytes([ + self.read_buffer[data_len_offset], + self.read_buffer[data_len_offset + 1], + ]) as usize; + + // Check if we have all the data + let total_frame_len = data_len_offset + 2 + data_len; + if self.read_buffer.len() < total_frame_len { + log::debug!( + "[XUDP READ] Buffer too short for complete frame: have {}, need {}", + self.read_buffer.len(), + total_frame_len + ); + return Ok(None); + } + + // Now we know we have a complete frame - consume it all + let metadata = FrameMetadata::decode(&mut self.read_buffer)? + .expect("metadata decode should succeed after length check"); + + log::debug!("[XUDP READ] Decoded frame: session_id={}, status={:?}, has_data={}, target={:?}, network={:?}", + metadata.session_id, metadata.status, metadata.option.has_data(), metadata.target, metadata.network); + + // Check for TCP destination - we don't support TCP over XUDP + if let Some(TargetNetwork::Tcp) = metadata.network { + return Err(std::io::Error::new( + std::io::ErrorKind::Unsupported, + "XUDP with TCP destinations is not supported. Only UDP destinations are supported.", + )); + } + + // Check for ERROR option bit - remote side is signaling an error + if metadata.option.has_error() { + log::error!( + "[XUDP READ] Received frame with ERROR option set for session {}", + metadata.session_id + ); + return Err(std::io::Error::new( + std::io::ErrorKind::ConnectionReset, + "XUDP session closed by remote with error", + )); + } + + // Store/update session mapping for session_id → destination + if let Some(ref target) = metadata.target { + log::debug!( + "[XUDP READ] Updating session {} mapping to target {}", + metadata.session_id, + target + ); + self.session_to_destination + .insert(metadata.session_id, target.clone()); + self.session_to_original_destination + .insert(metadata.session_id, target.clone()); + } + + // Consume data length (already verified we have it) + self.read_buffer.advance(2); + + if data_len == 0 { + // Empty data, try to decode next frame + return self.try_decode_one_frame(); + } + + // Extract data (already verified we have it) + let data = self.read_buffer[..data_len].to_vec(); + self.read_buffer.advance(data_len); + + // Determine destination + let destination = if let Some(ref target) = metadata.target { + target.clone() + } else { + // Look up in session map + self.session_to_destination + .get(&metadata.session_id) + .cloned() + .ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("unknown session ID: {}", metadata.session_id), + ) + })? + }; + + log::debug!( + "[XUDP READ] Decoded complete frame with {} bytes for destination {}", + data.len(), + destination + ); + Ok(Some((data, destination))) + } +} + +impl AsyncFlushMessage for XudpMessageStream { + fn poll_flush_message(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + + while !this.write_buffer.is_empty() { + let n = ready!(Pin::new(&mut this.inner_stream).poll_write(cx, &this.write_buffer))?; + this.write_buffer.advance(n); + } + + Pin::new(&mut this.inner_stream).poll_flush(cx) + } +} + +impl AsyncShutdownMessage for XudpMessageStream { + fn poll_shutdown_message( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + ready!(self.as_mut().poll_flush_message(cx))?; + Pin::new(&mut self.get_mut().inner_stream).poll_shutdown(cx) + } +} + +impl AsyncPing for XudpMessageStream { + fn supports_ping(&self) -> bool { + self.inner_stream.supports_ping() + } + + fn poll_write_ping(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.get_mut().inner_stream).poll_write_ping(cx) + } +} + +impl AsyncReadSessionMessage for XudpMessageStream { + fn poll_read_session_message( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + let this = self.get_mut(); + + // Return buffered message if available + if let Some((data, original_destination, resolved_opt)) = this.incoming_message.take() { + if data.len() > buf.remaining() { + // Re-buffer it + this.incoming_message = Some((data, original_destination, resolved_opt)); + return Poll::Ready(Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "buffer too small for incoming message", + ))); + } + + // Resolve if not already resolved + let resolved_destination = if let Some(resolved) = resolved_opt { + resolved + } else { + match this + .resolver_cache + .poll_resolve_location(cx, &original_destination) + { + Poll::Ready(Ok(socket_addr)) => { + let resolved = match socket_addr { + SocketAddr::V4(addr) => { + NetLocation::new(Address::Ipv4(*addr.ip()), addr.port()) + } + SocketAddr::V6(addr) => { + NetLocation::new(Address::Ipv6(*addr.ip()), addr.port()) + } + }; + resolved + } + Poll::Ready(Err(e)) => { + return Poll::Ready(Err(e)); + } + Poll::Pending => { + // Re-buffer and wait for DNS + this.incoming_message = Some((data, original_destination, None)); + return Poll::Pending; + } + } + }; + + // Find session ID for resolved destination, preserving original + let (session_id, _is_new) = + this.get_or_create_session(&resolved_destination, &original_destination); + + // Convert to SocketAddr for return + let socket_addr = futures::ready!(this + .resolver_cache + .poll_resolve_location(cx, &original_destination))?; + + buf.put_slice(&data); + return Poll::Ready(Ok((session_id, socket_addr))); + } + + if this.is_eof { + return Poll::Ready(Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + "EOF reached", + ))); + } + + loop { + // Try to decode a complete frame from the read buffer + match this.try_decode_one_frame()? { + Some((data, destination)) => { + // Resolve hostname to IP if needed + log::debug!("[XUDP SESSION READ] Resolving destination: {}", destination); + let socket_addr = + match this.resolver_cache.poll_resolve_location(cx, &destination) { + Poll::Ready(Ok(addr)) => addr, + Poll::Ready(Err(e)) => { + return Poll::Ready(Err(e)); + } + Poll::Pending => { + // DNS resolution pending - buffer the frame and wait + this.incoming_message = Some((data, destination, None)); + return Poll::Pending; + } + }; + + let resolved_destination = match socket_addr { + SocketAddr::V4(addr) => { + NetLocation::new(Address::Ipv4(*addr.ip()), addr.port()) + } + SocketAddr::V6(addr) => { + NetLocation::new(Address::Ipv6(*addr.ip()), addr.port()) + } + }; + log::debug!( + "[XUDP SESSION READ] Resolved {} -> {}", + destination, + resolved_destination + ); + + // Get or create session, preserving original destination (may be hostname) + let (session_id, _is_new) = + this.get_or_create_session(&resolved_destination, &destination); + log::debug!( + "[XUDP SESSION READ] Session {} mapped to {}", + session_id, + resolved_destination + ); + + // Successfully decoded a frame + if data.len() > buf.remaining() { + // Buffer it with resolved destination for next read + this.incoming_message = + Some((data, destination, Some(resolved_destination))); + return Poll::Ready(Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "buffer too small for incoming message", + ))); + } + + buf.put_slice(&data); + log::debug!( + "[XUDP SESSION READ] Returning {} bytes for session {} to {}", + data.len(), + session_id, + socket_addr + ); + return Poll::Ready(Ok((session_id, socket_addr))); + } + None => { + // Buffer doesn't have a complete frame, need to read more data + } + } + + // Read more data from inner stream + let original_filled = this.read_buffer.len(); + this.read_buffer.resize(original_filled + 8192, 0); + let mut temp_buf = ReadBuf::new(&mut this.read_buffer[original_filled..]); + + log::debug!( + "[XUDP SESSION READ] Reading from inner stream, current buffer has {} bytes", + original_filled + ); + let poll_result = Pin::new(&mut this.inner_stream).poll_read(cx, &mut temp_buf); + + let n = temp_buf.filled().len(); + this.read_buffer.truncate(original_filled + n); + + match ready!(poll_result) { + Ok(()) => { + log::debug!( + "[XUDP SESSION READ] Got {} bytes from inner stream (total buffer: {})", + n, + this.read_buffer.len() + ); + + if n == 0 { + this.is_eof = true; + return Poll::Ready(Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + "EOF reached", + ))); + } + + // Got new data, continue loop to try decoding again + continue; + } + Err(e) => { + log::error!("[XUDP SESSION READ] Error reading from inner stream: {}", e); + return Poll::Ready(Err(e)); + } + } + } + } +} + +impl AsyncWriteSessionMessage for XudpMessageStream { + fn poll_write_session_message( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + session_id: u16, + buf: &[u8], + target: &SocketAddr, + ) -> Poll> { + // This is the reverse direction: UDP response from internet → XUDP client + // CRITICAL: Use ORIGINAL destination (may be hostname) in response frame, NOT resolved IP! + // WebRTC and other protocols expect responses to reference the hostname they requested. + + log::debug!( + "[XUDP SESSION WRITE] Writing {} bytes for session {} from source {}", + buf.len(), + session_id, + target + ); + + // Flush any pending write buffer first + if !self.write_buffer.is_empty() { + ready!(self.as_mut().poll_flush_message(cx))?; + } + + // CRITICAL FIX: Use original destination (hostname) instead of resolved IP + // Look up the original destination that the client requested + let target_location = self + .session_to_original_destination + .get(&session_id) + .cloned() + .ok_or_else(|| { + log::error!( + "[XUDP SESSION WRITE] No original destination found for session {}! \ + This should not happen - session was created but original dest not stored.", + session_id + ); + std::io::Error::new( + std::io::ErrorKind::NotFound, + format!("No original destination for session {}", session_id), + ) + })?; + + log::debug!("[XUDP SESSION WRITE] Using original destination {} for session {} (response came from {})", + target_location, session_id, target); + + // Build Keep frame with data + let metadata = FrameMetadata { + session_id, + status: SessionStatus::Keep, + option: FrameOption::new().with_data(), + target: Some(target_location.clone()), // Use ORIGINAL (hostname), not resolved IP! + network: Some(TargetNetwork::Udp), + }; + + log::debug!( + "[XUDP SESSION WRITE] Encoding Keep frame: session_id={}, target={}, data_len={}", + session_id, + target_location, + buf.len() + ); + + // Encode metadata + metadata.encode(&mut self.write_buffer)?; + + // Write data length + self.write_buffer.put_u16(buf.len() as u16); + + // Write data + self.write_buffer.extend_from_slice(buf); + + // Flush immediately + + self.poll_flush_message(cx) + } +} + +impl AsyncSessionMessageStream for XudpMessageStream {} diff --git a/src/xudp/mod.rs b/src/xudp/mod.rs new file mode 100644 index 00000000..68857dcc --- /dev/null +++ b/src/xudp/mod.rs @@ -0,0 +1,8 @@ +// XUDP (Extended UDP) protocol implementation +// Protocol-agnostic UDP multiplexing over TCP connections +// Used by both VLESS and VMess protocols + +pub mod frame; +pub mod message_stream; + +pub use message_stream::XudpMessageStream;