Merge pull request #54 from renegade-fi/joey/api-key-cache

auth-server: Add local api key cache

auth/auth-server/Cargo.toml

@@ -35,6 +35,7 @@ renegade-api = { package = "external-api", git = "https://github.com/renegade-fi
 # === Misc Dependencies === #
 base64 = "0.22.1"
 bytes = "1.0"
+cached = "0.53"
 chrono = { version = "0.4", features = ["serde"] }
 futures-util = "0.3"
 serde = { version = "1.0", features = ["derive"] }

auth/auth-server/src/models.rs

@@ -33,3 +33,15 @@ impl NewApiKey {
         Self { id, encrypted_key, description }
     }
 }
+
+impl From<NewApiKey> for ApiKey {
+    fn from(key: NewApiKey) -> Self {
+        Self {
+            id: key.id,
+            encrypted_key: key.encrypted_key,
+            description: key.description,
+            created_at: SystemTime::now(),
+            is_active: true,
+        }
+    }
+}

auth/auth-server/src/server/mod.rs

@@ -7,10 +7,11 @@ mod handle_key_management;
 mod helpers;
 mod queries;
 
-use crate::{error::AuthServerError, ApiError, Cli};
+use crate::{error::AuthServerError, models::ApiKey, ApiError, Cli};
 use base64::{engine::general_purpose, Engine};
 use bb8::{Pool, PooledConnection};
 use bytes::Bytes;
+use cached::{Cached, UnboundCache};
 use diesel::ConnectionError;
 use diesel_async::{
     pooled_connection::{AsyncDieselConnectionManager, ManagerConfig},
@@ -23,7 +24,9 @@ use renegade_api::auth::add_expiring_auth_to_headers;
 use renegade_common::types::wallet::keychain::HmacKey;
 use reqwest::Client;
 use std::{sync::Arc, time::Duration};
+use tokio::sync::RwLock;
 use tracing::error;
+use uuid::Uuid;
 
 /// The duration for which the admin authentication is valid
 const ADMIN_AUTH_DURATION_MS: u64 = 5_000; // 5 seconds
@@ -32,6 +35,8 @@ const ADMIN_AUTH_DURATION_MS: u64 = 5_000; // 5 seconds
 pub type DbConn<'a> = PooledConnection<'a, AsyncDieselConnectionManager<AsyncPgConnection>>;
 /// The DB pool type
 pub type DbPool = Pool<AsyncDieselConnectionManager<AsyncPgConnection>>;
+/// The API key cache type
+pub type ApiKeyCache = Arc<RwLock<UnboundCache<Uuid, ApiKey>>>;
 
 /// The server struct that holds all the necessary components
 pub struct Server {
@@ -45,6 +50,8 @@ pub struct Server {
     pub management_key: HmacKey,
     /// The encryption key for storing API secrets
     pub encryption_key: Vec<u8>,
+    /// The api key cache
+    pub api_key_cache: ApiKeyCache,
     /// The HTTP client
     pub client: Client,
 }
@@ -71,6 +78,7 @@ impl Server {
             relayer_admin_key,
             management_key,
             encryption_key,
+            api_key_cache: Arc::new(RwLock::new(UnboundCache::new())),
             client: Client::new(),
         })
     }
@@ -127,6 +135,28 @@ impl Server {
         add_expiring_auth_to_headers(path, headers, body, &key, expiration);
         Ok(())
     }
+
+    // --- Caching --- //
+
+    /// Check the cache for an API key
+    pub async fn get_cached_api_secret(&self, id: Uuid) -> Option<ApiKey> {
+        let cache = self.api_key_cache.read().await;
+        cache.get_store().get(&id).cloned()
+    }
+
+    /// Cache an API key
+    pub async fn cache_api_key(&self, api_key: ApiKey) {
+        let mut cache = self.api_key_cache.write().await;
+        cache.cache_set(api_key.id, api_key);
+    }
+
+    /// Mark a cached API key as expired
+    pub async fn mark_cached_key_expired(&self, id: Uuid) {
+        let mut cache = self.api_key_cache.write().await;
+        if let Some(key) = cache.cache_get_mut(&id) {
+            key.is_active = false;
+        }
+    }
 }
 
 /// Create a database pool

auth/auth-server/src/server/queries.rs

@@ -16,42 +16,67 @@ impl Server {
 
     /// Get the API key entry for a given key
     pub async fn get_api_key_entry(&self, api_key: Uuid) -> Result<ApiKey, AuthServerError> {
+        // Check the cache first
+        if let Some(key) = self.get_cached_api_secret(api_key).await {
+            return Ok(key);
+        }
+
+        // Fetch the key from the database
         let mut conn = self.get_db_conn().await?;
         let result = api_keys::table
             .filter(api_keys::id.eq(api_key))
             .limit(1)
             .load::<ApiKey>(&mut conn)
             .await
             .map_err(AuthServerError::db)?;
+        drop(conn); // Drop the connection to release the mutable borrow on `self`
 
-        if result.is_empty() {
-            Err(AuthServerError::unauthorized("API key not found"))
+        let key = if result.is_empty() {
+            return Err(AuthServerError::unauthorized("API key not found"));
         } else {
-            Ok(result[0].clone())
+            result[0].clone()
+        };
+
+        // Cache the key and return
+        self.cache_api_key(key.clone()).await;
+        if !key.is_active {
+            return Err(AuthServerError::ApiKeyInactive);
         }
+
+        Ok(key)
     }
 
     // --- Setters --- //
 
     /// Add a new API key to the database
     pub async fn add_key_query(&self, new_key: NewApiKey) -> Result<(), AuthServerError> {
+        // Write to the database
         let mut conn = self.get_db_conn().await?;
         diesel::insert_into(api_keys::table)
             .values(&new_key)
             .execute(&mut conn)
             .await
-            .map_err(AuthServerError::db)
-            .map(|_| ())
+            .map_err(AuthServerError::db)?;
+        drop(conn); // Drop the connection to release the mutable borrow on `self`
+
+        // Cache the key
+        self.cache_api_key(new_key.into()).await;
+        Ok(())
     }
 
     /// Expire an existing API key
     pub async fn expire_key_query(&self, key_id: Uuid) -> Result<(), AuthServerError> {
+        // Update the database
         let mut conn = self.get_db_conn().await?;
         diesel::update(api_keys::table.filter(api_keys::id.eq(key_id)))
             .set(api_keys::is_active.eq(false))
             .execute(&mut conn)
             .await
-            .map_err(AuthServerError::db)
-            .map(|_| ())
+            .map_err(AuthServerError::db)?;
+        drop(conn); // Drop the connection to release the mutable borrow on `self`
+
+        // Remove the key from the cache
+        self.mark_cached_key_expired(key_id).await;
+        Ok(())
     }
 }