oci_client/
token_cache.rs

1//! Token cache for OCI registry authentication
2
3use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
4use oci_spec::distribution::Reference;
5use serde::Deserialize;
6use std::collections::BTreeMap;
7use std::fmt;
8use std::sync::Arc;
9use std::time::{SystemTime, UNIX_EPOCH};
10use tokio::sync::RwLock;
11use tracing::{debug, warn};
12
13/// A token granted during the OAuth2-like workflow for OCI registries.
14#[derive(Deserialize, Clone)]
15#[serde(untagged)]
16#[serde(rename_all = "snake_case")]
17pub enum RegistryToken {
18    /// Token value
19    Token {
20        /// The string value of the token
21        token: String,
22    },
23    /// AccessToken value
24    AccessToken {
25        /// The string value of the access_token
26        access_token: String,
27    },
28}
29
30impl fmt::Debug for RegistryToken {
31    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
32        let redacted = String::from("<redacted>");
33        match self {
34            RegistryToken::Token { .. } => {
35                f.debug_struct("Token").field("token", &redacted).finish()
36            }
37            RegistryToken::AccessToken { .. } => f
38                .debug_struct("AccessToken")
39                .field("access_token", &redacted)
40                .finish(),
41        }
42    }
43}
44
45#[derive(Debug, Clone)]
46/// Type of registry auth token
47pub enum RegistryTokenType {
48    /// Bearer auth token type
49    Bearer(RegistryToken),
50    /// Basic auth token type
51    Basic(String, String),
52}
53
54impl RegistryToken {
55    /// Returns the bearer token in a form suitable to use for an Authorization header
56    pub fn bearer_token(&self) -> String {
57        format!("Bearer {}", self.token())
58    }
59
60    /// Returns the token value
61    pub fn token(&self) -> &str {
62        match self {
63            RegistryToken::Token { token } => token,
64            RegistryToken::AccessToken { access_token } => access_token,
65        }
66    }
67}
68
69/// Desired operation for registry authentication
70#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
71pub enum RegistryOperation {
72    /// Authenticate for push operations
73    Push,
74    /// Authenticate for pull operations
75    Pull,
76}
77
78#[derive(Debug, Deserialize)]
79struct BearerTokenClaims {
80    exp: Option<u64>,
81}
82
83#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
84struct TokenCacheKey {
85    registry: String,
86    repository: String,
87    operation: RegistryOperation,
88}
89
90struct TokenCacheValue {
91    token: RegistryTokenType,
92    expiration: u64,
93}
94
95#[derive(Clone)]
96/// A cache to hold authentication tokens
97pub struct TokenCache {
98    // (registry, repository, scope) -> (token, expiration)
99    tokens: Arc<RwLock<BTreeMap<TokenCacheKey, TokenCacheValue>>>,
100    /// Default token expiration in seconds, to use when claim doesn't specify a value
101    pub default_expiration_secs: usize,
102}
103
104impl TokenCache {
105    pub(crate) fn new(default_expiration_secs: usize) -> Self {
106        TokenCache {
107            tokens: Arc::new(RwLock::new(BTreeMap::new())),
108            default_expiration_secs,
109        }
110    }
111
112    /// Insert a token corresponding to reference and operation keys
113    pub async fn insert(
114        &self,
115        reference: &Reference,
116        op: RegistryOperation,
117        token: RegistryTokenType,
118    ) {
119        let expiration = match token {
120            RegistryTokenType::Basic(_, _) => u64::MAX,
121            RegistryTokenType::Bearer(ref t) => {
122                match bearer_token_cache_expiration(t.token(), self.default_expiration_secs) {
123                    Some(value) => value,
124                    None => return,
125                }
126            }
127        };
128        let registry = reference.resolve_registry().to_string();
129        let repository = reference.repository().to_string();
130        debug!(%registry, %repository, ?op, %expiration, "Inserting token");
131        self.tokens.write().await.insert(
132            TokenCacheKey {
133                registry,
134                repository,
135                operation: op,
136            },
137            TokenCacheValue { token, expiration },
138        );
139    }
140
141    pub(crate) async fn get(
142        &self,
143        reference: &Reference,
144        op: RegistryOperation,
145    ) -> Option<RegistryTokenType> {
146        let registry = reference.resolve_registry().to_string();
147        let repository = reference.repository().to_string();
148        let key = TokenCacheKey {
149            registry,
150            repository,
151            operation: op,
152        };
153        match self.tokens.read().await.get(&key) {
154            Some(TokenCacheValue {
155                ref token,
156                expiration,
157            }) => {
158                let now = SystemTime::now();
159                let epoch = now
160                    .duration_since(UNIX_EPOCH)
161                    .expect("Time went backwards")
162                    .as_secs();
163                if epoch > *expiration {
164                    debug!(%key.registry, %key.repository, ?key.operation, %expiration, miss=false, expired=true, "Fetching token");
165                    None
166                } else {
167                    debug!(%key.registry, %key.repository, ?key.operation, %expiration, miss=false, expired=false, "Fetching token");
168                    Some(token.clone())
169                }
170            }
171            None => {
172                debug!(%key.registry, %key.repository, ?key.operation, miss = true, "Fetching token");
173                None
174            }
175        }
176    }
177}
178
179/// The longest time we keep a bearer token in the cache, even when its claimed
180/// expiration is further in the future.
181const MAX_TOKEN_CACHE_TTL_SECS: u64 = 24 * 60 * 60;
182
183/// Picks a cache eviction time for a bearer token.
184///
185/// This function reads the unverified `exp` claim from the token payload. It
186/// does not check the token signature, and it must not be used to decide
187/// whether the token is valid. The registry that issued the token is the
188/// only party that can make that decision; this cache only avoids
189/// requesting a new token before the old one expires.
190///
191/// A bearer token can be a JWT (a signed token with three base64 parts) or
192/// an opaque string, as GHCR issues. This function returns `None` only when
193/// the token looks like a JWT but its payload is not valid base64 or valid
194/// JSON, because we then cannot tell how long the token lasts.
195fn bearer_token_cache_expiration(token_str: &str, default_expiration_secs: usize) -> Option<u64> {
196    let mut parts = token_str.split('.');
197    let (Some(_header), Some(payload), Some(_signature), None) =
198        (parts.next(), parts.next(), parts.next(), parts.next())
199    else {
200        // The token is not a JWT (e.g., an opaque token issued by registries
201        // like GHCR). Use the default expiration as a best-effort assumption,
202        // mirroring the behaviour for JWT tokens that carry no `exp` claim.
203        debug!(
204            "Bearer token is not a JWT, assuming a {} seconds validity",
205            default_expiration_secs
206        );
207        return Some(default_expiration(default_expiration_secs));
208    };
209
210    // Registry tokens are opaque credentials. We only inspect the untrusted payload
211    // to choose a cache eviction time; signature verification is neither required nor
212    // useful here because the registry that issued the token also controls its lifetime.
213    let payload = match URL_SAFE_NO_PAD.decode(payload) {
214        Ok(payload) => payload,
215        Err(error) => {
216            warn!(?error, "Invalid bearer token payload encoding");
217            return None;
218        }
219    };
220    let claims: BearerTokenClaims = match serde_json::from_slice(&payload) {
221        Ok(claims) => claims,
222        Err(error) => {
223            warn!(?error, "Invalid bearer token payload");
224            return None;
225        }
226    };
227
228    let exp = match claims.exp {
229        Some(exp) => exp,
230        None => {
231            // The token doesn't have a claim that states a value for the expiration.
232            // The registry auth specification defaults such tokens to 60 seconds:
233            // https://distribution.github.io/distribution/spec/auth/token/
234            debug!(
235                "Cannot extract expiration from token's claims, assuming a {} seconds validity",
236                default_expiration_secs
237            );
238            default_expiration(default_expiration_secs)
239        }
240    };
241
242    // Cap the cache TTL so that a token with a far-future or malicious `exp`
243    // claim cannot pin a stale token in the cache indefinitely. The registry
244    // still rejects an expired or revoked token on the next real request.
245    let max_exp = default_expiration(MAX_TOKEN_CACHE_TTL_SECS as usize);
246    Some(exp.min(max_exp))
247}
248
249fn default_expiration(default_expiration_secs: usize) -> u64 {
250    SystemTime::now()
251        .duration_since(UNIX_EPOCH)
252        .expect("Time went backwards")
253        .as_secs()
254        + default_expiration_secs as u64
255}
256
257#[cfg(test)]
258mod tests {
259    use super::*;
260    use oci_spec::distribution::Reference;
261    use rstest::rstest;
262    use serde::Serialize;
263
264    // An opaque token as issued by registries like GHCR — not a JWT.
265    const OPAQUE_TOKEN: &str = "ghs_exampleOpaqueTokenFromGHCR1234567890";
266
267    #[derive(Serialize)]
268    struct ClaimsWithExp {
269        exp: u64,
270    }
271
272    #[derive(Serialize)]
273    struct ClaimsWithoutExp {
274        sub: &'static str,
275    }
276
277    fn make_jwt_with_exp(exp: u64) -> String {
278        make_jwt(&ClaimsWithExp { exp })
279    }
280
281    fn make_jwt_without_exp() -> String {
282        make_jwt(&ClaimsWithoutExp { sub: "test" })
283    }
284
285    fn make_jwt(claims: &impl Serialize) -> String {
286        let payload = serde_json::to_vec(claims).expect("failed to serialize JWT claims");
287        format!("e30.{}.signature", URL_SAFE_NO_PAD.encode(payload))
288    }
289
290    fn now_secs() -> u64 {
291        SystemTime::now()
292            .duration_since(UNIX_EPOCH)
293            .unwrap()
294            .as_secs()
295    }
296
297    #[test]
298    fn jwt_with_near_exp_uses_claims_expiration() {
299        let exp = now_secs() + 3600;
300        let token = make_jwt_with_exp(exp);
301        let cached_exp = bearer_token_cache_expiration(&token, 60)
302            .expect("should return Some for valid JWT with exp");
303        assert_eq!(cached_exp, exp);
304    }
305
306    #[test]
307    fn jwt_with_far_future_exp_is_capped() {
308        // A claimed expiration far in the future (year 2286) must not pin the
309        // cached token for that long. We cap the cache TTL instead.
310        let token = make_jwt_with_exp(9999999999);
311        let before = now_secs();
312        let cached_exp = bearer_token_cache_expiration(&token, 60)
313            .expect("should return Some for valid JWT with exp");
314        let after = now_secs();
315        assert!(cached_exp < 9999999999);
316        assert!(cached_exp >= before + MAX_TOKEN_CACHE_TTL_SECS);
317        assert!(cached_exp <= after + MAX_TOKEN_CACHE_TTL_SECS);
318    }
319
320    /// Tokens whose `exp` claim cannot be read fall back to the default
321    /// expiration: a JWT with no `exp` claim, an opaque token (as GHCR
322    /// issues), and a JWE-style token (5 dot-separated parts, which is not a
323    /// JWT we can read a claim from).
324    #[rstest]
325    #[case::jwt_without_exp(make_jwt_without_exp())]
326    #[case::opaque_token(OPAQUE_TOKEN.to_string())]
327    #[case::five_part_jwe("a.b.c.d.e".to_string())]
328    fn token_without_readable_exp_uses_default_expiration(#[case] token: String) {
329        let before = now_secs();
330        let exp = bearer_token_cache_expiration(&token, 60)
331            .expect("should return Some with default expiration");
332        let after = now_secs();
333        assert!(exp >= before + 60);
334        assert!(exp <= after + 60);
335    }
336
337    /// A token that looks like a JWT (three dot-separated parts) but whose
338    /// payload cannot be read as a bearer token's claims returns `None`, so
339    /// the caller does not cache it.
340    #[rstest]
341    #[case::invalid_base64("not-valid-base64!!!".to_string())]
342    #[case::empty_segment("".to_string())]
343    // `URL_SAFE_NO_PAD` rejects a payload that carries `=` padding.
344    #[case::padded_base64("eyJzdWIiOiJ0ZXN0In0=".to_string())]
345    #[case::not_json(URL_SAFE_NO_PAD.encode(b"not json"))]
346    #[case::exp_as_string(URL_SAFE_NO_PAD.encode(br#"{"exp":"9999999999"}"#))]
347    #[case::exp_negative(URL_SAFE_NO_PAD.encode(br#"{"exp":-1}"#))]
348    #[case::exp_float(URL_SAFE_NO_PAD.encode(br#"{"exp":1.5}"#))]
349    fn malformed_jwt_payload_returns_none(#[case] payload: String) {
350        let token = format!("e30.{payload}.signature");
351        assert!(bearer_token_cache_expiration(&token, 60).is_none());
352    }
353
354    #[tokio::test]
355    async fn opaque_token_is_cached() {
356        let cache = TokenCache::new(60);
357        let reference: Reference = "ghcr.io/kubewarden/policies/pod-privileged:v1.0.10"
358            .parse()
359            .unwrap();
360        let token = RegistryTokenType::Bearer(RegistryToken::Token {
361            token: OPAQUE_TOKEN.to_string(),
362        });
363
364        cache
365            .insert(&reference, RegistryOperation::Pull, token)
366            .await;
367
368        assert!(
369            cache
370                .get(&reference, RegistryOperation::Pull)
371                .await
372                .is_some(),
373            "opaque bearer token should be cached"
374        );
375    }
376}
oci_client/token_cache.rs

oci_client/
token_cache.rs
