jjkiers/layer4-proxy
1
0
Fork 0
Code Issues 3 Pull Requests Packages Projects Releases 5 Activity

Wait for the entire TLS header to become available, even if it takes

Browse Source 
multiple packets.

Closes: #10
This commit is contained in:
Jacob Kiers
2025-01-21 21:22:53 +01:00
parent 4c2711fc81
commit aecffa0d14
1 changed files with 152 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

163
src/servers/protocol/tls.rs
View File

@@ -1,12 +1,15 @@
use crate::servers::Proxy;
use log::{debug, error, trace, warn};
use std::error::Error;
use std::io; // Import io for ErrorKind
use std::sync::Arc;
use std::time::Duration; // For potential delays
use tls_parser::{
parse_tls_extensions, parse_tls_raw_record, parse_tls_record_with_header, TlsMessage,
TlsMessageHandshake,
};
use tokio::net::TcpStream;
use tokio::time::timeout; // Use timeout for peek operations
fn get_sni(buf: &[u8]) -> Vec<String> {
let mut snis: Vec<String> = Vec::new();
@@ -57,6 +60,9 @@ fn get_sni(buf: &[u8]) -> Vec<String> {
snis
}
// Timeout duration for waiting for TLS Hello data
const TLS_PEEK_TIMEOUT: Duration = Duration::from_secs(5); // Adjust as needed
pub(crate) async fn determine_upstream_name(
inbound: &TcpStream,
proxy: &Arc<Proxy>,
@@ -64,35 +70,170 @@ pub(crate) async fn determine_upstream_name(
let default_upstream = proxy.default_action.clone();
let mut header = [0u8; 9];
inbound.peek(&mut header).await?;
let required_bytes = client_hello_buffer_size(&header)?;
// --- Step 1: Peek the initial header (9 bytes) with timeout ---
match timeout(TLS_PEEK_TIMEOUT, async {
loop {
match inbound.peek(&mut header).await {
Ok(n) if n >= header.len() => return Ok::<usize, io::Error>(n), // Got enough bytes
Ok(0) => {
// Connection closed cleanly before sending enough data
trace!("Connection closed while peeking for TLS header");
return Err(io::Error::new(
io::ErrorKind::UnexpectedEof,
"Connection closed while peeking for TLS header",
)
.into()); // Convert to Box<dyn Error>
}
Ok(_) => {
// Not enough bytes yet, yield and loop again
tokio::task::yield_now().await;
}
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
// Should not happen with await, but yield defensively
tokio::task::yield_now().await;
}
Err(e) => {
// Other I/O error
warn!("Error peeking for TLS header: {}", e);
return Err(e.into()); // Convert to Box<dyn Error>
}
}
}
})
.await
{
Ok(Ok(_)) => { /* Header peeked successfully */ }
Ok(Err(e)) => {
// Inner loop returned an error (e.g., EOF, IO error)
trace!("Failed to peek header (inner error): {}", e);
return Ok(default_upstream); // Fallback on error/EOF
}
Err(_) => {
// Timeout occurred
error!("Timeout waiting for TLS header");
return Ok(default_upstream); // Fallback on timeout
}
}
let mut hello_buf = vec![0; required_bytes];
let read_bytes = inbound.peek(&mut hello_buf).await?;
// --- Step 2: Calculate required size ---
let required_bytes = match client_hello_buffer_size(&header) {
Ok(size) => size,
Err(e) => {
// Header was invalid or not a ClientHello
trace!("Could not determine required buffer size: {}", e);
return Ok(default_upstream);
}
};
if read_bytes < required_bytes.into() {
error!("Could not read enough bytes to determine SNI");
// Basic sanity check on size
if required_bytes > 16384 + 9 {
// TLS max record size + header approx
error!(
"Calculated required TLS buffer size is too large: {}",
required_bytes
);
return Ok(default_upstream);
}
// --- Step 3: Peek the full ClientHello with timeout ---
let mut hello_buf = vec![0; required_bytes];
match timeout(TLS_PEEK_TIMEOUT, async {
let mut total_peeked = 0;
loop {
// Peek into the portion of the buffer that hasn't been filled yet.
match inbound.peek(&mut hello_buf[total_peeked..]).await {
Ok(0) => {
// Connection closed cleanly before sending full ClientHello
trace!(
"Connection closed while peeking for full ClientHello (peeked {}/{} bytes)",
total_peeked,
required_bytes
);
return Err::<usize, io::Error>(
io::Error::new(
io::ErrorKind::UnexpectedEof,
"Connection closed while peeking for full ClientHello",
)
.into(),
);
}
Ok(n) => {
total_peeked += n;
if total_peeked >= required_bytes {
trace!("Successfully peeked {} bytes for ClientHello", total_peeked);
return Ok(total_peeked); // Got enough
} else {
// Not enough bytes yet, yield and loop again
trace!(
"Peeked {}/{} bytes for ClientHello, waiting for more...",
total_peeked,
required_bytes
);
tokio::task::yield_now().await;
}
}
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
tokio::task::yield_now().await;
}
Err(e) => {
warn!("Error peeking for full ClientHello: {}", e);
return Err(e.into());
}
}
}
})
.await
{
Ok(Ok(_)) => { /* Full hello peeked successfully */ }
Ok(Err(e)) => {
error!("Could not peek full ClientHello (inner error): {}", e);
return Ok(default_upstream); // Fallback on error/EOF
}
Err(_) => {
error!(
"Timeout waiting for full ClientHello (needed {} bytes)",
required_bytes
);
return Ok(default_upstream); // Fallback on timeout
}
}
// --- Step 4: Parse SNI ---
let snis = get_sni(&hello_buf);
// --- Step 5: Determine upstream based on SNI ---
if snis.is_empty() {
debug!("No SNI found in ClientHello, using default upstream.");
return Ok(default_upstream);
} else {
match proxy.sni.clone() {
Some(sni_map) => {
let mut upstream = default_upstream;
let mut upstream = default_upstream.clone(); // Clone here for default case
let mut found_match = false;
for sni in snis {
let m = sni_map.get(&sni);
if m.is_some() {
upstream = m.unwrap().clone();
// snis is already Vec<String>
if let Some(target_upstream) = sni_map.get(&sni) {
debug!(
"Found matching SNI '{}', routing to upstream: {}",
sni, target_upstream
);
upstream = target_upstream.clone();
found_match = true;
break;
} else {
trace!("SNI '{}' not found in map.", sni);
}
}
if !found_match {
debug!("SNI(s) found but none matched configuration, using default upstream.");
}
Ok(upstream)
}
None => return Ok(default_upstream),
None => {
debug!("SNI found but no SNI map configured, using default upstream.");
Ok(default_upstream)
}
}
}
}