From 7d65f9d24ce9eff7d4158fde1c16e788069f0bc6 Mon Sep 17 00:00:00 2001 From: Alex Zenla Date: Wed, 13 Mar 2024 14:20:22 +0000 Subject: [PATCH] krata: implement image seeding backend --- crates/kratart/examples/squashify.rs | 5 +- crates/kratart/src/image/compiler.rs | 12 +- crates/kratart/src/image/fetch.rs | 158 ++++++++++++++++++++++++--- crates/kratart/src/launch/mod.rs | 2 +- 4 files changed, 156 insertions(+), 21 deletions(-) diff --git a/crates/kratart/examples/squashify.rs b/crates/kratart/examples/squashify.rs index 89ae2ac..7ccf987 100644 --- a/crates/kratart/examples/squashify.rs +++ b/crates/kratart/examples/squashify.rs @@ -10,12 +10,15 @@ async fn main() -> Result<()> { env_logger::Builder::from_env(Env::default().default_filter_or("warn")).init(); let image = ImageName::parse(&args().nth(1).unwrap())?; + let seed = args().nth(2).map(PathBuf::from); + let cache_dir = PathBuf::from("krata-cache"); if !cache_dir.exists() { fs::create_dir(&cache_dir).await?; } + let cache = ImageCache::new(&cache_dir)?; - let compiler = ImageCompiler::new(&cache)?; + let compiler = ImageCompiler::new(&cache, seed)?; let info = compiler.compile(&image).await?; println!( "generated squashfs of {} to {}", diff --git a/crates/kratart/src/image/compiler.rs b/crates/kratart/src/image/compiler.rs index dffe1ac..cc3ca9a 100644 --- a/crates/kratart/src/image/compiler.rs +++ b/crates/kratart/src/image/compiler.rs @@ -55,11 +55,12 @@ impl ImageInfo { pub struct ImageCompiler<'a> { cache: &'a ImageCache, + seed: Option, } impl ImageCompiler<'_> { - pub fn new(cache: &ImageCache) -> Result { - Ok(ImageCompiler { cache }) + pub fn new(cache: &ImageCache, seed: Option) -> Result { + Ok(ImageCompiler { cache, seed }) } pub async fn compile(&self, image: &ImageName) -> Result { @@ -91,8 +92,11 @@ impl ImageCompiler<'_> { image_dir: &Path, squash_file: &Path, ) -> Result { - let downloader = - OciImageDownloader::new(layer_dir.to_path_buf(), OciRegistryPlatform::current()); + let downloader = OciImageDownloader::new( + self.seed.clone(), + layer_dir.to_path_buf(), + OciRegistryPlatform::current(), + ); let resolved = downloader.resolve(image.clone()).await?; let cache_key = format!( "manifest={}:squashfs-version={}\n", diff --git a/crates/kratart/src/image/fetch.rs b/crates/kratart/src/image/fetch.rs index 3fe3a72..7cf8f30 100644 --- a/crates/kratart/src/image/fetch.rs +++ b/crates/kratart/src/image/fetch.rs @@ -1,13 +1,20 @@ -use std::{path::PathBuf, pin::Pin}; +use std::{ + path::{Path, PathBuf}, + pin::Pin, +}; use anyhow::{anyhow, Result}; use async_compression::tokio::bufread::{GzipDecoder, ZstdDecoder}; use log::debug; -use oci_spec::image::{Descriptor, ImageConfiguration, ImageManifest, MediaType, ToDockerV2S2}; +use oci_spec::image::{ + Descriptor, ImageConfiguration, ImageIndex, ImageManifest, MediaType, ToDockerV2S2, +}; +use serde::de::DeserializeOwned; use tokio::{ fs::File, - io::{AsyncRead, BufReader}, + io::{AsyncRead, AsyncReadExt, BufReader}, }; +use tokio_stream::StreamExt; use tokio_tar::Archive; use super::{ @@ -16,6 +23,7 @@ use super::{ }; pub struct OciImageDownloader { + seed: Option, storage: PathBuf, platform: OciRegistryPlatform, } @@ -67,12 +75,122 @@ pub struct OciLocalImage { } impl OciImageDownloader { - pub fn new(storage: PathBuf, platform: OciRegistryPlatform) -> OciImageDownloader { - OciImageDownloader { storage, platform } + pub fn new( + seed: Option, + storage: PathBuf, + platform: OciRegistryPlatform, + ) -> OciImageDownloader { + OciImageDownloader { + seed, + storage, + platform, + } + } + + async fn load_seed_json_blob( + &self, + descriptor: &Descriptor, + ) -> Result> { + let digest = descriptor.digest(); + let Some((digest_type, digest_content)) = digest.split_once(':') else { + return Err(anyhow!("digest content was not properly formatted")); + }; + let want = format!("blobs/{}/{}", digest_type, digest_content); + self.load_seed_json(&want).await + } + + async fn load_seed_json(&self, want: &str) -> Result> { + let Some(ref seed) = self.seed else { + return Ok(None); + }; + + let file = File::open(seed).await?; + let mut archive = Archive::new(file); + let mut entries = archive.entries()?; + while let Some(entry) = entries.next().await { + let mut entry = entry?; + let path = String::from_utf8(entry.path_bytes().to_vec())?; + if path == want { + let mut content = String::new(); + entry.read_to_string(&mut content).await?; + let data = serde_json::from_str::(&content)?; + return Ok(Some(data)); + } + } + Ok(None) + } + + async fn extract_seed_blob(&self, descriptor: &Descriptor, to: &Path) -> Result { + let Some(ref seed) = self.seed else { + return Ok(false); + }; + + let digest = descriptor.digest(); + let Some((digest_type, digest_content)) = digest.split_once(':') else { + return Err(anyhow!("digest content was not properly formatted")); + }; + let want = format!("blobs/{}/{}", digest_type, digest_content); + + let seed = File::open(seed).await?; + let mut archive = Archive::new(seed); + let mut entries = archive.entries()?; + while let Some(entry) = entries.next().await { + let mut entry = entry?; + let path = String::from_utf8(entry.path_bytes().to_vec())?; + if path == want { + let mut file = File::create(to).await?; + tokio::io::copy(&mut entry, &mut file).await?; + return Ok(true); + } + } + Ok(false) } pub async fn resolve(&self, image: ImageName) -> Result { - debug!("download manifest image={}", image); + debug!("resolve manifest image={}", image); + + if let Some(index) = self.load_seed_json::("index.json").await? { + let mut found: Option<&Descriptor> = None; + for manifest in index.manifests() { + let Some(annotations) = manifest.annotations() else { + continue; + }; + + let Some(image_name) = annotations.get("io.containerd.image.name") else { + continue; + }; + + if *image_name != image.to_string() { + continue; + } + + if let Some(platform) = manifest.platform() { + if *platform.architecture() != self.platform.arch + || *platform.os() != self.platform.os + { + continue; + } + } + found = Some(manifest); + break; + } + + if let Some(found) = found { + if let Some(manifest) = self.load_seed_json_blob(found).await? { + debug!( + "found seeded manifest image={} manifest={}", + image, + found.digest() + ); + return Ok(OciResolvedImage { + name: image, + digest: found.digest().clone(), + manifest, + }); + } + } + } + let mut client = OciRegistryClient::new(image.registry_url()?, self.platform.clone())?; let (manifest, digest) = client .get_manifest_with_digest(&image.name, &image.reference) @@ -85,14 +203,23 @@ impl OciImageDownloader { } pub async fn download(&self, image: OciResolvedImage) -> Result { + let config: ImageConfiguration; + let mut client = OciRegistryClient::new(image.name.registry_url()?, self.platform.clone())?; - let config_bytes = client - .get_blob(&image.name.name, image.manifest.config()) - .await?; - let config: ImageConfiguration = serde_json::from_slice(&config_bytes)?; + if let Some(seeded) = self + .load_seed_json_blob::(image.manifest.config()) + .await? + { + config = seeded; + } else { + let config_bytes = client + .get_blob(&image.name.name, image.manifest.config()) + .await?; + config = serde_json::from_slice(&config_bytes)?; + } let mut layers = Vec::new(); for layer in image.manifest.layers() { - layers.push(self.download_layer(&image.name, layer, &mut client).await?); + layers.push(self.acquire_layer(&image.name, layer, &mut client).await?); } Ok(OciLocalImage { image, @@ -101,22 +228,23 @@ impl OciImageDownloader { }) } - async fn download_layer( + async fn acquire_layer( &self, image: &ImageName, layer: &Descriptor, client: &mut OciRegistryClient, ) -> Result { debug!( - "download layer digest={} size={}", + "acquire layer digest={} size={}", layer.digest(), layer.size() ); let mut layer_path = self.storage.clone(); layer_path.push(format!("{}.layer", layer.digest())); - { - let file = tokio::fs::File::create(&layer_path).await?; + let seeded = self.extract_seed_blob(layer, &layer_path).await?; + if !seeded { + let file = File::create(&layer_path).await?; let size = client.write_blob_to_file(&image.name, layer, file).await?; if layer.size() as u64 != size { return Err(anyhow!( diff --git a/crates/kratart/src/launch/mod.rs b/crates/kratart/src/launch/mod.rs index 16f685a..240f8c2 100644 --- a/crates/kratart/src/launch/mod.rs +++ b/crates/kratart/src/launch/mod.rs @@ -218,7 +218,7 @@ impl GuestLauncher { async fn compile(&self, image: &str, image_cache: &ImageCache) -> Result { let image = ImageName::parse(image)?; - let compiler = ImageCompiler::new(image_cache)?; + let compiler = ImageCompiler::new(image_cache, None)?; compiler.compile(&image).await }