diff --git a/Cargo.lock b/Cargo.lock index 655e6c1..e0d720b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11,6 +11,56 @@ dependencies = [ "memchr", ] +[[package]] +name = "anstream" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + +[[package]] +name = "anstyle-parse" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys 0.61.2", +] + [[package]] name = "anyhow" version = "1.0.102" @@ -177,6 +227,46 @@ dependencies = [ "rand_core 0.10.1", ] +[[package]] +name = "clap" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2ce8604710f6733aa641a2b3731eaa1e8b3d9973d5e3565da11800813f997a9" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + [[package]] name = "cmake" version = "0.1.58" @@ -192,6 +282,12 @@ version = "0.5.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" +[[package]] +name = "colorchoice" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" + [[package]] name = "combine" version = "4.6.7" @@ -758,6 +854,12 @@ version = "2.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" +[[package]] +name = "is_terminal_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" + [[package]] name = "itoa" version = "1.0.18" @@ -956,6 +1058,12 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "once_cell_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" + [[package]] name = "openssl" version = "0.10.80" @@ -1026,6 +1134,7 @@ dependencies = [ name = "peek-client" version = "0.1.2" dependencies = [ + "clap", "futures-util", "peek-proto", "reqwest", @@ -1647,6 +1756,12 @@ version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + [[package]] name = "subtle" version = "2.6.1" @@ -1993,6 +2108,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + [[package]] name = "valuable" version = "0.1.1" diff --git a/README.md b/README.md index 678dde4..c1f84ae 100644 --- a/README.md +++ b/README.md @@ -72,15 +72,26 @@ peek prints a public URL. | Option | Use | | --- | --- | -| `--token` | same value as `PEEK_AUTH_TOKEN` | | `--domain` | hosted peek domain, like `example.com` | +| `--token` | same value as `PEEK_AUTH_TOKEN` | | `--subdomain` | public URL name, like `myapp` | | `--password` | require a password for visitors | +| `--server` | full relay WebSocket URL, like `wss://example.com/tunnel` | `--token` creates the tunnel. `--password` protects the public URL and is optional. Without `--subdomain`, peek creates a random name. Without `--password`, the URL is public. +Environment variables: + +```bash +PEEK_DOMAIN=example.com +PEEK_AUTH_TOKEN=change-me +PEEK_PASSWORD=optional-visitor-password +``` + +`PEEK_SERVER` can replace `PEEK_DOMAIN` when you need a full WebSocket URL. `PEEK_TOKEN` is still accepted as an old alias for `PEEK_AUTH_TOKEN`. + --- ### Uninstall the CLI diff --git a/crates/peek-client/Cargo.toml b/crates/peek-client/Cargo.toml index cf9e2d7..4c98420 100644 --- a/crates/peek-client/Cargo.toml +++ b/crates/peek-client/Cargo.toml @@ -9,6 +9,7 @@ path = "src/main.rs" [dependencies] peek-proto = { path = "../peek-proto" } +clap = { version = "4", features = ["derive"] } tokio = { version = "1", features = ["full"] } tokio-tungstenite = { version = "0.29", features = ["native-tls"] } futures-util = "0.3" diff --git a/crates/peek-client/src/main.rs b/crates/peek-client/src/main.rs index e50422c..3e50be4 100644 --- a/crates/peek-client/src/main.rs +++ b/crates/peek-client/src/main.rs @@ -1,52 +1,85 @@ use std::env; +use clap::Parser; use peek_client::TunnelClient; #[tokio::main] async fn main() { if let Err(error) = run().await { eprintln!("error: {error}"); - eprintln!(); - print_usage(); + eprintln!("run `peek --help` for usage"); std::process::exit(1); } } -async fn run() -> Result<(), String> { - let raw_args: Vec = env::args().skip(1).collect(); - if raw_args.iter().any(|arg| arg == "--help" || arg == "-h") { - print_usage(); - std::process::exit(0); - } +#[derive(Debug, Parser)] +#[command( + name = "peek", + version, + about = "Create a public URL for a local server through your own peek relay." +)] +struct Cli { + #[arg( + value_name = "LOCAL_ADDRESS", + help = "Local address, like localhost:3000 or 3000" + )] + local: String, + + #[arg( + long, + value_name = "URL", + help = "Full WebSocket URL, like wss://example.com/tunnel" + )] + server: Option, + + #[arg( + long, + value_name = "DOMAIN", + help = "Hosted peek domain, like example.com" + )] + domain: Option, + + #[arg( + long, + value_name = "TOKEN", + help = "Server token used to create a tunnel" + )] + token: Option, + + #[arg( + long, + value_name = "PASSWORD", + help = "Require this password for visitors" + )] + password: Option, + + #[arg(long, value_name = "NAME", help = "Public URL name, like myapp")] + subdomain: Option, +} - let mut args = raw_args.into_iter(); - - let local = args.next().ok_or("missing local address")?; - let port = parse_local_port(&local)?; - let local_url = local_url(&local, port); - - let mut server_url = env::var("PEEK_SERVER").ok(); - let mut domain = env::var("PEEK_DOMAIN").ok(); - let mut token = env::var("PEEK_TOKEN") - .ok() - .or_else(|| env::var("PEEK_AUTH_TOKEN").ok()); - let mut password = env::var("PEEK_PASSWORD").ok(); - let mut subdomain = None; - - while let Some(arg) = args.next() { - match arg.as_str() { - "--server" => server_url = Some(next_option_value(&mut args, "--server")?), - "--domain" => domain = Some(next_option_value(&mut args, "--domain")?), - "--token" => token = Some(next_option_value(&mut args, "--token")?), - "--password" => password = Some(next_option_value(&mut args, "--password")?), - "--subdomain" => subdomain = Some(next_option_value(&mut args, "--subdomain")?), - _ => return Err(format!("unknown option: {arg}")), - } - } +async fn run() -> Result<(), String> { + let cli = Cli::parse(); + let port = parse_local_port(&cli.local)?; + let local_url = local_url(&cli.local, port); + + let server_url = cli + .server + .or_else(|| env::var("PEEK_SERVER").ok()) + .or_else(|| { + cli.domain + .or_else(|| env::var("PEEK_DOMAIN").ok()) + .map(|domain| format!("wss://{domain}/tunnel")) + }) + .ok_or("missing relay: pass --domain, --server, or set PEEK_DOMAIN")?; + let token = cli + .token + .or_else(|| env::var("PEEK_AUTH_TOKEN").ok()) + .or_else(|| env::var("PEEK_TOKEN").ok()); + let password = cli.password.or_else(|| env::var("PEEK_PASSWORD").ok()); + let password_enabled = password + .as_ref() + .is_some_and(|password| !password.is_empty()); - let server_url = server_url - .or_else(|| domain.map(|domain| format!("wss://{domain}/tunnel"))) - .ok_or("missing --server, --domain, PEEK_SERVER, or PEEK_DOMAIN")?; let mut client = TunnelClient::new(&server_url).map_err(|error| error.to_string())?; if let Some(token) = token { client = client.with_token(token); @@ -56,11 +89,11 @@ async fn run() -> Result<(), String> { } let handle = client - .connect_with_subdomain(port, subdomain) + .connect_with_subdomain(port, cli.subdomain) .await .map_err(|error| error.to_string())?; - print_tunnel_summary(&local_url, handle.url()); + print_tunnel_summary(&local_url, handle.url(), password_enabled); tokio::signal::ctrl_c() .await .map_err(|error| error.to_string())?; @@ -69,40 +102,6 @@ async fn run() -> Result<(), String> { Ok(()) } -fn print_usage() { - eprintln!("usage: peek [options]"); - eprintln!(); - eprintln!("example:"); - eprintln!(" peek localhost:3000 --domain example.com --token change-me"); - eprintln!(); - eprintln!("options:"); - eprintln!(" --server full WebSocket URL, like wss://example.com/tunnel"); - eprintln!(" --domain hosted peek domain, like example.com"); - eprintln!(" --token server token used to create a tunnel"); - eprintln!(" --password require this password for visitors"); - eprintln!(" --subdomain public URL name, like https://name.example.com"); - eprintln!(); - eprintln!("environment:"); - eprintln!(" PEEK_SERVER full WebSocket URL"); - eprintln!(" PEEK_DOMAIN hosted peek domain"); - eprintln!(" PEEK_TOKEN server token used to create a tunnel"); - eprintln!(" PEEK_AUTH_TOKEN same as PEEK_TOKEN"); - eprintln!(" PEEK_PASSWORD require this password for visitors"); -} - -fn next_option_value( - args: &mut impl Iterator, - option: &str, -) -> Result { - let value = args - .next() - .ok_or_else(|| format!("{option} needs a value"))?; - if value.starts_with("--") { - return Err(format!("{option} needs a value")); - } - Ok(value) -} - fn parse_local_port(local: &str) -> Result { if let Ok(port) = local.parse::() { return Ok(port); @@ -138,9 +137,14 @@ fn local_url(local: &str, port: u16) -> String { format!("http://{local}") } -fn print_tunnel_summary(local_url: &str, public_url: &str) { +fn print_tunnel_summary(local_url: &str, public_url: &str, password_enabled: bool) { let label_width = "Public URL".len(); - let value_width = local_url.len().max(public_url.len()); + let password = if password_enabled { + "enabled" + } else { + "disabled" + }; + let value_width = local_url.len().max(public_url.len()).max(password.len()); let border = format!( "+{}+{}+", "-".repeat(label_width + 2), @@ -159,8 +163,14 @@ fn print_tunnel_summary(local_url: &str, public_url: &str) { "| {:label_width$} | {:value_width$} |", "Public URL", public_url ); + println!( + "| {:label_width$} | {:value_width$} |", + "Password", password + ); println!("{border}"); println!(); + println!("Press Ctrl+C to stop."); + println!(); } #[cfg(test)] diff --git a/crates/peek-relay/src/lib.rs b/crates/peek-relay/src/lib.rs new file mode 100644 index 0000000..5ed6965 --- /dev/null +++ b/crates/peek-relay/src/lib.rs @@ -0,0 +1,53 @@ +mod handler; +mod pages; +mod rate_limit; +mod registry; + +use std::sync::Arc; +use std::time::Duration; + +use axum::{Router, extract::DefaultBodyLimit, routing::get}; +use tower_http::trace::TraceLayer; + +use handler::{public_handler, ws_handler}; +use rate_limit::RateLimiter; +use registry::Registry; + +pub struct AppConfig { + pub domain: String, + pub auth_token: Option, + pub max_tunnels: usize, + pub max_body_size: usize, + pub rate_limit_rpm: u32, + pub trust_proxy_headers: bool, +} + +pub fn build_app(config: AppConfig) -> Router { + let rate_limiter = RateLimiter::new(config.rate_limit_rpm, Duration::from_secs(60)); + let registry = Arc::new(Registry::new( + config.domain, + config.auth_token, + config.max_tunnels, + config.max_body_size, + config.trust_proxy_headers, + rate_limiter, + )); + + { + let registry = registry.clone(); + tokio::spawn(async move { + let mut interval = tokio::time::interval(Duration::from_secs(60)); + loop { + interval.tick().await; + registry.rate_limiter.cleanup(); + } + }); + } + + Router::new() + .route("/tunnel", get(ws_handler)) + .fallback(public_handler) + .layer(DefaultBodyLimit::max(config.max_body_size)) + .layer(TraceLayer::new_for_http()) + .with_state(registry) +} diff --git a/crates/peek-relay/src/main.rs b/crates/peek-relay/src/main.rs index 006e7fb..4606244 100644 --- a/crates/peek-relay/src/main.rs +++ b/crates/peek-relay/src/main.rs @@ -1,21 +1,11 @@ -mod handler; -mod pages; -mod rate_limit; -mod registry; - use std::error::Error; use std::net::SocketAddr; -use std::sync::Arc; use std::time::Duration; -use axum::{Router, extract::DefaultBodyLimit, routing::get}; -use tower_http::trace::TraceLayer; use tracing::{info, warn}; use tracing_subscriber::EnvFilter; -use handler::{public_handler, ws_handler}; -use rate_limit::RateLimiter; -use registry::Registry; +use peek_relay::{AppConfig, build_app}; #[tokio::main] async fn main() { @@ -70,33 +60,14 @@ async fn run() -> Result<(), Box> { "starting peek" ); - let rate_limiter = RateLimiter::new(rate_limit_rpm, Duration::from_secs(60)); - let registry = Arc::new(Registry::new( + let app = build_app(AppConfig { domain, - Some(auth_token), + auth_token: Some(auth_token), max_tunnels, max_body_size, + rate_limit_rpm, trust_proxy_headers, - rate_limiter, - )); - - { - let registry = registry.clone(); - tokio::spawn(async move { - let mut interval = tokio::time::interval(Duration::from_secs(60)); - loop { - interval.tick().await; - registry.rate_limiter.cleanup(); - } - }); - } - - let app = Router::new() - .route("/tunnel", get(ws_handler)) - .fallback(public_handler) - .layer(DefaultBodyLimit::max(max_body_size)) - .layer(TraceLayer::new_for_http()) - .with_state(registry); + }); let addr = format!("0.0.0.0:{port}"); let listener = tokio::net::TcpListener::bind(&addr).await?; diff --git a/crates/peek-relay/tests/integration.rs b/crates/peek-relay/tests/integration.rs index ace1f40..db598d8 100644 --- a/crates/peek-relay/tests/integration.rs +++ b/crates/peek-relay/tests/integration.rs @@ -1,226 +1,36 @@ use std::fmt::Write as _; use std::net::SocketAddr; -use std::sync::Arc; use std::time::Duration; use axum::{Router, routing::get}; -use peek_proto::{RegistrationRequest, RegistrationResponse}; +use peek_relay::{AppConfig, build_app}; use tokio::net::TcpListener; async fn start_relay(domain: &str) -> SocketAddr { - let registry = Arc::new(TestRegistry::new(domain.to_string())); - let reg = registry.clone(); - - let app = Router::new() - .route( - "/tunnel", - get(move |ws: axum::extract::ws::WebSocketUpgrade| { - let registry = reg.clone(); - async move { - ws.max_frame_size(10 * 1024 * 1024) - .max_message_size(10 * 1024 * 1024) - .on_upgrade(move |socket| test_handle_tunnel(socket, registry)) - } - }), - ) - .fallback(move |req: axum::extract::Request| { - let registry = registry.clone(); - async move { test_public_handler(registry, req).await } - }); + let app = build_app(AppConfig { + domain: domain.to_string(), + auth_token: None, + max_tunnels: 100, + max_body_size: 10 * 1024 * 1024, + rate_limit_rpm: 10_000, + trust_proxy_headers: false, + }); let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); tokio::spawn(async move { - axum::serve(listener, app).await.unwrap(); + axum::serve( + listener, + app.into_make_service_with_connect_info::(), + ) + .await + .unwrap(); }); addr } -use std::collections::HashMap; -use std::sync::atomic::{AtomicU32, Ordering}; -use tokio::sync::{Mutex, RwLock, mpsc, oneshot}; - -use axum::extract::ws::{Message, WebSocket}; -use axum::response::{IntoResponse, Response}; -use futures_util::{SinkExt, StreamExt}; - -struct TestRegistry { - domain: String, - tunnels: RwLock>>, -} - -struct TestTunnelConn { - write_tx: mpsc::Sender, - pending: Mutex>>>, - next_request_id: AtomicU32, -} - -impl TestRegistry { - fn new(domain: String) -> Self { - Self { - domain, - tunnels: RwLock::new(HashMap::new()), - } - } -} - -impl TestTunnelConn { - fn new(write_tx: mpsc::Sender) -> Self { - Self { - write_tx, - pending: Mutex::new(HashMap::new()), - next_request_id: AtomicU32::new(1), - } - } -} - -async fn test_handle_tunnel(socket: WebSocket, registry: Arc) { - let (mut sink, mut stream) = socket.split(); - - let Ok(Some(Ok(Message::Text(msg)))) = - tokio::time::timeout(Duration::from_secs(5), stream.next()).await - else { - return; - }; - let reg_req: RegistrationRequest = match serde_json::from_str(&msg) { - Ok(r) => r, - Err(_) => return, - }; - - let subdomain = reg_req - .subdomain - .unwrap_or_else(|| "testsubdomain".to_string()); - - let (write_tx, mut write_rx) = mpsc::channel::(256); - - let writer_handle = tokio::spawn(async move { - while let Some(msg) = write_rx.recv().await { - if sink.send(msg).await.is_err() { - break; - } - } - let _ = SinkExt::close(&mut sink).await; - }); - - let conn = Arc::new(TestTunnelConn::new(write_tx.clone())); - registry - .tunnels - .write() - .await - .insert(subdomain.clone(), conn.clone()); - - let resp = RegistrationResponse { - ok: true, - url: format!("http://{subdomain}.{}", registry.domain), - subdomain: subdomain.clone(), - error: None, - }; - let json = serde_json::to_string(&resp).unwrap(); - if write_tx.send(Message::Text(json.into())).await.is_err() { - return; - } - - while let Some(msg) = stream.next().await { - match msg { - Ok(Message::Binary(data)) => { - if let Ok((request_id, payload)) = peek_proto::decode_frame(&data) { - let mut pending = conn.pending.lock().await; - if let Some(tx) = pending.remove(&request_id) { - let _ = tx.send(payload.to_vec()); - } - } - } - Ok(Message::Ping(data)) => { - let _ = write_tx.send(Message::Pong(data)).await; - } - Ok(Message::Close(_)) | Err(_) => break, - _ => {} - } - } - - registry.tunnels.write().await.remove(&subdomain); - drop(write_tx); - let _ = writer_handle.await; -} - -async fn test_public_handler( - registry: Arc, - request: axum::extract::Request, -) -> Response { - use axum::body::Body; - use http_body_util::BodyExt; - - let host = request - .headers() - .get("host") - .and_then(|v| v.to_str().ok()) - .unwrap_or("") - .to_string(); - - let host_no_port = host.split(':').next().unwrap_or(&host); - let suffix = format!(".{}", registry.domain); - let subdomain = if host_no_port.ends_with(&suffix) { - let sub = &host_no_port[..host_no_port.len() - suffix.len()]; - if !sub.is_empty() && !sub.contains('.') { - sub.to_string() - } else { - return (axum::http::StatusCode::NOT_FOUND, "not found").into_response(); - } - } else { - return (axum::http::StatusCode::NOT_FOUND, "not found").into_response(); - }; - - let conn = registry.tunnels.read().await.get(&subdomain).cloned(); - let Some(conn) = conn else { - return (axum::http::StatusCode::NOT_FOUND, "tunnel not found").into_response(); - }; - - let method = request.method().to_string(); - let uri = request.uri().to_string(); - let headers: Vec<(String, String)> = request - .headers() - .iter() - .map(|(k, v)| (k.to_string(), v.to_str().unwrap_or("").to_string())) - .collect(); - let body_bytes = request - .into_body() - .collect() - .await - .map(|c| c.to_bytes().to_vec()) - .unwrap_or_default(); - - let serialized = peek_proto::serialize_request(&method, &uri, &headers, &body_bytes); - let request_id = conn.next_request_id.fetch_add(1, Ordering::Relaxed); - let (tx, rx) = oneshot::channel::>(); - conn.pending.lock().await.insert(request_id, tx); - - let frame = peek_proto::encode_frame(request_id, &serialized); - if conn - .write_tx - .send(Message::Binary(frame.into())) - .await - .is_err() - { - return (axum::http::StatusCode::BAD_GATEWAY, "ws send failed").into_response(); - } - - match tokio::time::timeout(Duration::from_secs(10), rx).await { - Ok(Ok(data)) => match peek_proto::deserialize_response(&data) { - Ok(resp) => { - let mut builder = Response::builder().status(resp.status); - for (k, v) in &resp.headers { - builder = builder.header(k, v); - } - builder.body(Body::from(resp.body)).unwrap() - } - Err(_) => (axum::http::StatusCode::BAD_GATEWAY, "bad response").into_response(), - }, - _ => (axum::http::StatusCode::GATEWAY_TIMEOUT, "timeout").into_response(), - } -} - async fn start_local_server() -> (SocketAddr, &'static str) { let expected_body = "Hello from local server!"; @@ -284,7 +94,10 @@ async fn test_tunnel_end_to_end() { let relay_addr = start_relay("test.local").await; let client = peek_client::TunnelClient::new(&format!("ws://{relay_addr}/tunnel")).unwrap(); - let handle = client.connect(local_addr.port()).await.unwrap(); + let handle = client + .connect_with_subdomain(local_addr.port(), Some("testsubdomain".into())) + .await + .unwrap(); assert!(handle.url().contains("test.local")); tokio::time::sleep(Duration::from_millis(100)).await; @@ -336,7 +149,10 @@ async fn test_tunnel_unreachable_local_server() { let relay_addr = start_relay("test3.local").await; let client = peek_client::TunnelClient::new(&format!("ws://{relay_addr}/tunnel")).unwrap(); - let handle = client.connect(19999).await.unwrap(); + let handle = client + .connect_with_subdomain(19999, Some("testsubdomain".into())) + .await + .unwrap(); tokio::time::sleep(Duration::from_millis(100)).await; @@ -359,7 +175,10 @@ async fn test_concurrent_requests_through_tunnel() { let relay_addr = start_relay("test-concurrent.local").await; let client = peek_client::TunnelClient::new(&format!("ws://{relay_addr}/tunnel")).unwrap(); - let handle = client.connect(local_addr.port()).await.unwrap(); + let handle = client + .connect_with_subdomain(local_addr.port(), Some("testsubdomain".into())) + .await + .unwrap(); tokio::time::sleep(Duration::from_millis(100)).await; let http_client = reqwest::Client::new(); @@ -393,7 +212,10 @@ async fn test_large_request_body() { let relay_addr = start_relay("test-large.local").await; let client = peek_client::TunnelClient::new(&format!("ws://{relay_addr}/tunnel")).unwrap(); - let handle = client.connect(local_addr.port()).await.unwrap(); + let handle = client + .connect_with_subdomain(local_addr.port(), Some("testsubdomain".into())) + .await + .unwrap(); tokio::time::sleep(Duration::from_millis(100)).await; let http_client = reqwest::Client::new(); @@ -460,7 +282,10 @@ async fn test_various_http_methods() { let relay_addr = start_relay("test-methods.local").await; let client = peek_client::TunnelClient::new(&format!("ws://{relay_addr}/tunnel")).unwrap(); - let handle = client.connect(local_addr.port()).await.unwrap(); + let handle = client + .connect_with_subdomain(local_addr.port(), Some("testsubdomain".into())) + .await + .unwrap(); tokio::time::sleep(Duration::from_millis(100)).await; let http_client = reqwest::Client::new(); @@ -490,7 +315,10 @@ async fn test_non_200_status_codes_forwarded() { let relay_addr = start_relay("test-status.local").await; let client = peek_client::TunnelClient::new(&format!("ws://{relay_addr}/tunnel")).unwrap(); - let handle = client.connect(local_addr.port()).await.unwrap(); + let handle = client + .connect_with_subdomain(local_addr.port(), Some("testsubdomain".into())) + .await + .unwrap(); tokio::time::sleep(Duration::from_millis(100)).await; let http_client = reqwest::Client::new(); @@ -532,7 +360,10 @@ async fn test_slow_local_server_responds() { let relay_addr = start_relay("test-slow.local").await; let client = peek_client::TunnelClient::new(&format!("ws://{relay_addr}/tunnel")).unwrap(); - let handle = client.connect(local_addr.port()).await.unwrap(); + let handle = client + .connect_with_subdomain(local_addr.port(), Some("testsubdomain".into())) + .await + .unwrap(); tokio::time::sleep(Duration::from_millis(100)).await; let http_client = reqwest::Client::builder()