use std::collections::VecDeque;
use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, MutexGuard, OnceLock};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
#[derive(Debug, Clone)]
pub enum Reply {
Full {
status: u16,
content_type: &'static str,
body: Vec<u8>,
},
Sse { frames: Vec<(Vec<u8>, u64)> },
}
impl Reply {
pub fn full(status: u16, content_type: &'static str, body: Vec<u8>) -> Self {
Reply::Full {
status,
content_type,
body,
}
}
pub fn sse(frames: Vec<(Vec<u8>, u64)>) -> Self {
Reply::Sse { frames }
}
}
#[derive(Debug, Clone)]
pub struct Recorded {
pub method: String,
pub path: String,
pub query: Option<String>,
pub headers: Vec<(String, String)>,
pub body: Vec<u8>,
}
impl Recorded {
pub fn header(&self, name: &str) -> Option<String> {
let name = name.to_ascii_lowercase();
self.headers
.iter()
.find(|(k, _)| *k == name)
.map(|(_, v)| v.clone())
}
}
#[derive(Debug)]
struct Route {
method: String,
path: String,
queue: VecDeque<Reply>,
}
pub struct StubServer {
addr: SocketAddr,
routes: Arc<Mutex<Vec<Route>>>,
requests: Arc<Mutex<Vec<Recorded>>>,
}
impl StubServer {
pub async fn start() -> StubServer {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind stub");
let addr = listener.local_addr().unwrap();
let routes: Arc<Mutex<Vec<Route>>> = Arc::new(Mutex::new(Vec::new()));
let requests: Arc<Mutex<Vec<Recorded>>> = Arc::new(Mutex::new(Vec::new()));
let accept_routes = Arc::clone(&routes);
let accept_requests = Arc::clone(&requests);
tokio::spawn(async move {
loop {
let Ok((stream, _)) = listener.accept().await else {
break;
};
let routes = Arc::clone(&accept_routes);
let requests = Arc::clone(&accept_requests);
tokio::spawn(async move {
let _ = handle_connection(stream, routes, requests).await;
});
}
});
StubServer {
addr,
routes,
requests,
}
}
pub fn base_url(&self) -> String {
format!("http://{}", self.addr)
}
pub fn port(&self) -> u16 {
self.addr.port()
}
pub fn route(&self, method: &str, path: &str, reply: Reply) {
self.route_seq(method, path, vec![reply]);
}
pub fn route_seq(&self, method: &str, path: &str, replies: Vec<Reply>) {
let mut routes = self.routes.lock().unwrap_or_else(|e| e.into_inner());
routes.push(Route {
method: method.to_string(),
path: path.to_string(),
queue: replies.into(),
});
}
pub fn requests(&self) -> Vec<Recorded> {
self.requests
.lock()
.unwrap_or_else(|e| e.into_inner())
.clone()
}
}
async fn handle_connection(
mut stream: tokio::net::TcpStream,
routes: Arc<Mutex<Vec<Route>>>,
requests: Arc<Mutex<Vec<Recorded>>>,
) -> std::io::Result<()> {
let mut buf: Vec<u8> = Vec::with_capacity(4096);
let head_end = loop {
let mut chunk = [0u8; 4096];
let n = stream.read(&mut chunk).await?;
if n == 0 {
return Ok(()); }
buf.extend_from_slice(&chunk[..n]);
if let Some(pos) = find_subsequence(&buf, b"\r\n\r\n") {
break pos;
}
if buf.len() > 1 << 20 {
return Ok(()); }
};
let head = String::from_utf8_lossy(&buf[..head_end]).to_string();
let mut lines = head.split("\r\n");
let request_line = lines.next().unwrap_or_default();
let mut parts = request_line.split_whitespace();
let method = parts.next().unwrap_or_default().to_string();
let target = parts.next().unwrap_or_default().to_string();
let (path, query) = match target.split_once('?') {
Some((p, q)) => (p.to_string(), Some(q.to_string())),
None => (target.clone(), None),
};
let headers: Vec<(String, String)> = lines
.filter_map(|l| l.split_once(':'))
.map(|(k, v)| (k.trim().to_ascii_lowercase(), v.trim().to_string()))
.collect();
let content_length: usize = headers
.iter()
.find(|(k, _)| k == "content-length")
.and_then(|(_, v)| v.parse().ok())
.unwrap_or(0);
let mut body: Vec<u8> = buf[head_end + 4..].to_vec();
while body.len() < content_length {
let mut chunk = [0u8; 4096];
let n = stream.read(&mut chunk).await?;
if n == 0 {
break;
}
body.extend_from_slice(&chunk[..n]);
}
requests
.lock()
.unwrap_or_else(|e| e.into_inner())
.push(Recorded {
method: method.clone(),
path: path.clone(),
query,
headers,
body,
});
let reply = {
let mut routes = routes.lock().unwrap_or_else(|e| e.into_inner());
routes
.iter_mut()
.find(|r| r.method == method && r.path == path)
.map(|r| {
if r.queue.len() > 1 {
r.queue.pop_front().unwrap()
} else {
r.queue
.front()
.cloned()
.expect("route with empty reply queue")
}
})
};
match reply {
Some(Reply::Full {
status,
content_type,
body,
}) => {
let head = format!(
"HTTP/1.1 {status} {}\r\ncontent-type: {content_type}\r\ncontent-length: {}\r\nconnection: close\r\n\r\n",
reason(status),
body.len()
);
stream.write_all(head.as_bytes()).await?;
stream.write_all(&body).await?;
}
Some(Reply::Sse { frames }) => {
let head =
"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncache-control: no-store\r\nconnection: close\r\n\r\n";
stream.write_all(head.as_bytes()).await?;
stream.flush().await?;
for (bytes, delay_ms) in frames {
if delay_ms > 0 {
tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await;
}
stream.write_all(&bytes).await?;
stream.flush().await?;
}
}
None => {
let body = format!(
"{{\"error\":\"stub: no route for {method} {path}\",\"code\":\"not_found\"}}"
);
let head = format!(
"HTTP/1.1 404 Not Found\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n",
body.len()
);
stream.write_all(head.as_bytes()).await?;
stream.write_all(body.as_bytes()).await?;
}
}
stream.shutdown().await.ok();
Ok(())
}
fn find_subsequence(haystack: &[u8], needle: &[u8]) -> Option<usize> {
haystack.windows(needle.len()).position(|w| w == needle)
}
fn reason(status: u16) -> &'static str {
match status {
200 => "OK",
202 => "Accepted",
400 => "Bad Request",
401 => "Unauthorized",
404 => "Not Found",
409 => "Conflict",
422 => "Unprocessable Entity",
423 => "Locked",
429 => "Too Many Requests",
500 => "Internal Server Error",
_ => "Status",
}
}
pub fn env_lock() -> MutexGuard<'static, ()> {
static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
LOCK.get_or_init(|| Mutex::new(()))
.lock()
.unwrap_or_else(|e| e.into_inner())
}
pub fn temp_dir(name: &str) -> PathBuf {
static COUNTER: AtomicU64 = AtomicU64::new(0);
let n = COUNTER.fetch_add(1, Ordering::Relaxed);
let dir = std::env::temp_dir().join(format!(
"writ-client-test-{}-{}-{}",
std::process::id(),
name,
n
));
std::fs::create_dir_all(&dir).expect("create temp dir");
dir
}