object_store/client/
token.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::future::Future;
19use std::time::{Duration, Instant};
20use tokio::sync::Mutex;
21
22/// A temporary authentication token with an associated expiry
23#[derive(Debug, Clone)]
24pub struct TemporaryToken<T> {
25    /// The temporary credential
26    pub token: T,
27    /// The instant at which this credential is no longer valid
28    /// None means the credential does not expire
29    pub expiry: Option<Instant>,
30}
31
32/// Provides [`TokenCache::get_or_insert_with`] which can be used to cache a
33/// [`TemporaryToken`] based on its expiry
34#[derive(Debug)]
35pub struct TokenCache<T> {
36    cache: Mutex<Option<TemporaryToken<T>>>,
37    min_ttl: Duration,
38}
39
40impl<T> Default for TokenCache<T> {
41    fn default() -> Self {
42        Self {
43            cache: Default::default(),
44            min_ttl: Duration::from_secs(300),
45        }
46    }
47}
48
49impl<T: Clone + Send> TokenCache<T> {
50    /// Override the minimum remaining TTL for a cached token to be used
51    #[cfg(feature = "aws")]
52    pub fn with_min_ttl(self, min_ttl: Duration) -> Self {
53        Self { min_ttl, ..self }
54    }
55
56    pub async fn get_or_insert_with<F, Fut, E>(&self, f: F) -> Result<T, E>
57    where
58        F: FnOnce() -> Fut + Send,
59        Fut: Future<Output = Result<TemporaryToken<T>, E>> + Send,
60    {
61        let now = Instant::now();
62        let mut locked = self.cache.lock().await;
63
64        if let Some(cached) = locked.as_ref() {
65            match cached.expiry {
66                Some(ttl) if ttl.checked_duration_since(now).unwrap_or_default() > self.min_ttl => {
67                    return Ok(cached.token.clone());
68                }
69                None => return Ok(cached.token.clone()),
70                _ => (),
71            }
72        }
73
74        let cached = f().await?;
75        let token = cached.token.clone();
76        *locked = Some(cached);
77
78        Ok(token)
79    }
80}
object_store/client/token.rs

object_store/client/
token.rs
