lazaret/crates/lazaret-prepared-lower/src/transport.rs

use std::path::{Path, PathBuf};
use std::str::FromStr;
use std::{fs::OpenOptions, io::Write, os::unix::fs::PermissionsExt};

use anyhow::{Context, Result, bail};
use oci_client::Reference;
use oci_client::client::{Client, ClientConfig};
use oci_client::manifest::{IMAGE_LAYER_GZIP_MEDIA_TYPE, IMAGE_LAYER_MEDIA_TYPE, OciImageManifest};
use oci_client::secrets::RegistryAuth;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};

const OCI_LAYER_ZSTD: &str = "application/vnd.oci.image.layer.v1.tar+zstd";
const MAX_CONFIG_BYTES: usize = 16 * 1024 * 1024;
const MAX_LAYER_BYTES: u64 = 32 * 1024 * 1024 * 1024;

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Platform {
    pub os: String,
    pub architecture: String,
    pub variant: Option<String>,
}

impl Platform {
    #[must_use]
    pub fn linux(architecture: impl Into<String>) -> Self {
        Self {
            os: "linux".to_owned(),
            architecture: architecture.into(),
            variant: None,
        }
    }
}

#[derive(Clone, Eq, PartialEq)]
pub enum RegistryCredentials {
    Anonymous,
    Basic { username: String, password: String },
    Bearer(String),
}

impl std::fmt::Debug for RegistryCredentials {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Anonymous => formatter.write_str("Anonymous"),
            Self::Basic { username, .. } => formatter
                .debug_struct("Basic")
                .field("username", username)
                .field("password", &"[redacted]")
                .finish(),
            Self::Bearer(_) => formatter
                .debug_tuple("Bearer")
                .field(&"[redacted]")
                .finish(),
        }
    }
}

#[derive(Clone, Debug)]
pub struct TransportedImage {
    pub layout_path: PathBuf,
    pub source_identity: Vec<u8>,
    pub selected_manifest_digest: [u8; 32],
}

#[derive(Clone, Debug, Default)]
pub struct OciRegistryTransport;

impl OciRegistryTransport {
    /// Pulls one explicitly selected platform into a digest-verified private OCI layout.
    ///
    /// # Errors
    ///
    /// Refuses invalid references, authentication failures, unsupported platforms, malformed
    /// descriptors, resource limits, and transport or filesystem errors.
    pub async fn pull(
        &self,
        reference: &str,
        credentials: &RegistryCredentials,
        platform: &Platform,
        layout_path: &Path,
    ) -> Result<TransportedImage> {
        self.pull_with_config(
            reference,
            credentials,
            platform,
            layout_path,
            ClientConfig::default(),
        )
        .await
    }

    async fn pull_with_config(
        &self,
        reference: &str,
        credentials: &RegistryCredentials,
        platform: &Platform,
        layout_path: &Path,
        mut config: ClientConfig,
    ) -> Result<TransportedImage> {
        validate_credentials(credentials)?;
        let reference = Reference::from_str(reference).context("invalid OCI image reference")?;
        let selected_platform = platform.clone();
        config.max_concurrent_download = 1;
        config.platform_resolver = Some(Box::new(move |entries| {
            entries
                .iter()
                .find(|entry| {
                    entry.platform.as_ref().is_some_and(|candidate| {
                        candidate.os.to_string() == selected_platform.os
                            && candidate.architecture.to_string() == selected_platform.architecture
                            && candidate.variant == selected_platform.variant
                    })
                })
                .map(|entry| entry.digest.clone())
        }));
        let client = Client::new(config);
        let auth = match credentials {
            RegistryCredentials::Anonymous => RegistryAuth::Anonymous,
            RegistryCredentials::Basic { username, password } => {
                RegistryAuth::Basic(username.clone(), password.clone())
            }
            RegistryCredentials::Bearer(token) => RegistryAuth::Bearer(token.clone()),
        };
        let (manifest, manifest_digest, config_json) = client
            .pull_manifest_and_config(&reference, &auth)
            .await
            .context("pulling selected OCI manifest and config")?;
        if config_json.len() > MAX_CONFIG_BYTES {
            bail!("OCI config exceeds {MAX_CONFIG_BYTES} bytes");
        }
        verify_platform(&config_json, platform)?;
        validate_manifest(&manifest)?;
        create_layout(layout_path)?;
        let blobs = layout_path.join("blobs/sha256");
        write_verified_blob(&blobs, &manifest.config.digest, config_json.as_bytes())?;
        for descriptor in &manifest.layers {
            let size = u64::try_from(descriptor.size).context("negative OCI layer size")?;
            if size > MAX_LAYER_BYTES {
                bail!("OCI layer exceeds {MAX_LAYER_BYTES} bytes");
            }
            let digest_component = digest_component(&descriptor.digest)?;
            let path = blobs.join(digest_component);
            if path.exists() {
                require_existing_blob(&path, &descriptor.digest, size)?;
                continue;
            }
            let file = tokio::fs::OpenOptions::new()
                .create_new(true)
                .write(true)
                .open(&path)
                .await?;
            if let Err(error) = client.pull_blob(&reference, descriptor, file).await {
                let _ = std::fs::remove_file(&path);
                return Err(error)
                    .with_context(|| format!("pulling OCI layer {}", descriptor.digest));
            }
            let observed = std::fs::metadata(&path)?.len();
            if observed != size {
                bail!("OCI layer {} size mismatch", descriptor.digest);
            }
        }
        let local_manifest = serde_json::to_vec(&manifest)?;
        let local_digest = format!("sha256:{}", hex(&Sha256::digest(&local_manifest)));
        write_verified_blob(&blobs, &local_digest, &local_manifest)?;
        write_layout_index(layout_path, &local_digest, local_manifest.len())?;
        let selected_manifest_digest = parse_sha256(&local_digest)?;
        let source_identity = format!(
            "oci-v1\0{}\0{}\0{manifest_digest}",
            reference.registry(),
            reference.repository()
        )
        .into_bytes();
        Ok(TransportedImage {
            layout_path: layout_path.to_owned(),
            source_identity,
            selected_manifest_digest,
        })
    }
}

#[derive(Deserialize)]
struct ImageConfigPlatform {
    os: String,
    architecture: String,
    #[serde(default)]
    variant: Option<String>,
}

fn verify_platform(config: &str, expected: &Platform) -> Result<()> {
    let observed: ImageConfigPlatform = serde_json::from_str(config)?;
    if observed.os != expected.os
        || observed.architecture != expected.architecture
        || observed.variant != expected.variant
    {
        bail!(
            "selected image platform {}/{} does not match requested {}/{}",
            observed.os,
            observed.architecture,
            expected.os,
            expected.architecture
        );
    }
    Ok(())
}

fn validate_credentials(credentials: &RegistryCredentials) -> Result<()> {
    match credentials {
        RegistryCredentials::Anonymous => Ok(()),
        RegistryCredentials::Basic { username, password }
            if !username.is_empty() && !password.is_empty() =>
        {
            Ok(())
        }
        RegistryCredentials::Bearer(token) if !token.is_empty() => Ok(()),
        RegistryCredentials::Basic { .. } | RegistryCredentials::Bearer(_) => {
            bail!("registry credentials are incomplete")
        }
    }
}

fn validate_manifest(manifest: &OciImageManifest) -> Result<()> {
    if manifest.schema_version != 2 || manifest.layers.len() > 256 {
        bail!("unsupported OCI manifest schema or layer count");
    }
    for descriptor in &manifest.layers {
        if !matches!(
            descriptor.media_type.as_str(),
            IMAGE_LAYER_MEDIA_TYPE | IMAGE_LAYER_GZIP_MEDIA_TYPE | OCI_LAYER_ZSTD
        ) {
            bail!("unsupported OCI layer media type {}", descriptor.media_type);
        }
        digest_component(&descriptor.digest)?;
        u64::try_from(descriptor.size).context("negative OCI layer size")?;
    }
    digest_component(&manifest.config.digest)?;
    Ok(())
}

fn create_layout(path: &Path) -> Result<()> {
    std::fs::create_dir(path).context("creating private OCI layout")?;
    std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o700))?;
    std::fs::create_dir_all(path.join("blobs/sha256"))?;
    std::fs::set_permissions(path.join("blobs"), std::fs::Permissions::from_mode(0o700))?;
    std::fs::set_permissions(
        path.join("blobs/sha256"),
        std::fs::Permissions::from_mode(0o700),
    )?;
    std::fs::write(
        path.join("oci-layout"),
        b"{\"imageLayoutVersion\":\"1.0.0\"}\n",
    )?;
    Ok(())
}

fn write_verified_blob(blobs: &Path, digest: &str, bytes: &[u8]) -> Result<()> {
    if parse_sha256(digest)? != Sha256::digest(bytes).as_slice() {
        bail!("OCI descriptor digest mismatch for {digest}");
    }
    let path = blobs.join(digest_component(digest)?);
    match OpenOptions::new().create_new(true).write(true).open(&path) {
        Ok(mut file) => file.write_all(bytes)?,
        Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => {
            require_existing_blob(&path, digest, bytes.len() as u64)?;
        }
        Err(error) => return Err(error.into()),
    }
    Ok(())
}

fn require_existing_blob(path: &Path, digest: &str, size: u64) -> Result<()> {
    let bytes = std::fs::read(path)?;
    if bytes.len() as u64 != size || parse_sha256(digest)? != Sha256::digest(&bytes).as_slice() {
        bail!("existing OCI blob does not match its descriptor");
    }
    Ok(())
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct LayoutIndex<'a> {
    schema_version: u8,
    manifests: [LayoutDescriptor<'a>; 1],
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct LayoutDescriptor<'a> {
    media_type: &'a str,
    digest: &'a str,
    size: usize,
}

fn write_layout_index(path: &Path, digest: &str, size: usize) -> Result<()> {
    let index = LayoutIndex {
        schema_version: 2,
        manifests: [LayoutDescriptor {
            media_type: "application/vnd.oci.image.manifest.v1+json",
            digest,
            size,
        }],
    };
    std::fs::write(path.join("index.json"), serde_json::to_vec(&index)?)?;
    Ok(())
}

fn digest_component(digest: &str) -> Result<&str> {
    let component = digest
        .strip_prefix("sha256:")
        .context("only sha256 OCI descriptors are supported")?;
    if component.len() != 64 || !component.bytes().all(|byte| byte.is_ascii_hexdigit()) {
        bail!("invalid sha256 OCI descriptor");
    }
    Ok(component)
}

fn parse_sha256(digest: &str) -> Result<[u8; 32]> {
    let component = digest_component(digest)?;
    let mut bytes = [0; 32];
    for (index, pair) in component.as_bytes().chunks_exact(2).enumerate() {
        bytes[index] = u8::from_str_radix(std::str::from_utf8(pair)?, 16)?;
    }
    Ok(bytes)
}

fn hex(bytes: &[u8]) -> String {
    use std::fmt::Write as _;

    bytes.iter().fold(
        String::with_capacity(bytes.len() * 2),
        |mut output, byte| {
            write!(output, "{byte:02x}").expect("writing to String cannot fail");
            output
        },
    )
}

#[cfg(test)]
mod tests {
    use std::io::{Read, Write};
    use std::net::TcpListener;

    use oci_client::client::ClientProtocol;

    use super::*;

    #[test]
    fn refuses_mismatched_and_incomplete_platforms() {
        assert!(
            verify_platform(
                r#"{"os":"linux","architecture":"amd64"}"#,
                &Platform::linux("amd64")
            )
            .is_ok()
        );
        assert!(
            verify_platform(
                r#"{"os":"linux","architecture":"arm64"}"#,
                &Platform::linux("amd64")
            )
            .is_err()
        );
        assert!(
            verify_platform(
                r#"{"os":"linux","architecture":"amd64","variant":"v2"}"#,
                &Platform::linux("amd64")
            )
            .is_err()
        );
        assert!(verify_platform(r#"{"architecture":"amd64"}"#, &Platform::linux("amd64")).is_err());
    }

    #[test]
    fn incomplete_authentication_refuses_without_layout_or_secret_disclosure() {
        let runtime = tokio::runtime::Builder::new_current_thread()
            .build()
            .unwrap();
        let temp = tempfile::tempdir().unwrap();
        let result = runtime.block_on(OciRegistryTransport.pull_with_config(
            "registry.invalid/private/image:latest",
            &RegistryCredentials::Basic {
                username: "developer".to_owned(),
                password: String::new(),
            },
            &Platform::linux("amd64"),
            &temp.path().join("layout"),
            ClientConfig::default(),
        ));
        let error = result.unwrap_err().to_string();
        assert!(error.contains("credentials are incomplete"));
        assert!(!error.contains("developer"));
        assert!(!temp.path().join("layout").exists());
        assert_eq!(
            format!(
                "{:?}",
                RegistryCredentials::Bearer("not-a-real-token".to_owned())
            ),
            "Bearer(\"[redacted]\")"
        );
    }

    #[test]
    fn registry_authentication_failure_refuses_before_layout_creation() {
        let listener = match TcpListener::bind("127.0.0.1:0") {
            Ok(listener) => listener,
            Err(error) if error.kind() == std::io::ErrorKind::PermissionDenied => return,
            Err(error) => panic!("binding test registry: {error}"),
        };
        let address = listener.local_addr().unwrap();
        let server = std::thread::spawn(move || {
            let (mut stream, _) = listener.accept().unwrap();
            let mut request = [0; 4096];
            assert!(stream.read(&mut request).unwrap() > 0);
            stream
                .write_all(
                    b"HTTP/1.1 401 Unauthorized\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
                )
                .unwrap();
        });
        let runtime = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .unwrap();
        let temp = tempfile::tempdir().unwrap();
        let result = runtime.block_on(OciRegistryTransport.pull_with_config(
            &format!("{address}/private/image:latest"),
            &RegistryCredentials::Basic {
                username: "wrong".to_owned(),
                password: "wrong".to_owned(),
            },
            &Platform::linux("amd64"),
            &temp.path().join("layout"),
            ClientConfig {
                protocol: ClientProtocol::Http,
                ..ClientConfig::default()
            },
        ));
        assert!(result.is_err());
        assert!(!temp.path().join("layout").exists());
        server.join().unwrap();
    }
}