Skip to main content

alibi_plugins/jwt/
keyring.rs

1use super::JwtAlgorithm;
2use super::JwtKeyPairConfig;
3use super::JwtKeyringContext;
4use super::JwtPlugin;
5use super::JwtSignOptions;
6use super::ResolvedJwtSigningKey;
7use super::crypto;
8use crate::token_crypto::decrypt_with_config;
9use crate::token_crypto::encrypt_with_config;
10use alibi_core::AuthContext;
11use alibi_core::AuthError;
12use alibi_core::AuthRequest;
13use alibi_core::AuthResult;
14use alibi_core::AuthSchema;
15use alibi_core::CreateJwk;
16use alibi_core::Jwk;
17use chrono::Utc;
18use std::str::FromStr;
19impl JwtPlugin {
20    pub(in crate::jwt) async fn keys(
21        &self,
22        request: Option<&AuthRequest>,
23        ctx: &AuthContext<impl AuthSchema>,
24    ) -> AuthResult<Vec<Jwk>> {
25        let endpoint = alibi_core::endpoint::current_endpoint_call_context();
26        let path = endpoint
27            .as_ref()
28            .and_then(alibi_core::endpoint::EndpointCall::path)
29            .unwrap_or_else(|| request.map_or("virtual:", AuthRequest::path));
30        self.keys_at_path(request, path, ctx).await
31    }
32
33    pub(in crate::jwt) async fn keys_in_transaction<S: AuthSchema>(
34        &self,
35        request: Option<&AuthRequest>,
36        ctx: &AuthContext<S>,
37        transaction: Option<&dyn alibi_core::store::AuthTransaction<S>>,
38    ) -> AuthResult<Vec<Jwk>> {
39        if self.config.keyring.is_none()
40            && let Some(transaction) = transaction
41        {
42            transaction.list_jwks().await
43        } else {
44            self.keys(request, ctx).await
45        }
46    }
47
48    pub(in crate::jwt) async fn keys_at_path(
49        &self,
50        request: Option<&AuthRequest>,
51        path: &str,
52        ctx: &AuthContext<impl AuthSchema>,
53    ) -> AuthResult<Vec<Jwk>> {
54        let endpoint = alibi_core::endpoint::current_endpoint_call_context();
55        match &self.config.keyring {
56            Some(keyring) => {
57                keyring
58                    .keys(&JwtKeyringContext {
59                        path,
60                        request,
61                        endpoint: endpoint.as_ref(),
62                    })
63                    .await
64            }
65            None => ctx.database.list_jwks().await,
66        }
67    }
68
69    /// Provision a private signing key and its public JWK in persistent storage.
70    pub async fn create_jwk<S: AuthSchema>(
71        &self,
72        config: Option<&JwtKeyPairConfig>,
73        request: Option<&AuthRequest>,
74        ctx: &AuthContext<S>,
75    ) -> AuthResult<Jwk> {
76        self.create_jwk_in_transaction(config, request, ctx, None)
77            .await
78    }
79
80    pub(in crate::jwt) async fn create_jwk_in_transaction<S: AuthSchema>(
81        &self,
82        config: Option<&JwtKeyPairConfig>,
83        request: Option<&AuthRequest>,
84        ctx: &AuthContext<S>,
85        transaction: Option<&dyn alibi_core::store::AuthTransaction<S>>,
86    ) -> AuthResult<Jwk> {
87        let config = config.unwrap_or(&self.config.key_pair);
88        let (public, private) = crypto::generate(config)?;
89        let private = serde_json::to_string(&private)?;
90        let private_key = if self.config.disable_private_key_encryption {
91            private
92        } else {
93            serde_json::to_string(&encrypt_with_config(&private, &ctx.config)?)?
94        };
95        let now = Utc::now();
96        let data = CreateJwk {
97            id: None,
98            public_key: serde_json::to_string(&public)?,
99            private_key,
100            created_at: now,
101            expires_at: self
102                .config
103                .rotation_interval
104                .filter(|duration| !duration.is_zero())
105                .map(|duration| now + duration),
106            alg: Some(config.algorithm.as_str().to_owned()),
107            crv: config.algorithm.curve().map(str::to_owned),
108        };
109        let endpoint = alibi_core::endpoint::current_endpoint_call_context();
110        match &self.config.keyring {
111            Some(keyring) => {
112                keyring
113                    .create_key(
114                        data,
115                        &JwtKeyringContext {
116                            path: endpoint
117                                .as_ref()
118                                .and_then(alibi_core::endpoint::EndpointCall::path)
119                                .unwrap_or_else(|| request.map_or("virtual:", AuthRequest::path)),
120                            request,
121                            endpoint: endpoint.as_ref(),
122                        },
123                    )
124                    .await
125            }
126            None => match transaction {
127                Some(transaction) => transaction.create_jwk(data).await,
128                None => ctx.database.create_jwk(data).await,
129            },
130        }
131    }
132
133    /// Select a live key, with explicit key or algorithm pinning when requested.
134    pub async fn resolve_signing_key<S: AuthSchema>(
135        &self,
136        options: &JwtSignOptions,
137        request: Option<&AuthRequest>,
138        ctx: &AuthContext<S>,
139    ) -> AuthResult<Option<ResolvedJwtSigningKey>> {
140        self.resolve_signing_key_in_transaction(options, request, ctx, None)
141            .await
142    }
143
144    pub(in crate::jwt) async fn resolve_signing_key_in_transaction<S: AuthSchema>(
145        &self,
146        options: &JwtSignOptions,
147        request: Option<&AuthRequest>,
148        ctx: &AuthContext<S>,
149        transaction: Option<&dyn alibi_core::store::AuthTransaction<S>>,
150    ) -> AuthResult<Option<ResolvedJwtSigningKey>> {
151        if self.config.remote_signer.is_some() {
152            return Ok(None);
153        }
154        // Pinned IDs use findOne independently of the adapter's findMany
155        // limit. Custom keyrings supply the raw full set once for this lookup.
156        let mut keys = if options.signing_key_id.is_some() {
157            Vec::new()
158        } else {
159            self.keys_in_transaction(request, ctx, transaction).await?
160        };
161        keys.sort_by_key(|key| std::cmp::Reverse(key.created_at));
162        let primary = self.config.key_pair.algorithm;
163        let key_alg = |key: &Jwk| {
164            key.alg
165                .as_deref()
166                .map_or(Ok(primary), JwtAlgorithm::from_str)
167        };
168        let now = Utc::now();
169        let live = |key: &&Jwk| key.expires_at.is_none_or(|expiry| expiry > now);
170        let mut minted_unpinned_key = false;
171        let mut key = if let Some(id) = &options.signing_key_id {
172            let key = match &self.config.keyring {
173                Some(_) => self.keys_in_transaction(request, ctx, transaction).await?.into_iter().find(|key| &key.id == id),
174                None => match transaction {Some(transaction)=>transaction.get_jwk_by_id(id).await?,None=>ctx.database.get_jwk_by_id(id).await?},
175            }.ok_or_else(|| AuthError::config(format!("signJWT: signingKeyId \"{id}\" not found in JWKS. The key must be provisioned before it can be referenced.")))?;
176            if let Some(algorithm) = options.signing_algorithm
177                && key_alg(&key)? != algorithm
178            {
179                return Err(AuthError::config(format!(
180                    "signJWT: signingKeyId \"{id}\" has a different algorithm than {}",
181                    algorithm.as_str()
182                )));
183            }
184            key
185        } else if let Some(algorithm) = options.signing_algorithm {
186            if let Some(key) = keys
187                .iter()
188                .filter(live)
189                .find(|key| key_alg(key).is_ok_and(|alg| alg == algorithm))
190                .cloned()
191            {
192                key
193            } else {
194                let config = self
195                    .config
196                    .additional_key_pairs
197                    .iter()
198                    .find(|config| config.algorithm == algorithm)
199                    .or_else(|| (primary == algorithm).then_some(&self.config.key_pair))
200                    .ok_or_else(|| {
201                        AuthError::config(format!(
202                            "No signing key configured for {}",
203                            algorithm.as_str()
204                        ))
205                    })?;
206                self.create_jwk_in_transaction(Some(config), request, ctx, transaction)
207                    .await?
208            }
209        } else {
210            if let Some(key) = keys
211                .iter()
212                .filter(live)
213                .find(|key| key_alg(key).is_ok_and(|alg| alg == primary))
214                .cloned()
215            {
216                key
217            } else {
218                // Source performs a separate fallback lookup. An application
219                // keyring can observe it or return a changed key set.
220                let mut fallback = self.keys_in_transaction(request, ctx, transaction).await?;
221                fallback.sort_by_key(|key| std::cmp::Reverse(key.created_at));
222                if let Some(key) = fallback
223                    .into_iter()
224                    .find(|key| key.expires_at.is_none_or(|expiry| expiry > Utc::now()))
225                {
226                    key
227                } else {
228                    minted_unpinned_key = true;
229                    self.create_jwk_in_transaction(None, request, ctx, transaction)
230                        .await?
231                }
232            }
233        };
234        if !minted_unpinned_key && key.expires_at.is_some_and(|expiry| expiry < Utc::now()) {
235            if options.signing_key_id.is_some() || options.signing_algorithm.is_some() {
236                return Err(AuthError::config(
237                    "signJWT: requested signing key is expired and an explicit kid/alg was provided; not auto-minting a replacement. Rotate the key explicitly.",
238                ));
239            }
240            key = self
241                .create_jwk_in_transaction(None, request, ctx, transaction)
242                .await?;
243        }
244        let private = if self.config.disable_private_key_encryption {
245            key.private_key
246        } else {
247            let encrypted: String = serde_json::from_str(&key.private_key)?;
248            decrypt_with_config(&encrypted, &ctx.config).map_err(|_error| AuthError::config("Failed to decrypt private key. Make sure the secret currently in use is the same as the one used to encrypt the private key. If you are using a different secret, either clean up your JWKS or disable private key encryption."))?
249        };
250        Ok(Some(ResolvedJwtSigningKey {
251            algorithm: key
252                .alg
253                .as_deref()
254                .map(JwtAlgorithm::from_str)
255                .unwrap_or(Ok(primary))?,
256            key_id: key.id,
257            private_key: serde_json::from_str(&private)?,
258        }))
259    }
260}