From d6421db20c6561fcd57e2cecfba13360e65dd99c Mon Sep 17 00:00:00 2001
From: "David E. Wheeler" <david.wheeler@tembo.io>
Date: Tue, 16 Jan 2024 13:26:32 -0500
Subject: [PATCH] Validate tarbal checksum on upload

---
 registry/src/routes/extensions.rs | 9 ++++-----
 registry/src/uploader.rs          | 4 ++++
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/registry/src/routes/extensions.rs b/registry/src/routes/extensions.rs
index 5545032d..513a4eb4 100644
--- a/registry/src/routes/extensions.rs
+++ b/registry/src/routes/extensions.rs
@@ -291,17 +291,15 @@ pub async fn publish(
     // The uploaded contents in .tar.gz
     let gzipped_archive = file.freeze();
 
-    let digest = sha256::digest(&*gzipped_archive);
-
-    // Extract the .tar.gz and its relevant contentss
+    // Extract the .tar.gz and its relevant contents
     let (extension_views, pg_version) =
         extractor::extract_extension_view(&gzipped_archive, &new_extension).map_err(|err| {
             tracing::error!("Failed to decompress archive: {err}");
             ExtensionRegistryError::ArchiveError
         })?;
 
-    // TODO(ianstanton) Generate checksum
-    let file_byte_stream = ByteStream::from(gzipped_archive.clone());
+    let digest = sha256::digest(&*gzipped_archive);
+    let file_byte_stream = ByteStream::from(gzipped_archive);
     let client = aws_sdk_s3::Client::new(&aws_config);
     let uploaded_path = upload_extension(
         &cfg.bucket_name,
@@ -310,6 +308,7 @@ pub async fn publish(
         &new_extension,
         &new_extension.vers,
         pg_version,
+        &digest,
     )
     .await?;
 
diff --git a/registry/src/uploader.rs b/registry/src/uploader.rs
index 760c77ae..de2e2809 100644
--- a/registry/src/uploader.rs
+++ b/registry/src/uploader.rs
@@ -37,6 +37,7 @@ pub async fn upload(
     path: &str,
     content: ByteStream,
     content_type: &str,
+    sha256: &str,
 ) -> Result<PutObjectOutput, SdkError<PutObjectError>> {
     let obj = s3_client
         .put_object()
@@ -46,6 +47,7 @@ pub async fn upload(
         .key(path)
         .cache_control(CACHE_CONTROL_IMMUTABLE)
         .set_server_side_encryption(Some(Aes256))
+        .checksum_sha256(sha256)
         .send()
         .await;
     debug!("OBJECT: {:?}", obj);
@@ -62,6 +64,7 @@ pub async fn upload_extension(
     extension: &ExtensionUpload,
     extension_version: &semver::Version,
     pg_version: u8,
+    sha256: &str,
 ) -> Result<String, ExtensionRegistryError> {
     let path_in_bucket =
         extension_path(&extension.name, &extension_version.to_string(), pg_version);
@@ -72,6 +75,7 @@ pub async fn upload_extension(
         &path_in_bucket,
         file,
         "application/gzip",
+        sha256,
     )
     .await?;