From f338e650d96f62c21fe1a846da927e62e78997c6 Mon Sep 17 00:00:00 2001 From: Joey Kraut Date: Thu, 24 Oct 2024 13:17:55 -0700 Subject: [PATCH] auth-server: Add local api key cache --- Cargo.lock | 53 +++++++++++++++++++++++++- auth/auth-server/Cargo.toml | 1 + auth/auth-server/src/models.rs | 12 ++++++ auth/auth-server/src/server/mod.rs | 32 +++++++++++++++- auth/auth-server/src/server/queries.rs | 39 +++++++++++++++---- 5 files changed, 128 insertions(+), 9 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 553456f..1e93eb3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -94,6 +94,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "allocator-api2" +version = "0.2.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c6cb57a04249c6480766f7f7cef5467412af1490f8d1e243141daddada3264f" + [[package]] name = "alloy-primitives" version = "0.7.7" @@ -722,6 +728,7 @@ dependencies = [ "base64 0.22.1", "bb8", "bytes", + "cached", "chrono", "clap", "common", @@ -1528,6 +1535,39 @@ dependencies = [ "pkg-config", ] +[[package]] +name = "cached" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4d73155ae6b28cf5de4cfc29aeb02b8a1c6dab883cb015d15cd514e42766846" +dependencies = [ + "ahash 0.8.11", + "cached_proc_macro", + "cached_proc_macro_types", + "hashbrown 0.14.5", + "once_cell", + "thiserror", + "web-time 1.1.0", +] + +[[package]] +name = "cached_proc_macro" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f42a145ed2d10dce2191e1dcf30cfccfea9026660e143662ba5eec4017d5daa" +dependencies = [ + "darling", + "proc-macro2", + "quote", + "syn 2.0.83", +] + +[[package]] +name = "cached_proc_macro_types" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ade8366b8bd5ba243f0a58f036cc0ca8a2f069cff1a2351ef1cac6b083e16fc0" + [[package]] name = "cadence" version = "1.5.0" @@ -3546,6 +3586,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" dependencies = [ "ahash 0.8.11", + "allocator-api2", ] [[package]] @@ -7979,7 +8020,7 @@ dependencies = [ "tracing-core", "tracing-log", "tracing-subscriber 0.3.18", - "web-time", + "web-time 0.2.4", ] [[package]] @@ -8493,6 +8534,16 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + [[package]] name = "web3" version = "0.18.0" diff --git a/auth/auth-server/Cargo.toml b/auth/auth-server/Cargo.toml index b635043..d22065b 100644 --- a/auth/auth-server/Cargo.toml +++ b/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"] } diff --git a/auth/auth-server/src/models.rs b/auth/auth-server/src/models.rs index 9d6c2b4..411454b 100644 --- a/auth/auth-server/src/models.rs +++ b/auth/auth-server/src/models.rs @@ -33,3 +33,15 @@ impl NewApiKey { Self { id, encrypted_key, description } } } + +impl From 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, + } + } +} diff --git a/auth/auth-server/src/server/mod.rs b/auth/auth-server/src/server/mod.rs index 0700880..c78c390 100644 --- a/auth/auth-server/src/server/mod.rs +++ b/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>; /// The DB pool type pub type DbPool = Pool>; +/// The API key cache type +pub type ApiKeyCache = Arc>>; /// 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, + /// 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 { + 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 diff --git a/auth/auth-server/src/server/queries.rs b/auth/auth-server/src/server/queries.rs index 6032696..d50dc4c 100644 --- a/auth/auth-server/src/server/queries.rs +++ b/auth/auth-server/src/server/queries.rs @@ -16,6 +16,12 @@ impl Server { /// Get the API key entry for a given key pub async fn get_api_key_entry(&self, api_key: Uuid) -> Result { + // 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)) @@ -23,35 +29,54 @@ impl Server { .load::(&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(()) } }