Skip to main content

andysu_cli/
serve.rs

1//! `andysu serve`: a tiny static HTTP/1.1 server for local remotes and tests. Supports
2//! `Range: bytes=N-`, `ETag`/`If-None-Match` (304), and fault injection: dropping the connection
3//! after N body bytes (once per path) to simulate network loss, and failing paths with 404/500.
4
5use std::collections::HashSet;
6use std::io::{BufRead, BufReader, Read, Write};
7use std::net::{TcpListener, TcpStream};
8use std::path::{Path, PathBuf};
9use std::sync::{Arc, Mutex};
10
11#[derive(Clone, Default)]
12pub struct Faults {
13    /// Close the connection after this many body bytes, the first time each path is served.
14    pub drop_after: Option<u64>,
15    /// Only drop paths starting with this prefix (empty = all).
16    pub drop_prefix: String,
17    /// Paths (prefix match, relative to the root) answered with this status.
18    pub fail: Vec<(String, u16)>,
19    /// Limit each response body to about this many bytes per second.
20    pub throttle_bps: Option<u64>,
21}
22
23#[derive(Clone)]
24pub struct Server {
25    pub port: u16,
26    pub log: Arc<Mutex<Vec<String>>>,
27    pub faults: Arc<Mutex<Faults>>,
28}
29
30impl Server {
31    pub fn url(&self) -> String {
32        format!("http://127.0.0.1:{}", self.port)
33    }
34    pub fn requests(&self) -> Vec<String> {
35        self.log.lock().unwrap().clone()
36    }
37    pub fn set_faults(&self, f: Faults) {
38        *self.faults.lock().unwrap() = f;
39    }
40}
41
42fn etag_of(p: &Path) -> Option<String> {
43    let md = std::fs::metadata(p).ok()?;
44    let mt = md.modified().ok()?.duration_since(std::time::UNIX_EPOCH).ok()?.as_nanos();
45    Some(format!("\"{:x}-{:x}\"", md.len(), mt))
46}
47
48fn handle(
49    mut s: TcpStream,
50    root: &Path,
51    log: &Mutex<Vec<String>>,
52    faults: &Mutex<Faults>,
53    dropped: &Mutex<HashSet<String>>,
54) -> std::io::Result<()> {
55    let mut r = BufReader::new(s.try_clone()?);
56    let mut line = String::new();
57    r.read_line(&mut line)?;
58    let mut parts = line.split_whitespace();
59    let (method, target) = (parts.next().unwrap_or(""), parts.next().unwrap_or("/").to_string());
60    let (mut range, mut inm) = (None, None);
61    loop {
62        let mut h = String::new();
63        if r.read_line(&mut h)? == 0 || h == "\r\n" || h == "\n" {
64            break;
65        }
66        let (k, v) = h.split_once(':').unwrap_or((&h, ""));
67        match k.trim().to_ascii_lowercase().as_str() {
68            "range" => {
69                range = v
70                    .trim()
71                    .strip_prefix("bytes=")
72                    .and_then(|x| x.strip_suffix('-'))
73                    .and_then(|x| x.parse::<u64>().ok())
74            }
75            "if-none-match" => inm = Some(v.trim().to_string()),
76            _ => {}
77        }
78    }
79    let path = target.split('?').next().unwrap_or("/").trim_start_matches('/').to_string();
80    let f = faults.lock().unwrap().clone();
81    let reply = |s: &mut TcpStream, code: u16, msg: &str| -> std::io::Result<()> {
82        write!(s, "HTTP/1.1 {code} {msg}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")
83    };
84    let mut entry = format!("{method} /{path}");
85    if let Some(o) = range {
86        entry.push_str(&format!(" range={o}"));
87    }
88    if let Some((_, code)) = f.fail.iter().find(|(pre, _)| path.starts_with(pre.as_str())) {
89        log.lock().unwrap().push(format!("{entry} -> {code}"));
90        return reply(&mut s, *code, "Injected");
91    }
92    let fp: PathBuf = root.join(&path);
93    if path.split('/').any(|c| c == "..") || !fp.is_file() {
94        log.lock().unwrap().push(format!("{entry} -> 404"));
95        return reply(&mut s, 404, "Not Found");
96    }
97    let tag = etag_of(&fp).unwrap_or_default();
98    if inm.as_deref() == Some(tag.as_str()) {
99        log.lock().unwrap().push(format!("{entry} -> 304"));
100        return write!(s, "HTTP/1.1 304 Not Modified\r\nETag: {tag}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n");
101    }
102    let mut file = std::fs::File::open(&fp)?;
103    let len = file.metadata()?.len();
104    let start = range.unwrap_or(0).min(len);
105    let mut body = vec![];
106    std::io::Seek::seek(&mut file, std::io::SeekFrom::Start(start))?;
107    file.read_to_end(&mut body)?;
108    let code = if range.is_some() { 206 } else { 200 };
109    log.lock().unwrap().push(format!("{entry} -> {code}"));
110    let mut head = format!(
111        "HTTP/1.1 {code} OK\r\nContent-Length: {}\r\nETag: {tag}\r\nAccept-Ranges: bytes\r\nConnection: close\r\n",
112        body.len()
113    );
114    if code == 206 {
115        head.push_str(&format!("Content-Range: bytes {start}-{}/{len}\r\n", len.saturating_sub(1)));
116    }
117    head.push_str("\r\n");
118    s.write_all(head.as_bytes())?;
119    if method == "HEAD" {
120        return Ok(());
121    }
122    if let Some(n) = f.drop_after {
123        if (n as usize) < body.len()
124            && path.starts_with(f.drop_prefix.as_str())
125            && dropped.lock().unwrap().insert(path.clone())
126        {
127            s.write_all(&body[..n as usize])?;
128            s.flush()?;
129            log.lock().unwrap().push(format!("{entry} dropped after {n} bytes"));
130            return s.shutdown(std::net::Shutdown::Both);
131        }
132    }
133    match f.throttle_bps {
134        Some(bps) if bps > 0 => {
135            let chunk = (bps / 20).max(1024) as usize;
136            for c in body.chunks(chunk) {
137                s.write_all(c)?;
138                std::thread::sleep(std::time::Duration::from_millis(50));
139            }
140            Ok(())
141        }
142        _ => s.write_all(&body),
143    }
144}
145
146/// Serve `root` on `addr` (use port 0 for an ephemeral port) in background threads.
147pub fn spawn(root: &Path, addr: &str, faults: Faults) -> std::io::Result<Server> {
148    let l = TcpListener::bind(addr)?;
149    let port = l.local_addr()?.port();
150    let srv = Server { port, log: Arc::new(Mutex::new(vec![])), faults: Arc::new(Mutex::new(faults)) };
151    let (root, s2) = (root.to_path_buf(), srv.clone());
152    let dropped = Arc::new(Mutex::new(HashSet::new()));
153    std::thread::spawn(move || {
154        for c in l.incoming().flatten() {
155            let (root, s3, d) = (root.clone(), s2.clone(), dropped.clone());
156            std::thread::spawn(move || {
157                let _ = handle(c, &root, &s3.log, &s3.faults, &d);
158            });
159        }
160    });
161    Ok(srv)
162}