reqwest_connect_rpc/
token_source.rs

1// Copyright 2025 Anapaya Systems
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//   http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14//! Token source trait for the connect RPC client.
15
16use std::sync::Arc;
17
18use async_trait::async_trait;
19use tokio::sync::watch;
20
21pub mod mock;
22pub mod refresh;
23pub mod static_token;
24
25/// The cause of a [`TokenSourceError`].
26pub type TokenSourceCause = Arc<dyn std::error::Error + Sync + Send>;
27
28/// Why a token source could not provide a token.
29///
30/// A source knows what it authenticates against, so it is the only place that can tell whether
31/// asking again may work. The variant it picks is how it says so; callers read it back through
32/// [`TokenSourceError::is_transient`] instead of interpreting the cause themselves.
33#[derive(Debug, Clone, thiserror::Error)]
34pub enum TokenSourceError {
35    /// No token right now, but a later call may produce one, e.g. because the request to the
36    /// token service did not get through.
37    #[error("token temporarily unavailable: {0}")]
38    Unavailable(TokenSourceCause),
39    /// The credential the source authenticates with was refused, or the answer it got back cannot
40    /// be used as a token. Asking again yields the same answer.
41    #[error("token rejected: {0}")]
42    Rejected(TokenSourceCause),
43    /// The source itself stopped working, e.g. the task that refreshes its token is gone, so it
44    /// will not produce a token again.
45    #[error("token source broken: {0}")]
46    Broken(TokenSourceCause),
47}
48
49impl TokenSourceError {
50    /// The token could not be obtained now, but a later call may succeed.
51    pub fn unavailable(cause: impl Into<Box<dyn std::error::Error + Sync + Send>>) -> Self {
52        Self::Unavailable(cause.into().into())
53    }
54
55    /// The token cannot be obtained, no matter how often it is asked for.
56    pub fn rejected(cause: impl Into<Box<dyn std::error::Error + Sync + Send>>) -> Self {
57        Self::Rejected(cause.into().into())
58    }
59
60    /// The source can no longer hand out tokens at all.
61    pub fn broken(cause: impl Into<Box<dyn std::error::Error + Sync + Send>>) -> Self {
62        Self::Broken(cause.into().into())
63    }
64
65    /// Returns whether the failure is transient, so that a retry may help.
66    ///
67    /// Prefer this over matching the variants: a new variant would silently fall into a caller's
68    /// wildcard arm.
69    #[must_use]
70    pub fn is_transient(&self) -> bool {
71        match self {
72            Self::Unavailable(_) => true,
73            Self::Rejected(_) | Self::Broken(_) => false,
74        }
75    }
76}
77
78/// A watch receiver for token source updates.
79pub type TokenSourceWatch = watch::Receiver<Option<Result<String, TokenSourceError>>>;
80
81/// A source for authentication tokens.
82#[async_trait]
83pub trait TokenSource: Send + Sync + 'static {
84    /// Returns a watch receiver that always holds the latest valid token.
85    ///
86    /// The receiver allows both grabbing the current value immediately
87    /// and awaiting updates.
88    fn watch(&self) -> TokenSourceWatch;
89
90    /// Gets a token, possibly refreshing it.
91    ///
92    /// If the token cannot be obtained, returns the [`TokenSourceError`] the source published,
93    /// classification included, so the caller can tell a source that is momentarily out of reach
94    /// from one that will not produce a token again.
95    ///
96    /// Prefer using `watch` if a subscription to token updates is needed.
97    ///
98    /// ### Implementation Note
99    ///
100    /// The default implementation uses the watch channel to get the latest token.
101    ///``
102    /// - Should be efficient to call multiple times.
103    /// - Errors should be returned if no valid token can be obtained.
104    /// - Should try to not return errors as long as a valid token is available.
105    async fn get_token(&self) -> Result<String, TokenSourceError> {
106        let mut watch = self.watch();
107
108        // First, try to get the current value without waiting. and return immediately if available.
109        // Cloning keeps the published classification; the cause is shared, not copied.
110        match watch.borrow_and_update().as_ref() {
111            Some(Ok(token)) => return Ok(token.clone()),
112            Some(Err(e)) => return Err(e.clone()),
113            None => {}
114        }
115
116        // If there is no current value, wait for an update.
117        watch.changed().await.map_err(TokenSourceError::broken)?;
118
119        // After being notified, get the updated value.
120        match watch.borrow().as_ref() {
121            Some(Ok(token)) => Ok(token.clone()),
122            Some(Err(e)) => Err(e.clone()),
123            None => {
124                Err(TokenSourceError::broken(
125                    "token source watch channel has no value",
126                ))
127            }
128        }
129    }
130
131    /// Formats the token for use in an `Authorization` header.
132    ///
133    /// The default implementation formats the token as a Bearer token.
134    /// Override this method if a different format is required.
135    fn format_header(&self, token: String) -> String {
136        format!("Bearer {token}")
137    }
138}
reqwest_connect_rpc/token_source.rs

reqwest_connect_rpc/
token_source.rs
