many changes

src/uploader.rs

@@ -1,36 +1,139 @@
 use anyhow::Result;
 use async_compression::{Level, tokio::bufread::ZstdEncoder};
-use ed25519_dalek;
+use aws_sdk_s3::{
+    self as s3,
+    types::{CompletedMultipartUpload, CompletedPart},
+};
+use futures::future::join_all;
+use log::debug;
 use nix_compat::{
-    narinfo::{self, NarInfo},
+    narinfo::{self, NarInfo, SigningKey},
     nixbase32,
     store_path::StorePath,
 };
 use sha2::{Digest, Sha256};
-use std::fs;
 use tokio::{io::AsyncReadExt, process::Command};
 
 use crate::path_info::PathInfo;
 
-pub struct Uploader {
-    signing_key: narinfo::SigningKey<ed25519_dalek::SigningKey>,
+const MULTIPART_CUTOFF: usize = 1024 * 1024 * 5;
+
+pub struct Uploader<'a> {
+    signing_key: &'a SigningKey<ed25519_dalek::SigningKey>,
     path: PathInfo,
-    compression: Option<String>,
+    s3_client: &'a s3::Client,
+    bucket: String,
 }
 
-impl Uploader {
-    pub fn new(key_file: &str, path: PathInfo) -> Result<Self> {
-        let key = fs::read_to_string(key_file)?;
-        let signing_key = narinfo::parse_keypair(key.as_str())?.0;
+impl<'a> Uploader<'a> {
+    pub fn new(
+        signing_key: &'a SigningKey<ed25519_dalek::SigningKey>,
+        path: PathInfo,
+        s3_client: &'a s3::Client,
+        bucket: String,
+    ) -> Result<Self> {
         Ok(Self {
             signing_key,
             path,
-            // TODO: support other algorithms
-            compression: Some("zstd".to_string()),
+            s3_client,
+            bucket,
         })
     }
 
-    pub async fn make_nar(&self) -> Result<Vec<u8>> {
+    pub async fn upload(&self) -> Result<()> {
+        let nar = self.make_nar().await?;
+        let mut nar_info = self.narinfo_from_nar(&nar)?;
+        let nar = self.compress_nar(&nar).await;
+
+        // update fields that we know after compression
+        nar_info.file_size = Some(nar.len() as u64);
+        let mut hasher = Sha256::new();
+        hasher.update(&nar);
+        nar_info.file_hash = Some(hasher.finalize().into());
+        let nar_url = self.nar_url(&nar);
+
+        if nar.len() < MULTIPART_CUTOFF {
+            let put_object = self
+                .s3_client
+                .put_object()
+                .bucket(&self.bucket)
+                .key(&nar_url)
+                .body(nar.into())
+                .send()
+                .await?;
+            debug!("put object: {:#?}", put_object);
+        } else {
+            let multipart = self
+                .s3_client
+                .create_multipart_upload()
+                .bucket(&self.bucket)
+                .key(&nar_url)
+                .send()
+                .await?;
+            let upload_id = multipart.upload_id().unwrap();
+
+            let mut parts = Vec::with_capacity(nar.len() / MULTIPART_CUTOFF);
+            let chunks = nar.array_chunks::<MULTIPART_CUTOFF>();
+            for (i, chunk) in chunks.enumerate() {
+                parts.push(tokio::task::spawn(
+                    self.s3_client
+                        .upload_part()
+                        .bucket(&self.bucket)
+                        .key(&nar_url)
+                        .upload_id(upload_id)
+                        .part_number(i as i32 + 1)
+                        .body(chunk.to_vec().into())
+                        .send(),
+                ));
+            }
+
+            let completed_parts = join_all(parts)
+                .await
+                .into_iter()
+                .flatten()
+                .collect::<Result<Vec<_>, _>>()?
+                .into_iter()
+                .enumerate()
+                .map(|(i, part)| {
+                    CompletedPart::builder()
+                        .set_e_tag(part.e_tag().map(ToString::to_string))
+                        .set_part_number(Some(i as i32 + 1))
+                        .set_checksum_sha256(part.checksum_sha256().map(ToString::to_string))
+                        .build()
+                })
+                .collect::<Vec<_>>();
+
+            let completed_mp_upload = CompletedMultipartUpload::builder()
+                .set_parts(Some(completed_parts))
+                .build();
+
+            let complete_mp_upload = self
+                .s3_client
+                .complete_multipart_upload()
+                .bucket(&self.bucket)
+                .key(&nar_url)
+                .upload_id(upload_id)
+                .multipart_upload(completed_mp_upload)
+                .send()
+                .await?;
+
+            debug!("complete multipart upload: {:#?}", complete_mp_upload);
+        }
+
+        nar_info.add_signature(self.signing_key);
+
+        self.s3_client
+            .put_object()
+            .bucket(&self.bucket)
+            .key(format!("{}.narinfo", self.path.digest()))
+            .body(nar_info.to_string().as_bytes().to_vec().into())
+            .send()
+            .await?;
+
+        Ok(())
+    }
+
+    async fn make_nar(&self) -> Result<Vec<u8>> {
         Ok(Command::new("nix")
             .arg("nar")
             .arg("dump-path")
@@ -40,7 +143,7 @@ impl Uploader {
             .stdout)
     }
 
-    pub fn narinfo_from_nar(&self, nar: &[u8]) -> Result<NarInfo> {
+    fn narinfo_from_nar(&self, nar: &[u8]) -> Result<NarInfo> {
         let mut hasher = Sha256::new();
         hasher.update(nar);
         let nar_hash: [u8; 32] = hasher.finalize().into();
@@ -54,7 +157,7 @@ impl Uploader {
             ca: self.path.ca.clone(),
             system: None,
             deriver: Some(self.path.deriver.as_ref()),
-            compression: self.compression.as_ref().map(String::as_str),
+            compression: Some("zstd"),
             file_hash: None,
             file_size: None,
             url: "",