#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
use crate::net::events::NetEvent;
use crate::net::observer::NetObserver;
use crate::net::types::BlockReason;
use dashmap::DashMap;
use std::collections::HashMap;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio_util::sync::CancellationToken;
use url::Url;
fn reason(code: u16) -> &'static str {
match code {
200 => "OK",
301 => "Moved Permanently",
302 => "Found",
400 => "Bad Request",
401 => "Unauthorized",
403 => "Forbidden",
404 => "Not Found",
500 => "Internal Server Error",
_ => "Unknown",
}
}
#[derive(Clone)]
pub enum RouteConfig {
Ok(Vec<u8>),
OkWithHeaders(Vec<(String, String)>, Vec<u8>),
Status(u16, Vec<u8>),
Delay(Duration, Vec<u8>),
StallMidBody {
initial: usize,
stall: Duration,
},
DropMidBody {
prefix: usize,
total: usize,
},
RedirectTo(String),
Redirect307(String),
RedirectAbsolute(String),
RedirectSelf,
NoLocationRedirect,
RedirectWithCookie {
target: String,
cookie: String,
},
EchoCookieHeader,
EchoRefererHeader,
RedirectWithReferrerPolicy {
target: String,
policy: String,
},
HangAfterConnect,
Chunked(Vec<Vec<u8>>),
ChunkedWithDelay {
chunks: Vec<Vec<u8>>,
delay: Duration,
},
GzipOk(Vec<u8>),
EchoBody,
}
impl RouteConfig {
pub fn ok(body: impl Into<Vec<u8>>) -> Self {
Self::Ok(body.into())
}
pub fn ok_with_headers(headers: &[(&str, &str)], body: impl Into<Vec<u8>>) -> Self {
Self::OkWithHeaders(
headers
.iter()
.map(|(k, v)| (k.to_string(), v.to_string()))
.collect(),
body.into(),
)
}
pub fn status(code: u16, body: impl Into<Vec<u8>>) -> Self {
Self::Status(code, body.into())
}
pub fn delay(d: Duration, body: impl Into<Vec<u8>>) -> Self {
Self::Delay(d, body.into())
}
pub fn stall_mid_body(initial: usize, stall: Duration) -> Self {
Self::StallMidBody { initial, stall }
}
pub fn drop_mid_body(prefix: usize, total: usize) -> Self {
Self::DropMidBody { prefix, total }
}
pub fn redirect_to(path: impl Into<String>) -> Self {
Self::RedirectTo(path.into())
}
pub fn redirect_307(path: impl Into<String>) -> Self {
Self::Redirect307(path.into())
}
pub fn redirect_absolute(target: impl Into<String>) -> Self {
Self::RedirectAbsolute(target.into())
}
pub fn gzip_ok(body: impl Into<Vec<u8>>) -> Self {
Self::GzipOk(body.into())
}
pub fn echo_body() -> Self {
Self::EchoBody
}
pub fn redirect_self() -> Self {
Self::RedirectSelf
}
pub fn no_location_redirect() -> Self {
Self::NoLocationRedirect
}
pub fn redirect_with_cookie(target: impl Into<String>, cookie: impl Into<String>) -> Self {
Self::RedirectWithCookie {
target: target.into(),
cookie: cookie.into(),
}
}
pub fn echo_cookie_header() -> Self {
Self::EchoCookieHeader
}
pub fn echo_referer_header() -> Self {
Self::EchoRefererHeader
}
pub fn redirect_with_referrer_policy(
target: impl Into<String>,
policy: impl Into<String>,
) -> Self {
Self::RedirectWithReferrerPolicy {
target: target.into(),
policy: policy.into(),
}
}
pub fn hang_after_connect() -> Self {
Self::HangAfterConnect
}
pub fn chunked(chunks: Vec<&[u8]>) -> Self {
Self::Chunked(chunks.into_iter().map(|c| c.to_vec()).collect())
}
pub fn chunked_with_delay(chunks: Vec<&[u8]>, delay: Duration) -> Self {
Self::ChunkedWithDelay {
chunks: chunks.into_iter().map(|c| c.to_vec()).collect(),
delay,
}
}
}
async fn send_response<S: AsyncWrite + Unpin>(stream: &mut S, code: u16, body: &[u8]) {
let hdr = format!(
"HTTP/1.1 {} {}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
code,
reason(code),
body.len()
);
let _ = stream.write_all(hdr.as_bytes()).await;
let _ = stream.write_all(body).await;
}
async fn handle_conn<S: AsyncRead + AsyncWrite + Unpin>(
mut stream: S,
routes: Arc<HashMap<String, RouteConfig>>,
default: Arc<RouteConfig>,
hits: Arc<DashMap<String, AtomicUsize>>,
base: Arc<String>,
) {
let mut buf = [0u8; 4096];
let n = stream.read(&mut buf).await.unwrap_or(0);
let req = std::str::from_utf8(&buf[..n]).unwrap_or("");
let raw_path = req
.lines()
.next()
.and_then(|l| l.split_whitespace().nth(1))
.unwrap_or("/");
let path = raw_path.split('?').next().unwrap_or(raw_path).to_string();
hits.entry(path.clone())
.or_insert_with(|| AtomicUsize::new(0))
.fetch_add(1, Ordering::Relaxed);
let cfg = routes
.get(&path)
.cloned()
.unwrap_or_else(|| (*default).clone());
match cfg {
RouteConfig::Ok(body) => {
send_response(&mut stream, 200, &body).await;
}
RouteConfig::OkWithHeaders(headers, body) => {
let extra: String = headers
.iter()
.map(|(k, v)| format!("{k}: {v}\r\n"))
.collect();
let hdr = format!(
"HTTP/1.1 200 OK\r\n{}Content-Length: {}\r\nConnection: close\r\n\r\n",
extra,
body.len()
);
let _ = stream.write_all(hdr.as_bytes()).await;
let _ = stream.write_all(&body).await;
}
RouteConfig::Status(code, body) => {
send_response(&mut stream, code, &body).await;
}
RouteConfig::Delay(d, body) => {
tokio::time::sleep(d).await;
send_response(&mut stream, 200, &body).await;
}
RouteConfig::StallMidBody { initial, stall } => {
let declared = initial + 8192;
let hdr = format!(
"HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
declared
);
let _ = stream.write_all(hdr.as_bytes()).await;
if initial > 0 {
let _ = stream.write_all(&vec![b'X'; initial]).await;
let _ = stream.flush().await;
}
tokio::time::sleep(stall).await;
}
RouteConfig::DropMidBody { prefix, total } => {
let hdr = format!(
"HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
total
);
let _ = stream.write_all(hdr.as_bytes()).await;
if prefix > 0 {
let _ = stream.write_all(&vec![b'X'; prefix]).await;
let _ = stream.flush().await;
}
}
RouteConfig::RedirectTo(target) => {
let hdr = format!(
"HTTP/1.1 302 Found\r\nLocation: {}{}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
base, target
);
let _ = stream.write_all(hdr.as_bytes()).await;
}
RouteConfig::Redirect307(target) => {
let hdr = format!(
"HTTP/1.1 307 Temporary Redirect\r\nLocation: {}{}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
base, target
);
let _ = stream.write_all(hdr.as_bytes()).await;
}
RouteConfig::RedirectAbsolute(target) => {
let hdr = format!(
"HTTP/1.1 302 Found\r\nLocation: {}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
target
);
let _ = stream.write_all(hdr.as_bytes()).await;
}
RouteConfig::RedirectSelf => {
let hdr = format!(
"HTTP/1.1 302 Found\r\nLocation: {}{}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
base, path
);
let _ = stream.write_all(hdr.as_bytes()).await;
}
RouteConfig::NoLocationRedirect => {
let _ = stream
.write_all(b"HTTP/1.1 302 Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")
.await;
}
RouteConfig::RedirectWithCookie { target, cookie } => {
let hdr = format!(
"HTTP/1.1 302 Found\r\nLocation: {}{}\r\nSet-Cookie: {}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
base, target, cookie
);
let _ = stream.write_all(hdr.as_bytes()).await;
}
RouteConfig::EchoCookieHeader => {
let cookie = req
.lines()
.find(|l| l.to_ascii_lowercase().starts_with("cookie:"))
.and_then(|l| l.split_once(':').map(|(_, v)| v.trim().to_string()))
.unwrap_or_default();
send_response(&mut stream, 200, cookie.as_bytes()).await;
}
RouteConfig::RedirectWithReferrerPolicy { target, policy } => {
let hdr = format!(
"HTTP/1.1 302 Found\r\nLocation: {}{}\r\nReferrer-Policy: {}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
base, target, policy
);
let _ = stream.write_all(hdr.as_bytes()).await;
}
RouteConfig::EchoRefererHeader => {
let referer = req
.lines()
.find(|l| l.to_ascii_lowercase().starts_with("referer:"))
.and_then(|l| l.split_once(':').map(|(_, v)| v.trim().to_string()))
.unwrap_or_else(|| "<absent>".to_string());
send_response(&mut stream, 200, referer.as_bytes()).await;
}
RouteConfig::HangAfterConnect => {
tokio::time::sleep(Duration::from_secs(3600)).await;
}
RouteConfig::Chunked(chunks) => {
let _ = stream
.write_all(
b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n",
)
.await;
for chunk in &chunks {
let _ = stream
.write_all(format!("{:x}\r\n", chunk.len()).as_bytes())
.await;
let _ = stream.write_all(chunk).await;
let _ = stream.write_all(b"\r\n").await;
}
let _ = stream.write_all(b"0\r\n\r\n").await;
}
RouteConfig::ChunkedWithDelay { chunks, delay } => {
let _ = stream
.write_all(
b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n",
)
.await;
let _ = stream.flush().await;
for chunk in &chunks {
tokio::time::sleep(delay).await;
let _ = stream
.write_all(format!("{:x}\r\n", chunk.len()).as_bytes())
.await;
let _ = stream.write_all(chunk).await;
let _ = stream.write_all(b"\r\n").await;
let _ = stream.flush().await;
}
let _ = stream.write_all(b"0\r\n\r\n").await;
}
RouteConfig::EchoBody => {
let content_length: usize = req
.lines()
.find(|l| l.to_ascii_lowercase().starts_with("content-length:"))
.and_then(|l| l.split_once(':').map(|(_, v)| v))
.and_then(|v| v.trim().parse().ok())
.unwrap_or(0);
let header_end = buf[..n]
.windows(4)
.position(|w| w == b"\r\n\r\n")
.map(|p| p + 4)
.unwrap_or(n);
let mut body = buf[header_end..n].to_vec();
while body.len() < content_length {
let mut extra = [0u8; 4096];
let k = stream.read(&mut extra).await.unwrap_or(0);
if k == 0 {
break;
}
body.extend_from_slice(&extra[..k]);
}
body.truncate(content_length);
send_response(&mut stream, 200, &body).await;
}
RouteConfig::GzipOk(body) => {
use flate2::{write::GzEncoder, Compression};
use std::io::Write as _;
let mut enc = GzEncoder::new(Vec::new(), Compression::default());
enc.write_all(&body).unwrap();
let compressed = enc.finish().unwrap();
let hdr = format!(
"HTTP/1.1 200 OK\r\nContent-Encoding: gzip\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
compressed.len()
);
let _ = stream.write_all(hdr.as_bytes()).await;
let _ = stream.write_all(&compressed).await;
}
}
}
pub struct TestServer {
routes: HashMap<String, RouteConfig>,
default: RouteConfig,
tls_domain: Option<String>,
}
impl Default for TestServer {
fn default() -> Self {
Self::new()
}
}
impl TestServer {
pub fn new() -> Self {
Self {
routes: HashMap::new(),
default: RouteConfig::Ok(b"hello".to_vec()),
tls_domain: None,
}
}
pub fn tls(mut self, domain: &str) -> Self {
self.tls_domain = Some(domain.to_string());
self
}
pub fn route(mut self, path: &str, config: RouteConfig) -> Self {
self.routes.insert(path.to_string(), config);
self
}
pub fn default_route(mut self, config: RouteConfig) -> Self {
self.default = config;
self
}
pub async fn start(self) -> TestServerHandle {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let hits: Arc<DashMap<String, AtomicUsize>> = Arc::new(DashMap::new());
let shutdown = CancellationToken::new();
let tls = self.tls_domain.as_ref().map(|domain| {
let ck = rcgen::generate_simple_self_signed(vec![domain.clone()]).unwrap();
let cert_pem = ck.cert.pem().into_bytes();
let key = tokio_rustls::rustls::pki_types::PrivateKeyDer::try_from(
rcgen::KeyPair::serialize_der(&ck.signing_key),
)
.unwrap();
let server_config = tokio_rustls::rustls::ServerConfig::builder()
.with_no_client_auth()
.with_single_cert(vec![ck.cert.der().clone()], key)
.unwrap();
(
domain.clone(),
cert_pem,
tokio_rustls::TlsAcceptor::from(Arc::new(server_config)),
)
});
let base = Arc::new(match &tls {
Some((domain, _, _)) => format!("https://{}:{}", domain, addr.port()),
None => format!("http://127.0.0.1:{}", addr.port()),
});
let routes = Arc::new(self.routes);
let default = Arc::new(self.default);
let hits_srv = hits.clone();
let shutdown_srv = shutdown.clone();
let acceptor = tls.as_ref().map(|(_, _, a)| a.clone());
let base_srv = base.clone();
tokio::spawn(async move {
loop {
tokio::select! {
_ = shutdown_srv.cancelled() => break,
result = listener.accept() => {
let Ok((stream, _)) = result else { break };
let routes = routes.clone();
let default = default.clone();
let hits = hits_srv.clone();
let base = base_srv.clone();
let acceptor = acceptor.clone();
tokio::spawn(async move {
match acceptor {
Some(acceptor) => {
if let Ok(tls_stream) = acceptor.accept(stream).await {
handle_conn(tls_stream, routes, default, hits, base).await
}
}
None => handle_conn(stream, routes, default, hits, base).await,
}
});
}
}
}
});
TestServerHandle {
addr,
hits,
shutdown,
tls: tls.map(|(domain, cert_pem, _)| TlsInfo { domain, cert_pem }),
}
}
}
struct TlsInfo {
domain: String,
cert_pem: Vec<u8>,
}
pub struct TestServerHandle {
addr: std::net::SocketAddr,
hits: Arc<DashMap<String, AtomicUsize>>,
shutdown: CancellationToken,
tls: Option<TlsInfo>,
}
impl TestServerHandle {
pub fn url(&self, path: &str) -> Url {
let base = match &self.tls {
Some(t) => format!("https://{}:{}", t.domain, self.addr.port()),
None => format!("http://127.0.0.1:{}", self.addr.port()),
};
Url::parse(&format!("{base}{path}")).unwrap()
}
pub fn base_url(&self) -> Url {
self.url("/")
}
pub fn socket_addr(&self) -> std::net::SocketAddr {
self.addr
}
pub fn cert_pem(&self) -> Option<&[u8]> {
self.tls.as_ref().map(|t| t.cert_pem.as_slice())
}
pub fn tls_domain(&self) -> Option<&str> {
self.tls.as_ref().map(|t| t.domain.as_str())
}
pub fn hit_count(&self, path: &str) -> usize {
self.hits
.get(path)
.map(|e| e.load(Ordering::Relaxed))
.unwrap_or(0)
}
}
impl Drop for TestServerHandle {
fn drop(&mut self) {
self.shutdown.cancel();
}
}
#[derive(Default)]
pub struct RecordingObserver {
events: Mutex<Vec<NetEvent>>,
}
impl RecordingObserver {
pub fn new() -> Self {
Self::default()
}
pub fn blocked_reason(&self) -> Option<BlockReason> {
self.events.lock().unwrap().iter().find_map(|e| match e {
NetEvent::Blocked { reason, .. } => Some(*reason),
_ => None,
})
}
pub fn warnings(&self) -> Vec<String> {
self.events
.lock()
.unwrap()
.iter()
.filter_map(|e| match e {
NetEvent::Warning { message, .. } => Some(message.clone()),
_ => None,
})
.collect()
}
pub fn len(&self) -> usize {
self.events.lock().unwrap().len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
impl NetObserver for RecordingObserver {
fn on_event(&self, ev: NetEvent) {
self.events.lock().unwrap().push(ev);
}
}