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.

#![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,
    /// Requests served (every status).
    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 log.
    pub fn reset(&self) {
        self.shared.requests.store(0, Ordering::SeqCst);
        self.shared.bytes.store(0, Ordering::SeqCst);
        self.shared.log.lock().unwrap().clear();
    }

    /// 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());
            }
        }
        s.requests.fetch_add(1, Ordering::SeqCst);
        let delay = s.delay_ms.load(Ordering::SeqCst);
        if delay > 0 {
            std::thread::sleep(Duration::from_millis(delay));
        }
        let close = headers
            .get("connection")
            .is_some_and(|v| v.eq_ignore_ascii_case("close"));
        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 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 {
            write!(out, "HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\n\r\n")?;
            continue;
        };
        let len = 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 (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",
                &data[a as usize..=b as usize],
                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
        };
        let mut response = head.into_bytes();
        response.extend_from_slice(body);
        out.write_all(&response)?;
        out.flush()?;
        s.bytes.fetch_add(body.len() as u64, Ordering::SeqCst);
        if truncate || close {
            let _ = out.shutdown(std::net::Shutdown::Both);
            return Ok(());
        }
    }
}