feat: aws s3 as a file storage option

Cargo.toml

@@ -106,6 +106,9 @@ anyhow = { version = "1.0", default-features = false }
 app-memory-usage-fetcher = { version = "0.2.1", default-features = false }
 argon2 = { version = "0.5.3", default-features = false }
 axum = { version = "0.8.4", default-features = false }
+aws-config = { version = "1.8.5", default-features = false }
+aws-sdk-s3 = { version = "1.103.0", default-features = false }
+aws-smithy-async = { version = "1.2.5", default-features = false }
 axum-server = { version = "0.7.2", default-features = false }
 base64 = { version = "0.22.1", default-features = false }
 bytecount = { version = "0.6.9", default-features = false }

crates/server/Cargo.toml

@@ -28,6 +28,9 @@ aes-gcm = { workspace = true, features = ["aes", "alloc", "getrandom", "std"] }
 anyhow = { workspace = true, features = ["std"] }
 app-memory-usage-fetcher = { workspace = true }
 argon2 = { workspace = true, features = ["alloc", "password-hash", "std"] }
+aws-config = { workspace = true, features = ["behavior-version-latest"] }
+aws-sdk-s3 = { workspace = true }
+aws-smithy-async = { workspace = true, features = ["rt-tokio"] }
 axum = { workspace = true, features = ["http1", "http2", "json", "macros", "tokio"] }
 axum-server = { workspace = true, features = ["tls-rustls", "rustls"] }
 base64 = { workspace = true, features = ["alloc", "std"] }

crates/server/src/cloud.rs

@@ -0,0 +1,112 @@
+// SPDX-License-Identifier: Apache-2.0
+
+use std::fmt::{Debug, Formatter};
+
+use anyhow::Result;
+use aws_config::{AppName, BehaviorVersion, SdkConfig};
+use aws_sdk_s3::config::{Credentials, SharedCredentialsProvider};
+use aws_sdk_s3::primitives::ByteStream;
+use aws_sdk_s3::Client;
+use aws_smithy_async::rt::sleep::{SharedAsyncSleep, TokioSleep};
+use base64::prelude::BASE64_STANDARD;
+use base64::Engine;
+use sha2::{Digest, Sha256};
+
+/// Cloud storage options
+pub enum CloudStorage {
+    /// AWS S3
+    S3(AWS),
+}
+
+impl CloudStorage {
+    pub(crate) async fn write(&self, data: Vec<u8>, hashed_path: &str) -> Result<()> {
+        match self {
+            CloudStorage::S3(aws) => aws.write(data, hashed_path).await,
+        }
+    }
+
+    pub(crate) async fn read(&self, hashed_path: &str) -> Result<Vec<u8>> {
+        match self {
+            CloudStorage::S3(aws) => aws.read(hashed_path).await,
+        }
+    }
+}
+
+impl Debug for CloudStorage {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            CloudStorage::S3(_) => write!(f, "S3"),
+        }
+    }
+}
+
+/// AWS S3 connection
+#[derive(Debug)]
+pub struct AWS {
+    /// AWS Client object
+    client: Client,
+
+    /// S3 bucket name
+    bucket: String,
+}
+
+impl AWS {
+    /// Create a new AWS connection
+    ///
+    /// # Panics
+    ///
+    /// The code will not panic despite the call to `unwrap()`
+    #[must_use]
+    pub fn new(access_key_id: String, secret_access_key: String, bucket: String) -> Self {
+        let creds = Credentials::new(access_key_id, secret_access_key, None, None, "malwaredb");
+        let provider = SharedCredentialsProvider::new(creds);
+        let sleep_impl = SharedAsyncSleep::new(TokioSleep::new());
+        let config = SdkConfig::builder()
+            .credentials_provider(provider)
+            .behavior_version(BehaviorVersion::latest())
+            .sleep_impl(sleep_impl)
+            .app_name(AppName::new(constcat::concat!("malwaredb-v", crate::MDB_VERSION)).unwrap())
+            .build();
+
+        Self {
+            client: Client::new(&config),
+            bucket,
+        }
+    }
+
+    #[inline]
+    async fn write(&self, data: Vec<u8>, hashed_path: &str) -> Result<()> {
+        // Data might be compressed and/or encrypted, so re-hash for transmission to AWS
+        let mut hasher = Sha256::new();
+        hasher.update(&data);
+
+        let sha256_b64 = BASE64_STANDARD.encode(hasher.finalize());
+        let bytes = ByteStream::from(data);
+
+        self.client
+            .put_object()
+            .bucket(&self.bucket)
+            .checksum_sha256(sha256_b64)
+            .key(hashed_path)
+            .body(bytes)
+            .send()
+            .await?;
+
+        Ok(())
+    }
+
+    #[inline]
+    async fn read(&self, hashed_path: &str) -> Result<Vec<u8>> {
+        Ok(self
+            .client
+            .get_object()
+            .bucket(&self.bucket)
+            .key(hashed_path)
+            .send()
+            .await?
+            .body
+            .collect()
+            .await?
+            .to_vec())
+    }
+}

crates/server/src/db/mod.rs

@@ -1010,6 +1010,7 @@ mod tests {
                 }),
                 cert: None,
                 key: None,
+                cloud: None,
             };
 
             let vt: VtUpdater = state.try_into().expect("failed to create VtUpdater");
@@ -1071,6 +1072,7 @@ mod tests {
                 }),
                 cert: None,
                 key: None,
+                cloud: None,
             };
 
             let sqlite_second = Sqlite::new(DB_FILE)

crates/server/src/http/mod.rs

@@ -455,6 +455,7 @@ mod tests {
             vt_client: None,
             cert: None,
             key: None,
+            cloud: None,
         };
 
         state
@@ -516,6 +517,7 @@ mod tests {
                 vt_client: None,
                 cert: Some("../../testdata/server_ca_cert.pem".into()),
                 key: Some("../../testdata/server_key.pem".into()),
+                cloud: None,
             }
         } else {
             State {
@@ -536,6 +538,7 @@ mod tests {
                 vt_client: None,
                 cert: Some("../../testdata/server_cert.der".into()),
                 key: Some("../../testdata/server_key.der".into()),
+                cloud: None,
             }
         };
 

crates/server/src/lib.rs

@@ -5,6 +5,9 @@
 #![deny(clippy::all)]
 #![deny(clippy::pedantic)]
 
+/// Cloud storage logic
+pub mod cloud;
+
 /// Cryptographic functionality for file storage
 pub mod crypto;
 
@@ -21,6 +24,7 @@ pub mod utils;
 #[cfg(feature = "vt")]
 pub mod vt;
 
+use crate::cloud::CloudStorage;
 use crate::crypto::FileEncryption;
 use crate::db::MDBConfig;
 use malwaredb_api::ServerInfo;
@@ -87,6 +91,9 @@ pub struct State {
 
     /// Https private key file
     key: Option<PathBuf>,
+
+    /// AWS client and S3 bucket
+    cloud: Option<CloudStorage>,
 }
 
 impl State {
@@ -107,6 +114,7 @@ impl State {
         key: Option<PathBuf>,
         pg_cert: Option<PathBuf>,
         #[cfg(feature = "vt")] vt_client: Option<malwaredb_virustotal::VirusTotalClient>,
+        cloud: Option<CloudStorage>,
     ) -> Result<Self> {
         if let Some(dir) = &directory {
             if !dir.exists() {
@@ -147,6 +155,7 @@ impl State {
             started: SystemTime::now(),
             cert,
             key,
+            cloud,
         })
     }
 
@@ -157,10 +166,13 @@ impl State {
     /// * If the file can't be written.
     /// * If a necessary sub-directory can't be created.
     pub async fn store_bytes(&self, data: &[u8]) -> Result<bool> {
-        if let Some(dest_path) = &self.directory {
+        if self.directory.is_none() && self.cloud.is_none() {
+            Ok(false)
+        } else {
             let mut hasher = Sha256::new();
             hasher.update(data);
-            let sha256 = hex::encode(hasher.finalize());
+            let sha256_bytes = hasher.finalize();
+            let sha256 = hex::encode(sha256_bytes);
 
             // Trait `HashPath` needs to be re-worked so it can work with Strings.
             // This code below ends up making the String into ASCII representations of the hash
@@ -173,16 +185,6 @@ impl State {
                 sha256
             );
 
-            // The path which has the file name included, with the storage directory prepended.
-            //let hashed_path = result.hashed_path(3);
-            let mut dest_path = dest_path.clone();
-            dest_path.push(hashed_path);
-
-            // Remove the file name so we can just have the directory path.
-            let mut just_the_dir = dest_path.clone();
-            just_the_dir.pop();
-            std::fs::create_dir_all(just_the_dir)?;
-
             let data = if self.db_config.compression {
                 let buff = Cursor::new(data);
                 let mut compressed = Vec::with_capacity(data.len() / 2);
@@ -206,11 +208,23 @@ impl State {
                 data
             };
 
-            std::fs::write(dest_path, data)?;
+            if let Some(dest_path) = &self.directory {
+                // The path which has the file name included, with the storage directory prepended.
+                //let hashed_path = result.hashed_path(3);
+                let mut dest_path = dest_path.clone();
+                dest_path.push(hashed_path);
+
+                // Remove the file name so we can just have the directory path.
+                let mut just_the_dir = dest_path.clone();
+                just_the_dir.pop();
+                std::fs::create_dir_all(just_the_dir)?;
+
+                std::fs::write(dest_path, data)?;
+            } else if let Some(cloud) = &self.cloud {
+                cloud.write(data, &hashed_path).await?;
+            }
 
             Ok(true)
-        } else {
-            Ok(false)
         }
     }
 
@@ -222,53 +236,61 @@ impl State {
     /// * The file could not be read, maybe because it doesn't exist.
     /// * Failure to decrypt or decompress (corruption).
     pub async fn retrieve_bytes(&self, sha256: &String) -> Result<Vec<u8>> {
-        if let Some(dest_path) = &self.directory {
-            let path = format!(
-                "{}/{}/{}/{}",
-                &sha256[0..2],
-                &sha256[2..4],
-                &sha256[4..6],
-                sha256
-            );
-            // Trait `HashPath` needs to be re-worked so it can work with Strings.
-            // This code below ends up making the String into ASCII representations of the hash
-            // See: https://github.com/malwaredb/malwaredb-rs/issues/60
-            //let path = sha256.as_bytes().iter().hashed_path(3);
-            let contents = std::fs::read(dest_path.join(path))?;
+        if self.directory.is_none() && self.cloud.is_none() {
+            bail!("files are not saved")
+        }
 
-            let contents = if self.keys.is_empty() {
-                // We don't have file encryption enabled
-                contents
-            } else {
-                let (key_id, nonce) = self.db_type.get_file_encryption_key_id(sha256).await?;
-                if let Some(key_id) = key_id {
-                    if let Some(key) = self.keys.get(&key_id) {
-                        key.decrypt(&contents, nonce)?
-                    } else {
-                        bail!("File was encrypted but we don't have tke key!")
-                    }
+        // Trait `HashPath` needs to be re-worked so it can work with Strings.
+        // This code below ends up making the String into ASCII representations of the hash
+        // See: https://github.com/malwaredb/malwaredb-rs/issues/60
+        //let path = sha256.as_bytes().iter().hashed_path(3);
+
+        let path = format!(
+            "{}/{}/{}/{}",
+            &sha256[0..2],
+            &sha256[2..4],
+            &sha256[4..6],
+            sha256
+        );
+
+        let contents = if let Some(dest_path) = &self.directory {
+            std::fs::read(dest_path.join(path))?
+        } else if let Some(cloud) = &self.cloud {
+            cloud.read(&path).await?
+        } else {
+            bail!("No AWS config and no local directory, shouldn't happen!")
+        };
+
+        let contents = if self.keys.is_empty() {
+            // We don't have file encryption enabled
+            contents
+        } else {
+            let (key_id, nonce) = self.db_type.get_file_encryption_key_id(sha256).await?;
+            if let Some(key_id) = key_id {
+                if let Some(key) = self.keys.get(&key_id) {
+                    key.decrypt(&contents, nonce)?
                 } else {
-                    // File was not encrypted
-                    contents
+                    bail!("File was encrypted but we don't have tke key!")
                 }
-            };
-
-            if contents.starts_with(&GZIP_MAGIC) {
-                let buff = Cursor::new(contents);
-                let mut decompressor = GzDecoder::new(buff);
-                let mut decompressed: Vec<u8> = vec![];
-                decompressor.read_to_end(&mut decompressed)?;
-                Ok(decompressed)
-            } else if contents.starts_with(&ZSTD_MAGIC) {
-                let buff = Cursor::new(contents);
-                let mut decompressed: Vec<u8> = vec![];
-                zstd::stream::copy_decode(buff, &mut decompressed)?;
-                Ok(decompressed)
             } else {
-                Ok(contents)
+                // File was not encrypted
+                contents
             }
+        };
+
+        if contents.starts_with(&GZIP_MAGIC) {
+            let buff = Cursor::new(contents);
+            let mut decompressor = GzDecoder::new(buff);
+            let mut decompressed: Vec<u8> = vec![];
+            decompressor.read_to_end(&mut decompressed)?;
+            Ok(decompressed)
+        } else if contents.starts_with(&ZSTD_MAGIC) {
+            let buff = Cursor::new(contents);
+            let mut decompressed: Vec<u8> = vec![];
+            zstd::stream::copy_decode(buff, &mut decompressed)?;
+            Ok(decompressed)
         } else {
-            bail!("files are not saved")
+            Ok(contents)
         }
     }