use std::path::{Path, PathBuf};
use std::sync::Arc;
use a3s_box_core::error::{BoxError, Result};
use oci_distribution::client::{ClientConfig, ClientProtocol, Config, ImageLayer, PushResponse};
use oci_distribution::manifest::{ImageIndexEntry, OciImageManifest};
use oci_distribution::secrets::RegistryAuth as OciRegistryAuth;
use oci_distribution::{Client, Reference};
use super::credentials::CredentialStore;
use super::reference::ImageReference;
use super::signing::{verify_image_signature, SignaturePolicy, VerifyResult};
const REGISTRY_PROTOCOL_ENV: &str = "A3S_REGISTRY_PROTOCOL";
fn registry_protocol_from_env() -> ClientProtocol {
match std::env::var(REGISTRY_PROTOCOL_ENV) {
Ok(value) if value.eq_ignore_ascii_case("http") => ClientProtocol::Http,
_ => ClientProtocol::Https,
}
}
struct HashingFileWriter {
file: tokio::fs::File,
hasher: sha2::Sha256,
}
impl HashingFileWriter {
fn new(file: tokio::fs::File) -> Self {
use sha2::Digest;
Self {
file,
hasher: sha2::Sha256::new(),
}
}
fn finalize_hex(self) -> String {
use sha2::Digest;
format!("{:x}", self.hasher.finalize())
}
}
impl tokio::io::AsyncWrite for HashingFileWriter {
fn poll_write(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
buf: &[u8],
) -> std::task::Poll<std::io::Result<usize>> {
use sha2::Digest;
let this = self.get_mut();
match std::pin::Pin::new(&mut this.file).poll_write(cx, buf) {
std::task::Poll::Ready(Ok(n)) => {
this.hasher.update(&buf[..n]);
std::task::Poll::Ready(Ok(n))
}
other => other,
}
}
fn poll_flush(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<std::io::Result<()>> {
std::pin::Pin::new(&mut self.get_mut().file).poll_flush(cx)
}
fn poll_shutdown(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<std::io::Result<()>> {
std::pin::Pin::new(&mut self.get_mut().file).poll_shutdown(cx)
}
}
async fn stream_and_verify_blob(
client: &Client,
oci_ref: &Reference,
descriptor: &oci_distribution::manifest::OciDescriptor,
dest: &Path,
what: &str,
registry: &str,
) -> Result<()> {
use tokio::io::AsyncWriteExt;
let tmp = dest.with_extension("partial");
let file = tokio::fs::File::create(&tmp)
.await
.map_err(|e| BoxError::RegistryError {
registry: registry.to_string(),
message: format!("Failed to create {what} file: {e}"),
})?;
let mut writer = HashingFileWriter::new(file);
if let Err(e) = client.pull_blob(oci_ref, descriptor, &mut writer).await {
let _ = tokio::fs::remove_file(&tmp).await;
return Err(BoxError::RegistryError {
registry: registry.to_string(),
message: format!("Failed to pull {what}: {e}"),
});
}
let _ = writer.flush().await;
let _ = writer.shutdown().await;
let actual_hex = writer.finalize_hex();
match descriptor.digest.strip_prefix("sha256:") {
Some(expected_hex) if actual_hex.eq_ignore_ascii_case(expected_hex) => {}
Some(expected_hex) => {
let _ = tokio::fs::remove_file(&tmp).await;
return Err(BoxError::RegistryError {
registry: registry.to_string(),
message: format!(
"{what} digest mismatch: expected sha256:{expected_hex}, computed sha256:{actual_hex}"
),
});
}
None => {
let _ = tokio::fs::remove_file(&tmp).await;
return Err(BoxError::RegistryError {
registry: registry.to_string(),
message: format!(
"{what} uses an unsupported digest algorithm ({}); refusing to store \
unverifiable content (only sha256 is supported)",
descriptor.digest
),
});
}
}
tokio::fs::rename(&tmp, dest)
.await
.map_err(|e| BoxError::RegistryError {
registry: registry.to_string(),
message: format!("Failed to store {what} blob: {e}"),
})
}
type PullProgressFn = Arc<dyn Fn(usize, usize, &str, i64) + Send + Sync>;
#[derive(Debug, Clone)]
pub struct RegistryAuth {
username: Option<String>,
password: Option<String>,
}
impl RegistryAuth {
pub fn anonymous() -> Self {
Self {
username: None,
password: None,
}
}
pub fn basic(username: impl Into<String>, password: impl Into<String>) -> Self {
Self {
username: Some(username.into()),
password: Some(password.into()),
}
}
pub fn from_env() -> Self {
let username = std::env::var("REGISTRY_USERNAME").ok();
let password = std::env::var("REGISTRY_PASSWORD").ok();
if username.is_some() && password.is_some() {
Self { username, password }
} else {
Self::anonymous()
}
}
pub fn from_credential_store(registry: &str) -> Self {
if let Ok(store) = CredentialStore::default_path() {
if let Ok(Some((username, password))) = store.get(registry) {
return Self::basic(username, password);
}
}
Self::from_env()
}
fn to_oci_auth(&self) -> OciRegistryAuth {
match (&self.username, &self.password) {
(Some(u), Some(p)) => OciRegistryAuth::Basic(u.clone(), p.clone()),
_ => OciRegistryAuth::Anonymous,
}
}
}
pub(crate) struct RegistryPuller {
client: Client,
auth: RegistryAuth,
signature_policy: SignaturePolicy,
progress_fn: Option<PullProgressFn>,
}
impl Default for RegistryPuller {
fn default() -> Self {
Self::new()
}
}
impl RegistryPuller {
pub fn new() -> Self {
Self::with_auth(RegistryAuth::anonymous())
}
pub fn with_auth(auth: RegistryAuth) -> Self {
let config = ClientConfig {
protocol: registry_protocol_from_env(),
platform_resolver: Some(Box::new(linux_platform_resolver)),
..Default::default()
};
let client = Client::new(config);
Self {
client,
auth,
signature_policy: SignaturePolicy::default(),
progress_fn: None,
}
}
pub fn with_auth_and_platform(auth: RegistryAuth, platform: Option<String>) -> Self {
let Some(platform) = platform else {
return Self::with_auth(auth);
};
let arch = resolve_target_arch(Some(&platform));
let config = ClientConfig {
protocol: registry_protocol_from_env(),
platform_resolver: Some(Box::new(platform_resolver_for(arch))),
..Default::default()
};
let client = Client::new(config);
Self {
client,
auth,
signature_policy: SignaturePolicy::default(),
progress_fn: None,
}
}
pub fn with_signature_policy(mut self, policy: SignaturePolicy) -> Self {
self.signature_policy = policy;
self
}
pub fn with_progress_fn(mut self, f: PullProgressFn) -> Self {
self.progress_fn = Some(f);
self
}
pub async fn pull(&self, reference: &ImageReference, target_dir: &Path) -> Result<PathBuf> {
let oci_ref = self.to_oci_reference(reference)?;
tracing::info!(
reference = %reference,
target = %target_dir.display(),
"Pulling image from registry"
);
let blobs_dir = target_dir.join("blobs").join("sha256");
std::fs::create_dir_all(&blobs_dir).map_err(|e| BoxError::RegistryError {
registry: reference.registry.clone(),
message: format!("Failed to create blobs directory: {}", e),
})?;
let auth = self.auth.to_oci_auth();
let (image_manifest, manifest_digest) = self
.client
.pull_image_manifest(&oci_ref, &auth)
.await
.map_err(|e| BoxError::RegistryError {
registry: reference.registry.clone(),
message: format!("Failed to pull manifest: {}", e),
})?;
let verify_result = verify_image_signature(
&self.signature_policy,
&reference.registry,
&reference.repository,
&manifest_digest,
)
.await;
if !verify_result.is_ok() {
return Err(BoxError::RegistryError {
registry: reference.registry.clone(),
message: match verify_result {
VerifyResult::NoSignature => format!(
"Image {}:{} has no signature and policy requires verification",
reference.repository,
reference.tag.as_deref().unwrap_or("latest")
),
VerifyResult::Failed(msg) => format!(
"Image signature verification failed for {}:{}: {}",
reference.repository,
reference.tag.as_deref().unwrap_or("latest"),
msg
),
_ => "Signature verification failed".to_string(),
},
});
}
let manifest_json = serde_json::to_vec(&image_manifest)?;
let manifest_digest_hex = manifest_digest
.strip_prefix("sha256:")
.unwrap_or(&manifest_digest);
std::fs::write(blobs_dir.join(manifest_digest_hex), &manifest_json).map_err(|e| {
BoxError::RegistryError {
registry: reference.registry.clone(),
message: format!("Failed to write manifest: {}", e),
}
})?;
self.pull_image_content(&oci_ref, &image_manifest, &blobs_dir, &reference.registry)
.await?;
std::fs::write(
target_dir.join("oci-layout"),
r#"{"imageLayoutVersion":"1.0.0"}"#,
)
.map_err(|e| BoxError::RegistryError {
registry: reference.registry.clone(),
message: format!("Failed to write oci-layout: {}", e),
})?;
let index = serde_json::json!({
"schemaVersion": 2,
"manifests": [{
"mediaType": "application/vnd.oci.image.manifest.v1+json",
"digest": manifest_digest,
"size": manifest_json.len()
}]
});
std::fs::write(
target_dir.join("index.json"),
serde_json::to_string_pretty(&index)?,
)
.map_err(|e| BoxError::RegistryError {
registry: reference.registry.clone(),
message: format!("Failed to write index.json: {}", e),
})?;
tracing::info!(
reference = %reference,
digest = %manifest_digest,
"Image pulled successfully"
);
Ok(target_dir.to_path_buf())
}
pub async fn pull_manifest_digest(&self, reference: &ImageReference) -> Result<String> {
let oci_ref = self.to_oci_reference(reference)?;
let auth = self.auth.to_oci_auth();
let (_manifest, digest) =
self.client
.pull_manifest(&oci_ref, &auth)
.await
.map_err(|e| BoxError::RegistryError {
registry: reference.registry.clone(),
message: format!("Failed to pull manifest: {}", e),
})?;
Ok(digest)
}
async fn pull_image_content(
&self,
oci_ref: &Reference,
manifest: &OciImageManifest,
blobs_dir: &Path,
registry: &str,
) -> Result<()> {
let config_descriptor = &manifest.config;
let config_digest_hex = config_descriptor
.digest
.strip_prefix("sha256:")
.unwrap_or(&config_descriptor.digest);
stream_and_verify_blob(
&self.client,
oci_ref,
config_descriptor,
&blobs_dir.join(config_digest_hex),
"config blob",
registry,
)
.await?;
let total = manifest.layers.len();
for (idx, layer) in manifest.layers.iter().enumerate() {
tracing::debug!(
digest = %layer.digest,
size = layer.size,
"Pulling layer"
);
if let Some(ref f) = self.progress_fn {
f(idx + 1, total, &layer.digest, layer.size);
}
let layer_digest_hex = layer
.digest
.strip_prefix("sha256:")
.unwrap_or(&layer.digest);
stream_and_verify_blob(
&self.client,
oci_ref,
layer,
&blobs_dir.join(layer_digest_hex),
"layer",
registry,
)
.await?;
if let Some(ref f) = self.progress_fn {
f(idx + 1, total, &layer.digest, -(layer.size));
}
}
Ok(())
}
fn to_oci_reference(&self, reference: &ImageReference) -> Result<Reference> {
let ref_str = if let Some(ref digest) = reference.digest {
format!("{}/{}@{}", reference.registry, reference.repository, digest)
} else if let Some(ref tag) = reference.tag {
format!("{}/{}:{}", reference.registry, reference.repository, tag)
} else {
format!("{}/{}:latest", reference.registry, reference.repository)
};
ref_str.parse::<Reference>().map_err(|e| {
BoxError::OciImageError(format!("Invalid OCI reference '{}': {}", ref_str, e))
})
}
}
#[derive(Debug, Clone)]
pub struct PushResult {
pub config_url: String,
pub manifest_url: String,
pub manifest_digest: String,
}
pub struct RegistryPusher {
client: Client,
auth: RegistryAuth,
}
impl Default for RegistryPusher {
fn default() -> Self {
Self::new()
}
}
impl RegistryPusher {
pub fn new() -> Self {
Self::with_auth(RegistryAuth::anonymous())
}
pub fn with_auth(auth: RegistryAuth) -> Self {
let config = ClientConfig {
protocol: registry_protocol_from_env(),
..Default::default()
};
let client = Client::new(config);
Self { client, auth }
}
pub async fn push(&self, reference: &ImageReference, image_dir: &Path) -> Result<PushResult> {
let oci_ref = self.to_oci_reference(reference)?;
tracing::info!(
reference = %reference,
source = %image_dir.display(),
"Pushing image to registry"
);
let index_path = image_dir.join("index.json");
let index_data = std::fs::read_to_string(&index_path)
.map_err(|e| BoxError::OciImageError(format!("Failed to read index.json: {}", e)))?;
let index: serde_json::Value = serde_json::from_str(&index_data)?;
let manifest_digest = index["manifests"][0]["digest"].as_str().ok_or_else(|| {
BoxError::OciImageError("No manifest digest in index.json".to_string())
})?;
let manifest_digest_hex = manifest_digest
.strip_prefix("sha256:")
.unwrap_or(manifest_digest);
let blobs_dir = image_dir.join("blobs").join("sha256");
let manifest_data = std::fs::read(blobs_dir.join(manifest_digest_hex))
.map_err(|e| BoxError::OciImageError(format!("Failed to read manifest blob: {}", e)))?;
let manifest: OciImageManifest = serde_json::from_slice(&manifest_data)?;
let config_digest_hex = manifest
.config
.digest
.strip_prefix("sha256:")
.unwrap_or(&manifest.config.digest);
let config_data = std::fs::read(blobs_dir.join(config_digest_hex))
.map_err(|e| BoxError::OciImageError(format!("Failed to read config blob: {}", e)))?;
let config = Config::new(config_data, manifest.config.media_type.clone(), None);
let mut layers = Vec::new();
for layer_desc in &manifest.layers {
let layer_digest_hex = layer_desc
.digest
.strip_prefix("sha256:")
.unwrap_or(&layer_desc.digest);
let layer_data = std::fs::read(blobs_dir.join(layer_digest_hex)).map_err(|e| {
BoxError::OciImageError(format!(
"Failed to read layer blob {}: {}",
layer_desc.digest, e
))
})?;
tracing::debug!(
digest = %layer_desc.digest,
size = layer_data.len(),
"Read layer for push"
);
layers.push(ImageLayer::new(
layer_data,
layer_desc.media_type.clone(),
None,
));
}
let auth = self.auth.to_oci_auth();
let response: PushResponse = self
.client
.push(&oci_ref, &layers, config, &auth, Some(manifest))
.await
.map_err(|e| BoxError::RegistryError {
registry: reference.registry.clone(),
message: format!("Failed to push image: {}", e),
})?;
tracing::info!(
reference = %reference,
manifest_url = %response.manifest_url,
"Image pushed successfully"
);
Ok(PushResult {
config_url: response.config_url,
manifest_url: response.manifest_url,
manifest_digest: manifest_digest.to_string(),
})
}
fn to_oci_reference(&self, reference: &ImageReference) -> Result<Reference> {
let ref_str = if let Some(ref tag) = reference.tag {
format!("{}/{}:{}", reference.registry, reference.repository, tag)
} else {
format!("{}/{}:latest", reference.registry, reference.repository)
};
ref_str.parse::<Reference>().map_err(|e| {
BoxError::OciImageError(format!("Invalid OCI reference '{}': {}", ref_str, e))
})
}
}
fn resolve_target_arch(platform: Option<&str>) -> String {
let raw = match platform {
Some(p) => p
.split('/')
.nth(1)
.or_else(|| p.split('/').next())
.unwrap_or(p)
.to_string(),
None => std::env::consts::ARCH.to_string(),
};
match raw.as_str() {
"x86_64" => "amd64".to_string(),
"aarch64" => "arm64".to_string(),
other => other.to_string(),
}
}
fn platform_resolver_for(arch: String) -> impl Fn(&[ImageIndexEntry]) -> Option<String> {
move |manifests: &[ImageIndexEntry]| {
manifests
.iter()
.find(|entry| {
entry
.platform
.as_ref()
.is_some_and(|p| p.os == "linux" && p.architecture == arch)
})
.map(|entry| entry.digest.clone())
}
}
fn linux_platform_resolver(manifests: &[ImageIndexEntry]) -> Option<String> {
platform_resolver_for(resolve_target_arch(None))(manifests)
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::{Mutex, OnceLock};
fn env_lock() -> std::sync::MutexGuard<'static, ()> {
static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
LOCK.get_or_init(|| Mutex::new(())).lock().unwrap()
}
#[test]
fn test_registry_auth_anonymous() {
let auth = RegistryAuth::anonymous();
assert!(auth.username.is_none());
assert!(auth.password.is_none());
}
#[test]
fn test_resolve_target_arch() {
assert_eq!(resolve_target_arch(Some("linux/arm64")), "arm64");
assert_eq!(resolve_target_arch(Some("linux/amd64")), "amd64");
assert_eq!(resolve_target_arch(Some("linux/arm64/v8")), "arm64");
assert_eq!(resolve_target_arch(Some("arm64")), "arm64");
assert_eq!(resolve_target_arch(Some("linux/x86_64")), "amd64");
assert_eq!(resolve_target_arch(Some("aarch64")), "arm64");
assert!(!resolve_target_arch(None).is_empty());
}
#[tokio::test]
async fn test_hashing_file_writer_matches_sha256() {
use sha2::{Digest, Sha256};
use tokio::io::AsyncWriteExt;
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("blob");
let payload = b"a3s-box streaming blob hash test payload";
let mut writer = HashingFileWriter::new(tokio::fs::File::create(&path).await.unwrap());
writer.write_all(&payload[..10]).await.unwrap();
writer.write_all(&payload[10..]).await.unwrap();
writer.flush().await.unwrap();
writer.shutdown().await.unwrap();
let streamed = writer.finalize_hex();
let expected = format!("{:x}", Sha256::digest(payload));
assert_eq!(
streamed, expected,
"streamed hash must equal sha256(payload)"
);
assert_eq!(std::fs::read(&path).unwrap(), payload);
}
#[test]
fn test_registry_auth_basic() {
let auth = RegistryAuth::basic("user", "pass");
assert_eq!(auth.username, Some("user".to_string()));
assert_eq!(auth.password, Some("pass".to_string()));
}
#[test]
fn test_registry_auth_to_oci_anonymous() {
let auth = RegistryAuth::anonymous();
let oci_auth = auth.to_oci_auth();
assert!(matches!(oci_auth, OciRegistryAuth::Anonymous));
}
#[test]
fn test_registry_auth_to_oci_basic() {
let auth = RegistryAuth::basic("user", "pass");
let oci_auth = auth.to_oci_auth();
assert!(matches!(oci_auth, OciRegistryAuth::Basic(_, _)));
}
#[test]
fn test_registry_protocol_defaults_to_https() {
let _guard = env_lock();
std::env::remove_var(REGISTRY_PROTOCOL_ENV);
assert!(matches!(
registry_protocol_from_env(),
ClientProtocol::Https
));
}
#[test]
fn test_registry_protocol_can_use_http_for_local_testing() {
let _guard = env_lock();
std::env::set_var(REGISTRY_PROTOCOL_ENV, "http");
assert!(matches!(registry_protocol_from_env(), ClientProtocol::Http));
std::env::remove_var(REGISTRY_PROTOCOL_ENV);
}
#[test]
fn test_registry_protocol_rejects_unknown_values_to_https() {
let _guard = env_lock();
std::env::set_var(REGISTRY_PROTOCOL_ENV, "ftp");
assert!(matches!(
registry_protocol_from_env(),
ClientProtocol::Https
));
std::env::remove_var(REGISTRY_PROTOCOL_ENV);
}
#[test]
fn test_to_oci_reference_with_tag() {
let puller = RegistryPuller::new();
let img_ref = ImageReference {
registry: "ghcr.io".to_string(),
repository: "a3s-box/code".to_string(),
tag: Some("v0.1.0".to_string()),
digest: None,
};
let oci_ref = puller.to_oci_reference(&img_ref).unwrap();
assert_eq!(oci_ref.to_string(), "ghcr.io/a3s-box/code:v0.1.0");
}
#[test]
fn test_to_oci_reference_with_digest() {
let puller = RegistryPuller::new();
let img_ref = ImageReference {
registry: "ghcr.io".to_string(),
repository: "a3s-box/code".to_string(),
tag: None,
digest: Some(
"sha256:abcdef1234567890abcdef1234567890abcdef1234567890abcdef1234567890"
.to_string(),
),
};
let oci_ref = puller.to_oci_reference(&img_ref).unwrap();
let ref_str = oci_ref.to_string();
assert!(ref_str.contains("sha256:"));
}
#[test]
fn test_to_oci_reference_default_tag() {
let puller = RegistryPuller::new();
let img_ref = ImageReference {
registry: "docker.io".to_string(),
repository: "library/nginx".to_string(),
tag: None,
digest: None,
};
let oci_ref = puller.to_oci_reference(&img_ref).unwrap();
assert!(oci_ref.to_string().contains("latest"));
}
}