mozambigue 0.1.2

JWT validation library with JWKS caching for Kubernetes and Okta tokens
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112

use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use std::time::Instant;

use jsonwebtoken::jwk::JwkSet;
use reqwest::Client;
use serde::Deserialize;
use tokio::sync::RwLock;

use crate::error::fetch_jwks_error;
use crate::error::openid_jwks_error;
use crate::error::Result;

#[derive(Debug, Deserialize)]
struct OpenIdConfig {
    jwks_uri: String,
}

struct CachedJwks {
    jwks: JwkSet,
    fetched_at: Instant,
}

impl CachedJwks {
    fn new(jwks: JwkSet) -> Self {
        Self {
            jwks,
            fetched_at: Instant::now(),
        }
    }

    fn is_expired(&self, ttl: Duration) -> bool {
        self.fetched_at.elapsed() >= ttl
    }
}

/// Cache for JWKS (JSON Web Key Sets) with automatic expiration and refresh
pub(crate) struct JwksCache {
    cache: Arc<RwLock<HashMap<String, CachedJwks>>>,
    ttl: Duration,
    client: Client,
}

impl JwksCache {
    /// Create a new JWKS cache with the given TTL and HTTP client
    pub(crate) fn new(ttl: Duration, client: Client) -> Self {
        Self {
            cache: Arc::new(RwLock::new(HashMap::new())),
            ttl,
            client,
        }
    }

    /// Get JWKS for the given issuer, fetching from the network if not cached or expired
    pub(crate) async fn get_jwks(&self, issuer: &str) -> Result<JwkSet> {
        if let Some(jwks) = self.try_get_cached(issuer).await {
            return Ok(jwks);
        }

        // Cache miss or expired, refresh
        self.refresh(issuer).await
    }

    /// Try to get JWKS from cache if present and not expired
    async fn try_get_cached(&self, issuer: &str) -> Option<JwkSet> {
        let cache = self.cache.read().await;
        let cached = cache.get(issuer)?;

        if cached.is_expired(self.ttl) {
            return None;
        }

        Some(cached.jwks.clone())
    }

    /// Force refresh the JWKS for the given issuer
    async fn refresh(&self, issuer: &str) -> Result<JwkSet> {
        let jwks = self.fetch_jwks(issuer).await?;

        let mut cache = self.cache.write().await;
        cache.insert(issuer.to_string(), CachedJwks::new(jwks.clone()));

        Ok(jwks)
    }

    /// Fetch JWKS from the OpenID configuration endpoint
    async fn fetch_jwks(&self, issuer: &str) -> Result<JwkSet> {
        let openid_url = format!("{issuer}/.well-known/openid-configuration");
        let OpenIdConfig { jwks_uri, .. } = self
            .client
            .get(&openid_url)
            .send()
            .await
            .map_err(openid_jwks_error)?
            .json()
            .await
            .map_err(openid_jwks_error)?;

        let jwks: JwkSet = self
            .client
            .get(&jwks_uri)
            .send()
            .await
            .map_err(fetch_jwks_error)?
            .json()
            .await
            .map_err(fetch_jwks_error)?;

        Ok(jwks)
    }
}