Skip to content

Commit 52fb65a

Browse files
authored
feat: Improve introspection cache (#567)
Add introspection cache to axum and use `moka` as cache implementation. BREAKING CHANGE: This removes the "roles" enum from the introspected user. It is possible to achieve the same mechanism with the following work around: ```rust enum Role { Admin, Client } trait MyExtIntrospectedUser { fn role(&self, role: Role) -> Option<..>; } impl MyExtIntrospectedUser for IntrospectedUser { fn role(&self, role: Role) -> Option<..> { // convenience impl here } } ```
1 parent 3040d32 commit 52fb65a

File tree

9 files changed

+379
-117
lines changed

9 files changed

+379
-117
lines changed

Cargo.toml

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ interceptors = ["credentials", "dep:time", "dep:tokio"]
7272
## By default, only the in-memory cache is available. To use a different cache,
7373
## enable specific features of this crate, or implement your own cache with
7474
## the trait.
75-
introspection_cache = ["dep:async-trait", "dep:time"]
75+
introspection_cache = ["dep:async-trait", "dep:time", "dep:moka"]
7676

7777
## The OIDC module enables basic OIDC (OpenID Connect) features to communicate
7878
## with ZITADEL. Two examples are the `discover` and `introspect` functions.
@@ -146,6 +146,7 @@ base64-compat = { version = "1", optional = true }
146146
custom_error = "1.9.2"
147147
document-features = { version = "0.2.8", optional = true }
148148
jsonwebtoken = { version = "9.3.0", optional = true }
149+
moka = { version = "0.12.8", features = ["future"], optional = true }
149150
openidconnect = { version = "3.5.0", optional = true }
150151
pbjson-types = { version = "0.7.0", optional = true }
151152
prost = { version = "0.13.1", optional = true }
@@ -170,6 +171,7 @@ tonic-types = { version = "0.12.1", optional = true }
170171
chrono = "0.4.38"
171172
tokio = { version = "1.37.0", features = ["macros", "rt-multi-thread"] }
172173
tower = { version = "0.4.13" }
174+
http-body-util = "0.1.0"
173175

174176
[package.metadata.docs.rs]
175177
all-features = true

src/axum/introspection/mod.rs

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,12 @@
1-
//! The intropection module allows you to use the OAuth 2.0 Token Introspection flow to authenticate users against ZITADEL.
1+
//! The introspection module allows you to use the OAuth 2.0 Token Introspection flow to authenticate users against ZITADEL.
22
//!
3-
//! Axum uses "extracters" and "middlewares" to intercept calls. To authenticate a user against ZITADEL, you can use the [IntrospectedUser].
3+
//! Axum uses "extractors" and "middlewares" to intercept calls. To authenticate a user against ZITADEL, you can use the [IntrospectedUser].
44
//! Which enables an extractor workflow: [extractor](https://docs.rs/axum/latest/axum/extract/index.html)
55
//!
66
//! #### Configure Axum
77
//!
88
//! 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).
9+
//! When a custom state is used, [FromRef](axum::extract::FromRef) must be implemented. See [IntrospectionState] for more Details.
910
//!
1011
//! ```no_run
1112
//! #
@@ -60,6 +61,6 @@ mod state;
6061
mod state_builder;
6162
mod user;
6263

63-
pub use state::{IntrospectionConfig, IntrospectionState};
64+
pub use state::IntrospectionState;
6465
pub use state_builder::{IntrospectionStateBuilder, IntrospectionStateBuilderError};
6566
pub use user::{IntrospectedUser, IntrospectionGuardError};

src/axum/introspection/state.rs

Lines changed: 27 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,30 +1,39 @@
1-
use axum::extract::FromRef;
21
use openidconnect::IntrospectionUrl;
2+
use std::sync::Arc;
33

4+
#[cfg(feature = "introspection_cache")]
5+
use crate::oidc::introspection::cache::IntrospectionCache;
46
use crate::oidc::introspection::AuthorityAuthentication;
57

8+
/// State which must be present for extractor to work,
9+
/// compare [axum's official documentation](https://docs.rs/axum/0.6.4/axum/extract/struct.State.html#for-library-authors).
10+
/// Use [IntrospectionStateBuilder](super::IntrospectionStateBuilder) to configure the respective parameters.
11+
///
12+
/// If a custom state is used, then [FromRef](axum::extract::FromRef) must be implemented,
13+
/// to make the necessary state available.
14+
///
15+
/// ```
16+
/// use axum::extract::FromRef;
17+
/// use zitadel::axum::introspection::IntrospectionState;
18+
/// struct UserState {
19+
/// introspection_state: IntrospectionState
20+
/// }
21+
///
22+
/// impl FromRef<UserState> for IntrospectionState {
23+
/// fn from_ref(input: &UserState) -> Self {
24+
/// input.introspection_state.clone()
25+
/// }
26+
/// }
627
#[derive(Clone, Debug)]
728
pub struct IntrospectionState {
8-
pub(crate) config: IntrospectionConfig,
29+
pub(crate) config: Arc<IntrospectionConfig>,
930
}
1031

11-
impl IntrospectionState {
12-
pub fn config(&self) -> &IntrospectionConfig {
13-
&self.config
14-
}
15-
}
16-
17-
/// Configuration that must be inject into the axum application state. Used by the
18-
/// [IntrospectionStateBuilder](super::IntrospectionStateBuilder). This struct is also used to create the [IntrospectionState](IntrospectionState)
19-
#[derive(Debug, Clone)]
20-
pub struct IntrospectionConfig {
32+
#[derive(Debug)]
33+
pub(crate) struct IntrospectionConfig {
2134
pub(crate) authority: String,
2235
pub(crate) authentication: AuthorityAuthentication,
2336
pub(crate) introspection_uri: IntrospectionUrl,
24-
}
25-
26-
impl FromRef<IntrospectionState> for IntrospectionConfig {
27-
fn from_ref(input: &IntrospectionState) -> Self {
28-
input.config.clone()
29-
}
37+
#[cfg(feature = "introspection_cache")]
38+
pub(crate) cache: Option<Box<dyn IntrospectionCache>>,
3039
}

src/axum/introspection/state_builder.rs

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,14 @@
11
use custom_error::custom_error;
2+
use std::sync::Arc;
23

34
use crate::axum::introspection::state::IntrospectionConfig;
45
use crate::credentials::Application;
56
use crate::oidc::discovery::{discover, DiscoveryError};
67
use crate::oidc::introspection::AuthorityAuthentication;
78

9+
#[cfg(feature = "introspection_cache")]
10+
use crate::oidc::introspection::cache::IntrospectionCache;
11+
812
use super::state::IntrospectionState;
913

1014
custom_error! {
@@ -18,6 +22,8 @@ custom_error! {
1822
pub struct IntrospectionStateBuilder {
1923
authority: String,
2024
authentication: Option<AuthorityAuthentication>,
25+
#[cfg(feature = "introspection_cache")]
26+
cache: Option<Box<dyn IntrospectionCache>>,
2127
}
2228

2329
/// Builder for [IntrospectionConfig]
@@ -26,6 +32,8 @@ impl IntrospectionStateBuilder {
2632
Self {
2733
authority: authority.to_string(),
2834
authentication: None,
35+
#[cfg(feature = "introspection_cache")]
36+
cache: None,
2937
}
3038
}
3139

@@ -48,6 +56,17 @@ impl IntrospectionStateBuilder {
4856
self
4957
}
5058

59+
/// Set the [IntrospectionCache] to use for caching introspection responses.
60+
#[cfg(feature = "introspection_cache")]
61+
pub fn with_introspection_cache(
62+
&mut self,
63+
cache: impl IntrospectionCache + 'static,
64+
) -> &mut IntrospectionStateBuilder {
65+
self.cache = Some(Box::new(cache));
66+
67+
self
68+
}
69+
5170
pub async fn build(&mut self) -> Result<IntrospectionState, IntrospectionStateBuilderError> {
5271
if self.authentication.is_none() {
5372
return Err(IntrospectionStateBuilderError::NoAuthSchema);
@@ -67,11 +86,13 @@ impl IntrospectionStateBuilder {
6786
}
6887

6988
Ok(IntrospectionState {
70-
config: IntrospectionConfig {
89+
config: Arc::new(IntrospectionConfig {
7190
authority: self.authority.clone(),
7291
introspection_uri: introspection_uri.unwrap(),
7392
authentication: self.authentication.as_ref().unwrap().clone(),
74-
},
93+
#[cfg(feature = "introspection_cache")]
94+
cache: self.cache.take(),
95+
}),
7596
})
7697
}
7798
}

0 commit comments

Comments
 (0)