From 28436b5e9a8e080980e57d9820a5ceec58b37928 Mon Sep 17 00:00:00 2001 From: Nilstrieb <48135649+Nilstrieb@users.noreply.github.com> Date: Fri, 15 Mar 2024 22:29:02 +0100 Subject: [PATCH] make download better --- Cargo.lock | 146 +++++++++++++++++++++++++++++++++++++++++++++++++++ Cargo.toml | 4 ++ src/image.rs | 39 ++++++++++++++ src/util.rs | 22 +++++++- 4 files changed, 210 insertions(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index 697d27c..8bcd892 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -95,10 +95,14 @@ dependencies = [ "cfg-if", "clap", "eyre", + "futures", + "hex", + "humansize", "monostate", "reqwest", "serde", "serde_json", + "sha2", "tokio", "tracing", "tracing-subscriber", @@ -137,6 +141,15 @@ version = "1.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + [[package]] name = "bumpalo" version = "3.15.4" @@ -223,6 +236,15 @@ version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "06ea2b9bc92be3c2baa9334a323ebca2d6f074ff852cd1d7b11064035cd3868f" +[[package]] +name = "cpufeatures" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53fe5e26ff1b7aef8bca9c6080520cfb8d9333c7568e1829cef191a9723e5504" +dependencies = [ + "libc", +] + [[package]] name = "crc32fast" version = "1.4.0" @@ -232,6 +254,26 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "crypto-common" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bfb12502f3fc46cca1bb51ac28df9d618d813cdc3d2f25b9fe775a34af26bb3" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + [[package]] name = "encoding_rs" version = "0.8.33" @@ -282,6 +324,21 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "futures" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "645c6916888f6cb6350d2550b80fb63e734897a8498abe35cfb732b6487804b0" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.30" @@ -289,6 +346,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eac8f7d7865dcb88bd4373ab671c8cf4508703796caa2b1985a9ca867b3fcb78" dependencies = [ "futures-core", + "futures-sink", ] [[package]] @@ -297,6 +355,34 @@ version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dfc6580bb841c5a68e9ef15c77ccc837b40a7504914d52e47b8b0e9bbda25a1d" +[[package]] +name = "futures-executor" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a576fc72ae164fca6b9db127eaa9a9dda0d61316034f33a0a0d4eda41f02b01d" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a44623e20b9681a318efdd71c299b6b222ed6f231972bfe2f224ebad6311f0c1" + +[[package]] +name = "futures-macro" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87750cf4b7a4c0625b1529e4c543c2182106e4dedc60a2a6455e00d212c489ac" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "futures-sink" version = "0.3.30" @@ -315,10 +401,26 @@ version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d6401deb83407ab3da39eba7e33987a73c3df0c82b4bb5813ee871c19c41d48" dependencies = [ + "futures-channel", "futures-core", + "futures-io", + "futures-macro", + "futures-sink", "futures-task", + "memchr", "pin-project-lite", "pin-utils", + "slab", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", ] [[package]] @@ -375,6 +477,12 @@ version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024" +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + [[package]] name = "http" version = "0.2.12" @@ -409,6 +517,15 @@ version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" +[[package]] +name = "humansize" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6cb51c9a029ddc91b07a787f1d86b53ccfa49b0e86688c946ebe8d3555685dd7" +dependencies = [ + "libm", +] + [[package]] name = "hyper" version = "0.14.28" @@ -506,6 +623,12 @@ version = "0.2.153" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9c198f91728a82281a64e1f4f9eeb25d82cb32a5de251c6bd1b5154d63a8e7bd" +[[package]] +name = "libm" +version = "0.2.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ec2a862134d2a7d32d7983ddcdd1c4923530833c9f2ea1a44fc5fa473989058" + [[package]] name = "lock_api" version = "0.4.11" @@ -897,6 +1020,17 @@ dependencies = [ "serde", ] +[[package]] +name = "sha2" +version = "0.10.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "793db75ad2bcafc3ffa7c68b215fee268f537982cd901d132f89c6343f3a3dc8" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + [[package]] name = "sharded-slab" version = "0.1.7" @@ -1142,6 +1276,12 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "typenum" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825" + [[package]] name = "unicode-bidi" version = "0.3.15" @@ -1192,6 +1332,12 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "830b7e5d4d90034032940e4ace0d9a9a057e7a45cd94e6c007832e39edb82f6d" +[[package]] +name = "version_check" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f" + [[package]] name = "want" version = "0.3.1" diff --git a/Cargo.toml b/Cargo.toml index 656864d..6452a08 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,10 +8,14 @@ bytes = "1.5.0" cfg-if = "1.0.0" clap = { version = "4.5.3", features = ["derive"] } eyre = "0.6.12" +futures = "0.3.30" +hex = "0.4.3" +humansize = "2.1.3" monostate = "0.1.11" reqwest = { version = "0.11.26", default-features = false, features = ["json", "rustls-tls", "gzip"] } serde = { version = "1.0.197", features = ["derive"] } serde_json = "1.0.114" +sha2 = "0.10.8" tokio = { version = "1.36.0", features = ["full"] } tracing = "0.1.40" tracing-subscriber = { version = "0.3.18", features = ["env-filter"] } diff --git a/src/image.rs b/src/image.rs index 46b01c6..79750bf 100644 --- a/src/image.rs +++ b/src/image.rs @@ -10,6 +10,7 @@ use reqwest::{ }; use serde::{Deserialize, Serialize}; use tokio::fs; +use tracing::info; use crate::util; @@ -215,6 +216,27 @@ impl Client { .write_blob(&manifest.config.digest, &config_blob) .await?; + for layer in manifest.layers { + if !writer.already_exists(&layer.digest).await? { + info!( + "{} Downloading... ({})", + layer.digest, + humansize::format_size(layer.size, humansize::DECIMAL) + ); + + let content = self + .get_blob(image, &layer.digest) + .await + .wrap_err("getting layer")?; + writer + .write_blob(&layer.digest, &content) + .await + .wrap_err("writing blob")?; + } else { + info!("{} Skipping", layer.digest); + } + } + Ok(()) } } @@ -276,6 +298,7 @@ pub struct OciImageConfigConfig { // tty: bool, } +#[derive(Clone)] pub struct ImageLayoutWriter { // look, this could be a `Dir`` but that's too annoying right now dir: PathBuf, @@ -301,6 +324,22 @@ impl ImageLayoutWriter { Ok(Self { dir }) } + pub async fn already_exists(&self, digest: &str) -> Result { + let (alg, encoded) = digest + .split_once(":") + .wrap_err_with(|| format!("digest {digest} does not have ALG:ENCODED format"))?; + let blob = self.dir.join("blobs").join(alg).join(encoded); + let content = fs::read(blob).await; + Ok(match content { + Ok(content) => { + let digest = util::digest(alg, content).await?; + // if the digest doesn't match, try again. + digest.eq_ignore_ascii_case(encoded) + } + _ => false, + }) + } + pub async fn write_blob(&self, digest: &str, blob_content: &[u8]) -> Result<()> { let (alg, encoded) = digest .split_once(":") diff --git a/src/util.rs b/src/util.rs index b26c3b0..45d6539 100644 --- a/src/util.rs +++ b/src/util.rs @@ -1,6 +1,7 @@ use std::path::Path; -use eyre::Context; +use eyre::{bail, Context}; +use sha2::Digest; pub fn host_oci_arch() -> &'static str { cfg_if::cfg_if! { @@ -30,3 +31,22 @@ pub async fn write_file(path: impl AsRef, content: impl AsRef<[u8]>) -> ey .await .wrap_err_with(|| format!("writing to {}", path.display())) } + +pub async fn digest(alg: &str, content: Vec) -> eyre::Result { + let alg = alg.to_owned(); + tokio::task::spawn_blocking(move || match alg.as_str() { + "sha256" => { + let mut hasher = sha2::Sha256::new(); + hasher.update(content); + Ok(hex::encode(hasher.finalize())) + } + "sha512" => { + let mut hasher = sha2::Sha512::new(); + hasher.update(content); + Ok(hex::encode(hasher.finalize())) + } + _ => bail!("unrecognized hashing algorithm '{alg}'"), + }) + .await + .wrap_err("failed to spawn blocking task")? +}