Websocket parser

Cargo.toml

@@ -13,6 +13,9 @@ keywords = ["http", "proxy", "http-proxy"]
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
+[features]
+websocket = ["dep:winnow"]
+
 [dependencies]
 tokio = { version = "1.39.3", features = [
     "macros",
@@ -35,10 +38,15 @@ hyper-util = { version = "0.1.7", features = ["tokio"] }
 native-tls = { version = "0.2.12", features = ["alpn"] }
 thiserror = "1.0.62"
 moka = { version = "0.12.8", features = ["sync"] }
+winnow = { version = "0.6.20", optional = true }
 
 [dev-dependencies]
 axum = { version = "0.7.2", features = ["http2"] }
 clap = { version = "4.5.16", features = ["derive"] }
 rcgen = { version = "0.13.1", features = ["x509-parser"] }
 reqwest = { version = "0.12.7", features = ["native-tls-alpn"] }
 tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
+
+[[example]]
+name = "websocket"
+required-features = ["websocket"]

README.md

@@ -97,32 +97,9 @@ async fn main() {
                 // You can modify request here
                 // or You can just return response anywhere
 
-                let (res, upgrade) = client.send_request(req).await?;
+                let (res, _upgrade) = client.send_request(req).await?;
 
                 println!("{} -> {}", uri, res.status());
-                if let Some(upgrade) = upgrade {
-                    // If the response is an upgrade, e.g. Websocket, you can see traffic.
-                    // Modifying upgraded traffic is not supported yet.
-
-                    // You can try https://echo.websocket.org/.ws to test websocket.
-                    println!("Upgrade connection");
-                    let Upgrade {
-                        mut client_to_server,
-                        mut server_to_client,
-                    } = upgrade;
-                    let url = uri.to_string();
-                    tokio::spawn(async move {
-                        while let Some(data) = client_to_server.next().await {
-                            println!("Client -> Server: {} {:?}", url, data);
-                        }
-                    });
-                    let url = uri.to_string();
-                    tokio::spawn(async move {
-                        while let Some(data) = server_to_client.next().await {
-                            println!("Server -> Client: {} {:?}", url, data);
-                        }
-                    });
-                }
 
                 // You can modify response here
 

examples/proxy.rs

@@ -1,8 +1,7 @@
 use std::path::PathBuf;
 
 use clap::{Args, Parser};
-use futures::StreamExt;
-use http_mitm_proxy::{default_client::Upgrade, DefaultClient, MitmProxy};
+use http_mitm_proxy::{DefaultClient, MitmProxy};
 use moka::sync::Cache;
 use tracing_subscriber::EnvFilter;
 
@@ -84,32 +83,9 @@ async fn main() {
                 // You can modify request here
                 // or You can just return response anywhere
 
-                let (res, upgrade) = client.send_request(req).await?;
+                let (res, _upgrade) = client.send_request(req).await?;
 
                 println!("{} -> {}", uri, res.status());
-                if let Some(upgrade) = upgrade {
-                    // If the response is an upgrade, e.g. Websocket, you can see traffic.
-                    // Modifying upgraded traffic is not supported yet.
-
-                    // You can try https://echo.websocket.org/.ws to test websocket.
-                    println!("Upgrade connection");
-                    let Upgrade {
-                        mut client_to_server,
-                        mut server_to_client,
-                    } = upgrade;
-                    let url = uri.to_string();
-                    tokio::spawn(async move {
-                        while let Some(data) = client_to_server.next().await {
-                            println!("Client -> Server: {} {:?}", url, data);
-                        }
-                    });
-                    let url = uri.to_string();
-                    tokio::spawn(async move {
-                        while let Some(data) = server_to_client.next().await {
-                            println!("Server -> Client: {} {:?}", url, data);
-                        }
-                    });
-                }
 
                 // You can modify response here
 

examples/websocket.rs

@@ -0,0 +1,170 @@
+use std::path::PathBuf;
+
+use clap::{Args, Parser};
+use futures::StreamExt;
+use http_mitm_proxy::{
+    default_client::{websocket, Upgrade},
+    DefaultClient, MitmProxy,
+};
+use moka::sync::Cache;
+use tracing_subscriber::EnvFilter;
+
+#[derive(Parser)]
+struct Opt {
+    #[clap(flatten)]
+    external_cert: Option<ExternalCert>,
+}
+
+#[derive(Args, Debug)]
+struct ExternalCert {
+    #[arg(required = false)]
+    cert: PathBuf,
+    #[arg(required = false)]
+    private_key: PathBuf,
+}
+
+fn make_root_cert() -> rcgen::CertifiedKey {
+    let mut param = rcgen::CertificateParams::default();
+
+    param.distinguished_name = rcgen::DistinguishedName::new();
+    param.distinguished_name.push(
+        rcgen::DnType::CommonName,
+        rcgen::DnValue::Utf8String("<HTTP-MITM-PROXY CA>".to_string()),
+    );
+    param.key_usages = vec![
+        rcgen::KeyUsagePurpose::KeyCertSign,
+        rcgen::KeyUsagePurpose::CrlSign,
+    ];
+    param.is_ca = rcgen::IsCa::Ca(rcgen::BasicConstraints::Unconstrained);
+
+    let key_pair = rcgen::KeyPair::generate().unwrap();
+    let cert = param.self_signed(&key_pair).unwrap();
+
+    rcgen::CertifiedKey { cert, key_pair }
+}
+
+#[tokio::main]
+async fn main() {
+    let opt = Opt::parse();
+
+    tracing_subscriber::fmt()
+        .with_env_filter(EnvFilter::from_default_env())
+        .init();
+
+    let root_cert = if let Some(external_cert) = opt.external_cert {
+        // Use existing key
+        let param = rcgen::CertificateParams::from_ca_cert_pem(
+            &std::fs::read_to_string(&external_cert.cert).unwrap(),
+        )
+        .unwrap();
+        let key_pair =
+            rcgen::KeyPair::from_pem(&std::fs::read_to_string(&external_cert.private_key).unwrap())
+                .unwrap();
+
+        let cert = param.self_signed(&key_pair).unwrap();
+
+        rcgen::CertifiedKey { cert, key_pair }
+    } else {
+        make_root_cert()
+    };
+
+    let root_cert_pem = root_cert.cert.pem();
+    let root_cert_key = root_cert.key_pair.serialize_pem();
+
+    let proxy = MitmProxy::new(
+        // This is the root cert that will be used to sign the fake certificates
+        Some(root_cert),
+        Some(Cache::new(128)),
+    );
+
+    let client = DefaultClient::new().unwrap();
+    let server = proxy
+        .bind(("127.0.0.1", 3003), move |_client_addr, req| {
+            let client = client.clone();
+            async move {
+                let uri = req.uri().clone();
+
+                // You can modify request here
+                // or You can just return response anywhere
+
+                let (res, upgrade) = client.send_request(req).await?;
+
+                // println!("{} -> {}", uri, res.status());
+                if let Some(upgrade) = upgrade {
+                    // If the response is an upgrade, e.g. Websocket, you can see traffic.
+                    // Modifying upgraded traffic is not supported yet.
+
+                    // You can try https://echo.websocket.org/.ws to test websocket.
+                    println!("Upgrade connection");
+                    let Upgrade {
+                        mut client_to_server,
+                        mut server_to_client,
+                    } = upgrade;
+                    let url = uri.to_string();
+                    tokio::spawn(async move {
+                        let mut buf = Vec::new();
+                        while let Some(data) = client_to_server.next().await {
+                            buf.extend(data);
+                            loop {
+                                let input = &mut buf.as_slice();
+                                if let Ok(frame) = websocket::frame(input) {
+                                    println!(
+                                        "Client -> Server: {} {:?}",
+                                        url,
+                                        String::from_utf8_lossy(&frame.payload_data)
+                                    );
+                                    buf = input.to_vec();
+                                } else {
+                                    break;
+                                }
+                            }
+                        }
+                    });
+                    let url = uri.to_string();
+                    tokio::spawn(async move {
+                        let mut buf = Vec::new();
+                        while let Some(data) = server_to_client.next().await {
+                            buf.extend(data);
+                            loop {
+                                let input = &mut buf.as_slice();
+                                if let Ok(frame) = websocket::frame(input) {
+                                    println!(
+                                        "Server -> Client: {} {:?}",
+                                        url,
+                                        String::from_utf8_lossy(&frame.payload_data)
+                                    );
+                                    buf = input.to_vec();
+                                } else {
+                                    break;
+                                }
+                            }
+                        }
+                    });
+                }
+
+                // You can modify response here
+
+                Ok::<_, http_mitm_proxy::default_client::Error>(res)
+            }
+        })
+        .await
+        .unwrap();
+
+    println!("HTTP Proxy is listening on http://127.0.0.1:3003");
+
+    println!();
+    println!("Trust this cert if you want to use HTTPS");
+    println!();
+    println!("{}", root_cert_pem);
+    println!();
+
+    /*
+        Save this cert to ca.crt and use it with curl like this:
+        curl https://www.google.com -x http://127.0.0.1:3003 --cacert ca.crt
+    */
+
+    println!("Private key");
+    println!("{}", root_cert_key);
+
+    server.await;
+}

src/default_client.rs

@@ -272,3 +272,85 @@ fn remove_authority<B>(req: &mut Request<B>) {
     parts.authority = None;
     *req.uri_mut() = Uri::from_parts(parts).unwrap();
 }
+
+#[cfg(feature = "websocket")]
+pub mod websocket {
+    /*
+    https://developer.mozilla.org/en-US/docs/Web/API/WebSockets_API/Writing_WebSocket_servers
+    Frame format:
+
+          0                   1                   2                   3
+          0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
+         +-+-+-+-+-------+-+-------------+-------------------------------+
+         |F|R|R|R| opcode|M| Payload len |    Extended payload length    |
+         |I|S|S|S|  (4)  |A|     (7)     |             (16/64)           |
+         |N|V|V|V|       |S|             |   (if payload len==126/127)   |
+         | |1|2|3|       |K|             |                               |
+         +-+-+-+-+-------+-+-------------+ - - - - - - - - - - - - - - - +
+         |     Extended payload length continued, if payload len == 127  |
+         + - - - - - - - - - - - - - - - +-------------------------------+
+         |                               |Masking-key, if MASK set to 1  |
+         +-------------------------------+-------------------------------+
+         | Masking-key (continued)       |          Payload Data         |
+         +-------------------------------- - - - - - - - - - - - - - - - +
+         :                     Payload Data continued ...                :
+         + - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - +
+         |                     Payload Data continued ...                |
+         +---------------------------------------------------------------+
+
+    */
+
+    use winnow::{
+        binary::{be_u16, be_u64, u8},
+        prelude::*,
+        token::take,
+    };
+
+    pub struct Frame {
+        pub b0: u8,
+        pub b1: u8,
+        pub payload_len: usize,
+        pub masking_key: Option<[u8; 4]>,
+        pub payload_data: Vec<u8>,
+    }
+
+    pub fn frame(input: &mut &[u8]) -> PResult<Frame> {
+        let b0 = u8(input)?;
+        let b1 = u8(input)?;
+
+        let payload_len = match b1 & 0b0111_1111 {
+            126 => {
+                let len = be_u16(input)?;
+                len as usize
+            }
+            127 => {
+                let len = be_u64(input)?;
+                len as usize
+            }
+            _ => (b1 & 0b0111_1111) as usize,
+        };
+
+        let mask = b1 & 0b1000_0000 != 0;
+        let masking_key = if mask {
+            Some([u8(input)?, u8(input)?, u8(input)?, u8(input)?])
+        } else {
+            None
+        };
+
+        let mut payload_data = take(payload_len).parse_next(input)?.to_vec();
+
+        if let Some(mask) = masking_key {
+            for (i, byte) in payload_data.iter_mut().enumerate() {
+                *byte ^= mask[i % 4];
+            }
+        }
+
+        Ok(Frame {
+            b0,
+            b1,
+            payload_len,
+            masking_key,
+            payload_data,
+        })
+    }
+}