oci-client 0.18.0

// Tests for validating digests of different types and for malicious servers
use std::{net::SocketAddr, sync::LazyLock};

use axum::{
    extract::{Path, State},
    http::{HeaderMap, StatusCode},
    routing::get,
    Router,
};
use oci_client::{
    client::{linux_amd64_resolver, ClientConfig, ClientProtocol},
    errors::{OciDistributionError, OciErrorCode},
    manifest::OciDescriptor,
    Client, Reference,
};
use rstest::rstest;
use sha2::{Digest, Sha256, Sha512};
use tokio::{net::TcpListener, task::JoinHandle};

const DIGEST_HEADER: &str = "Docker-Content-Digest";

static MANIFEST: &[u8] = include_bytes!("./fixtures/manifest.json");
static BLOB: &[u8] = include_bytes!("./fixtures/blob.tar.gz");
static CONFIG: &[u8] = include_bytes!("./fixtures/config.json");
static BLOB_UNKNOWN_BODY: &[u8] =
    br#"{"errors":[{"code":"BLOB_UNKNOWN","message":"blob unknown to registry"}]}"#;

static MANIFEST_DIGEST: LazyLock<String> = LazyLock::new(|| digest(MANIFEST));
static MANIFEST_DIGEST_SHA512: LazyLock<String> = LazyLock::new(|| digest_sha512(MANIFEST));
static BLOB_DIGEST: LazyLock<String> = LazyLock::new(|| digest(BLOB));
static BLOB_DIGEST_SHA512: LazyLock<String> = LazyLock::new(|| digest_sha512(BLOB));
static CONFIG_DIGEST: LazyLock<String> = LazyLock::new(|| digest(CONFIG));

fn digest(data: &[u8]) -> String {
    format!("sha256:{}", hex::encode(Sha256::digest(data)))
}

fn digest_sha512(data: &[u8]) -> String {
    format!("sha512:{}", hex::encode(Sha512::digest(data)))
}

async fn manifest_handler(
    State(state): State<ServerConfig>,
    Path(digest): Path<String>,
) -> (HeaderMap, &'static [u8]) {
    let resp_digest = if digest.starts_with("sha256:") && state.bad_manifest {
        digest
    } else {
        MANIFEST_DIGEST.clone()
    };

    let mut headers = HeaderMap::new();
    if state.empty_digest {
        headers.insert(DIGEST_HEADER, "".parse().unwrap());
    } else {
        headers.insert(DIGEST_HEADER, resp_digest.parse().unwrap());
    }
    headers.insert(
        "Content-Type",
        "application/vnd.docker.distribution.manifest.v2+json"
            .parse()
            .unwrap(),
    );

    (headers, MANIFEST)
}

async fn blob_handler(
    State(state): State<ServerConfig>,
    Path(digest): Path<String>,
) -> Result<(HeaderMap, &'static [u8]), (StatusCode, &'static [u8])> {
    let (content, resp_digest) = match digest.as_str() {
        d if d == CONFIG_DIGEST.as_str() => (
            CONFIG,
            if state.bad_config {
                "sha256:deadbeef"
            } else {
                CONFIG_DIGEST.as_str()
            },
        ),
        d if state.blob_sha512 && d == BLOB_DIGEST.as_str() => (
            BLOB,
            if state.bad_blob {
                "sha256:deadbeef"
            } else {
                BLOB_DIGEST_SHA512.as_str()
            },
        ),
        d if d == BLOB_DIGEST.as_str() => (
            BLOB,
            if state.bad_blob {
                "sha256:deadbeef"
            } else {
                BLOB_DIGEST.as_str()
            },
        ),
        _ => return Err((StatusCode::NOT_FOUND, BLOB_UNKNOWN_BODY)),
    };

    let mut headers = HeaderMap::new();
    headers.insert(DIGEST_HEADER, resp_digest.parse().unwrap());

    Ok((headers, content))
}

#[derive(Clone, Copy)]
struct ServerConfig {
    bad_manifest: bool,
    bad_config: bool,
    bad_blob: bool,
    blob_sha512: bool,
    empty_digest: bool,
}

struct BadServer {
    handle: JoinHandle<()>,
    pub server: String,
}

impl Drop for BadServer {
    fn drop(&mut self) {
        self.handle.abort()
    }
}

impl BadServer {
    pub async fn new(config: ServerConfig) -> Self {
        let app = Router::new()
            .route("/v2/busybox/manifests/{digest}", get(manifest_handler))
            .route("/v2/busybox/blobs/{digest}", get(blob_handler))
            .with_state(config);

        let addr = SocketAddr::from(([127, 0, 0, 1], 0));
        let listener = TcpListener::bind(addr).await.unwrap();
        let server_addr = listener.local_addr().unwrap();
        let port = server_addr.port();
        let server = format!("127.0.0.1:{port}");
        let handle = tokio::spawn(async move {
            axum::serve(listener, app).await.unwrap();
        });
        Self { handle, server }
    }
}
#[tokio::test]
async fn test_bad_manifest() {
    let server = BadServer::new(ServerConfig {
        bad_manifest: true,
        bad_config: false,
        bad_blob: false,
        blob_sha512: false,
        empty_digest: false,
    })
    .await;

    let client = Client::new(ClientConfig {
        protocol: ClientProtocol::Http,
        platform_resolver: Some(Box::new(linux_amd64_resolver)),
        ..Default::default()
    });
    let auth = &oci_client::secrets::RegistryAuth::Anonymous;

    let reference = Reference::try_from(format!(
        "{}/busybox@sha256:deadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeef",
        server.server
    ))
    .expect("failed to parse reference");

    client
        .pull_manifest(&reference, auth)
        .await
        .expect_err("Expected an error with a mismatched sha");
}

#[tokio::test]
async fn test_bad_config() {
    let server = BadServer::new(ServerConfig {
        bad_manifest: false,
        bad_config: true,
        bad_blob: false,
        blob_sha512: false,
        empty_digest: false,
    })
    .await;

    let client = Client::new(ClientConfig {
        protocol: ClientProtocol::Http,
        platform_resolver: Some(Box::new(linux_amd64_resolver)),
        ..Default::default()
    });
    let auth = &oci_client::secrets::RegistryAuth::Anonymous;

    let reference = Reference::try_from(format!(
        "{}/busybox@{}",
        server.server,
        MANIFEST_DIGEST.as_str()
    ))
    .expect("failed to parse reference");

    assert!(
        client
            .pull(
                &reference,
                auth,
                vec!["application/vnd.docker.image.rootfs.diff.tar.gzip"],
            )
            .await
            .is_err(),
        "Expected an error with a bad config"
    );
}

#[tokio::test]
async fn test_bad_blob() {
    let server = BadServer::new(ServerConfig {
        bad_manifest: false,
        bad_config: false,
        bad_blob: true,
        blob_sha512: false,
        empty_digest: false,
    })
    .await;
    let client = Client::new(ClientConfig {
        protocol: ClientProtocol::Http,
        platform_resolver: Some(Box::new(linux_amd64_resolver)),
        ..Default::default()
    });
    let auth = &oci_client::secrets::RegistryAuth::Anonymous;

    let reference = Reference::try_from(format!(
        "{}/busybox@{}",
        server.server,
        MANIFEST_DIGEST.as_str()
    ))
    .expect("failed to parse reference");

    assert!(
        client
            .pull(
                &reference,
                auth,
                vec!["application/vnd.docker.image.rootfs.diff.tar.gzip"],
            )
            .await
            .is_err(),
        "Expected an error with a bad blob"
    );
}

#[tokio::test]
async fn test_good_pull() {
    let server = BadServer::new(ServerConfig {
        bad_manifest: false,
        bad_config: false,
        bad_blob: false,
        blob_sha512: false,
        empty_digest: false,
    })
    .await;

    let client = Client::new(ClientConfig {
        protocol: ClientProtocol::Http,
        platform_resolver: Some(Box::new(linux_amd64_resolver)),
        ..Default::default()
    });
    let auth = &oci_client::secrets::RegistryAuth::Anonymous;

    let reference = Reference::try_from(format!(
        "{}/busybox@{}",
        server.server,
        MANIFEST_DIGEST.as_str()
    ))
    .expect("failed to parse reference");

    client
        .pull(
            &reference,
            auth,
            vec!["application/vnd.docker.image.rootfs.diff.tar.gzip"],
        )
        .await
        .expect("Expected a good pull");
}

#[tokio::test]
async fn test_different_reference_sha() {
    let server = BadServer::new(ServerConfig {
        bad_manifest: false,
        bad_config: false,
        bad_blob: false,
        blob_sha512: false,
        empty_digest: false,
    })
    .await;

    let client = Client::new(ClientConfig {
        protocol: ClientProtocol::Http,
        platform_resolver: Some(Box::new(linux_amd64_resolver)),
        ..Default::default()
    });
    let auth = &oci_client::secrets::RegistryAuth::Anonymous;

    let reference = Reference::try_from(format!(
        "{}/busybox@sha256:deadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeef",
        server.server
    ))
    .expect("failed to parse reference");

    client
        .pull_manifest(&reference, auth)
        .await
        .expect_err("Expected an error with a mismatched reference sha");

    // Also try using a sha512 digest
    let reference = Reference::try_from(format!(
        "{}/busybox@sha256:deadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeef",
        server.server
    ))
    .expect("failed to parse reference");

    client
        .pull_manifest(&reference, auth)
        .await
        .expect_err("Expected an error with a mismatched reference sha");
}

#[tokio::test]
async fn test_different_manifest_algos() {
    let server = BadServer::new(ServerConfig {
        bad_manifest: false,
        bad_config: false,
        bad_blob: false,
        blob_sha512: false,
        empty_digest: false,
    })
    .await;
    let client = Client::new(ClientConfig {
        protocol: ClientProtocol::Http,
        platform_resolver: Some(Box::new(linux_amd64_resolver)),
        ..Default::default()
    });
    let auth = &oci_client::secrets::RegistryAuth::Anonymous;
    let reference = Reference::try_from(format!(
        "{}/busybox@{}",
        server.server,
        MANIFEST_DIGEST_SHA512.as_str()
    ))
    .expect("failed to parse reference");

    client
        .pull_manifest(&reference, auth)
        .await
        .expect("Expected a good pull with two different algorithms");
}

#[tokio::test]
async fn test_different_blob_algos() {
    let server = BadServer::new(ServerConfig {
        bad_manifest: false,
        bad_config: false,
        bad_blob: false,
        blob_sha512: true,
        empty_digest: false,
    })
    .await;

    let client = Client::new(ClientConfig {
        protocol: ClientProtocol::Http,
        platform_resolver: Some(Box::new(linux_amd64_resolver)),
        ..Default::default()
    });
    let auth = &oci_client::secrets::RegistryAuth::Anonymous;

    let reference = Reference::try_from(format!(
        "{}/busybox@{}",
        server.server,
        MANIFEST_DIGEST.as_str()
    ))
    .expect("failed to parse reference");

    client
        .pull(
            &reference,
            auth,
            vec!["application/vnd.docker.image.rootfs.diff.tar.gzip"],
        )
        .await
        .expect("Expected a good pull");
}

#[tokio::test]
async fn test_empty_digest_header() {
    let server = BadServer::new(ServerConfig {
        bad_manifest: false,
        bad_config: false,
        bad_blob: false,
        blob_sha512: false,
        empty_digest: true,
    })
    .await;

    let client = Client::new(ClientConfig {
        protocol: ClientProtocol::Http,
        platform_resolver: Some(Box::new(linux_amd64_resolver)),
        ..Default::default()
    });
    let auth = &oci_client::secrets::RegistryAuth::Anonymous;

    let reference = Reference::try_from(format!(
        "{}/busybox@{}",
        server.server,
        MANIFEST_DIGEST.as_str()
    ))
    .expect("failed to parse reference");

    client
        .pull_manifest(&reference, auth)
        .await
        .expect("Expected empty digest header to be treated as missing");
}

// Regression test for https://github.com/oras-project/rust-oci-client/issues/64:
// pulling a blob that the registry doesn't have (404 with a BLOB_UNKNOWN envelope) must return an
// error and must not write the error body into the caller's output buffer.
#[tokio::test]
async fn test_pull_blob_not_found() {
    let server = BadServer::new(ServerConfig {
        bad_manifest: false,
        bad_config: false,
        bad_blob: false,
        blob_sha512: false,
        empty_digest: false,
    })
    .await;

    let client = Client::new(ClientConfig {
        protocol: ClientProtocol::Http,
        platform_resolver: Some(Box::new(linux_amd64_resolver)),
        ..Default::default()
    });

    let reference = Reference::try_from(format!(
        "{}/busybox@{}",
        server.server,
        MANIFEST_DIGEST.as_str()
    ))
    .expect("failed to parse reference");

    // A layer digest the mock registry doesn't know about, so the blob handler returns 404.
    let bad_layer = OciDescriptor {
        digest: "sha256:deadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeef"
            .to_string(),
        ..Default::default()
    };

    // pull_blob must return an error and must not leak the error body into the output buffer.
    let mut buf: Vec<u8> = Vec::new();
    let err = client
        .pull_blob(&reference, &bad_layer, &mut buf)
        .await
        .expect_err("Expected an error pulling a missing blob");
    assert!(
        matches!(&err, OciDistributionError::RegistryError { envelope, .. }
            if envelope.errors.iter().any(|e| e.code == OciErrorCode::BlobUnknown)),
        "Expected a BlobUnknown registry error, got: {err}"
    );
    assert!(
        buf.is_empty(),
        "Expected no body to be written for a missing blob, got {} bytes",
        buf.len()
    );
}

enum StreamVariant {
    Stream,
    Partial,
}

// Regression test for https://github.com/oras-project/rust-oci-client/issues/64 (streaming path):
// pulling a blob that the registry doesn't have (404 with a BLOB_UNKNOWN envelope) must return a
// BlobUnknown registry error for both streaming variants.
#[rstest]
#[case::stream(StreamVariant::Stream)]
#[case::partial(StreamVariant::Partial)]
#[tokio::test]
async fn test_pull_blob_stream_not_found(#[case] variant: StreamVariant) {
    let server = BadServer::new(ServerConfig {
        bad_manifest: false,
        bad_config: false,
        bad_blob: false,
        blob_sha512: false,
        empty_digest: false,
    })
    .await;

    let client = Client::new(ClientConfig {
        protocol: ClientProtocol::Http,
        platform_resolver: Some(Box::new(linux_amd64_resolver)),
        ..Default::default()
    });

    let reference = Reference::try_from(format!(
        "{}/busybox@{}",
        server.server,
        MANIFEST_DIGEST.as_str()
    ))
    .expect("failed to parse reference");

    // A layer digest the mock registry doesn't know about, so the blob handler returns 404.
    let bad_layer = OciDescriptor {
        digest: "sha256:deadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeef"
            .to_string(),
        ..Default::default()
    };

    let err = match variant {
        StreamVariant::Stream => client
            .pull_blob_stream(&reference, &bad_layer)
            .await
            .err()
            .expect("Expected an error streaming a missing blob"),
        StreamVariant::Partial => client
            .pull_blob_stream_partial(&reference, &bad_layer, 0, None)
            .await
            .err()
            .expect("Expected an error partially streaming a missing blob"),
    };
    assert!(
        matches!(&err, OciDistributionError::RegistryError { envelope, .. }
            if envelope.errors.iter().any(|e| e.code == OciErrorCode::BlobUnknown)),
        "Expected a BlobUnknown registry error, got: {err}"
    );
}