udb_client/
auth.rs

1//! Login and the token lifecycle.
2//!
3//! The broker's refresh token is SINGLE-USE: every successful refresh mints a new
4//! one and invalidates the presented one atomically. A client that keeps sending
5//! its original refresh token authenticates, refreshes once, and then fails with
6//! `Unauthenticated: invalid credential` at the second refresh boundary — an hour
7//! or a day later, far from the code that caused it. [`TokenManager`] persists the
8//! rotated token, which is the whole reason it exists rather than leaving refresh
9//! to callers.
10
11use std::sync::Arc;
12use std::time::{Duration, SystemTime, UNIX_EPOCH};
13
14use tokio::sync::Mutex;
15use tonic::transport::Channel;
16use tonic::Status;
17
18use crate::metadata::Metadata;
19use crate::proto::udb::core::authn::services::v1 as authn;
20use crate::proto::udb::core::authn::services::v1::authn_service_client::AuthnServiceClient;
21
22/// A stored credential set. `expires_at_unix` of 0 means "unknown"; the manager
23/// then treats the token as valid until an explicit refresh.
24#[derive(Clone, Debug, Default, PartialEq, Eq)]
25pub struct Token {
26    pub access_token: String,
27    pub refresh_token: String,
28    pub session_id: String,
29    pub expires_at_unix: u64,
30}
31
32impl Token {
33    /// Non-empty and not within `skew` of expiry.
34    pub fn is_valid(&self, now_unix: u64, skew: Duration) -> bool {
35        if self.access_token.is_empty() {
36            return false;
37        }
38        if self.expires_at_unix == 0 {
39            return true;
40        }
41        now_unix.saturating_add(skew.as_secs()) < self.expires_at_unix
42    }
43}
44
45fn now_unix() -> u64 {
46    SystemTime::now()
47        .duration_since(UNIX_EPOCH)
48        .map(|d| d.as_secs())
49        .unwrap_or(0)
50}
51
52/// Logs in and keeps the token fresh.
53///
54/// Refresh is single-flighted: concurrent callers that all observe a stale token
55/// share one `RefreshToken` round-trip. Without that, N tasks would each present
56/// the same single-use refresh token, one would win, and the rest would be told
57/// their credential is invalid.
58#[derive(Clone)]
59pub struct TokenManager {
60    inner: Arc<Inner>,
61}
62
63struct Inner {
64    client: Mutex<AuthnServiceClient<Channel>>,
65    token: Mutex<Token>,
66    skew: Duration,
67    meta: Metadata,
68}
69
70/// Hand-written rather than derived: a derived `Debug` would print the stored
71/// access and refresh tokens, and this type is exactly the sort of thing that
72/// ends up in a `tracing` field or a panic message.
73impl std::fmt::Debug for TokenManager {
74    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
75        f.debug_struct("TokenManager")
76            .field("tenant_id", &self.inner.meta.tenant_id)
77            .field("project_id", &self.inner.meta.project_id)
78            .field("skew", &self.inner.skew)
79            .field("token", &"<redacted>")
80            .finish()
81    }
82}
83
84impl TokenManager {
85    /// Default refresh skew: refresh this long before actual expiry.
86    pub const DEFAULT_SKEW: Duration = Duration::from_secs(30);
87
88    pub fn new(channel: Channel, meta: Metadata) -> Self {
89        Self::with_skew(channel, meta, Self::DEFAULT_SKEW)
90    }
91
92    pub fn with_skew(channel: Channel, meta: Metadata, skew: Duration) -> Self {
93        Self {
94            inner: Arc::new(Inner {
95                client: Mutex::new(AuthnServiceClient::new(channel)),
96                token: Mutex::new(Token::default()),
97                skew,
98                meta,
99            }),
100        }
101    }
102
103    /// The currently stored token, without refreshing.
104    pub async fn peek(&self) -> Token {
105        self.inner.token.lock().await.clone()
106    }
107
108    /// Authenticate with username and password, storing the result.
109    pub async fn login(
110        &self,
111        username: impl Into<String>,
112        password: impl Into<String>,
113    ) -> Result<Token, Status> {
114        let mut req = tonic::Request::new(authn::LoginRequest {
115            username: username.into(),
116            password: password.into(),
117            ..Default::default()
118        });
119        self.inner.meta.apply(&mut req)?;
120
121        let resp = {
122            let mut client = self.inner.client.lock().await;
123            client.login(req).await?.into_inner()
124        };
125
126        if resp.mfa_required {
127            return Err(Status::unauthenticated(
128                "login requires a second factor; re-call Login with the MFA credential",
129            ));
130        }
131
132        let token = Token {
133            access_token: resp.access_token,
134            refresh_token: resp.refresh_token,
135            session_id: resp.session_id,
136            expires_at_unix: expiry_from(resp.access_token_expires_in),
137        };
138        *self.inner.token.lock().await = token.clone();
139        Ok(token)
140    }
141
142    /// The access token, refreshing first if it is stale.
143    pub async fn access_token(&self) -> Result<String, Status> {
144        {
145            let token = self.inner.token.lock().await;
146            if token.is_valid(now_unix(), self.inner.skew) {
147                return Ok(token.access_token.clone());
148            }
149        }
150        self.refresh().await.map(|t| t.access_token)
151    }
152
153    /// Force a refresh. Holding the token lock across the RPC is what makes this
154    /// single-flighted: a second caller waits and then observes the new token
155    /// rather than presenting the spent one.
156    pub async fn refresh(&self) -> Result<Token, Status> {
157        let mut stored = self.inner.token.lock().await;
158
159        // Someone refreshed while we waited for the lock.
160        if stored.is_valid(now_unix(), self.inner.skew) {
161            return Ok(stored.clone());
162        }
163        if stored.refresh_token.is_empty() && stored.session_id.is_empty() {
164            return Err(Status::unauthenticated(
165                "no refresh token or session id stored; call login() first",
166            ));
167        }
168
169        let mut req = tonic::Request::new(authn::RefreshTokenRequest {
170            refresh_token: stored.refresh_token.clone(),
171            session_id: stored.session_id.clone(),
172        });
173        self.inner.meta.apply(&mut req)?;
174
175        let resp = {
176            let mut client = self.inner.client.lock().await;
177            client.refresh_token(req).await?.into_inner()
178        };
179
180        stored.access_token = resp.access_token;
181        stored.expires_at_unix = expiry_from(resp.access_token_expires_in);
182        // Persist the ROTATED refresh token. Guarded on non-empty because the
183        // response omits it when the caller refreshed with a legacy server-side
184        // session id rather than a token-family credential; assigning blindly
185        // would erase a working credential.
186        if !resp.refresh_token.is_empty() {
187            stored.refresh_token = resp.refresh_token;
188        }
189        Ok(stored.clone())
190    }
191
192    /// Metadata carrying the current access token, ready for a data-plane client.
193    pub async fn authenticated_metadata(&self) -> Result<Metadata, Status> {
194        let token = self.access_token().await?;
195        Ok(self.inner.meta.clone().with_bearer_token(token))
196    }
197}
198
199fn expiry_from(expires_in_secs: i32) -> u64 {
200    if expires_in_secs <= 0 {
201        return 0;
202    }
203    now_unix().saturating_add(expires_in_secs as u64)
204}
205
206#[cfg(test)]
207mod tests {
208    use super::*;
209
210    #[test]
211    fn empty_access_token_is_never_valid() {
212        assert!(!Token::default().is_valid(1_000, Duration::from_secs(30)));
213    }
214
215    #[test]
216    fn unknown_expiry_is_treated_as_valid() {
217        let t = Token {
218            access_token: "a".into(),
219            expires_at_unix: 0,
220            ..Default::default()
221        };
222        assert!(t.is_valid(u64::MAX, Duration::from_secs(30)));
223    }
224
225    #[test]
226    fn skew_expires_the_token_early() {
227        let t = Token {
228            access_token: "a".into(),
229            expires_at_unix: 1_000,
230            ..Default::default()
231        };
232        assert!(t.is_valid(900, Duration::from_secs(30)), "930 < 1000");
233        assert!(
234            !t.is_valid(980, Duration::from_secs(30)),
235            "1010 >= 1000: inside the skew window, must refresh"
236        );
237    }
238
239    #[test]
240    fn expiry_from_ignores_non_positive() {
241        assert_eq!(expiry_from(0), 0);
242        assert_eq!(expiry_from(-5), 0);
243        assert!(expiry_from(60) >= now_unix());
244    }
245}
udb_client/auth.rs

udb_client/
auth.rs
