🤖 AI: Fix Caching of Pub Keys

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "mauth-client"
-version = "0.7.2"
+version = "0.8.0"
 authors = ["Mason Gup <mgup@mdsol.com>"]
 edition = "2024"
 rust-version = "1.88"
@@ -31,15 +31,17 @@ axum = { version = ">= 0.8", optional = true }
 futures-core = { version = "0.3", optional = true }
 http = "1"
 bytes = { version = "1", optional = true }
+lru = { version = "0.18", default-features = false, optional = true }
 thiserror = "2"
 mauth-core = "0.6"
 tracing = { version = "0.1", optional = true }
 
 [dev-dependencies]
 tokio = { version = "1", features = ["rt-multi-thread", "macros"] }
+rsa = "0.9"
 
 [features]
-axum-service = ["tower", "futures-core", "axum", "bytes", "tracing"]
+axum-service = ["tower", "futures-core", "axum", "bytes", "lru", "tracing"]
 tracing-otel-26 = ["reqwest-tracing/opentelemetry_0_26"]
 tracing-otel-27 = ["reqwest-tracing/opentelemetry_0_27"]
 tracing-otel-28 = ["reqwest-tracing/opentelemetry_0_28"]

README.md

@@ -180,6 +180,33 @@ from there in order to take further actions based on the error type. This error
 implements Axum's `OptionalFromRequestParts`, so you can more easily retrieve it using
 `Option<MAuthValidationError>` anywhere that supports extractors.
 
+Both layers also record the app UUID that the rejected request *claimed*, as an
+`AttemptedMAuthIdentity`, so a rejection can be attributed to a caller. The Required layer
+attaches it to the 401 **response** extensions, since the request itself is gone by that
+point; the Optional layer attaches it to the **request** extensions alongside the error.
+It is included on the logged error either way.
+
+`AttemptedMAuthIdentity` is deliberately a separate type from `ValidatedRequestDetails`,
+and it is not an authenticated identity — it is read straight off the signature header of
+a request that failed to validate, so anyone could have sent it. Use it for logging and
+diagnostics; anything making a trust decision must use `ValidatedRequestDetails`.
+
+### Public Key Caching
+
+Verifying an incoming signature needs the sending app's public key, which is fetched from
+MAuth and then cached in process. Entries expire according to the `Cache-Control` headers
+on MAuth's own response, so rotating an application's key takes effect within the lifetime
+MAuth states, without restarting anything. A response arriving without a usable directive
+is cached for a conservative 60 seconds, so that MAuth does not become a synchronous
+dependency of every authenticated request. Lookups that MAuth answers
+with a 404 are not cached at all, so an app registered moments ago authenticates
+immediately.
+
+The cache holds 10,000 keys by default, evicting the least recently used, which is sized
+for a service with a wide range of callers rather than a handful. Set
+`pubkey_cache_capacity` on the `ConfigFileSection` to change it. As with the shared HTTP
+client, only the first configuration loaded in a process takes effect.
+
 ### OpenTelemetry Integration
 
 There are also optional features `tracing-otel-26` through `tracing-otel-31`

src/axum_service.rs

@@ -13,7 +13,9 @@ use std::task::{Context, Poll};
 use tower::{Layer, Service};
 use tracing::error;
 
-use crate::validate_incoming::{MAuthValidationError, ValidatedRequestDetails};
+use crate::validate_incoming::{
+    AttemptedMAuthIdentity, MAuthValidationError, ValidatedRequestDetails,
+};
 use crate::{
     MAuthInfo,
     config::{ConfigFileSection, ConfigReadError},
@@ -51,12 +53,22 @@ where
                     Ok(response) => Ok(response.into()),
                     Err(err) => Err(err),
                 },
-                Err(err) => {
+                Err(rejection) => {
                     error!(
-                        error = ?err,
+                        error = ?rejection.error,
+                        app_uuid = rejection.app_uuid.map(tracing::field::display),
                         "Failed to validate MAuth signature, rejecting request"
                     );
-                    Ok(StatusCode::UNAUTHORIZED.into_response())
+                    let mut response = StatusCode::UNAUTHORIZED.into_response();
+                    // The request is gone by the time we reject it, so the
+                    // claimed identity rides out on the response instead, for
+                    // any logging layer sitting outside this one.
+                    if let Some(app_uuid) = rejection.app_uuid {
+                        response
+                            .extensions_mut()
+                            .insert(AttemptedMAuthIdentity { app_uuid });
+                    }
+                    Ok(response)
                 }
             }
         })

src/config.rs

@@ -60,6 +60,9 @@ impl MAuthInfo {
             signer: Signer::new(section.app_uuid.clone(), pk_data.unwrap())?,
         };
 
+        #[cfg(feature = "axum-service")]
+        crate::init_pubkey_cache(section.pubkey_cache_capacity);
+
         CLIENT.get_or_init(|| {
             let builder = ClientBuilder::new(Client::new()).with(mauth_info.clone());
             #[cfg(any(
@@ -89,6 +92,13 @@ pub struct ConfigFileSection {
     pub private_key_data: Option<String>,
     pub v2_only_sign_requests: Option<bool>,
     pub v2_only_authenticate: Option<bool>,
+    /// How many verifying keys to retain when validating incoming requests.
+    /// `None` uses a default sized for a wide-fanout service; raise it only if
+    /// more distinct applications than that call this service within the
+    /// lifetime MAuth puts on a key. Ignored without the `axum-service` feature.
+    ///
+    /// Only the first configuration loaded in a process takes effect.
+    pub pubkey_cache_capacity: Option<usize>,
 }
 
 impl Default for ConfigFileSection {
@@ -101,6 +111,7 @@ impl Default for ConfigFileSection {
             private_key_data: None,
             v2_only_sign_requests: Some(true),
             v2_only_authenticate: Some(true),
+            pubkey_cache_capacity: None,
         }
     }
 }
@@ -155,6 +166,7 @@ mod test {
             private_key_data: None,
             v2_only_sign_requests: None,
             v2_only_authenticate: None,
+            pubkey_cache_capacity: None,
         };
         let load_result = MAuthInfo::from_config_section(&bad_config);
         assert!(matches!(load_result, Err(ConfigReadError::InvalidUri(_))));
@@ -170,6 +182,7 @@ mod test {
             private_key_data: None,
             v2_only_sign_requests: None,
             v2_only_authenticate: None,
+            pubkey_cache_capacity: None,
         };
         let load_result = MAuthInfo::from_config_section(&bad_config);
         assert!(matches!(
@@ -190,6 +203,7 @@ mod test {
             private_key_data: None,
             v2_only_sign_requests: None,
             v2_only_authenticate: None,
+            pubkey_cache_capacity: None,
         };
         let load_result = MAuthInfo::from_config_section(&bad_config);
         fs::remove_file(&filename).await.unwrap();
@@ -211,6 +225,7 @@ mod test {
             private_key_data: None,
             v2_only_sign_requests: None,
             v2_only_authenticate: None,
+            pubkey_cache_capacity: None,
         };
         let load_result = MAuthInfo::from_config_section(&bad_config);
         fs::remove_file(&filename).await.unwrap();

src/lib.rs

@@ -3,30 +3,60 @@
 
 use ::reqwest_middleware::ClientWithMiddleware;
 use mauth_core::signer::Signer;
-use mauth_core::verifier::Verifier;
 use reqwest::Url;
-use std::collections::HashMap;
-use std::sync::{LazyLock, OnceLock, RwLock};
+use std::sync::OnceLock;
 use uuid::Uuid;
 
+#[cfg(feature = "axum-service")]
+use crate::validate_incoming::{DEFAULT_PUBKEY_CACHE_CAPACITY, PubkeyCache};
+
 /// This is the primary struct of this class. It contains all of the information
 /// required to sign requests using the MAuth protocol and verify the responses.
 ///
-/// Note that it contains a cache of response keys for verifying response signatures. This cache
-/// makes the struct non-Sync.
+/// The cache of keys used to verify incoming signatures is process-wide and is
+/// shared by all instances of this struct.
 #[derive(Clone)]
 pub struct MAuthInfo {
     app_id: Uuid,
     sign_with_v1_also: bool,
     signer: Signer,
+    // These two are only consulted when validating incoming requests, which the
+    // `axum-service` feature gates. They are always populated regardless, so the
+    // struct shape does not vary between feature sets.
+    #[cfg_attr(not(feature = "axum-service"), allow(dead_code))]
     mauth_uri_base: Url,
+    #[cfg_attr(not(feature = "axum-service"), allow(dead_code))]
     allow_v1_auth: bool,
 }
 
 static CLIENT: OnceLock<ClientWithMiddleware> = OnceLock::new();
 
-static PUBKEY_CACHE: LazyLock<RwLock<HashMap<Uuid, Verifier>>> =
-    LazyLock::new(|| RwLock::new(HashMap::new()));
+#[cfg(feature = "axum-service")]
+static PUBKEY_CACHE: OnceLock<PubkeyCache> = OnceLock::new();
+
+/// The process-wide cache of verifier keys fetched from MAuth.
+///
+/// The validation services clone themselves for every request. Their custom
+/// `Clone` implementations rebuild `MAuthInfo` from configuration; `MAuthInfo`
+/// itself derives `Clone`. Keeping the cache global preserves entries across
+/// those service clones without changing the existing service implementation.
+///
+/// Falls back to [`DEFAULT_PUBKEY_CACHE_CAPACITY`] if a lookup somehow happens
+/// before any config has been loaded.
+#[cfg(feature = "axum-service")]
+pub(crate) fn pubkey_cache() -> &'static PubkeyCache {
+    PUBKEY_CACHE.get_or_init(|| PubkeyCache::new(DEFAULT_PUBKEY_CACHE_CAPACITY))
+}
+
+/// Size the key cache from configuration.
+///
+/// First caller wins, matching the behavior of `CLIENT`: a later
+/// `ConfigFileSection` naming a different capacity is silently ignored.
+#[cfg(feature = "axum-service")]
+pub(crate) fn init_pubkey_cache(capacity: Option<usize>) {
+    PUBKEY_CACHE
+        .get_or_init(|| PubkeyCache::new(capacity.unwrap_or(DEFAULT_PUBKEY_CACHE_CAPACITY)));
+}
 
 /// Tower Service and Layer to allow Tower-integrated servers to validate incoming request
 #[cfg(feature = "axum-service")]

src/protocol_test_suite.rs

@@ -30,6 +30,7 @@ async fn setup_mauth_info() -> (MAuthInfo, u64) {
         private_key_data: None,
         v2_only_sign_requests: None,
         v2_only_authenticate: None,
+        pubkey_cache_capacity: None,
     };
     (
         MAuthInfo::from_config_section(&mock_config_section).unwrap(),