rig_core/providers/internal/
auth.rs

1//! Shared OAuth errors, device-code callbacks, token stores, and
2//! authentication transport helpers.
3//!
4//! ```
5//! use rig_core::providers::chatgpt::auth::DeviceCodeHandler;
6//! let handler = DeviceCodeHandler::new(|prompt| {
7//!     println!("Visit {} and enter {}", prompt.verification_uri, prompt.user_code);
8//! });
9//! ```
10
11use crate::http_client::{self, HttpClientExt};
12use crate::wasm_compat::{WasmCompatSend, WasmCompatSync};
13use std::sync::Arc;
14
15/// Device authorization details surfaced to a provider callback.
16#[derive(Debug, Clone)]
17pub struct DeviceCodePrompt {
18    /// URL where the user authorizes the device.
19    pub verification_uri: String,
20    /// Short code the user enters at the verification URL.
21    pub user_code: String,
22}
23
24/// Device-code callback with thread-safety bounds outside browser WASM.
25/// Browser WASM authenticators do not invoke it.
26#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
27pub(crate) type DeviceCodeCallback = dyn Fn(DeviceCodePrompt) + Send + Sync;
28#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
29pub(crate) type DeviceCodeCallback = dyn Fn(DeviceCodePrompt);
30
31/// Optional callback invoked when an OAuth device flow needs user action.
32#[derive(Clone, Default)]
33pub struct DeviceCodeHandler(pub(crate) Option<Arc<DeviceCodeCallback>>);
34
35impl DeviceCodeHandler {
36    /// Wraps a device-code callback.
37    pub fn new<F>(handler: F) -> Self
38    where
39        F: Fn(DeviceCodePrompt) + WasmCompatSend + WasmCompatSync + 'static,
40    {
41        Self(Some(Arc::new(handler)))
42    }
43}
44
45impl DeviceCodeHandler {
46    /// Show `prompt` with the callback, or print `fallback_message` to
47    /// stdout when there is none.
48    #[cfg(not(target_family = "wasm"))]
49    pub(crate) fn show(&self, prompt: DeviceCodePrompt, fallback_message: &str) {
50        match &self.0 {
51            Some(callback) => callback(prompt),
52            None => println!("{fallback_message}"),
53        }
54    }
55}
56
57impl std::fmt::Debug for DeviceCodeHandler {
58    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
59        if self.0.is_some() {
60            f.write_str("DeviceCodeHandler(<callback>)")
61        } else {
62            f.write_str("DeviceCodeHandler(None)")
63        }
64    }
65}
66
67#[derive(Debug, thiserror::Error)]
68pub enum AuthError {
69    #[error("{0}")]
70    Message(String),
71    #[error(transparent)]
72    Io(#[from] std::io::Error),
73    #[error(transparent)]
74    Json(#[from] serde_json::Error),
75    /// The HTTP transport failed. Non-success responses arrive as the
76    /// status-bearing [`http_client::Error`] variants (so the status is still
77    /// inspectable); response-less failures as [`http_client::Error::Instance`].
78    #[error(transparent)]
79    Http(#[from] http_client::Error),
80    /// The [`TokenStore`] could not load or save a credential.
81    #[error("token store failed: {0}")]
82    Store(#[source] crate::error::BoxError),
83}
84
85/// The credential a signed-in request is sent with.
86#[derive(Debug, Clone, PartialEq, Eq)]
87pub struct AuthContext {
88    /// The resolved key. Use [`Secret::expose`](crate::wire::Secret::expose)
89    /// only when raw bytes are required.
90    pub api_key: crate::wire::Secret,
91    /// The account the key belongs to, when the provider names one
92    /// (ChatGPT's `ChatGPT-Account-Id`).
93    pub account_id: Option<String>,
94    /// The API root the sign-in names, when it names one (Copilot's
95    /// session endpoint).
96    pub api_base: Option<String>,
97}
98
99/// What every provider's authenticator holds: where its credential comes
100/// from, the store that keeps it, how a device code is shown, whether a
101/// device flow may start, the lock a refresh holds and the credential last
102/// resolved. Clones share the store, the lock and the remembered credential.
103#[derive(Clone)]
104#[cfg_attr(target_family = "wasm", allow(dead_code))]
105pub(crate) struct Shell<S> {
106    pub(crate) source: S,
107    pub(crate) store: Arc<dyn TokenStore>,
108    pub(crate) device_code_handler: DeviceCodeHandler,
109    pub(crate) allow_device_flow: bool,
110    /// Held across a cache refresh to prevent concurrent updates.
111    pub(crate) refresh_lock: Arc<futures::lock::Mutex<()>>,
112    pub(crate) remembered: Remembered,
113}
114
115impl<S> Shell<S> {
116    pub(crate) fn new(
117        source: S,
118        store: impl TokenStore + 'static,
119        device_code_handler: DeviceCodeHandler,
120        allow_device_flow: bool,
121    ) -> Self {
122        Self {
123            source,
124            store: Arc::new(store),
125            device_code_handler,
126            allow_device_flow,
127            refresh_lock: Arc::default(),
128            remembered: Remembered::default(),
129        }
130    }
131}
132
133impl<S> Shell<S> {
134    /// Holds the record under `key` for a refresh, a sign-in or a sign-out:
135    /// the refresh lock its clones share, then the store's lock on `key`
136    /// ([`TokenStore::lock`]), so one runs at a time across clones and
137    /// processes, and each reads what the last one saved.
138    #[cfg_attr(target_family = "wasm", allow(dead_code))]
139    pub(crate) async fn hold(&self, key: &str) -> Result<Held<'_>, AuthError> {
140        let refresh = self.refresh_lock.lock().await;
141        let store = self.store.lock(key).await.map_err(AuthError::Store)?;
142        Ok(Held {
143            _refresh: refresh,
144            _store: store,
145        })
146    }
147}
148
149/// What [`Shell::hold`] holds, let go when dropped: the store's lock first.
150#[cfg_attr(target_family = "wasm", allow(dead_code))]
151pub(crate) struct Held<'a> {
152    _store: StoreLock,
153    _refresh: futures::lock::MutexGuard<'a, ()>,
154}
155
156impl<S: std::fmt::Debug> std::fmt::Debug for Shell<S> {
157    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
158        f.debug_struct("Authenticator")
159            .field("source", &self.source)
160            .field("device_code_handler", &self.device_code_handler)
161            .field("allow_device_flow", &self.allow_device_flow)
162            .finish_non_exhaustive()
163    }
164}
165
166/// The last credential an authenticator resolved and its expiry, shared by
167/// its clones: a live one is served without the store, its lock or a parse.
168#[derive(Clone, Default)]
169pub(crate) struct Remembered(Arc<std::sync::Mutex<Option<(AuthContext, i64)>>>);
170
171impl Remembered {
172    /// The credential while it is live, `skew_seconds` before its expiry.
173    pub(crate) fn live(&self, skew_seconds: i64) -> Option<AuthContext> {
174        let held = self.0.lock().ok()?;
175        let (context, expires_at) = held.as_ref()?;
176        (!token_expired(Some(*expires_at), skew_seconds)).then(|| context.clone())
177    }
178
179    /// Remember `context` until `expires_at`; one with no expiry is not
180    /// remembered.
181    pub(crate) fn keep(&self, context: &AuthContext, expires_at: Option<i64>) {
182        if let Ok(mut held) = self.0.lock() {
183            *held = expires_at.map(|expires_at| (context.clone(), expires_at));
184        }
185    }
186
187    /// Forget the remembered credential.
188    pub(crate) fn forget(&self) {
189        if let Ok(mut held) = self.0.lock() {
190            *held = None;
191        }
192    }
193}
194
195/// Build a request to an absolute authentication URL, independent of the API base.
196pub(crate) fn request(method: http::Method, url: &str) -> http::request::Builder {
197    http::Request::builder().method(method).uri(url)
198}
199
200/// Send `req` and decode its JSON response.
201/// Request and transport failures return [`AuthError::Http`], preserving HTTP
202/// status when available. Invalid JSON returns [`AuthError::Json`].
203pub(crate) async fn send_json<H, T>(
204    http: &H,
205    req: http::Result<http::Request<bytes::Bytes>>,
206) -> Result<T, AuthError>
207where
208    H: HttpClientExt,
209    T: serde::de::DeserializeOwned,
210{
211    let bytes = send_bytes(http, req).await?;
212    Ok(serde_json::from_slice(&bytes)?)
213}
214
215/// Send `req` through the transport and return the raw success body.
216pub(crate) async fn send_bytes<H>(
217    http: &H,
218    req: http::Result<http::Request<bytes::Bytes>>,
219) -> Result<bytes::Bytes, AuthError>
220where
221    H: HttpClientExt,
222{
223    let req = req.map_err(http_client::Error::Protocol)?;
224    let response = http.send::<_, bytes::Bytes>(req).await?;
225    Ok(response.into_body().await?)
226}
227
228/// Returns true when the token is expired (or has no expiry), treating the
229/// token as expired `skew_seconds` before its actual `expires_at`.
230pub(crate) fn token_expired(expires_at: Option<i64>, skew_seconds: i64) -> bool {
231    let now = web_time::SystemTime::now()
232        .duration_since(web_time::UNIX_EPOCH)
233        .map(|duration| duration.as_secs() as i64)
234        .unwrap_or_default();
235
236    expires_at.is_none_or(|exp| now >= exp - skew_seconds)
237}
238
239mod store;
240
241#[cfg(not(target_family = "wasm"))]
242pub use store::FileTokenStore;
243pub use store::{MemoryTokenStore, StoreLock, TokenStore};
244pub(crate) use store::{load_record, save_record};
rig_core/providers/internal/auth.rs

rig_core/providers/internal/
auth.rs
