zsync-rs 0.1.4

use std::io::Read;

use crate::control::ControlFile;

#[derive(Debug, thiserror::Error)]
pub enum HttpError {
    #[error("HTTP error: {0}")]
    Http(String),
    #[error("IO error: {0}")]
    Io(#[from] std::io::Error),
    #[error("Invalid URL: {0}")]
    InvalidUrl(String),
    #[error("No URLs available")]
    NoUrls,
}

pub struct HttpRangeReader {
    reader: Box<dyn std::io::Read + Send + Sync>,
}

impl std::io::Read for HttpRangeReader {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        self.reader.read(buf)
    }
}

/// HTTP transport for fetching control files and byte ranges.
pub struct HttpClient {
    agent: ureq::Agent,
}

impl Default for HttpClient {
    fn default() -> Self {
        Self::new()
    }
}

impl HttpClient {
    /// A client with ureq's default policy, permitting plain HTTP.
    ///
    /// zsync is routinely served over HTTP and across mirror redirects, so
    /// the defaults are permissive on purpose. An embedder with a stricter
    /// policy, such as one that must refuse a downgrade to HTTP or pin its
    /// own trust roots, should build an agent and pass it to
    /// [`HttpClient::with_agent`] rather than relying on these.
    pub fn new() -> Self {
        Self {
            agent: ureq::Agent::config_builder()
                .https_only(false)
                .build()
                .new_agent(),
        }
    }

    /// A client using a caller-supplied agent.
    ///
    /// The agent carries the whole transport policy: TLS roots, whether a
    /// redirect may downgrade to HTTP, timeouts and proxies. Supplying one
    /// keeps that decision with the embedder, which is the only place that
    /// knows what the transferred bytes are trusted for.
    pub fn with_agent(agent: ureq::Agent) -> Self {
        Self { agent }
    }

    pub fn fetch_control_file(&self, url: &str) -> Result<ControlFile, HttpError> {
        let response = self
            .agent
            .get(url)
            .call()
            .map_err(|e| HttpError::Http(e.to_string()))?;

        let mut reader = response.into_body().into_reader();
        ControlFile::parse(&mut reader).map_err(|e| HttpError::Http(e.to_string()))
    }

    /// A reader over `start..=end` of `url`.
    ///
    /// The reader is capped at the requested length. Without that cap the
    /// origin decides how much the client reads: a response longer than
    /// the range is consumed in full, and callers that accumulate it, such
    /// as block assembly, grow a buffer to match. Bytes past the range
    /// were never asked for and are of no use, so they are not read.
    pub fn fetch_range_reader(
        &self,
        url: &str,
        start: u64,
        end: u64,
    ) -> Result<HttpRangeReader, HttpError> {
        let range_header = format!("bytes={}-{}", start, end);
        let requested = end.saturating_sub(start).saturating_add(1);

        let response = self
            .agent
            .get(url)
            .header("Range", &range_header)
            .call()
            .map_err(|e| HttpError::Http(e.to_string()))?;

        let status = response.status();
        if status != 206 && status != 200 {
            return Err(HttpError::Http(format!(
                "Expected 206 Partial Content, got {}",
                status
            )));
        }

        Ok(HttpRangeReader {
            reader: Box::new(response.into_body().into_reader().take(requested)),
        })
    }

    pub fn fetch_range(&self, url: &str, start: u64, end: u64) -> Result<Vec<u8>, HttpError> {
        let mut reader = self.fetch_range_reader(url, start, end)?;
        let mut buf = Vec::new();
        reader.read_to_end(&mut buf)?;
        Ok(buf)
    }

    pub fn fetch_ranges(
        &self,
        url: &str,
        ranges: &[(u64, u64)],
        blocksize: usize,
    ) -> Result<Vec<(u64, Vec<u8>)>, HttpError> {
        let mut results = Vec::new();

        for &(start, end) in ranges {
            let data = self.fetch_range(url, start, end)?;
            let aligned_start = (start / blocksize as u64) * blocksize as u64;
            results.push((aligned_start, data));
        }

        Ok(results)
    }
}

/// Default gap threshold for range merging (256 KiB, same as zsync2).
pub const DEFAULT_RANGE_GAP_THRESHOLD: u64 = 256 * 1024;

/// Merge byte ranges to minimize HTTP requests.
/// Gaps smaller than the threshold are merged to save HTTP round-trips.
pub fn merge_byte_ranges(ranges: &[(u64, u64)], gap_threshold: u64) -> Vec<(u64, u64)> {
    if ranges.len() <= 1 {
        return ranges.to_vec();
    }

    let mut merged = vec![ranges[0]];
    for &(start, end) in &ranges[1..] {
        let last = merged.last_mut().unwrap();
        let gap = start.saturating_sub(last.1 + 1);
        if gap <= gap_threshold {
            last.1 = end;
        } else {
            merged.push((start, end));
        }
    }
    merged
}

pub fn byte_ranges_from_block_ranges(
    block_ranges: &[(usize, usize)],
    blocksize: usize,
    file_length: u64,
) -> Vec<(u64, u64)> {
    block_ranges
        .iter()
        .map(|&(start_block, end_block)| {
            let start = start_block as u64 * blocksize as u64;
            let end =
                ((end_block as u64 * blocksize as u64).saturating_sub(1)).min(file_length - 1);
            (start, end)
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    /// A server that answers any request with far more data than asked
    /// for, which is what a hostile or broken origin does.
    fn overlong_server(body_len: usize) -> String {
        use std::io::Write;
        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();
        std::thread::spawn(move || {
            if let Ok((mut sock, _)) = listener.accept() {
                let mut req = [0u8; 1024];
                let _ = std::io::Read::read(&mut sock, &mut req);
                let hdr = format!(
                    "HTTP/1.1 206 Partial Content\r\nContent-Range: bytes 0-1023/{}\r\nContent-Length: {}\r\n\r\n",
                    body_len, body_len
                );
                let _ = sock.write_all(hdr.as_bytes());
                let chunk = vec![0u8; 64 * 1024];
                let mut sent = 0;
                while sent < body_len {
                    let n = chunk.len().min(body_len - sent);
                    if sock.write_all(&chunk[..n]).is_err() {
                        break;
                    }
                    sent += n;
                }
            }
        });
        format!("http://{}/f", addr)
    }

    #[test]
    fn a_range_response_is_capped_at_what_was_requested() {
        // 1 KiB asked for, 8 MiB offered. Accepting the surplus lets any
        // origin decide how much memory the client spends.
        let url = overlong_server(8 * 1024 * 1024);
        let client = HttpClient::new();
        let data = client.fetch_range(&url, 0, 1023).expect("fetch");
        assert_eq!(
            data.len(),
            1024,
            "a range response must be capped at the requested length"
        );
    }

    #[test]
    fn a_supplied_agent_carries_its_own_policy() {
        // The point of `with_agent` is that the embedder's policy reaches
        // the transfer. A client built with an HTTPS-only agent must refuse
        // a plain-HTTP URL, where the default client would allow it.
        let strict = HttpClient::with_agent(
            ureq::Agent::config_builder()
                .https_only(true)
                .build()
                .new_agent(),
        );
        let err = strict
            .fetch_control_file("http://127.0.0.1:1/nothing.zsync")
            .expect_err("an HTTPS-only agent must refuse a plain-HTTP URL");
        let msg = err.to_string();
        assert!(
            msg.to_lowercase().contains("http"),
            "expected a scheme rejection, got: {msg}"
        );
    }

    #[test]
    fn test_byte_ranges_from_block_ranges() {
        let block_ranges = vec![(0, 2), (4, 6)];
        let byte_ranges = byte_ranges_from_block_ranges(&block_ranges, 1024, 10000);
        assert_eq!(byte_ranges, vec![(0, 2047), (4096, 6143)]);
    }

    #[test]
    fn test_byte_ranges_clamped_to_file_length() {
        let block_ranges = vec![(9, 10)];
        let byte_ranges = byte_ranges_from_block_ranges(&block_ranges, 1024, 9500);
        assert_eq!(byte_ranges, vec![(9216, 9499)]);
    }
}