worker/jobs/downloads: Extract CdnLogStorageConfig enum

This allows us to use in-memory storage for testing purposes and fall back to the local file system for development.

Author: Turbo87

src/config.rs

@@ -1,11 +1,13 @@
 mod balance_capacity;
 mod base;
+mod cdn_log_storage;
 mod database_pools;
 mod sentry;
 mod server;
 
 pub use self::balance_capacity::BalanceCapacityConfig;
 pub use self::base::Base;
+pub use self::cdn_log_storage::CdnLogStorageConfig;
 pub use self::database_pools::{DatabasePools, DbPoolConfig};
 pub use self::sentry::SentryConfig;
 pub use self::server::Server;

src/config/cdn_log_storage.rs

@@ -0,0 +1,53 @@
+use anyhow::Context;
+use crates_io_env_vars::{required_var, var};
+use secrecy::SecretString;
+use std::path::PathBuf;
+
+#[derive(Debug, Clone)]
+pub enum CdnLogStorageConfig {
+    S3 {
+        access_key: String,
+        secret_key: SecretString,
+    },
+    Local {
+        path: PathBuf,
+    },
+    Memory,
+}
+
+impl CdnLogStorageConfig {
+    pub fn s3(access_key: String, secret_key: SecretString) -> Self {
+        Self::S3 {
+            access_key,
+            secret_key,
+        }
+    }
+
+    pub fn local(path: PathBuf) -> Self {
+        Self::Local { path }
+    }
+
+    pub fn memory() -> Self {
+        Self::Memory
+    }
+
+    pub fn from_env() -> anyhow::Result<Self> {
+        if let Some(access_key) = var("AWS_ACCESS_KEY")? {
+            let secret_key = required_var("AWS_SECRET_KEY")?.into();
+            return Ok(Self::s3(access_key, secret_key));
+        }
+
+        let current_dir = std::env::current_dir();
+        let current_dir = current_dir.context("Failed to read the current directory")?;
+
+        let path = current_dir.join("local_uploads");
+        let path_display = path.display();
+        if path.exists() {
+            info!("Falling back to local CDN log storage at {path_display}");
+            return Ok(Self::local(path));
+        }
+
+        warn!("Falling back to in-memory CDN log storage because {path_display} does not exist");
+        Ok(Self::memory())
+    }
+}

src/config/server.rs

@@ -8,6 +8,7 @@ use crate::Env;
 use super::base::Base;
 use super::database_pools::DatabasePools;
 use crate::config::balance_capacity::BalanceCapacityConfig;
+use crate::config::cdn_log_storage::CdnLogStorageConfig;
 use crate::middleware::cargo_compat::StatusCodeConfig;
 use crate::storage::StorageConfig;
 use crates_io_env_vars::{list, list_parsed, required_var, var, var_parsed};
@@ -34,6 +35,7 @@ pub struct Server {
     pub max_blocking_threads: Option<usize>,
     pub db: DatabasePools,
     pub storage: StorageConfig,
+    pub cdn_log_storage: CdnLogStorageConfig,
     pub session_key: cookie::Key,
     pub gh_client_id: ClientId,
     pub gh_client_secret: ClientSecret,
@@ -172,6 +174,7 @@ impl Server {
         Ok(Server {
             db: DatabasePools::full_from_environment(&base)?,
             storage,
+            cdn_log_storage: CdnLogStorageConfig::from_env()?,
             base,
             ip,
             port,

src/tests/util/test_app.rs

@@ -2,7 +2,9 @@ use super::{MockAnonymousUser, MockCookieUser, MockTokenUser};
 use crate::util::chaosproxy::ChaosProxy;
 use crate::util::github::{MockGitHubClient, MOCK_GITHUB_DATA};
 use anyhow::Context;
-use crates_io::config::{self, BalanceCapacityConfig, Base, DatabasePools, DbPoolConfig};
+use crates_io::config::{
+    self, BalanceCapacityConfig, Base, CdnLogStorageConfig, DatabasePools, DbPoolConfig,
+};
 use crates_io::middleware::cargo_compat::StatusCodeConfig;
 use crates_io::models::token::{CrateScope, EndpointScope};
 use crates_io::rate_limiter::{LimitedAction, RateLimiterConfig};
@@ -422,6 +424,7 @@ fn simple_config() -> config::Server {
         max_blocking_threads: None,
         db,
         storage,
+        cdn_log_storage: CdnLogStorageConfig::memory(),
         session_key: cookie::Key::derive_from("test this has to be over 32 bytes long".as_bytes()),
         gh_client_id: ClientId::new(dotenvy::var("GH_CLIENT_ID").unwrap_or_default()),
         gh_client_secret: ClientSecret::new(dotenvy::var("GH_CLIENT_SECRET").unwrap_or_default()),

src/worker/jobs/downloads.rs

@@ -1,9 +1,11 @@
+use crate::config::CdnLogStorageConfig;
 use crate::worker::Environment;
 use anyhow::Context;
 use crates_io_cdn_logs::{count_downloads, Decompressor};
-use crates_io_env_vars::required_var;
 use crates_io_worker::BackgroundJob;
 use object_store::aws::AmazonS3Builder;
+use object_store::local::LocalFileSystem;
+use object_store::memory::InMemory;
 use object_store::ObjectStore;
 use std::cmp::Reverse;
 use std::collections::HashSet;
@@ -33,16 +35,9 @@ impl BackgroundJob for ProcessCdnLog {
 
     type Context = Arc<Environment>;
 
-    async fn run(&self, _ctx: Self::Context) -> anyhow::Result<()> {
-        let access_key = required_var("AWS_ACCESS_KEY")?;
-        let secret_key = required_var("AWS_SECRET_KEY")?;
-
-        let store = AmazonS3Builder::new()
-            .with_region(&self.region)
-            .with_bucket_name(&self.bucket)
-            .with_access_key_id(access_key)
-            .with_secret_access_key(secret_key)
-            .build()
+    async fn run(&self, ctx: Self::Context) -> anyhow::Result<()> {
+        let store = self
+            .build_store(&ctx.config.cdn_log_storage)
             .context("Failed to build object store")?;
 
         let path = object_store::path::Path::parse(&self.path)
@@ -104,3 +99,29 @@ impl BackgroundJob for ProcessCdnLog {
         Ok(())
     }
 }
+
+impl ProcessCdnLog {
+    fn build_store(&self, config: &CdnLogStorageConfig) -> anyhow::Result<Box<dyn ObjectStore>> {
+        match config {
+            CdnLogStorageConfig::S3 {
+                access_key,
+                secret_key,
+            } => {
+                use secrecy::ExposeSecret;
+
+                let store = AmazonS3Builder::new()
+                    .with_region(&self.region)
+                    .with_bucket_name(&self.bucket)
+                    .with_access_key_id(access_key)
+                    .with_secret_access_key(secret_key.expose_secret())
+                    .build()?;
+
+                Ok(Box::new(store))
+            }
+            CdnLogStorageConfig::Local { path } => {
+                Ok(Box::new(LocalFileSystem::new_with_prefix(path)?))
+            }
+            CdnLogStorageConfig::Memory => Ok(Box::new(InMemory::new())),
+        }
+    }
+}