Update based on lints

Signed-off-by: Jacob Kiers <code@kiers.eu>
Fix SNI header parsing
2025-01-09 20:38:05 +01:00 · 2025-01-09 20:23:02 +01:00
12 changed files with 56 additions and 437 deletions
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,15 +7,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0

 ## [Unreleased]

-## [0.1.10] - 2025-01-09
-
-### Fixed
-
-* The ClientHello TLS header is now read in full before it is parsed, solving
-  an error where there was not enough data to fully read it. In those cases
-  it was not possible to determine the upstream address and therefore the proxy
-  would go the the default action instead.
-
 ### Changed

 * Updated some dependencies to prevent the build from breaking.
@@ -41,10 +32,10 @@ The ability to run `l4p` without arguments is now deprecated. Please use
 ## Previous versions

 [unreleased]: https://code.kiers.eu/jjkiers/layer4-proxy/compare/v0.1.9...HEAD
-[0.1.10]: https://code.kiers.eu/jjkiers/layer4-proxy/compare/v0.1.9...v0.1.10
 [0.1.9]: https://code.kiers.eu/jjkiers/layer4-proxy/compare/v0.1.8...v0.1.9


+
 Types of changes:

 * `Added` for new features.
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1,6 +1,6 @@
 # This file is automatically @generated by Cargo.
 # It is not intended for manual editing.
-version = 4
+version = 3

 [[package]]
 name = "addr2line"
@@ -26,12 +26,6 @@ dependencies = [
 "memchr",
 ]

-[[package]]
-name = "anyhow"
-version = "1.0.102"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c"
-
 [[package]]
 name = "async-trait"
 version = "0.1.77"
@@ -55,28 +49,6 @@ version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa"

-[[package]]
-name = "aws-lc-rs"
-version = "1.16.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d9a7b350e3bb1767102698302bc37256cbd48422809984b98d292c40e2579aa9"
-dependencies = [
- "aws-lc-sys",
- "zeroize",
-]
-
-[[package]]
-name = "aws-lc-sys"
-version = "0.37.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b092fe214090261288111db7a2b2c2118e5a7f30dc2569f1732c4069a6840549"
-dependencies = [
- "cc",
- "cmake",
- "dunce",
- "fs_extra",
-]
-
 [[package]]
 name = "backtrace"
 version = "0.3.69"
@@ -130,15 +102,9 @@ checksum = "a2bd12c1caf447e69cd4528f47f94d203fd2582878ecb9e9465484c4148a8223"

 [[package]]
 name = "cc"
-version = "1.2.56"
+version = "1.0.86"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "aebf35691d1bfb0ac386a69bac2fde4dd276fb618cf8bf4f5318fe285e821bb2"
-dependencies = [
- "find-msvc-tools",
- "jobserver",
- "libc",
- "shlex",
-]
+checksum = "7f9fa1897e4325be0d68d48df6aa1a71ac2ed4d27723887e7754192705350730"

 [[package]]
 name = "cfg-if"
@@ -146,15 +112,6 @@ version = "1.0.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"

-[[package]]
-name = "cmake"
-version = "0.1.57"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "75443c44cd6b379beb8c5b45d85d0773baf31cce901fe7bb252f4eff3008ef7d"
-dependencies = [
- "cc",
-]
-
 [[package]]
 name = "console"
 version = "0.15.8"
@@ -177,12 +134,6 @@ dependencies = [
 "powerfmt",
 ]

-[[package]]
-name = "dunce"
-version = "1.0.5"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813"
-
 [[package]]
 name = "encode_unicode"
 version = "0.3.6"
@@ -233,12 +184,6 @@ version = "2.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9fc0510504f03c51ada170672ac806f1f105a88aa97a5281117e1ddc3368e51a"

-[[package]]
-name = "find-msvc-tools"
-version = "0.1.9"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582"
-
 [[package]]
 name = "fnv"
 version = "1.0.7"
@@ -254,12 +199,6 @@ dependencies = [
 "percent-encoding",
 ]

-[[package]]
-name = "fs_extra"
-version = "1.3.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c"
-
 [[package]]
 name = "futures"
 version = "0.3.30"
@@ -566,15 +505,6 @@ version = "1.0.10"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b1a46d1a171d865aa5f83f92695765caa047a9b4cbae2cbf37dbd613a793fd4c"

-[[package]]
-name = "jobserver"
-version = "0.1.32"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "48d1dbcbbeb6a7fec7e059840aa538bd62aaccf972c7346c4d9d2059312853d0"
-dependencies = [
- "libc",
-]
-
 [[package]]
 name = "js-sys"
 version = "0.3.69"
@@ -586,9 +516,8 @@ dependencies = [

 [[package]]
 name = "l4p"
-version = "0.1.11"
+version = "0.1.9"
 dependencies = [
- "anyhow",
 "async-trait",
 "byte_string",
 "bytes",
@@ -596,15 +525,12 @@ dependencies = [
 "log",
 "pico-args",
 "pretty_env_logger",
- "rustls",
- "rustls-pemfile",
 "self_update",
 "serde",
 "serde_yaml",
 "time",
 "tls-parser",
 "tokio",
- "tokio-rustls",
 "url",
 ]

@@ -1158,12 +1084,10 @@ dependencies = [

 [[package]]
 name = "rustls"
-version = "0.23.36"
+version = "0.23.10"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c665f33d38cea657d9614f766881e4d510e0eda4239891eea56b4cadcf01801b"
+checksum = "05cff451f60db80f490f3c182b77c35260baace73209e9cdbbe526bfe3a4d402"
 dependencies = [
- "aws-lc-rs",
- "log",
 "once_cell",
 "ring",
 "rustls-pki-types",
@@ -1174,29 +1098,26 @@ dependencies = [

 [[package]]
 name = "rustls-pemfile"
-version = "2.2.0"
+version = "2.1.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "dce314e5fee3f39953d46bb63bb8a46d40c2f8fb7cc5a3b6cab2bde9721d6e50"
+checksum = "29993a25686778eb88d4189742cd713c9bce943bc54251a33509dc63cbacf73d"
 dependencies = [
+ "base64",
 "rustls-pki-types",
 ]

 [[package]]
 name = "rustls-pki-types"
-version = "1.14.0"
+version = "1.7.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "be040f8b0a225e40375822a563fa9524378b9d63112f53e19ffff34df5d33fdd"
-dependencies = [
- "zeroize",
-]
+checksum = "976295e77ce332211c0d24d92c0e83e50f5c5f046d11082cea19f3df13a3562d"

 [[package]]
 name = "rustls-webpki"
-version = "0.103.9"
+version = "0.102.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d7df23109aa6c1567d1c575b9952556388da57401e4ace1d15f79eedad0d8f53"
+checksum = "ff448f7e92e913c4b7d4c6d8e4540a1724b319b4152b8aef6d4cf8339712b33e"
 dependencies = [
- "aws-lc-rs",
 "ring",
 "rustls-pki-types",
 "untrusted",
@@ -1312,12 +1233,6 @@ dependencies = [
 "unsafe-libyaml",
 ]

-[[package]]
-name = "shlex"
-version = "1.3.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64"
-
 [[package]]
 name = "signal-hook-registry"
 version = "1.4.1"
@@ -1533,11 +1448,12 @@ dependencies = [

 [[package]]
 name = "tokio-rustls"
-version = "0.26.4"
+version = "0.26.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61"
+checksum = "0c7bc40d0e5a97695bb96e27995cd3a08538541b0a846f65bba7a359f36700d4"
 dependencies = [
 "rustls",
+ "rustls-pki-types",
 "tokio",
 ]

--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "l4p"
-version = "0.1.11"
+version = "0.1.9"
 edition = "2021"
 authors = ["Jacob Kiers <code@kiers.eu>"]
 license = "Apache-2.0"
@@ -20,7 +20,6 @@ name = "l4p"
 path = "src/main.rs"

 [dependencies]
-anyhow = "1.0.102"
 async-trait = "0.1.73"
 byte_string = "1"
 bytes = "1.1"
@@ -28,14 +27,11 @@ futures = "0.3"
 log = "0.4"
 pico-args = "0.5.0"
 pretty_env_logger = "0.5"
-rustls = "0.23"
-rustls-pemfile = "2.2.0"
 serde = { version = "1.0", features = ["derive"] }
 serde_yaml = "0.9.21"
 time = { version = "0.3.37", features = ["local-offset", "formatting"] }
 tls-parser = "0.12.2"
 tokio = { version = "1.0", features = ["full"] }
-tokio-rustls = "0.26.4"
 url = "2.2.2"

 [dependencies.self_update]
--- a/PLAN.md
+++ b/PLAN.md
@@ -1,30 +0,0 @@
-## Plan for TLS Termination and Dynamic Port Handling
-
-### Task
-Modify the `l4p` (layer 4 proxy) to perform TLS termination and handle dynamic/random ports for backend services. The backend services are on a specific IPv6 address, and the user can dynamically determine the hostname of these services in the format "https://{port}.my-host". The proxy should listen on port 443 and use SNI for routing.
-
-### Completed Actions
-
- Added `tokio-rustls`, `rustls-pemfile`, `anyhow` dependencies to `Cargo.toml`.
- Modified `src/config/config_v1.rs` to include `TlsTerminationConfig`, `CertificateConfig`, `SniCertificateConfig` structs and updated `ServerConfig`.
- Created `src/tls.rs` (multiple iterations due to compilation issues).
- Integrated `anyhow::Result` into various functions and imported `Context` to `src/servers/mod.rs`.
- Corrected imports in `src/main.rs`, `src/servers/mod.rs`, `src/servers/protocol/tcp.rs`.
- Removed `mod tls;` from `src/servers/protocol/mod.rs`.
- Attempted to fix various compilation errors related to `rustls` API changes, lifetime issues, and `tokio` task handling.
- Changed `handle.await??;` to explicit match for debugging purposes.
-
-### Current State (with persistent errors)
-
-The code currently has compilation errors, primarily related to:
-
-1.  **`src/servers/mod.rs`**: Still showing an error for `map_err` not found for unit type `()`. This arises from the complex double `Result` handling (`Result<anyhow::Result<()>, JoinError>`) when awaiting spawned tasks.
-2.  **`src/tls.rs`**: Facing issues with `rustls::pki_types::PrivateKeyDer` and `CertificateDer` conversions, specifically for ensuring `'static` lifetimes and incorrect method usages like `to_vec()` or `as_ref()`, or `into_owned()` methods not existing for certain types. The `borrowed data escapes outside of function` error indicates deeper lifetime mismatches.
-
-### Next Steps (Requires Manual Intervention)
-
- **Refactor `src/servers/mod.rs` error handling**: The current `match handle.await` block needs to be carefully reviewed to ensure correct unwraping of the nested `Result` types and proper error propagation from `tokio::task::JoinError` to `anyhow::Error`.
- **Re-evaluate `src/tls.rs` `rustls::pki_types` usage**: A deeper understanding of `rustls-pki-types` crate and its `CertificateDer` and `PrivateKeyDer` lifetimes and conversion methods is needed. The specific error message `no method named to_vec found for struct PrivatePkcs8KeyDer` is a key indicator of incorrect usage.
- **Review `rustls` version and documentation**: It might be helpful to review the `rustls` and `tokio-rustls` documentation for version-specific changes and best practices regarding `pki_types` and asynchronous error handling.
-
-This commit contains the work in progress as of the current session, including these unresolved errors, to allow for external review and debugging.
--- a/src/config/config_v1.rs
+++ b/src/config/config_v1.rs
@@ -35,9 +35,6 @@ pub struct ServerConfig {
    pub tls: Option<bool>,
    pub sni: Option<HashMap<String, String>>,
    pub default: Option<String>,
-    pub termination_certs: Option<TlsTerminationConfig>,
-    pub dynamic_backend_pattern: Option<String>,
-    pub fixed_backend_ipv6: Option<String>,
 }
 impl TryInto<ProxyToUpstream> for &str {
    type Error = ConfigError;
@@ -204,25 +201,6 @@ fn verify_config(config: ParsedConfigV1) -> Result<ParsedConfigV1, ConfigError>
    Ok(config)
 }

-#[derive(Debug, Default, Deserialize, Clone)]
-pub struct TlsTerminationConfig {
-    pub default_certificate: CertificateConfig,
-    pub sni_certificates: Option<Vec<SniCertificateConfig>>,
-}
-
-#[derive(Debug, Default, Deserialize, Clone)]
-pub struct CertificateConfig {
-    pub certificate_path: String,
-    pub private_key_path: String,
-}
-
-#[derive(Debug, Default, Deserialize, Clone)]
-pub struct SniCertificateConfig {
-    pub hostname: String,
-    pub certificate_path: String,
-    pub private_key_path: String,
-}
-
 impl From<IOError> for ConfigError {
    fn from(err: IOError) -> ConfigError {
        ConfigError::IO(err)
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -1,3 +1,3 @@
-pub mod config_v1;
+mod config_v1;
 pub(crate) use config_v1::ConfigV1;
 pub(crate) use config_v1::ParsedConfigV1;
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,7 +2,6 @@ mod config;
 mod servers;
 mod update;
 mod upstreams;
-mod tls; // NEW: Declare the new TLS module

 use crate::config::ConfigV1;
 use crate::servers::Server;
--- a/src/servers/mod.rs
+++ b/src/servers/mod.rs
@@ -1,4 +1,3 @@
-use anyhow::{anyhow, Result};
 use log::{error, info};
 use std::collections::{HashMap, HashSet};
 use std::net::SocketAddr;
@@ -9,15 +8,10 @@ use tokio::task::JoinHandle;
 mod protocol;
 pub(crate) mod upstream_address;

-use crate::config::{ParsedConfigV1, config_v1::TlsTerminationConfig};
-        // use crate::tls;
+use crate::config::ParsedConfigV1;
 use crate::upstreams::Upstream;
 use protocol::tcp;

-// A helper to convert Box<dyn Error> to anyhow::Error
-fn unhandled_error_for_box_error(e: Box<dyn std::error::Error>) -> anyhow::Error {
-    anyhow!("{}", e)
-}
 #[derive(Debug)]
 pub(crate) struct Server {
    pub proxies: Vec<Arc<Proxy>>,
@@ -32,9 +26,6 @@ pub(crate) struct Proxy {
    pub sni: Option<HashMap<String, String>>,
    pub default_action: String,
    pub upstream: HashMap<String, Upstream>,
-    pub termination_certs: Option<TlsTerminationConfig>,
-    pub dynamic_backend_pattern: Option<String>,
-    pub fixed_backend_ipv6: Option<String>,
 }

 impl Server {
@@ -73,9 +64,6 @@ impl Server {
                    sni: sni.clone(),
                    default_action: default.clone(),
                    upstream: upstream.clone(),
-                    termination_certs: proxy.termination_certs.clone(),
-                    dynamic_backend_pattern: proxy.dynamic_backend_pattern.clone(),
-                    fixed_backend_ipv6: proxy.fixed_backend_ipv6.clone(),
                };
                new_server.proxies.push(Arc::new(proxy));
            }
@@ -85,7 +73,7 @@ impl Server {
    }

    #[tokio::main]
-    pub async fn run(&mut self) -> Result<()> {
+    pub async fn run(&mut self) -> Result<(), Box<dyn std::error::Error>> {
        let proxies = self.proxies.clone();
        let mut handles: Vec<JoinHandle<()>> = Vec::new();

@@ -95,14 +83,6 @@ impl Server {
                config.protocol, config.name, config.listen
            );
            let handle = tokio::spawn(async move {
-                if config.tls && config.termination_certs.is_some() {
-                    // New TLS termination handling
-                    let res = tcp::tls_proxy(config.clone()).await;
-                    if res.is_err() {
-                        error!("Failed to start TLS server {}: {}", config.name, res.err().unwrap());
-                    }
-                } else {
-                    // Existing plain TCP handling
                match config.protocol.as_ref() {
                    "tcp" | "tcp4" | "tcp6" => {
                        let res = tcp::proxy(config.clone()).await;
@@ -114,13 +94,13 @@ impl Server {
                        error!("Invalid protocol: {}", config.protocol)
                    }
                }
-                }
            });
            handles.push(handle);
        }

        for handle in handles {
-            handle.await.map_err(anyhow::Error::from)?.map_err(anyhow::Error::from)?;        }
+            handle.await?;
+        }
        Ok(())
    }
 }
--- a/src/servers/protocol/mod.rs
+++ b/src/servers/protocol/mod.rs
@@ -1 +1,2 @@
 pub mod tcp;
+pub mod tls;
--- a/src/servers/protocol/tcp.rs
+++ b/src/servers/protocol/tcp.rs
@@ -1,48 +1,11 @@
-use anyhow::Result;
+use crate::servers::protocol::tls::determine_upstream_name;
 use crate::servers::Proxy;
 use log::{debug, error, info, warn};
+use std::error::Error;
 use std::sync::Arc;
 use tokio::net::{TcpListener, TcpStream};
-use tokio_rustls::TlsAcceptor;
-use crate::tls;

-pub(crate) async fn tls_proxy(config: Arc<Proxy>) -> Result<()> {
-    let listener = TcpListener::bind(config.listen).await?;
-    let config = config.clone();
-
-    let acceptor = tls::build_tls_acceptor(
-        config.termination_certs.as_ref().expect("TLS termination config missing"),
-    )?;
-
-    loop {
-        let _config = config.clone();
-        let thread_acceptor = acceptor.clone();
-        match listener.accept().await {
-            Err(err) => {
-                error!("Failed to accept TLS connection: {}", err);
-            }
-            Ok((stream, _)) => {
-                tokio::spawn(async move {
-                    let res = match thread_acceptor.accept(stream).await {
-                        Ok(tls_stream) => {
-                            info!("TLS handshake successful with {:?}", tls_stream.into_inner().0.peer_addr().ok().map(|s| s.to_string()).unwrap_or_else(|| "unknown".to_string()));
-                            Ok(()) // Return Ok(()) for now
-                        }
-                        Err(err) => {
-                            error!("TLS handshake failed: {}", err);
-                            Err(anyhow::anyhow!("{}", err))
-                        }
-                    };
-                    if res.is_err() {
-                        error!("TLS handling error: {}", res.unwrap_err());
-                    }
-                });
-            }
-        }
-    }
-}
-
-pub(crate) async fn proxy(config: Arc<Proxy>) -> Result<()> {
+pub(crate) async fn proxy(config: Arc<Proxy>) -> Result<(), Box<dyn Error>> {
    let listener = TcpListener::bind(config.listen).await?;
    let config = config.clone();

@@ -51,7 +14,7 @@ pub(crate) async fn proxy(config: Arc<Proxy>) -> Result<()> {
        match listener.accept().await {
            Err(err) => {
                error!("Failed to accept connection: {}", err);
-                return Err(anyhow::Error::new(err)); // Convert to anyhow::Error
+                return Err(Box::new(err));
            }
            Ok((stream, _)) => {
                tokio::spawn(async move {
@@ -67,10 +30,13 @@ pub(crate) async fn proxy(config: Arc<Proxy>) -> Result<()> {
    }
 }

-async fn accept(inbound: TcpStream, proxy: Arc<Proxy>) -> Result<()> {
+async fn accept(inbound: TcpStream, proxy: Arc<Proxy>) -> Result<(), Box<dyn Error>> {
    info!("New connection from {:?}", inbound.peer_addr()?);

-    let upstream_name = proxy.default_action.clone();
+    let upstream_name = match proxy.tls {
+        false => proxy.default_action.clone(),
+        true => determine_upstream_name(&inbound, &proxy).await?,
+    };

    debug!("Upstream: {}", upstream_name);

@@ -81,9 +47,9 @@ async fn accept(inbound: TcpStream, proxy: Arc<Proxy>) -> Result<()> {
                "No upstream named {:?} on server {:?}",
                proxy.default_action, proxy.name
            );
-            proxy.upstream.get(&proxy.default_action).expect("Default upstream must exist")
+            proxy.upstream.get(&proxy.default_action).unwrap()
        }
    };

-    upstream.process(inbound).await.map_err(|e| anyhow::anyhow!("{}", e))
+    upstream.process(inbound).await
 }
--- a/src/servers/protocol/tls.rs
+++ b/src/servers/protocol/tls.rs
@@ -1,15 +1,12 @@
 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();
@@ -60,9 +57,6 @@ 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>,
@@ -70,170 +64,35 @@ pub(crate) async fn determine_upstream_name(
    let default_upstream = proxy.default_action.clone();

    let mut header = [0u8; 9];
+    inbound.peek(&mut header).await?;

-    // --- 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 required_bytes = client_hello_buffer_size(&header)?;

-    // --- 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);
-        }
-    };
-
-    // 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
-        }
+    let read_bytes = inbound.peek(&mut hello_buf).await?;
+
+    if read_bytes < required_bytes.into() {
+        error!("Could not read enough bytes to determine SNI");
+        return Ok(default_upstream);
    }

-    // --- 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.clone(); // Clone here for default case
-                let mut found_match = false;
+                let mut upstream = default_upstream;
                for sni in snis {
-                    // 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;
+                    let m = sni_map.get(&sni);
+                    if m.is_some() {
+                        upstream = m.unwrap().clone();
                        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 => {
-                debug!("SNI found but no SNI map configured, using default upstream.");
-                Ok(default_upstream)
-            }
+            None => return Ok(default_upstream),
        }
    }
 }
--- a/src/tls.rs
+++ b/src/tls.rs
@@ -1,37 +0,0 @@
-use anyhow::{anyhow, Result};
-use rustls::pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer};
-use rustls::ServerConfig;
-use rustls_pemfile::{certs, pkcs8_private_keys};
-use std::fs::File;
-use std::io::BufReader;
-use std::sync::Arc;
-
-pub fn load_certs(path: &str) -> Result<Vec<CertificateDer>> {
-    let mut reader = BufReader::new(File::open(path)?);
-    certs(&mut reader)
-        .collect::<Result<Vec<CertificateDer>, std::io::Error>>()
-        .map_err(|e| anyhow!("failed to load certificates from {}: {}", path, e))
-}
-
-pub fn load_private_key(path: &str) -> Result<PrivateKeyDer> {
-    let mut reader = BufReader::new(File::open(path)?);
-    let mut keys = pkcs8_private_keys(&mut reader)
-        .collect::<Result<Vec<PrivatePkcs8KeyDer>, std::io::Error>>()
-        .map_err(|e| anyhow!("failed to load private keys from {}: {}", path, e))?;
-    keys.pop()
-        .map(|k| PrivateKeyDer::Pkcs8(k.to_vec().into()))
-        .ok_or_else(|| anyhow!("no private keys found for {}", path))
-}
-
-pub fn build_tls_acceptor(
-    config: &crate::config::config_v1::TlsTerminationConfig,
-) -> Result<tokio_rustls::TlsAcceptor> {
-    let certs = load_certs(&config.default_certificate.certificate_path)?;
-    let key = load_private_key(&config.default_certificate.private_key_path)?;
-
-    let tls_config = ServerConfig::builder()
-        .with_no_client_auth()
-        .with_single_cert(certs, key)?;
-
-    Ok(tokio_rustls::TlsAcceptor::from(Arc::new(tls_config)))
-}