diff --git a/Cargo.lock b/Cargo.lock index 5f8c4b9..f5c4ca0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -746,6 +746,7 @@ dependencies = [ "tokio-rustls", "tracing", "tracing-subscriber", + "winnow", ] [[package]] @@ -2292,6 +2293,15 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "winnow" +version = "0.6.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "36c1fec1a2bb5866f07c25f68c26e565c4c200aebb96d7e55710c19d3e8ac49b" +dependencies = [ + "memchr", +] + [[package]] name = "x509-parser" version = "0.16.0" diff --git a/Cargo.toml b/Cargo.toml index 0e0b610..7bfd830 100644 --- a/Cargo.toml +++ b/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,6 +38,7 @@ 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"] } @@ -42,3 +46,7 @@ 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"] diff --git a/README.md b/README.md index ede1a0f..93fdbf5 100644 --- a/README.md +++ b/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 diff --git a/examples/proxy.rs b/examples/proxy.rs index ca9fc01..015e305 100644 --- a/examples/proxy.rs +++ b/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 diff --git a/examples/websocket.rs b/examples/websocket.rs new file mode 100644 index 0000000..e488b03 --- /dev/null +++ b/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, +} + +#[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("".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; +} diff --git a/src/default_client.rs b/src/default_client.rs index 8f90878..3fdb59f 100644 --- a/src/default_client.rs +++ b/src/default_client.rs @@ -272,3 +272,85 @@ fn remove_authority(req: &mut Request) { 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, + } + + pub fn frame(input: &mut &[u8]) -> PResult { + 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, + }) + } +} diff --git a/src/lib.rs b/src/lib.rs index 7d4582c..238706d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -206,17 +206,15 @@ impl + Send + Sync + 'static> MitmProxy { } fn get_certified_key(&self, host: String) -> Option { - if let Some(root_cert) = self.root_cert.as_ref() { - Some(if let Some(cache) = self.cert_cache.as_ref() { + self.root_cert.as_ref().map(|root_cert| { + if let Some(cache) = self.cert_cache.as_ref() { cache.get_with(host.clone(), move || { generate_cert(host, root_cert.borrow()) }) } else { generate_cert(host, root_cert.borrow()) - }) - } else { - None - } + } + }) } fn server_config(