clawhdf5/crates/clawhdf5-remote/tests/common/server.rs

//! A small HTTP/1.1 file server on 127.0.0.1 for tests and the example:
//! `Range: bytes=a-b` requests (206, 416), `HEAD`, keep-alive, strong ETags
//! and Last-Modified with `If-Match` / `If-Unmodified-Since` (412), and
//! switches to misbehave — ignore ranges (200 with the whole file), cut
//! bodies short, answer 503, respond slowly, send no validators. It counts
//! requests and body bytes, and logs every range asked for — only for the
//! paths it serves: a request for any other path (a local port scanner's
//! `GET /`, say) is answered 404 and not counted, so request budgets in
//! tests stay exact.

#![allow(dead_code)]

use std::collections::HashMap;
use std::io::{BufRead, BufReader, Write};
use std::net::{TcpListener, TcpStream};
use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
use std::sync::{Arc, Mutex, RwLock};
use std::time::Duration;

/// A logged GET: the path and the range asked for (`None`: whole file).
pub type LogEntry = (String, Option<(u64, u64)>);

struct Resource {
    data: Arc<Vec<u8>>,
    etag: String,
    last_modified: String,
}

/// Switches and counters shared with the connection threads.
#[derive(Default)]
pub struct Shared {
    files: RwLock<HashMap<String, Resource>>,
    version: AtomicU64,
    /// Answer every request with 200 and the whole file.
    pub ignore_range: AtomicBool,
    /// Send no ETag and no Last-Modified.
    pub no_validators: AtomicBool,
    /// Send a weak ETag only (no Last-Modified).
    pub weak_etag: AtomicBool,
    /// Send Last-Modified but no ETag.
    pub no_etag: AtomicBool,
    /// Label bodies `Content-Encoding: gzip` (they are not).
    pub gzip_label: AtomicBool,
    /// Cut the body of the next N ranged responses in half (then close the
    /// connection).
    pub truncate_next: AtomicU32,
    /// Answer the next N requests with 503.
    pub fail_next: AtomicU32,
    /// Sleep this long before answering each request.
    pub delay_ms: AtomicU64,
    /// When non-zero, claim the file is this long (in `Content-Range` and
    /// when checking ranges) and serve zeros past its real end: a hostile
    /// server lying about the length.
    pub fake_total: AtomicU64,
    /// When non-zero, answer every request for a served file with this
    /// status (and an empty body).
    pub force_status: AtomicU32,
    /// When non-zero, send bodies at about this many bytes per second.
    pub throttle_bps: AtomicU64,
    /// When non-zero, stop this many milliseconds after the headers and
    /// half the body (a stalled connection).
    pub stall_ms: AtomicU64,
    /// Answer ranges with a `Content-Range` one byte off.
    pub wrong_range: AtomicBool,
    /// Path → `Location`: answered `302 Found` (counted as a request).
    pub redirects: RwLock<HashMap<String, String>>,
    /// (path, headers) of every counted request, header names lowercase.
    pub seen: Mutex<Vec<(String, HashMap<String, String>)>>,
    /// Requests for a served path (every status); requests for other
    /// paths are not counted.
    pub requests: AtomicU64,
    /// Body bytes sent.
    pub bytes: AtomicU64,
    /// (path, range) of every GET.
    pub log: Mutex<Vec<LogEntry>>,
    stop: AtomicBool,
}

/// A running server; stops accepting when dropped.
pub struct Server {
    pub addr: std::net::SocketAddr,
    pub shared: Arc<Shared>,
}

impl Server {
    /// Serve `files` (URL path such as `/a.h5` → bytes) on `127.0.0.1`, on
    /// a free port.
    pub fn start(files: Vec<(String, Vec<u8>)>) -> Server {
        Server::bind("127.0.0.1:0", files)
    }

    /// Serve `files` on `addr`.
    pub fn bind(addr: &str, files: Vec<(String, Vec<u8>)>) -> Server {
        let listener = TcpListener::bind(addr).expect("bind");
        let addr = listener.local_addr().unwrap();
        let shared = Arc::new(Shared::default());
        for (path, data) in files {
            shared.put(&path, data);
        }
        let s = shared.clone();
        std::thread::spawn(move || {
            for conn in listener.incoming() {
                if s.stop.load(Ordering::SeqCst) {
                    break;
                }
                let Ok(conn) = conn else { continue };
                let s = s.clone();
                std::thread::spawn(move || {
                    let _ = serve(conn, &s);
                });
            }
        });
        Server { addr, shared }
    }

    /// The URL of `path` (which starts with `/`).
    pub fn url(&self, path: &str) -> String {
        format!("http://{}{path}", self.addr)
    }

    /// Requests served so far.
    pub fn requests(&self) -> u64 {
        self.shared.requests.load(Ordering::SeqCst)
    }

    /// Body bytes sent so far.
    pub fn bytes(&self) -> u64 {
        self.shared.bytes.load(Ordering::SeqCst)
    }

    /// Zero the counters and the logs.
    pub fn reset(&self) {
        self.shared.requests.store(0, Ordering::SeqCst);
        self.shared.bytes.store(0, Ordering::SeqCst);
        self.shared.log.lock().unwrap().clear();
        self.shared.seen.lock().unwrap().clear();
    }

    /// Answer `path` with a redirect to `location`.
    pub fn redirect(&self, path: &str, location: &str) {
        self.shared
            .redirects
            .write()
            .unwrap()
            .insert(path.to_string(), location.to_string());
    }

    /// Whether any request the server counted carried header `name`.
    pub fn saw_header(&self, name: &str) -> bool {
        self.shared
            .seen
            .lock()
            .unwrap()
            .iter()
            .any(|(_, h)| h.contains_key(name))
    }

    /// The ranges asked for so far.
    pub fn log(&self) -> Vec<LogEntry> {
        self.shared.log.lock().unwrap().clone()
    }
}

impl Drop for Server {
    fn drop(&mut self) {
        self.shared.stop.store(true, Ordering::SeqCst);
        let _ = TcpStream::connect(self.addr);
    }
}

impl Shared {
    /// Add or replace a file (a replacement gets a new ETag and
    /// Last-Modified).
    pub fn put(&self, path: &str, data: Vec<u8>) {
        let v = self.version.fetch_add(1, Ordering::SeqCst) + 1;
        let secs = 1_700_000_000 + v;
        self.files.write().unwrap().insert(
            path.to_string(),
            Resource {
                data: Arc::new(data),
                etag: format!("\"v{v}-{path}\""),
                last_modified: http_date(secs),
            },
        );
    }
}

/// An IMF-fixdate for `secs` since the epoch (enough for distinct values).
fn http_date(secs: u64) -> String {
    const DAYS: [&str; 7] = ["Thu", "Fri", "Sat", "Sun", "Mon", "Tue", "Wed"];
    const MONTHS: [&str; 12] = [
        "Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec",
    ];
    let days = secs / 86_400;
    let rem = secs % 86_400;
    // Civil-from-days (Howard Hinnant).
    let z = days as i64 + 719_468;
    let era = z.div_euclid(146_097);
    let doe = z - era * 146_097;
    let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
    let y = yoe + era * 400;
    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
    let mp = (5 * doy + 2) / 153;
    let d = doy - (153 * mp + 2) / 5 + 1;
    let m = if mp < 10 { mp + 3 } else { mp - 9 };
    let y = if m <= 2 { y + 1 } else { y };
    format!(
        "{}, {:02} {} {} {:02}:{:02}:{:02} GMT",
        DAYS[(days % 7) as usize],
        d,
        MONTHS[(m - 1) as usize],
        y,
        rem / 3600,
        rem % 3600 / 60,
        rem % 60
    )
}

fn parse_range(v: &str, len: u64) -> Option<Result<(u64, u64), ()>> {
    let spec = v.trim().strip_prefix("bytes=")?;
    if spec.contains(',') {
        return None; // multiple ranges: serve the whole file
    }
    let (a, b) = spec.split_once('-')?;
    let a: u64 = a.trim().parse().ok()?;
    let b: u64 = match b.trim() {
        "" => u64::MAX,
        b => b.parse().ok()?,
    };
    if b < a {
        return None;
    }
    if a >= len {
        return Some(Err(()));
    }
    Some(Ok((a, b.min(len - 1))))
}

fn serve(conn: TcpStream, s: &Shared) -> std::io::Result<()> {
    conn.set_read_timeout(Some(Duration::from_secs(30)))?;
    // One write per response and no Nagle delay: otherwise every request
    // waits for a delayed ACK (~40 ms).
    conn.set_nodelay(true)?;
    let mut reader = BufReader::new(conn.try_clone()?);
    let mut out = conn;
    loop {
        let mut line = String::new();
        if reader.read_line(&mut line)? == 0 {
            return Ok(());
        }
        let mut parts = line.split_whitespace();
        let method = parts.next().unwrap_or("").to_string();
        let path = parts.next().unwrap_or("").to_string();
        let mut headers = HashMap::new();
        loop {
            let mut h = String::new();
            if reader.read_line(&mut h)? == 0 {
                return Ok(());
            }
            let h = h.trim_end();
            if h.is_empty() {
                break;
            }
            if let Some((k, v)) = h.split_once(':') {
                headers.insert(k.trim().to_ascii_lowercase(), v.trim().to_string());
            }
        }
        // The query string (a presigned URL's signature, say) is not part
        // of the file's name.
        let path = path.split('?').next().unwrap_or("").to_string();
        let close = headers
            .get("connection")
            .is_some_and(|v| v.eq_ignore_ascii_case("close"));
        let redirect = s.redirects.read().unwrap().get(&path).cloned();
        if let Some(location) = redirect {
            s.requests.fetch_add(1, Ordering::SeqCst);
            s.seen.lock().unwrap().push((path.clone(), headers.clone()));
            write!(
                out,
                "HTTP/1.1 302 Found\r\nLocation: {location}\r\nContent-Length: 0\r\n\r\n"
            )?;
            continue;
        }
        let res = {
            let files = s.files.read().unwrap();
            files
                .get(&path)
                .map(|r| (r.data.clone(), r.etag.clone(), r.last_modified.clone()))
        };
        let Some((data, etag, lm)) = res else {
            // Not ours: not counted, not delayed, no failure injected.
            write!(out, "HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\n\r\n")?;
            if close {
                return Ok(());
            }
            continue;
        };
        s.requests.fetch_add(1, Ordering::SeqCst);
        s.seen.lock().unwrap().push((path.clone(), headers.clone()));
        let delay = s.delay_ms.load(Ordering::SeqCst);
        if delay > 0 {
            std::thread::sleep(Duration::from_millis(delay));
        }
        let forced = s.force_status.load(Ordering::SeqCst);
        if forced != 0 {
            write!(out, "HTTP/1.1 {forced} Forced\r\nContent-Length: 0\r\n\r\n")?;
            continue;
        }
        if s.fail_next
            .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |n| n.checked_sub(1))
            .is_ok()
        {
            write!(
                out,
                "HTTP/1.1 503 Service Unavailable\r\nContent-Length: 0\r\n\r\n"
            )?;
            continue;
        }
        let fake = s.fake_total.load(Ordering::SeqCst);
        let len = if fake > 0 { fake } else { data.len() as u64 };
        let mut validators = String::new();
        if s.weak_etag.load(Ordering::SeqCst) {
            validators = format!("ETag: W/{etag}\r\n");
        } else if s.no_etag.load(Ordering::SeqCst) {
            validators = format!("Last-Modified: {lm}\r\n");
        } else if !s.no_validators.load(Ordering::SeqCst) {
            validators = format!("ETag: {etag}\r\nLast-Modified: {lm}\r\n");
        }
        let precondition_failed = headers.get("if-match").is_some_and(|v| v != &etag)
            || headers.get("if-unmodified-since").is_some_and(|v| v != &lm);
        if precondition_failed {
            write!(
                out,
                "HTTP/1.1 412 Precondition Failed\r\n{validators}Content-Length: 0\r\n\r\n"
            )?;
            continue;
        }
        let range = if s.ignore_range.load(Ordering::SeqCst) {
            None
        } else {
            headers.get("range").and_then(|v| parse_range(v, len))
        };
        if method == "GET" {
            s.log
                .lock()
                .unwrap()
                .push((path.clone(), range.and_then(Result::ok)));
        }
        let mut padded = Vec::new();
        let (status, body, extra) = match range {
            Some(Err(())) => {
                write!(
                    out,
                    "HTTP/1.1 416 Range Not Satisfiable\r\nContent-Range: bytes */{len}\r\n\
                     {validators}Content-Length: 0\r\n\r\n"
                )?;
                continue;
            }
            Some(Ok((a, b))) => (
                "206 Partial Content",
                slice_or_zeros(&data, a, b, &mut padded),
                if s.wrong_range.load(Ordering::SeqCst) {
                    format!("Content-Range: bytes {}-{}/{len}\r\n", a + 1, b + 1)
                } else {
                    format!("Content-Range: bytes {a}-{b}/{len}\r\n")
                },
            ),
            None => ("200 OK", &data[..], String::new()),
        };
        let extra = if s.gzip_label.load(Ordering::SeqCst) {
            format!("{extra}Content-Encoding: gzip\r\n")
        } else {
            extra
        };
        let head = format!(
            "HTTP/1.1 {status}\r\nContent-Length: {}\r\nAccept-Ranges: bytes\r\n{extra}{validators}\r\n",
            body.len()
        );
        if method == "HEAD" {
            out.write_all(head.as_bytes())?;
            continue;
        }
        let truncate = status.starts_with("206")
            && s.truncate_next
                .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |n| n.checked_sub(1))
                .is_ok();
        let body = if truncate {
            &body[..body.len() / 2]
        } else {
            body
        };
        // Counted before the client can have the bytes, so a test that
        // resets the counters after a read never sees them arrive late.
        s.bytes.fetch_add(body.len() as u64, Ordering::SeqCst);
        let bps = s.throttle_bps.load(Ordering::SeqCst);
        let stall = s.stall_ms.load(Ordering::SeqCst);
        if bps > 0 || stall > 0 {
            out.write_all(head.as_bytes())?;
            let (first, rest) = body.split_at(if stall > 0 { body.len() / 2 } else { 0 });
            out.write_all(first)?;
            out.flush()?;
            if stall > 0 {
                std::thread::sleep(Duration::from_millis(stall));
            }
            for piece in rest.chunks(4096) {
                out.write_all(piece)?;
                out.flush()?;
                if bps > 0 {
                    std::thread::sleep(Duration::from_micros(4096 * 1_000_000 / bps));
                }
            }
        } else {
            let mut response = head.into_bytes();
            response.extend_from_slice(body);
            out.write_all(&response)?;
            out.flush()?;
        }
        if truncate || close {
            let _ = out.shutdown(std::net::Shutdown::Both);
            return Ok(());
        }
    }
}

/// `data[a..=b]`, padded with zeros past its end (into `padded`).
fn slice_or_zeros<'a>(data: &'a [u8], a: u64, b: u64, padded: &'a mut Vec<u8>) -> &'a [u8] {
    let real = data.len() as u64;
    if b < real {
        return &data[a as usize..=b as usize];
    }
    let n = usize::try_from(b - a + 1).expect("range fits in memory");
    padded.resize(n, 0);
    if a < real {
        let have = (real - a) as usize;
        padded[..have].copy_from_slice(&data[a as usize..]);
    }
    padded
}