Improve introspection cache

Cargo.toml

@@ -72,7 +72,7 @@ interceptors = ["credentials", "dep:time", "dep:tokio"]
 ## By default, only the in-memory cache is available. To use a different cache,
 ## enable specific features of this crate, or implement your own cache with
 ## the trait.
-introspection_cache = ["dep:async-trait", "dep:time"]
+introspection_cache = ["dep:async-trait", "dep:time", "dep:moka"]
 
 ## The OIDC module enables basic OIDC (OpenID Connect) features to communicate
 ## with ZITADEL. Two examples are the `discover` and `introspect` functions.
@@ -146,6 +146,7 @@ base64-compat = { version = "1", optional = true }
 custom_error = "1.9.2"
 document-features = { version = "0.2.8", optional = true }
 jsonwebtoken = { version = "9.3.0", optional = true }
+moka = { version = "0.12.8", features = ["future"], optional = true }
 openidconnect = { version = "3.5.0", optional = true }
 pbjson-types = { version = "0.7.0", optional = true }
 prost = { version = "0.13.1", optional = true }
@@ -170,6 +171,7 @@ tonic-types = { version = "0.12.1", optional = true }
 chrono = "0.4.38"
 tokio = { version = "1.37.0", features = ["macros", "rt-multi-thread"] }
 tower = { version = "0.4.13" }
+http-body-util = "0.1.0"
 
 [package.metadata.docs.rs]
 all-features = true

src/axum/introspection/mod.rs

@@ -1,11 +1,12 @@
-//! The intropection module allows you to use the OAuth 2.0 Token Introspection flow to authenticate users against ZITADEL.
+//! The introspection module allows you to use the OAuth 2.0 Token Introspection flow to authenticate users against ZITADEL.
 //!
-//! Axum uses "extracters" and "middlewares" to intercept calls. To authenticate a user against ZITADEL, you can use the [IntrospectedUser].
+//! Axum uses "extractors" and "middlewares" to intercept calls. To authenticate a user against ZITADEL, you can use the [IntrospectedUser].
 //! Which enables an extractor workflow: [extractor](https://docs.rs/axum/latest/axum/extract/index.html)
 //!
 //! #### Configure Axum
 //!
 //! To use the introspection flow, you need to configure the [IntrospectionState] and add it to your [Router](https://docs.rs/axum/latest/axum/routing/struct.Router.html).
+//! When a custom state is used, [FromRef](axum::extract::FromRef) must be implemented. See [IntrospectionState] for more Details.
 //!
 //! ```no_run
 //! #
@@ -60,6 +61,6 @@ mod state;
 mod state_builder;
 mod user;
 
-pub use state::{IntrospectionConfig, IntrospectionState};
+pub use state::IntrospectionState;
 pub use state_builder::{IntrospectionStateBuilder, IntrospectionStateBuilderError};
 pub use user::{IntrospectedUser, IntrospectionGuardError};

src/axum/introspection/state.rs

@@ -1,30 +1,39 @@
-use axum::extract::FromRef;
 use openidconnect::IntrospectionUrl;
+use std::sync::Arc;
 
+#[cfg(feature = "introspection_cache")]
+use crate::oidc::introspection::cache::IntrospectionCache;
 use crate::oidc::introspection::AuthorityAuthentication;
 
+/// State which must be present for extractor to work,
+/// compare [axum's official documentation](https://docs.rs/axum/0.6.4/axum/extract/struct.State.html#for-library-authors).
+/// Use [IntrospectionStateBuilder](super::IntrospectionStateBuilder) to configure the respective parameters.
+///
+/// If a custom state is used, then [FromRef](axum::extract::FromRef) must be implemented,
+/// to make the necessary state available.
+///
+/// ```
+/// use axum::extract::FromRef;
+/// use zitadel::axum::introspection::IntrospectionState;
+/// struct UserState {
+///     introspection_state: IntrospectionState
+/// }
+///
+/// impl FromRef<UserState> for IntrospectionState {
+///     fn from_ref(input: &UserState) -> Self {
+///         input.introspection_state.clone()
+///     }
+/// }
 #[derive(Clone, Debug)]
 pub struct IntrospectionState {
-    pub(crate) config: IntrospectionConfig,
+    pub(crate) config: Arc<IntrospectionConfig>,
 }
 
-impl IntrospectionState {
-    pub fn config(&self) -> &IntrospectionConfig {
-        &self.config
-    }
-}
-
-/// Configuration that must be inject into the axum application state. Used by the
-/// [IntrospectionStateBuilder](super::IntrospectionStateBuilder). This struct is also used to create the [IntrospectionState](IntrospectionState)
-#[derive(Debug, Clone)]
-pub struct IntrospectionConfig {
+#[derive(Debug)]
+pub(crate) struct IntrospectionConfig {
     pub(crate) authority: String,
     pub(crate) authentication: AuthorityAuthentication,
     pub(crate) introspection_uri: IntrospectionUrl,
-}
-
-impl FromRef<IntrospectionState> for IntrospectionConfig {
-    fn from_ref(input: &IntrospectionState) -> Self {
-        input.config.clone()
-    }
+    #[cfg(feature = "introspection_cache")]
+    pub(crate) cache: Option<Box<dyn IntrospectionCache>>,
 }

src/axum/introspection/state_builder.rs

@@ -1,10 +1,14 @@
 use custom_error::custom_error;
+use std::sync::Arc;
 
 use crate::axum::introspection::state::IntrospectionConfig;
 use crate::credentials::Application;
 use crate::oidc::discovery::{discover, DiscoveryError};
 use crate::oidc::introspection::AuthorityAuthentication;
 
+#[cfg(feature = "introspection_cache")]
+use crate::oidc::introspection::cache::IntrospectionCache;
+
 use super::state::IntrospectionState;
 
 custom_error! {
@@ -18,6 +22,8 @@ custom_error! {
 pub struct IntrospectionStateBuilder {
     authority: String,
     authentication: Option<AuthorityAuthentication>,
+    #[cfg(feature = "introspection_cache")]
+    cache: Option<Box<dyn IntrospectionCache>>,
 }
 
 /// Builder for [IntrospectionConfig]
@@ -26,6 +32,8 @@ impl IntrospectionStateBuilder {
         Self {
             authority: authority.to_string(),
             authentication: None,
+            #[cfg(feature = "introspection_cache")]
+            cache: None,
         }
     }
 
@@ -48,6 +56,17 @@ impl IntrospectionStateBuilder {
         self
     }
 
+    /// Set the [IntrospectionCache] to use for caching introspection responses.
+    #[cfg(feature = "introspection_cache")]
+    pub fn with_introspection_cache(
+        &mut self,
+        cache: impl IntrospectionCache + 'static,
+    ) -> &mut IntrospectionStateBuilder {
+        self.cache = Some(Box::new(cache));
+
+        self
+    }
+
     pub async fn build(&mut self) -> Result<IntrospectionState, IntrospectionStateBuilderError> {
         if self.authentication.is_none() {
             return Err(IntrospectionStateBuilderError::NoAuthSchema);
@@ -67,11 +86,13 @@ impl IntrospectionStateBuilder {
         }
 
         Ok(IntrospectionState {
-            config: IntrospectionConfig {
+            config: Arc::new(IntrospectionConfig {
                 authority: self.authority.clone(),
                 introspection_uri: introspection_uri.unwrap(),
                 authentication: self.authentication.as_ref().unwrap().clone(),
-            },
+                #[cfg(feature = "introspection_cache")]
+                cache: self.cache.take(),
+            }),
         })
     }
 }