oci-image 0.1.0

//! Anonymous registry pull (feature `pull`): manifest, platform pick, config
//! and layer blobs into a [`LayerStore`], every blob checked by digest.

use std::collections::HashMap;

use crate::{Digest, Error, ImageLayer, ImageRef, LayerStore, Result};

const MANIFEST_ACCEPT: &str = "application/vnd.oci.image.index.v1+json, \
    application/vnd.oci.image.manifest.v1+json, \
    application/vnd.docker.distribution.manifest.list.v2+json, \
    application/vnd.docker.distribution.manifest.v2+json";

/// What [`pull`] fetched.
#[derive(Debug, Clone)]
pub struct Pulled {
    pub reference: ImageRef,
    /// The config digest, Docker's image ID.
    pub id: Digest,
    /// The image config JSON.
    pub config: serde_json::Value,
    /// The layers, bottom first, with their diff_ids from the config. Each
    /// is in the store as an unpacked layer or as a blob; pass them to
    /// [`LayerStore::add_image`].
    pub layers: Vec<ImageLayer>,
}

impl Pulled {
    /// The layer blob digests, bottom first (for [`LayerStore::rootfs`]).
    pub fn blobs(&self) -> Vec<Digest> {
        self.layers.iter().map(|l| l.blob.clone()).collect()
    }
}

/// Pull `image` for this machine's platform into `store`, skipping layers it
/// already has unpacked and blobs it already holds. `progress` gets status
/// lines.
pub async fn pull(
    client: &reqwest::Client,
    image: &str,
    store: &LayerStore,
    progress: &mut (dyn FnMut(String) + Send),
) -> Result<Pulled> {
    let r = ImageRef::parse(image)?;
    progress(format!("Pulling {}", r.canonical()));
    let mut reg = Registry {
        client,
        image: &r,
        token: None,
    };
    let (mut manifest, mut media_type) = reg.manifest(&r.reference).await?;
    if is_index(&media_type, &manifest) {
        let digest = pick_platform(&manifest)?;
        (manifest, media_type) = reg.manifest(&digest).await?;
    }
    if media_type.contains("v1+prettyjws") || manifest.get("fsLayers").is_some() {
        return Err(Error::Invalid(format!(
            "schema 1 image manifests are not supported ({image})"
        )));
    }
    let id: Digest = manifest["config"]["digest"]
        .as_str()
        .ok_or_else(|| Error::Invalid("image manifest has no config".into()))?
        .parse()?;
    let config_bytes = reg.blob_bytes(&id).await?;
    let config: serde_json::Value = serde_json::from_slice(&config_bytes)
        .map_err(|e| Error::Invalid(format!("image config: {e}")))?;
    let diff_ids = config["rootfs"]["diff_ids"]
        .as_array()
        .into_iter()
        .flatten()
        .map(|d| d.as_str().unwrap_or_default().parse::<Digest>())
        .collect::<Result<Vec<_>>>()?;

    let entries = manifest["layers"].as_array().cloned().unwrap_or_default();
    if diff_ids.len() != entries.len() {
        return Err(Error::Invalid(format!(
            "{} layers in the manifest, {} diff_ids in the config",
            entries.len(),
            diff_ids.len()
        )));
    }
    let mut layers = Vec::with_capacity(entries.len());
    for (i, layer) in entries.iter().enumerate() {
        let digest: Digest = layer["digest"]
            .as_str()
            .ok_or_else(|| Error::Invalid("layer without digest".into()))?
            .parse()?;
        let diff_id = diff_ids[i].clone();
        // An unpacked layer needs no blob; a stored blob no download.
        if !store.has_layer(&diff_id) && !store.has_blob(&digest) {
            progress(format!("Layer {}/{} {}", i + 1, entries.len(), digest));
            reg.blob_to_store(&digest, store).await?;
        }
        layers.push(ImageLayer {
            blob: digest,
            diff_id,
        });
    }
    progress(format!("Pulled {} ({})", r.canonical(), id));
    Ok(Pulled {
        reference: r,
        id,
        config,
        layers,
    })
}

fn is_index(media_type: &str, manifest: &serde_json::Value) -> bool {
    media_type.contains("index")
        || media_type.contains("manifest.list")
        || manifest.get("manifests").is_some()
}

/// The manifest digest for this machine in an image index.
pub fn pick_platform(index: &serde_json::Value) -> Result<String> {
    let arch = match std::env::consts::ARCH {
        "x86_64" => "amd64",
        "aarch64" => "arm64",
        "powerpc64" => "ppc64le",
        other => other,
    };
    index["manifests"]
        .as_array()
        .into_iter()
        .flatten()
        .find(|m| m["platform"]["os"] == "linux" && m["platform"]["architecture"] == arch)
        .and_then(|m| m["digest"].as_str())
        .map(str::to_string)
        .ok_or_else(|| Error::Invalid(format!("no linux/{arch} image in the index")))
}

struct Registry<'a> {
    client: &'a reqwest::Client,
    image: &'a ImageRef,
    token: Option<String>,
}

impl Registry<'_> {
    async fn get(&mut self, url: &str, accept: &str) -> Result<reqwest::Response> {
        for _ in 0..2 {
            let mut req = self.client.get(url).header("Accept", accept);
            if let Some(t) = &self.token {
                req = req.bearer_auth(t);
            }
            let resp = req.send().await?;
            if resp.status() == reqwest::StatusCode::UNAUTHORIZED && self.token.is_none() {
                let challenge = resp
                    .headers()
                    .get("www-authenticate")
                    .and_then(|v| v.to_str().ok())
                    .unwrap_or("")
                    .to_string();
                self.token = Some(self.fetch_token(&challenge).await?);
                continue;
            }
            if !resp.status().is_success() {
                return Err(Error::Invalid(format!(
                    "registry returned {} for {url}",
                    resp.status()
                )));
            }
            return Ok(resp);
        }
        Err(Error::Invalid(format!("registry denied access to {url}")))
    }

    /// Anonymous bearer token for `WWW-Authenticate: Bearer realm=...`.
    async fn fetch_token(&self, challenge: &str) -> Result<String> {
        let params = parse_challenge(challenge)
            .ok_or_else(|| Error::Invalid(format!("unsupported auth challenge: {challenge}")))?;
        let realm = params
            .get("realm")
            .ok_or_else(|| Error::Invalid("auth challenge has no realm".into()))?;
        let mut query: Vec<(&str, String)> = Vec::new();
        if let Some(s) = params.get("service") {
            query.push(("service", s.clone()));
        }
        query.push((
            "scope",
            params
                .get("scope")
                .cloned()
                .unwrap_or_else(|| format!("repository:{}:pull", self.image.repository)),
        ));
        let url = reqwest::Url::parse_with_params(realm, &query)
            .map_err(|e| Error::Invalid(format!("auth realm: {e}")))?;
        let bytes = self
            .client
            .get(url)
            .send()
            .await?
            .error_for_status()?
            .bytes()
            .await?;
        let body: serde_json::Value = serde_json::from_slice(&bytes)
            .map_err(|e| Error::Invalid(format!("token response: {e}")))?;
        body["token"]
            .as_str()
            .or_else(|| body["access_token"].as_str())
            .map(str::to_string)
            .ok_or_else(|| Error::Invalid("token endpoint returned no token".into()))
    }

    async fn manifest(&mut self, reference: &str) -> Result<(serde_json::Value, String)> {
        let url = format!("{}/manifests/{reference}", self.image.api_base());
        let resp = self.get(&url, MANIFEST_ACCEPT).await?;
        let media_type = resp
            .headers()
            .get("content-type")
            .and_then(|v| v.to_str().ok())
            .unwrap_or("")
            .to_string();
        let bytes = resp.bytes().await?;
        if reference.starts_with("sha256:") {
            reference.parse::<Digest>()?.verify(&Digest::of(&bytes))?;
        }
        let v =
            serde_json::from_slice(&bytes).map_err(|e| Error::Invalid(format!("manifest: {e}")))?;
        Ok((v, media_type))
    }

    async fn blob_bytes(&mut self, digest: &Digest) -> Result<Vec<u8>> {
        let url = format!("{}/blobs/{digest}", self.image.api_base());
        let bytes = self.get(&url, "*/*").await?.bytes().await?.to_vec();
        digest.verify(&Digest::of(&bytes))?;
        Ok(bytes)
    }

    async fn blob_to_store(&mut self, digest: &Digest, store: &LayerStore) -> Result<()> {
        let url = format!("{}/blobs/{digest}", self.image.api_base());
        let mut resp = self.get(&url, "*/*").await?;
        let mut w = store.blob_writer()?;
        while let Some(chunk) = resp.chunk().await? {
            w.write(&chunk)?;
        }
        w.commit(digest).map(drop)
    }
}

/// Parse `Bearer realm="...",service="...",scope="..."`.
pub fn parse_challenge(header: &str) -> Option<HashMap<String, String>> {
    let rest = header.strip_prefix("Bearer ")?;
    let mut out = HashMap::new();
    let mut s = rest.trim();
    while !s.is_empty() {
        let (key, after) = s.split_once('=')?;
        let after = after.strip_prefix('"')?;
        let end = after.find('"')?;
        out.insert(key.trim().to_string(), after[..end].to_string());
        s = after[end + 1..].trim_start_matches(',').trim();
    }
    Some(out)
}

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

    #[test]
    fn parses_bearer_challenge() {
        let p = parse_challenge(
            r#"Bearer realm="https://auth.docker.io/token",service="registry.docker.io",scope="repository:library/busybox:pull""#,
        )
        .unwrap();
        assert_eq!(p["realm"], "https://auth.docker.io/token");
        assert_eq!(p["service"], "registry.docker.io");
        assert_eq!(p["scope"], "repository:library/busybox:pull");
        assert!(parse_challenge("Basic realm=\"x\"").is_none());
    }

    #[test]
    fn picks_this_platform() {
        let arch = match std::env::consts::ARCH {
            "x86_64" => "amd64",
            "aarch64" => "arm64",
            o => o,
        };
        let idx = serde_json::json!({"manifests": [
            {"digest": "sha256:other", "platform": {"os": "linux", "architecture": "s390x-not"}},
            {"digest": "sha256:mine", "platform": {"os": "linux", "architecture": arch}},
        ]});
        assert_eq!(pick_platform(&idx).unwrap(), "sha256:mine");
    }
}