feat: auth cache

auth/Cargo.toml

@@ -9,28 +9,30 @@ anyhow = { workspace = true }
 async-trait = { workspace = true }
 axum = { workspace = true, features = ["headers"] }
 axum-sessions = "0.4.1"
+chrono = { workspace = true, features = ["serde"] }
 clap = { workspace = true }
 http = { workspace = true }
-jsonwebtoken = { workspace = true}
+hyper = { workspace = true }
+jsonwebtoken = { workspace = true }
 opentelemetry = { workspace = true }
 opentelemetry-datadog = { workspace = true }
 rand = { workspace = true }
 ring = { workspace = true }
-serde = { workspace = true, features = [ "derive" ] }
-sqlx = { version = "0.6.2", features = [ "sqlite", "json", "runtime-tokio-native-tls", "migrate" ] }
+serde = { workspace = true, features = ["derive"] }
+serde_json = { workspace = true }
+sqlx = { version = "0.6.2", features = ["sqlite", "json", "runtime-tokio-native-tls", "migrate"] }
 strum = { workspace = true }
 thiserror = { workspace = true }
-tokio = { version = "1.22.0", features = [ "full" ] }
+tokio = { version = "1.22.0", features = ["full"] }
+tower = { workspace = true, features = ["util"] }
 tracing = { workspace = true }
 tracing-opentelemetry = { workspace = true }
 tracing-subscriber = { workspace = true, features = ["env-filter"] }
+ttl_cache = "0.5.1"
 
 [dependencies.shuttle-common]
 workspace = true
 features = ["backend", "models"]
 
 [dev-dependencies]
 axum-extra = { version = "0.5.0", features = ["cookie"] }
-hyper = { workspace = true }
-serde_json = { workspace = true }
-tower = { workspace = true, features = ["util"] }

auth/src/api/builder.rs

@@ -13,24 +13,33 @@ use shuttle_common::{
     request_span,
 };
 use sqlx::SqlitePool;
+use std::sync::RwLock;
 use tracing::field;
+use ttl_cache::TtlCache;
 
 use crate::{
     secrets::{EdDsaManager, KeyManager},
     user::{UserManagement, UserManager},
+    COOKIE_EXPIRATION,
 };
 
-use super::handlers::{
-    convert_cookie, convert_key, get_public_key, get_user, login, logout, post_user, refresh_token,
+use super::{
+    cache::{CacheLayer, CacheManagement, CacheManager},
+    handlers::{
+        convert_cookie, convert_key, get_public_key, get_user, login, logout, post_user,
+        refresh_token,
+    },
 };
 
 pub type UserManagerState = Arc<Box<dyn UserManagement>>;
 pub type KeyManagerState = Arc<Box<dyn KeyManager>>;
+pub type CacheManagerState = Arc<Box<dyn CacheManagement>>;
 
 #[derive(Clone)]
 pub struct RouterState {
     pub user_manager: UserManagerState,
     pub key_manager: KeyManagerState,
+    pub cache_manager: CacheManagerState,
 }
 
 // Allow getting a user management state directly
@@ -47,10 +56,18 @@ impl FromRef<RouterState> for KeyManagerState {
     }
 }
 
+// Allow getting a cache state directly
+impl FromRef<RouterState> for CacheManagerState {
+    fn from_ref(router_state: &RouterState) -> Self {
+        router_state.cache_manager.clone()
+    }
+}
+
 pub struct ApiBuilder {
     router: Router<RouterState>,
     pool: Option<SqlitePool>,
     session_layer: Option<SessionLayer<MemoryStore>>,
+    cache: Option<Arc<RwLock<TtlCache<String, String>>>>,
 }
 
 impl Default for ApiBuilder {
@@ -87,6 +104,7 @@ impl ApiBuilder {
             router,
             pool: None,
             session_layer: None,
+            cache: None,
         }
     }
 
@@ -102,24 +120,42 @@ impl ApiBuilder {
         self.session_layer = Some(
             SessionLayer::new(store, &secret)
                 .with_cookie_name("shuttle.sid")
-                .with_session_ttl(Some(std::time::Duration::from_secs(60 * 60 * 24))) // One day
+                .with_session_ttl(Some(COOKIE_EXPIRATION)) // One day
                 .with_secure(true),
         );
 
         self
     }
 
+    pub fn with_cache(mut self) -> Self {
+        let cache = Arc::new(RwLock::new(TtlCache::new(1000)));
+
+        self.cache = Some(cache);
+
+        self
+    }
+
     pub fn into_router(self) -> Router {
         let pool = self.pool.expect("an sqlite pool is required");
         let session_layer = self.session_layer.expect("a session layer is required");
+        let cache = self.cache.expect("a cache is required");
 
+        let cache_manager = CacheManager { cache };
         let user_manager = UserManager { pool };
         let key_manager = EdDsaManager::new();
 
-        self.router.layer(session_layer).with_state(RouterState {
+        let state = RouterState {
             user_manager: Arc::new(Box::new(user_manager)),
             key_manager: Arc::new(Box::new(key_manager)),
-        })
+            cache_manager: Arc::new(Box::new(cache_manager)),
+        };
+
+        self.router
+            .layer(session_layer)
+            .layer(CacheLayer {
+                state: state.clone(),
+            })
+            .with_state(state)
     }
 }