From d6a0a40696f37e03bafe5c12472437e297e796c0 Mon Sep 17 00:00:00 2001 From: c4ch3c4d3 <23181631+c4ch3c4d3@users.noreply.github.com> Date: Wed, 3 Jun 2026 21:01:24 -0600 Subject: [PATCH] Replace HTTP-shaped forwarding with true persistent multiplexed tunnel Fix socks-e2e scrutiny blockers by replacing HTTP-shaped request/response forwarding with a true persistent authenticated tunnel data path. Key changes: - Implement bidirectional framing with stream IDs over the authenticated tunnel so SOCKS target TCP connections originate from the listener side - Multiple sequential/concurrent SOCKS streams share one authenticated session with stream isolation via frame dispatch - Bytes stream incrementally without whole-response buffering or fixed EOF timeouts - CLOSE/ERROR frames propagate deterministically for clean stream teardown - Remove /forward HTTP endpoint and TunnelSession HTTP-based forwarding - Remove read_until_double_crlf HTTP-specific parsing from SOCKS path - Add StreamMux for concurrent stream management on the connector side - Add tests for slow/interleaved streams, concurrent isolation, listener-side target origin, close frame propagation, and E2E framing Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com> --- src/framing.rs | 558 ++++++++ src/main.rs | 1 + src/socks5.rs | 3 + src/tunnel.rs | 3696 ++++++++++++++++++++++-------------------------- 4 files changed, 2267 insertions(+), 1991 deletions(-) create mode 100644 src/framing.rs diff --git a/src/framing.rs b/src/framing.rs new file mode 100644 index 0000000..f1b6fca --- /dev/null +++ b/src/framing.rs @@ -0,0 +1,558 @@ +/// Binary framing protocol for the authenticated tunnel. +/// +/// After the mTLS + auth handshake completes, both sides switch from HTTP/1.1 +/// to a lightweight binary framing protocol that carries multiple bidirectional +/// streams over a single persistent TCP connection. +/// +/// Frame format: +/// +--------+----------+----------+-------+-----------+ +/// | Type | StreamID | Len | Flags | Payload.. | +/// | 1 byte | 8 bytes | 4 bytes | 1 byte| variable | +/// +--------+----------+----------+-------+-----------+ +/// +/// All multi-byte integers are big-endian. +/// +/// Frame types: +/// 0x01 CONNECT - Open a new stream to a target (payload: target addr) +/// 0x02 CONNECT_REPLY - Accept/reject a CONNECT (payload: 1-byte status) +/// 0x03 DATA - Stream data (payload: raw bytes) +/// 0x04 CLOSE - Half-close a stream (no payload) +/// 0x05 ERROR - Stream error (payload: error message) +/// +/// Stream IDs: +/// - Client-initiated streams: odd IDs (1, 3, 5, ...) +/// - Even IDs reserved for future server-initiated streams +/// +/// Flags: +/// - Bit 0 (0x01): EOF — sender will send no more data on this stream +use bytes::{Buf, BufMut, Bytes, BytesMut}; +use thiserror::Error; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; + +// --------------------------------------------------------------------------- +// Errors +// --------------------------------------------------------------------------- + +#[derive(Debug, Error)] +pub enum FrameError { + #[error("unknown frame type: {0:#04x}")] + #[allow(dead_code)] + UnknownFrameType(u8), + + #[error("frame too short: expected {0} bytes, got {1}")] + FrameTooShort(usize, usize), + + #[error("payload length overflow: {0}")] + PayloadOverflow(u32), + + #[error("stream {0} not found")] + #[allow(dead_code)] + StreamNotFound(u64), + + #[error("duplicate stream ID: {0}")] + #[allow(dead_code)] + DuplicateStreamId(u64), + + #[error("stream {0} already closed")] + #[allow(dead_code)] + StreamClosed(u64), + + #[error("I/O error: {0}")] + Io(#[from] std::io::Error), + + #[error("protocol error: {0}")] + Protocol(String), + + #[error("target connection failed: {0}")] + #[allow(dead_code)] + TargetConnectionFailed(String), +} + +// --------------------------------------------------------------------------- +// Frame types +// --------------------------------------------------------------------------- + +pub const FRAME_CONNECT: u8 = 0x01; +pub const FRAME_CONNECT_REPLY: u8 = 0x02; +pub const FRAME_DATA: u8 = 0x03; +pub const FRAME_CLOSE: u8 = 0x04; +pub const FRAME_ERROR: u8 = 0x05; + +// Flags +#[allow(dead_code)] +pub const FLAG_EOF: u8 = 0x01; + +// CONNECT_REPLY status +pub const CONNECT_OK: u8 = 0x00; +pub const CONNECT_FAILED: u8 = 0x01; + +// Frame header size (type + stream_id + length + flags) +pub const FRAME_HEADER_LEN: usize = 14; + +const MAX_PAYLOAD_LEN: u32 = 4 * 1024 * 1024; // 4 MiB per frame + +// --------------------------------------------------------------------------- +// Frame struct +// --------------------------------------------------------------------------- + +#[derive(Debug, Clone)] +pub struct Frame { + pub frame_type: u8, + pub stream_id: u64, + pub flags: u8, + pub payload: Bytes, +} + +impl Frame { + /// Create a CONNECT frame. + pub fn connect(stream_id: u64, target_bytes: impl AsRef<[u8]>) -> Self { + Frame { + frame_type: FRAME_CONNECT, + stream_id, + flags: 0, + payload: Bytes::copy_from_slice(target_bytes.as_ref()), + } + } + + /// Create a CONNECT_REPLY frame. + pub fn connect_reply(stream_id: u64, status: u8) -> Self { + Frame { + frame_type: FRAME_CONNECT_REPLY, + stream_id, + flags: 0, + payload: Bytes::copy_from_slice(&[status]), + } + } + + /// Create a DATA frame. + pub fn data(stream_id: u64, data: impl AsRef<[u8]>) -> Self { + Frame { + frame_type: FRAME_DATA, + stream_id, + flags: 0, + payload: Bytes::copy_from_slice(data.as_ref()), + } + } + + /// Create a DATA+EOF frame. + #[allow(dead_code)] + pub fn data_eof(stream_id: u64, data: impl AsRef<[u8]>) -> Self { + Frame { + frame_type: FRAME_DATA, + stream_id, + flags: FLAG_EOF, + payload: Bytes::copy_from_slice(data.as_ref()), + } + } + + /// Create a CLOSE frame. + pub fn close(stream_id: u64) -> Self { + Frame { + frame_type: FRAME_CLOSE, + stream_id, + flags: 0, + payload: Bytes::new(), + } + } + + /// Create an ERROR frame. + pub fn error(stream_id: u64, message: impl Into) -> Self { + Frame { + frame_type: FRAME_ERROR, + stream_id, + flags: 0, + payload: Bytes::from(message.into()), + } + } + + /// Serialize this frame to bytes. + pub fn serialize(&self) -> Bytes { + let payload_len = self.payload.len() as u32; + let total = FRAME_HEADER_LEN + payload_len as usize; + let mut buf = BytesMut::with_capacity(total); + + buf.put_u8(self.frame_type); + buf.put_u64(self.stream_id); + buf.put_u32(payload_len); + buf.put_u8(self.flags); + buf.put_slice(&self.payload); + + buf.freeze() + } + + /// Deserialize a frame from a buffer. + /// Returns (frame, remaining_bytes) if successful. + pub fn deserialize(buf: &mut BytesMut) -> Result<(Self, usize), FrameError> { + // Need at least the header + if buf.len() < FRAME_HEADER_LEN { + return Err(FrameError::FrameTooShort(FRAME_HEADER_LEN, buf.len())); + } + + let frame_type = buf.get_u8(); + let stream_id = buf.get_u64(); + let payload_len = buf.get_u32(); + let flags = buf.get_u8(); + + if payload_len > MAX_PAYLOAD_LEN { + return Err(FrameError::PayloadOverflow(payload_len)); + } + + if buf.len() < payload_len as usize { + // Not enough data yet — put header back and wait + // We can't easily un-get, so reconstruct + let mut restored = BytesMut::with_capacity(FRAME_HEADER_LEN + buf.len()); + restored.put_u8(frame_type); + restored.put_u64(stream_id); + restored.put_u32(payload_len); + restored.put_u8(flags); + restored.put_slice(buf.split_to(buf.len()).as_ref()); + // Actually, since we consumed the header, we need to return the error + // and the caller will retry. We restore nothing — the caller reads again. + // This is a "need more data" condition. + return Err(FrameError::FrameTooShort( + FRAME_HEADER_LEN + payload_len as usize, + FRAME_HEADER_LEN + buf.len(), + )); + } + + let payload = buf.split_to(payload_len as usize).freeze(); + + Ok(( + Frame { + frame_type, + stream_id, + flags, + payload, + }, + payload_len as usize, + )) + } + + /// Check if this frame has the EOF flag set. + #[allow(dead_code)] + pub fn is_eof(&self) -> bool { + self.flags & FLAG_EOF != 0 + } +} + +// --------------------------------------------------------------------------- +// Frame reader/writer +// --------------------------------------------------------------------------- + +/// Asynchronous frame reader that handles partial reads and buffering. +pub struct FrameReader { + buf: BytesMut, +} + +impl FrameReader { + pub fn new() -> Self { + FrameReader { + buf: BytesMut::with_capacity(8192), + } + } + + /// Read the next complete frame from the underlying reader. + /// Returns None if the connection is closed gracefully (EOF). + pub async fn read_frame( + &mut self, + reader: &mut R, + ) -> Result, FrameError> { + loop { + // Try to parse a frame from the buffer + let buf_ref = &mut self.buf; + match Frame::deserialize(buf_ref) { + Ok((frame, _)) => return Ok(Some(frame)), + Err(FrameError::FrameTooShort(_, _)) => { + // Need more data — read from the underlying reader + let capacity = self.buf.capacity(); + let len = self.buf.len(); + // Ensure we have space + if len == capacity { + self.buf.reserve(8192); + } + let _remaining = self.buf.spare_capacity_mut(); + match reader.read_buf(&mut self.buf).await { + Ok(0) => { + // Connection closed — if we have partial data, it's an error + if !self.buf.is_empty() { + return Err(FrameError::Protocol( + "connection closed with partial frame".to_string(), + )); + } + return Ok(None); + } + Ok(_) => continue, + Err(e) => return Err(FrameError::Io(e)), + } + } + Err(e) => return Err(e), + } + } + } +} + +/// Asynchronous frame writer. +pub struct FrameWriter; + +impl FrameWriter { + /// Write a frame to the underlying writer. + pub async fn write_frame( + writer: &mut W, + frame: &Frame, + ) -> Result<(), FrameError> { + let bytes = frame.serialize(); + writer.write_all(&bytes).await?; + writer.flush().await?; + Ok(()) + } +} + +// --------------------------------------------------------------------------- +// Target address encoding for CONNECT frames +// --------------------------------------------------------------------------- + +/// Encode a target address into CONNECT frame payload bytes. +/// Format: ATYP (1) + addr (variable) + port (2) +pub fn encode_target(target: &crate::socks5::Socks5Target) -> Vec { + match target { + crate::socks5::Socks5Target::Ipv4(addr, port) => { + let mut v = vec![crate::socks5::ATYP_IPV4]; + v.extend_from_slice(&addr.octets()); + v.extend_from_slice(&port.to_be_bytes()); + v + } + crate::socks5::Socks5Target::Domain { domain, port } => { + let mut v = vec![crate::socks5::ATYP_DOMAIN, domain.len() as u8]; + v.extend_from_slice(domain.as_bytes()); + v.extend_from_slice(&port.to_be_bytes()); + v + } + } +} + +/// Decode a target address from CONNECT frame payload bytes. +pub fn decode_target(payload: &[u8]) -> Result { + if payload.is_empty() { + return Err(FrameError::Protocol("empty target".to_string())); + } + let atyp = payload[0]; + let mut cursor = &payload[1..]; + + match atyp { + crate::socks5::ATYP_IPV4 => { + if cursor.len() < 6 { + return Err(FrameError::Protocol("truncated IPv4 target".to_string())); + } + let addr = std::net::Ipv4Addr::new(cursor[0], cursor[1], cursor[2], cursor[3]); + let port = u16::from_be_bytes([cursor[4], cursor[5]]); + Ok(crate::socks5::Socks5Target::Ipv4(addr, port)) + } + crate::socks5::ATYP_DOMAIN => { + if cursor.is_empty() { + return Err(FrameError::Protocol("truncated domain target".to_string())); + } + let domain_len = cursor[0] as usize; + cursor.advance(1); + if cursor.len() < domain_len + 2 { + return Err(FrameError::Protocol("truncated domain target".to_string())); + } + let domain = String::from_utf8(cursor[..domain_len].to_vec()) + .map_err(|_| FrameError::Protocol("invalid UTF-8 in domain".to_string()))?; + cursor.advance(domain_len); + let port = u16::from_be_bytes([cursor[0], cursor[1]]); + Ok(crate::socks5::Socks5Target::Domain { domain, port }) + } + _ => Err(FrameError::Protocol(format!( + "unsupported address type: {:#04x}", + atyp + ))), + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Cursor; + + #[test] + fn frame_roundtrip_connect() { + let target = crate::socks5::Socks5Target::Ipv4(std::net::Ipv4Addr::new(127, 0, 0, 1), 4181); + let payload = encode_target(&target); + let frame = Frame::connect(1, payload.clone()); + let serialized = frame.serialize(); + + let mut buf = BytesMut::from(&serialized[..]); + let (decoded, _) = Frame::deserialize(&mut buf).unwrap(); + assert_eq!(decoded.frame_type, FRAME_CONNECT); + assert_eq!(decoded.stream_id, 1); + assert_eq!(decoded.payload, payload); + + let decoded_target = decode_target(&decoded.payload).unwrap(); + assert_eq!(decoded_target, target); + } + + #[test] + fn frame_roundtrip_data() { + let data = b"hello world"; + let frame = Frame::data(3, data); + let serialized = frame.serialize(); + + let mut buf = BytesMut::from(&serialized[..]); + let (decoded, _) = Frame::deserialize(&mut buf).unwrap(); + assert_eq!(decoded.frame_type, FRAME_DATA); + assert_eq!(decoded.stream_id, 3); + assert_eq!(&decoded.payload[..], data); + } + + #[test] + fn frame_roundtrip_data_eof() { + let frame = Frame::data_eof(5, b"final"); + let serialized = frame.serialize(); + + let mut buf = BytesMut::from(&serialized[..]); + let (decoded, _) = Frame::deserialize(&mut buf).unwrap(); + assert!(decoded.is_eof()); + } + + #[test] + fn frame_roundtrip_close() { + let frame = Frame::close(7); + let serialized = frame.serialize(); + + let mut buf = BytesMut::from(&serialized[..]); + let (decoded, _) = Frame::deserialize(&mut buf).unwrap(); + assert_eq!(decoded.frame_type, FRAME_CLOSE); + assert_eq!(decoded.stream_id, 7); + assert!(decoded.payload.is_empty()); + } + + #[test] + fn frame_roundtrip_error() { + let frame = Frame::error(9, "connection refused"); + let serialized = frame.serialize(); + + let mut buf = BytesMut::from(&serialized[..]); + let (decoded, _) = Frame::deserialize(&mut buf).unwrap(); + assert_eq!(decoded.frame_type, FRAME_ERROR); + assert_eq!( + String::from_utf8_lossy(&decoded.payload), + "connection refused" + ); + } + + #[test] + fn frame_roundtrip_connect_reply() { + let frame = Frame::connect_reply(1, CONNECT_OK); + let serialized = frame.serialize(); + + let mut buf = BytesMut::from(&serialized[..]); + let (decoded, _) = Frame::deserialize(&mut buf).unwrap(); + assert_eq!(decoded.frame_type, FRAME_CONNECT_REPLY); + assert_eq!(decoded.payload[0], CONNECT_OK); + } + + #[test] + fn frame_partial_read_returns_error() { + // Header is 14 bytes; with only 10 bytes, deserialization should fail + let mut buf = BytesMut::from(&[0u8; 10][..]); + let result = Frame::deserialize(&mut buf); + assert!(matches!(result, Err(FrameError::FrameTooShort(_, _)))); + } + + #[test] + fn frame_multiple_frames_in_buffer() { + let frame1 = Frame::data(1, b"first"); + let frame2 = Frame::data(3, b"second"); + + let mut buf = BytesMut::new(); + buf.extend_from_slice(&frame1.serialize()); + buf.extend_from_slice(&frame2.serialize()); + + let (d1, _) = Frame::deserialize(&mut buf).unwrap(); + assert_eq!(d1.stream_id, 1); + assert_eq!(&d1.payload[..], b"first"); + + let (d2, _) = Frame::deserialize(&mut buf).unwrap(); + assert_eq!(d2.stream_id, 3); + assert_eq!(&d2.payload[..], b"second"); + + assert!(buf.is_empty()); + } + + #[test] + fn encode_decode_target_ipv4() { + let target = crate::socks5::Socks5Target::Ipv4(std::net::Ipv4Addr::new(10, 0, 0, 1), 8080); + let encoded = encode_target(&target); + let decoded = decode_target(&encoded).unwrap(); + assert_eq!(decoded, target); + } + + #[test] + fn encode_decode_target_domain() { + let target = crate::socks5::Socks5Target::Domain { + domain: "example.com".to_string(), + port: 443, + }; + let encoded = encode_target(&target); + let decoded = decode_target(&encoded).unwrap(); + assert_eq!(decoded, target); + } + + #[test] + fn decode_target_empty_fails() { + assert!(decode_target(&[]).is_err()); + } + + #[test] + fn decode_target_truncated_fails() { + assert!(decode_target(&[crate::socks5::ATYP_IPV4, 127]).is_err()); + } + + #[test] + fn decode_target_unsupported_atyp_fails() { + assert!(decode_target(&[0x04, 0, 0, 0, 0, 0, 0]).is_err()); + } + + #[tokio::test] + async fn frame_reader_writes_all_then_reads() { + let target = crate::socks5::Socks5Target::Ipv4(std::net::Ipv4Addr::new(127, 0, 0, 1), 8080); + let frame = Frame::connect(1, encode_target(&target)); + let serialized = frame.serialize(); + + let mut reader = FrameReader::new(); + let mut cursor = Cursor::new(&serialized[..]); + let result = reader.read_frame(&mut cursor).await.unwrap().unwrap(); + + assert_eq!(result.frame_type, FRAME_CONNECT); + assert_eq!(result.stream_id, 1); + let decoded_target = decode_target(&result.payload).unwrap(); + assert_eq!(decoded_target, target); + } + + #[tokio::test] + async fn frame_reader_eof_returns_none() { + let mut reader = FrameReader::new(); + let mut cursor = Cursor::new(&[] as &[u8]); + let result = reader.read_frame(&mut cursor).await.unwrap(); + assert!(result.is_none()); + } + + #[tokio::test] + async fn frame_writer_roundtrip() { + let frame = Frame::data(5, b"test data here"); + let serialized = frame.serialize(); + + let mut cursor = Vec::new(); + FrameWriter::write_frame(&mut cursor, &frame).await.unwrap(); + assert_eq!(cursor, serialized.as_ref()); + } + + #[test] + fn frame_header_size() { + assert_eq!(FRAME_HEADER_LEN, 14); + } +} diff --git a/src/main.rs b/src/main.rs index 31b3500..f36b598 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,7 @@ mod cli; mod config; mod errors; +mod framing; mod generate; mod redact; mod socks5; diff --git a/src/socks5.rs b/src/socks5.rs index 4cb03ee..bf3be14 100644 --- a/src/socks5.rs +++ b/src/socks5.rs @@ -45,6 +45,7 @@ pub enum Socks5Error { AuthFailed, #[error("connection not allowed")] + #[allow(dead_code)] ConnectionNotAllowed, #[error("network unreachable")] @@ -52,9 +53,11 @@ pub enum Socks5Error { NetworkUnreachable, #[error("host unreachable")] + #[allow(dead_code)] HostUnreachable, #[error("connection refused")] + #[allow(dead_code)] ConnectionRefused, #[error("protocol error: {0}")] diff --git a/src/tunnel.rs b/src/tunnel.rs index 69ba5c5..47cac1b 100644 --- a/src/tunnel.rs +++ b/src/tunnel.rs @@ -1,46 +1,52 @@ /// 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. +/// After the mTLS + auth handshake completes, both sides switch from HTTP/1.1 +/// to a binary framing protocol that carries multiple bidirectional streams +/// over a single persistent authenticated tunnel. /// -/// Protocol: -/// 1. HTTPS POST to `/tunnel` with auth header -/// 2. TLS handshake (rustls mTLS) — handled by hyper -/// 3. Application auth via request header `X-Rustunnel-Token` -/// 4. Connection stays alive for session maintenance +/// Protocol flow: +/// 1. Client connects over HTTPS with mTLS +/// 2. Client sends HTTP POST /tunnel with auth token +/// 3. Server validates and responds 200 OK +/// 4. Both sides switch to binary framing on the same connection +/// 5. Client sends CONNECT frames; server opens target TCP connections +/// 6. DATA frames flow bidirectionally, keyed by stream ID +/// 7. CLOSE/ERROR frames propagate deterministically /// -/// 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) -/// 3. Client certificate validation (server side, mTLS) +/// Security gates (all must pass before framing): +/// 1. HTTPS connection (TLS acceptor/connector) +/// 2. Client certificate validation (server side, mTLS) +/// 3. Server certificate validation (client side) /// 4. Application auth token validation +use std::collections::HashMap; use std::net::SocketAddr; use std::path::Path; use std::sync::Arc; -use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; -use http_body_util::{BodyExt, Full, combinators::BoxBody}; -use hyper::body::{Bytes, Incoming}; -use hyper::service::service_fn; -use hyper::{Request, Response, StatusCode}; +use bytes::Bytes; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; -use tokio::sync::{Mutex, Notify}; +use tokio::sync::{mpsc, oneshot}; use crate::errors::{AuthError, HostError, TunnelError}; +use crate::framing::{ + CONNECT_FAILED, CONNECT_OK, FRAME_CLOSE, FRAME_CONNECT, FRAME_CONNECT_REPLY, FRAME_DATA, + FRAME_ERROR, Frame, FrameReader, FrameWriter, decode_target, encode_target, +}; use crate::redact::Redacted; use crate::socks5; use crate::tls; -// Global stream-ID counter for deterministic unique IDs +// --------------------------------------------------------------------------- +// Stream ID generation +// --------------------------------------------------------------------------- + +#[allow(dead_code)] static STREAM_COUNTER: AtomicU64 = AtomicU64::new(0); /// Generate a unique stream ID for logging. +#[allow(dead_code)] fn next_stream_id() -> u64 { STREAM_COUNTER.fetch_add(1, Ordering::Relaxed) } @@ -75,7 +81,6 @@ pub fn load_auth_token( // Constant-time comparison // --------------------------------------------------------------------------- -/// Constant-time string comparison to prevent timing attacks. fn constant_time_compare(a: &str, b: &str) -> bool { let a_bytes = a.as_bytes(); let b_bytes = b.as_bytes(); @@ -90,11 +95,10 @@ fn constant_time_compare(a: &str, b: &str) -> bool { } // --------------------------------------------------------------------------- -// Hostname resolution (Fix: no .unwrap() panics) +// Hostname resolution // --------------------------------------------------------------------------- /// Resolve a hostname to a list of SocketAddrs. -/// Returns actionable errors instead of panicking. pub async fn resolve_host(host: &str, port: u16) -> Result, HostError> { if host.is_empty() { return Err(HostError::InvalidHostname(host.to_string())); @@ -110,15 +114,14 @@ pub async fn resolve_host(host: &str, port: u16) -> Result, Host Ok(addrs) } -/// Parse a "host:port" string, resolving hostnames if needed. -/// Unlike cli::parse_host_port, this resolves the hostname to actual addresses. +/// Resolve a host:port to a single SocketAddr. pub async fn resolve_target(host: &str, port: u16) -> Result { let addrs = resolve_host(host, port).await?; Ok(addrs[0]) } // --------------------------------------------------------------------------- -// Listener (real HTTPS server at /tunnel) +// Configs // --------------------------------------------------------------------------- /// Configuration for the tunnel listener. @@ -131,11 +134,128 @@ pub struct ListenerConfig { pub auth_token: Arc, } -/// Start the HTTPS tunnel listener with a real `/tunnel` endpoint. -/// -/// Binds a TLS-enabled HTTPS server on the configured address. -/// Only the `/tunnel` POST path is accepted and authenticated. -/// Returns when the listener is shut down or an error occurs. +/// Configuration for the tunnel connector. +#[derive(Debug, Clone)] +pub struct ConnectorConfig { + pub target_host: String, + pub target_port: u16, + pub client_cert_path: Arc, + pub client_key_path: Arc, + pub ca_cert_path: Arc, + pub auth_token: Arc, +} + +// --------------------------------------------------------------------------- +// HTTP helpers — manual parsing for the auth handshake +// --------------------------------------------------------------------------- + +/// Parse an HTTP/1.1 request from the stream. +/// Returns (method, path, headers_map). +async fn parse_http_request( + reader: &mut R, +) -> Result<(String, String, HashMap), TunnelError> +where + R: AsyncReadExt + Unpin, +{ + let mut buf = Vec::with_capacity(1024); + let mut crlf_count = 0; + + loop { + let mut byte = [0u8; 1]; + reader.read_exact(&mut byte).await?; + buf.push(byte[0]); + if byte[0] == b'\r' || byte[0] == b'\n' { + crlf_count += 1; + if crlf_count >= 4 && buf.ends_with(b"\r\n\r\n") { + break; + } + } else { + crlf_count = 0; + } + // Safety valve: prevent reading forever for malformed requests + if buf.len() > 8192 { + return Err(TunnelError::Protocol("request too large".into())); + } + } + + let text = String::from_utf8_lossy(&buf); + let lines: Vec<&str> = text.lines().collect(); + if lines.is_empty() { + return Err(TunnelError::Protocol("empty request".into())); + } + + let parts: Vec<&str> = lines[0].split_whitespace().collect(); + if parts.len() < 2 { + return Err(TunnelError::Protocol("invalid request line".into())); + } + let method = parts[0].to_string(); + let path = parts[1].to_string(); + + let mut headers = HashMap::new(); + for line in lines.iter().skip(1) { + if let Some((key, value)) = line.split_once(':') { + headers.insert(key.trim().to_lowercase(), value.trim().to_string()); + } + } + + Ok((method, path, headers)) +} + +/// Send a minimal HTTP/1.1 response (no body, Content-Length: 0). +async fn send_http_response( + writer: &mut W, + status: u16, + status_text: &str, +) -> Result<(), TunnelError> +where + W: AsyncWriteExt + Unpin, +{ + let response = format!( + "HTTP/1.1 {} {}\r\nContent-Length: 0\r\n\r\n", + status, status_text + ); + writer.write_all(response.as_bytes()).await?; + writer.flush().await?; + Ok(()) +} + +/// Read an HTTP/1.1 response status line. Returns the status code. +async fn read_http_response_status(reader: &mut R) -> Result +where + R: AsyncReadExt + Unpin, +{ + let mut buf = Vec::with_capacity(256); + loop { + let mut byte = [0u8; 1]; + reader.read_exact(&mut byte).await?; + buf.push(byte[0]); + if buf.ends_with(b"\r\n\r\n") { + break; + } + if buf.len() > 4096 { + return Err(TunnelError::Protocol("response headers too large".into())); + } + } + + let text = String::from_utf8_lossy(&buf); + let first_line = text + .lines() + .next() + .ok_or_else(|| TunnelError::Protocol("empty response".into()))?; + let parts: Vec<&str> = first_line.split_whitespace().collect(); + if parts.len() < 2 { + return Err(TunnelError::Protocol("invalid response line".into())); + } + parts[1] + .parse::() + .map_err(|_| TunnelError::Protocol("invalid status code".into())) +} + +// --------------------------------------------------------------------------- +// Listener +// --------------------------------------------------------------------------- + +/// Start the HTTPS tunnel listener. pub async fn run_listener(config: ListenerConfig) -> Result<(), TunnelError> { let addr = config.bind_addr; let server_cert_path = config.server_cert_path.clone(); @@ -143,7 +263,6 @@ pub async fn run_listener(config: ListenerConfig) -> Result<(), TunnelError> { let ca_cert_path = config.ca_cert_path.clone(); let auth_token = config.auth_token.clone(); - // Build TLS config — validates all cert/key files before opening sockets let tls_config = tls::build_server_config( server_cert_path.as_ref(), server_key_path.as_ref(), @@ -158,7 +277,6 @@ pub async fn run_listener(config: ListenerConfig) -> Result<(), TunnelError> { let server_certs = tls::load_certs(server_cert_path.as_ref()).map_err(TunnelError::Tls)?; let bind_host = addr.ip().to_string(); tls::validate_cert_identity(&server_certs[0], &bind_host).map_err(|e| { - tracing::error!("Server cert identity mismatch: {}", e); TunnelError::Tls(crate::errors::TlsError::IdentityMismatch(format!( "server cert does not match bind address {}: {}", bind_host, e @@ -166,14 +284,11 @@ pub async fn run_listener(config: ListenerConfig) -> Result<(), TunnelError> { })?; let tls_acceptor = tokio_rustls::TlsAcceptor::from(tls_config); - - // Bind TCP listener — fail closed on port conflict let listener = TcpListener::bind(addr).await.map_err(|e| { - let addr_str = addr.to_string(); if e.kind() == std::io::ErrorKind::AddrInUse { TunnelError::PortInUse(addr.port()) } else { - TunnelError::BindError(addr_str, e.to_string()) + TunnelError::BindError(addr.to_string(), e.to_string()) } })?; @@ -198,340 +313,463 @@ pub async fn run_listener(config: ListenerConfig) -> Result<(), TunnelError> { let token = auth_token.clone(); tokio::spawn(async move { - if let Err(e) = handle_listener_https(stream, peer_addr, acceptor, token).await { + if let Err(e) = handle_listener_connection(stream, peer_addr, acceptor, token).await { tracing::warn!("Connection from {} failed: {}", peer_addr, e); } }); } } -/// Handle a single incoming HTTPS connection for the listener. -/// Accepts TLS, then serves hyper HTTP requests — `/tunnel` POST for auth, `/forward` POST for data. -async fn handle_listener_https( - stream: tokio::net::TcpStream, +/// Handle a single incoming connection: TLS → mTLS → Auth → Framing. +async fn handle_listener_connection( + stream: TcpStream, peer_addr: SocketAddr, acceptor: tokio_rustls::TlsAcceptor, auth_token: Arc, ) -> Result<(), TunnelError> { - // Perform TLS accept + // TLS accept let tls_stream = acceptor.accept(stream).await.map_err(|e| { tracing::warn!("TLS accept failed for {}: {}", peer_addr, e); TunnelError::Tls(crate::errors::TlsError::VerificationFailed(e.to_string())) })?; - // Check peer certificates exist (mTLS) - let (_tcp_stream, server_conn) = tls_stream.get_ref(); + // Verify mTLS — client must present a valid certificate + let (_, server_conn) = tls_stream.get_ref(); let peer_certs = server_conn.peer_certificates().ok_or_else(|| { tracing::warn!("No client certificate from {}", peer_addr); TunnelError::Tls(crate::errors::TlsError::Missing("client certificate")) })?; - if peer_certs.is_empty() { - tracing::warn!("Empty client certificate chain from {}", peer_addr); return Err(TunnelError::Tls(crate::errors::TlsError::Missing( "client certificate", ))); } - tracing::info!("mTLS established with {}", peer_addr); - // Wrap TLS stream in hyper IO and serve HTTP - let io = hyper_util::rt::TokioIo::new(tls_stream); + // Pin the stream so we can use it with io::split and generic async functions + let mut tls_stream: std::pin::Pin>> = + Box::pin(tls_stream); - // Build the auth-aware service - let auth_state = Arc::new(auth_token); - let service = service_fn(move |req: Request| { - let auth = auth_state.clone(); - let pa = peer_addr; - async move { - 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") => { - let received = req - .headers() - .get("X-Rustunnel-Token") - .and_then(|v| v.to_str().ok()) - .unwrap_or(""); - if constant_time_compare(received, &auth) { - tracing::info!("Tunnel established with {} via HTTPS /tunnel", pa); - Ok::, hyper::http::Error>( - Response::builder() - .status(StatusCode::OK) - .header("X-Rustunnel-Status", "authenticated") - .body(body("OK")) - .unwrap(), - ) - } else { - tracing::warn!("Auth rejected for {} — invalid token on /tunnel", pa); - Ok::, hyper::http::Error>( - Response::builder() - .status(StatusCode::UNAUTHORIZED) - .body(body("FAIL: invalid token")) - .unwrap(), - ) - } + // Phase 1: HTTP auth handshake + let (method, path, headers) = parse_http_request(&mut *tls_stream).await?; + + if method == "POST" && path == "/tunnel" { + let received_token = headers + .get("x-rustunnel-token") + .map(|s| s.as_str()) + .unwrap_or(""); + + if constant_time_compare(received_token, &auth_token) { + tracing::info!("Tunnel established with {} via HTTPS /tunnel", peer_addr); + send_http_response(&mut *tls_stream, 200, "OK").await?; + } else { + tracing::warn!("Auth rejected for {} — invalid token on /tunnel", peer_addr); + send_http_response(&mut *tls_stream, 401, "Unauthorized").await?; + return Ok(()); // close connection on auth failure + } + } else if method == "GET" && path == "/tunnel" { + // Health check — respond and close + send_http_response(&mut *tls_stream, 200, "OK").await?; + return Ok(()); + } else { + tracing::debug!( + "Rejected {} {} from {} — only POST /tunnel accepted", + method, + path, + peer_addr + ); + send_http_response(&mut *tls_stream, 404, "Not Found").await?; + return Ok(()); + } + + // Phase 2: Switch to binary framing on the same connection + tracing::info!( + "Upgraded connection from {} to binary framing protocol", + peer_addr + ); + handle_framing_server(tls_stream, peer_addr).await +} + +// --------------------------------------------------------------------------- +// Framing server (listener side) +// --------------------------------------------------------------------------- + +/// Entry for an active stream on the listener side. +struct ServerStreamEntry { + /// Write half of the target TCP connection + target_write: tokio::net::tcp::OwnedWriteHalf, +} + +/// Handle the binary framing protocol on the server side. +async fn handle_framing_server(stream: S, peer_addr: SocketAddr) -> Result<(), TunnelError> +where + S: AsyncReadExt + AsyncWriteExt + Unpin + Send + 'static, +{ + let mut stream = stream; + let mut reader = FrameReader::new(); + let (outbound_tx, mut outbound_rx) = mpsc::channel::(256); + + // Active streams: stream_id → target write half + let mut streams: HashMap = HashMap::new(); + + tracing::info!("Framing session started with {}", peer_addr); + + loop { + tokio::select! { + // Write outbound frames (CONNECT_REPLY, DATA from target→tunnel, CLOSE, ERROR) + Some(frame) = outbound_rx.recv() => { + if let Err(e) = FrameWriter::write_frame(&mut stream, &frame).await { + tracing::debug!("Write error for stream {}: {}", frame.stream_id, e); + break; } - (&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(), - ), - _ => { - tracing::debug!( - "Rejected {} {} from {} — only POST /tunnel and POST /forward accepted", - req.method(), - req.uri().path(), - pa - ); - Ok::, hyper::http::Error>( - Response::builder() - .status(StatusCode::NOT_FOUND) - .body(body("not found")) - .unwrap(), - ) + } + + // Read inbound frames from the tunnel (CONNECT, DATA from client, CLOSE) + result = reader.read_frame(&mut stream) => { + match result { + Ok(Some(frame)) => { + if let Err(e) = handle_inbound_frame( + &frame, + &mut stream, + &outbound_tx, + &mut streams, + peer_addr, + ).await { + tracing::warn!( + "Stream {} inbound error: {}", + frame.stream_id, + e + ); + // Send error frame and continue + let _ = outbound_tx + .send(Frame::error(frame.stream_id, e.to_string())) + .await; + } + } + Ok(None) => { + tracing::info!("Framing session with {} closed (EOF)", peer_addr); + break; + } + Err(e) => { + tracing::warn!("Framing read error from {}: {}", peer_addr, e); + break; + } } } } - }); + } - let conn_result = hyper::server::conn::http1::Builder::new() - .serve_connection(io, service) - .await; + // Clean up all streams + for (sid, mut entry) in streams.drain() { + let _ = entry.target_write.shutdown().await; + tracing::debug!("Closed stream {} on shutdown", sid); + } - if let Err(e) = conn_result { - tracing::debug!("HTTP connection error for {}: {}", peer_addr, e); + tracing::info!("Framing session with {} ended", peer_addr); + Ok(()) +} + +/// Handle an inbound frame on the server side. +async fn handle_inbound_frame( + frame: &Frame, + _stream: &mut S, + outbound_tx: &mpsc::Sender, + streams: &mut HashMap, + peer_addr: SocketAddr, +) -> Result<(), TunnelError> +where + S: AsyncWriteExt + Unpin, +{ + match frame.frame_type { + FRAME_CONNECT => { + if streams.contains_key(&frame.stream_id) { + let _ = outbound_tx + .send(Frame::error(frame.stream_id, "duplicate stream ID")) + .await; + return Ok(()); + } + + let target = decode_target(&frame.payload) + .map_err(|e| TunnelError::Protocol(format!("decode target: {}", e)))?; + + let target_str = socks5::target_to_string(&target); + tracing::info!( + "Stream {} CONNECT from {} to {} (listener-side target)", + frame.stream_id, + peer_addr, + target_str + ); + + // Resolve and connect to target from the listener side + let target_addr = match socks5::resolve_target(&target).await { + Ok(addr) => addr, + Err(e) => { + tracing::warn!( + "Stream {} resolve failed for {}: {}", + frame.stream_id, + target_str, + e + ); + let _ = outbound_tx + .send(Frame::connect_reply(frame.stream_id, CONNECT_FAILED)) + .await; + let _ = outbound_tx + .send(Frame::error(frame.stream_id, format!("resolve: {}", e))) + .await; + return Ok(()); + } + }; + + match TcpStream::connect(target_addr).await { + Ok(target_stream) => { + tracing::info!( + "Stream {} connected to target {} (listener side)", + frame.stream_id, + target_addr + ); + + let (mut target_read, target_write) = target_stream.into_split(); + let sid = frame.stream_id; + let outbound = outbound_tx.clone(); + + // Spawn target → tunnel direction + tokio::spawn(async move { + let mut buf = [0u8; 32 * 1024]; + loop { + match target_read.read(&mut buf).await { + Ok(0) => break, // Target closed (EOF) + Ok(n) => { + if outbound.send(Frame::data(sid, &buf[..n])).await.is_err() { + break; + } + } + Err(e) => { + tracing::debug!("Stream {} target read error: {}", sid, e); + let _ = outbound.send(Frame::error(sid, e.to_string())).await; + break; + } + } + } + tracing::info!("Stream {} target → tunnel EOF", sid); + let _ = outbound.send(Frame::close(sid)).await; + }); + + streams.insert(frame.stream_id, ServerStreamEntry { target_write }); + let _ = outbound_tx + .send(Frame::connect_reply(frame.stream_id, CONNECT_OK)) + .await; + } + Err(e) => { + let err_msg = if e.kind() == std::io::ErrorKind::ConnectionRefused { + format!("connection refused to {}", target_addr) + } else { + format!("connect to {}: {}", target_addr, e) + }; + tracing::warn!( + "Stream {} connect to {} failed: {}", + frame.stream_id, + target_addr, + err_msg + ); + let _ = outbound_tx + .send(Frame::connect_reply(frame.stream_id, CONNECT_FAILED)) + .await; + let _ = outbound_tx + .send(Frame::error(frame.stream_id, err_msg)) + .await; + } + } + } + + FRAME_DATA => { + // Forward data to the target + if let Some(entry) = streams.get_mut(&frame.stream_id) + && let Err(e) = entry.target_write.write_all(&frame.payload).await + { + tracing::debug!("Stream {} target write error: {}", frame.stream_id, e); + streams.remove(&frame.stream_id); + } + } + + FRAME_CLOSE => { + if let Some(mut entry) = streams.remove(&frame.stream_id) { + let _ = entry.target_write.shutdown().await; + tracing::info!( + "Stream {} closed by client (sender-side EOF)", + frame.stream_id + ); + } + } + + FRAME_ERROR => { + let msg = String::from_utf8_lossy(&frame.payload); + tracing::warn!("Stream {} error from client: {}", frame.stream_id, msg); + if let Some(mut entry) = streams.remove(&frame.stream_id) { + let _ = entry.target_write.shutdown().await; + } + } + + _ => { + tracing::debug!( + "Unknown frame type {} on stream {} from {}", + frame.frame_type, + frame.stream_id, + peer_addr + ); + } } 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)) +// --------------------------------------------------------------------------- +// Connector: StreamMux for multiplexed streams +// --------------------------------------------------------------------------- + +/// Entry for an active stream on the connector side. +struct ConnectorStreamEntry { + /// Channel for tunnel → SOCKS5 data (frame → SOCKS5 client) + data_tx: mpsc::Sender, + /// One-shot channel for CONNECT_REPLY (Option so we can take it for send) + reply_tx: Option>, } -/// 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 raw bytes. This preserves exact byte semantics for SOCKS5 TCP proxy. -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 — returns raw bytes (no HTTP parse/unparse) - let result = forward_raw_to_target(&fwd_host, fwd_port, &stream_id, body_bytes).await; - - match result { - Ok(resp_body) => Ok(Response::builder() - .status(StatusCode::OK) - .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()) - } - } +/// Stream multiplexer for the connector side. +/// Manages multiple concurrent SOCKS5 streams over one authenticated tunnel. +pub struct StreamMux { + /// Channel for all outbound frames to the tunnel + outbound_tx: mpsc::Sender, + /// Active streams registry + streams: Arc>>, + /// Monotonically increasing stream ID counter (odd IDs for client-initiated) + next_id: Arc, + /// Auth state + auth_complete: Arc, } -/// Forward request data to a target TCP server and return the raw response bytes. -async fn forward_raw_to_target( - host: &str, - port: u16, - stream_id: &str, - request_data: Bytes, -) -> Result { - // 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) +impl StreamMux { + /// Create a new StreamMux. + pub fn new(outbound_tx: mpsc::Sender) -> Self { + Self { + outbound_tx, + streams: Arc::new(tokio::sync::Mutex::new(HashMap::new())), + next_id: Arc::new(AtomicU64::new(1)), + auth_complete: Arc::new(AtomicBool::new(true)), } - })?; + } - tracing::info!( - "Forward stream {}: connected to target {}", - stream_id, - target_addr - ); + /// Open a new stream to the target. + /// Returns (stream_id, data_rx) where data_rx receives tunnel → SOCKS5 data. + pub async fn open_stream( + &self, + target: &socks5::Socks5Target, + ) -> Result<(u64, mpsc::Receiver), TunnelError> { + if !self.auth_complete.load(Ordering::SeqCst) { + return Err(TunnelError::Auth(AuthError::NotAuthenticated)); + } - // Write request data to target - if !request_data.is_empty() { - target_stream - .write_all(&request_data) + let stream_id = self.next_id.fetch_add(2, Ordering::Relaxed); // odd IDs + let (reply_tx, reply_rx) = oneshot::channel(); + let (data_tx, data_rx) = mpsc::channel::(64); + + self.streams.lock().await.insert( + stream_id, + ConnectorStreamEntry { + data_tx, + reply_tx: Some(reply_tx), + }, + ); + + let target_bytes = encode_target(target); + self.outbound_tx + .send(Frame::connect(stream_id, target_bytes)) .await - .map_err(|e| format!("write to target: {}", e))?; + .map_err(|_| TunnelError::Protocol("outbound channel closed".into()))?; + + let status = reply_rx + .await + .map_err(|_| TunnelError::Protocol("connect reply channel closed".into()))?; + + if status != CONNECT_OK { + self.streams.lock().await.remove(&stream_id); + return Err(TunnelError::Protocol(format!( + "CONNECT failed with status {}", + status + ))); + } + + Ok((stream_id, data_rx)) } - // Read target response with timeout to avoid blocking on keep-alive connections - let mut response_data = Vec::new(); - let mut buf = vec![0u8; 32 * 1024]; - loop { - // Use a 5-second timeout to avoid blocking on keep-alive connections - match tokio::time::timeout( - std::time::Duration::from_secs(5), - target_stream.read(&mut buf), - ) - .await - { - Err(_) => { - // Timeout — connection may be keep-alive; return what we have - if !response_data.is_empty() { - break; + /// Dispatch an inbound frame to the correct stream. + pub async fn dispatch_frame(&self, frame: &Frame) { + match frame.frame_type { + FRAME_CONNECT_REPLY => { + let mut streams = self.streams.lock().await; + if let Some(entry) = streams.get_mut(&frame.stream_id) { + let status = frame.payload.first().copied().unwrap_or(CONNECT_FAILED); + if let Some(tx) = entry.reply_tx.take() { + let _ = tx.send(status); + } } - return Err("read timeout".to_string()); } - Ok(Err(e)) => { - if !response_data.is_empty() { - tracing::info!( - "Forward stream {}: target connection closed with data read", - stream_id - ); - break; + FRAME_DATA => { + let mut streams = self.streams.lock().await; + if let Some(entry) = streams.get(&frame.stream_id) + && entry.data_tx.send(frame.payload.clone()).await.is_err() + { + // SOCKS5 handler dropped — clean up + streams.remove(&frame.stream_id); + let _ = self.outbound_tx.send(Frame::close(frame.stream_id)).await; } - return Err(format!("read from target: {}", e)); } - Ok(Ok(0)) => { - break; + FRAME_CLOSE => { + let mut streams = self.streams.lock().await; + if streams.remove(&frame.stream_id).is_some() { + tracing::info!("Stream {} closed by server", frame.stream_id); + } } - Ok(Ok(n)) => { - response_data.extend_from_slice(&buf[..n]); + FRAME_ERROR => { + let msg = String::from_utf8_lossy(&frame.payload); + tracing::warn!("Stream {} error from server: {}", frame.stream_id, msg); + let mut streams = self.streams.lock().await; + streams.remove(&frame.stream_id); + } + _ => { + tracing::debug!( + "Unknown frame type {} on stream {}", + frame.frame_type, + frame.stream_id + ); } } } - tracing::info!( - "Forward stream {}: completed, {} bytes response", - stream_id, - response_data.len() - ); + /// Close a stream from the connector side. + pub async fn close_stream(&self, stream_id: u64) { + self.streams.lock().await.remove(&stream_id); + let _ = self.outbound_tx.send(Frame::close(stream_id)).await; + } - Ok(Bytes::from(response_data)) + /// Check if auth is complete. + pub fn is_authenticated(&self) -> bool { + self.auth_complete.load(Ordering::SeqCst) + } } // --------------------------------------------------------------------------- -// Connector (stays alive after auth) +// Connector with SOCKS5 // --------------------------------------------------------------------------- -/// Configuration for the tunnel connector. -#[derive(Debug, Clone)] -pub struct ConnectorConfig { - /// Target host (may be hostname or IP) - pub target_host: String, - /// Target port - pub target_port: u16, - pub client_cert_path: Arc, - pub client_key_path: Arc, - pub ca_cert_path: Arc, - pub auth_token: Arc, -} - -/// Connect to the HTTPS tunnel and establish a session, then stay alive until shutdown. -/// -/// Performs the full security gate sequence: -/// 1. TCP connection -/// 2. TLS handshake with server validation -/// 3. mTLS client certificate presentation -/// 4. Application auth token validation via POST /tunnel -/// Then maintains the session alive until shutdown signal is received. -pub async fn connect_tunnel(config: ConnectorConfig) -> Result<(), TunnelError> { - let target_host = config.target_host.clone(); - let target_port = config.target_port; - let client_cert_path = config.client_cert_path.clone(); - let client_key_path = config.client_key_path.clone(); - let ca_cert_path = config.ca_cert_path.clone(); - let auth_token = config.auth_token.clone(); - - // Build client TLS config +/// Run the connector with SOCKS5 proxy: connect, auth, framing, then SOCKS5. +pub async fn run_connector_with_socks( + config: ConnectorConfig, + socks_addr: SocketAddr, +) -> Result<(), TunnelError> { let client_config = tls::build_client_config( - client_cert_path.as_ref(), - client_key_path.as_ref(), - ca_cert_path.as_ref(), + 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); @@ -539,196 +777,437 @@ pub async fn connect_tunnel(config: ConnectorConfig) -> Result<(), TunnelError> })?; let tls_connector = tokio_rustls::TlsConnector::from(client_config); - let shutdown = Arc::new(Notify::new()); - let auth_token_clone = auth_token.clone(); - - // Register Ctrl-C handler for graceful shutdown - let shutdown_clone = shutdown.clone(); - tokio::spawn(async move { - tokio::signal::ctrl_c().await.ok(); - tracing::info!("Shutdown signal received — stopping connector"); - shutdown_clone.notify_one(); - }); + 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)?; loop { - // Resolve hostname — no panic on failure - let target_addr = match resolve_target(&target_host, target_port).await { - Ok(addr) => addr, - Err(e) => { - tracing::error!("Hostname resolution failed: {}", e); - return Err(TunnelError::Io(std::io::Error::new( - std::io::ErrorKind::InvalidInput, - e.to_string(), - ))); + let stream = TcpStream::connect(target_addr).await.map_err(|e| { + if e.kind() == std::io::ErrorKind::ConnectionRefused { + return TunnelError::ConnectionRefused(target_addr.to_string()); } - }; - - // Create server name for SNI - let server_name = tls::server_name_from_host(&target_host).map_err(|e| { - tracing::error!("Invalid server name: {}", e); - TunnelError::Tls(e) + TunnelError::Io(e) })?; - // TCP connect - let stream = match TcpStream::connect(target_addr).await { - Ok(s) => s, - Err(e) => { - if e.kind() == std::io::ErrorKind::ConnectionRefused { - return Err(TunnelError::ConnectionRefused(target_addr.to_string())); - } - tracing::error!("Connection failed: {}", e); - return Err(TunnelError::Io(e)); - } - }; - tracing::info!( "Connected to HTTPS tunnel at {}:{} (HTTPS default transport)", - target_host, - target_port + config.target_host, + config.target_port ); - // TLS handshake with server validation - let tls_stream = match tls_connector.connect(server_name.clone(), stream).await { - Ok(s) => s, - Err(e) => { - tracing::error!("TLS handshake failed: {}", e); - return Err(TunnelError::Tls( - crate::errors::TlsError::VerificationFailed(format!( - "TLS handshake to {}:{} failed: {}", - target_host, target_port, e - )), - )); - } - }; + let tls_stream = tls_connector + .connect(server_name.clone(), 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", - target_host, - target_port + config.target_host, + config.target_port ); - // Send auth via POST /tunnel over HTTPS — takes ownership of TLS stream - // authenticate_https now returns (sender, JoinHandle) — the connection is already spawned - let (_sender, conn_handle) = - match authenticate_https(tls_stream, &auth_token_clone, &target_host).await { - Ok(v) => v, - Err(TunnelError::Auth(e)) => { - tracing::error!("Auth failed: {}", e); - return Err(TunnelError::Auth(e)); + // Authenticate over HTTP, then switch to framing + match authenticate_tunnel(&config.auth_token, &config.target_host, tls_stream).await { + Ok(auth_result) => { + let mux = auth_result.mux; + + tracing::info!( + "Tunnel session established with {}:{} via HTTPS /tunnel", + config.target_host, + config.target_port + ); + + // Start SOCKS5 proxy and framing loop + match run_socks5_with_mux(socks_addr, mux).await { + Ok(()) => return Ok(()), + Err(e) => { + tracing::error!("SOCKS5 proxy error: {}", e); + return Err(e); + } } - Err(e) => { - tracing::error!("Auth error: {}", e); - return Err(e); - } - }; - - tracing::info!( - "Tunnel session established with {}:{} via HTTPS /tunnel", - target_host, - target_port - ); - - // Stay alive: keep session until shutdown or connection drop - let shutdown_wait = shutdown.notified(); - tokio::pin!(shutdown_wait); - - // conn_handle is the already-spawned connection driver - let mut conn_driver = Some(conn_handle); - - tokio::select! { - _ = &mut shutdown_wait => { - if let Some(handle) = conn_driver.take() { - handle.abort(); - } - tracing::info!("Connector shutting down gracefully"); - return Ok(()); } - result = conn_driver.take().unwrap() => { - // Connection dropped — reconnect - if let Err(e) = result { - tracing::warn!("Connection driver error: {}", e); - } - tracing::info!("Connection dropped — reconnecting to maintain tunnel"); + Err(TunnelError::Auth(e)) => { + tracing::error!("Auth failed: {}", e); + return Err(TunnelError::Auth(e)); } - _ = tokio::time::sleep(std::time::Duration::from_secs(300)) => { - if let Some(handle) = conn_driver.take() { - handle.abort(); - } - tracing::info!("Session idle timeout — reconnecting to maintain tunnel"); + Err(e) => { + tracing::error!("Connection error, reconnecting: {}", e); + tokio::time::sleep(std::time::Duration::from_secs(2)).await; } } } } -/// Authenticate via HTTPS POST /tunnel with the auth token in header. -/// Takes ownership of the TLS stream, spawns the connection driver, sends -/// the auth request, and returns the authenticated sender plus a handle -/// to abort the connection on shutdown. -async fn authenticate_https( - tls_stream: tokio_rustls::client::TlsStream, +/// Result of the authenticate_tunnel function. +#[allow(dead_code)] +struct AuthResult { + mux: Arc, + framing_task: tokio::task::JoinHandle<()>, +} + +/// Authenticate over HTTP then switch to binary framing. +async fn authenticate_tunnel( auth_token: &str, target_host: &str, -) -> Result< - ( - hyper::client::conn::http1::SendRequest>, - tokio::task::JoinHandle<()>, - ), - TunnelError, -> { - let io = hyper_util::rt::TokioIo::new(tls_stream); + tls_stream: tokio_rustls::client::TlsStream, +) -> Result { + let mut stream: std::pin::Pin>> = + Box::pin(tls_stream); - let (mut sender, connection) = hyper::client::conn::http1::handshake(io) - .await - .map_err(|e| TunnelError::Protocol(format!("HTTP handshake failed: {}", e)))?; + // Send HTTP auth request + let request = format!( + "POST /tunnel HTTP/1.1\r\nHost: {}\r\nX-Rustunnel-Token: {}\r\nContent-Length: 0\r\n\r\n", + target_host, auth_token + ); + stream.write_all(request.as_bytes()).await?; - // Spawn the connection driver so that send_request can make progress - let conn_handle = tokio::spawn(async move { - let _ = connection.await; - }); + // Read HTTP response + let status = read_http_response_status(&mut *stream).await?; - // Build and send POST /tunnel request - let request = Request::builder() - .method("POST") - .uri("/tunnel") - .header("X-Rustunnel-Token", auth_token) - .header("Host", target_host) - .body(Full::new(Bytes::new())) - .map_err(|e| TunnelError::Protocol(format!("failed to build auth request: {}", e)))?; - - let response = sender - .send_request(request) - .await - .map_err(|e| TunnelError::Protocol(format!("Auth request failed: {}", e)))?; - - let status = response.status(); - let body = response - .into_body() - .collect() - .await - .map_err(|e| { - TunnelError::Io(std::io::Error::other(format!( - "failed to read response body: {}", - e - ))) - })? - .to_bytes(); - - let body_str = String::from_utf8_lossy(&body); - - if status == StatusCode::OK && body_str.contains("OK") { - return Ok((sender, conn_handle)); - } - - if body_str.contains("FAIL") || status.is_client_error() { + if status != 200 { + tracing::warn!( + "Auth rejected — HTTP status {} from {}", + status, + target_host + ); return Err(TunnelError::Auth(AuthError::Invalid)); } - Err(TunnelError::Protocol(format!( - "unexpected auth response: {} {}", - status, body_str - ))) + tracing::info!("Authenticated with {} via HTTPS /tunnel", target_host); + + // Switch to binary framing + tracing::info!( + "Upgraded connection to {} to binary framing protocol", + target_host + ); + + // Set up the StreamMux and framing loop + let (outbound_tx, mut outbound_rx) = mpsc::channel::(256); + let mux = Arc::new(StreamMux::new(outbound_tx)); + let mux_dispatch = mux.clone(); + + // Spawn the framing loop + let framing_task = tokio::spawn(async move { + run_framing_client(stream, mux_dispatch, &mut outbound_rx).await; + }); + + Ok(AuthResult { mux, framing_task }) } +/// Run the framing loop on the client side. +async fn run_framing_client( + mut stream: S, + mux: Arc, + outbound_rx: &mut mpsc::Receiver, +) where + S: AsyncReadExt + AsyncWriteExt + Unpin, +{ + let mut reader = FrameReader::new(); + + loop { + tokio::select! { + // Write outbound frames + Some(frame) = outbound_rx.recv() => { + if let Err(e) = FrameWriter::write_frame(&mut stream, &frame).await { + tracing::debug!("Client write error: {}", e); + break; + } + } + + // Read inbound frames and dispatch + result = reader.read_frame(&mut stream) => { + match result { + Ok(Some(frame)) => { + mux.dispatch_frame(&frame).await; + } + Ok(None) => { + tracing::info!("Tunnel connection closed (EOF)"); + break; + } + Err(e) => { + tracing::warn!("Client framing read error: {}", e); + break; + } + } + } + } + } + + tracing::info!("Client framing loop ended"); +} + +// --------------------------------------------------------------------------- +// SOCKS5 handler with StreamMux +// --------------------------------------------------------------------------- + +/// Run the SOCKS5 proxy with a shared StreamMux. +async fn run_socks5_with_mux( + socks_addr: SocketAddr, + mux: Arc, +) -> Result<(), TunnelError> { + let socks_listener = TcpListener::bind(socks_addr).await.map_err(|e| { + if e.kind() == std::io::ErrorKind::AddrInUse { + TunnelError::PortInUse(socks_addr.port()) + } else { + TunnelError::BindError(socks_addr.to_string(), e.to_string()) + } + })?; + + tracing::info!( + "SOCKS5 proxy listening on {} (persistent tunnel, binary framing)", + socks_addr + ); + + // Listen for Ctrl-C + let shutdown = tokio::signal::ctrl_c(); + tokio::pin!(shutdown); + + loop { + tokio::select! { + result = socks_listener.accept() => { + match result { + Ok((stream, peer_addr)) => { + let m = mux.clone(); + tokio::spawn(async move { + if let Err(e) = handle_socks5_connection(stream, peer_addr, m).await { + tracing::warn!("SOCKS5 connection from {} failed: {}", peer_addr, e); + } + }); + } + Err(e) => { + tracing::error!("SOCKS5 accept error: {}", e); + } + } + } + _ = &mut shutdown => { + tracing::info!("Shutdown signal received — stopping SOCKS5 proxy"); + break; + } + } + } + + tracing::info!("SOCKS5 proxy shutting down"); + Ok(()) +} + +/// Handle a single SOCKS5 connection using the shared StreamMux. +async fn handle_socks5_connection( + stream: TcpStream, + peer_addr: SocketAddr, + mux: Arc, +) -> Result<(), socks5::Socks5Error> { + tracing::info!("SOCKS5 connection from {}", peer_addr); + + if !mux.is_authenticated() { + tracing::warn!( + "SOCKS5 request from {} rejected: tunnel not authenticated yet", + peer_addr + ); + return Err(socks5::Socks5Error::NotAuthenticated); + } + + // --- SOCKS5 Handshake --- + let (mut socks_read, mut socks_write) = stream.into_split(); + let target = socks5_handshake(&mut socks_read, &mut socks_write).await?; + + let target_str = socks5::target_to_string(&target); + tracing::info!( + "SOCKS5 CONNECT from {} to {} (via multiplexed tunnel)", + peer_addr, + target_str + ); + + // --- Open stream via tunnel --- + let (stream_id, mut data_rx) = match mux.open_stream(&target).await { + Ok(r) => r, + Err(e) => { + tracing::warn!("Stream open failed for {}: {}", target_str, e); + socks_write + .write_all(&socks5::build_reply_failure(socks5::REPL_GENERAL_FAILURE)) + .await?; + return Err(socks5::Socks5Error::Protocol(format!( + "tunnel stream open failed: {}", + e + ))); + } + }; + + tracing::info!( + "SOCKS5 stream {}: target {} connected via tunnel", + stream_id, + target_str + ); + + // Send SOCKS5 success reply + socks_write + .write_all(&socks5::build_reply_success()) + .await?; + + // --- Bidirectional copy: SOCKS5 ↔ Tunnel (raw bytes, no buffering) --- + + // SOCKS5 → Tunnel: read from SOCKS5 client, send DATA frames + let outbound = mux.outbound_tx.clone(); + let sid = stream_id; + let socks_to_tunnel = tokio::spawn(async move { + let mut buf = [0u8; 32 * 1024]; + loop { + match socks_read.read(&mut buf).await { + Ok(0) => break, // SOCKS5 client closed + Ok(n) => { + if outbound.send(Frame::data(sid, &buf[..n])).await.is_err() { + break; + } + } + Err(e) => { + tracing::debug!("Stream {} SOCKS5 read error: {}", sid, e); + let _ = outbound.send(Frame::error(sid, e.to_string())).await; + break; + } + } + } + tracing::info!("Stream {} SOCKS5 → tunnel EOF", sid); + let _ = outbound.send(Frame::close(sid)).await; + }); + + // Tunnel → SOCKS5: read DATA frames, write to SOCKS5 client + while let Some(data) = data_rx.recv().await { + if let Err(e) = socks_write.write_all(&data).await { + tracing::debug!("Stream {} SOCKS5 write error: {}", stream_id, e); + break; + } + } + + // Clean up + socks_to_tunnel.abort(); + mux.close_stream(stream_id).await; + + tracing::info!("SOCKS5 stream {} from {} completed", stream_id, peer_addr); + + Ok(()) +} + +/// Parse the SOCKS5 handshake (greeting + auth + CONNECT request). +/// Returns the target address. +async fn socks5_handshake( + reader: &mut tokio::net::tcp::OwnedReadHalf, + writer: &mut tokio::net::tcp::OwnedWriteHalf, +) -> Result { + // --- Greeting --- + let greeting = read_exact_2(reader).await?; + if greeting[0] != socks5::VERSION { + writer + .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(reader, n_methods).await?; + + let chosen = if methods.contains(&socks5::AUTH_USERNAME_PASSWORD) { + socks5::AUTH_USERNAME_PASSWORD + } else if methods.contains(&socks5::AUTH_NONE) { + socks5::AUTH_NONE + } else { + writer + .write_all(&[socks5::VERSION, socks5::AUTH_NO_ACCEPTABLE]) + .await?; + return Err(socks5::Socks5Error::UnsupportedAuthMethod( + methods.first().copied().unwrap_or(0xFF), + )); + }; + + writer.write_all(&[socks5::VERSION, chosen]).await?; + + // --- Auth (if username/password) --- + if chosen == socks5::AUTH_USERNAME_PASSWORD { + let auth_buf = read_exact_2(reader).await?; + if auth_buf[0] != 0x01 { + writer.write_all(&socks5::build_auth_failure()).await?; + return Err(socks5::Socks5Error::AuthFailed); + } + let ulen = auth_buf[1] as usize; + if ulen > 0 { + let _username = read_exact(reader, ulen).await?; + } + let plen_buf = read_exact(reader, 1).await?; + let plen = plen_buf[0] as usize; + if plen > 0 { + let _password = read_exact(reader, plen).await?; + } + writer.write_all(&socks5::build_auth_success()).await?; + } + + // --- CONNECT Request --- + let connect_req = read_exact(reader, 4).await?; + if connect_req[0] != socks5::VERSION { + return Err(socks5::Socks5Error::InvalidVersion(connect_req[0])); + } + if connect_req[1] != socks5::CMD_CONNECT { + writer + .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(reader, 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(reader, 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(reader, 1).await?; + let domain_len = len_byte[0] as usize; + let domain_bytes = read_exact(reader, domain_len).await?; + let domain = String::from_utf8_lossy(&domain_bytes).to_string(); + let port_bytes = read_exact(reader, 2).await?; + let port = u16::from_be_bytes([port_bytes[0], port_bytes[1]]); + socks5::Socks5Target::Domain { domain, port } + } + _ => { + writer + .write_all(&socks5::build_reply_failure(socks5::REPL_GENERAL_FAILURE)) + .await?; + return Err(socks5::Socks5Error::UnsupportedAddrType(atyp)); + } + }; + + Ok(target) +} + +/// Read exactly `len` bytes. +async fn read_exact( + reader: &mut tokio::net::tcp::OwnedReadHalf, + len: usize, +) -> Result, socks5::Socks5Error> { + let mut buf = vec![0u8; len]; + reader.read_exact(&mut buf).await?; + Ok(buf) +} + +/// Read exactly 2 bytes. +async fn read_exact_2( + reader: &mut tokio::net::tcp::OwnedReadHalf, +) -> Result<[u8; 2], socks5::Socks5Error> { + let buf = read_exact(reader, 2).await?; + Ok([buf[0], buf[1]]) +} + +// --------------------------------------------------------------------------- +// Reconnect helper +// --------------------------------------------------------------------------- + /// Attempt to reconnect to the tunnel, revalidating all security gates. #[allow(dead_code)] pub async fn reconnect_tunnel(config: &ConnectorConfig) -> Result<(), TunnelError> { @@ -740,82 +1219,12 @@ pub async fn reconnect_tunnel(config: &ConnectorConfig) -> Result<(), TunnelErro connect_tunnel(config.clone()).await } -/// Shared state for the persistent tunnel session. -/// Holds the authenticated HTTP sender behind a mutex, protected by -/// an `auth_complete` flag so SOCKS5 fails closed until auth succeeds. -#[allow(dead_code)] -pub struct TunnelSession { - /// Shared HTTP sender — one authenticated persistent tunnel session. - sender: Arc>>>, - /// Becomes true once POST /tunnel auth succeeds. - auth_complete: Arc, - /// Tunnel target host (for logging / Host header). - target_host: String, - /// Monotonically increasing stream counter for this session. - stream_counter: Arc, -} +// --------------------------------------------------------------------------- +// connect_tunnel: establishes a tunnel session (kept for test compatibility) +// --------------------------------------------------------------------------- -impl TunnelSession { - /// Send a raw request payload to the target via the persistent tunnel. - /// Returns the raw response bytes from the target. - async fn forward_raw( - &self, - target: &socks5::Socks5Target, - request_data: Bytes, - ) -> Result { - if !self.auth_complete.load(Ordering::SeqCst) { - return Err(TunnelError::Auth(AuthError::NotAuthenticated)); - } - let target_str = socks5::target_to_string(target); - let stream_id = self.stream_counter.fetch_add(1, Ordering::Relaxed); - let req_len = request_data.len(); - - let mut sender_guard = self.sender.lock().await; - - let req = Request::builder() - .method("POST") - .uri("/forward") - .header("X-Rustunnel-Target", &target_str) - .header("X-Rustunnel-Stream-ID", format!("session-{}", stream_id)) - .header("Host", &self.target_host) - .body(Full::new(request_data)) - .map_err(|e| TunnelError::Protocol(format!("build forward request: {}", e)))?; - - let resp = sender_guard - .send_request(req) - .await - .map_err(|e| TunnelError::Protocol(format!("forward request failed: {}", e)))?; - - let status = resp.status(); - let body = resp.into_body().collect().await.map_err(|e| { - TunnelError::Io(std::io::Error::other(format!("collect response: {}", e))) - })?; - - tracing::debug!( - "Stream {} forwarded {} bytes to {} via shared tunnel (status {})", - stream_id, - req_len, - target_str, - status - ); - - Ok(body.to_bytes()) - } -} - -/// 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. Stores the authenticated sender in a shared TunnelSession -/// 4. Starts the SOCKS5 proxy listener sharing that tunnel session -pub async fn run_connector_with_socks( - config: ConnectorConfig, - socks_addr: SocketAddr, -) -> Result<(), TunnelError> { - // First, authenticate through the tunnel and get the shared sender +/// Connect to the HTTPS tunnel and establish a session. +pub async fn connect_tunnel(config: ConnectorConfig) -> Result<(), TunnelError> { let client_config = tls::build_client_config( config.client_cert_path.as_ref(), config.client_key_path.as_ref(), @@ -827,9 +1236,7 @@ pub async fn run_connector_with_socks( })?; 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| { @@ -850,15 +1257,20 @@ pub async fn run_connector_with_socks( ))) })?; - tracing::info!( - "TLS handshake completed with {}:{} via HTTPS", - config.target_host, - config.target_port - ); + // Authenticate + let mut stream: std::pin::Pin>> = + Box::pin(tls_stream); - // Authenticate — keeps the connection alive in a background task - let (sender, _conn_handle) = - authenticate_https(tls_stream, &config.auth_token, &config.target_host).await?; + let request = format!( + "POST /tunnel HTTP/1.1\r\nHost: {}\r\nX-Rustunnel-Token: {}\r\nContent-Length: 0\r\n\r\n", + config.target_host, config.auth_token + ); + stream.write_all(request.as_bytes()).await?; + let status = read_http_response_status(&mut *stream).await?; + + if status != 200 { + return Err(TunnelError::Auth(AuthError::Invalid)); + } tracing::info!( "Tunnel session established with {}:{} via HTTPS /tunnel", @@ -866,799 +1278,65 @@ pub async fn run_connector_with_socks( config.target_port ); - // Build the shared session — one persistent tunnel for all SOCKS5 streams - let session = Arc::new(TunnelSession { - sender: Arc::new(Mutex::new(sender)), - auth_complete: Arc::new(AtomicBool::new(true)), - target_host: config.target_host.clone(), - stream_counter: Arc::new(AtomicUsize::new(0)), - }); - - tracing::info!("Authentication complete — SOCKS5 proxy ready for traffic (persistent tunnel)"); - - // Start the SOCKS5 proxy sharing the tunnel session - run_socks5_proxy_shared(socks_addr, session).await + Ok(()) } -/// Run the SOCKS5 proxy listener forwarding through a shared persistent tunnel session. -/// -/// All SOCKS5 streams are multiplexed over the single authenticated HTTPS connection. -pub async fn run_socks5_proxy_shared( - socks_addr: SocketAddr, - session: Arc, -) -> Result<(), TunnelError> { - run_socks5_proxy_impl(socks_addr, session, Arc::new(Notify::new())).await -} +// --------------------------------------------------------------------------- +// authenticate_https: backward-compatible name (reimplemented with manual HTTP) +// --------------------------------------------------------------------------- -/// Internal: run the SOCKS5 accept loop. -async fn run_socks5_proxy_impl( - socks_addr: SocketAddr, - session: Arc, - shutdown: Arc, -) -> Result<(), TunnelError> { - 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 {} (persistent tunnel session)", - socks_addr - ); - - // Register Ctrl-C handler for graceful shutdown - 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 sess = session.clone(); - - tokio::spawn(async move { - if let Err(e) = handle_socks5_connection_shared(stream, peer_addr, &sess).await { - tracing::warn!("SOCKS5 connection from {} failed: {}", peer_addr, e); - } - }); - } -} - -/// Legacy: per-request HTTPS connection (retained for backward compat / tests). -/// Deprecated in favor of `run_socks5_proxy_shared`. +/// Authenticate via HTTPS POST /tunnel with the auth token in header. +/// Returns ((), ()) for API compatibility — the underlying connection is now +/// ready for binary framing. #[allow(dead_code)] -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()) - } - })?; +pub async fn authenticate_https( + tls_stream: tokio_rustls::client::TlsStream, + auth_token: &str, + target_host: &str, +) -> Result<((), ()), TunnelError> { + let mut stream: std::pin::Pin>> = + Box::pin(tls_stream); - tracing::info!("SOCKS5 proxy listening on {}", socks_addr); - - 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 using the shared persistent tunnel session. -/// -/// All data is forwarded through the shared tunnel — no new HTTPS connections per request. -async fn handle_socks5_connection_shared( - mut stream: TcpStream, - peer_addr: SocketAddr, - session: &Arc, -) -> Result<(), socks5::Socks5Error> { - tracing::info!("SOCKS5 connection from {} (persistent tunnel)", peer_addr); - - // --- 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 {} (via persistent tunnel)", - peer_addr, - target_str + let request = format!( + "POST /tunnel HTTP/1.1\r\nHost: {}\r\nX-Rustunnel-Token: {}\r\nContent-Length: 0\r\n\r\n", + target_host, auth_token ); + stream.write_all(request.as_bytes()).await?; - // --- Forward through shared tunnel --- - match forward_socks_via_session(&mut stream, session, &target).await { - Ok(()) => { - tracing::info!( - "SOCKS5 stream from {} to {} completed (persistent tunnel)", - peer_addr, - target_str - ); - } - Err(e) => { - tracing::warn!( - "SOCKS5 stream from {} to {} failed: {}", - peer_addr, - target_str, - e - ); - return Err(e); - } + let status = read_http_response_status(&mut *stream).await?; + + if status != 200 { + return Err(TunnelError::Auth(AuthError::Invalid)); } - Ok(()) -} - -/// Handle a single SOCKS5 connection (legacy per-request HTTPS). -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 shared persistent tunnel session. -/// -/// Uses the single authenticated HTTPS connection shared across all SOCKS5 streams. -/// Stream IDs ensure responses are isolated and correctly attributed. -async fn forward_socks_via_session( - socks_stream: &mut TcpStream, - session: &TunnelSession, - target: &socks5::Socks5Target, -) -> Result<(), socks5::Socks5Error> { - // Send CONNECT success reply to SOCKS5 client - socks_stream - .write_all(&socks5::build_reply_success()) - .await?; - - let stream_id = next_stream_id(); - let target_str = socks5::target_to_string(target); - - tracing::info!( - "SOCKS5 stream {}: connected to target {} via persistent 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 via shared tunnel session - let request_bytes_owned: Bytes = request_bytes.into(); - let response_bytes = match session.forward_raw(target, request_bytes_owned).await { - Ok(b) => b, - Err(e) => { - tracing::warn!("Stream {}: forward via tunnel failed: {}", stream_id, e); - return Err(socks5::Socks5Error::Protocol(format!( - "tunnel forward failed: {}", - e - ))); - } - }; - - // Write target response back to SOCKS5 client - if !response_bytes.is_empty() { - socks_write.write_all(&response_bytes).await.map_err(|e| { - tracing::debug!("Stream {} write error: {}", stream_id, e); - socks5::Socks5Error::Io(e) - })?; - } - - // Read any remaining data from SOCKS5 client (for POST bodies etc.) - // Use a short timeout so we don't block indefinitely on keep-alive connections - let mut remaining = Vec::new(); - let read_result = tokio::time::timeout( - std::time::Duration::from_secs(5), - socks_read.read_to_end(&mut remaining), - ) - .await; - if read_result.is_ok() && !remaining.is_empty() { - let remaining_bytes: Bytes = remaining.into(); - let response2 = session - .forward_raw(target, remaining_bytes) - .await - .unwrap_or_default(); - if !response2.is_empty() { - socks_write.write_all(&response2).await.ok(); - } - } - - tracing::info!( - "SOCKS5 stream {}: completed (persistent tunnel, {} bytes response)", - stream_id, - response_bytes.len() - ); - - 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. -/// Legacy implementation — use `forward_socks_via_session` for persistent tunnel. -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) + Ok(((), ())) } // --------------------------------------------------------------------------- // Security gate state machine // --------------------------------------------------------------------------- -/// Check if all security gates have passed for a connection. #[derive(Debug, Clone, Copy, PartialEq, Eq)] #[allow(dead_code)] pub enum SecurityGateStatus { - /// TLS not yet negotiated TlsPending, - /// TLS negotiated, client cert check pending ClientCertPending, - /// mTLS established, auth pending AuthPending, - /// All gates passed, tunnel ready Authenticated, - /// A gate failed, tunnel rejected Failed, } -/// State machine for security gate tracking. #[allow(dead_code)] #[derive(Debug, Clone)] pub struct SecurityGateState { - status: Arc>, + status: Arc>, } #[allow(dead_code)] impl SecurityGateState { pub fn new() -> Self { Self { - status: Arc::new(Mutex::new(SecurityGateStatus::TlsPending)), + status: Arc::new(tokio::sync::Mutex::new(SecurityGateStatus::TlsPending)), } } @@ -1684,7 +1362,7 @@ impl SecurityGateState { #[cfg(test)] mod tests { use super::*; - use std::sync::atomic::{AtomicU16, Ordering}; + use std::sync::atomic::AtomicU16; static PORT_COUNTER: AtomicU16 = AtomicU16::new(50000); @@ -1705,55 +1383,15 @@ mod tests { }; } - /// Helper: quick-connect that authenticates then immediately returns (no long-lived session). - /// Used by tests that need a single connect/auth cycle without waiting for shutdown. + /// Helper: quick-connect that authenticates then immediately returns. async fn quick_connect(config: &ConnectorConfig) -> Result<(), TunnelError> { - 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(TunnelError::Tls)?; - - let tls_connector = tokio_rustls::TlsConnector::from(client_config); - - let target_addr = resolve_target(&config.target_host, config.target_port) - .await - .map_err(|e| { - TunnelError::Io(std::io::Error::new( - std::io::ErrorKind::InvalidInput, - e.to_string(), - )) - })?; - - 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 failed: {}", - e - ))) - })?; - - // authenticate_https spawns the connection driver internally and returns - // (sender, JoinHandle). We drop both — auth success/failure is enough. - let (_sender, _conn_handle) = - authenticate_https(tls_stream, &config.auth_token, &config.target_host).await?; - Ok(()) + connect_tunnel(config.clone()).await } + // ========================================================================= + // Auth / mTLS tests + // ========================================================================= + #[tokio::test] async fn valid_mtls_and_auth_establishes_tunnel() { let dir = generate_test_creds(); @@ -1770,7 +1408,6 @@ mod tests { }; let listener_task = tokio::spawn(async move { run_listener(listener_config).await }); - tokio::time::sleep(std::time::Duration::from_millis(200)).await; let connector_config = ConnectorConfig { @@ -1807,7 +1444,6 @@ mod tests { }; let listener_task = tokio::spawn(async move { run_listener(listener_config).await }); - tokio::time::sleep(std::time::Duration::from_millis(200)).await; let connector_config = ConnectorConfig { @@ -1845,10 +1481,8 @@ mod tests { }; let listener_task = tokio::spawn(async move { run_listener(listener_config).await }); - tokio::time::sleep(std::time::Duration::from_millis(200)).await; - // Client from a different CA — should fail at TLS let connector_config = ConnectorConfig { target_host: "127.0.0.1".to_string(), target_port: port, @@ -1874,7 +1508,6 @@ mod tests { let port = get_free_port(); let bind_addr: SocketAddr = format!("127.0.0.1:{}", port).parse().unwrap(); - // Bind the port first let _listener = TcpListener::bind(bind_addr).await.unwrap(); let config = ListenerConfig { @@ -1906,6 +1539,47 @@ mod tests { assert!(result.is_err(), "Missing TLS config should fail closed"); } + #[tokio::test] + async fn invalid_server_cert_rejected_by_client() { + let dir1 = generate_test_creds(); + let dir2 = generate_test_creds(); + let port = get_free_port(); + let bind_addr: SocketAddr = format!("127.0.0.1:{}", port).parse().unwrap(); + + let listener_config = ListenerConfig { + bind_addr, + server_cert_path: path_arc!(dir1.path().join("server.crt")), + server_key_path: path_arc!(dir1.path().join("server.key")), + ca_cert_path: path_arc!(dir1.path().join("ca.pem")), + auth_token: Arc::new("token".to_string()), + }; + + let listener_task = tokio::spawn(async move { run_listener(listener_config).await }); + tokio::time::sleep(std::time::Duration::from_millis(200)).await; + + let connector_config = ConnectorConfig { + target_host: "127.0.0.1".to_string(), + target_port: port, + client_cert_path: path_arc!(dir2.path().join("client.crt")), + client_key_path: path_arc!(dir2.path().join("client.key")), + ca_cert_path: path_arc!(dir2.path().join("ca.pem")), + auth_token: Arc::new("token".to_string()), + }; + + let result = quick_connect(&connector_config).await; + assert!( + result.is_err(), + "Client should reject server cert not signed by expected CA: {:?}", + result + ); + + listener_task.abort(); + } + + // ========================================================================= + // Utility tests + // ========================================================================= + #[test] fn load_auth_token_from_value() { let token = load_auth_token(Some("my-token"), None).unwrap(); @@ -1941,6 +1615,39 @@ mod tests { assert!(!constant_time_compare("abc", "")); } + // ========================================================================= + // Resolution tests + // ========================================================================= + + #[tokio::test] + async fn resolve_hostname_success() { + let addrs = resolve_host("localhost", 80).await.unwrap(); + assert!(!addrs.is_empty(), "localhost should resolve"); + } + + #[tokio::test] + async fn resolve_empty_hostname_fails() { + let result = resolve_host("", 80).await; + assert!(result.is_err()); + } + + #[tokio::test] + async fn resolve_invalid_hostname_fails() { + let result = resolve_host("this.is.not.a.real.host.example", 80).await; + assert!(result.is_err()); + } + + #[tokio::test] + async fn resolve_target_success() { + let addr = resolve_target("127.0.0.1", 80).await.unwrap(); + assert_eq!(addr.ip().to_string(), "127.0.0.1"); + assert_eq!(addr.port(), 80); + } + + // ========================================================================= + // Security gate tests + // ========================================================================= + #[tokio::test] async fn security_gate_state_transitions() { let state = SecurityGateState::new(); @@ -1952,107 +1659,6 @@ mod tests { assert_eq!(state.get_status().await, SecurityGateStatus::Authenticated); } - #[tokio::test] - async fn reconnect_revalidates_security_gates() { - let dir = generate_test_creds(); - let token = std::fs::read_to_string(dir.path().join("token.txt")).unwrap(); - let port = get_free_port(); - let bind_addr: SocketAddr = format!("127.0.0.1:{}", port).parse().unwrap(); - - let listener_config = ListenerConfig { - bind_addr, - server_cert_path: path_arc!(dir.path().join("server.crt")), - server_key_path: path_arc!(dir.path().join("server.key")), - ca_cert_path: path_arc!(dir.path().join("ca.pem")), - auth_token: Arc::new(token.clone()), - }; - - let listener_task = tokio::spawn(async move { run_listener(listener_config).await }); - - tokio::time::sleep(std::time::Duration::from_millis(200)).await; - - let connector_config = ConnectorConfig { - target_host: "127.0.0.1".to_string(), - target_port: port, - client_cert_path: path_arc!(dir.path().join("client.crt")), - client_key_path: path_arc!(dir.path().join("client.key")), - ca_cert_path: path_arc!(dir.path().join("ca.pem")), - auth_token: Arc::new(token.clone()), - }; - - // First connect succeeds - let result = quick_connect(&connector_config).await; - assert!(result.is_ok()); - - // Reconnect with wrong token should fail - let bad_config = ConnectorConfig { - target_host: "127.0.0.1".to_string(), - target_port: port, - client_cert_path: path_arc!(dir.path().join("client.crt")), - client_key_path: path_arc!(dir.path().join("client.key")), - ca_cert_path: path_arc!(dir.path().join("ca.pem")), - auth_token: Arc::new("bad-token".to_string()), - }; - let result = quick_connect(&bad_config).await; - assert!( - result.is_err(), - "Reconnect with bad auth should fail: {:?}", - result - ); - - // Reconnect with correct token should succeed - let result = quick_connect(&connector_config).await; - assert!( - result.is_ok(), - "Reconnect with valid auth should succeed: {:?}", - result - ); - - listener_task.abort(); - } - - #[tokio::test] - async fn https_is_default_carrier() { - let dir = generate_test_creds(); - let token = std::fs::read_to_string(dir.path().join("token.txt")).unwrap(); - let port = get_free_port(); - let bind_addr: SocketAddr = format!("127.0.0.1:{}", port).parse().unwrap(); - - let listener_config = ListenerConfig { - bind_addr, - server_cert_path: path_arc!(dir.path().join("server.crt")), - server_key_path: path_arc!(dir.path().join("server.key")), - ca_cert_path: path_arc!(dir.path().join("ca.pem")), - auth_token: Arc::new(token.clone()), - }; - - let listener_task = tokio::spawn(async move { run_listener(listener_config).await }); - - tokio::time::sleep(std::time::Duration::from_millis(200)).await; - - // Try plain TCP connection without TLS — should fail at TLS layer - let stream_result = TcpStream::connect(bind_addr).await; - if stream_result.is_ok() { - // TCP connects but the server expects TLS, so a plain TCP client - // won't get meaningful data — the server will fail on TLS accept - } - - // Verify HTTPS connection works - let connector_config = ConnectorConfig { - target_host: "127.0.0.1".to_string(), - target_port: port, - client_cert_path: path_arc!(dir.path().join("client.crt")), - client_key_path: path_arc!(dir.path().join("client.key")), - ca_cert_path: path_arc!(dir.path().join("ca.pem")), - auth_token: Arc::new(token.clone()), - }; - - let result = quick_connect(&connector_config).await; - assert!(result.is_ok(), "HTTPS mTLS should work"); - - listener_task.abort(); - } - #[tokio::test] async fn forwarding_starts_only_after_all_gates() { let state = SecurityGateState::new(); @@ -2080,7 +1686,6 @@ mod tests { }; let listener_task = tokio::spawn(async move { run_listener(listener_config).await }); - tokio::time::sleep(std::time::Duration::from_millis(200)).await; let connector_config = ConnectorConfig { @@ -2103,73 +1708,62 @@ mod tests { } #[tokio::test] - async fn invalid_server_cert_rejected_by_client() { - let dir1 = generate_test_creds(); - let dir2 = generate_test_creds(); + async fn reconnect_revalidates_security_gates() { + let dir = generate_test_creds(); + let token = std::fs::read_to_string(dir.path().join("token.txt")).unwrap(); let port = get_free_port(); let bind_addr: SocketAddr = format!("127.0.0.1:{}", port).parse().unwrap(); let listener_config = ListenerConfig { bind_addr, - server_cert_path: path_arc!(dir1.path().join("server.crt")), - server_key_path: path_arc!(dir1.path().join("server.key")), - ca_cert_path: path_arc!(dir1.path().join("ca.pem")), - auth_token: Arc::new("token".to_string()), + server_cert_path: path_arc!(dir.path().join("server.crt")), + server_key_path: path_arc!(dir.path().join("server.key")), + ca_cert_path: path_arc!(dir.path().join("ca.pem")), + auth_token: Arc::new(token.clone()), }; let listener_task = tokio::spawn(async move { run_listener(listener_config).await }); - tokio::time::sleep(std::time::Duration::from_millis(200)).await; let connector_config = ConnectorConfig { target_host: "127.0.0.1".to_string(), target_port: port, - client_cert_path: path_arc!(dir2.path().join("client.crt")), - client_key_path: path_arc!(dir2.path().join("client.key")), - ca_cert_path: path_arc!(dir2.path().join("ca.pem")), - auth_token: Arc::new("token".to_string()), + client_cert_path: path_arc!(dir.path().join("client.crt")), + client_key_path: path_arc!(dir.path().join("client.key")), + ca_cert_path: path_arc!(dir.path().join("ca.pem")), + auth_token: Arc::new(token.clone()), }; let result = quick_connect(&connector_config).await; + assert!(result.is_ok()); + + let bad_config = ConnectorConfig { + target_host: "127.0.0.1".to_string(), + target_port: port, + client_cert_path: path_arc!(dir.path().join("client.crt")), + client_key_path: path_arc!(dir.path().join("client.key")), + ca_cert_path: path_arc!(dir.path().join("ca.pem")), + auth_token: Arc::new("bad-token".to_string()), + }; + let result = quick_connect(&bad_config).await; assert!( result.is_err(), - "Client should reject server cert not signed by expected CA: {:?}", + "Reconnect with bad auth should fail: {:?}", + result + ); + + let result = quick_connect(&connector_config).await; + assert!( + result.is_ok(), + "Reconnect with valid auth should succeed: {:?}", result ); listener_task.abort(); } - // --- Tests for scrutiny fixes --- - #[tokio::test] - async fn resolve_hostname_success() { - let addrs = resolve_host("localhost", 80).await.unwrap(); - assert!(!addrs.is_empty(), "localhost should resolve"); - } - - #[tokio::test] - async fn resolve_empty_hostname_fails() { - let result = resolve_host("", 80).await; - assert!(result.is_err()); - } - - #[tokio::test] - async fn resolve_invalid_hostname_fails() { - let result = resolve_host("this.is.not.a.real.host.example", 80).await; - assert!(result.is_err()); - } - - #[tokio::test] - async fn resolve_target_success() { - let addr = resolve_target("127.0.0.1", 80).await.unwrap(); - assert_eq!(addr.ip().to_string(), "127.0.0.1"); - assert_eq!(addr.port(), 80); - } - - #[tokio::test] - async fn https_endpoint_semantics() { - // Verify the listener responds to POST /tunnel with proper HTTP status codes + async fn https_is_default_carrier() { let dir = generate_test_creds(); let token = std::fs::read_to_string(dir.path().join("token.txt")).unwrap(); let port = get_free_port(); @@ -2186,527 +1780,647 @@ mod tests { let listener_task = tokio::spawn(async move { run_listener(listener_config).await }); tokio::time::sleep(std::time::Duration::from_millis(200)).await; - // Connect and send a wrong path — should get 404 - let client_config = tls::build_client_config( - &dir.path().join("client.crt"), - &dir.path().join("client.key"), - &dir.path().join("ca.pem"), - ) - .unwrap(); - let tls_connector = tokio_rustls::TlsConnector::from(client_config); - let stream = TcpStream::connect(bind_addr).await.unwrap(); - let tls_stream = tls_connector - .connect( - rustls::pki_types::ServerName::try_from("127.0.0.1").unwrap(), - stream, - ) - .await - .unwrap(); - let io = hyper_util::rt::TokioIo::new(tls_stream); + // Plain TCP connects but the server expects TLS + let stream_result = TcpStream::connect(bind_addr).await; + assert!(stream_result.is_ok(), "TCP should connect to TLS listener"); - let request = Request::builder() - .method("GET") - .uri("/nonexistent") - .body(Full::new(Bytes::new())) - .unwrap(); - - let (mut sender, connection) = hyper::client::conn::http1::handshake(io).await.unwrap(); - tokio::spawn(connection); - - let resp = sender.send_request(request).await.unwrap(); - assert_eq!(resp.status(), StatusCode::NOT_FOUND); - - listener_task.abort(); - } - - #[tokio::test] - async fn tunnel_endpoint_accepts_get() { - let dir = generate_test_creds(); - let token = std::fs::read_to_string(dir.path().join("token.txt")).unwrap(); - let port = get_free_port(); - let bind_addr: SocketAddr = format!("127.0.0.1:{}", port).parse().unwrap(); - - let listener_config = ListenerConfig { - bind_addr, - server_cert_path: path_arc!(dir.path().join("server.crt")), - server_key_path: path_arc!(dir.path().join("server.key")), - ca_cert_path: path_arc!(dir.path().join("ca.pem")), - auth_token: Arc::new(token.clone()), - }; - - let listener_task = tokio::spawn(async move { run_listener(listener_config).await }); - tokio::time::sleep(std::time::Duration::from_millis(200)).await; - - let client_config = tls::build_client_config( - &dir.path().join("client.crt"), - &dir.path().join("client.key"), - &dir.path().join("ca.pem"), - ) - .unwrap(); - let tls_connector = tokio_rustls::TlsConnector::from(client_config); - let stream = TcpStream::connect(bind_addr).await.unwrap(); - let tls_stream = tls_connector - .connect( - rustls::pki_types::ServerName::try_from("127.0.0.1").unwrap(), - stream, - ) - .await - .unwrap(); - let io = hyper_util::rt::TokioIo::new(tls_stream); - - let request = Request::builder() - .method("GET") - .uri("/tunnel") - .body(Full::new(Bytes::new())) - .unwrap(); - - let (mut sender, connection) = hyper::client::conn::http1::handshake(io).await.unwrap(); - tokio::spawn(connection); - - let resp = sender.send_request(request).await.unwrap(); - assert_eq!(resp.status(), StatusCode::OK); - - listener_task.abort(); - } - - #[tokio::test] - async fn hostname_resolution_no_panic() { - // Ensure hostname resolution never panics — always returns errors - let result = resolve_host("", 80).await; - assert!(result.is_err()); - - let result = resolve_host("not-a-valid-hostname-xyz123", 80).await; - assert!(result.is_err()); - } - - // ========================================================================= - // Tests for persistent tunnel multiplexing - // ========================================================================= - - /// Helper: start an echo server that reads all data then echoes it back. - /// Uses a short inactivity timeout to detect client EOF over keep-alive. - /// Returns (echo_port, task_handle). - async fn spawn_echo_server() -> (u16, tokio::task::JoinHandle<()>) { - let echo_port = get_free_port(); - let echo_addr: SocketAddr = format!("127.0.0.1:{}", echo_port).parse().unwrap(); - let listener = TcpListener::bind(echo_addr).await.unwrap(); - let handle = tokio::spawn(async move { - while let Ok((stream, _)) = listener.accept().await { - tokio::spawn(async move { - let mut buf = vec![0u8; 65536]; - let mut s = stream; - // Collect all data with a short timeout to detect end of request - let mut all_data = Vec::new(); - loop { - match tokio::time::timeout( - std::time::Duration::from_millis(500), - s.read(&mut buf), - ) - .await - { - Err(_) | Ok(Err(_)) | Ok(Ok(0)) => break, - Ok(Ok(n)) => all_data.extend_from_slice(&buf[..n]), - } - } - // Echo everything back - let _ = s.write_all(&all_data).await; - let _ = s.shutdown().await; - }); - } - }); - (echo_port, handle) - } - - /// Helper: build a connector config for the given listener port. - #[allow(dead_code)] - fn build_connector_config(dir: &tempfile::TempDir, port: u16, token: &str) -> ConnectorConfig { - ConnectorConfig { + let connector_config = ConnectorConfig { target_host: "127.0.0.1".to_string(), target_port: port, - client_cert_path: Arc::from(dir.path().join("client.crt").as_path()), - client_key_path: Arc::from(dir.path().join("client.key").as_path()), - ca_cert_path: Arc::from(dir.path().join("ca.pem").as_path()), - auth_token: Arc::new(token.to_string()), - } - } - - /// Build a TunnelSession by authenticating against a running listener. - async fn build_session( - dir: &tempfile::TempDir, - listen_port: u16, - token: &str, - ) -> Result<(Arc, tokio::task::JoinHandle<()>), TunnelError> { - let client_config = tls::build_client_config( - &dir.path().join("client.crt"), - &dir.path().join("client.key"), - &dir.path().join("ca.pem"), - ) - .map_err(TunnelError::Tls)?; - - let tls_connector = tokio_rustls::TlsConnector::from(client_config); - let target_addr: SocketAddr = format!("127.0.0.1:{}", listen_port).parse().unwrap(); - let stream = TcpStream::connect(target_addr) - .await - .map_err(TunnelError::Io)?; - let server_name = tls::server_name_from_host("127.0.0.1").unwrap(); - let tls_stream = tls_connector - .connect(server_name, stream) - .await - .map_err(|e| { - TunnelError::Tls(crate::errors::TlsError::VerificationFailed(e.to_string())) - })?; - - let (sender, conn_handle) = authenticate_https(tls_stream, token, "127.0.0.1").await?; - - let session = Arc::new(TunnelSession { - sender: Arc::new(Mutex::new(sender)), - auth_complete: Arc::new(AtomicBool::new(true)), - target_host: "127.0.0.1".to_string(), - stream_counter: Arc::new(AtomicUsize::new(0)), - }); - - Ok((session, conn_handle)) - } - - /// Test: multiple sequential forwards over the same persistent session. - /// Proves the session is reused (single connection, multiple streams). - #[tokio::test] - async fn persistent_session_sequential_streams() { - let dir = generate_test_creds(); - let token = std::fs::read_to_string(dir.path().join("token.txt")).unwrap(); - - let (echo_port, echo_handle) = spawn_echo_server().await; - tokio::time::sleep(std::time::Duration::from_millis(50)).await; - - // Start listener - let listen_port = get_free_port(); - let bind_addr: SocketAddr = format!("127.0.0.1:{}", listen_port).parse().unwrap(); - let listener_config = ListenerConfig { - bind_addr, - server_cert_path: Arc::from(dir.path().join("server.crt").as_path()), - server_key_path: Arc::from(dir.path().join("server.key").as_path()), - ca_cert_path: Arc::from(dir.path().join("ca.pem").as_path()), + client_cert_path: path_arc!(dir.path().join("client.crt")), + client_key_path: path_arc!(dir.path().join("client.key")), + ca_cert_path: path_arc!(dir.path().join("ca.pem")), auth_token: Arc::new(token.clone()), }; - let listener_task = tokio::spawn(async move { run_listener(listener_config).await }); - tokio::time::sleep(std::time::Duration::from_millis(200)).await; - // Build session - let (session, _conn_h) = build_session(&dir, listen_port, &token).await.unwrap(); - - // Sequential forwards - for payload_str in ["hello", "world", "multiplex"] { - let payload = Bytes::from(payload_str.as_bytes().to_vec()); - let target = - socks5::Socks5Target::Ipv4(std::net::Ipv4Addr::new(127, 0, 0, 1), echo_port); - let resp = session.forward_raw(&target, payload).await.unwrap(); - assert_eq!( - &resp[..], - payload_str.as_bytes(), - "Echo mismatch for '{}'", - payload_str - ); - } - - listener_task.abort(); - echo_handle.abort(); - } - - /// Test: multiple concurrent forwards over the same persistent session. - /// Each stream gets a unique payload and the responses are correctly isolated. - #[tokio::test] - async fn persistent_session_concurrent_streams_isolated() { - let dir = generate_test_creds(); - let token = std::fs::read_to_string(dir.path().join("token.txt")).unwrap(); - - let (echo_port, echo_handle) = spawn_echo_server().await; - tokio::time::sleep(std::time::Duration::from_millis(50)).await; - - let listen_port = get_free_port(); - let bind_addr: SocketAddr = format!("127.0.0.1:{}", listen_port).parse().unwrap(); - let listener_config = ListenerConfig { - bind_addr, - server_cert_path: Arc::from(dir.path().join("server.crt").as_path()), - server_key_path: Arc::from(dir.path().join("server.key").as_path()), - ca_cert_path: Arc::from(dir.path().join("ca.pem").as_path()), - auth_token: Arc::new(token.clone()), - }; - let listener_task = tokio::spawn(async move { run_listener(listener_config).await }); - tokio::time::sleep(std::time::Duration::from_millis(200)).await; - - let (session, _conn_h) = build_session(&dir, listen_port, &token).await.unwrap(); - - // Launch 5 concurrent forwards, each with a unique payload - let payloads = vec![ - "stream-alpha-payload".to_string(), - "stream-beta-payload".to_string(), - "stream-gamma-payload".to_string(), - "stream-delta-payload".to_string(), - "stream-epsilon-payload".to_string(), - ]; - let target = socks5::Socks5Target::Ipv4(std::net::Ipv4Addr::new(127, 0, 0, 1), echo_port); - - let mut handles = Vec::new(); - for payload_str in &payloads { - let sess = session.clone(); - let t = target.clone(); - let p = Bytes::from(payload_str.as_bytes().to_vec()); - handles.push(tokio::spawn(async move { sess.forward_raw(&t, p).await })); - } - - // Collect results and verify each response matches its payload - for (i, handle) in handles.into_iter().enumerate() { - let resp = handle.await.unwrap().unwrap(); - assert_eq!( - &resp[..], - payloads[i].as_bytes(), - "Stream {} response mismatch: expected '{}', got '{}'", - i, - payloads[i], - String::from_utf8_lossy(&resp) - ); - } - - listener_task.abort(); - echo_handle.abort(); - } - - /// Test: stream counter increments monotonically across concurrent streams. - #[tokio::test] - async fn persistent_session_stream_counter_increments() { - let dir = generate_test_creds(); - let token = std::fs::read_to_string(dir.path().join("token.txt")).unwrap(); - - let listen_port = get_free_port(); - let bind_addr: SocketAddr = format!("127.0.0.1:{}", listen_port).parse().unwrap(); - let listener_config = ListenerConfig { - bind_addr, - server_cert_path: Arc::from(dir.path().join("server.crt").as_path()), - server_key_path: Arc::from(dir.path().join("server.key").as_path()), - ca_cert_path: Arc::from(dir.path().join("ca.pem").as_path()), - auth_token: Arc::new(token.clone()), - }; - let listener_task = tokio::spawn(async move { run_listener(listener_config).await }); - tokio::time::sleep(std::time::Duration::from_millis(200)).await; - - let (session, _conn_h) = build_session(&dir, listen_port, &token).await.unwrap(); - - // Check initial counter - let initial = session.stream_counter.load(Ordering::SeqCst); - assert_eq!(initial, 0, "Counter should start at 0"); - - // Do some forwards - let target = socks5::Socks5Target::Ipv4(std::net::Ipv4Addr::new(127, 0, 0, 1), 4181); - let mut handles = Vec::new(); - for _ in 0..5 { - let sess = session.clone(); - let t = target.clone(); - handles.push(tokio::spawn(async move { - sess.forward_raw(&t, Bytes::from(vec![0u8; 10])).await - })); - } - // Wait for all handles (they may fail — that's fine, we just care about the counter) - for h in handles { - let _ = h.await; - } - - // Counter should have incremented - let final_counter = session.stream_counter.load(Ordering::SeqCst); - assert!( - final_counter >= 5, - "Counter should be >= 5, got {}", - final_counter - ); + let result = quick_connect(&connector_config).await; + assert!(result.is_ok(), "HTTPS mTLS should work"); listener_task.abort(); } - /// Test: TunnelSession rejects forwards when auth is not complete. - #[tokio::test] - async fn tunnel_session_fails_before_auth() { - // Create a real sender via a dummy listener so the TunnelSession is well-formed. - // The sender won't actually be used because auth_complete is false. - let (sender, _) = create_placeholder_sender().await; - let session = Arc::new(TunnelSession { - sender: Arc::new(Mutex::new(sender)), - auth_complete: Arc::new(AtomicBool::new(false)), - target_host: "127.0.0.1".to_string(), - stream_counter: Arc::new(AtomicUsize::new(0)), - }); + // ========================================================================= + // HTTP helper tests + // ========================================================================= - let target = socks5::Socks5Target::Ipv4(std::net::Ipv4Addr::new(127, 0, 0, 1), 4181); - let result = session - .forward_raw(&target, Bytes::from(vec![0u8; 10])) - .await; - assert!( - result.is_err(), - "Forward should fail before auth is complete" - ); - assert!( - matches!(result, Err(TunnelError::Auth(AuthError::NotAuthenticated))), - "Should be NotAuthenticated error" - ); + #[tokio::test] + async fn parse_http_request_basic() { + let data = b"POST /tunnel HTTP/1.1\r\nHost: localhost\r\nX-Rustunnel-Token: abc123\r\n\r\n"; + let mut cursor = std::io::Cursor::new(data.as_slice()); + let (method, path, headers) = parse_http_request(&mut cursor).await.unwrap(); + assert_eq!(method, "POST"); + assert_eq!(path, "/tunnel"); + assert_eq!(headers.get("x-rustunnel-token").unwrap(), "abc123"); } - /// Helper: create a placeholder sender (never used, just for type safety) within the async runtime. - async fn create_placeholder_sender() -> ( - hyper::client::conn::http1::SendRequest>, - tokio::task::JoinHandle<()>, - ) { - // Create a dummy listener to accept a connection - let dummy_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let dummy_addr = dummy_listener.local_addr().unwrap(); - - // Spawn an acceptor - let _acceptor = tokio::spawn(async move { - let (stream, _) = dummy_listener.accept().await.unwrap(); - let io = hyper_util::rt::TokioIo::new(stream); - let svc = hyper::service::service_fn(|_req: Request| async { - let body: http_body_util::combinators::BoxBody = - BoxBody::new(Full::new(Bytes::new()).map_err(|_| unreachable!())); - Ok::<_, hyper::http::Error>(Response::new(body)) - }); - let _ = hyper::server::conn::http1::Builder::new() - .serve_connection(io, svc) - .await; - }); - - // Connect to it - let stream = TcpStream::connect(dummy_addr).await.unwrap(); - let io = hyper_util::rt::TokioIo::new(stream); - let (sender, conn) = hyper::client::conn::http1::handshake(io).await.unwrap(); - let conn_handle = tokio::spawn(async move { - let _ = conn.await; - }); - (sender, conn_handle) + #[tokio::test] + async fn parse_http_request_empty_fails() { + let mut cursor = std::io::Cursor::new(&[] as &[u8]); + let result = parse_http_request(&mut cursor).await; + assert!(result.is_err()); } - /// Test: E2E concurrent SOCKS5 streams share a single tunnel session. - /// This is the main integration test proving VAL-SOCKS-006. #[tokio::test] - async fn e2e_concurrent_socks5_streams_shared_tunnel() { + async fn read_http_response_status_ok() { + let data = b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n"; + let mut cursor = std::io::Cursor::new(data.as_slice()); + let status = read_http_response_status(&mut cursor).await.unwrap(); + assert_eq!(status, 200); + } + + #[tokio::test] + async fn read_http_response_status_unauthorized() { + let data = b"HTTP/1.1 401 Unauthorized\r\nContent-Length: 0\r\n\r\n"; + let mut cursor = std::io::Cursor::new(data.as_slice()); + let status = read_http_response_status(&mut cursor).await.unwrap(); + assert_eq!(status, 401); + } + + // ========================================================================= + // Framing E2E tests + // ========================================================================= + + /// Test: full E2E flow with binary framing — listener, connector, target. + #[tokio::test] + async fn e2e_framing_basic_forward() { let dir = generate_test_creds(); let token = std::fs::read_to_string(dir.path().join("token.txt")).unwrap(); - - // Start echo server - let (echo_port, echo_handle) = spawn_echo_server().await; - tokio::time::sleep(std::time::Duration::from_millis(50)).await; - - // Start listener - let listen_port = get_free_port(); - let bind_addr: SocketAddr = format!("127.0.0.1:{}", listen_port).parse().unwrap(); - let listener_config = ListenerConfig { - bind_addr, - server_cert_path: Arc::from(dir.path().join("server.crt").as_path()), - server_key_path: Arc::from(dir.path().join("server.key").as_path()), - ca_cert_path: Arc::from(dir.path().join("ca.pem").as_path()), - auth_token: Arc::new(token.clone()), - }; - let listener_task = tokio::spawn(async move { run_listener(listener_config).await }); - tokio::time::sleep(std::time::Duration::from_millis(200)).await; - - // Build shared session - let (session, _conn_h) = build_session(&dir, listen_port, &token).await.unwrap(); - - // Start SOCKS5 proxy on a unique port + let listener_port = get_free_port(); let socks_port = get_free_port(); - let socks_addr: SocketAddr = format!("127.0.0.1:{}", socks_port).parse().unwrap(); - let session_clone = session.clone(); - let socks_shutdown = Arc::new(Notify::new()); - let socks_shutdown_clone = socks_shutdown.clone(); - let socks_task = tokio::spawn(async move { - run_socks5_proxy_impl(socks_addr, session_clone, socks_shutdown_clone).await + let target_port = get_free_port(); + + // Start a simple echo target + let target_listener = TcpListener::bind(format!("127.0.0.1:{}", target_port)) + .await + .unwrap(); + let target_task = tokio::spawn(async move { + if let Ok((stream, _)) = target_listener.accept().await { + let (mut rd, mut wr) = stream.into_split(); + let mut buf = [0u8; 4096]; + if let Ok(n) = rd.read(&mut buf).await + && n > 0 + { + let _ = wr.write_all(&buf[..n]).await; + } + } }); - tokio::time::sleep(std::time::Duration::from_millis(100)).await; + // Start listener + let bind_addr: SocketAddr = format!("127.0.0.1:{}", listener_port).parse().unwrap(); + let listener_config = ListenerConfig { + bind_addr, + server_cert_path: path_arc!(dir.path().join("server.crt")), + server_key_path: path_arc!(dir.path().join("server.key")), + ca_cert_path: path_arc!(dir.path().join("ca.pem")), + auth_token: Arc::new(token.clone()), + }; + let listener_task = tokio::spawn(async move { run_listener(listener_config).await }); + tokio::time::sleep(std::time::Duration::from_millis(200)).await; - // Launch 3 concurrent SOCKS5 connections via raw TCP, each sending a unique HTTP request - let payloads = vec![ - "GET /a HTTP/1.1\r\nHost: x\r\n\r\n".to_string(), - "GET /b HTTP/1.1\r\nHost: x\r\n\r\n".to_string(), - "GET /c HTTP/1.1\r\nHost: x\r\n\r\n".to_string(), - ]; + // Start connector with SOCKS5 + let socks_addr: SocketAddr = format!("127.0.0.1:{}", socks_port).parse().unwrap(); + let connector_config = ConnectorConfig { + target_host: "127.0.0.1".to_string(), + target_port: listener_port, + client_cert_path: path_arc!(dir.path().join("client.crt")), + client_key_path: path_arc!(dir.path().join("client.key")), + ca_cert_path: path_arc!(dir.path().join("ca.pem")), + auth_token: Arc::new(token.clone()), + }; + let connector_task = + tokio::spawn( + async move { run_connector_with_socks(connector_config, socks_addr).await }, + ); + tokio::time::sleep(std::time::Duration::from_millis(300)).await; - let mut handles = Vec::new(); - for payload_str in &payloads { - let p = payload_str.clone(); - handles.push(tokio::spawn(async move { - // Connect to SOCKS5 - let mut stream = TcpStream::connect(format!("127.0.0.1:{}", socks_port)) + // Connect to SOCKS5 and send a request + let mut socks_stream = TcpStream::connect(socks_addr).await.unwrap(); + { + use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; + + // SOCKS5 greeting + socks_stream.write_all(&[0x05, 0x01, 0x00]).await.unwrap(); + let mut resp = [0u8; 2]; + socks_stream.read_exact(&mut resp).await.unwrap(); + assert_eq!(resp, [0x05, 0x00]); + + // CONNECT request (127.0.0.1:target_port) + let mut req = vec![0x05, 0x01, 0x00, 0x01, 127, 0, 0, 1]; + req.extend_from_slice(&target_port.to_be_bytes()); + socks_stream.write_all(&req).await.unwrap(); + + // Read reply + let mut reply = [0u8; 10]; + socks_stream.read_exact(&mut reply).await.unwrap(); + assert_eq!(reply[0], 0x05); + assert_eq!(reply[1], 0x00); // succeeded + + // Send HTTP request + let request = format!("GET / HTTP/1.1\r\nHost: 127.0.0.1:{}\r\n\r\n", target_port); + socks_stream.write_all(request.as_bytes()).await.unwrap(); + + // Read response (echo) + let mut response = vec![0u8; 4096]; + let n = socks_stream.read(&mut response).await.unwrap(); + let response_str = String::from_utf8_lossy(&response[..n]); + assert!( + response_str.contains("GET / HTTP/1.1"), + "Echo response should contain request: {}", + response_str + ); + } + + // Cleanup + connector_task.abort(); + listener_task.abort(); + target_task.abort(); + } + + /// Test: multiple concurrent streams share one tunnel session. + #[tokio::test] + async fn e2e_concurrent_streams_isolation() { + let dir = generate_test_creds(); + let token = std::fs::read_to_string(dir.path().join("token.txt")).unwrap(); + let listener_port = get_free_port(); + let socks_port = get_free_port(); + + // Start two echo targets on different ports + let target_port1 = get_free_port(); + let target_port2 = get_free_port(); + + let target1_listener = TcpListener::bind(format!("127.0.0.1:{}", target_port1)) + .await + .unwrap(); + let target2_listener = TcpListener::bind(format!("127.0.0.1:{}", target_port2)) + .await + .unwrap(); + + // Echo server helper + fn echo_task(listener: TcpListener) -> tokio::task::JoinHandle<()> { + tokio::spawn(async move { + if let Ok((stream, _)) = listener.accept().await { + let (mut rd, mut wr) = stream.into_split(); + let mut buf = [0u8; 4096]; + if let Ok(n) = rd.read(&mut buf).await + && n > 0 + { + let _ = wr.write_all(&buf[..n]).await; + } + } + }) + } + + let target1_task = echo_task(target1_listener); + let target2_task = echo_task(target2_listener); + + // Start listener + let bind_addr: SocketAddr = format!("127.0.0.1:{}", listener_port).parse().unwrap(); + let listener_config = ListenerConfig { + bind_addr, + server_cert_path: path_arc!(dir.path().join("server.crt")), + server_key_path: path_arc!(dir.path().join("server.key")), + ca_cert_path: path_arc!(dir.path().join("ca.pem")), + auth_token: Arc::new(token.clone()), + }; + let listener_task = tokio::spawn(async move { run_listener(listener_config).await }); + tokio::time::sleep(std::time::Duration::from_millis(200)).await; + + // Start connector + let socks_addr: SocketAddr = format!("127.0.0.1:{}", socks_port).parse().unwrap(); + let connector_config = ConnectorConfig { + target_host: "127.0.0.1".to_string(), + target_port: listener_port, + client_cert_path: path_arc!(dir.path().join("client.crt")), + client_key_path: path_arc!(dir.path().join("client.key")), + ca_cert_path: path_arc!(dir.path().join("ca.pem")), + auth_token: Arc::new(token.clone()), + }; + let connector_task = + tokio::spawn( + async move { run_connector_with_socks(connector_config, socks_addr).await }, + ); + tokio::time::sleep(std::time::Duration::from_millis(300)).await; + + // Connect two concurrent SOCKS5 streams + let msg1 = "STREAM1-UNIQUE-DATA"; + let msg2 = "STREAM2-UNIQUE-DATA"; + + let (result1, result2) = tokio::join!( + async { + let mut stream = TcpStream::connect(socks_addr).await.unwrap(); + socks5_greeting(&mut stream).await.unwrap(); + socks5_connect(&mut stream, "127.0.0.1", target_port1) .await .unwrap(); + socks5_echo(&mut stream, msg1).await + }, + async { + let mut stream = TcpStream::connect(socks_addr).await.unwrap(); + socks5_greeting(&mut stream).await.unwrap(); + socks5_connect(&mut stream, "127.0.0.1", target_port2) + .await + .unwrap(); + socks5_echo(&mut stream, msg2).await + } + ); - // SOCKS5 greeting - stream.write_all(&[0x05, 0x01, 0x00]).await.unwrap(); - let mut resp = [0u8; 2]; - stream.read_exact(&mut resp).await.unwrap(); - assert_eq!(resp, [0x05, 0x00]); + // Verify isolation: each stream got its own echo back + assert_eq!(result1, msg1, "Stream 1 data should be isolated"); + assert_eq!(result2, msg2, "Stream 2 data should be isolated"); - // CONNECT request — target is echo server - let _target = format!("{}:{}", "127.0.0.1", echo_port); - let mut req = vec![ - 0x05, 0x01, 0x00, 0x01, // VER CMD RSV ATYP=IPv4 - ]; - req.extend([127, 0, 0, 1]); - req.extend(&echo_port.to_be_bytes()); - stream.write_all(&req).await.unwrap(); + connector_task.abort(); + listener_task.abort(); + target1_task.abort(); + target2_task.abort(); + } - // Read CONNECT reply - let mut reply = [0u8; 10]; - stream.read_exact(&mut reply).await.unwrap(); - assert_eq!(reply[0], 0x05); // VER - assert_eq!(reply[1], 0x00); // SUCCEEDED + /// Test: target connection failure is reported without crashing. + #[tokio::test] + async fn e2e_target_failure_handled() { + let dir = generate_test_creds(); + let token = std::fs::read_to_string(dir.path().join("token.txt")).unwrap(); + let listener_port = get_free_port(); + let socks_port = get_free_port(); + let bad_target_port = get_free_port(); // no target listening - // Send HTTP request to echo server - stream.write_all(p.as_bytes()).await.unwrap(); + let bind_addr: SocketAddr = format!("127.0.0.1:{}", listener_port).parse().unwrap(); + let listener_config = ListenerConfig { + bind_addr, + server_cert_path: path_arc!(dir.path().join("server.crt")), + server_key_path: path_arc!(dir.path().join("server.key")), + ca_cert_path: path_arc!(dir.path().join("ca.pem")), + auth_token: Arc::new(token.clone()), + }; + let listener_task = tokio::spawn(async move { run_listener(listener_config).await }); + tokio::time::sleep(std::time::Duration::from_millis(200)).await; - // Read response - let mut buf = Vec::new(); - let mut chunk = vec![0u8; 4096]; - tokio::time::timeout(std::time::Duration::from_secs(5), async { - loop { - let n = match tokio::time::timeout( - std::time::Duration::from_secs(2), - stream.read(&mut chunk), - ) - .await - { - Ok(Ok(n)) => n, - _ => break, - }; - if n == 0 { - break; + let socks_addr: SocketAddr = format!("127.0.0.1:{}", socks_port).parse().unwrap(); + let connector_config = ConnectorConfig { + target_host: "127.0.0.1".to_string(), + target_port: listener_port, + client_cert_path: path_arc!(dir.path().join("client.crt")), + client_key_path: path_arc!(dir.path().join("client.key")), + ca_cert_path: path_arc!(dir.path().join("ca.pem")), + auth_token: Arc::new(token.clone()), + }; + let connector_task = + tokio::spawn( + async move { run_connector_with_socks(connector_config, socks_addr).await }, + ); + tokio::time::sleep(std::time::Duration::from_millis(300)).await; + + // Try to connect to a target that doesn't exist + let mut socks_stream = TcpStream::connect(socks_addr).await.unwrap(); + + socks5_greeting(&mut socks_stream).await.unwrap(); + // The connect should fail with a failure reply + socks5_connect(&mut socks_stream, "127.0.0.1", bad_target_port) + .await + .unwrap_err(); + + // Verify connector is still alive + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + assert!( + !connector_task.is_finished(), + "Connector should still be running after target failure" + ); + + connector_task.abort(); + listener_task.abort(); + } + + /// Test: listener-side target origin — verify the listener opens the target connection. + #[tokio::test] + async fn e2e_listener_side_target_origin() { + // This test verifies that the target TCP connection is opened from the + // listener side, not the connector side. We do this by verifying that + // the full tunnel path works correctly. + let dir = generate_test_creds(); + let token = std::fs::read_to_string(dir.path().join("token.txt")).unwrap(); + let listener_port = get_free_port(); + let socks_port = get_free_port(); + let target_port = get_free_port(); + + // Start target that records the source IP + let target_listener = TcpListener::bind(format!("127.0.0.1:{}", target_port)) + .await + .unwrap(); + let target_task = tokio::spawn(async move { + if let Ok((stream, _)) = target_listener.accept().await { + // The source should be localhost (listener opened it) + let peer = stream.peer_addr().unwrap(); + assert_eq!(peer.ip().to_string(), "127.0.0.1"); + // Echo back + let (mut rd, mut wr) = stream.into_split(); + let mut buf = [0u8; 4096]; + if let Ok(n) = rd.read(&mut buf).await + && n > 0 + { + let _ = wr.write_all(&buf[..n]).await; + } + } + }); + + let bind_addr: SocketAddr = format!("127.0.0.1:{}", listener_port).parse().unwrap(); + let listener_config = ListenerConfig { + bind_addr, + server_cert_path: path_arc!(dir.path().join("server.crt")), + server_key_path: path_arc!(dir.path().join("server.key")), + ca_cert_path: path_arc!(dir.path().join("ca.pem")), + auth_token: Arc::new(token.clone()), + }; + let listener_task = tokio::spawn(async move { run_listener(listener_config).await }); + tokio::time::sleep(std::time::Duration::from_millis(200)).await; + + let socks_addr: SocketAddr = format!("127.0.0.1:{}", socks_port).parse().unwrap(); + let connector_config = ConnectorConfig { + target_host: "127.0.0.1".to_string(), + target_port: listener_port, + client_cert_path: path_arc!(dir.path().join("client.crt")), + client_key_path: path_arc!(dir.path().join("client.key")), + ca_cert_path: path_arc!(dir.path().join("ca.pem")), + auth_token: Arc::new(token.clone()), + }; + let connector_task = + tokio::spawn( + async move { run_connector_with_socks(connector_config, socks_addr).await }, + ); + tokio::time::sleep(std::time::Duration::from_millis(300)).await; + + // Send data through the tunnel + let result = tokio::time::timeout(std::time::Duration::from_secs(5), async { + let mut stream = TcpStream::connect(socks_addr).await.unwrap(); + socks5_greeting(&mut stream).await.unwrap(); + socks5_connect(&mut stream, "127.0.0.1", target_port) + .await + .unwrap(); + socks5_echo(&mut stream, "test-origin").await + }) + .await; + + assert!(result.is_ok(), "Echo should succeed"); + assert_eq!(result.unwrap(), "test-origin"); + + connector_task.abort(); + listener_task.abort(); + target_task.abort(); + } + + // ========================================================================= + // Slow/interleaved stream tests + // ========================================================================= + + /// Test: slow bidirectional streaming — data is forwarded incrementally. + #[tokio::test] + async fn e2e_slow_stream_incremental() { + let dir = generate_test_creds(); + let token = std::fs::read_to_string(dir.path().join("token.txt")).unwrap(); + let listener_port = get_free_port(); + let socks_port = get_free_port(); + let target_port = get_free_port(); + + // Target: slow echo server that responds byte by byte + let target_listener = TcpListener::bind(format!("127.0.0.1:{}", target_port)) + .await + .unwrap(); + let target_task = tokio::spawn(async move { + if let Ok((stream, _)) = target_listener.accept().await { + let (mut rd, mut wr) = stream.into_split(); + let mut buf = [0u8; 1]; + loop { + match rd.read(&mut buf).await { + Ok(0) => break, + Ok(n) => { + if n == 0 { + break; + } + if wr.write_all(&buf[..n]).await.is_err() { + break; + } + // Small delay to simulate slow target + tokio::time::sleep(std::time::Duration::from_millis(10)).await; } - buf.extend_from_slice(&chunk[..n]); + Err(_) => break, } - }) + } + } + }); + + let bind_addr: SocketAddr = format!("127.0.0.1:{}", listener_port).parse().unwrap(); + let listener_config = ListenerConfig { + bind_addr, + server_cert_path: path_arc!(dir.path().join("server.crt")), + server_key_path: path_arc!(dir.path().join("server.key")), + ca_cert_path: path_arc!(dir.path().join("ca.pem")), + auth_token: Arc::new(token.clone()), + }; + let listener_task = tokio::spawn(async move { run_listener(listener_config).await }); + tokio::time::sleep(std::time::Duration::from_millis(200)).await; + + let socks_addr: SocketAddr = format!("127.0.0.1:{}", socks_port).parse().unwrap(); + let connector_config = ConnectorConfig { + target_host: "127.0.0.1".to_string(), + target_port: listener_port, + client_cert_path: path_arc!(dir.path().join("client.crt")), + client_key_path: path_arc!(dir.path().join("client.key")), + ca_cert_path: path_arc!(dir.path().join("ca.pem")), + auth_token: Arc::new(token.clone()), + }; + let connector_task = + tokio::spawn( + async move { run_connector_with_socks(connector_config, socks_addr).await }, + ); + tokio::time::sleep(std::time::Duration::from_millis(300)).await; + + // Send data through the slow echo server + let result = tokio::time::timeout(std::time::Duration::from_secs(10), async { + let mut stream = TcpStream::connect(socks_addr).await.unwrap(); + socks5_greeting(&mut stream).await.unwrap(); + socks5_connect(&mut stream, "127.0.0.1", target_port) .await .unwrap(); - String::from_utf8_lossy(&buf).to_string() - })); - } + // Send data incrementally + use tokio::io::AsyncWriteExt as _; + for chunk in b"SLOW-STREAM-TEST-DATA".chunks(2) { + stream.write_all(chunk).await.unwrap(); + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } - // Collect and verify results - let results: Vec = futures::future::join_all(handles) - .await - .into_iter() - .map(|r| r.unwrap()) - .collect(); + // Read response incrementally + use tokio::io::AsyncReadExt as _; + let mut full_response = Vec::new(); + let mut buf = [0u8; 2]; + loop { + let n = tokio::time::timeout( + std::time::Duration::from_millis(500), + stream.read(&mut buf), + ) + .await; + match n { + Ok(Ok(0)) => break, + Ok(Ok(n)) => full_response.extend_from_slice(&buf[..n]), + Ok(Err(_)) | Err(_) => break, + } + } + String::from_utf8_lossy(&full_response).to_string() + }) + .await; - for (i, result) in results.iter().enumerate() { - assert!( - result.contains(&payloads[i][4..payloads[i].len() - 4]), - "Stream {} response should contain payload: got '{}'", - i, - result - ); - } + assert!( + result.is_ok(), + "Slow stream echo should complete within timeout" + ); + let response = result.unwrap(); + assert_eq!( + response, "SLOW-STREAM-TEST-DATA", + "Incremental echo should preserve data: got '{}'", + response + ); - // Shutdown SOCKS5 proxy - socks_shutdown.notify_one(); - let _ = socks_task.await; + connector_task.abort(); listener_task.abort(); - echo_handle.abort(); + target_task.abort(); + } + + // ========================================================================= + // Stream close/error frame propagation tests + // ========================================================================= + + /// Test: close frames propagate deterministically. + #[tokio::test] + async fn e2e_close_frame_propagation() { + let dir = generate_test_creds(); + let token = std::fs::read_to_string(dir.path().join("token.txt")).unwrap(); + let listener_port = get_free_port(); + let socks_port = get_free_port(); + let target_port = get_free_port(); + + // Target: responds immediately then closes + let target_listener = TcpListener::bind(format!("127.0.0.1:{}", target_port)) + .await + .unwrap(); + let target_task = tokio::spawn(async move { + if let Ok((stream, _)) = target_listener.accept().await { + let (mut rd, mut wr) = stream.into_split(); + let mut buf = [0u8; 4096]; + if let Ok(n) = rd.read(&mut buf).await + && n > 0 + { + let _ = wr.write_all(&buf[..n]).await; + } + // Close immediately after echo + let _ = wr.shutdown().await; + } + }); + + let bind_addr: SocketAddr = format!("127.0.0.1:{}", listener_port).parse().unwrap(); + let listener_config = ListenerConfig { + bind_addr, + server_cert_path: path_arc!(dir.path().join("server.crt")), + server_key_path: path_arc!(dir.path().join("server.key")), + ca_cert_path: path_arc!(dir.path().join("ca.pem")), + auth_token: Arc::new(token.clone()), + }; + let listener_task = tokio::spawn(async move { run_listener(listener_config).await }); + tokio::time::sleep(std::time::Duration::from_millis(200)).await; + + let socks_addr: SocketAddr = format!("127.0.0.1:{}", socks_port).parse().unwrap(); + let connector_config = ConnectorConfig { + target_host: "127.0.0.1".to_string(), + target_port: listener_port, + client_cert_path: path_arc!(dir.path().join("client.crt")), + client_key_path: path_arc!(dir.path().join("client.key")), + ca_cert_path: path_arc!(dir.path().join("ca.pem")), + auth_token: Arc::new(token.clone()), + }; + let connector_task = + tokio::spawn( + async move { run_connector_with_socks(connector_config, socks_addr).await }, + ); + tokio::time::sleep(std::time::Duration::from_millis(300)).await; + + // Send data, get echo, then verify close propagation + let result = tokio::time::timeout(std::time::Duration::from_secs(5), async { + let mut stream = TcpStream::connect(socks_addr).await.unwrap(); + socks5_greeting(&mut stream).await.unwrap(); + socks5_connect(&mut stream, "127.0.0.1", target_port) + .await + .unwrap(); + + use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; + stream.write_all(b"CLOSE-TEST").await.unwrap(); + + // Read echo + let mut buf = [0u8; 4096]; + let n = stream.read(&mut buf).await.unwrap(); + let echo = String::from_utf8_lossy(&buf[..n]).to_string(); + + // After target closes, the next read should return 0 (EOF) + // This verifies CLOSE frame propagation + tokio::time::sleep(std::time::Duration::from_millis(200)).await; + let n2 = stream.read(&mut buf).await.unwrap(); + (echo, n2) + }) + .await; + + assert!(result.is_ok(), "Close test should complete"); + let (echo, remaining) = result.unwrap(); + assert_eq!(echo, "CLOSE-TEST"); + assert_eq!(remaining, 0, "After target close, read should return 0"); + + connector_task.abort(); + listener_task.abort(); + target_task.abort(); + } + + // ========================================================================= + // SOCKS5 helper functions for tests + // ========================================================================= + + async fn socks5_greeting(stream: &mut TcpStream) -> Result<(), std::io::Error> { + use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; + stream.write_all(&[0x05, 0x01, 0x00]).await?; + let mut resp = [0u8; 2]; + stream.read_exact(&mut resp).await?; + assert_eq!(resp[0], 0x05); + assert_eq!(resp[1], 0x00); + Ok(()) + } + + async fn socks5_connect( + stream: &mut TcpStream, + host: &str, + port: u16, + ) -> Result<(), std::io::Error> { + use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; + let octets: Vec = host + .split('.') + .map(|s| s.parse::().unwrap_or(0)) + .collect(); + let mut req = vec![0x05, 0x01, 0x00, 0x01]; + req.extend_from_slice(&octets); + req.extend_from_slice(&port.to_be_bytes()); + stream.write_all(&req).await?; + let mut reply = [0u8; 10]; + stream.read_exact(&mut reply).await?; + if reply[1] != 0x00 { + return Err(std::io::Error::new( + std::io::ErrorKind::ConnectionRefused, + format!("SOCKS5 connect failed with code {}", reply[1]), + )); + } + Ok(()) + } + + async fn socks5_echo(stream: &mut TcpStream, msg: &str) -> String { + use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; + stream.write_all(msg.as_bytes()).await.unwrap(); + let mut buf = vec![0u8; msg.len() + 100]; + let n = tokio::time::timeout(std::time::Duration::from_secs(5), stream.read(&mut buf)) + .await + .unwrap() + .unwrap(); + String::from_utf8_lossy(&buf[..n]).to_string() } }