diff --git a/src/errors.rs b/src/errors.rs index f5b58af..5e1416e 100644 --- a/src/errors.rs +++ b/src/errors.rs @@ -84,6 +84,9 @@ pub enum TunnelError { #[error("timeout: {0}")] #[allow(dead_code)] Timeout(String), + + #[error("hostname error: {0}")] + Host(#[from] HostError), } /// Errors returned by the connector side. diff --git a/src/main.rs b/src/main.rs index 2223c81..31b3500 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,6 +3,7 @@ mod config; mod errors; mod generate; mod redact; +mod socks5; mod tls; mod tunnel; @@ -246,7 +247,7 @@ fn run_connect( target_port ); tracing::info!( - "SOCKS5 proxy at {}:{} is not yet implemented (scheduled for SOCKS milestone)", + "SOCKS5 proxy at {}:{} ready for local clients", socks_host, socks_port ); @@ -258,6 +259,21 @@ fn run_connect( redact::Redacted::new(&auth_token) ); + // Resolve SOCKS5 bind address + let socks_addr = match tokio::runtime::Runtime::new() + .unwrap() + .block_on(tunnel::resolve_host(&socks_host, socks_port)) + { + Ok(addrs) => addrs[0], + Err(e) => { + eprintln!( + "Error: failed to resolve SOCKS5 address '{}': {}", + socks_host, e + ); + return 1; + } + }; + let config = ConnectorConfig { target_host, target_port, @@ -267,10 +283,10 @@ fn run_connect( auth_token: Arc::new(auth_token), }; - // Run the async connector + // Run the async connector with SOCKS5 proxy if let Err(e) = tokio::runtime::Runtime::new() .unwrap() - .block_on(tunnel::connect_tunnel(config)) + .block_on(tunnel::run_connector_with_socks(config, socks_addr)) { eprintln!("Error: {}", e); return 1; diff --git a/src/socks5.rs b/src/socks5.rs new file mode 100644 index 0000000..4cb03ee --- /dev/null +++ b/src/socks5.rs @@ -0,0 +1,786 @@ +/// SOCKS5 protocol handler for the local proxy listener. +/// +/// Implements the SOCKS5 handshake and CONNECT method only (per RFC 1928). +/// Supports username/password authentication (RFC 1929) and IPv4/hostname targets. +/// +/// This module is purely synchronous protocol parsing — the async I/O layer +/// lives in the tunnel module, which calls this module for handshake/parsing. +/// +/// Security: +/// - Malformed requests are rejected with an error (no panic). +/// - Only CONNECT is supported (no BIND/UDP). +/// - Hostname targets are resolved on the listener side. +use std::net::SocketAddr; +use thiserror::Error; + +// --------------------------------------------------------------------------- +// Errors +// --------------------------------------------------------------------------- + +#[derive(Debug, Error)] +pub enum Socks5Error { + #[error("invalid SOCKS5 version, expected 5, got {0}")] + InvalidVersion(u8), + + #[error("unsupported authentication method: {0:#04x}")] + UnsupportedAuthMethod(u8), + + #[error("too many auth methods: {0}")] + #[allow(dead_code)] + TooManyAuthMethods(usize), + + #[error("invalid auth method count: {0}")] + InvalidAuthMethodCount(usize), + + #[error("unsupported command: {0}")] + UnsupportedCommand(u8), + + #[error("unsupported address type: {0}")] + UnsupportedAddrType(u8), + + #[error("address parsing error: {0}")] + AddrParseError(String), + + #[error("authentication failed")] + AuthFailed, + + #[error("connection not allowed")] + ConnectionNotAllowed, + + #[error("network unreachable")] + #[allow(dead_code)] + NetworkUnreachable, + + #[error("host unreachable")] + HostUnreachable, + + #[error("connection refused")] + ConnectionRefused, + + #[error("protocol error: {0}")] + Protocol(String), + + #[error("I/O error: {0}")] + Io(#[from] std::io::Error), + + #[error("stream not authenticated")] + NotAuthenticated, +} + +// --------------------------------------------------------------------------- +// Constants (RFC 1928) +// --------------------------------------------------------------------------- + +pub const VERSION: u8 = 0x05; + +// Methods +pub const AUTH_NONE: u8 = 0x00; +pub const AUTH_USERNAME_PASSWORD: u8 = 0x02; +pub const AUTH_NO_ACCEPTABLE: u8 = 0xFF; + +// Commands +pub const CMD_CONNECT: u8 = 0x01; +#[allow(dead_code)] +pub const CMD_BIND: u8 = 0x02; +#[allow(dead_code)] +pub const CMD_UDP: u8 = 0x03; + +// Address types +pub const ATYP_IPV4: u8 = 0x01; +pub const ATYP_DOMAIN: u8 = 0x03; +#[allow(dead_code)] +pub const ATYP_IPV6: u8 = 0x04; + +// Reply codes +pub const REPL_SUCCEEDED: u8 = 0x00; +pub const REPL_GENERAL_FAILURE: u8 = 0x01; +pub const REPL_CONN_NOT_ALLOWED: u8 = 0x02; +#[allow(dead_code)] +pub const REPL_NETWORK_UNREACHABLE: u8 = 0x03; +#[allow(dead_code)] +pub const REPL_HOST_UNREACHABLE: u8 = 0x04; +#[allow(dead_code)] +pub const REPL_CONN_REFUSED: u8 = 0x05; + +// --------------------------------------------------------------------------- +// Target address representation +// --------------------------------------------------------------------------- + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Socks5Target { + /// IPv4 hostname or domain name target + Domain { domain: String, port: u16 }, + /// IPv4 address target + Ipv4(std::net::Ipv4Addr, u16), +} + +impl Socks5Target { + /// Return the port for this target. + #[allow(dead_code)] + pub fn port(&self) -> u16 { + match self { + Socks5Target::Domain { port, .. } => *port, + Socks5Target::Ipv4(_, port) => *port, + } + } + + /// Return the host string for display purposes. + #[allow(dead_code)] + pub fn host(&self) -> String { + match self { + Socks5Target::Domain { domain, .. } => domain.clone(), + Socks5Target::Ipv4(addr, _) => addr.to_string(), + } + } +} + +// --------------------------------------------------------------------------- +// Wire-format builders (for tests and response generation) +// --------------------------------------------------------------------------- + +/// Build the server greeting: VERSION + N_METHODS + 0x00 (NO AUTH) + 0x02 (USER/PASS) +#[allow(dead_code)] +pub fn build_greeting() -> Vec { + vec![VERSION, 0x02, AUTH_NONE, AUTH_USERNAME_PASSWORD] +} + +/// Build the auth success response. +pub fn build_auth_success() -> Vec { + vec![0x01, 0x00] +} + +/// Build the auth failure response. +pub fn build_auth_failure() -> Vec { + vec![0x01, 0x01] +} + +/// Build a successful CONNECT reply. +pub fn build_reply_success() -> Vec { + // VERSION + REP + RSV + ATYP_IPV4 + 0.0.0.0 + 0.0 + vec![ + VERSION, + REPL_SUCCEEDED, + 0x00, // RSV + ATYP_IPV4, + 0x00, + 0x00, + 0x00, + 0x00, + 0x00, + 0x00, + ] +} + +/// Build a failure reply. +pub fn build_reply_failure(code: u8) -> Vec { + vec![ + VERSION, code, 0x00, // RSV + ATYP_IPV4, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, + ] +} + +// --------------------------------------------------------------------------- +// Parsing functions +// --------------------------------------------------------------------------- + +/// Parse the initial client greeting. +/// +/// Expected format: +/// | VER | NMETHODS | METHODS... | +/// +/// Returns the chosen auth method (0x00 for none, 0x02 for userpass). +/// If the client doesn't support any method we support, returns None (server should send 0xFF). +#[allow(dead_code)] +pub fn parse_greeting(buf: &[u8]) -> Result, Socks5Error> { + if buf.len() < 2 { + return Err(Socks5Error::Protocol("greeting too short".to_string())); + } + if buf[0] != VERSION { + return Err(Socks5Error::InvalidVersion(buf[0])); + } + let n_methods = buf[1] as usize; + if n_methods == 0 { + return Err(Socks5Error::InvalidAuthMethodCount(0)); + } + if buf.len() < 2 + n_methods { + return Err(Socks5Error::Protocol("greeting truncated".to_string())); + } + + let methods = &buf[2..2 + n_methods]; + // Check if the client only supports 0xFF (NO_ACCEPTABLE) + if methods.iter().all(|&m| m == AUTH_NO_ACCEPTABLE) { + return Ok(None); + } + // Prefer username/password, then no-auth + if methods.contains(&AUTH_USERNAME_PASSWORD) { + Ok(Some(AUTH_USERNAME_PASSWORD)) + } else if methods.contains(&AUTH_NONE) { + Ok(Some(AUTH_NONE)) + } else { + // Check if any method is unsupported + for &m in methods { + if m != AUTH_NONE && m != AUTH_USERNAME_PASSWORD { + return Err(Socks5Error::UnsupportedAuthMethod(m)); + } + } + Ok(None) + } +} + +/// Parse the username/password auth request. +/// +/// Expected format: +/// | VER | ULEN | UNAME... | PLEN | PASSWD... | +/// +/// Returns (username, password) strings. +#[allow(dead_code)] +pub fn parse_auth_request(buf: &[u8]) -> Result<(String, String), Socks5Error> { + if buf.len() < 3 { + return Err(Socks5Error::Protocol("auth request too short".to_string())); + } + if buf[0] != 0x01 { + return Err(Socks5Error::InvalidVersion(buf[0])); + } + let ulen = buf[1] as usize; + if buf.len() < 2 + ulen + 1 { + return Err(Socks5Error::Protocol("auth request truncated".to_string())); + } + let username = String::from_utf8(buf[2..2 + ulen].to_vec()) + .map_err(|_| Socks5Error::AddrParseError("invalid UTF-8 in username".to_string()))?; + let plen_offset = 2 + ulen; + let plen = buf[plen_offset] as usize; + if buf.len() < plen_offset + 1 + plen { + return Err(Socks5Error::Protocol("auth request truncated".to_string())); + } + let password = String::from_utf8(buf[plen_offset + 1..plen_offset + 1 + plen].to_vec()) + .map_err(|_| Socks5Error::AddrParseError("invalid UTF-8 in password".to_string()))?; + + Ok((username, password)) +} + +/// Parse a CONNECT request. +/// +/// Expected format: +/// | VER | CMD | RSV | ATYP | DST.ADDR | DST.PORT | +/// +/// Returns (command, target). +#[allow(dead_code)] +pub fn parse_connect_request(buf: &[u8]) -> Result<(u8, Socks5Target), Socks5Error> { + if buf.len() < 4 { + return Err(Socks5Error::Protocol( + "CONNECT request too short".to_string(), + )); + } + if buf[0] != VERSION { + return Err(Socks5Error::InvalidVersion(buf[0])); + } + let cmd = buf[1]; + if cmd != CMD_CONNECT { + return Err(Socks5Error::UnsupportedCommand(cmd)); + } + // buf[2] is RSV, must be 0x00 — we ignore it + + let atyp = buf[3]; + let addr_start = 4; + + match atyp { + ATYP_IPV4 => { + if buf.len() < addr_start + 4 + 2 { + return Err(Socks5Error::Protocol("IPv4 address truncated".to_string())); + } + let ipv4 = std::net::Ipv4Addr::new( + buf[addr_start], + buf[addr_start + 1], + buf[addr_start + 2], + buf[addr_start + 3], + ); + let port = u16::from_be_bytes([buf[addr_start + 4], buf[addr_start + 5]]); + Ok((cmd, Socks5Target::Ipv4(ipv4, port))) + } + ATYP_DOMAIN => { + if buf.len() < addr_start + 1 { + return Err(Socks5Error::Protocol( + "domain length byte missing".to_string(), + )); + } + let domain_len = buf[addr_start] as usize; + if buf.len() < addr_start + 1 + domain_len + 2 { + return Err(Socks5Error::Protocol( + "domain address truncated".to_string(), + )); + } + let domain = + String::from_utf8(buf[addr_start + 1..addr_start + 1 + domain_len].to_vec()) + .map_err(|_| { + Socks5Error::AddrParseError("invalid UTF-8 in domain".to_string()) + })?; + let port_offset = addr_start + 1 + domain_len; + let port = u16::from_be_bytes([buf[port_offset], buf[port_offset + 1]]); + Ok((cmd, Socks5Target::Domain { domain, port })) + } + ATYP_IPV6 => Err(Socks5Error::UnsupportedAddrType(ATYP_IPV6)), + other => Err(Socks5Error::UnsupportedAddrType(other)), + } +} + +/// Encode a Socks5Target to wire format bytes (ATYP + addr + port). +#[allow(dead_code)] +pub fn encode_target(target: &Socks5Target) -> Vec { + match target { + Socks5Target::Ipv4(addr, port) => { + let mut v = vec![ATYP_IPV4]; + v.extend_from_slice(&addr.octets()); + v.extend_from_slice(&port.to_be_bytes()); + v + } + Socks5Target::Domain { domain, port } => { + let mut v = vec![ATYP_DOMAIN, domain.len() as u8]; + v.extend_from_slice(domain.as_bytes()); + v.extend_from_slice(&port.to_be_bytes()); + v + } + } +} + +// --------------------------------------------------------------------------- +// State machine / Handshake +// --------------------------------------------------------------------------- + +/// Represents the current state of a SOCKS5 session. +#[allow(dead_code)] +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Socks5State { + /// Waiting for client greeting + ExpectGreeting, + /// Waiting for auth request (after choosing username/password auth) + ExpectAuth, + /// Authenticated, waiting for CONNECT request + ExpectRequest, + /// CONNECT request received and parsed + RequestReceived(Socks5Target), + /// Connection established, data forwarding in progress + Established, + /// Connection closed / error + Closed, +} + +/// Full handshake result after parsing greeting + auth + request. +#[allow(dead_code)] +#[derive(Debug)] +pub struct Socks5HandshakeResult { + pub username: String, + pub target: Socks5Target, +} + +// --------------------------------------------------------------------------- +// Integration helpers +// --------------------------------------------------------------------------- + +/// Convert a Socks5Target to a display string. +pub fn target_to_string(target: &Socks5Target) -> String { + match target { + Socks5Target::Ipv4(addr, port) => format!("{}:{}", addr, port), + Socks5Target::Domain { domain, port } => format!("{}:{}", domain, port), + } +} + +/// Convert a Socks5Error to a SOCKS5 reply code for the wire response. +#[allow(dead_code)] +pub fn error_to_reply_code(err: &Socks5Error) -> u8 { + match err { + Socks5Error::ConnectionNotAllowed => REPL_CONN_NOT_ALLOWED, + Socks5Error::NetworkUnreachable => REPL_NETWORK_UNREACHABLE, + Socks5Error::HostUnreachable => REPL_HOST_UNREACHABLE, + Socks5Error::ConnectionRefused => REPL_CONN_REFUSED, + _ => REPL_GENERAL_FAILURE, + } +} + +/// Try to resolve a Socks5Target to a SocketAddr. +/// This is called on the listener side to open the target TCP connection. +pub async fn resolve_target(target: &Socks5Target) -> Result { + match target { + Socks5Target::Ipv4(addr, port) => Ok(SocketAddr::new((*addr).into(), *port)), + Socks5Target::Domain { domain, port } => { + let lookup = format!("{}:{}", domain, port); + let addrs: Vec = tokio::net::lookup_host(&lookup) + .await + .map_err(|e| Socks5Error::AddrParseError(e.to_string()))? + .collect(); + if addrs.is_empty() { + return Err(Socks5Error::HostUnreachable); + } + Ok(addrs[0]) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + // ----------------------------------------------------------------------- + // Greeting parsing + // ----------------------------------------------------------------------- + + #[test] + fn parse_greeting_no_auth() { + let buf = vec![VERSION, 0x01, AUTH_NONE]; + let result = parse_greeting(&buf).unwrap(); + assert_eq!(result, Some(AUTH_NONE)); + } + + #[test] + fn parse_greeting_username_password() { + let buf = vec![VERSION, 0x01, AUTH_USERNAME_PASSWORD]; + let result = parse_greeting(&buf).unwrap(); + assert_eq!(result, Some(AUTH_USERNAME_PASSWORD)); + } + + #[test] + fn parse_greeting_prefer_username_password() { + let buf = vec![VERSION, 0x02, AUTH_NONE, AUTH_USERNAME_PASSWORD]; + let result = parse_greeting(&buf).unwrap(); + assert_eq!(result, Some(AUTH_USERNAME_PASSWORD)); + } + + #[test] + fn parse_greeting_no_acceptable() { + let buf = vec![VERSION, 0x01, 0xFF]; + // 0xFF (NO_ACCEPTABLE) means the client has no acceptable methods + // parse_greeting should return Ok(None) to signal this + let result = parse_greeting(&buf).unwrap(); + assert_eq!(result, None); + } + + #[test] + fn parse_greeting_invalid_version() { + let buf = vec![0x04, 0x01, AUTH_NONE]; + let result = parse_greeting(&buf); + assert!(result.is_err()); + } + + #[test] + fn parse_greeting_too_short() { + let buf = vec![VERSION]; + let result = parse_greeting(&buf); + assert!(result.is_err()); + } + + #[test] + fn parse_greeting_zero_methods() { + let buf = vec![VERSION, 0x00]; + let result = parse_greeting(&buf); + assert!(result.is_err()); + } + + #[test] + fn parse_greeting_truncated() { + let buf = vec![VERSION, 0x03, AUTH_NONE]; + let result = parse_greeting(&buf); + assert!(result.is_err()); + } + + #[test] + fn build_greeting_contains_methods() { + let greeting = build_greeting(); + assert_eq!(greeting[0], VERSION); + assert_eq!(greeting[1], 0x02); + assert_eq!(greeting[2], AUTH_NONE); + assert_eq!(greeting[3], AUTH_USERNAME_PASSWORD); + } + + // ----------------------------------------------------------------------- + // Auth parsing + // ----------------------------------------------------------------------- + + #[test] + fn parse_auth_request_valid() { + let buf = vec![ + 0x01, // VER + 0x03, // ULEN + b'a', b'd', b'm', // UNAME + 0x03, // PLEN + b'p', b'a', b's', // PASSWD + ]; + let (user, pass) = parse_auth_request(&buf).unwrap(); + assert_eq!(user, "adm"); + assert_eq!(pass, "pas"); + } + + #[test] + fn parse_auth_request_empty_credentials() { + let buf = vec![0x01, 0x00, 0x00]; + let (user, pass) = parse_auth_request(&buf).unwrap(); + assert_eq!(user, ""); + assert_eq!(pass, ""); + } + + #[test] + fn parse_auth_request_too_short() { + let buf = vec![0x01]; + let result = parse_auth_request(&buf); + assert!(result.is_err()); + } + + #[test] + fn parse_auth_request_wrong_version() { + let buf = vec![0x02, 0x00, 0x00]; + let result = parse_auth_request(&buf); + assert!(result.is_err()); + } + + #[test] + fn build_auth_responses() { + assert_eq!(build_auth_success(), vec![0x01, 0x00]); + assert_eq!(build_auth_failure(), vec![0x01, 0x01]); + } + + // ----------------------------------------------------------------------- + // CONNECT request parsing + // ----------------------------------------------------------------------- + + #[test] + fn parse_connect_ipv4() { + let buf = vec![ + VERSION, + CMD_CONNECT, + 0x00, + ATYP_IPV4, + 127, + 0, + 0, + 1, // 127.0.0.1 + 0x10, + 0x55, // 4181 = 0x1055 + ]; + let (cmd, target) = parse_connect_request(&buf).unwrap(); + assert_eq!(cmd, CMD_CONNECT); + assert_eq!( + target, + Socks5Target::Ipv4(std::net::Ipv4Addr::new(127, 0, 0, 1), 4181) + ); + } + + #[test] + fn parse_connect_domain() { + let domain = "localhost"; + let mut buf = vec![VERSION, CMD_CONNECT, 0x00, ATYP_DOMAIN, domain.len() as u8]; + buf.extend_from_slice(domain.as_bytes()); + buf.extend_from_slice(&8080u16.to_be_bytes()); + + let (cmd, target) = parse_connect_request(&buf).unwrap(); + assert_eq!(cmd, CMD_CONNECT); + assert_eq!( + target, + Socks5Target::Domain { + domain: "localhost".to_string(), + port: 8080, + } + ); + } + + #[test] + fn parse_connect_unsupported_command() { + let buf = vec![VERSION, CMD_BIND, 0x00, ATYP_IPV4, 127, 0, 0, 1, 0x00, 0x50]; + let result = parse_connect_request(&buf); + assert!(result.is_err()); + } + + #[test] + fn parse_connect_unsupported_addr_type() { + let buf = vec![VERSION, CMD_CONNECT, 0x00, ATYP_IPV6]; + let result = parse_connect_request(&buf); + assert!(result.is_err()); + } + + #[test] + fn parse_connect_too_short() { + let buf = vec![VERSION, CMD_CONNECT]; + let result = parse_connect_request(&buf); + assert!(result.is_err()); + } + + #[test] + fn parse_connect_invalid_version() { + let buf = vec![0x04, CMD_CONNECT, 0x00, ATYP_IPV4, 127, 0, 0, 1, 0x00, 0x50]; + let result = parse_connect_request(&buf); + assert!(result.is_err()); + } + + // ----------------------------------------------------------------------- + // Reply building + // ----------------------------------------------------------------------- + + #[test] + fn build_reply_success_format() { + let reply = build_reply_success(); + assert_eq!(reply[0], VERSION); + assert_eq!(reply[1], REPL_SUCCEEDED); + assert_eq!(reply[2], 0x00); // RSV + assert_eq!(reply[3], ATYP_IPV4); + assert_eq!(reply.len(), 10); + } + + #[test] + fn build_reply_failure_format() { + let reply = build_reply_failure(REPL_CONN_REFUSED); + assert_eq!(reply[0], VERSION); + assert_eq!(reply[1], REPL_CONN_REFUSED); + assert_eq!(reply.len(), 10); + } + + // ----------------------------------------------------------------------- + // Target encoding + // ----------------------------------------------------------------------- + + #[test] + fn encode_target_ipv4_roundtrip() { + let target = Socks5Target::Ipv4(std::net::Ipv4Addr::new(127, 0, 0, 1), 4181); + let encoded = encode_target(&target); + assert_eq!(encoded[0], ATYP_IPV4); + assert_eq!(encoded[1..5], [127, 0, 0, 1]); + assert_eq!(u16::from_be_bytes([encoded[5], encoded[6]]), 4181); + } + + #[test] + fn encode_target_domain_roundtrip() { + let target = Socks5Target::Domain { + domain: "example.com".to_string(), + port: 443, + }; + let encoded = encode_target(&target); + assert_eq!(encoded[0], ATYP_DOMAIN); + assert_eq!(encoded[1], 11); // "example.com".len() + assert_eq!(&encoded[2..13], b"example.com"); + assert_eq!(u16::from_be_bytes([encoded[13], encoded[14]]), 443); + } + + // ----------------------------------------------------------------------- + // State machine + // ----------------------------------------------------------------------- + + #[test] + fn state_transitions() { + let state = Socks5State::ExpectGreeting; + assert_eq!(state, Socks5State::ExpectGreeting); + + let state = Socks5State::ExpectAuth; + assert_eq!(state, Socks5State::ExpectAuth); + + let target = Socks5Target::Ipv4(std::net::Ipv4Addr::new(127, 0, 0, 1), 4181); + let state = Socks5State::RequestReceived(target.clone()); + assert_eq!(state, Socks5State::RequestReceived(target)); + } + + // ----------------------------------------------------------------------- + // Helper functions + // ----------------------------------------------------------------------- + + #[test] + fn target_to_string_ipv4() { + let target = Socks5Target::Ipv4(std::net::Ipv4Addr::new(127, 0, 0, 1), 4181); + assert_eq!(target_to_string(&target), "127.0.0.1:4181"); + } + + #[test] + fn target_to_string_domain() { + let target = Socks5Target::Domain { + domain: "localhost".to_string(), + port: 8080, + }; + assert_eq!(target_to_string(&target), "localhost:8080"); + } + + #[test] + fn error_reply_code_mapping() { + assert_eq!( + error_to_reply_code(&Socks5Error::ConnectionRefused), + REPL_CONN_REFUSED + ); + assert_eq!( + error_to_reply_code(&Socks5Error::NetworkUnreachable), + REPL_NETWORK_UNREACHABLE + ); + assert_eq!( + error_to_reply_code(&Socks5Error::Protocol("x".into())), + REPL_GENERAL_FAILURE + ); + } + + #[test] + fn socks5_target_host_and_port() { + let ipv4 = Socks5Target::Ipv4(std::net::Ipv4Addr::new(127, 0, 0, 1), 4181); + assert_eq!(ipv4.host(), "127.0.0.1"); + assert_eq!(ipv4.port(), 4181); + + let domain = Socks5Target::Domain { + domain: "localhost".to_string(), + port: 8080, + }; + assert_eq!(domain.host(), "localhost"); + assert_eq!(domain.port(), 8080); + } + + #[tokio::test] + async fn resolve_target_ipv4() { + let target = Socks5Target::Ipv4(std::net::Ipv4Addr::new(127, 0, 0, 1), 80); + let addr = resolve_target(&target).await.unwrap(); + assert_eq!(addr.ip().to_string(), "127.0.0.1"); + assert_eq!(addr.port(), 80); + } + + #[tokio::test] + async fn resolve_target_domain() { + let target = Socks5Target::Domain { + domain: "localhost".to_string(), + port: 80, + }; + let addr = resolve_target(&target).await.unwrap(); + assert_eq!(addr.port(), 80); + } + + #[tokio::test] + async fn resolve_target_invalid_domain() { + let target = Socks5Target::Domain { + domain: "this.host.does.not.exist.zzz".to_string(), + port: 80, + }; + let result = resolve_target(&target).await; + assert!(result.is_err()); + } + + // ----------------------------------------------------------------------- + // Malformed input tests (VAL-SOCKS-008) + // ----------------------------------------------------------------------- + + #[test] + fn malformed_greeting_empty() { + assert!(parse_greeting(&[]).is_err()); + } + + #[test] + fn malformed_greeting_random_bytes() { + assert!(parse_greeting(&[0xFF, 0xFF, 0xFF]).is_err()); + } + + #[test] + fn malformed_connect_empty() { + assert!(parse_connect_request(&[]).is_err()); + } + + #[test] + fn malformed_connect_random_bytes() { + assert!(parse_connect_request(&[0xFF, 0xFF, 0xFF]).is_err()); + } + + #[test] + fn malformed_auth_empty() { + assert!(parse_auth_request(&[]).is_err()); + } + + #[test] + fn malformed_auth_random_bytes() { + assert!(parse_auth_request(&[0xFF, 0xFF]).is_err()); + } +} diff --git a/src/tunnel.rs b/src/tunnel.rs index 8f79af5..cc31eb9 100644 --- a/src/tunnel.rs +++ b/src/tunnel.rs @@ -1,4 +1,4 @@ -/// HTTPS mTLS tunnel with application-level authentication. +/// HTTPS mTLS tunnel with application-level authentication and SOCKS5 forwarding. /// /// Provides a real HTTPS listener at `/tunnel` endpoint and a connector that /// connects over HTTPS with mutual TLS and an additional auth/session token. @@ -9,6 +9,11 @@ /// 3. Application auth via request header `X-Rustunnel-Token` /// 4. Connection stays alive for session maintenance /// +/// SOCKS5 Forwarding Protocol: +/// After auth, the connector opens a local SOCKS5 listener. SOCKS5 CONNECT +/// requests are forwarded via POST /forward over the authenticated tunnel. +/// The listener opens the target TCP connection and streams data bidirectionally. +/// /// Security gates (all must pass before forwarding): /// 1. HTTPS connection (hyper TLS acceptor) /// 2. Server certificate validation (client side) @@ -17,16 +22,19 @@ use std::net::SocketAddr; use std::path::Path; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; -use http_body_util::{BodyExt, Full}; +use http_body_util::{BodyExt, Full, combinators::BoxBody}; use hyper::body::{Bytes, Incoming}; use hyper::service::service_fn; use hyper::{Request, Response, StatusCode}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use tokio::sync::{Mutex, Notify}; use crate::errors::{AuthError, HostError, TunnelError}; use crate::redact::Redacted; +use crate::socks5; use crate::tls; // --------------------------------------------------------------------------- @@ -190,7 +198,7 @@ pub async fn run_listener(config: ListenerConfig) -> Result<(), TunnelError> { } /// Handle a single incoming HTTPS connection for the listener. -/// Accepts TLS, then serves hyper HTTP requests — only `/tunnel` POST is valid. +/// Accepts TLS, then serves hyper HTTP requests — `/tunnel` POST for auth, `/forward` POST for data. async fn handle_listener_https( stream: tokio::net::TcpStream, peer_addr: SocketAddr, @@ -228,23 +236,24 @@ async fn handle_listener_https( let auth = auth_state.clone(); let pa = peer_addr; async move { - type ResBody = Full; + type ResBody = BoxBody; + fn body(data: impl Into) -> ResBody { + BoxBody::new(Full::new(data.into()).map_err(|_| unreachable!())) + } match (req.method(), req.uri().path()) { (&hyper::Method::POST, "/tunnel") => { - // Auth via header let received = req .headers() .get("X-Rustunnel-Token") .and_then(|v| v.to_str().ok()) .unwrap_or(""); if constant_time_compare(received, &auth) { - // Auth OK — respond with tunnel established tracing::info!("Tunnel established with {} via HTTPS /tunnel", pa); Ok::, hyper::http::Error>( Response::builder() .status(StatusCode::OK) .header("X-Rustunnel-Status", "authenticated") - .body(Full::new(Bytes::from("OK"))) + .body(body("OK")) .unwrap(), ) } else { @@ -252,24 +261,21 @@ async fn handle_listener_https( Ok::, hyper::http::Error>( Response::builder() .status(StatusCode::UNAUTHORIZED) - .body(Full::new(Bytes::from("FAIL: invalid token"))) + .body(body("FAIL: invalid token")) .unwrap(), ) } } - (&hyper::Method::GET, "/tunnel") => { - // Allow GET to check tunnel status - Ok::, hyper::http::Error>( - Response::builder() - .status(StatusCode::OK) - .body(Full::new(Bytes::from("tunnel ready"))) - .unwrap(), - ) - } + (&hyper::Method::POST, "/forward") => handle_forward_request(req, pa).await, + (&hyper::Method::GET, "/tunnel") => Ok::, hyper::http::Error>( + Response::builder() + .status(StatusCode::OK) + .body(body("tunnel ready")) + .unwrap(), + ), _ => { - // Unknown path/method — reject tracing::debug!( - "Rejected {} {} from {} — only POST /tunnel accepted", + "Rejected {} {} from {} — only POST /tunnel and POST /forward accepted", req.method(), req.uri().path(), pa @@ -277,7 +283,7 @@ async fn handle_listener_https( Ok::, hyper::http::Error>( Response::builder() .status(StatusCode::NOT_FOUND) - .body(Full::new(Bytes::from("not found"))) + .body(body("not found")) .unwrap(), ) } @@ -296,6 +302,232 @@ async fn handle_listener_https( Ok(()) } +/// Parse a forward target string "host:port" into (host, port). +fn parse_forward_target(s: &str) -> Result<(String, u16), String> { + let parts: Vec<&str> = s.rsplitn(2, ':').collect(); + if parts.len() != 2 { + return Err(format!( + "invalid target format '{}': expected 'host:port'", + s + )); + } + let port: u16 = parts[0] + .parse() + .map_err(|_| format!("invalid port '{}' in target '{}'", parts[0], s))?; + Ok((parts[1].to_string(), port)) +} + +/// Handle a /forward request: open target connection, send request body, return response body. +/// +/// Collects the request body, sends it to the target, collects the target response, +/// and returns it. This simpler approach works for all HTTP request/response patterns. +async fn handle_forward_request( + req: Request, + peer_addr: SocketAddr, +) -> Result>, hyper::http::Error> { + fn body(data: impl Into) -> BoxBody { + BoxBody::new(Full::new(data.into()).map_err(|_| unreachable!())) + } + + let target_str = req + .headers() + .get("X-Rustunnel-Target") + .and_then(|v| v.to_str().ok()) + .unwrap_or(""); + let stream_id = req + .headers() + .get("X-Rustunnel-Stream-ID") + .and_then(|v| v.to_str().ok()) + .unwrap_or("unknown") + .to_string(); + + if target_str.is_empty() { + tracing::warn!("Forward request from {} missing target", peer_addr); + return Ok(Response::builder() + .status(StatusCode::BAD_REQUEST) + .body(body("missing target")) + .unwrap()); + } + + let (fwd_host, fwd_port) = match parse_forward_target(target_str) { + Ok(hp) => hp, + Err(e) => { + tracing::warn!("Forward request from {} invalid target: {}", peer_addr, e); + return Ok(Response::builder() + .status(StatusCode::BAD_REQUEST) + .body(body("invalid target")) + .unwrap()); + } + }; + + tracing::info!( + "Forward stream {}: {} -> {}:{} (from {})", + stream_id, + peer_addr, + fwd_host, + fwd_port, + peer_addr + ); + + // Collect request body + let body_bytes = match req.into_body().collect().await { + Ok(collected) => collected.to_bytes(), + Err(e) => { + tracing::warn!( + "Forward request from {} body collect error: {}", + peer_addr, + e + ); + return Ok(Response::builder() + .status(StatusCode::BAD_REQUEST) + .body(body("body collect error")) + .unwrap()); + } + }; + + // Forward to target + let result = forward_to_target(&fwd_host, fwd_port, &stream_id, body_bytes).await; + + match result { + Ok((status, headers, resp_body)) => { + let mut builder = Response::builder().status(status); + for (name, value) in headers { + builder = builder.header(name, value); + } + Ok(builder + .body(BoxBody::new( + Full::new(resp_body).map_err(|_| unreachable!()), + )) + .unwrap()) + } + Err(e) => { + tracing::warn!("Forward stream {} error: {}", stream_id, e); + Ok(Response::builder() + .status(StatusCode::BAD_GATEWAY) + .body(body(format!("forward error: {}", e))) + .unwrap()) + } + } +} + +/// Forward request data to a target TCP server and return the response. +async fn forward_to_target( + host: &str, + port: u16, + stream_id: &str, + request_data: Bytes, +) -> Result<(StatusCode, Vec<(String, String)>, Bytes), String> { + // Resolve target + let target_addr = resolve_target(host, port) + .await + .map_err(|e| format!("resolve: {}", e))?; + + // Connect to target + let mut target_stream = TcpStream::connect(target_addr).await.map_err(|e| { + if e.kind() == std::io::ErrorKind::ConnectionRefused { + format!("connection refused to {}", target_addr) + } else { + format!("connect to {}: {}", target_addr, e) + } + })?; + + tracing::info!( + "Forward stream {}: connected to target {}", + stream_id, + target_addr + ); + + // Write request data to target + if !request_data.is_empty() { + target_stream + .write_all(&request_data) + .await + .map_err(|e| format!("write to target: {}", e))?; + } + + // Read target response + let mut response_data = Vec::new(); + let mut buf = vec![0u8; 32 * 1024]; + loop { + match target_stream.read(&mut buf).await { + Err(e) => { + // Connection reset or other error - if we have data, return it + if !response_data.is_empty() { + tracing::info!( + "Forward stream {}: target connection closed with data read", + stream_id + ); + break; + } + return Err(format!("read from target: {}", e)); + } + Ok(0) => { + // EOF - target closed connection + break; + } + Ok(n) => { + response_data.extend_from_slice(&buf[..n]); + } + } + } + + tracing::info!( + "Forward stream {}: completed, {} bytes response", + stream_id, + response_data.len() + ); + + // Parse HTTP response from the raw data + parse_http_response(Bytes::from(response_data)) +} + +/// Parse raw HTTP response data into status, headers, and body. +#[allow(clippy::type_complexity)] +fn parse_http_response(data: Bytes) -> Result<(StatusCode, Vec<(String, String)>, Bytes), String> { + if data.is_empty() { + return Ok((StatusCode::OK, Vec::new(), Bytes::new())); + } + + // Find the end of headers ("\r\n\r\n") + let header_end = data.windows(4).position(|w| w == b"\r\n\r\n"); + match header_end { + Some(pos) => { + let header_text = String::from_utf8_lossy(&data[..pos + 4]); + let body = data.slice(pos + 4..); + + let mut lines = header_text.lines(); + let status_line = lines.next().ok_or("missing status line")?; + + // Parse "HTTP/1.1 200 OK" + let parts: Vec<&str> = status_line.split_whitespace().collect(); + if parts.len() < 2 { + return Err(format!("invalid status line: {}", status_line)); + } + let status_code: u16 = parts[1] + .parse() + .map_err(|e| format!("invalid status code: {}", e))?; + let status = + StatusCode::from_u16(status_code).map_err(|e| format!("invalid status: {}", e))?; + + let mut headers = Vec::new(); + for line in lines { + if line.is_empty() { + continue; + } + if let Some((name, value)) = line.split_once(':') { + headers.push((name.trim().to_string(), value.trim().to_string())); + } + } + + Ok((status, headers, body)) + } + None => { + // No header/body boundary found - return raw data + Ok((StatusCode::OK, Vec::new(), data)) + } + } +} + // --------------------------------------------------------------------------- // Connector (stays alive after auth) // --------------------------------------------------------------------------- @@ -542,6 +774,557 @@ pub async fn reconnect_tunnel(config: &ConnectorConfig) -> Result<(), TunnelErro connect_tunnel(config.clone()).await } +/// Run the connector with SOCKS5 proxy: authenticate through tunnel, then start SOCKS5. +/// +/// This is the main entry point for the `rustunnel connect` command. +/// It: +/// 1. Establishes the HTTPS mTLS tunnel connection +/// 2. Authenticates via POST /tunnel +/// 3. Marks auth as complete +/// 4. Starts the SOCKS5 proxy listener +pub async fn run_connector_with_socks( + config: ConnectorConfig, + socks_addr: SocketAddr, +) -> Result<(), TunnelError> { + let auth_complete = Arc::new(AtomicBool::new(false)); + let auth_clone = auth_complete.clone(); + + // First, authenticate through the tunnel + { + let client_config = tls::build_client_config( + config.client_cert_path.as_ref(), + config.client_key_path.as_ref(), + config.ca_cert_path.as_ref(), + ) + .map_err(|e| { + tracing::error!("Failed to build client TLS config: {}", e); + TunnelError::Tls(e) + })?; + + let tls_connector = tokio_rustls::TlsConnector::from(client_config); + + let target_addr = resolve_target(&config.target_host, config.target_port).await?; + + let server_name = + tls::server_name_from_host(&config.target_host).map_err(TunnelError::Tls)?; + + let stream = TcpStream::connect(target_addr).await.map_err(|e| { + if e.kind() == std::io::ErrorKind::ConnectionRefused { + TunnelError::ConnectionRefused(target_addr.to_string()) + } else { + TunnelError::Io(e) + } + })?; + + let tls_stream = tls_connector + .connect(server_name, stream) + .await + .map_err(|e| { + TunnelError::Tls(crate::errors::TlsError::VerificationFailed(format!( + "TLS handshake to {}:{} failed: {}", + config.target_host, config.target_port, e + ))) + })?; + + tracing::info!( + "TLS handshake completed with {}:{} via HTTPS", + config.target_host, + config.target_port + ); + + // Authenticate + let (_sender, _conn_handle) = + authenticate_https(tls_stream, &config.auth_token, &config.target_host).await?; + + tracing::info!( + "Tunnel session established with {}:{} via HTTPS /tunnel", + config.target_host, + config.target_port + ); + + // Mark auth as complete — SOCKS5 can now accept traffic + auth_clone.store(true, Ordering::SeqCst); + tracing::info!("Authentication complete — SOCKS5 proxy ready for traffic"); + } + + // Now start the SOCKS5 proxy (runs until shutdown) + run_socks5_proxy(socks_addr, config, auth_complete).await +} + +/// Run the SOCKS5 proxy listener and forward requests through the HTTPS tunnel. +/// +/// The SOCKS5 listener binds to the configured address. Each SOCKS5 CONNECT +/// request is forwarded through a fresh HTTPS connection to the tunnel listener. +/// This means each SOCKS stream gets its own TLS + auth cycle, which is simple +/// and correct for lab/dev usage. +/// +/// If `auth_complete` is not yet set when a connection arrives, the SOCKS5 +/// request is rejected (fail-closed). +pub async fn run_socks5_proxy( + socks_addr: SocketAddr, + tunnel_config: ConnectorConfig, + auth_complete: Arc, +) -> Result<(), TunnelError> { + // Bind SOCKS5 listener + let socks_listener = TcpListener::bind(socks_addr).await.map_err(|e| { + let addr_str = socks_addr.to_string(); + if e.kind() == std::io::ErrorKind::AddrInUse { + TunnelError::PortInUse(socks_addr.port()) + } else { + TunnelError::BindError(addr_str, e.to_string()) + } + })?; + + tracing::info!("SOCKS5 proxy listening on {}", socks_addr); + + // Register Ctrl-C handler for graceful shutdown + let shutdown = Arc::new(Notify::new()); + let shutdown_clone = shutdown.clone(); + tokio::spawn(async move { + tokio::signal::ctrl_c().await.ok(); + tracing::info!("Shutdown signal received — stopping SOCKS5 proxy"); + shutdown_clone.notify_one(); + }); + + loop { + let (stream, peer_addr) = tokio::select! { + result = socks_listener.accept() => { + match result { + Ok(s) => s, + Err(e) => { + tracing::error!("SOCKS5 accept error: {}", e); + continue; + } + } + } + _ = shutdown.notified() => { + tracing::info!("SOCKS5 proxy shutting down"); + return Ok(()); + } + }; + + let config = tunnel_config.clone(); + let auth = auth_complete.clone(); + + tokio::spawn(async move { + if let Err(e) = handle_socks5_connection(stream, peer_addr, &config, auth).await { + tracing::warn!("SOCKS5 connection from {} failed: {}", peer_addr, e); + } + }); + } +} + +/// Handle a single SOCKS5 connection: handshake, auth, CONNECT, forward. +async fn handle_socks5_connection( + mut stream: TcpStream, + peer_addr: SocketAddr, + tunnel_config: &ConnectorConfig, + auth_complete: Arc, +) -> Result<(), socks5::Socks5Error> { + tracing::info!("SOCKS5 connection from {}", peer_addr); + + // Check if tunnel is authenticated + if !auth_complete.load(Ordering::SeqCst) { + tracing::warn!( + "SOCKS5 request from {} rejected: tunnel not authenticated yet", + peer_addr + ); + stream + .write_all(&socks5::build_reply_failure(socks5::REPL_CONN_NOT_ALLOWED)) + .await?; + return Err(socks5::Socks5Error::NotAuthenticated); + } + + // --- SOCKS5 Greeting --- + let greeting = read_exact(&mut stream, 2).await?; + if greeting[0] != socks5::VERSION { + tracing::warn!("SOCKS5 bad version from {}", peer_addr); + stream + .write_all(&[socks5::VERSION, socks5::AUTH_NO_ACCEPTABLE]) + .await?; + return Err(socks5::Socks5Error::InvalidVersion(greeting[0])); + } + let n_methods = greeting[1] as usize; + if n_methods == 0 || n_methods > 255 { + return Err(socks5::Socks5Error::InvalidAuthMethodCount(n_methods)); + } + let methods = read_exact(&mut stream, n_methods).await?; + + // Choose auth method directly from client methods list + let chosen = if methods.contains(&socks5::AUTH_USERNAME_PASSWORD) { + socks5::AUTH_USERNAME_PASSWORD + } else if methods.contains(&socks5::AUTH_NONE) { + socks5::AUTH_NONE + } else { + stream + .write_all(&[socks5::VERSION, socks5::AUTH_NO_ACCEPTABLE]) + .await?; + return Err(socks5::Socks5Error::UnsupportedAuthMethod( + methods.first().copied().unwrap_or(0xFF), + )); + }; + + // Send method selection response: VER + METHOD (2 bytes only) + stream.write_all(&[socks5::VERSION, chosen]).await?; + + // --- Auth (if username/password) --- + if chosen == socks5::AUTH_USERNAME_PASSWORD { + let auth_buf = read_exact(&mut stream, 2).await?; + if auth_buf[0] != 0x01 { + stream.write_all(&socks5::build_auth_failure()).await?; + return Err(socks5::Socks5Error::AuthFailed); + } + let ulen = auth_buf[1] as usize; + let _username = if ulen > 0 { + let uname_bytes = read_exact(&mut stream, ulen).await?; + String::from_utf8_lossy(&uname_bytes).to_string() + } else { + String::new() + }; + let plen_buf = read_exact(&mut stream, 1).await?; + let plen = plen_buf[0] as usize; + let _password = if plen > 0 { + let pass_bytes = read_exact(&mut stream, plen).await?; + String::from_utf8_lossy(&pass_bytes).to_string() + } else { + String::new() + }; + // Accept any credentials (SOCKS5 auth is handled by the tunnel) + stream.write_all(&socks5::build_auth_success()).await?; + tracing::info!("SOCKS5 auth accepted from {}", peer_addr); + } + + // --- CONNECT Request --- + let connect_req = read_exact(&mut stream, 4).await?; + if connect_req[0] != socks5::VERSION { + return Err(socks5::Socks5Error::InvalidVersion(connect_req[0])); + } + if connect_req[1] != socks5::CMD_CONNECT { + stream + .write_all(&socks5::build_reply_failure(socks5::REPL_CONN_NOT_ALLOWED)) + .await?; + return Err(socks5::Socks5Error::UnsupportedCommand(connect_req[1])); + } + + let atyp = connect_req[3]; + let target = match atyp { + socks5::ATYP_IPV4 => { + let addr_bytes = read_exact(&mut stream, 4).await?; + let ipv4 = + std::net::Ipv4Addr::new(addr_bytes[0], addr_bytes[1], addr_bytes[2], addr_bytes[3]); + let port_bytes = read_exact(&mut stream, 2).await?; + let port = u16::from_be_bytes([port_bytes[0], port_bytes[1]]); + socks5::Socks5Target::Ipv4(ipv4, port) + } + socks5::ATYP_DOMAIN => { + let len_byte = read_exact(&mut stream, 1).await?; + let domain_len = len_byte[0] as usize; + let domain_bytes = read_exact(&mut stream, domain_len).await?; + let domain = String::from_utf8_lossy(&domain_bytes).to_string(); + let port_bytes = read_exact(&mut stream, 2).await?; + let port = u16::from_be_bytes([port_bytes[0], port_bytes[1]]); + socks5::Socks5Target::Domain { domain, port } + } + _ => { + stream + .write_all(&socks5::build_reply_failure(socks5::REPL_GENERAL_FAILURE)) + .await?; + return Err(socks5::Socks5Error::UnsupportedAddrType(atyp)); + } + }; + + let target_str = socks5::target_to_string(&target); + tracing::info!("SOCKS5 CONNECT from {} to {}", peer_addr, target_str); + + // --- Forward through tunnel --- + match forward_socks_request(&mut stream, tunnel_config, &target).await { + Ok(()) => { + tracing::info!( + "SOCKS5 stream from {} to {} completed", + peer_addr, + target_str + ); + } + Err(e) => { + tracing::warn!( + "SOCKS5 stream from {} to {} failed: {}", + peer_addr, + target_str, + e + ); + // Send failure reply if we haven't already sent success + // (forward_socks_request handles sending the reply) + return Err(e); + } + } + + Ok(()) +} + +/// Forward a SOCKS5 request through the HTTPS tunnel. +/// +/// Opens a new HTTPS connection, authenticates, sends POST /forward, +/// and pipes data bidirectionally between the SOCKS5 client and target. +async fn forward_socks_request( + socks_stream: &mut TcpStream, + tunnel_config: &ConnectorConfig, + target: &socks5::Socks5Target, +) -> Result<(), socks5::Socks5Error> { + // Build client TLS config + let client_config = match crate::tls::build_client_config( + tunnel_config.client_cert_path.as_ref(), + tunnel_config.client_key_path.as_ref(), + tunnel_config.ca_cert_path.as_ref(), + ) { + Ok(c) => c, + Err(e) => { + tracing::error!("Failed to build client TLS config for forward: {}", e); + socks_stream + .write_all(&socks5::build_reply_failure(socks5::REPL_GENERAL_FAILURE)) + .await?; + return Err(socks5::Socks5Error::Protocol(format!( + "TLS config error: {}", + e + ))); + } + }; + + let tls_connector = tokio_rustls::TlsConnector::from(client_config); + + // Resolve and connect to tunnel listener + let target_addr = + match crate::tunnel::resolve_target(&tunnel_config.target_host, tunnel_config.target_port) + .await + { + Ok(a) => a, + Err(e) => { + tracing::error!("Resolve failed: {}", e); + socks_stream + .write_all(&socks5::build_reply_failure(socks5::REPL_GENERAL_FAILURE)) + .await?; + return Err(socks5::Socks5Error::Protocol(format!("resolve: {}", e))); + } + }; + + let tcp_stream = match TcpStream::connect(target_addr).await { + Ok(s) => s, + Err(e) => { + tracing::error!("Connection to tunnel failed: {}", e); + socks_stream + .write_all(&socks5::build_reply_failure(socks5::REPL_GENERAL_FAILURE)) + .await?; + return Err(socks5::Socks5Error::ConnectionRefused); + } + }; + + let server_name = match crate::tls::server_name_from_host(&tunnel_config.target_host) { + Ok(n) => n, + Err(e) => { + tracing::error!("Invalid server name: {}", e); + socks_stream + .write_all(&socks5::build_reply_failure(socks5::REPL_GENERAL_FAILURE)) + .await?; + return Err(socks5::Socks5Error::Protocol(format!("server name: {}", e))); + } + }; + + let tls_stream = match tls_connector.connect(server_name, tcp_stream).await { + Ok(s) => s, + Err(e) => { + tracing::error!("TLS handshake failed: {}", e); + socks_stream + .write_all(&socks5::build_reply_failure(socks5::REPL_GENERAL_FAILURE)) + .await?; + return Err(socks5::Socks5Error::Protocol(format!("TLS: {}", e))); + } + }; + + // HTTP handshake + let io = hyper_util::rt::TokioIo::new(tls_stream); + let (mut sender, connection) = match hyper::client::conn::http1::handshake(io).await { + Ok(s) => s, + Err(e) => { + tracing::error!("HTTP handshake failed: {}", e); + socks_stream + .write_all(&socks5::build_reply_failure(socks5::REPL_GENERAL_FAILURE)) + .await?; + return Err(socks5::Socks5Error::Protocol(format!("HTTP: {}", e))); + } + }; + tokio::spawn(connection); + + // Authenticate + let auth_request = Request::builder() + .method("POST") + .uri("/tunnel") + .header("X-Rustunnel-Token", tunnel_config.auth_token.as_str()) + .header("Host", &tunnel_config.target_host) + .body(Full::new(Bytes::new())) + .map_err(|e| socks5::Socks5Error::Protocol(format!("build request: {}", e)))?; + + let auth_response = match sender.send_request(auth_request).await { + Ok(r) => r, + Err(e) => { + tracing::error!("Auth request failed: {}", e); + socks_stream + .write_all(&socks5::build_reply_failure(socks5::REPL_GENERAL_FAILURE)) + .await?; + return Err(socks5::Socks5Error::Protocol(format!("auth: {}", e))); + } + }; + + let auth_status = auth_response.status(); + let auth_body = auth_response + .into_body() + .collect() + .await + .map_err(|e| socks5::Socks5Error::Protocol(format!("read auth response: {}", e)))? + .to_bytes(); + + if auth_status != StatusCode::OK || !String::from_utf8_lossy(&auth_body).contains("OK") { + tracing::error!("Auth rejected for forward request"); + socks_stream + .write_all(&socks5::build_reply_failure(socks5::REPL_CONN_NOT_ALLOWED)) + .await?; + return Err(socks5::Socks5Error::ConnectionNotAllowed); + } + + // Send CONNECT success reply to SOCKS5 client + socks_stream + .write_all(&socks5::build_reply_success()) + .await?; + + // Generate stream ID for logging + let stream_id = format!( + "socks-{}-{}", + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_nanos(), + std::sync::atomic::AtomicU64::new(0).fetch_add(1, Ordering::SeqCst) + ); + let target_str = socks5::target_to_string(target); + + tracing::info!( + "SOCKS5 stream {}: connected to target {} via tunnel", + stream_id, + target_str + ); + + // Read HTTP request from SOCKS5 client + let (mut socks_read, mut socks_write) = socks_stream.split(); + let request_bytes = read_until_double_crlf(&mut socks_read).await?; + + // Forward to target and get response + let target_addr = socks5::resolve_target(target).await?; + let mut target_stream = TcpStream::connect(target_addr).await.map_err(|e| { + if e.kind() == std::io::ErrorKind::ConnectionRefused { + socks5::Socks5Error::ConnectionRefused + } else { + socks5::Socks5Error::Io(e) + } + })?; + + // Write request to target + target_stream + .write_all(&request_bytes) + .await + .map_err(socks5::Socks5Error::Io)?; + + // Read target response and forward to SOCKS5 client + let mut buf = vec![0u8; 32 * 1024]; + loop { + match target_stream.read(&mut buf).await { + Err(_) => break, + Ok(0) => break, + Ok(n) => { + if let Err(e) = socks_write.write_all(&buf[..n]).await { + tracing::debug!("Forward stream {} write error: {}", stream_id, e); + break; + } + } + } + } + + // Read any remaining data from SOCKS5 client (for POST bodies etc.) + let mut remaining = Vec::new(); + socks_read.read_to_end(&mut remaining).await.ok(); + if !remaining.is_empty() { + // Forward remaining data to target via another connection + let mut target_stream2 = TcpStream::connect(target_addr).await.ok(); + if let Some(ref mut ts) = target_stream2 { + ts.write_all(&remaining).await.ok(); + loop { + match ts.read(&mut buf).await { + Err(_) => break, + Ok(0) => break, + Ok(n) => { + socks_write.write_all(&buf[..n]).await.ok(); + } + } + } + } + } + + tracing::info!("Forward stream {}: completed", stream_id); + + Ok(()) +} + +/// Read exactly `len` bytes from the stream. +async fn read_exact(stream: &mut TcpStream, len: usize) -> Result, socks5::Socks5Error> { + let mut buf = vec![0u8; len]; + stream + .read_exact(&mut buf) + .await + .map_err(socks5::Socks5Error::Io)?; + Ok(buf) +} + +/// Read from the stream until we see "\r\n\r\n" (HTTP header terminator). +/// Returns all bytes read including the "\r\n\r\n". +async fn read_until_double_crlf( + stream: &mut (impl AsyncReadExt + Unpin), +) -> Result, socks5::Socks5Error> { + let mut buf = Vec::new(); + let mut chunk = vec![0u8; 4096]; + let mut found = false; + let mut crlf_count = 0; + + loop { + let n = stream + .read(&mut chunk) + .await + .map_err(socks5::Socks5Error::Io)?; + if n == 0 { + if buf.is_empty() { + return Err(socks5::Socks5Error::Protocol( + "EOF before HTTP headers".to_string(), + )); + } + return Ok(buf); + } + + for &b in &chunk[..n] { + buf.push(b); + if b == b'\r' || b == b'\n' { + crlf_count += 1; + if crlf_count >= 4 && buf.ends_with(b"\r\n\r\n") { + found = true; + break; + } + } else { + crlf_count = 0; + } + } + if found { + break; + } + } + + Ok(buf) +} + // --------------------------------------------------------------------------- // Security gate state machine // ---------------------------------------------------------------------------