article_scraper 3.0.0-alpha1

//! A minimal HTTP server on 127.0.0.1 for tests that need a response the web won't give on
//! request (redirect loops, lying headers, bodies that never end).

use reqwest::Url;
use std::io::{BufRead, BufReader, Write};
use std::net::{TcpListener, TcpStream};
#[cfg(feature = "image-downloader")]
use std::sync::Arc;

/// Starts the server and returns its base URL. `respond` gets the request line
/// (`GET /a HTTP/1.1`) and writes the whole response. After `max_requests` connections the
/// server closes, so a client that keeps requesting fails instead of hanging.
pub fn serve(max_requests: usize, respond: impl Fn(&str, &mut TcpStream) + Send + 'static) -> Url {
    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let base = Url::parse(&format!("http://{}/", listener.local_addr().unwrap())).unwrap();

    std::thread::spawn(move || {
        for stream in listener.incoming().take(max_requests) {
            let Ok(mut stream) = stream else { continue };
            let request_line = read_request(&stream);
            respond(&request_line, &mut stream);
        }
    });

    base
}

/// Like `serve`, but answers every connection on its own thread, so requests overlap.
#[cfg(feature = "image-downloader")]
pub fn serve_parallel(
    max_requests: usize,
    respond: impl Fn(&str, &mut TcpStream) + Send + Sync + 'static,
) -> Url {
    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let base = Url::parse(&format!("http://{}/", listener.local_addr().unwrap())).unwrap();
    let respond = Arc::new(respond);

    std::thread::spawn(move || {
        for stream in listener.incoming().take(max_requests) {
            let Ok(mut stream) = stream else { continue };
            let respond = respond.clone();
            std::thread::spawn(move || {
                let request_line = read_request(&stream);
                respond(&request_line, &mut stream);
            });
        }
    });

    base
}

/// Reads the request headers and returns the request line.
fn read_request(stream: &TcpStream) -> String {
    let mut reader = BufReader::new(stream.try_clone().unwrap());
    let mut request_line = String::new();
    _ = reader.read_line(&mut request_line);
    let mut line = String::new();
    while reader.read_line(&mut line).is_ok_and(|n| n > 2) {
        line.clear();
    }
    request_line
}

/// Writes a complete `200 OK` response.
pub fn write_response(stream: &mut TcpStream, content_type: &str, body: &str) {
    write_status_response(stream, "200 OK", content_type, body);
}

/// Writes a complete response with the given `status` (`404 Not Found`).
pub fn write_status_response(stream: &mut TcpStream, status: &str, content_type: &str, body: &str) {
    _ = write!(
        stream,
        "HTTP/1.1 {status}\r\nContent-Type: {content_type}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
        body.len()
    );
}

/// Writes a `200 OK` response whose `Content-Length` claims a terabyte, without a body.
pub fn write_huge_content_length(stream: &mut TcpStream, content_type: &str) {
    _ = write!(
        stream,
        "HTTP/1.1 200 OK\r\nContent-Type: {content_type}\r\nContent-Length: 1000000000000\r\nConnection: close\r\n\r\n"
    );
}

/// Writes a complete `200 OK` response with a binary body.
#[cfg(feature = "image-downloader")]
pub fn write_bytes_response(stream: &mut TcpStream, content_type: &str, body: &[u8]) {
    _ = write!(
        stream,
        "HTTP/1.1 200 OK\r\nContent-Type: {content_type}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
        body.len()
    );
    _ = stream.write_all(body);
}

/// Writes a `200 OK` response with a chunked body and no `Content-Length`.
#[cfg(feature = "image-downloader")]
pub fn write_chunked_response(stream: &mut TcpStream, content_type: &str, body: &[u8]) {
    _ = write!(
        stream,
        "HTTP/1.1 200 OK\r\nContent-Type: {content_type}\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n{:x}\r\n",
        body.len()
    );
    _ = stream.write_all(body);
    _ = stream.write_all(b"\r\n0\r\n\r\n");
}

/// Writes a chunked `200 OK` response that never ends, until the client hangs up.
pub fn write_endless_body(stream: &mut TcpStream, content_type: &str) {
    if write!(
        stream,
        "HTTP/1.1 200 OK\r\nContent-Type: {content_type}\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n"
    )
    .is_err()
    {
        return;
    }

    let chunk = vec![b'a'; 64 * 1024];
    while write!(stream, "{:x}\r\n", chunk.len()).is_ok()
        && stream.write_all(&chunk).is_ok()
        && stream.write_all(b"\r\n").is_ok()
    {}
}