olai_http/
credential.rs

1use std::fmt::Debug;
2use std::sync::Arc;
3use std::time::Duration;
4
5use async_trait::async_trait;
6use futures::future::BoxFuture;
7use reqwest::Client;
8
9use crate::service::HttpService;
10use crate::{Result, RetryConfig, TemporaryToken, TokenCache};
11
12/// Provides credentials for use when signing requests.
13///
14/// Implementors are responsible for fetching, refreshing, and returning a
15/// credential of an associated type (e.g. an AWS SigV4 key set, an Azure bearer
16/// token, or a GCP service-account token).  The cloud-client crate ships
17/// ready-made implementations for every supported provider; see the per-cloud
18/// `credential` modules for details.
19#[async_trait]
20pub trait CredentialProvider: std::fmt::Debug + Send + Sync {
21    /// The type of credential returned by this provider
22    type Credential;
23
24    /// Return a credential
25    async fn get_credential(&self) -> Result<Arc<Self::Credential>>;
26}
27
28/// A static set of credentials
29#[derive(Debug)]
30pub struct StaticCredentialProvider<T> {
31    credential: Arc<T>,
32}
33
34impl<T> StaticCredentialProvider<T> {
35    /// A [`CredentialProvider`] for a static credential of type `T`
36    pub fn new(credential: T) -> Self {
37        Self {
38            credential: Arc::new(credential),
39        }
40    }
41}
42
43#[async_trait]
44impl<T> CredentialProvider for StaticCredentialProvider<T>
45where
46    T: std::fmt::Debug + Send + Sync,
47{
48    type Credential = T;
49
50    async fn get_credential(&self) -> Result<Arc<T>> {
51        Ok(Arc::clone(&self.credential))
52    }
53}
54
55/// A [`CredentialProvider`] that fetches temporary tokens via HTTP
56#[derive(Debug)]
57pub(crate) struct TokenCredentialProvider<T: TokenProvider> {
58    inner: T,
59    client: Client,
60    service: Arc<dyn HttpService>,
61    retry: RetryConfig,
62    cache: TokenCache<Arc<T::Credential>>,
63}
64
65impl<T: TokenProvider> TokenCredentialProvider<T> {
66    pub(crate) fn new(
67        inner: T,
68        client: Client,
69        service: Arc<dyn HttpService>,
70        retry: RetryConfig,
71    ) -> Self {
72        Self {
73            inner,
74            client,
75            service,
76            retry,
77            cache: Default::default(),
78        }
79    }
80
81    /// Override the minimum remaining TTL for a cached token to be used
82    pub(crate) fn with_min_ttl(mut self, min_ttl: Duration) -> Self {
83        self.cache = self.cache.with_min_ttl(min_ttl);
84        self
85    }
86}
87
88#[async_trait]
89impl<T: TokenProvider> CredentialProvider for TokenCredentialProvider<T> {
90    type Credential = T::Credential;
91
92    async fn get_credential(&self) -> Result<Arc<Self::Credential>> {
93        self.cache
94            .get_or_insert_with(|| {
95                self.inner
96                    .fetch_token(&self.client, &self.service, &self.retry)
97            })
98            .await
99    }
100}
101
102/// Applies authentication to a built reqwest request.
103///
104/// Each cloud provider and Databricks implements this trait.
105/// `send()` in `CloudClient` calls `self.signer.sign(builder).await?`.
106///
107/// # References
108/// - <https://docs.rs/reqwest/latest/reqwest/struct.RequestBuilder.html>
109pub trait RequestSigner: Debug + Send + Sync {
110    /// Sign a [`reqwest::RequestBuilder`], returning the modified builder.
111    fn sign<'a>(
112        &'a self,
113        req: reqwest::RequestBuilder,
114    ) -> BoxFuture<'a, Result<reqwest::RequestBuilder>>;
115}
116
117/// Tries signers in order; uses the first that succeeds without error.
118///
119/// This enables Databricks-style fallback chains without bespoke builder logic.
120#[derive(Debug)]
121pub struct SignerChain {
122    signers: Vec<Arc<dyn RequestSigner>>,
123}
124
125impl SignerChain {
126    pub fn new(signers: Vec<Arc<dyn RequestSigner>>) -> Self {
127        Self { signers }
128    }
129}
130
131impl RequestSigner for SignerChain {
132    fn sign<'a>(
133        &'a self,
134        req: reqwest::RequestBuilder,
135    ) -> BoxFuture<'a, Result<reqwest::RequestBuilder>> {
136        Box::pin(async move {
137            let mut current = req;
138            for signer in &self.signers {
139                current = signer.sign(current).await?;
140            }
141            Ok(current)
142        })
143    }
144}
145
146#[async_trait]
147pub(crate) trait TokenProvider: std::fmt::Debug + Send + Sync {
148    type Credential: std::fmt::Debug + Send + Sync;
149
150    async fn fetch_token(
151        &self,
152        client: &Client,
153        service: &Arc<dyn HttpService>,
154        retry: &RetryConfig,
155    ) -> Result<TemporaryToken<Arc<Self::Credential>>>;
156}
olai_http/credential.rs

olai_http/
credential.rs
