Add support for downloading prebuilt blobs from Buildomat (#51)

gjcolombo · web-flow · commit a425676edcde · 2023-04-10T13:37:54.000-07:00
Allow packaging TOMLs to define one or more `buildomat_blobs` that refer to a
Buildomat artifact. This artifact need not be a zone image.

The existing S3 blob type is not changed, and buildomat blobs are optional, so
this should be a non-breaking change for existing package manifests.

Tested: cargo test; picked up the changes into Propolis's packaging utility
and verified that they can be used to pick up guest firmware images produced
by Buildomat.
diff --git a/Cargo.toml b/Cargo.toml
@@ -19,7 +19,9 @@ chrono = "0.4.24"
 filetime = "0.2"
 flate2 = "1.0.25"
 futures-util = "0.3"
+hex = "0.4.3"
 reqwest = { version = "0.11", default-features = false, features = ["rustls-tls", "stream"] }
+ring = "0.16.20"
 semver = { version = "1.0.17", features = ["std", "serde"] }
 serde = { version = "1.0", features = [ "derive" ] }
 serde_derive = "1.0"
diff --git a/src/blob.rs b/src/blob.rs
@@ -4,85 +4,138 @@
 
 //! Tools for downloading blobs
 
-use anyhow::{anyhow, Result};
+use anyhow::{anyhow, Context, Result};
 use chrono::{DateTime, FixedOffset, Utc};
 use futures_util::StreamExt;
 use reqwest::header::{CONTENT_LENGTH, LAST_MODIFIED};
-use std::path::Path;
+use ring::digest::{Context as DigestContext, Digest, SHA256};
+use std::path::{Path, PathBuf};
 use std::str::FromStr;
-use tokio::io::AsyncWriteExt;
+use tokio::io::{AsyncReadExt, AsyncWriteExt, BufReader};
 
-use crate::progress::Progress;
+use crate::progress::{NoProgress, Progress};
 
 // Path to the blob S3 Bucket.
 const S3_BUCKET: &str = "https://oxide-omicron-build.s3.amazonaws.com";
 // Name for the directory component where downloaded blobs are stored.
 pub(crate) const BLOB: &str = "blob";
 
+#[derive(Debug)]
+pub enum Source<'a> {
+    S3(&'a PathBuf),
+    Buildomat(&'a crate::package::PrebuiltBlob),
+}
+
+impl<'a> Source<'a> {
+    pub(crate) fn get_url(&self) -> String {
+        match self {
+            Self::S3(s) => format!("{}/{}", S3_BUCKET, s.to_string_lossy()),
+            Self::Buildomat(spec) => {
+                format!(
+                    "https://buildomat.eng.oxide.computer/public/file/oxidecomputer/{}/{}/{}/{}",
+                    spec.repo, spec.series, spec.commit, spec.artifact
+                )
+            }
+        }
+    }
+
+    async fn download_required(
+        &self,
+        url: &str,
+        client: &reqwest::Client,
+        destination: &Path,
+    ) -> Result<bool> {
+        if !destination.exists() {
+            return Ok(true);
+        }
+
+        match self {
+            Self::S3(_) => {
+                // Issue a HEAD request to get the blob's size and last modified
+                // time. If these match what's on disk, assume the blob is
+                // current and don't re-download it.
+                let head_response = client
+                    .head(url)
+                    .send()
+                    .await?
+                    .error_for_status()
+                    .with_context(|| format!("HEAD failed for {}", url))?;
+                let headers = head_response.headers();
+                let content_length = headers
+                    .get(CONTENT_LENGTH)
+                    .ok_or_else(|| anyhow!("no content length on {} HEAD response!", url))?;
+                let content_length: u64 = u64::from_str(content_length.to_str()?)?;
+
+                // From S3, header looks like:
+                //
+                //    "Last-Modified: Fri, 27 May 2022 20:50:17 GMT"
+                let last_modified = headers
+                    .get(LAST_MODIFIED)
+                    .ok_or_else(|| anyhow!("no last modified on {} HEAD response!", url))?;
+                let last_modified: DateTime<FixedOffset> =
+                    chrono::DateTime::parse_from_rfc2822(last_modified.to_str()?)?;
+
+                let metadata = tokio::fs::metadata(&destination).await?;
+                let metadata_modified: DateTime<Utc> = metadata.modified()?.into();
+
+                Ok(metadata.len() != content_length || metadata_modified != last_modified)
+            }
+            Self::Buildomat(blob_spec) => {
+                let digest = get_sha256_digest(destination).await?;
+                let expected_digest = hex::decode(&blob_spec.sha256)?;
+                Ok(digest.as_ref() != expected_digest)
+            }
+        }
+    }
+}
+
 // Downloads "source" from S3_BUCKET to "destination".
-pub async fn download(progress: &impl Progress, source: &str, destination: &Path) -> Result<()> {
+pub async fn download<'a>(
+    progress: &impl Progress,
+    source: &Source<'a>,
+    destination: &Path,
+) -> Result<()> {
     let blob = destination
         .file_name()
         .ok_or_else(|| anyhow!("missing blob filename"))?;
 
-    let url = format!("{}/{}", S3_BUCKET, source);
+    let url = source.get_url();
     let client = reqwest::Client::new();
-
-    let head_response = client.head(&url).send().await?.error_for_status()?;
-    let headers = head_response.headers();
-
-    // From S3, header looks like:
-    //
-    //    "Content-Length: 49283072"
-    let content_length = headers
-        .get(CONTENT_LENGTH)
-        .ok_or_else(|| anyhow!("no content length on {} HEAD response!", url))?;
-    let mut content_length: u64 = u64::from_str(content_length.to_str()?)?;
-
-    if destination.exists() {
-        // If destination exists, check against size and last modified time. If
-        // both are the same, then return Ok
-
-        // From S3, header looks like:
-        //
-        //    "Last-Modified: Fri, 27 May 2022 20:50:17 GMT"
-        let last_modified = headers
-            .get(LAST_MODIFIED)
-            .ok_or_else(|| anyhow!("no last modified on {} HEAD response!", url))?;
-        let last_modified: DateTime<FixedOffset> =
-            chrono::DateTime::parse_from_rfc2822(last_modified.to_str()?)?;
-        let metadata = tokio::fs::metadata(&destination).await?;
-        let metadata_modified: DateTime<Utc> = metadata.modified()?.into();
-
-        if metadata.len() == content_length && metadata_modified == last_modified {
-            return Ok(());
-        }
+    if !source.download_required(&url, &client, destination).await? {
+        return Ok(());
     }
 
     let response = client.get(url).send().await?.error_for_status()?;
     let response_headers = response.headers();
 
     // Grab update Content-Length from response headers, if present.
     // We only use it as a hint for the progress so no need to fail.
-    if let Some(Ok(Ok(resp_len))) = response_headers
+    let content_length = if let Some(Ok(Ok(resp_len))) = response_headers
         .get(CONTENT_LENGTH)
         .map(|c| c.to_str().map(u64::from_str))
     {
-        content_length = resp_len;
-    }
-
-    // Store modified time from HTTPS response
-    let last_modified = response_headers
-        .get(LAST_MODIFIED)
-        .ok_or_else(|| anyhow!("no last modified on GET response!"))?;
-    let last_modified: DateTime<FixedOffset> =
-        chrono::DateTime::parse_from_rfc2822(last_modified.to_str()?)?;
+        Some(resp_len)
+    } else {
+        None
+    };
+
+    // If the server advertised a last-modified time for the blob, save it here
+    // so that the downloaded blob's last-modified time can be set to it.
+    let last_modified = if let Some(time) = response_headers.get(LAST_MODIFIED) {
+        Some(chrono::DateTime::parse_from_rfc2822(time.to_str()?)?)
+    } else {
+        None
+    };
 
     // Write file bytes to destination
     let mut file = tokio::fs::File::create(destination).await?;
 
     // Create a sub-progress for the blob download
-    let blob_progress = progress.sub_progress(content_length);
+    let blob_progress = if let Some(length) = content_length {
+        progress.sub_progress(length)
+    } else {
+        Box::new(NoProgress)
+    };
     blob_progress.set_message(blob.to_string_lossy().into_owned().into());
 
     let mut stream = response.bytes_stream();
@@ -107,14 +160,39 @@ pub async fn download(progress: &impl Progress, source: &str, destination: &Path
     drop(file);
 
     // Set destination file's modified time based on HTTPS response
-    filetime::set_file_mtime(
-        destination,
-        filetime::FileTime::from_system_time(last_modified.into()),
-    )?;
+    if let Some(last_modified) = last_modified {
+        filetime::set_file_mtime(
+            destination,
+            filetime::FileTime::from_system_time(last_modified.into()),
+        )?;
+    }
 
     Ok(())
 }
 
+async fn get_sha256_digest(path: &Path) -> Result<Digest> {
+    let mut reader = BufReader::new(
+        tokio::fs::File::open(path)
+            .await
+            .with_context(|| format!("could not open {path:?}"))?,
+    );
+    let mut context = DigestContext::new(&SHA256);
+    let mut buffer = [0; 1024];
+
+    loop {
+        let count = reader
+            .read(&mut buffer)
+            .await
+            .with_context(|| format!("failed to read {path:?}"))?;
+        if count == 0 {
+            break;
+        } else {
+            context.update(&buffer[..count]);
+        }
+    }
+    Ok(context.finish())
+}
+
 #[test]
 fn test_converts() {
     let content_length = "1966080";
diff --git a/src/package.rs b/src/package.rs
@@ -172,6 +172,19 @@ async fn add_package_to_zone_archive(
     Ok(())
 }
 
+/// Describes a path to a Buildomat-generated artifact that should reside at
+/// the following path:
+///
+/// <https://buildomat.eng.oxide.computer/public/file/oxidecomputer/REPO/SERIES/COMMIT/ARTIFACT>
+#[derive(Deserialize, Debug)]
+pub struct PrebuiltBlob {
+    pub repo: String,
+    pub series: String,
+    pub commit: String,
+    pub artifact: String,
+    pub sha256: String,
+}
+
 /// Describes the origin of an externally-built package.
 #[derive(Deserialize, Debug)]
 #[serde(tag = "type", rename_all = "lowercase")]
@@ -182,6 +195,9 @@ pub enum PackageSource {
         /// within this package.
         blobs: Option<Vec<PathBuf>>,
 
+        /// A list of Buildomat blobs that should be placed in this package.
+        buildomat_blobs: Option<Vec<PrebuiltBlob>>,
+
         /// Configuration for packages containing Rust binaries.
         rust: Option<RustPackage>,
 
@@ -228,6 +244,16 @@ impl PackageSource {
             _ => None,
         }
     }
+
+    fn buildomat_blobs(&self) -> Option<&[PrebuiltBlob]> {
+        match self {
+            PackageSource::Local {
+                buildomat_blobs: Some(buildomat_blobs),
+                ..
+            } => Some(buildomat_blobs),
+            _ => None,
+        }
+    }
 }
 
 /// Describes the output format of the package.
@@ -439,11 +465,19 @@ impl Package {
         //
         // - 1 tick for each included path
         // - 1 tick per rust binary
-        // - 1 tick per blob + 1 tick for appending blob dir to archive
+        // - 1 tick per blob
+        // - 1 tick for appending the blob directory to the archive, but only if
+        //   there is at least one blob
         let progress_total = match &self.source {
-            PackageSource::Local { blobs, rust, paths } => {
-                let blob_work = blobs.as_ref().map(|b| b.len() + 1).unwrap_or(0);
-
+            PackageSource::Local {
+                blobs,
+                buildomat_blobs,
+                rust,
+                paths,
+            } => {
+                let blob_work = blobs.as_ref().map(|b| b.len()).unwrap_or(0);
+                let buildomat_work = buildomat_blobs.as_ref().map(|b| b.len()).unwrap_or(0);
+                let blob_dir_work = (blob_work != 0 || buildomat_work != 0) as usize;
                 let rust_work = rust.as_ref().map(|r| r.binary_names.len()).unwrap_or(0);
 
                 let mut paths_work = 0;
@@ -455,7 +489,7 @@ impl Package {
                         .count();
                 }
 
-                rust_work + blob_work + paths_work
+                rust_work + blob_work + buildomat_work + paths_work + blob_dir_work
             }
             _ => 1,
         };
@@ -631,19 +665,32 @@ impl Package {
         download_directory: &Path,
         destination_path: &Path,
     ) -> Result<()> {
+        let mut all_blobs = Vec::new();
         if let Some(blobs) = self.source.blobs() {
+            all_blobs.extend(blobs.iter().map(crate::blob::Source::S3));
+        }
+
+        if let Some(buildomat_blobs) = self.source.buildomat_blobs() {
+            all_blobs.extend(buildomat_blobs.iter().map(crate::blob::Source::Buildomat));
+        }
+
+        if !all_blobs.is_empty() {
             progress.set_message("downloading blobs".into());
             let blobs_path = download_directory.join(&self.service_name);
             std::fs::create_dir_all(&blobs_path)?;
-            stream::iter(blobs.iter())
+            stream::iter(all_blobs.iter())
                 .map(Ok)
                 .try_for_each_concurrent(None, |blob| {
-                    let blob_path = blobs_path.join(blob);
+                    let blob_path = match blob {
+                        blob::Source::S3(s) => blobs_path.join(s),
+                        blob::Source::Buildomat(spec) => blobs_path.join(&spec.artifact),
+                    };
+
                     async move {
-                        blob::download(progress, &blob.to_string_lossy(), &blob_path)
+                        blob::download(progress, blob, &blob_path)
                             .await
                             .with_context(|| {
-                                format!("failed to download blob: {}", blob.to_string_lossy())
+                                format!("failed to download blob: {}", blob.get_url())
                             })?;
                         progress.increment(1);
                         Ok::<_, anyhow::Error>(())
diff --git a/tests/mod.rs b/tests/mod.rs
@@ -231,11 +231,12 @@ mod test {
     async fn test_download() -> Result<()> {
         let out = tempfile::tempdir()?;
 
-        let url = "OVMF_CODE.fd";
-        let dst = out.path().join(url);
+        let path = PathBuf::from("OVMF_CODE.fd");
+        let src = omicron_zone_package::blob::Source::S3(&path);
+        let dst = out.path().join(&path);
 
-        download(&NoProgress, &url, &dst).await?;
-        download(&NoProgress, &url, &dst).await?;
+        download(&NoProgress, &src, &dst).await?;
+        download(&NoProgress, &src, &dst).await?;
 
         Ok(())
     }
