feat: selective MITM based on SNI

zu1k · Jul 21, 2023 · e0ad482 · e0ad482 · vercel · Jul 21, 2023
1 parent 9f385e3
commit e0ad482
Show file tree

Hide file tree

Showing 6 changed files with 446 additions and 9 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/core/Cargo.toml b/crates/core/Cargo.toml
@@ -10,6 +10,7 @@ license = "MIT"
 [dependencies]
 async-trait = "0.1"
 bytes = { version = "1", features = ["serde"] }
+byteorder = "1.4"
 cfg-if = "1"
 http = "0.2"
 hyper = { version = "0.14", features = ["http1", "http2", "server", "stream", "tcp", "runtime"]  }
@@ -19,6 +20,7 @@ hyper-tls = { version = "0.5", optional = true }
 log = "0.4"
 moka = { version = "0.11", features = ["future"] }
 openssl = { version = "0.10", features = ["vendored"], optional = true }
+pin-project = "1"
 rcgen = { version = "0.10", features = ["x509-parser"] }
 serde = { version = "1.0", features = ["derive"] }
 thiserror = "1"

diff --git a/crates/core/src/handler.rs b/crates/core/src/handler.rs
@@ -45,7 +45,7 @@ impl<D: CustomContextData> MitmFilter<D> {
         }
     }
 
-    pub async fn filter(&self, _ctx: &HttpContext<D>, req: &Request<Body>) -> bool {
+    pub async fn filter_req(&self, _ctx: &HttpContext<D>, req: &Request<Body>) -> bool {
         let host = req.uri().host().unwrap_or_default();
         let list = self.filters.read().unwrap();
         for m in list.iter() {
@@ -55,4 +55,14 @@ impl<D: CustomContextData> MitmFilter<D> {
         }
         false
     }
+
+    pub async fn filter(&self, host: &str) -> bool {
+        let list = self.filters.read().unwrap();
+        for m in list.iter() {
+            if m.matches(host) {
+                return true;
+            }
+        }
+        false
+    }
 }
diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs
@@ -17,6 +17,7 @@ mod error;
 pub mod handler;
 mod http_client;
 pub mod mitm;
+mod sni_reader;
 
 #[derive(TypedBuilder)]
 pub struct Proxy<F, H, D>

diff --git a/crates/core/src/mitm.rs b/crates/core/src/mitm.rs
@@ -2,17 +2,21 @@ use crate::{
     ca::CertificateAuthority,
     handler::{CustomContextData, HttpHandler, MitmFilter},
     http_client::HttpClient,
+    sni_reader::{
+        read_sni_host_name_from_client_hello, HandshakeRecordReader, PrefixedReaderWriter,
+        RecordingBufReader,
+    },
 };
 use http::{header, uri::Scheme, HeaderValue, Uri};
 use hyper::{
-    body::HttpBody, server::conn::Http, service::service_fn, upgrade::Upgraded, Body, Method,
-    Request, Response,
+    body::HttpBody, server::conn::Http, service::service_fn, Body, Method, Request, Response,
 };
 use log::*;
-use std::{marker::PhantomData, sync::Arc};
+use std::{marker::PhantomData, sync::Arc, time::Duration};
 use tokio::{
     io::{AsyncRead, AsyncWrite},
     net::TcpStream,
+    pin,
 };
 use tokio_rustls::TlsAcceptor;
 
@@ -150,7 +154,7 @@ where
             ..Default::default()
         };
 
-        if self.mitm_filter.filter(&ctx, &req).await {
+        if self.mitm_filter.filter_req(&ctx, &req).await {
             tokio::task::spawn(async move {
                 let authority = req
                     .uri()
@@ -175,10 +179,34 @@ where
         Ok(Response::new(Body::empty()))
     }
 
-    pub async fn serve_tls<IO: AsyncRead + AsyncWrite + Unpin + Send + 'static>(self, stream: IO) {
+    pub async fn serve_tls<IO: AsyncRead + AsyncWrite + Unpin + Send + 'static>(
+        self,
+        mut stream: IO,
+    ) {
+        // Read SNI hostname.
+        let mut recording_reader = RecordingBufReader::new(&mut stream);
+        let reader = HandshakeRecordReader::new(&mut recording_reader);
+        pin!(reader);
+        let sni_hostname = tokio::time::timeout(
+            Duration::from_secs(5),
+            read_sni_host_name_from_client_hello(reader),
+        )
+        .await
+        .unwrap()
+        .unwrap();
+
+        let read_buf = recording_reader.buf();
+        let client_stream = PrefixedReaderWriter::new(stream, read_buf);
+
+        if !self.mitm_filter.filter(&sni_hostname).await {
+            let remote_addr = format!("{sni_hostname}:443");
+            tokio::task::spawn(async move { tunnel(client_stream, remote_addr).await });
+            return;
+        }
+
         let server_config = self.ca.clone().gen_server_config();
 
-        match TlsAcceptor::from(server_config).accept(stream).await {
+        match TlsAcceptor::from(server_config).accept(client_stream).await {
             Ok(stream) => {
                 if let Err(e) = Http::new()
                     .http1_preserve_header_case(true)
@@ -239,8 +267,11 @@ fn host_addr(uri: &http::Uri) -> Option<String> {
     uri.authority().map(|auth| auth.to_string())
 }
 
-async fn tunnel(mut upgraded: Upgraded, addr: String) -> std::io::Result<()> {
+async fn tunnel<A>(mut client_stream: A, addr: String) -> std::io::Result<()>
+where
+    A: AsyncRead + AsyncWrite + Unpin,
+{
     let mut server = TcpStream::connect(addr).await?;
-    tokio::io::copy_bidirectional(&mut upgraded, &mut server).await?;
+    tokio::io::copy_bidirectional(&mut client_stream, &mut server).await?;
     Ok(())
 }