mirror of
https://github.com/lancedb/lancedb.git
synced 2026-09-22 04:55:39 +00:00
feat: add persistent OAuth token cache and session APIs (#4182)
Stacked on #4173 (diff includes it until that merges; will rebase after). Addresses the token-cache part of [Colin's review](https://github.com/lancedb/lancedb/pull/4173#issuecomment-5674048100). Adds an explicit, opt-in persistent OAuth token cache shared by Rust, Python, and Node clients, plus `login` / `status` / `logout` session APIs, so short-lived processes (CLIs, scripts, notebooks) reuse one session instead of restarting a browser or device flow on every start. - **Opt-in and minimal**: existing callers stay memory-only and lazy. Only refresh tokens are persisted (never access tokens, never client secrets), so there are no local token-expiry decisions to get wrong when clocks move. Each process start performs one silent refresh grant. - **Hardened file backend**: private directory (`0700`), per-record files (`0600`), owner validation, symlink rejection, and atomic `rename` replacement. Corrupt, truncated, unknown-version, or permission-invalid records fail with actionable errors naming the file. Native keyring backends were evaluated (keyring crate routes Linux through D-Bus/zbus: heavy deps, headless/CI flakiness) and are deferred; the file store is the explicit opt-in, not a downgrade from a keyring. - **Cache key**: SHA-256 of the canonical identity (issuer, client ID, sorted/de-duplicated scopes, flow, public/confidential), so no secret appears in a filename and distinct identities never collide. Versioned record schema (`version: 1`). One record per identity: last login wins, documented. - **Cross-process rotation locking**: per-key `fs4` file lock (`flock` / `LockFileEx`) around the refresh critical section — acquire, reread the durable record, refresh exactly once, atomically store the rotated refresh token, release. The OS releases locks on process death, so crashes cannot strand stale locks. Only confirmed `invalid_grant`/`invalid_token` deletes a record and reauthenticates; transport, 5xx, 429, and parse failures retain it. - **Session APIs**: `OAuthSession::login/status/logout` in Rust, `lancedb.remote.OAuthSession` (async) in Python, `OAuthSession` class in Node. `status` returns non-secret metadata only. `logout` removes only the local credential — provider revocation (RFC 7009) is a deliberate follow-up, and local logout never terminates browser SSO. Azure managed identity is rejected for persistence (machine identity stays in memory); client credentials have nothing refreshable to persist and stay memory-only. - No CLI binary exists in this repo, so this ships library APIs plus doc examples in all three languages. Tests: Rust unit + mock-IdP integration (cache-key canonicalization/separation, record versioning/corruption/truncation/symlink/owner/perms, lock serialization + release, two concurrent providers proving no `invalid_grant` and correct rotation, transient-failure retention, `invalid_grant` delete + reauthenticate, login/status/logout lifecycle, client-credentials no-op, IMDS rejection, secret redaction); Python lifecycle + a true two-subprocess cross-process reuse test (second process refreshes once, never hits the device endpoint); Node lifecycle + device-flow login test. Local builds were skipped in development; CI validates all bindings. --------- Co-authored-by: Xuanwo <github@xuanwo.io>
This commit is contained in:
Jack Ye
Xuanwo
+222
-2
@@ -6,6 +6,8 @@ use std::collections::HashMap;
|
||||
use lancedb::error::Error;
|
||||
use napi_derive::*;
|
||||
|
||||
use crate::error::NapiErrorExt;
|
||||
|
||||
/// Timeout configuration for remote HTTP client.
|
||||
#[napi(object)]
|
||||
#[derive(Debug)]
|
||||
@@ -141,6 +143,34 @@ impl From<TlsConfig> for lancedb::remote::TlsConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// Options for the persistent OAuth token cache.
|
||||
///
|
||||
/// The cache is opt-in: it is only used when set as `tokenCache` on
|
||||
/// `OAuthConfig`. Only refresh tokens are persisted, in a private directory
|
||||
/// with owner-only permissions, so short-lived processes can reuse an
|
||||
/// authenticated session instead of re-prompting on every start.
|
||||
#[napi(object)]
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct TokenCacheOptions {
|
||||
/// Directory that holds cached credentials. Defaults to
|
||||
/// `$XDG_CACHE_HOME/lancedb/oauth`, `$HOME/.cache/lancedb/oauth` on Unix,
|
||||
/// or `%LOCALAPPDATA%\lancedb\oauth` on Windows. The directory is created
|
||||
/// with owner-only permissions (`0700`) when missing.
|
||||
pub cache_dir: Option<String>,
|
||||
/// How long to wait for the cross-process refresh lock before failing,
|
||||
/// in seconds (default: 30).
|
||||
pub lock_timeout_secs: Option<u32>,
|
||||
}
|
||||
|
||||
impl From<TokenCacheOptions> for lancedb::remote::TokenCacheOptions {
|
||||
fn from(options: TokenCacheOptions) -> Self {
|
||||
Self {
|
||||
cache_dir: options.cache_dir.map(std::path::PathBuf::from),
|
||||
lock_timeout_secs: options.lock_timeout_secs.map(|secs| secs as u64),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// OAuth configuration for LanceDB authentication.
|
||||
///
|
||||
/// This is the generated napi-rs binding shape. TypeScript users should prefer
|
||||
@@ -158,16 +188,26 @@ pub struct OAuthConfig {
|
||||
/// OAuth scopes to request. For Azure managed identity, exactly one scope
|
||||
/// or resource is required. For example: `["api://{app_id}/.default"]`
|
||||
pub scopes: Vec<String>,
|
||||
/// Authentication flow: "client_credentials" or "azure_managed_identity"
|
||||
/// Authentication flow: "client_credentials", "authorization_code",
|
||||
/// "device_code", or "azure_managed_identity"
|
||||
pub flow: Option<String>,
|
||||
/// Client secret (required for client_credentials).
|
||||
pub client_secret: Option<String>,
|
||||
/// Loopback redirect URI for authorization_code.
|
||||
pub redirect_uri: Option<String>,
|
||||
/// Port for the authorization_code loopback callback server.
|
||||
pub callback_port: Option<u16>,
|
||||
/// Whether authorization_code uses S256 PKCE (default: true).
|
||||
pub use_pkce: Option<bool>,
|
||||
/// Client ID for user-assigned managed identity (azure_managed_identity).
|
||||
pub managed_identity_client_id: Option<String>,
|
||||
/// Seconds before expiry to trigger proactive refresh (default: 300).
|
||||
/// Keep this well below the token TTL; if it is greater than or equal to
|
||||
/// the TTL, each request refreshes the token.
|
||||
pub refresh_buffer_secs: Option<u32>,
|
||||
/// Opt in to the persistent token cache so short-lived processes reuse
|
||||
/// one session. Only refresh tokens are persisted.
|
||||
pub token_cache: Option<TokenCacheOptions>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for OAuthConfig {
|
||||
@@ -181,11 +221,15 @@ impl std::fmt::Debug for OAuthConfig {
|
||||
"client_secret",
|
||||
&self.client_secret.as_deref().map(|_| "<redacted>"),
|
||||
)
|
||||
.field("redirect_uri", &self.redirect_uri)
|
||||
.field("callback_port", &self.callback_port)
|
||||
.field("use_pkce", &self.use_pkce)
|
||||
.field(
|
||||
"managed_identity_client_id",
|
||||
&self.managed_identity_client_id,
|
||||
)
|
||||
.field("refresh_buffer_secs", &self.refresh_buffer_secs)
|
||||
.field("token_cache", &self.token_cache)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
@@ -194,10 +238,22 @@ impl TryFrom<OAuthConfig> for lancedb::remote::oauth::OAuthConfig {
|
||||
type Error = Error;
|
||||
|
||||
fn try_from(config: OAuthConfig) -> Result<Self, Self::Error> {
|
||||
use lancedb::remote::oauth::OAuthFlow;
|
||||
use lancedb::remote::oauth::{AuthorizationCodeOptions, OAuthFlow};
|
||||
|
||||
let flow = match config.flow.as_deref().unwrap_or("client_credentials") {
|
||||
"client_credentials" => OAuthFlow::ClientCredentials,
|
||||
"authorization_code" => {
|
||||
let mut options =
|
||||
AuthorizationCodeOptions::new().use_pkce(config.use_pkce.unwrap_or(true));
|
||||
if let Some(redirect_uri) = config.redirect_uri {
|
||||
options = options.redirect_uri(redirect_uri);
|
||||
}
|
||||
if let Some(callback_port) = config.callback_port {
|
||||
options = options.callback_port(callback_port);
|
||||
}
|
||||
OAuthFlow::AuthorizationCode(options)
|
||||
}
|
||||
"device_code" => OAuthFlow::DeviceCode,
|
||||
"azure_managed_identity" => OAuthFlow::AzureManagedIdentity {
|
||||
client_id: config.managed_identity_client_id,
|
||||
},
|
||||
@@ -215,10 +271,115 @@ impl TryFrom<OAuthConfig> for lancedb::remote::oauth::OAuthConfig {
|
||||
scopes: config.scopes,
|
||||
flow,
|
||||
refresh_buffer_secs: config.refresh_buffer_secs.map(|v| v as u64),
|
||||
token_cache: config.token_cache.map(Into::into),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Safe, non-secret view of a cached OAuth session, returned by
|
||||
/// `OAuthSession.status()` and `OAuthSession.login()`.
|
||||
#[napi(object)]
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct SessionStatus {
|
||||
/// Whether a cached session exists that can obtain tokens without
|
||||
/// interactive authentication.
|
||||
pub refreshable: bool,
|
||||
/// Canonical issuer URL of the cached session.
|
||||
pub issuer_url: String,
|
||||
/// Client ID of the cached session.
|
||||
pub client_id: String,
|
||||
/// Canonical (sorted, de-duplicated) scopes of the cached session.
|
||||
pub scopes: Vec<String>,
|
||||
/// Flow that produced the cached session.
|
||||
pub flow: String,
|
||||
/// When the cached session was obtained, as Unix seconds.
|
||||
pub obtained_at: Option<f64>,
|
||||
}
|
||||
|
||||
/// Result of `OAuthSession.logout()`.
|
||||
#[napi(object)]
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct SessionLogout {
|
||||
/// Whether a cached credential was removed. `false` means no matching
|
||||
/// session was cached; logout is idempotent.
|
||||
pub removed: bool,
|
||||
}
|
||||
|
||||
/// Explicit OAuth session lifecycle for the persistent token cache: eager
|
||||
/// `login`, non-secret `status`, and local `logout`.
|
||||
///
|
||||
/// A session is built from the same `OAuthConfig` used to connect (including
|
||||
/// its `tokenCache` options). A connection created with the same
|
||||
/// configuration shares the cache, so logging in here prepares tokens for
|
||||
/// later processes without any database request.
|
||||
#[napi]
|
||||
pub struct OAuthSession {
|
||||
inner: lancedb::remote::OAuthSession,
|
||||
}
|
||||
|
||||
#[napi]
|
||||
impl OAuthSession {
|
||||
/// Create a session manager for the given OAuth configuration.
|
||||
///
|
||||
/// The configuration must enable `tokenCache` options and use a flow that
|
||||
/// supports persistent sessions (authorization code or device code).
|
||||
#[napi(constructor)]
|
||||
pub fn new(config: OAuthConfig) -> napi::Result<Self> {
|
||||
let config: lancedb::remote::oauth::OAuthConfig = config.try_into().default_error()?;
|
||||
let inner = lancedb::remote::OAuthSession::new(config).default_error()?;
|
||||
Ok(Self { inner })
|
||||
}
|
||||
|
||||
/// Eagerly run the configured authentication flow and store the session.
|
||||
///
|
||||
/// A successful login always replaces any prior cached session for this
|
||||
/// identity; if the provider does not issue a refresh token (for example
|
||||
/// without `offline_access`), the previous record is removed and the
|
||||
/// status reports `refreshable == false`.
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn login(&self) -> napi::Result<SessionStatus> {
|
||||
let status = self.inner.login().await.default_error()?;
|
||||
Ok(SessionStatus::from(status))
|
||||
}
|
||||
|
||||
/// Report whether a matching cached session exists, with safe metadata.
|
||||
///
|
||||
/// This never contacts the identity provider and never exposes token
|
||||
/// values.
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn status(&self) -> napi::Result<SessionStatus> {
|
||||
let status = self.inner.status().await.default_error()?;
|
||||
Ok(SessionStatus::from(status))
|
||||
}
|
||||
|
||||
/// Remove the matching local cached credential.
|
||||
///
|
||||
/// This only deletes the local cache entry. It does not revoke the
|
||||
/// refresh token with the provider and does not sign out of a browser
|
||||
/// SSO session. Repeated calls succeed; `removed` reports whether a
|
||||
/// credential existed.
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn logout(&self) -> napi::Result<SessionLogout> {
|
||||
let logout = self.inner.logout().await.default_error()?;
|
||||
Ok(SessionLogout {
|
||||
removed: logout.removed,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl From<lancedb::remote::SessionStatus> for SessionStatus {
|
||||
fn from(status: lancedb::remote::SessionStatus) -> Self {
|
||||
Self {
|
||||
refreshable: status.refreshable,
|
||||
issuer_url: status.issuer_url,
|
||||
client_id: status.client_id,
|
||||
scopes: status.scopes,
|
||||
flow: status.flow,
|
||||
obtained_at: status.obtained_at.map(|secs| secs as f64),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<ClientConfig> for lancedb::remote::ClientConfig {
|
||||
fn from(config: ClientConfig) -> Self {
|
||||
Self {
|
||||
@@ -252,8 +413,12 @@ mod tests {
|
||||
scopes: vec!["scope".to_string()],
|
||||
flow: Some("typo".to_string()),
|
||||
client_secret: None,
|
||||
redirect_uri: None,
|
||||
callback_port: None,
|
||||
use_pkce: None,
|
||||
managed_identity_client_id: None,
|
||||
refresh_buffer_secs: None,
|
||||
token_cache: None,
|
||||
};
|
||||
|
||||
let err = lancedb::remote::oauth::OAuthConfig::try_from(config).unwrap_err();
|
||||
@@ -272,12 +437,67 @@ mod tests {
|
||||
scopes: vec!["scope".to_string()],
|
||||
flow: Some("client_credentials".to_string()),
|
||||
client_secret: Some("super-secret".to_string()),
|
||||
redirect_uri: None,
|
||||
callback_port: None,
|
||||
use_pkce: None,
|
||||
managed_identity_client_id: None,
|
||||
refresh_buffer_secs: None,
|
||||
token_cache: None,
|
||||
};
|
||||
|
||||
let debug = format!("{config:?}");
|
||||
assert!(!debug.contains("super-secret"));
|
||||
assert!(debug.contains("client_secret: Some(\"<redacted>\")"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_authorization_code_conversion_preserves_options() {
|
||||
let config = OAuthConfig {
|
||||
issuer_url: "https://issuer.example.com".to_string(),
|
||||
client_id: "client-id".to_string(),
|
||||
scopes: vec!["openid".to_string()],
|
||||
flow: Some("authorization_code".to_string()),
|
||||
client_secret: Some("secret".to_string()),
|
||||
redirect_uri: Some("http://127.0.0.1:9000/callback".to_string()),
|
||||
callback_port: Some(9000),
|
||||
use_pkce: Some(false),
|
||||
managed_identity_client_id: None,
|
||||
refresh_buffer_secs: None,
|
||||
token_cache: None,
|
||||
};
|
||||
|
||||
let converted = lancedb::remote::oauth::OAuthConfig::try_from(config).unwrap();
|
||||
let lancedb::remote::oauth::OAuthFlow::AuthorizationCode(options) = converted.flow else {
|
||||
panic!("expected authorization code flow");
|
||||
};
|
||||
assert_eq!(
|
||||
options.redirect_uri.as_deref(),
|
||||
Some("http://127.0.0.1:9000/callback")
|
||||
);
|
||||
assert_eq!(options.callback_port, Some(9000));
|
||||
assert!(!options.use_pkce);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_device_code_conversion() {
|
||||
let config = OAuthConfig {
|
||||
issuer_url: "https://issuer.example.com".to_string(),
|
||||
client_id: "client-id".to_string(),
|
||||
scopes: vec!["openid".to_string()],
|
||||
flow: Some("device_code".to_string()),
|
||||
client_secret: None,
|
||||
redirect_uri: None,
|
||||
callback_port: None,
|
||||
use_pkce: None,
|
||||
managed_identity_client_id: None,
|
||||
refresh_buffer_secs: None,
|
||||
token_cache: None,
|
||||
};
|
||||
|
||||
let converted = lancedb::remote::oauth::OAuthConfig::try_from(config).unwrap();
|
||||
assert!(matches!(
|
||||
converted.flow,
|
||||
lancedb::remote::oauth::OAuthFlow::DeviceCode
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user