From 417bc38288cc7f50dc616f7abf8fdaded0d7846b Mon Sep 17 00:00:00 2001
From: mdecimus <mauro@stalw.art>
Date: Sat, 2 Mar 2024 09:33:36 +0100
Subject: [PATCH] Support for S3 bucket prefixes

---
 crates/store/src/backend/s3/mod.rs      | 26 +++++++++++++++++--------
 crates/utils/src/codec/base32_custom.rs | 10 +++++++++-
 2 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/crates/store/src/backend/s3/mod.rs b/crates/store/src/backend/s3/mod.rs
index d683ec92..547abc86 100644
--- a/crates/store/src/backend/s3/mod.rs
+++ b/crates/store/src/backend/s3/mod.rs
@@ -21,7 +21,7 @@
  * for more details.
 */
 
-use std::{ops::Range, time::Duration};
+use std::{io::Write, ops::Range, time::Duration};
 
 use s3::{
     creds::{error::CredentialsError, Credentials},
@@ -35,6 +35,7 @@ use utils::{
 
 pub struct S3Store {
     bucket: Bucket,
+    prefix: Option<String>,
 }
 
 impl S3Store {
@@ -67,6 +68,7 @@ impl S3Store {
             )?
             .with_path_style()
             .with_request_timeout(timeout),
+            prefix: config.value((&prefix, "key-prefix")).map(|s| s.to_string()),
         })
     }
 
@@ -75,7 +77,7 @@ impl S3Store {
         key: &[u8],
         range: Range<u32>,
     ) -> crate::Result<Option<Vec<u8>>> {
-        let path = Base32Writer::from_bytes(key).finalize();
+        let path = self.build_key(key);
         let response = if range.start != 0 || range.end != u32::MAX {
             self.bucket
                 .get_object_range(
@@ -102,11 +104,7 @@ impl S3Store {
     }
 
     pub(crate) async fn put_blob(&self, key: &[u8], data: &[u8]) -> crate::Result<()> {
-        match self
-            .bucket
-            .put_object(Base32Writer::from_bytes(key).finalize(), data)
-            .await
-        {
+        match self.bucket.put_object(self.build_key(key), data).await {
             Ok(response) if (200..300).contains(&response.status_code()) => Ok(()),
             Ok(response) => Err(crate::Error::InternalError(format!(
                 "S3 error code {}: {}",
@@ -119,11 +117,23 @@ impl S3Store {
 
     pub(crate) async fn delete_blob(&self, key: &[u8]) -> crate::Result<bool> {
         self.bucket
-            .delete_object(Base32Writer::from_bytes(key).finalize())
+            .delete_object(self.build_key(key))
             .await
             .map(|response| (200..300).contains(&response.status_code()))
             .map_err(|e| e.into())
     }
+
+    fn build_key(&self, key: &[u8]) -> String {
+        if let Some(prefix) = &self.prefix {
+            let mut writer =
+                Base32Writer::with_raw_capacity(prefix.len() + ((key.len() + 3) / 4 * 5));
+            writer.push_string(prefix);
+            writer.write_all(key).unwrap();
+            writer.finalize()
+        } else {
+            Base32Writer::from_bytes(key).finalize()
+        }
+    }
 }
 
 impl From<S3Error> for crate::Error {
diff --git a/crates/utils/src/codec/base32_custom.rs b/crates/utils/src/codec/base32_custom.rs
index 1c35a580..d9cba474 100644
--- a/crates/utils/src/codec/base32_custom.rs
+++ b/crates/utils/src/codec/base32_custom.rs
@@ -57,8 +57,12 @@ impl Base32Writer {
     }
 
     pub fn with_capacity(capacity: usize) -> Self {
+        Self::with_raw_capacity((capacity + 3) / 4 * 5)
+    }
+
+    pub fn with_raw_capacity(capacity: usize) -> Self {
         Base32Writer {
-            result: String::with_capacity((capacity + 3) / 4 * 5),
+            result: String::with_capacity(capacity),
             last_byte: 0,
             pos: 0,
         }
@@ -68,6 +72,10 @@ impl Base32Writer {
         self.result.push(ch);
     }
 
+    pub fn push_string(&mut self, string: &str) {
+        self.result.push_str(string);
+    }
+
     fn push_byte(&mut self, byte: u8, is_remainder: bool) {
         let (ch1, ch2) = match self.pos % 5 {
             0 => ((byte & 0xF8) >> 3, u8::MAX),