make download better

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"

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"] }

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<bool> {
+        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(":")

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<Path>, content: impl AsRef<[u8]>) -> ey
         .await
         .wrap_err_with(|| format!("writing to {}", path.display()))
 }
+
+pub async fn digest(alg: &str, content: Vec<u8>) -> eyre::Result<String> {
+    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")?
+}