uv_test/packse/

server.rs

//! A per-test wiremock server that serves packse scenario packages.
//!
//! Each [`PackseServer`] reads a single scenario TOML file and serves:
//! - PEP 691 Simple API at `/simple/{package}/`
//! - Distribution downloads at `/files/{filename}`
//!
//! Cached build dependencies are exposed through the same `/simple/*` and
//! `/files/*` routes as scenario packages.

use std::collections::HashMap;
use std::path::Path;
use std::str::FromStr;
use std::sync::Arc;

use serde_json::json;
use wiremock::{
    Mock, MockServer, Request, ResponseTemplate,
    matchers::{header_exists, method, path},
};

use uv_distribution_filename::WheelFilename;
use uv_normalize::PackageName;
use uv_pep440::VersionSpecifiers;

use crate::http_server::{HttpServer, content_type_for_filename};
use crate::vendor::{VendorArtifact, vendor_artifacts};

use super::scenario::{Scenario, WheelTag};
use super::scenarios_dir;
use super::wheel::{generate_sdist, generate_wheel, sha256_hex};

const PACKSE_UPLOAD_TIME: &str = "2024-03-24T00:00:00Z";

/// Information about a single distribution file (metadata for the Simple API).
struct DistInfo {
    filename: String,
    sha256: String,
    requires_python: Option<VersionSpecifiers>,
    upload_time: &'static str,
    yanked: bool,
}

/// All distributions for a given package name, across versions.
struct PackageEntry {
    dists: Vec<DistInfo>,
}

enum FileData {
    Bytes(Arc<[u8]>),
    Vendor(&'static VendorArtifact),
}

impl FileData {
    fn bytes(&self) -> anyhow::Result<Arc<[u8]>> {
        match self {
            Self::Bytes(bytes) => Ok(Arc::clone(bytes)),
            Self::Vendor(artifact) => artifact.bytes(),
        }
    }
}

/// The complete pre-indexed database for a server instance.
struct ServerIndex {
    /// Simple API: normalized package name → list of distribution metadata.
    packages: HashMap<PackageName, PackageEntry>,
    /// File downloads: filename → generated bytes or a lazy vendored artifact.
    files: HashMap<String, FileData>,
}

/// A running mock PyPI server for a single packse scenario.
///
/// The server runs on a background thread with its own single-threaded tokio runtime.
/// When [`PackseServer`] is dropped, the background thread and server are shut down.
pub struct PackseServer {
    server: HttpServer,
}

impl PackseServer {
    /// Load a scenario from a TOML path (relative to the vendored scenarios directory)
    /// and start a mock server for it.
    pub fn new(scenario_path: &str) -> Self {
        let full_path = scenarios_dir().join(scenario_path);
        let scenario =
            Scenario::from_path(&full_path).expect("vendored Packse scenario should parse");
        Self::from_scenario(&scenario)
    }

    /// Start a mock server with no packages (only cached build dependencies).
    ///
    /// Useful as a dummy index that will 404 for any non-cached package lookup.
    pub fn empty() -> Self {
        Self::from_scenario(&Scenario::empty())
    }

    /// Start a mock server for the given scenario.
    pub fn from_scenario(scenario: &Scenario) -> Self {
        let index = Arc::new(build_server_index(scenario));
        let server = HttpServer::start(move |request, server_uri| {
            handle_request(request, server_uri, &index)
        });

        Self { server }
    }

    /// The Simple API index URL (e.g., `http://127.0.0.1:PORT/simple/`).
    pub fn index_url(&self) -> String {
        format!("{}/simple/", self.server.url())
    }

    /// Return the URL for a generated distribution file.
    pub fn file_url(&self, filename: &str) -> String {
        format!("{}/files/{filename}", self.server.url())
    }
}

/// Build the complete [`ServerIndex`] from a scenario and cached build dependencies.
fn build_server_index(scenario: &Scenario) -> ServerIndex {
    let mut packages = HashMap::new();
    let mut files: HashMap<String, FileData> = HashMap::new();

    for (package_name, package) in &scenario.packages {
        let mut dists = Vec::new();

        for (version, meta) in &package.versions {
            if meta.wheel {
                let tags = if meta.wheel_tags.is_empty() {
                    vec!["py3-none-any"]
                } else {
                    meta.wheel_tags.iter().map(WheelTag::as_str).collect()
                };

                for tag in tags {
                    let (filename, bytes) = generate_wheel(
                        package_name,
                        version,
                        &meta.requires,
                        &meta.extras,
                        meta.requires_python.as_ref(),
                        tag,
                    );
                    let sha256 = sha256_hex(&bytes);
                    files.insert(filename.clone(), FileData::Bytes(bytes.into()));
                    dists.push(DistInfo {
                        filename,
                        sha256,
                        requires_python: meta.requires_python.clone(),
                        upload_time: PACKSE_UPLOAD_TIME,
                        yanked: meta.yanked,
                    });
                }
            }

            if meta.sdist {
                let (filename, bytes) = generate_sdist(
                    package_name,
                    version,
                    &meta.requires,
                    &meta.extras,
                    meta.requires_python.as_ref(),
                );
                let sha256 = sha256_hex(&bytes);
                files.insert(filename.clone(), FileData::Bytes(bytes.into()));
                dists.push(DistInfo {
                    filename,
                    sha256,
                    requires_python: meta.requires_python.clone(),
                    upload_time: PACKSE_UPLOAD_TIME,
                    yanked: meta.yanked,
                });
            }
        }

        packages.insert(package_name.clone(), PackageEntry { dists });
    }

    for artifact in vendor_artifacts() {
        if !Path::new(artifact.filename)
            .extension()
            .is_some_and(|extension| extension.eq_ignore_ascii_case("whl"))
        {
            continue;
        }

        let wheel_filename =
            WheelFilename::from_str(artifact.filename).expect("invalid vendor wheel filename");

        files.insert(artifact.filename.to_string(), FileData::Vendor(artifact));
        packages
            .entry(wheel_filename.name)
            .or_insert_with(|| PackageEntry { dists: Vec::new() })
            .dists
            .push(DistInfo {
                filename: artifact.filename.to_string(),
                sha256: artifact.sha256.to_string(),
                requires_python: None,
                upload_time: PACKSE_UPLOAD_TIME,
                yanked: false,
            });
    }

    ServerIndex { packages, files }
}

fn handle_request(req: &Request, server_uri: &str, index: &ServerIndex) -> ResponseTemplate {
    let path = req.url.path();

    if let Some(pkg) = extract_package_name(path) {
        let Ok(package_name) = PackageName::from_str(pkg) else {
            return ResponseTemplate::new(404);
        };

        if let Some(entry) = index.packages.get(&package_name) {
            return build_simple_api_response(pkg, entry, server_uri);
        }
        return ResponseTemplate::new(404);
    }

    if let Some(filename) = path.strip_prefix("/files/") {
        if let Some(file) = index.files.get(filename) {
            return match file.bytes() {
                Ok(bytes) => distribution_file_response(req, filename, &bytes),
                Err(error) => ResponseTemplate::new(500).set_body_string(format!("{error:#}")),
            };
        }
        return ResponseTemplate::new(404);
    }

    ResponseTemplate::new(404)
}

/// Build a response for a distribution file, including support for single byte ranges.
pub fn distribution_file_response(req: &Request, filename: &str, bytes: &[u8]) -> ResponseTemplate {
    let content_type = content_type_for_filename(filename);
    let Some(range) = req.headers.get("range") else {
        return ResponseTemplate::new(200)
            .insert_header("Accept-Ranges", "bytes")
            .set_body_raw(bytes.to_vec(), content_type);
    };

    let Some((start, end)) = range
        .to_str()
        .ok()
        .and_then(|range| parse_byte_range(range, bytes.len()))
    else {
        return ResponseTemplate::new(416)
            .insert_header("Accept-Ranges", "bytes")
            .insert_header("Content-Range", format!("bytes */{}", bytes.len()));
    };

    ResponseTemplate::new(206)
        .insert_header("Accept-Ranges", "bytes")
        .insert_header(
            "Content-Range",
            format!("bytes {start}-{end}/{}", bytes.len()),
        )
        .set_body_raw(bytes[start..=end].to_vec(), content_type)
}

/// Mount a distribution that serves different bytes to range and full-file requests.
///
/// `HEAD` responses advertise range support. Ranged `GET` requests receive the ranged bytes,
/// while full `GET` requests receive the full-file bytes instead.
pub async fn mount_mismatched_distribution(
    server: &MockServer,
    file_path: &str,
    filename: &str,
    ranged_bytes: Vec<u8>,
    full_bytes: Vec<u8>,
) {
    Mock::given(method("HEAD"))
        .and(path(file_path))
        .respond_with(
            ResponseTemplate::new(200)
                .insert_header("Accept-Ranges", "bytes")
                .set_body_bytes(ranged_bytes.clone()),
        )
        .mount(server)
        .await;

    let filename = filename.to_string();
    Mock::given(method("GET"))
        .and(path(file_path))
        .and(header_exists("range"))
        .respond_with(move |request: &Request| {
            distribution_file_response(request, &filename, &ranged_bytes)
        })
        .with_priority(1)
        .expect(1..)
        .mount(server)
        .await;

    Mock::given(method("GET"))
        .and(path(file_path))
        .respond_with(ResponseTemplate::new(200).set_body_bytes(full_bytes))
        .mount(server)
        .await;
}

/// Parse a single HTTP byte range and return its inclusive bounds.
fn parse_byte_range(range: &str, length: usize) -> Option<(usize, usize)> {
    let range = range.strip_prefix("bytes=")?;
    let (start, end) = range.split_once('-')?;

    if start.is_empty() {
        let suffix = end.parse::<usize>().ok()?;
        if suffix == 0 || length == 0 {
            return None;
        }
        return Some((length.saturating_sub(suffix), length - 1));
    }

    let start = start.parse::<usize>().ok()?;
    if start >= length {
        return None;
    }
    let end = if end.is_empty() {
        length - 1
    } else {
        end.parse::<usize>().ok()?.min(length - 1)
    };
    if start > end {
        return None;
    }
    Some((start, end))
}

/// Build PEP 691 JSON response for a package.
fn build_simple_api_response(
    package_name: &str,
    entry: &PackageEntry,
    server_uri: &str,
) -> ResponseTemplate {
    let files: Vec<serde_json::Value> = entry
        .dists
        .iter()
        .map(|dist| {
            let url = format!("{server_uri}/files/{}", dist.filename);
            let mut file_obj = json!({
                "filename": dist.filename,
                "url": url,
                "hashes": {
                    "sha256": dist.sha256,
                },
                "upload-time": dist.upload_time,
            });
            if let Some(rp) = &dist.requires_python {
                file_obj["requires-python"] = json!(rp);
            }
            if dist.yanked {
                file_obj["yanked"] = json!(true);
            }
            file_obj
        })
        .collect();

    let body = json!({
        "meta": { "api-version": "1.1" },
        "name": package_name,
        "files": files,
    });

    let body_str = body.to_string();
    ResponseTemplate::new(200)
        .insert_header("Content-Type", "application/vnd.pypi.simple.v1+json")
        .set_body_raw(body_str, "application/vnd.pypi.simple.v1+json")
}

/// Extract the package name from `/simple/{package}` or `/simple/{package}/`.
fn extract_package_name(path: &str) -> Option<&str> {
    let rest = path.strip_prefix("/simple/")?;
    let pkg = rest.strip_suffix('/').unwrap_or(rest);
    if pkg.is_empty() || pkg.contains('/') {
        return None;
    }
    Some(pkg)
}

#[cfg(test)]
mod tests {
    use anyhow::Result;
    use reqwest::StatusCode;
    use reqwest::header::{ACCEPT_RANGES, CONTENT_RANGE, RANGE};
    use wiremock::MockServer;

    use crate::vendor::vendor_artifacts;

    use super::{
        PackseServer, Scenario, build_server_index, extract_package_name,
        mount_mismatched_distribution,
    };

    #[test]
    fn extract_package_name_accepts_with_or_without_trailing_slash() {
        assert_eq!(extract_package_name("/simple/foo/"), Some("foo"));
        assert_eq!(extract_package_name("/simple/foo"), Some("foo"));
    }

    #[test]
    fn extract_package_name_rejects_invalid_paths() {
        assert_eq!(extract_package_name("/simple/"), None);
        assert_eq!(extract_package_name("/simple"), None);
        assert_eq!(extract_package_name("/simple/foo/bar"), None);
    }

    #[test]
    fn server_index_construction_does_not_load_vendor_artifacts() {
        let _index = build_server_index(&Scenario::empty());

        assert!(
            vendor_artifacts()
                .iter()
                .all(|artifact| !artifact.is_loaded())
        );
    }

    #[tokio::test]
    async fn file_requests_support_byte_ranges() -> Result<()> {
        let scenario = toml::from_str::<Scenario>(
            r#"
name = "range-requests"

[root]
requires = ["a"]

[expected]
satisfiable = true

[packages.a.versions."1.0.0"]
sdist = false
"#,
        )?;
        let server = PackseServer::from_scenario(&scenario);
        let url = server.file_url("a-1.0.0-py3-none-any.whl");
        let client = reqwest::Client::new();

        let response = client.get(&url).send().await?;
        assert_eq!(response.status(), StatusCode::OK);
        assert_eq!(
            response.headers().get(ACCEPT_RANGES),
            Some(&"bytes".parse()?)
        );
        let bytes = response.bytes().await?;
        let length = bytes.len();

        let response = client.get(&url).header(RANGE, "bytes=-8").send().await?;
        assert_eq!(response.status(), StatusCode::PARTIAL_CONTENT);
        assert_eq!(
            response.headers().get(CONTENT_RANGE),
            Some(&format!("bytes {}-{}/{}", length - 8, length - 1, length).parse()?)
        );
        assert_eq!(response.bytes().await?, bytes[length - 8..]);

        let response = client
            .get(&url)
            .header(RANGE, "bytes=-999999")
            .send()
            .await?;
        assert_eq!(response.status(), StatusCode::PARTIAL_CONTENT);
        assert_eq!(
            response.headers().get(CONTENT_RANGE),
            Some(&format!("bytes 0-{}/{}", length - 1, length).parse()?)
        );
        assert_eq!(response.bytes().await?, bytes);

        let response = client.get(&url).header(RANGE, "bytes=3-9").send().await?;
        assert_eq!(response.status(), StatusCode::PARTIAL_CONTENT);
        assert_eq!(
            response.headers().get(CONTENT_RANGE),
            Some(&format!("bytes 3-9/{length}").parse()?)
        );
        assert_eq!(response.bytes().await?, bytes[3..=9]);

        let response = client.get(&url).header(RANGE, "bytes=10-").send().await?;
        assert_eq!(response.status(), StatusCode::PARTIAL_CONTENT);
        assert_eq!(
            response.headers().get(CONTENT_RANGE),
            Some(&format!("bytes 10-{}/{}", length - 1, length).parse()?)
        );
        assert_eq!(response.bytes().await?, bytes[10..]);

        let response = client
            .get(&url)
            .header(RANGE, format!("bytes={length}-"))
            .send()
            .await?;
        assert_eq!(response.status(), StatusCode::RANGE_NOT_SATISFIABLE);
        assert_eq!(
            response.headers().get(CONTENT_RANGE),
            Some(&format!("bytes */{length}").parse()?)
        );

        Ok(())
    }

    #[tokio::test]
    async fn mismatched_distribution_serves_ranged_and_full_bytes() -> Result<()> {
        let server = MockServer::start().await;
        mount_mismatched_distribution(
            &server,
            "/files/example.whl",
            "example.whl",
            b"forged".to_vec(),
            b"authentic".to_vec(),
        )
        .await;
        let url = format!("{}/files/example.whl", server.uri());
        let client = reqwest::Client::new();

        let response = client.head(&url).send().await?;
        assert_eq!(response.status(), StatusCode::OK);
        assert_eq!(
            response.headers().get(ACCEPT_RANGES),
            Some(&"bytes".parse()?)
        );

        let response = client.get(&url).header(RANGE, "bytes=1-3").send().await?;
        assert_eq!(response.status(), StatusCode::PARTIAL_CONTENT);
        assert_eq!(response.bytes().await?, b"org".as_slice());

        let response = client.get(&url).send().await?;
        assert_eq!(response.status(), StatusCode::OK);
        assert_eq!(response.bytes().await?, b"authentic".as_slice());

        Ok(())
    }
}