diff --git a/Cargo.lock b/Cargo.lock index 24889deec..090e2c654 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3088,6 +3088,16 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "dispatch2" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e0e367e4e7da84520dedcac1901e4da967309406d1e51017ae1abfb97adbd38" +dependencies = [ + "bitflags 2.11.1", + "objc2", +] + [[package]] name = "displaydoc" version = "0.2.5" @@ -3483,6 +3493,16 @@ version = "1.20260804.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "82eb03a32a1d50555353c85a7b9d3279a6f1e91af9890b789acdf544ed57c8d7" +[[package]] +name = "fs4" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8640e34b88f7652208ce9e88b1a37a2ae95227d84abec377ccd3c5cfeb141ed4" +dependencies = [ + "rustix", + "windows-sys 0.59.0", +] + [[package]] name = "fs_extra" version = "1.3.0" @@ -5477,6 +5497,7 @@ dependencies = [ "aws-sdk-s3", "aws-smithy-runtime", "aws-smithy-types", + "base64 0.22.1", "bytes", "candle-core", "candle-nn", @@ -5491,6 +5512,7 @@ dependencies = [ "datafusion-physical-expr", "datafusion-physical-plan", "datafusion-sql", + "fs4", "futures", "half", "hf-hub", @@ -5543,6 +5565,7 @@ dependencies = [ "urlencoding", "uuid", "walkdir", + "webbrowser", ] [[package]] @@ -6227,6 +6250,12 @@ dependencies = [ "rawpointer", ] +[[package]] +name = "ndk-context" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27b02d87554356db9e9a873add8782d4ea6e3e58ea071a9adb9a2e8ddb884a8b" + [[package]] name = "nibble_vec" version = "0.1.0" @@ -6409,6 +6438,26 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" +[[package]] +name = "objc2" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a12a8ed07aefc768292f076dc3ac8c48f3781c8f2d5851dd3d98950e8c5a89f" +dependencies = [ + "objc2-encode", +] + +[[package]] +name = "objc2-app-kit" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d49e936b501e5c5bf01fda3a9452ff86dc3ea98ad5f283e1455153142d97518c" +dependencies = [ + "bitflags 2.11.1", + "objc2", + "objc2-foundation", +] + [[package]] name = "objc2-core-foundation" version = "0.3.2" @@ -6416,6 +6465,25 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" dependencies = [ "bitflags 2.11.1", + "dispatch2", + "objc2", +] + +[[package]] +name = "objc2-encode" +version = "4.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ef25abbcd74fb2609453eb695bd2f860d389e457f67dc17cafc8b8cbc89d0c33" + +[[package]] +name = "objc2-foundation" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3e0adef53c21f888deb4fa59fc59f7eb17404926ee8a6f59f5df0fd7f9f3272" +dependencies = [ + "bitflags 2.11.1", + "objc2", + "objc2-core-foundation", ] [[package]] @@ -10788,6 +10856,22 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "webbrowser" +version = "1.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62c35be770821a214dbc362fc26908c853e776c0004294d0b10b8a6bad582f94" +dependencies = [ + "jni", + "log", + "ndk-context", + "objc2", + "objc2-app-kit", + "objc2-foundation", + "url", + "web-sys", +] + [[package]] name = "webpki-root-certs" version = "1.0.7" diff --git a/docs/src/js/classes/OAuthSession.md b/docs/src/js/classes/OAuthSession.md new file mode 100644 index 000000000..e80202bea --- /dev/null +++ b/docs/src/js/classes/OAuthSession.md @@ -0,0 +1,105 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / OAuthSession + +# Class: OAuthSession + +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](../interfaces/OAuthConfig.md) 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. + +`login` always runs the configured interactive flow and replaces the cached +session (the most recent login wins). `logout` removes only the local +credential; it does not revoke anything with the provider and does not sign +out of a browser SSO session. + +## Example + +```typescript +const config: OAuthConfig = { + issuerUrl: "https://issuer.example.com", + clientId: "my-app", + scopes: ["openid", "offline_access"], + flow: OAuthFlowType.DeviceCode, + tokenCache: { cacheDir: "/tmp/my-app/oauth-cache" }, +}; +const session = new OAuthSession(config); +const status = await session.login(); +``` + +## Constructors + +### new OAuthSession() + +```ts +new OAuthSession(config): OAuthSession +``` + +Create a session manager for the given OAuth configuration. + +#### Parameters + +* **config**: [`OAuthConfig`](../interfaces/OAuthConfig.md) + +#### Returns + +[`OAuthSession`](OAuthSession.md) + +## Methods + +### login() + +```ts +login(): Promise +``` + +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`. + +#### Returns + +`Promise`<[`SessionStatus`](../interfaces/SessionStatus.md)> + +*** + +### logout() + +```ts +logout(): Promise +``` + +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. + +#### Returns + +`Promise`<[`SessionLogout`](../interfaces/SessionLogout.md)> + +*** + +### status() + +```ts +status(): Promise +``` + +Report whether a matching cached session exists, with safe metadata. + +This never contacts the identity provider and never exposes token values. + +#### Returns + +`Promise`<[`SessionStatus`](../interfaces/SessionStatus.md)> diff --git a/docs/src/js/enumerations/OAuthFlowType.md b/docs/src/js/enumerations/OAuthFlowType.md index fe546a140..7daae22fb 100644 --- a/docs/src/js/enumerations/OAuthFlowType.md +++ b/docs/src/js/enumerations/OAuthFlowType.md @@ -10,6 +10,16 @@ OAuth authentication flow types. ## Enumeration Members +### AuthorizationCode + +```ts +AuthorizationCode: "authorization_code"; +``` + +Interactive Authorization Code grant, using PKCE by default. + +*** + ### AzureManagedIdentity ```ts @@ -27,3 +37,13 @@ ClientCredentials: "client_credentials"; ``` Client Credentials grant (service-to-service / M2M). + +*** + +### DeviceCode + +```ts +DeviceCode: "device_code"; +``` + +Device Authorization grant for CLI and headless environments. diff --git a/docs/src/js/globals.md b/docs/src/js/globals.md index 4a5effaae..422c6c428 100644 --- a/docs/src/js/globals.md +++ b/docs/src/js/globals.md @@ -35,6 +35,7 @@ - [MultiMatchQuery](classes/MultiMatchQuery.md) - [NativeJsHeaderProvider](classes/NativeJsHeaderProvider.md) - [OAuthHeaderProvider](classes/OAuthHeaderProvider.md) +- [OAuthSession](classes/OAuthSession.md) - [PermutationBuilder](classes/PermutationBuilder.md) - [PhraseQuery](classes/PhraseQuery.md) - [Query](classes/Query.md) @@ -122,6 +123,8 @@ - [RestNamespaceConfig](interfaces/RestNamespaceConfig.md) - [RetryConfig](interfaces/RetryConfig.md) - [ScannableOptions](interfaces/ScannableOptions.md) +- [SessionLogout](interfaces/SessionLogout.md) +- [SessionStatus](interfaces/SessionStatus.md) - [ShuffleOptions](interfaces/ShuffleOptions.md) - [SplitCalculatedOptions](interfaces/SplitCalculatedOptions.md) - [SplitHashOptions](interfaces/SplitHashOptions.md) @@ -131,6 +134,7 @@ - [TableStatistics](interfaces/TableStatistics.md) - [TimeoutConfig](interfaces/TimeoutConfig.md) - [TlsConfig](interfaces/TlsConfig.md) +- [TokenCacheOptions](interfaces/TokenCacheOptions.md) - [TokenResponse](interfaces/TokenResponse.md) - [TokenizeOptions](interfaces/TokenizeOptions.md) - [UpdateFieldMetadataResult](interfaces/UpdateFieldMetadataResult.md) diff --git a/docs/src/js/interfaces/NativeOAuthConfig.md b/docs/src/js/interfaces/NativeOAuthConfig.md index 6959f17dd..afe05de9f 100644 --- a/docs/src/js/interfaces/NativeOAuthConfig.md +++ b/docs/src/js/interfaces/NativeOAuthConfig.md @@ -15,6 +15,16 @@ All token acquisition and refresh is handled in the Rust layer. ## Properties +### callbackPort? + +```ts +optional callbackPort: number; +``` + +Port for the authorization_code loopback callback server. + +*** + ### clientId ```ts @@ -41,7 +51,8 @@ Client secret (required for client_credentials). optional flow: string; ``` -Authentication flow: "client_credentials" or "azure_managed_identity" +Authentication flow: "client_credentials", "authorization_code", +"device_code", or "azure_managed_identity" *** @@ -66,6 +77,16 @@ Client ID for user-assigned managed identity (azure_managed_identity). *** +### redirectUri? + +```ts +optional redirectUri: string; +``` + +Loopback redirect URI for authorization_code. + +*** + ### refreshBufferSecs? ```ts @@ -86,3 +107,24 @@ scopes: string[]; OAuth scopes to request. For Azure managed identity, exactly one scope or resource is required. For example: `["api://{app_id}/.default"]` + +*** + +### tokenCache? + +```ts +optional tokenCache: TokenCacheOptions; +``` + +Opt in to the persistent token cache so short-lived processes reuse +one session. Only refresh tokens are persisted. + +*** + +### usePkce? + +```ts +optional usePkce: boolean; +``` + +Whether authorization_code uses S256 PKCE (default: true). diff --git a/docs/src/js/interfaces/OAuthConfig.md b/docs/src/js/interfaces/OAuthConfig.md index f9d5d1c7b..f7d61c663 100644 --- a/docs/src/js/interfaces/OAuthConfig.md +++ b/docs/src/js/interfaces/OAuthConfig.md @@ -35,8 +35,32 @@ const config: OAuthConfig = { }; ``` +The authorization URL is written to stderr before LanceDB tries to open a +browser, so it can be copied in headless environments. +```typescript +const config: OAuthConfig = { + issuerUrl: "https://login.microsoftonline.com/{tenant}/v2.0", + clientId: "app-id", + scopes: ["openid", "api://lancedb-api/access"], + flow: OAuthFlowType.AuthorizationCode, +}; +``` + +Device Authorization writes the verification URL and user code to stderr +before polling begins. + ## Properties +### callbackPort? + +```ts +optional callbackPort: number; +``` + +Port for the AuthorizationCode loopback callback server (default: 8400). + +*** + ### clientId ```ts @@ -88,6 +112,16 @@ Client ID for user-assigned managed identity (AzureManagedIdentity). *** +### redirectUri? + +```ts +optional redirectUri: string; +``` + +Loopback redirect URI for AuthorizationCode. + +*** + ### refreshBufferSecs? ```ts @@ -109,3 +143,26 @@ scopes: string[]; OAuth scopes to request. For Azure managed identity, exactly one scope or resource is required. For example: `["api://{app_id}/.default"]` + +*** + +### tokenCache? + +```ts +optional tokenCache: TokenCacheOptions; +``` + +Opt in to the persistent token cache so short-lived processes reuse one +session. Only refresh tokens are persisted. Only supported by +AuthorizationCode and DeviceCode; Azure managed identity is rejected. +Default: unset (memory only). + +*** + +### usePkce? + +```ts +optional usePkce: boolean; +``` + +Protect AuthorizationCode with S256 PKCE (default: true). diff --git a/docs/src/js/interfaces/SessionLogout.md b/docs/src/js/interfaces/SessionLogout.md new file mode 100644 index 000000000..c91fe5108 --- /dev/null +++ b/docs/src/js/interfaces/SessionLogout.md @@ -0,0 +1,20 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / SessionLogout + +# Interface: SessionLogout + +Result of [OAuthSession.logout](../classes/OAuthSession.md#logout). + +## Properties + +### removed + +```ts +removed: boolean; +``` + +Whether a cached credential was removed. `false` means no matching +session was cached; logout is idempotent. diff --git a/docs/src/js/interfaces/SessionStatus.md b/docs/src/js/interfaces/SessionStatus.md new file mode 100644 index 000000000..3844bfae6 --- /dev/null +++ b/docs/src/js/interfaces/SessionStatus.md @@ -0,0 +1,74 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / SessionStatus + +# Interface: SessionStatus + +Safe, non-secret view of a cached OAuth session, returned by +[OAuthSession.status](../classes/OAuthSession.md#status) and [OAuthSession.login](../classes/OAuthSession.md#login). + +## Properties + +### clientId + +```ts +clientId: string; +``` + +Client ID of the cached session. + +*** + +### flow + +```ts +flow: string; +``` + +Flow that produced the cached session. + +*** + +### issuerUrl + +```ts +issuerUrl: string; +``` + +Canonical issuer URL of the cached session. + +*** + +### obtainedAt? + +```ts +optional obtainedAt: number; +``` + +When the cached session was obtained, as Unix seconds. + +*** + +### refreshable + +```ts +refreshable: boolean; +``` + +Whether a cached session exists that can obtain tokens without +interactive authentication. Because access tokens are not persisted, +this is `true` exactly when a refresh token is cached; the next +connection refreshes with it rather than opening a browser or device +prompt. + +*** + +### scopes + +```ts +scopes: string[]; +``` + +Canonical (sorted, de-duplicated) scope set of the cached session. diff --git a/docs/src/js/interfaces/TokenCacheOptions.md b/docs/src/js/interfaces/TokenCacheOptions.md new file mode 100644 index 000000000..95b96ef2f --- /dev/null +++ b/docs/src/js/interfaces/TokenCacheOptions.md @@ -0,0 +1,41 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / TokenCacheOptions + +# Interface: TokenCacheOptions + +Options for the persistent OAuth token cache. + +The cache is opt-in: it is only used when set as `tokenCache` on +[OAuthConfig](OAuthConfig.md). 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. + +Multiple identities (issuer, client, scopes, flow, client authentication) +get separate cache entries. Within one identity the most recent login wins. + +## Properties + +### cacheDir? + +```ts +optional cacheDir: string; +``` + +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. + +*** + +### lockTimeoutSecs? + +```ts +optional lockTimeoutSecs: number; +``` + +How long to wait for the cross-process refresh lock before failing, in +seconds (default: 30). diff --git a/nodejs/__test__/oauth.test.ts b/nodejs/__test__/oauth.test.ts new file mode 100644 index 000000000..83ea2791b --- /dev/null +++ b/nodejs/__test__/oauth.test.ts @@ -0,0 +1,211 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import * as fs from "fs"; +import * as http from "http"; +import * as os from "os"; +import * as path from "path"; +import { OAuthConfig, OAuthFlowType, OAuthSession } from "../lancedb/oauth"; + +function tempCacheDir(): string { + return fs.mkdtempSync(path.join(os.tmpdir(), "lancedb-oauth-cache-")); +} + +function deviceConfig(issuerUrl: string, cacheDir: string): OAuthConfig { + return { + issuerUrl, + clientId: "client-id", + scopes: ["openid"], + flow: OAuthFlowType.DeviceCode, + tokenCache: { cacheDir }, + }; +} + +describe("OAuthSession", () => { + beforeAll(() => { + // Point the Rust browser helper at a no-op so device-flow logins never + // open a real browser window during tests. + process.env.LANCEDB_OAUTH_BROWSER = "/usr/bin/true"; + }); + + it("reports an absent session and logout is idempotent", async () => { + const cacheDir = tempCacheDir(); + const session = new OAuthSession( + deviceConfig("https://issuer.example.com", cacheDir), + ); + + const status = await session.status(); + expect(status.refreshable).toBe(false); + expect(status.issuerUrl).toBe("https://issuer.example.com"); + expect(status.clientId).toBe("client-id"); + expect(status.scopes).toEqual(["openid"]); + expect(status.flow).toBe("device_code"); + expect(status.obtainedAt).toBeUndefined(); + + const logout = await session.logout(); + expect(logout.removed).toBe(false); + }); + + it("requires token cache options", () => { + const config: OAuthConfig = { + issuerUrl: "https://issuer.example.com", + clientId: "client-id", + scopes: ["openid"], + flow: OAuthFlowType.DeviceCode, + }; + expect(() => new OAuthSession(config)).toThrow(/token/); + }); + + it("rejects azure managed identity persistence", () => { + const config: OAuthConfig = { + issuerUrl: "https://login.microsoftonline.com/tenant/v2.0", + clientId: "app-id", + scopes: ["api://app/.default"], + flow: OAuthFlowType.AzureManagedIdentity, + tokenCache: { cacheDir: tempCacheDir() }, + }; + expect(() => new OAuthSession(config)).toThrow(/AzureManagedIdentity/); + }); + + it("logs in via device flow, caches, and logs out", async () => { + const server = new MockIdp(); + await server.start(); + try { + const cacheDir = tempCacheDir(); + const issuerUrl = server.issuerUrl(); + + const session = new OAuthSession(deviceConfig(issuerUrl, cacheDir)); + const status = await session.login(); + expect(status.refreshable).toBe(true); + expect(status.obtainedAt).toBeGreaterThan(0); + expect(server.state.deviceAuthorizations).toBe(1); + + // An independent session (a fresh "process") sees the cached login. + const other = new OAuthSession(deviceConfig(issuerUrl, cacheDir)); + const cached = await other.status(); + expect(cached.refreshable).toBe(true); + + const logout = await other.logout(); + expect(logout.removed).toBe(true); + const again = await session.logout(); + expect(again.removed).toBe(false); + expect((await session.status()).refreshable).toBe(false); + + // Only the initial login used the interactive device flow. + expect(server.state.deviceAuthorizations).toBe(1); + expect(server.state.refreshGrants).toBe(0); + } finally { + server.close(); + } + }, 15000); +}); + +/** Mock IdP with discovery, device authorization, and rotating refresh. */ +class MockIdp { + readonly state = { + deviceAuthorizations: 0, + refreshGrants: 0, + accessTokensIssued: 0, + currentRefresh: null as string | null, + }; + private server?: http.Server; + private port = 0; + + issuerUrl(): string { + return `http://127.0.0.1:${this.port}`; + } + + async start(): Promise { + const server = http.createServer((req, res) => { + const chunks: Buffer[] = []; + req.on("data", (chunk) => chunks.push(chunk)); + req.on("end", () => { + const body = Buffer.concat(chunks).toString(); + const params = new URLSearchParams(body); + this.handle(req.url ?? "", params, res); + }); + }); + await new Promise((resolve) => { + server.listen(0, "127.0.0.1", () => resolve()); + }); + const address = server.address(); + if (address && typeof address === "object") { + this.port = address.port; + } + this.server = server; + } + + private handle( + url: string, + params: URLSearchParams, + res: http.ServerResponse, + ): void { + const respond = (status: number, payload: unknown): void => { + const body = JSON.stringify(payload); + res.writeHead(status, { + "Content-Type": "application/json", + "Content-Length": Buffer.byteLength(body), + }); + res.end(body); + }; + + if (url === "/.well-known/openid-configuration") { + respond(200, { + // biome-ignore lint/style/useNamingConvention: OAuth wire format + token_endpoint: `${this.issuerUrl()}/token`, + // biome-ignore lint/style/useNamingConvention: OAuth wire format + device_authorization_endpoint: `${this.issuerUrl()}/device`, + }); + return; + } + + if (url === "/device") { + this.state.deviceAuthorizations += 1; + respond(200, { + // biome-ignore lint/style/useNamingConvention: OAuth wire format + device_code: "device-code", + // biome-ignore lint/style/useNamingConvention: OAuth wire format + user_code: "ABCD-EFGH", + // biome-ignore lint/style/useNamingConvention: OAuth wire format + verification_uri: `${this.issuerUrl()}/verify`, + // biome-ignore lint/style/useNamingConvention: OAuth wire format + expires_in: 60, + interval: 1, + }); + return; + } + + if (url === "/token") { + const grantType = params.get("grant_type") ?? ""; + if (grantType === "refresh_token") { + this.state.refreshGrants += 1; + if (params.get("refresh_token") !== this.state.currentRefresh) { + respond(400, { error: "invalid_grant" }); + return; + } + } else if (!grantType.includes("device_code")) { + respond(400, { error: "unsupported_grant_type" }); + return; + } + this.state.accessTokensIssued += 1; + const number = this.state.accessTokensIssued; + const refresh = `refresh-${number}`; + this.state.currentRefresh = refresh; + respond(200, { + // biome-ignore lint/style/useNamingConvention: OAuth wire format + access_token: `access-${number}`, + // biome-ignore lint/style/useNamingConvention: OAuth wire format + refresh_token: refresh, + // biome-ignore lint/style/useNamingConvention: OAuth wire format + expires_in: 3600, + }); + return; + } + + respond(404, {}); + } + + close(): void { + this.server?.close(); + } +} diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index d94007a11..a58cb182e 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -170,7 +170,14 @@ export { TokenResponse, } from "./header"; -export { OAuthConfig, OAuthFlowType } from "./oauth"; +export { + OAuthConfig, + OAuthFlowType, + OAuthSession, + SessionLogout, + SessionStatus, + TokenCacheOptions, +} from "./oauth"; export { MergeInsertBuilder, WriteExecutionOptions } from "./merge"; diff --git a/nodejs/lancedb/oauth.ts b/nodejs/lancedb/oauth.ts index 345eda87a..c9c78afa6 100644 --- a/nodejs/lancedb/oauth.ts +++ b/nodejs/lancedb/oauth.ts @@ -1,16 +1,52 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors +import { + OAuthConfig as NativeOAuthConfig, + OAuthSession as NativeOAuthSession, +} from "./native"; + /** * OAuth authentication flow types. */ export enum OAuthFlowType { /** Client Credentials grant (service-to-service / M2M). */ ClientCredentials = "client_credentials", + /** Interactive Authorization Code grant, using PKCE by default. */ + AuthorizationCode = "authorization_code", + /** Device Authorization grant for CLI and headless environments. */ + DeviceCode = "device_code", /** Azure Managed Identity via IMDS. */ AzureManagedIdentity = "azure_managed_identity", } +/** + * Options for the persistent OAuth token cache. + * + * The cache is opt-in: it is only used when set as `tokenCache` on + * {@link 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. + * + * Multiple identities (issuer, client, scopes, flow, client authentication) + * get separate cache entries. Within one identity the most recent login wins. + */ +export interface 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. + */ + cacheDir?: string; + + /** + * How long to wait for the cross-process refresh lock before failing, in + * seconds (default: 30). + */ + lockTimeoutSecs?: number; +} + /** * OAuth configuration for LanceDB authentication. * @@ -40,6 +76,21 @@ export enum OAuthFlowType { * flow: OAuthFlowType.AzureManagedIdentity, * }; * ``` + * + * @example Authorization Code with PKCE: + * The authorization URL is written to stderr before LanceDB tries to open a + * browser, so it can be copied in headless environments. + * ```typescript + * const config: OAuthConfig = { + * issuerUrl: "https://login.microsoftonline.com/{tenant}/v2.0", + * clientId: "app-id", + * scopes: ["openid", "api://lancedb-api/access"], + * flow: OAuthFlowType.AuthorizationCode, + * }; + * ``` + * + * Device Authorization writes the verification URL and user code to stderr + * before polling begins. */ export interface OAuthConfig { /** @@ -64,6 +115,15 @@ export interface OAuthConfig { /** Client secret (required for ClientCredentials). */ clientSecret?: string; + /** Loopback redirect URI for AuthorizationCode. */ + redirectUri?: string; + + /** Port for the AuthorizationCode loopback callback server (default: 8400). */ + callbackPort?: number; + + /** Protect AuthorizationCode with S256 PKCE (default: true). */ + usePkce?: boolean; + /** Client ID for user-assigned managed identity (AzureManagedIdentity). */ managedIdentityClientId?: string; @@ -73,4 +133,119 @@ export interface OAuthConfig { * the TTL, each request refreshes the token. */ refreshBufferSecs?: number; + + /** + * Opt in to the persistent token cache so short-lived processes reuse one + * session. Only refresh tokens are persisted. Only supported by + * AuthorizationCode and DeviceCode; Azure managed identity is rejected. + * Default: unset (memory only). + */ + tokenCache?: TokenCacheOptions; +} + +/** + * Safe, non-secret view of a cached OAuth session, returned by + * {@link OAuthSession.status} and {@link OAuthSession.login}. + */ +export interface SessionStatus { + /** + * Whether a cached session exists that can obtain tokens without + * interactive authentication. Because access tokens are not persisted, + * this is `true` exactly when a refresh token is cached; the next + * connection refreshes with it rather than opening a browser or device + * prompt. + */ + refreshable: boolean; + + /** Canonical issuer URL of the cached session. */ + issuerUrl: string; + + /** Client ID of the cached session. */ + clientId: string; + + /** Canonical (sorted, de-duplicated) scope set of the cached session. */ + scopes: string[]; + + /** Flow that produced the cached session. */ + flow: string; + + /** When the cached session was obtained, as Unix seconds. */ + obtainedAt?: number; +} + +/** Result of {@link OAuthSession.logout}. */ +export interface SessionLogout { + /** + * Whether a cached credential was removed. `false` means no matching + * session was cached; logout is idempotent. + */ + removed: boolean; +} + +/** + * Explicit OAuth session lifecycle for the persistent token cache: eager + * `login`, non-secret `status`, and local `logout`. + * + * A session is built from the same {@link 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. + * + * `login` always runs the configured interactive flow and replaces the cached + * session (the most recent login wins). `logout` removes only the local + * credential; it does not revoke anything with the provider and does not sign + * out of a browser SSO session. + * + * @example + * ```typescript + * const config: OAuthConfig = { + * issuerUrl: "https://issuer.example.com", + * clientId: "my-app", + * scopes: ["openid", "offline_access"], + * flow: OAuthFlowType.DeviceCode, + * tokenCache: { cacheDir: "/tmp/my-app/oauth-cache" }, + * }; + * const session = new OAuthSession(config); + * const status = await session.login(); + * ``` + */ +export class OAuthSession { + private readonly inner: NativeOAuthSession; + + /** Create a session manager for the given OAuth configuration. */ + constructor(config: OAuthConfig) { + this.inner = new NativeOAuthSession(config as unknown as NativeOAuthConfig); + } + + /** + * 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`. + */ + async login(): Promise { + return this.inner.login(); + } + + /** + * Report whether a matching cached session exists, with safe metadata. + * + * This never contacts the identity provider and never exposes token values. + */ + async status(): Promise { + return this.inner.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. + */ + async logout(): Promise { + return this.inner.logout(); + } } diff --git a/nodejs/src/remote.rs b/nodejs/src/remote.rs index 4bdb5685e..c619aca35 100644 --- a/nodejs/src/remote.rs +++ b/nodejs/src/remote.rs @@ -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 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, + /// How long to wait for the cross-process refresh lock before failing, + /// in seconds (default: 30). + pub lock_timeout_secs: Option, +} + +impl From 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, - /// Authentication flow: "client_credentials" or "azure_managed_identity" + /// Authentication flow: "client_credentials", "authorization_code", + /// "device_code", or "azure_managed_identity" pub flow: Option, /// Client secret (required for client_credentials). pub client_secret: Option, + /// Loopback redirect URI for authorization_code. + pub redirect_uri: Option, + /// Port for the authorization_code loopback callback server. + pub callback_port: Option, + /// Whether authorization_code uses S256 PKCE (default: true). + pub use_pkce: Option, /// Client ID for user-assigned managed identity (azure_managed_identity). pub managed_identity_client_id: Option, /// 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, + /// Opt in to the persistent token cache so short-lived processes reuse + /// one session. Only refresh tokens are persisted. + pub token_cache: Option, } impl std::fmt::Debug for OAuthConfig { @@ -181,11 +221,15 @@ impl std::fmt::Debug for OAuthConfig { "client_secret", &self.client_secret.as_deref().map(|_| ""), ) + .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 for lancedb::remote::oauth::OAuthConfig { type Error = Error; fn try_from(config: OAuthConfig) -> Result { - 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 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, + /// Flow that produced the cached session. + pub flow: String, + /// When the cached session was obtained, as Unix seconds. + pub obtained_at: Option, +} + +/// 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 { + 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 { + 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 { + 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 { + let logout = self.inner.logout().await.default_error()?; + Ok(SessionLogout { + removed: logout.removed, + }) + } +} + +impl From 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 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(\"\")")); } + + #[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 + )); + } } diff --git a/nodejs/typedoc.json b/nodejs/typedoc.json index e46085cda..c601ee61e 100644 --- a/nodejs/typedoc.json +++ b/nodejs/typedoc.json @@ -4,7 +4,8 @@ "lancedb/native.d.ts:VectorQuery", "lancedb/native.d.ts:TakeQuery", "lancedb/native.d.ts:RecordBatchIterator", - "lancedb/native.d.ts:NativeMergeInsertBuilder" + "lancedb/native.d.ts:NativeMergeInsertBuilder", + "lancedb/native.d.ts:TokenCacheOptions" ], "useHTMLEncodedBrackets": true, "useCodeBlocks": true, diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 0f7b110ac..08d81bfd3 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -267,6 +267,33 @@ class JobInfo: @property def created_at_millis(self) -> int: ... +class SessionStatus: + @property + def refreshable(self) -> bool: ... + @property + def issuer_url(self) -> str: ... + @property + def client_id(self) -> str: ... + @property + def scopes(self) -> List[str]: ... + @property + def flow(self) -> str: ... + @property + def obtained_at(self) -> Optional[int]: ... + def __repr__(self) -> str: ... + +class SessionLogout: + @property + def removed(self) -> bool: ... + def __repr__(self) -> str: ... + +class OAuthSession: + def __init__(self, config: Any) -> None: ... + async def login(self) -> SessionStatus: ... + async def status(self) -> SessionStatus: ... + async def logout(self) -> SessionLogout: ... + def __repr__(self) -> str: ... + class JobFailureInfo: @property def phase(self) -> Optional[str]: ... diff --git a/python/python/lancedb/remote/__init__.py b/python/python/lancedb/remote/__init__.py index 1f255b991..602d8dbab 100644 --- a/python/python/lancedb/remote/__init__.py +++ b/python/python/lancedb/remote/__init__.py @@ -9,7 +9,7 @@ from typing import List, Optional from lancedb import __version__ from .header import HeaderProvider -from .oauth import OAuthConfig, OAuthFlowType +from .oauth import OAuthConfig, OAuthFlowType, OAuthSession, TokenCacheOptions # The API reference renders this module with a single mkdocstrings directive, # which only picks up names listed here. New public names must be added to this @@ -22,6 +22,8 @@ __all__ = [ "HeaderProvider", "OAuthConfig", "OAuthFlowType", + "OAuthSession", + "TokenCacheOptions", ] diff --git a/python/python/lancedb/remote/oauth.py b/python/python/lancedb/remote/oauth.py index 9175c3614..ad9fd2b0e 100644 --- a/python/python/lancedb/remote/oauth.py +++ b/python/python/lancedb/remote/oauth.py @@ -12,10 +12,49 @@ class OAuthFlowType(str, Enum): CLIENT_CREDENTIALS = "client_credentials" """Client Credentials grant (service-to-service / M2M).""" + AUTHORIZATION_CODE = "authorization_code" + """Interactive Authorization Code grant, using PKCE by default.""" + + DEVICE_CODE = "device_code" + """Device Authorization grant for CLI and headless environments.""" + AZURE_MANAGED_IDENTITY = "azure_managed_identity" """Azure Managed Identity via IMDS.""" +@dataclass +class TokenCacheOptions: + """Options for the persistent OAuth token cache. + + The cache is opt-in: it is only used when set as ``token_cache`` on + :class:`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. + + Parameters + ---------- + cache_dir : Optional[str] + 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. + lock_timeout_secs : Optional[int] + How long to wait for the cross-process refresh lock before failing + (default: 30 seconds). + + Examples + -------- + >>> opts = TokenCacheOptions(cache_dir="/tmp/my-app/oauth-cache") + + Multiple identities (issuer, client, scopes, flow, client + authentication) get separate cache entries. Within one identity the most + recent login wins. + """ + + cache_dir: Optional[str] = None + lock_timeout_secs: Optional[int] = None + + @dataclass class OAuthConfig: """OAuth configuration for LanceDB authentication. @@ -38,12 +77,23 @@ class OAuthConfig: Authentication flow to use. Default: CLIENT_CREDENTIALS. client_secret : Optional[str] Client secret (required for CLIENT_CREDENTIALS). + redirect_uri : Optional[str] + Loopback redirect URI for AUTHORIZATION_CODE. The default is + ``http://127.0.0.1:{callback_port}/callback``. + callback_port : Optional[int] + Port for the AUTHORIZATION_CODE loopback callback server (default: 8400). + use_pkce : bool + Protect AUTHORIZATION_CODE with S256 PKCE (default: True). managed_identity_client_id : Optional[str] Client ID for user-assigned managed identity (AZURE_MANAGED_IDENTITY). refresh_buffer_secs : Optional[int] 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. + token_cache : Optional[TokenCacheOptions] + Opt in to the persistent token cache so short-lived processes reuse + one session. Only supported by AUTHORIZATION_CODE and DEVICE_CODE; + azure managed identity is rejected. Default: None (memory only). Examples -------- @@ -64,6 +114,29 @@ class OAuthConfig: ... scopes=["api://lancedb-api/.default"], ... flow=OAuthFlowType.AZURE_MANAGED_IDENTITY, ... ) + + Authorization Code with PKCE: + + The authorization URL is written to standard error before LanceDB tries to + open a browser, so it can be copied in headless environments. + + >>> config = OAuthConfig( + ... issuer_url="https://login.microsoftonline.com/{tenant}/v2.0", + ... client_id="app-id", + ... scopes=["openid", "api://lancedb-api/access"], + ... flow=OAuthFlowType.AUTHORIZATION_CODE, + ... ) + + Device Authorization with a persistent cache, so later processes reuse + the session without a new device prompt: + + >>> config = OAuthConfig( + ... issuer_url="https://login.microsoftonline.com/{tenant}/v2.0", + ... client_id="app-id", + ... scopes=["openid", "offline_access", "api://lancedb-api/access"], + ... flow=OAuthFlowType.DEVICE_CODE, + ... token_cache=TokenCacheOptions(), + ... ) """ issuer_url: str @@ -71,5 +144,70 @@ class OAuthConfig: scopes: List[str] flow: OAuthFlowType = OAuthFlowType.CLIENT_CREDENTIALS client_secret: Optional[str] = field(default=None, repr=False) + redirect_uri: Optional[str] = None + callback_port: Optional[int] = None + use_pkce: bool = True managed_identity_client_id: Optional[str] = None refresh_buffer_secs: Optional[int] = None + token_cache: Optional[TokenCacheOptions] = None + + +class OAuthSession: + """Explicit OAuth session lifecycle for the persistent token cache. + + Built from the same :class:`OAuthConfig` used for + :func:`lancedb.connect_async` (including its ``token_cache`` options). + A connection created with the same configuration shares the cache, so + logging in here prepares tokens for later processes without any database + request. + + ``login`` always runs the configured interactive flow and replaces the + cached session (the most recent login wins). ``logout`` removes only the + local credential; it does not revoke anything with the provider and does + not sign out of a browser SSO session. + + Examples + -------- + >>> config = OAuthConfig( + ... issuer_url="https://issuer.example.com", + ... client_id="my-app", + ... scopes=["openid", "offline_access"], + ... flow=OAuthFlowType.DEVICE_CODE, + ... token_cache=TokenCacheOptions(), + ... ) + >>> session = OAuthSession(config) # doctest: +SKIP + >>> status = await session.login() # doctest: +SKIP + >>> status.refreshable # doctest: +SKIP + True + """ + + def __init__(self, config: OAuthConfig): + from lancedb._lancedb import OAuthSession as PyOAuthSession + + self._inner: PyOAuthSession = PyOAuthSession(config) + + async def login(self): + """Eagerly run the configured flow and store the session. + + Returns a :class:`lancedb._lancedb.SessionStatus` describing the + cached 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 ``refreshable`` is ``False``. + """ + return await self._inner.login() + + async def status(self): + """Report whether a cached session exists, with safe metadata. + + Never contacts the identity provider and never exposes token values. + """ + return await self._inner.status() + + async def logout(self): + """Remove the matching local cached credential. + + Returns a :class:`lancedb._lancedb.SessionLogout` whose ``removed`` + flag reports whether a credential existed. Logout is idempotent. + """ + return await self._inner.logout() diff --git a/python/src/lib.rs b/python/src/lib.rs index 06b31d033..f83776e7b 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -47,6 +47,9 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/python/src/oauth.rs b/python/src/oauth.rs index 11ea011e2..ff24a9566 100644 --- a/python/src/oauth.rs +++ b/python/src/oauth.rs @@ -1,10 +1,33 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -use pyo3::FromPyObject; +use std::path::PathBuf; +use std::sync::Arc; +use pyo3::{FromPyObject, PyResult, Python, pyclass, pymethods}; + +use crate::error::PythonErrorExt; +use crate::runtime::future_into_py; use lancedb::error::Error; -use lancedb::remote::oauth::{OAuthConfig, OAuthFlow}; +use lancedb::remote::oauth::{AuthorizationCodeOptions, OAuthConfig, OAuthFlow}; +use lancedb::remote::{OAuthSession, SessionLogout, SessionStatus, TokenCacheOptions}; + +/// Python-side persistent token cache options, extracted via FromPyObject. +/// Maps to `lancedb.remote.oauth.TokenCacheOptions` Python dataclass. +#[derive(FromPyObject, Default)] +pub struct PyTokenCacheOptions { + pub cache_dir: Option, + pub lock_timeout_secs: Option, +} + +impl From for TokenCacheOptions { + fn from(py: PyTokenCacheOptions) -> Self { + TokenCacheOptions { + cache_dir: py.cache_dir.map(PathBuf::from), + lock_timeout_secs: py.lock_timeout_secs, + } + } +} /// Python-side OAuth configuration, extracted via FromPyObject. /// Maps to `lancedb.remote.oauth.OAuthConfig` Python dataclass. @@ -15,8 +38,12 @@ pub struct PyOAuthConfig { pub scopes: Vec, pub flow: String, pub client_secret: Option, + pub redirect_uri: Option, + pub callback_port: Option, + pub use_pkce: bool, pub managed_identity_client_id: Option, pub refresh_buffer_secs: Option, + pub token_cache: Option, } impl TryFrom for OAuthConfig { @@ -25,6 +52,17 @@ impl TryFrom for OAuthConfig { fn try_from(py: PyOAuthConfig) -> Result { let flow = match py.flow.as_str() { "client_credentials" => OAuthFlow::ClientCredentials, + "authorization_code" => { + let mut options = AuthorizationCodeOptions::new().use_pkce(py.use_pkce); + if let Some(redirect_uri) = py.redirect_uri { + options = options.redirect_uri(redirect_uri); + } + if let Some(callback_port) = py.callback_port { + options = options.callback_port(callback_port); + } + OAuthFlow::AuthorizationCode(options) + } + "device_code" => OAuthFlow::DeviceCode, "azure_managed_identity" => OAuthFlow::AzureManagedIdentity { client_id: py.managed_identity_client_id, }, @@ -42,6 +80,148 @@ impl TryFrom for OAuthConfig { scopes: py.scopes, flow, refresh_buffer_secs: py.refresh_buffer_secs, + token_cache: py.token_cache.map(TokenCacheOptions::from), + }) + } +} + +/// Wrapper around [`lancedb::remote::SessionStatus`] exposing safe metadata. +#[pyclass(name = "SessionStatus", skip_from_py_object)] +#[derive(Clone)] +pub struct PySessionStatus { + inner: SessionStatus, +} + +#[pymethods] +impl PySessionStatus { + /// Whether a cached session exists that can obtain tokens without + /// interactive authentication. + #[getter] + pub fn refreshable(&self) -> bool { + self.inner.refreshable + } + + /// Canonical issuer URL of the cached session. + #[getter] + pub fn issuer_url(&self) -> String { + self.inner.issuer_url.clone() + } + + /// Client ID of the cached session. + #[getter] + pub fn client_id(&self) -> String { + self.inner.client_id.clone() + } + + /// Canonical (sorted, de-duplicated) scopes of the cached session. + #[getter] + pub fn scopes(&self) -> Vec { + self.inner.scopes.clone() + } + + /// Flow that produced the cached session. + #[getter] + pub fn flow(&self) -> String { + self.inner.flow.clone() + } + + /// When the cached session was obtained, as Unix seconds. + #[getter] + pub fn obtained_at(&self) -> Option { + self.inner.obtained_at + } + + pub fn __repr__(&self) -> String { + format!( + "SessionStatus(refreshable={}, issuer_url='{}', client_id='{}', flow='{}')", + self.inner.refreshable, self.inner.issuer_url, self.inner.client_id, self.inner.flow + ) + } +} + +impl From for PySessionStatus { + fn from(inner: SessionStatus) -> Self { + Self { inner } + } +} + +/// Wrapper around [`lancedb::remote::SessionLogout`]. +#[pyclass(name = "SessionLogout", skip_from_py_object)] +#[derive(Clone)] +pub struct PySessionLogout { + inner: SessionLogout, +} + +#[pymethods] +impl PySessionLogout { + /// Whether a cached credential was removed. + #[getter] + pub fn removed(&self) -> bool { + self.inner.removed + } + + pub fn __repr__(&self) -> String { + format!("SessionLogout(removed={})", self.inner.removed) + } +} + +impl From for PySessionLogout { + fn from(inner: SessionLogout) -> Self { + Self { inner } + } +} + +/// Wrapper around [`lancedb::remote::OAuthSession`]. +#[pyclass(name = "OAuthSession", skip_from_py_object)] +#[derive(Clone)] +pub struct PyOAuthSession { + inner: Arc, +} + +#[pymethods] +impl PyOAuthSession { + /// Create a session manager for the given OAuth configuration. + /// + /// The configuration must set ``token_cache`` options and use a flow that + /// supports persistent sessions (authorization code or device code). + #[new] + pub fn new(config: PyOAuthConfig) -> PyResult { + let config: OAuthConfig = config.try_into().infer_error()?; + let inner = OAuthSession::new(config).infer_error()?; + Ok(Self { + inner: Arc::new(inner), + }) + } + + /// Eagerly run the configured authentication flow and store the session. + pub fn login<'py>(&self, py: Python<'py>) -> PyResult> { + let inner = Arc::clone(&self.inner); + future_into_py(py, async move { + inner.login().await.map(PySessionStatus::from).infer_error() + }) + } + + /// Report whether a matching cached session exists, with safe metadata. + pub fn status<'py>(&self, py: Python<'py>) -> PyResult> { + let inner = Arc::clone(&self.inner); + future_into_py(py, async move { + inner + .status() + .await + .map(PySessionStatus::from) + .infer_error() + }) + } + + /// Remove the matching local cached credential. + pub fn logout<'py>(&self, py: Python<'py>) -> PyResult> { + let inner = Arc::clone(&self.inner); + future_into_py(py, async move { + inner + .logout() + .await + .map(PySessionLogout::from) + .infer_error() }) } } @@ -50,16 +230,27 @@ impl TryFrom for OAuthConfig { mod tests { use super::*; - #[test] - fn test_unknown_oauth_flow_returns_invalid_input() { - let config = PyOAuthConfig { + fn base_config() -> PyOAuthConfig { + PyOAuthConfig { issuer_url: "https://issuer.example.com".to_string(), client_id: "client-id".to_string(), scopes: vec!["scope".to_string()], - flow: "typo".to_string(), + flow: "device_code".to_string(), client_secret: None, + redirect_uri: None, + callback_port: None, + use_pkce: true, managed_identity_client_id: None, refresh_buffer_secs: None, + token_cache: None, + } + } + + #[test] + fn test_unknown_oauth_flow_returns_invalid_input() { + let config = PyOAuthConfig { + flow: "typo".to_string(), + ..base_config() }; let err = OAuthConfig::try_from(config).unwrap_err(); @@ -69,4 +260,53 @@ mod tests { if message == "Unknown OAuth flow type: typo" )); } + + #[test] + fn test_authorization_code_conversion_preserves_options() { + let config = PyOAuthConfig { + flow: "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: false, + ..base_config() + }; + + let converted = OAuthConfig::try_from(config).unwrap(); + let 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 = base_config(); + let converted = OAuthConfig::try_from(config).unwrap(); + assert!(matches!(converted.flow, OAuthFlow::DeviceCode)); + } + + #[test] + fn test_token_cache_conversion() { + let config = PyOAuthConfig { + token_cache: Some(PyTokenCacheOptions { + cache_dir: Some("/tmp/oauth-cache".to_string()), + lock_timeout_secs: Some(5), + }), + ..base_config() + }; + + let converted = OAuthConfig::try_from(config).unwrap(); + let cache = converted.token_cache.expect("token cache options"); + assert_eq!( + cache.cache_dir.as_deref(), + Some(std::path::Path::new("/tmp/oauth-cache")) + ); + assert_eq!(cache.lock_timeout_secs, Some(5)); + } } diff --git a/python/tests/test_oauth.py b/python/tests/test_oauth.py index 89f5b3f8d..36f687042 100644 --- a/python/tests/test_oauth.py +++ b/python/tests/test_oauth.py @@ -1,10 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The LanceDB Authors +import asyncio import importlib.util +import json +import os +import subprocess import sys +import threading +import urllib.parse +from http.server import BaseHTTPRequestHandler, HTTPServer from pathlib import Path +import pytest + def _load_oauth_module(): oauth_path = ( @@ -31,3 +40,289 @@ def test_oauth_config_repr_redacts_client_secret(): rendered = repr(config) assert "super-secret" not in rendered assert "client_secret" not in rendered + + +def test_authorization_code_uses_pkce_by_default(): + oauth = _load_oauth_module() + + config = oauth.OAuthConfig( + issuer_url="https://issuer.example.com", + client_id="client-id", + scopes=["openid"], + flow=oauth.OAuthFlowType.AUTHORIZATION_CODE, + ) + + assert config.use_pkce is True + assert config.redirect_uri is None + assert config.callback_port is None + + +def test_device_code_flow_value(): + oauth = _load_oauth_module() + + assert oauth.OAuthFlowType.DEVICE_CODE.value == "device_code" + + +def test_token_cache_options_default_to_memory_only(): + oauth = _load_oauth_module() + + config = oauth.OAuthConfig( + issuer_url="https://issuer.example.com", + client_id="client-id", + scopes=["openid"], + ) + assert config.token_cache is None + + options = oauth.TokenCacheOptions() + assert options.cache_dir is None + assert options.lock_timeout_secs is None + + +def _remote_oauth(): + pytest.importorskip("lancedb") + from lancedb.remote import oauth as remote_oauth + + return remote_oauth + + +def _device_config(remote_oauth, issuer_url, cache_dir): + return remote_oauth.OAuthConfig( + issuer_url=issuer_url, + client_id="client-id", + scopes=["openid"], + flow=remote_oauth.OAuthFlowType.DEVICE_CODE, + token_cache=remote_oauth.TokenCacheOptions(cache_dir=str(cache_dir)), + ) + + +def test_oauth_session_status_and_logout_without_cache_entry(tmp_path): + remote_oauth = _remote_oauth() + config = _device_config(remote_oauth, "https://issuer.example.com", tmp_path) + + session = remote_oauth.OAuthSession(config) + status = asyncio.run(session.status()) + assert status.refreshable is False + assert status.issuer_url == "https://issuer.example.com" + assert status.client_id == "client-id" + assert status.scopes == ["openid"] + assert status.flow == "device_code" + assert status.obtained_at is None + + logout = asyncio.run(session.logout()) + assert logout.removed is False + + +class _MockIdpState: + def __init__(self, port): + self.port = port + self.lock = threading.Lock() + self.device_authorizations = 0 + self.refresh_grants = 0 + self.invalid_grant_rejections = 0 + self.access_tokens_issued = 0 + self.current_refresh = None + + +class _MockIdpHandler(BaseHTTPRequestHandler): + @property + def state(self) -> _MockIdpState: + return self.server.state + + def log_message(self, fmt, *args): + pass + + def _respond(self, status, payload): + body = json.dumps(payload).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self): + if self.path.startswith("/.well-known/openid-configuration"): + base = f"http://127.0.0.1:{self.state.port}" + self._respond( + 200, + { + "token_endpoint": f"{base}/token", + "device_authorization_endpoint": f"{base}/device", + }, + ) + else: + self._respond(404, {}) + + def do_POST(self): + length = int(self.headers.get("Content-Length", 0)) + body = self.rfile.read(length).decode() + params = urllib.parse.parse_qs(body) + + if self.path == "/device": + with self.state.lock: + self.state.device_authorizations += 1 + base = f"http://127.0.0.1:{self.state.port}" + self._respond( + 200, + { + "device_code": "device-code", + "user_code": "ABCD-EFGH", + "verification_uri": f"{base}/verify", + "expires_in": 60, + "interval": 1, + }, + ) + return + + if self.path == "/token": + grant_type = params.get("grant_type", [""])[0] + with self.state.lock: + if grant_type == "refresh_token": + self.state.refresh_grants += 1 + offered = params.get("refresh_token", [""])[0] + if offered != self.state.current_refresh: + self.state.invalid_grant_rejections += 1 + self._respond(400, {"error": "invalid_grant"}) + return + elif "device_code" not in grant_type: + self._respond(400, {"error": "unsupported_grant_type"}) + return + self.state.access_tokens_issued += 1 + number = self.state.access_tokens_issued + refresh = f"refresh-{number}" + self.state.current_refresh = refresh + self._respond( + 200, + { + "access_token": f"access-{number}", + "refresh_token": refresh, + "expires_in": 3600, + }, + ) + return + + self._respond(404, {}) + + +def _start_mock_idp() -> tuple[_MockIdpState, HTTPServer]: + server = HTTPServer(("127.0.0.1", 0), _MockIdpHandler) + state = _MockIdpState(server.server_address[1]) + server.state = state + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + return state, server + + +def _run_subprocess(script: Path, issuer_url: str, cache_dir: Path): + env = dict(os.environ) + env["LANCEDB_OAUTH_BROWSER"] = "/usr/bin/true" + result = subprocess.run( + [sys.executable, str(script), issuer_url, str(cache_dir)], + capture_output=True, + text=True, + timeout=120, + env=env, + ) + assert result.returncode == 0, ( + f"subprocess failed:\nstdout: {result.stdout}\nstderr: {result.stderr}" + ) + return result + + +LOGIN_SCRIPT = """ +import asyncio +import sys + +from lancedb.remote import OAuthConfig, OAuthFlowType, OAuthSession, TokenCacheOptions + +issuer_url, cache_dir = sys.argv[1], sys.argv[2] +config = OAuthConfig( + issuer_url=issuer_url, + client_id="client-id", + scopes=["openid"], + flow=OAuthFlowType.DEVICE_CODE, + token_cache=TokenCacheOptions(cache_dir=cache_dir), +) +session = OAuthSession(config) +status = asyncio.run(session.login()) +assert status.refreshable, "login must cache a refresh token" +print("LOGIN-OK") +""" + +REUSE_SCRIPT = """ +import asyncio +import sys + +import lancedb +from lancedb.remote import OAuthConfig, OAuthFlowType, OAuthSession, TokenCacheOptions + +issuer_url, cache_dir = sys.argv[1], sys.argv[2] +config = OAuthConfig( + issuer_url=issuer_url, + client_id="client-id", + scopes=["openid"], + flow=OAuthFlowType.DEVICE_CODE, + token_cache=TokenCacheOptions(cache_dir=cache_dir), +) + +session = OAuthSession(config) +status = asyncio.run(session.status()) +assert status.refreshable, "second process must see the cached session" + + +async def main(): + # Point the database endpoint at a dead port. OAuth headers are fetched + # before the request is sent, so a successful refresh proves the second + # process reused the cached session; only the database call fails. + db = await lancedb.connect_async( + "db://e2e", + host_override="http://127.0.0.1:1", + client_config={"retry_config": {"retries": 0}}, + oauth_config=config, + ) + try: + await db.table_names() + except Exception: + print("DATABASE-UNREACHABLE-AS-EXPECTED") + else: + raise AssertionError("expected the database request to fail") + + +asyncio.run(main()) +print("REUSE-OK") +""" + + +def test_cross_process_session_reuse_without_new_prompt(tmp_path): + pytest.importorskip("lancedb") + state, server = _start_mock_idp() + try: + issuer_url = f"http://127.0.0.1:{state.port}" + login_script = tmp_path / "login.py" + login_script.write_text(LOGIN_SCRIPT) + reuse_script = tmp_path / "reuse.py" + reuse_script.write_text(REUSE_SCRIPT) + cache_dir = tmp_path / "oauth-cache" + + result = _run_subprocess(login_script, issuer_url, cache_dir) + assert "LOGIN-OK" in result.stdout + assert state.device_authorizations == 1 + + result = _run_subprocess(reuse_script, issuer_url, cache_dir) + assert "REUSE-OK" in result.stdout + assert "DATABASE-UNREACHABLE-AS-EXPECTED" in result.stdout + + # The second process refreshed exactly once and never started a new + # interactive device flow. + assert state.refresh_grants == 1 + assert state.device_authorizations == 1 + assert state.invalid_grant_rejections == 0 + + logout = asyncio.run( + _remote_oauth() + .OAuthSession(_device_config(_remote_oauth(), issuer_url, cache_dir)) + .logout() + ) + assert logout.removed is True + finally: + server.shutdown() + server.server_close() diff --git a/rust/lancedb/Cargo.toml b/rust/lancedb/Cargo.toml index 9bc4c4a8c..1934ff8ce 100644 --- a/rust/lancedb/Cargo.toml +++ b/rust/lancedb/Cargo.toml @@ -53,7 +53,7 @@ metrics = { workspace = true, optional = true } metrics-util = { workspace = true, optional = true } moka = { workspace = true } pin-project = { workspace = true } -tokio = { workspace = true } +tokio = { workspace = true, features = ["io-util", "net", "time"] } log.workspace = true async-trait = { workspace = true } bytes = { workspace = true } @@ -82,6 +82,9 @@ reqwest = { version = "0.12.0", default-features = false, features = [ tonic = { workspace = true, optional = true } http = { version = "1", optional = true } # Matching what is in reqwest urlencoding = { version = "2", optional = true } +base64 = { version = "0.22", optional = true } +fs4 = { version = "0.13", optional = true } +webbrowser = { version = "1", optional = true } uuid = { workspace = true, features = ["v5"] } polars-arrow = { version = ">=0.37,<0.40.0", optional = true } polars = { version = ">=0.37,<0.40.0", optional = true } @@ -159,6 +162,9 @@ remote = [ "dep:http", "dep:tonic", "dep:urlencoding", + "dep:base64", + "dep:webbrowser", + "dep:fs4", "lance-namespace-impls/rest", "lance-namespace-impls/rest-adapter", ] diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index df7da3d62..b0446919b 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -1546,6 +1546,7 @@ mod tests { scopes: vec!["scope".to_string()], flow: crate::remote::OAuthFlow::ClientCredentials, refresh_buffer_secs: None, + token_cache: None, }; let result = ConnectBuilder::new("db://my-container/my-prefix") @@ -1588,6 +1589,7 @@ mod tests { scopes: vec!["scope".to_string()], flow: crate::remote::OAuthFlow::ClientCredentials, refresh_buffer_secs: None, + token_cache: None, }; let client_config = crate::remote::ClientConfig { header_provider: Some( diff --git a/rust/lancedb/src/remote.rs b/rust/lancedb/src/remote.rs index 6c37ec6a0..db8d98f12 100644 --- a/rust/lancedb/src/remote.rs +++ b/rust/lancedb/src/remote.rs @@ -13,6 +13,7 @@ pub mod oauth; mod retry; pub(crate) mod sql; pub(crate) mod table; +pub(crate) mod token_cache; pub(crate) mod util; const ARROW_STREAM_CONTENT_TYPE: &str = "application/vnd.apache.arrow.stream"; @@ -31,4 +32,5 @@ fn extract_job_id(body: &str) -> Option { pub use client::{ClientConfig, HeaderProvider, RetryConfig, TimeoutConfig, TlsConfig}; pub use db::{RemoteDatabaseOptions, RemoteDatabaseOptionsBuilder}; -pub use oauth::{OAuthConfig, OAuthFlow, OAuthHeaderProvider}; +pub use oauth::{AuthorizationCodeOptions, OAuthConfig, OAuthFlow, OAuthHeaderProvider}; +pub use token_cache::{OAuthSession, SessionLogout, SessionStatus, TokenCacheOptions}; diff --git a/rust/lancedb/src/remote/oauth.rs b/rust/lancedb/src/remote/oauth.rs index 3ebe8f86e..75a8878a8 100644 --- a/rust/lancedb/src/remote/oauth.rs +++ b/rust/lancedb/src/remote/oauth.rs @@ -2,24 +2,144 @@ // SPDX-FileCopyrightText: Copyright The LanceDB Authors use std::collections::HashMap; -use std::net::IpAddr; +use std::net::{IpAddr, SocketAddr}; +use std::process::Command; use std::sync::Arc; use std::time::{Duration, Instant}; use async_trait::async_trait; -use log::debug; +use base64::Engine; +use log::{debug, warn}; +use rand::Rng; use reqwest::Client; use serde::Deserialize; +use sha2::{Digest, Sha256}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; use tokio::sync::RwLock; +use tokio::time::Instant as TokioInstant; +use url::Url; use crate::error::{Error, Result}; use crate::remote::client::HeaderProvider; const DEFAULT_REFRESH_BUFFER_SECS: u64 = 300; const DEFAULT_TOKEN_TTL_SECS: u64 = 3600; +const DEFAULT_CALLBACK_PORT: u16 = 8400; +const AUTHORIZATION_CALLBACK_TIMEOUT_SECS: u64 = 300; const AZURE_IMDS_ENDPOINT: &str = "http://169.254.169.254/metadata/identity/oauth2/token"; const AZURE_IMDS_API_VERSION: &str = "2018-02-01"; +fn oauth_url_uses_secure_transport(url: &Url) -> bool { + url.scheme() == "https" + || (url.scheme() == "http" + && match url.host() { + Some(url::Host::Domain(host)) => host.eq_ignore_ascii_case("localhost"), + Some(url::Host::Ipv4(ip)) => ip.is_loopback(), + Some(url::Host::Ipv6(ip)) => ip.is_loopback(), + None => false, + }) +} + +fn validate_oauth_url(value: &str, name: &str) -> Result { + let url = Url::parse(value).map_err(|e| Error::InvalidInput { + message: format!("Invalid OAuth {name}: {e}"), + })?; + if oauth_url_uses_secure_transport(&url) { + Ok(url) + } else { + Err(Error::InvalidInput { + message: format!("OAuth {name} must use https, except for http on a loopback host"), + }) + } +} + +fn authorization_prompt(url: &Url) -> String { + format!("Open this URL to authenticate with OAuth: {url}") +} + +fn device_prompt(verification_uri: &str, user_code: &str) -> String { + format!("To authenticate with OAuth, visit {verification_uri} and enter code {user_code}") +} + +fn write_oauth_prompt(mut output: impl std::io::Write, prompt: &str) { + let _ = writeln!(output, "{prompt}"); +} + +fn show_oauth_prompt(prompt: &str) { + write_oauth_prompt(std::io::stderr().lock(), prompt); +} + +/// Options for the interactive OAuth Authorization Code flow. +/// +/// The built-in callback server accepts only loopback HTTP redirect URIs. PKCE +/// with the S256 challenge method is enabled by default and should be disabled +/// only for providers that do not support it. +/// +/// # Example +/// +/// ``` +/// use lancedb::remote::{AuthorizationCodeOptions, OAuthFlow}; +/// +/// let flow = OAuthFlow::AuthorizationCode( +/// AuthorizationCodeOptions::new() +/// .callback_port(8400) +/// .use_pkce(true), +/// ); +/// ``` +#[derive(Debug, Clone)] +#[non_exhaustive] +pub struct AuthorizationCodeOptions { + /// Redirect URI registered with the identity provider. + /// + /// Defaults to `http://127.0.0.1:{callback_port}/callback`. + pub redirect_uri: Option, + + /// Port for the built-in loopback callback server. + /// + /// Defaults to 8400. When `redirect_uri` contains an explicit port, this + /// option must either be omitted or match that port. + pub callback_port: Option, + + /// Whether to protect the authorization code exchange with S256 PKCE. + pub use_pkce: bool, +} + +impl Default for AuthorizationCodeOptions { + fn default() -> Self { + Self { + redirect_uri: None, + callback_port: None, + use_pkce: true, + } + } +} + +impl AuthorizationCodeOptions { + /// Create authorization-code options with S256 PKCE enabled. + pub fn new() -> Self { + Self::default() + } + + /// Set the loopback redirect URI registered with the identity provider. + pub fn redirect_uri(mut self, redirect_uri: impl Into) -> Self { + self.redirect_uri = Some(redirect_uri.into()); + self + } + + /// Set the port for the built-in loopback callback server. + pub fn callback_port(mut self, callback_port: u16) -> Self { + self.callback_port = Some(callback_port); + self + } + + /// Enable or disable S256 PKCE. + pub fn use_pkce(mut self, use_pkce: bool) -> Self { + self.use_pkce = use_pkce; + self + } +} + /// OAuth authentication flow configuration. #[derive(Debug, Clone)] pub enum OAuthFlow { @@ -27,6 +147,15 @@ pub enum OAuthFlow { /// Requires `client_secret` in [`OAuthConfig`]. ClientCredentials, + /// Authorization Code grant using an interactive browser and a built-in + /// loopback callback server. The authorization URL is also written to + /// stderr so it remains available when the browser cannot be opened. + AuthorizationCode(AuthorizationCodeOptions), + + /// Device Authorization grant for CLI and headless environments. The + /// verification URI and user code are written to stderr. + DeviceCode, + /// Azure Managed Identity via IMDS. /// Works on Azure VMs, AKS, App Service, and Azure Functions. /// IMDS requests bypass proxy settings because the endpoint is link-local. @@ -65,6 +194,14 @@ pub struct OAuthConfig { /// 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, + + /// Opt in to the persistent token cache so short-lived processes can + /// reuse an authenticated session instead of re-prompting. + /// + /// When unset (the default), tokens stay in process memory only. Only + /// refresh tokens are persisted; see + /// [`TokenCacheOptions`](crate::remote::TokenCacheOptions). + pub token_cache: Option, } impl std::fmt::Debug for OAuthConfig { @@ -79,6 +216,7 @@ impl std::fmt::Debug for OAuthConfig { .field("scopes", &self.scopes) .field("flow", &self.flow) .field("refresh_buffer_secs", &self.refresh_buffer_secs) + .field("token_cache", &self.token_cache) .finish() } } @@ -88,26 +226,34 @@ impl std::fmt::Debug for OAuthConfig { #[derive(Clone, Debug, Deserialize)] struct OidcDiscovery { token_endpoint: String, + authorization_endpoint: Option, + device_authorization_endpoint: Option, } // -- Token Response -- #[derive(Deserialize)] -struct TokenResponse { - access_token: String, +pub(crate) struct TokenResponse { + pub(crate) access_token: String, + #[serde(default)] + pub(crate) refresh_token: Option, /// Token lifetime in seconds. /// Some providers (Azure IMDS) return this as a string, so we accept both. #[serde(default, deserialize_with = "deserialize_optional_u64_or_string")] - expires_in: Option, + pub(crate) expires_in: Option, #[serde(default)] #[allow(dead_code)] - token_type: Option, + pub(crate) token_type: Option, } impl std::fmt::Debug for TokenResponse { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("TokenResponse") .field("access_token", &"") + .field( + "refresh_token", + &self.refresh_token.as_ref().map(|_| ""), + ) .field("expires_in", &self.expires_in) .field("token_type", &self.token_type) .finish() @@ -168,6 +314,7 @@ where struct TokenState { access_token: Option, + refresh_token: Option, expires_at: Option, } @@ -175,6 +322,7 @@ impl TokenState { fn new() -> Self { Self { access_token: None, + refresh_token: None, expires_at: None, } } @@ -189,50 +337,72 @@ impl TokenState { fn update(&mut self, resp: &TokenResponse) { self.access_token = Some(resp.access_token.clone()); + if resp.refresh_token.is_some() { + self.refresh_token = resp.refresh_token.clone(); + } let expires_in = resp.expires_in.unwrap_or(DEFAULT_TOKEN_TTL_SECS); self.expires_at = Some(Instant::now() + Duration::from_secs(expires_in)); } } #[async_trait] -trait TokenSource: Send + Sync + std::fmt::Debug { +pub(crate) trait TokenSource: Send + Sync + std::fmt::Debug { async fn fetch_token(&self) -> Result; + + async fn refresh_token(&self, _refresh_token: &str) -> Result { + Ok(RefreshResult::Unsupported) + } } -struct ClientCredentialsSource { +#[derive(Debug)] +pub(crate) enum RefreshResult { + Refreshed(TokenResponse), + Reauthenticate, + Unsupported, +} + +struct OidcClient { issuer_url: String, client_id: String, - client_secret: String, + client_secret: Option, scopes: Vec, http_client: Client, discovery: RwLock>, } -impl std::fmt::Debug for ClientCredentialsSource { +impl std::fmt::Debug for OidcClient { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("ClientCredentialsSource") + f.debug_struct("OidcClient") .field("issuer_url", &self.issuer_url) .field("client_id", &self.client_id) - .field("client_secret", &"") + .field( + "client_secret", + &self.client_secret.as_ref().map(|_| ""), + ) .field("scopes", &self.scopes) .finish() } } -impl ClientCredentialsSource { +impl OidcClient { fn new( issuer_url: String, client_id: String, client_secret: Option, scopes: Vec, ) -> Result { - let client_secret = client_secret.ok_or(Error::InvalidInput { - message: "client_secret is required for ClientCredentials flow".to_string(), - })?; Self::validate_issuer_transport(&issuer_url)?; let http_client = Client::builder() .timeout(Duration::from_secs(30)) + .redirect(reqwest::redirect::Policy::custom(|attempt| { + if oauth_url_uses_secure_transport(attempt.url()) { + attempt.follow() + } else { + attempt + .error("OAuth redirects must use https, except for http on a loopback host") + } + })) .build() .map_err(|e| Error::Runtime { message: format!("Failed to create HTTP client for OAuth: {e}"), @@ -249,31 +419,7 @@ impl ClientCredentialsSource { } fn validate_issuer_transport(issuer_url: &str) -> Result<()> { - let issuer = url::Url::parse(issuer_url).map_err(|e| Error::InvalidInput { - message: format!("Invalid OAuth issuer_url: {e}"), - })?; - - match issuer.scheme() { - "https" => Ok(()), - "http" if Self::is_loopback_issuer(&issuer) => Ok(()), - _ => Err(Error::InvalidInput { - message: - "ClientCredentials OAuth issuer_url must use https, except for loopback hosts" - .to_string(), - }), - } - } - - fn is_loopback_issuer(issuer: &url::Url) -> bool { - let Some(host) = issuer.host_str() else { - return false; - }; - - host.eq_ignore_ascii_case("localhost") - || host - .parse::() - .map(|addr| addr.is_loopback()) - .unwrap_or(false) + validate_oauth_url(issuer_url, "issuer_url").map(drop) } async fn get_discovery(&self) -> Result { @@ -319,6 +465,13 @@ impl ClientCredentialsSource { let disc: OidcDiscovery = resp.json().await.map_err(|e| Error::Runtime { message: format!("Failed to parse OIDC discovery document: {e}"), })?; + validate_oauth_url(&disc.token_endpoint, "token_endpoint")?; + if let Some(endpoint) = disc.authorization_endpoint.as_deref() { + validate_oauth_url(endpoint, "authorization_endpoint")?; + } + if let Some(endpoint) = disc.device_authorization_endpoint.as_deref() { + validate_oauth_url(endpoint, "device_authorization_endpoint")?; + } let result = disc.clone(); @@ -337,7 +490,7 @@ impl ClientCredentialsSource { async fn post_token_request( &self, endpoint: &str, - params: &[(&str, &str)], + params: &[(String, String)], ) -> Result { let resp = self .http_client @@ -363,24 +516,690 @@ impl ClientCredentialsSource { message: format!("Failed to parse token response: {e}"), }) } + + async fn refresh_token(&self, refresh_token: &str) -> Result { + let endpoint = self.get_token_endpoint().await?; + let mut params = vec![ + ("grant_type".to_string(), "refresh_token".to_string()), + ("client_id".to_string(), self.client_id.clone()), + ("refresh_token".to_string(), refresh_token.to_string()), + ]; + if let Some(secret) = self.client_secret.as_ref() { + params.push(("client_secret".to_string(), secret.clone())); + } + let response = self + .http_client + .post(&endpoint) + .form(¶ms) + .send() + .await + .map_err(|e| Error::Runtime { + message: format!("Refresh token request to {endpoint} failed: {e}"), + })?; + if response.status().is_success() { + return response + .json() + .await + .map(RefreshResult::Refreshed) + .map_err(|e| Error::Runtime { + message: format!("Failed to parse refresh token response: {e}"), + }); + } + + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + let error_code = serde_json::from_str::(&body) + .ok() + .map(|error| error.error); + if matches!( + error_code.as_deref(), + Some("invalid_grant" | "invalid_token") + ) { + return Ok(RefreshResult::Reauthenticate); + } + Err(Error::Runtime { + message: format!("Refresh token request failed with status {status}: {body}"), + }) + } +} + +struct ClientCredentialsSource { + oidc: OidcClient, +} + +impl std::fmt::Debug for ClientCredentialsSource { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ClientCredentialsSource") + .field("oidc", &self.oidc) + .finish() + } +} + +impl ClientCredentialsSource { + fn new( + issuer_url: String, + client_id: String, + client_secret: Option, + scopes: Vec, + ) -> Result { + if client_secret.is_none() { + return Err(Error::InvalidInput { + message: "client_secret is required for ClientCredentials flow".to_string(), + }); + } + Ok(Self { + oidc: OidcClient::new(issuer_url, client_id, client_secret, scopes)?, + }) + } } #[async_trait] impl TokenSource for ClientCredentialsSource { async fn fetch_token(&self) -> Result { - let token_endpoint = self.get_token_endpoint().await?; - let scope = self.scopes_string(); + let token_endpoint = self.oidc.get_token_endpoint().await?; let params = [ - ("grant_type", "client_credentials"), - ("client_id", self.client_id.as_str()), - ("client_secret", self.client_secret.as_str()), - ("scope", scope.as_str()), + ("grant_type".to_string(), "client_credentials".to_string()), + ("client_id".to_string(), self.oidc.client_id.clone()), + ( + "client_secret".to_string(), + self.oidc.client_secret.clone().expect("validated in new"), + ), + ("scope".to_string(), self.oidc.scopes_string()), ]; - self.post_token_request(&token_endpoint, ¶ms).await + self.oidc.post_token_request(&token_endpoint, ¶ms).await } } +#[derive(Debug)] +struct ResolvedRedirect { + uri: String, + bind_addr: SocketAddr, + callback_path: String, +} + +impl ResolvedRedirect { + fn new(options: &AuthorizationCodeOptions) -> Result { + let uri = options.redirect_uri.clone().unwrap_or_else(|| { + format!( + "http://127.0.0.1:{}/callback", + options.callback_port.unwrap_or(DEFAULT_CALLBACK_PORT) + ) + }); + let parsed = Url::parse(&uri).map_err(|e| Error::InvalidInput { + message: format!("Invalid OAuth redirect_uri: {e}"), + })?; + + if parsed.scheme() != "http" { + return Err(Error::InvalidInput { + message: "OAuth redirect_uri must use http with a loopback host".to_string(), + }); + } + if parsed.query().is_some() || parsed.fragment().is_some() { + return Err(Error::InvalidInput { + message: "OAuth redirect_uri must not contain a query or fragment".to_string(), + }); + } + + let ip = match parsed.host() { + Some(url::Host::Domain(host)) if host.eq_ignore_ascii_case("localhost") => { + IpAddr::V4(std::net::Ipv4Addr::LOCALHOST) + } + Some(url::Host::Ipv4(ip)) if ip.is_loopback() => IpAddr::V4(ip), + Some(url::Host::Ipv6(ip)) if ip.is_loopback() => IpAddr::V6(ip), + Some(_) => { + return Err(Error::InvalidInput { + message: "OAuth redirect_uri must use a loopback host".to_string(), + }); + } + None => { + return Err(Error::InvalidInput { + message: "OAuth redirect_uri must include a loopback host".to_string(), + }); + } + }; + let port = parsed.port().ok_or(Error::InvalidInput { + message: "OAuth redirect_uri must include a port".to_string(), + })?; + if port == 0 { + return Err(Error::InvalidInput { + message: "OAuth redirect_uri port must be greater than zero".to_string(), + }); + } + if let Some(callback_port) = options.callback_port + && callback_port != port + { + return Err(Error::InvalidInput { + message: format!( + "OAuth callback_port {callback_port} does not match redirect_uri port {port}" + ), + }); + } + + Ok(Self { + uri, + bind_addr: SocketAddr::new(ip, port), + callback_path: parsed.path().to_string(), + }) + } +} + +#[derive(Debug)] +struct AuthorizationRequest { + url: Url, + state: String, + code_verifier: Option, +} + +#[derive(Debug, PartialEq)] +enum AuthorizationCallback { + Code(String), + ProviderError(String), +} + +struct AuthorizationCodeSource { + oidc: OidcClient, + options: AuthorizationCodeOptions, + redirect: ResolvedRedirect, +} + +impl std::fmt::Debug for AuthorizationCodeSource { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AuthorizationCodeSource") + .field("oidc", &self.oidc) + .field("options", &self.options) + .field("redirect", &self.redirect) + .finish() + } +} + +impl AuthorizationCodeSource { + fn new( + issuer_url: String, + client_id: String, + client_secret: Option, + scopes: Vec, + options: AuthorizationCodeOptions, + ) -> Result { + let redirect = ResolvedRedirect::new(&options)?; + Ok(Self { + oidc: OidcClient::new(issuer_url, client_id, client_secret, scopes)?, + options, + redirect, + }) + } + + async fn build_authorization_request(&self) -> Result { + let endpoint = self + .oidc + .get_discovery() + .await? + .authorization_endpoint + .ok_or(Error::Runtime { + message: "OIDC discovery did not provide authorization_endpoint".to_string(), + })?; + let mut url = validate_oauth_url(&endpoint, "authorization_endpoint")?; + let state = random_urlsafe_string(32); + let code_verifier = self.options.use_pkce.then(|| random_urlsafe_string(64)); + + { + let mut query = url.query_pairs_mut(); + query + .append_pair("response_type", "code") + .append_pair("client_id", &self.oidc.client_id) + .append_pair("redirect_uri", &self.redirect.uri) + .append_pair("scope", &self.oidc.scopes_string()) + .append_pair("state", &state); + if let Some(verifier) = code_verifier.as_ref() { + let challenge = base64::engine::general_purpose::URL_SAFE_NO_PAD + .encode(Sha256::digest(verifier.as_bytes())); + query + .append_pair("code_challenge", &challenge) + .append_pair("code_challenge_method", "S256"); + } + } + + Ok(AuthorizationRequest { + url, + state, + code_verifier, + }) + } + + async fn wait_for_callback( + &self, + listener: &TcpListener, + expected_state: &str, + ) -> Result { + let deadline = + TokioInstant::now() + Duration::from_secs(AUTHORIZATION_CALLBACK_TIMEOUT_SECS); + loop { + let (mut stream, _) = tokio::time::timeout_at(deadline, listener.accept()) + .await + .map_err(|_| Error::Runtime { + message: "Timed out waiting for the OAuth authorization callback".to_string(), + })? + .map_err(|e| Error::Runtime { + message: format!("Failed to accept OAuth callback connection: {e}"), + })?; + match read_authorization_callback( + &mut stream, + &self.redirect.callback_path, + expected_state, + deadline, + ) + .await + { + Ok(AuthorizationCallback::Code(code)) => { + write_callback_response(&mut stream, true).await; + return Ok(code); + } + Ok(AuthorizationCallback::ProviderError(message)) => { + write_callback_response(&mut stream, false).await; + return Err(Error::Runtime { message }); + } + Err(error) => { + if TokioInstant::now() >= deadline { + return Err(Error::Runtime { + message: "Timed out waiting for the OAuth authorization callback" + .to_string(), + }); + } + debug!("Ignoring unrelated OAuth callback connection: {error}"); + write_callback_response(&mut stream, false).await; + } + } + } + } + + async fn exchange_code( + &self, + code: &str, + code_verifier: Option<&str>, + ) -> Result { + let endpoint = self.oidc.get_token_endpoint().await?; + let mut params = vec![ + ("grant_type".to_string(), "authorization_code".to_string()), + ("client_id".to_string(), self.oidc.client_id.clone()), + ("code".to_string(), code.to_string()), + ("redirect_uri".to_string(), self.redirect.uri.clone()), + ]; + if let Some(verifier) = code_verifier { + params.push(("code_verifier".to_string(), verifier.to_string())); + } + if let Some(secret) = self.oidc.client_secret.as_ref() { + params.push(("client_secret".to_string(), secret.clone())); + } + self.oidc.post_token_request(&endpoint, ¶ms).await + } +} + +#[async_trait] +impl TokenSource for AuthorizationCodeSource { + async fn fetch_token(&self) -> Result { + let listener = TcpListener::bind(self.redirect.bind_addr) + .await + .map_err(|e| Error::Runtime { + message: format!( + "Failed to bind OAuth callback server at {}: {e}", + self.redirect.bind_addr + ), + })?; + let request = self.build_authorization_request().await?; + show_oauth_prompt(&authorization_prompt(&request.url)); + launch_browser(request.url.clone()); + let code = self.wait_for_callback(&listener, &request.state).await?; + self.exchange_code(&code, request.code_verifier.as_deref()) + .await + } + + async fn refresh_token(&self, refresh_token: &str) -> Result { + self.oidc.refresh_token(refresh_token).await + } +} + +#[derive(Deserialize)] +struct DeviceAuthorizationResponse { + device_code: String, + user_code: String, + verification_uri: String, + #[serde(default)] + verification_uri_complete: Option, + expires_in: u64, + #[serde(default)] + interval: Option, +} + +#[derive(Debug, Deserialize)] +struct OAuthErrorResponse { + error: String, + #[serde(default)] + error_description: Option, +} + +struct DeviceCodeSource { + oidc: OidcClient, +} + +impl std::fmt::Debug for DeviceCodeSource { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DeviceCodeSource") + .field("oidc", &self.oidc) + .finish() + } +} + +impl DeviceCodeSource { + fn new( + issuer_url: String, + client_id: String, + client_secret: Option, + scopes: Vec, + ) -> Result { + Ok(Self { + oidc: OidcClient::new(issuer_url, client_id, client_secret, scopes)?, + }) + } + + async fn request_device_authorization(&self) -> Result { + let endpoint = self + .oidc + .get_discovery() + .await? + .device_authorization_endpoint + .ok_or(Error::Runtime { + message: "OIDC discovery did not provide device_authorization_endpoint".to_string(), + })?; + let mut params = vec![ + ("client_id".to_string(), self.oidc.client_id.clone()), + ("scope".to_string(), self.oidc.scopes_string()), + ]; + if let Some(secret) = self.oidc.client_secret.as_ref() { + params.push(("client_secret".to_string(), secret.clone())); + } + let response = self + .oidc + .http_client + .post(&endpoint) + .form(¶ms) + .send() + .await + .map_err(|e| Error::Runtime { + message: format!("Device authorization request to {endpoint} failed: {e}"), + })?; + if !response.status().is_success() { + return Err(Error::Runtime { + message: format!( + "Device authorization request failed with status {}: {}", + response.status(), + response.text().await.unwrap_or_default() + ), + }); + } + let device: DeviceAuthorizationResponse = + response.json().await.map_err(|e| Error::Runtime { + message: format!("Failed to parse device authorization response: {e}"), + })?; + validate_oauth_url(&device.verification_uri, "verification_uri")?; + if let Some(uri) = device.verification_uri_complete.as_deref() { + validate_oauth_url(uri, "verification_uri_complete")?; + } + Ok(device) + } + + async fn poll_for_token(&self, device: &DeviceAuthorizationResponse) -> Result { + let endpoint = self.oidc.get_token_endpoint().await?; + let deadline = TokioInstant::now() + Duration::from_secs(device.expires_in); + let mut interval = Duration::from_secs(device.interval.unwrap_or(5).max(1)); + + loop { + let now = TokioInstant::now(); + if now >= deadline { + return Err(Error::Runtime { + message: "Device authorization expired before authentication completed" + .to_string(), + }); + } + tokio::time::sleep_until(std::cmp::min(now + interval, deadline)).await; + if TokioInstant::now() >= deadline { + return Err(Error::Runtime { + message: "Device authorization expired before authentication completed" + .to_string(), + }); + } + + let mut params = vec![ + ( + "grant_type".to_string(), + "urn:ietf:params:oauth:grant-type:device_code".to_string(), + ), + ("client_id".to_string(), self.oidc.client_id.clone()), + ("device_code".to_string(), device.device_code.clone()), + ]; + if let Some(secret) = self.oidc.client_secret.as_ref() { + params.push(("client_secret".to_string(), secret.clone())); + } + + let response = match self + .oidc + .http_client + .post(&endpoint) + .form(¶ms) + .send() + .await + { + Ok(response) => response, + Err(error) => { + warn!("Device token request to {endpoint} failed; retrying: {error}"); + continue; + } + }; + if response.status().is_success() { + return response.json().await.map_err(|e| Error::Runtime { + message: format!("Failed to parse device token response: {e}"), + }); + } + + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + let oauth_error = serde_json::from_str::(&body).ok(); + match oauth_error.as_ref().map(|error| error.error.as_str()) { + Some("authorization_pending") => continue, + Some("slow_down") => { + interval += Duration::from_secs(5); + continue; + } + Some("temporarily_unavailable") => continue, + Some("access_denied") => { + return Err(Error::Runtime { + message: "Device authorization was denied by the user".to_string(), + }); + } + Some("expired_token") => { + return Err(Error::Runtime { + message: "Device authorization expired before authentication completed" + .to_string(), + }); + } + _ if status == reqwest::StatusCode::TOO_MANY_REQUESTS + || status.is_server_error() => + { + warn!("Device token endpoint returned {status}; retrying"); + continue; + } + _ => { + let detail = oauth_error + .and_then(|error| error.error_description) + .unwrap_or(body); + return Err(Error::Runtime { + message: format!( + "Device token request failed with status {status}: {detail}" + ), + }); + } + } + } + } +} + +#[async_trait] +impl TokenSource for DeviceCodeSource { + async fn fetch_token(&self) -> Result { + let device = self.request_device_authorization().await?; + show_oauth_prompt(&device_prompt(&device.verification_uri, &device.user_code)); + let (browser_url, name) = device + .verification_uri_complete + .as_deref() + .map(|url| (url, "verification_uri_complete")) + .unwrap_or((&device.verification_uri, "verification_uri")); + launch_browser(validate_oauth_url(browser_url, name)?); + self.poll_for_token(&device).await + } + + async fn refresh_token(&self, refresh_token: &str) -> Result { + self.oidc.refresh_token(refresh_token).await + } +} + +fn random_urlsafe_string(length: usize) -> String { + const CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-._~"; + let mut rng = rand::rng(); + (0..length) + .map(|_| CHARSET[rng.random_range(0..CHARSET.len())] as char) + .collect() +} + +fn launch_browser(url: Url) { + drop(tokio::task::spawn_blocking(move || { + if let Some(browser) = std::env::var_os("LANCEDB_OAUTH_BROWSER") { + match Command::new(browser).arg(url.as_str()).status() { + Ok(status) if !status.success() => { + warn!("OAuth browser helper exited with status {status}"); + } + Err(error) => warn!("Could not run the OAuth browser helper: {error}"), + Ok(_) => {} + } + } else if let Err(error) = webbrowser::open(url.as_str()) { + warn!("Could not open an OAuth browser automatically: {error}"); + } + })); +} + +async fn read_authorization_callback( + stream: &mut TcpStream, + expected_path: &str, + expected_state: &str, + overall_deadline: TokioInstant, +) -> Result { + const MAX_CALLBACK_REQUEST_BYTES: usize = 16 * 1024; + let deadline = std::cmp::min( + overall_deadline, + TokioInstant::now() + Duration::from_secs(10), + ); + let mut request = Vec::with_capacity(1024); + loop { + let mut buffer = [0; 1024]; + let count = tokio::time::timeout_at(deadline, stream.read(&mut buffer)) + .await + .map_err(|_| Error::Runtime { + message: "Timed out reading the OAuth authorization callback".to_string(), + })? + .map_err(|e| Error::Runtime { + message: format!("Failed to read OAuth authorization callback: {e}"), + })?; + if count == 0 { + return Err(Error::Runtime { + message: "OAuth authorization callback closed before sending a request".to_string(), + }); + } + request.extend_from_slice(&buffer[..count]); + if request.windows(4).any(|window| window == b"\r\n\r\n") { + break; + } + if request.len() >= MAX_CALLBACK_REQUEST_BYTES { + return Err(Error::Runtime { + message: "OAuth authorization callback request was too large".to_string(), + }); + } + } + let request = std::str::from_utf8(&request).map_err(|e| Error::Runtime { + message: format!("OAuth authorization callback was not valid UTF-8: {e}"), + })?; + parse_authorization_callback(request, expected_path, expected_state) +} + +fn parse_authorization_callback( + request: &str, + expected_path: &str, + expected_state: &str, +) -> Result { + let request_target = request + .lines() + .next() + .and_then(|line| { + let mut parts = line.split_whitespace(); + (parts.next() == Some("GET")) + .then(|| parts.next()) + .flatten() + }) + .ok_or(Error::Runtime { + message: "OAuth authorization callback was not a valid HTTP GET request".to_string(), + })?; + let callback = + Url::parse(&format!("http://loopback{request_target}")).map_err(|e| Error::Runtime { + message: format!("OAuth authorization callback URL was invalid: {e}"), + })?; + if callback.path() != expected_path { + return Err(Error::Runtime { + message: format!( + "OAuth authorization callback used unexpected path {}", + callback.path() + ), + }); + } + let params: HashMap<_, _> = callback.query_pairs().into_owned().collect(); + if params.get("state").map(String::as_str) != Some(expected_state) { + return Err(Error::Runtime { + message: "OAuth authorization callback state did not match".to_string(), + }); + } + if let Some(error) = params.get("error") { + let description = params + .get("error_description") + .map(String::as_str) + .unwrap_or(error); + return Ok(AuthorizationCallback::ProviderError(format!( + "OAuth authorization failed: {description}" + ))); + } + params + .get("code") + .cloned() + .map(AuthorizationCallback::Code) + .ok_or(Error::Runtime { + message: "OAuth authorization callback did not contain a code".to_string(), + }) +} + +async fn write_callback_response(stream: &mut TcpStream, success: bool) { + let (status, body) = if success { + ( + "200 OK", + "

Authentication successful

You can close this window.

", + ) + } else { + ( + "400 Bad Request", + "

Authentication failed

Return to the application for details.

", + ) + }; + let response = format!( + "HTTP/1.1 {status}\r\nContent-Type: text/html; charset=utf-8\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ); + let _ = stream.write_all(response.as_bytes()).await; +} + struct AzureImdsSource { client_id: Option, resource: String, @@ -463,22 +1282,66 @@ impl TokenSource for AzureImdsSource { } } +/// Build the token source for a configuration. +/// +/// Shared by [`OAuthHeaderProvider`] and +/// [`OAuthSession`](crate::remote::OAuthSession). +pub(crate) fn build_token_source(config: &OAuthConfig) -> Result> { + if config.scopes.is_empty() { + return Err(Error::InvalidInput { + message: "At least one OAuth scope is required".to_string(), + }); + } + Ok(match &config.flow { + OAuthFlow::ClientCredentials => Box::new(ClientCredentialsSource::new( + config.issuer_url.clone(), + config.client_id.clone(), + config.client_secret.clone(), + config.scopes.clone(), + )?), + OAuthFlow::AuthorizationCode(options) => Box::new(AuthorizationCodeSource::new( + config.issuer_url.clone(), + config.client_id.clone(), + config.client_secret.clone(), + config.scopes.clone(), + options.clone(), + )?), + OAuthFlow::DeviceCode => Box::new(DeviceCodeSource::new( + config.issuer_url.clone(), + config.client_id.clone(), + config.client_secret.clone(), + config.scopes.clone(), + )?), + OAuthFlow::AzureManagedIdentity { client_id } => Box::new(AzureImdsSource::new( + config.scopes.clone(), + client_id.clone(), + )?), + }) +} + /// OAuth header provider that manages the full token lifecycle. /// /// Implements [`HeaderProvider`] to inject `Authorization: Bearer ` /// headers into every LanceDB request, with automatic token refresh. It also /// identifies the bearer credential as OIDC so LanceDB's SQL service selects /// OIDC validation instead of API-key validation. +/// +/// When the configuration enables +/// [`token_cache`](OAuthConfig::token_cache), tokens are additionally shared +/// through a hardened on-disk cache so separate processes reuse one session; +/// see [`crate::remote::token_cache`]. pub struct OAuthHeaderProvider { token_source: Box, token_state: Arc>, refresh_buffer: Duration, + token_cache: Option>, } impl std::fmt::Debug for OAuthHeaderProvider { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("OAuthHeaderProvider") .field("token_source", &self.token_source) + .field("token_cache", &self.token_cache) .finish() } } @@ -486,39 +1349,19 @@ impl std::fmt::Debug for OAuthHeaderProvider { impl OAuthHeaderProvider { /// Create a new OAuth header provider from configuration. pub fn new(config: OAuthConfig) -> Result { - let OAuthConfig { - issuer_url, - client_id, - client_secret, - scopes, - flow, - refresh_buffer_secs, - } = config; - - if scopes.is_empty() { - return Err(Error::InvalidInput { - message: "At least one OAuth scope is required".to_string(), - }); - } - - let refresh_buffer = - Duration::from_secs(refresh_buffer_secs.unwrap_or(DEFAULT_REFRESH_BUFFER_SECS)); - let token_source: Box = match flow { - OAuthFlow::ClientCredentials => Box::new(ClientCredentialsSource::new( - issuer_url, - client_id, - client_secret, - scopes, - )?), - OAuthFlow::AzureManagedIdentity { client_id } => { - Box::new(AzureImdsSource::new(scopes, client_id)?) - } - }; + let refresh_buffer = Duration::from_secs( + config + .refresh_buffer_secs + .unwrap_or(DEFAULT_REFRESH_BUFFER_SECS), + ); + let token_source = build_token_source(&config)?; + let token_cache = crate::remote::token_cache::token_cache_for_config(&config)?; Ok(Self { token_source, token_state: Arc::new(RwLock::new(TokenState::new())), refresh_buffer, + token_cache, }) } @@ -544,8 +1387,34 @@ impl OAuthHeaderProvider { return Ok(token.clone()); } - debug!("Acquiring new OAuth token via {:?}", self.token_source); - let resp = self.token_source.fetch_token().await?; + if let Some(cache) = &self.token_cache { + // Cross-process critical section: serialize with other processes, + // reread the durable record, refresh or acquire exactly once, and + // persist the rotated refresh token. + let resp = cache.refresh_or_acquire(self.token_source.as_ref()).await?; + state.update(&resp); + return Ok(resp.access_token); + } + + let refresh_token = state.refresh_token.clone(); + let resp = if let Some(refresh_token) = refresh_token.as_deref() { + debug!("Refreshing OAuth access token via {:?}", self.token_source); + match self.token_source.refresh_token(refresh_token).await? { + RefreshResult::Refreshed(response) => response, + RefreshResult::Unsupported => self.token_source.fetch_token().await?, + RefreshResult::Reauthenticate => { + warn!( + "OAuth refresh token was rejected; acquiring a new token via {:?}", + self.token_source + ); + state.refresh_token = None; + self.token_source.fetch_token().await? + } + } + } else { + debug!("Acquiring new OAuth token via {:?}", self.token_source); + self.token_source.fetch_token().await? + }; state.update(&resp); Ok(resp.access_token) @@ -591,6 +1460,7 @@ mod tests { let mut state = TokenState::new(); let response = TokenResponse { access_token: "tok".to_string(), + refresh_token: None, expires_in: None, token_type: None, }; @@ -601,6 +1471,25 @@ mod tests { assert!(state.is_expired(Duration::from_secs(DEFAULT_TOKEN_TTL_SECS + 1))); } + #[test] + fn test_token_state_retains_refresh_token_when_not_rotated() { + let mut state = TokenState::new(); + state.update(&TokenResponse { + access_token: "token-1".to_string(), + refresh_token: Some("refresh-1".to_string()), + expires_in: Some(60), + token_type: None, + }); + state.update(&TokenResponse { + access_token: "token-2".to_string(), + refresh_token: None, + expires_in: Some(60), + token_type: None, + }); + + assert_eq!(state.refresh_token.as_deref(), Some("refresh-1")); + } + #[test] fn test_token_response_accepts_float_expires_in() { let response: TokenResponse = @@ -622,12 +1511,14 @@ mod tests { fn test_token_response_debug_redacts_access_token() { let response = TokenResponse { access_token: "secret-token".to_string(), + refresh_token: Some("secret-refresh-token".to_string()), expires_in: Some(3600), token_type: Some("Bearer".to_string()), }; let debug = format!("{response:?}"); assert!(!debug.contains("secret-token")); + assert!(!debug.contains("secret-refresh-token")); assert!(debug.contains("access_token: \"\"")); } @@ -641,7 +1532,655 @@ mod tests { ) .unwrap(); - assert_eq!(source.scopes_string(), "scope1 scope2"); + assert_eq!(source.oidc.scopes_string(), "scope1 scope2"); + } + + #[test] + fn test_oauth_transport_requires_https_except_for_loopback() { + assert!(validate_oauth_url("https://idp.example.com/token", "endpoint").is_ok()); + assert!(validate_oauth_url("http://localhost:8080/token", "endpoint").is_ok()); + assert!(validate_oauth_url("http://127.0.0.1:8080/token", "endpoint").is_ok()); + assert!(validate_oauth_url("http://[::1]:8080/token", "endpoint").is_ok()); + + let err = validate_oauth_url("http://idp.example.com/token", "endpoint").unwrap_err(); + assert!(matches!( + err, + Error::InvalidInput { message } + if message == "OAuth endpoint must use https, except for http on a loopback host" + )); + } + + #[test] + fn test_interactive_prompts_use_default_visible_output() { + let authorization_url = Url::parse("https://idp.example.com/authorize?state=abc").unwrap(); + let authorization = authorization_prompt(&authorization_url); + let device = device_prompt("https://idp.example.com/device", "ABCD-EFGH"); + let mut output = Vec::new(); + + write_oauth_prompt(&mut output, &authorization); + write_oauth_prompt(&mut output, &device); + + let output = String::from_utf8(output).unwrap(); + assert!(output.contains(authorization_url.as_str())); + assert!(output.contains("https://idp.example.com/device")); + assert!(output.contains("ABCD-EFGH")); + } + + #[test] + fn test_authorization_code_options_default_to_pkce() { + let options = AuthorizationCodeOptions::new(); + + assert!(options.use_pkce); + assert!(options.redirect_uri.is_none()); + assert!(options.callback_port.is_none()); + } + + #[test] + fn test_authorization_redirect_defaults_to_ipv4_loopback() { + let redirect = ResolvedRedirect::new(&AuthorizationCodeOptions::new()).unwrap(); + + assert_eq!( + redirect.uri, + format!("http://127.0.0.1:{DEFAULT_CALLBACK_PORT}/callback") + ); + assert!(redirect.bind_addr.ip().is_loopback()); + assert_eq!(redirect.bind_addr.port(), DEFAULT_CALLBACK_PORT); + assert_eq!(redirect.callback_path, "/callback"); + } + + #[test] + fn test_authorization_redirect_rejects_non_loopback_host() { + let options = AuthorizationCodeOptions::new() + .redirect_uri("https://client.example.com/oauth/callback"); + + let err = ResolvedRedirect::new(&options).unwrap_err(); + assert!(matches!( + err, + Error::InvalidInput { message } + if message == "OAuth redirect_uri must use http with a loopback host" + )); + } + + #[test] + fn test_authorization_redirect_rejects_mismatched_port() { + let options = AuthorizationCodeOptions::new() + .redirect_uri("http://127.0.0.1:8401/callback") + .callback_port(8400); + + let err = ResolvedRedirect::new(&options).unwrap_err(); + assert!(matches!( + err, + Error::InvalidInput { message } + if message.contains("does not match redirect_uri port") + )); + } + + #[test] + fn test_authorization_redirect_requires_explicit_port() { + let options = AuthorizationCodeOptions::new().redirect_uri("http://127.0.0.1/callback"); + + let err = ResolvedRedirect::new(&options).unwrap_err(); + assert!(matches!( + err, + Error::InvalidInput { message } + if message == "OAuth redirect_uri must include a port" + )); + } + + #[test] + fn test_authorization_redirect_accepts_ipv6_loopback() { + let options = AuthorizationCodeOptions::new().redirect_uri("http://[::1]:8400/callback"); + + let redirect = ResolvedRedirect::new(&options).unwrap(); + assert_eq!( + redirect.bind_addr, + "[::1]:8400".parse::().unwrap() + ); + } + + #[test] + fn test_authorization_callback_parsing() { + let callback = parse_authorization_callback( + "GET /callback?code=abc%20123&state=expected HTTP/1.1\r\nHost: localhost\r\n", + "/callback", + "expected", + ) + .unwrap(); + + assert_eq!(callback, AuthorizationCallback::Code("abc 123".to_string())); + } + + #[test] + fn test_authorization_callback_rejects_state_mismatch() { + let err = parse_authorization_callback( + "GET /callback?code=abc&state=wrong HTTP/1.1\r\nHost: localhost\r\n", + "/callback", + "expected", + ) + .unwrap_err(); + + assert!(matches!( + err, + Error::Runtime { message } + if message == "OAuth authorization callback state did not match" + )); + } + + #[test] + fn test_authorization_callback_reports_provider_error() { + let callback = parse_authorization_callback( + "GET /callback?error=access_denied&error_description=user+cancelled&state=expected HTTP/1.1\r\nHost: localhost\r\n", + "/callback", + "expected", + ) + .unwrap(); + + assert_eq!( + callback, + AuthorizationCallback::ProviderError( + "OAuth authorization failed: user cancelled".to_string() + ) + ); + } + + #[tokio::test] + async fn test_authorization_callback_ignores_unrelated_connection_and_partial_read() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let source = AuthorizationCodeSource::new( + "http://127.0.0.1:1".to_string(), + "client-id".to_string(), + None, + vec!["openid".to_string()], + AuthorizationCodeOptions::new() + .redirect_uri(format!("http://127.0.0.1:{port}/callback")), + ) + .unwrap(); + + let browser = tokio::spawn(async move { + let mut unrelated = TcpStream::connect(("127.0.0.1", port)).await.unwrap(); + unrelated + .write_all(b"GET /favicon.ico HTTP/1.1\r\nHost: localhost\r\n\r\n") + .await + .unwrap(); + let mut response = Vec::new(); + unrelated.read_to_end(&mut response).await.unwrap(); + assert!( + String::from_utf8(response) + .unwrap() + .starts_with("HTTP/1.1 400") + ); + + let mut callback = TcpStream::connect(("127.0.0.1", port)).await.unwrap(); + callback + .write_all(b"GET /callback?code=auth") + .await + .unwrap(); + tokio::task::yield_now().await; + callback + .write_all(b"-code&state=expected HTTP/1.1\r\nHost: localhost\r\n\r\n") + .await + .unwrap(); + }); + + assert_eq!( + source + .wait_for_callback(&listener, "expected") + .await + .unwrap(), + "auth-code" + ); + browser.await.unwrap(); + } + + #[tokio::test] + async fn test_authorization_callback_read_respects_overall_deadline() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let client = tokio::spawn(async move { + let _stream = TcpStream::connect(("127.0.0.1", port)).await.unwrap(); + tokio::time::sleep(Duration::from_secs(1)).await; + }); + let (mut stream, _) = listener.accept().await.unwrap(); + + let err = read_authorization_callback( + &mut stream, + "/callback", + "expected", + TokioInstant::now() + Duration::from_millis(20), + ) + .await + .unwrap_err(); + + assert!(matches!( + err, + Error::Runtime { message } + if message == "Timed out reading the OAuth authorization callback" + )); + client.abort(); + } + + #[tokio::test] + async fn test_authorization_request_uses_pkce_by_default() { + let (issuer_url, server) = spawn_discovery_server(1).await; + let source = AuthorizationCodeSource::new( + issuer_url, + "client-id".to_string(), + None, + vec!["openid".to_string(), "profile".to_string()], + AuthorizationCodeOptions::new(), + ) + .unwrap(); + + let request = source.build_authorization_request().await.unwrap(); + let params: HashMap<_, _> = request.url.query_pairs().into_owned().collect(); + assert_eq!( + params.get("response_type").map(String::as_str), + Some("code") + ); + assert_eq!( + params.get("scope").map(String::as_str), + Some("openid profile") + ); + assert_eq!( + params.get("code_challenge_method").map(String::as_str), + Some("S256") + ); + assert!(params.contains_key("code_challenge")); + assert!(request.code_verifier.is_some()); + server.await.unwrap(); + } + + #[tokio::test] + async fn test_authorization_request_rejects_plaintext_provider_endpoint() { + let (issuer_url, server) = spawn_insecure_authorization_discovery_server().await; + let source = AuthorizationCodeSource::new( + issuer_url, + "client-id".to_string(), + None, + vec!["openid".to_string()], + AuthorizationCodeOptions::new(), + ) + .unwrap(); + + let err = source.build_authorization_request().await.unwrap_err(); + assert!(matches!( + err, + Error::InvalidInput { message } + if message.contains("authorization_endpoint must use https") + )); + server.await.unwrap(); + } + + #[tokio::test] + async fn test_authorization_request_can_disable_pkce() { + let (issuer_url, server) = spawn_discovery_server(1).await; + let source = AuthorizationCodeSource::new( + issuer_url, + "client-id".to_string(), + Some("secret".to_string()), + vec!["openid".to_string()], + AuthorizationCodeOptions::new().use_pkce(false), + ) + .unwrap(); + + let request = source.build_authorization_request().await.unwrap(); + let params: HashMap<_, _> = request.url.query_pairs().into_owned().collect(); + assert!(!params.contains_key("code_challenge")); + assert!(!params.contains_key("code_challenge_method")); + assert!(request.code_verifier.is_none()); + server.await.unwrap(); + } + + #[tokio::test] + async fn test_authorization_code_exchange_includes_optional_credentials() { + let (issuer_url, request_body, server) = spawn_token_exchange_server().await; + let source = AuthorizationCodeSource::new( + issuer_url, + "client-id".to_string(), + Some("secret".to_string()), + vec!["openid".to_string()], + AuthorizationCodeOptions::new(), + ) + .unwrap(); + + let response = source + .exchange_code("auth-code", Some("verifier")) + .await + .unwrap(); + assert_eq!(response.access_token, "token"); + let body = request_body.lock().unwrap().clone().unwrap(); + assert!(body.contains("grant_type=authorization_code")); + assert!(body.contains("code=auth-code")); + assert!(body.contains("code_verifier=verifier")); + assert!(body.contains("client_secret=secret")); + server.await.unwrap(); + } + + #[tokio::test] + async fn test_refresh_invalid_grant_requires_reauthentication() { + let (issuer_url, server) = + spawn_refresh_error_server("400 Bad Request", r#"{"error":"invalid_grant"}"#).await; + let source = AuthorizationCodeSource::new( + issuer_url, + "client-id".to_string(), + None, + vec!["openid".to_string()], + AuthorizationCodeOptions::new(), + ) + .unwrap(); + + assert!(matches!( + source.oidc.refresh_token("revoked").await.unwrap(), + RefreshResult::Reauthenticate + )); + server.await.unwrap(); + } + + #[tokio::test] + async fn test_refresh_transient_failure_remains_retryable() { + let (issuer_url, server) = spawn_refresh_error_server( + "503 Service Unavailable", + r#"{"error":"temporarily_unavailable"}"#, + ) + .await; + let source = AuthorizationCodeSource::new( + issuer_url, + "client-id".to_string(), + None, + vec!["openid".to_string()], + AuthorizationCodeOptions::new(), + ) + .unwrap(); + + let err = source.oidc.refresh_token("still-valid").await.unwrap_err(); + assert!(matches!( + err, + Error::Runtime { message } + if message.contains("503 Service Unavailable") + )); + server.await.unwrap(); + } + + #[tokio::test] + async fn test_device_authorization_polls_until_success() { + let (issuer_url, token_requests, server) = spawn_device_server().await; + let source = DeviceCodeSource::new( + issuer_url, + "client-id".to_string(), + Some("secret".to_string()), + vec!["openid".to_string()], + ) + .unwrap(); + + let device = source.request_device_authorization().await.unwrap(); + let response = source.poll_for_token(&device).await.unwrap(); + + assert_eq!(response.access_token, "device-token"); + assert_eq!(response.refresh_token.as_deref(), Some("device-refresh")); + assert_eq!(token_requests.load(Ordering::SeqCst), 3); + server.await.unwrap(); + } + + #[tokio::test] + async fn test_device_authorization_rejects_plaintext_verification_uri() { + let (issuer_url, server) = spawn_insecure_device_verification_server().await; + let source = DeviceCodeSource::new( + issuer_url, + "client-id".to_string(), + None, + vec!["openid".to_string()], + ) + .unwrap(); + + let Err(err) = source.request_device_authorization().await else { + panic!("expected insecure verification URI to be rejected"); + }; + assert!(matches!( + err, + Error::InvalidInput { message } + if message.contains("verification_uri must use https") + )); + server.await.unwrap(); + } + + #[tokio::test] + async fn test_device_authorization_retries_transient_failures() { + let (issuer_url, token_requests, server) = spawn_device_transient_server().await; + let source = DeviceCodeSource::new( + issuer_url, + "client-id".to_string(), + None, + vec!["openid".to_string()], + ) + .unwrap(); + let device = test_device_authorization_response(10, 1); + + let response = source.poll_for_token(&device).await.unwrap(); + + assert_eq!(response.access_token, "device-token"); + assert_eq!(token_requests.load(Ordering::SeqCst), 4); + server.await.unwrap(); + } + + #[tokio::test] + async fn test_device_authorization_reports_access_denied() { + let (issuer_url, server) = spawn_device_error_server("access_denied").await; + let source = DeviceCodeSource::new( + issuer_url, + "client-id".to_string(), + None, + vec!["openid".to_string()], + ) + .unwrap(); + let device = test_device_authorization_response(60, 1); + + let err = source.poll_for_token(&device).await.unwrap_err(); + assert!(matches!( + err, + Error::Runtime { message } + if message == "Device authorization was denied by the user" + )); + server.await.unwrap(); + } + + #[tokio::test] + async fn test_device_authorization_reports_provider_expiry() { + let (issuer_url, server) = spawn_device_error_server("expired_token").await; + let source = DeviceCodeSource::new( + issuer_url, + "client-id".to_string(), + None, + vec!["openid".to_string()], + ) + .unwrap(); + let device = test_device_authorization_response(60, 1); + + let err = source.poll_for_token(&device).await.unwrap_err(); + assert!(matches!( + err, + Error::Runtime { message } + if message == "Device authorization expired before authentication completed" + )); + server.await.unwrap(); + } + + #[tokio::test] + async fn test_device_authorization_stops_at_local_deadline() { + let (issuer_url, server) = spawn_discovery_server(1).await; + let source = DeviceCodeSource::new( + issuer_url, + "client-id".to_string(), + None, + vec!["openid".to_string()], + ) + .unwrap(); + let device = test_device_authorization_response(1, 5); + + let err = source.poll_for_token(&device).await.unwrap_err(); + assert!(matches!( + err, + Error::Runtime { message } + if message == "Device authorization expired before authentication completed" + )); + server.await.unwrap(); + } + + #[derive(Debug)] + struct RefreshingTokenSource { + fetches: Arc, + refreshes: Arc, + } + + #[async_trait] + impl TokenSource for RefreshingTokenSource { + async fn fetch_token(&self) -> Result { + self.fetches.fetch_add(1, Ordering::SeqCst); + Ok(TokenResponse { + access_token: "initial".to_string(), + refresh_token: Some("refresh".to_string()), + expires_in: Some(3600), + token_type: Some("Bearer".to_string()), + }) + } + + async fn refresh_token(&self, refresh_token: &str) -> Result { + assert_eq!(refresh_token, "refresh"); + self.refreshes.fetch_add(1, Ordering::SeqCst); + Ok(RefreshResult::Refreshed(TokenResponse { + access_token: "refreshed".to_string(), + refresh_token: None, + expires_in: Some(3600), + token_type: Some("Bearer".to_string()), + })) + } + } + + #[tokio::test] + async fn test_header_provider_uses_and_retains_refresh_token() { + let fetches = Arc::new(AtomicUsize::new(0)); + let refreshes = Arc::new(AtomicUsize::new(0)); + let provider = OAuthHeaderProvider { + token_source: Box::new(RefreshingTokenSource { + fetches: Arc::clone(&fetches), + refreshes: Arc::clone(&refreshes), + }), + token_state: Arc::new(RwLock::new(TokenState::new())), + refresh_buffer: Duration::ZERO, + token_cache: None, + }; + + assert_eq!(provider.get_valid_token().await.unwrap(), "initial"); + provider.token_state.write().await.expires_at = + Some(Instant::now() - Duration::from_secs(1)); + assert_eq!(provider.get_valid_token().await.unwrap(), "refreshed"); + assert_eq!(fetches.load(Ordering::SeqCst), 1); + assert_eq!(refreshes.load(Ordering::SeqCst), 1); + assert_eq!( + provider.token_state.read().await.refresh_token.as_deref(), + Some("refresh") + ); + } + + #[derive(Debug)] + struct FailedRefreshTokenSource { + fetches: Arc, + refreshes: Arc, + } + + #[async_trait] + impl TokenSource for FailedRefreshTokenSource { + async fn fetch_token(&self) -> Result { + self.fetches.fetch_add(1, Ordering::SeqCst); + Ok(TokenResponse { + access_token: "reauthenticated".to_string(), + refresh_token: Some("new-refresh".to_string()), + expires_in: Some(3600), + token_type: Some("Bearer".to_string()), + }) + } + + async fn refresh_token(&self, refresh_token: &str) -> Result { + assert_eq!(refresh_token, "revoked-refresh"); + self.refreshes.fetch_add(1, Ordering::SeqCst); + Ok(RefreshResult::Reauthenticate) + } + } + + #[tokio::test] + async fn test_header_provider_reauthenticates_after_refresh_failure() { + let fetches = Arc::new(AtomicUsize::new(0)); + let refreshes = Arc::new(AtomicUsize::new(0)); + let provider = OAuthHeaderProvider { + token_source: Box::new(FailedRefreshTokenSource { + fetches: Arc::clone(&fetches), + refreshes: Arc::clone(&refreshes), + }), + token_state: Arc::new(RwLock::new(TokenState { + access_token: Some("expired".to_string()), + refresh_token: Some("revoked-refresh".to_string()), + expires_at: Some(Instant::now() - Duration::from_secs(1)), + })), + refresh_buffer: Duration::ZERO, + token_cache: None, + }; + + assert_eq!(provider.get_valid_token().await.unwrap(), "reauthenticated"); + assert_eq!(fetches.load(Ordering::SeqCst), 1); + assert_eq!(refreshes.load(Ordering::SeqCst), 1); + assert_eq!( + provider.token_state.read().await.refresh_token.as_deref(), + Some("new-refresh") + ); + } + + #[derive(Debug)] + struct TransientRefreshFailureSource { + fetches: Arc, + } + + #[async_trait] + impl TokenSource for TransientRefreshFailureSource { + async fn fetch_token(&self) -> Result { + self.fetches.fetch_add(1, Ordering::SeqCst); + unreachable!("a transient refresh failure must not start an interactive flow") + } + + async fn refresh_token(&self, refresh_token: &str) -> Result { + assert_eq!(refresh_token, "valid-refresh"); + Err(Error::Runtime { + message: "token endpoint temporarily unavailable".to_string(), + }) + } + } + + #[tokio::test] + async fn test_header_provider_preserves_refresh_token_after_transient_failure() { + let fetches = Arc::new(AtomicUsize::new(0)); + let provider = OAuthHeaderProvider { + token_source: Box::new(TransientRefreshFailureSource { + fetches: Arc::clone(&fetches), + }), + token_state: Arc::new(RwLock::new(TokenState { + access_token: Some("expired".to_string()), + refresh_token: Some("valid-refresh".to_string()), + expires_at: Some(Instant::now() - Duration::from_secs(1)), + })), + refresh_buffer: Duration::ZERO, + token_cache: None, + }; + + let err = provider.get_valid_token().await.unwrap_err(); + assert!(matches!( + err, + Error::Runtime { message } + if message == "token endpoint temporarily unavailable" + )); + assert_eq!(fetches.load(Ordering::SeqCst), 0); + assert_eq!( + provider.token_state.read().await.refresh_token.as_deref(), + Some("valid-refresh") + ); } #[test] @@ -653,6 +2192,7 @@ mod tests { scopes: vec!["scope".to_string()], flow: OAuthFlow::ClientCredentials, refresh_buffer_secs: None, + token_cache: None, }; let debug = format!("{config:?}"); @@ -669,12 +2209,13 @@ mod tests { scopes: vec!["scope".to_string()], flow: OAuthFlow::ClientCredentials, refresh_buffer_secs: None, + token_cache: None, }; let provider = OAuthHeaderProvider::new(config).unwrap(); let debug = format!("{provider:?}"); assert!(!debug.contains("super-secret")); - assert!(debug.contains("client_secret: \"\"")); + assert!(debug.contains("client_secret: Some(\"\")")); } #[test] @@ -705,6 +2246,7 @@ mod tests { ], flow: OAuthFlow::AzureManagedIdentity { client_id: None }, refresh_buffer_secs: None, + token_cache: None, }; assert!(OAuthHeaderProvider::new(config).is_err()); } @@ -720,7 +2262,7 @@ mod tests { ) .unwrap(); - let err = source.get_token_endpoint().await.unwrap_err(); + let err = source.oidc.get_token_endpoint().await.unwrap_err(); assert!(matches!( err, Error::Runtime { message } @@ -738,6 +2280,7 @@ mod tests { scopes: vec!["scope".to_string()], flow: OAuthFlow::ClientCredentials, refresh_buffer_secs: None, + token_cache: None, }; assert!(OAuthHeaderProvider::new(config).is_err()); } @@ -751,13 +2294,15 @@ mod tests { scopes: vec!["scope".to_string()], flow: OAuthFlow::ClientCredentials, refresh_buffer_secs: None, + token_cache: None, }; let err = OAuthHeaderProvider::new(config).unwrap_err(); assert!(matches!( err, Error::InvalidInput { message } - if message == "ClientCredentials OAuth issuer_url must use https, except for loopback hosts" + if message + == "OAuth issuer_url must use https, except for http on a loopback host" )); } @@ -770,6 +2315,7 @@ mod tests { scopes: vec![], flow: OAuthFlow::AzureManagedIdentity { client_id: None }, refresh_buffer_secs: None, + token_cache: None, }; assert!(OAuthHeaderProvider::new(config).is_err()); } @@ -784,6 +2330,7 @@ mod tests { scopes: vec!["scope".to_string()], flow: OAuthFlow::ClientCredentials, refresh_buffer_secs: Some(0), + token_cache: None, }; let provider = OAuthHeaderProvider::new(config).unwrap(); @@ -805,6 +2352,297 @@ mod tests { server.await.unwrap(); } + async fn spawn_discovery_server(expected_requests: usize) -> (String, JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let issuer_url = format!("http://{addr}"); + + let server = tokio::spawn(async move { + for _ in 0..expected_requests { + let (mut stream, _) = listener.accept().await.unwrap(); + let (request_line, _) = read_http_request(&mut stream).await; + assert!(request_line.starts_with("GET /.well-known/openid-configuration ")); + let discovery = format!( + r#"{{"token_endpoint":"http://{addr}/token","authorization_endpoint":"http://{addr}/authorize","device_authorization_endpoint":"http://{addr}/device"}}"# + ); + write_json_response(&mut stream, "200 OK", &discovery).await; + } + }); + + (issuer_url, server) + } + + async fn spawn_insecure_authorization_discovery_server() -> (String, JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let issuer_url = format!("http://{addr}"); + let server = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let (request_line, _) = read_http_request(&mut stream).await; + assert!(request_line.starts_with("GET /.well-known/openid-configuration ")); + write_json_response( + &mut stream, + "200 OK", + r#"{"token_endpoint":"https://idp.example.com/token","authorization_endpoint":"http://idp.example.com/authorize"}"#, + ) + .await; + }); + + (issuer_url, server) + } + + async fn spawn_insecure_device_verification_server() -> (String, JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let issuer_url = format!("http://{addr}"); + let server = tokio::spawn(async move { + for _ in 0..2 { + let (mut stream, _) = listener.accept().await.unwrap(); + let (request_line, _) = read_http_request(&mut stream).await; + if request_line.starts_with("GET /.well-known/openid-configuration ") { + let discovery = format!( + r#"{{"token_endpoint":"http://{addr}/token","device_authorization_endpoint":"http://{addr}/device"}}"# + ); + write_json_response(&mut stream, "200 OK", &discovery).await; + } else { + assert!(request_line.starts_with("POST /device ")); + write_json_response( + &mut stream, + "200 OK", + r#"{"device_code":"device-code","user_code":"ABCD-EFGH","verification_uri":"http://idp.example.com/device","expires_in":60}"#, + ) + .await; + } + } + }); + + (issuer_url, server) + } + + async fn spawn_token_exchange_server() -> ( + String, + Arc>>, + JoinHandle<()>, + ) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let issuer_url = format!("http://{addr}"); + let request_body = Arc::new(std::sync::Mutex::new(None)); + let server_request_body = Arc::clone(&request_body); + + let server = tokio::spawn(async move { + for _ in 0..2 { + let (mut stream, _) = listener.accept().await.unwrap(); + let (request_line, body) = read_http_request(&mut stream).await; + if request_line.starts_with("GET /.well-known/openid-configuration ") { + let discovery = format!( + r#"{{"token_endpoint":"http://{addr}/token","authorization_endpoint":"http://{addr}/authorize"}}"# + ); + write_json_response(&mut stream, "200 OK", &discovery).await; + } else if request_line.starts_with("POST /token ") { + *server_request_body.lock().unwrap() = Some(body); + write_json_response( + &mut stream, + "200 OK", + r#"{"access_token":"token","refresh_token":"refresh","expires_in":3600}"#, + ) + .await; + } else { + write_json_response(&mut stream, "404 Not Found", "{}").await; + } + } + }); + + (issuer_url, request_body, server) + } + + async fn spawn_refresh_error_server( + status: &'static str, + response_body: &'static str, + ) -> (String, JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let issuer_url = format!("http://{addr}"); + + let server = tokio::spawn(async move { + for _ in 0..2 { + let (mut stream, _) = listener.accept().await.unwrap(); + let (request_line, body) = read_http_request(&mut stream).await; + if request_line.starts_with("GET /.well-known/openid-configuration ") { + let discovery = format!(r#"{{"token_endpoint":"http://{addr}/token"}}"#); + write_json_response(&mut stream, "200 OK", &discovery).await; + } else if request_line.starts_with("POST /token ") { + assert!(body.contains("grant_type=refresh_token")); + assert!(body.contains("refresh_token=")); + write_json_response(&mut stream, status, response_body).await; + } else { + write_json_response(&mut stream, "404 Not Found", "{}").await; + } + } + }); + + (issuer_url, server) + } + + async fn spawn_device_server() -> (String, Arc, JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let issuer_url = format!("http://{addr}"); + let token_requests = Arc::new(AtomicUsize::new(0)); + let server_token_requests = Arc::clone(&token_requests); + + let server = tokio::spawn(async move { + for _ in 0..5 { + let (mut stream, _) = listener.accept().await.unwrap(); + let (request_line, body) = read_http_request(&mut stream).await; + if request_line.starts_with("GET /.well-known/openid-configuration ") { + let discovery = format!( + r#"{{"token_endpoint":"http://{addr}/token","device_authorization_endpoint":"http://{addr}/device"}}"# + ); + write_json_response(&mut stream, "200 OK", &discovery).await; + } else if request_line.starts_with("POST /device ") { + assert!(body.contains("client_id=client-id")); + assert!(body.contains("client_secret=secret")); + assert!(body.contains("scope=openid")); + let device = format!( + r#"{{"device_code":"device-code","user_code":"ABCD-EFGH","verification_uri":"http://{addr}/verify","verification_uri_complete":"http://{addr}/verify?user_code=ABCD-EFGH","expires_in":60,"interval":1}}"# + ); + write_json_response(&mut stream, "200 OK", &device).await; + } else if request_line.starts_with("POST /token ") { + assert!(body.contains( + "grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Adevice_code" + )); + assert!(body.contains("device_code=device-code")); + assert!(body.contains("client_secret=secret")); + let request = server_token_requests.fetch_add(1, Ordering::SeqCst); + match request { + 0 => { + write_json_response( + &mut stream, + "400 Bad Request", + r#"{"error":"authorization_pending"}"#, + ) + .await; + } + 1 => { + write_json_response( + &mut stream, + "400 Bad Request", + r#"{"error":"slow_down"}"#, + ) + .await; + } + _ => { + write_json_response( + &mut stream, + "200 OK", + r#"{"access_token":"device-token","refresh_token":"device-refresh","expires_in":3600}"#, + ) + .await; + } + } + } else { + write_json_response(&mut stream, "404 Not Found", "{}").await; + } + } + }); + + (issuer_url, token_requests, server) + } + + async fn spawn_device_transient_server() -> (String, Arc, JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let issuer_url = format!("http://{addr}"); + let token_requests = Arc::new(AtomicUsize::new(0)); + let server_token_requests = Arc::clone(&token_requests); + + let server = tokio::spawn(async move { + for _ in 0..5 { + let (mut stream, _) = listener.accept().await.unwrap(); + let (request_line, _) = read_http_request(&mut stream).await; + if request_line.starts_with("GET /.well-known/openid-configuration ") { + let discovery = format!(r#"{{"token_endpoint":"http://{addr}/token"}}"#); + write_json_response(&mut stream, "200 OK", &discovery).await; + continue; + } + + assert!(request_line.starts_with("POST /token ")); + match server_token_requests.fetch_add(1, Ordering::SeqCst) { + 0 => drop(stream), + 1 => { + write_json_response( + &mut stream, + "503 Service Unavailable", + r#"{"error":"server_error"}"#, + ) + .await; + } + 2 => { + write_json_response( + &mut stream, + "400 Bad Request", + r#"{"error":"temporarily_unavailable"}"#, + ) + .await; + } + _ => { + write_json_response( + &mut stream, + "200 OK", + r#"{"access_token":"device-token","expires_in":3600}"#, + ) + .await; + } + } + } + }); + + (issuer_url, token_requests, server) + } + + fn test_device_authorization_response( + expires_in: u64, + interval: u64, + ) -> DeviceAuthorizationResponse { + DeviceAuthorizationResponse { + device_code: "device-code".to_string(), + user_code: "ABCD-EFGH".to_string(), + verification_uri: "http://127.0.0.1/verify".to_string(), + verification_uri_complete: None, + expires_in, + interval: Some(interval), + } + } + + async fn spawn_device_error_server(error: &'static str) -> (String, JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let issuer_url = format!("http://{addr}"); + + let server = tokio::spawn(async move { + for _ in 0..2 { + let (mut stream, _) = listener.accept().await.unwrap(); + let (request_line, _) = read_http_request(&mut stream).await; + if request_line.starts_with("GET /.well-known/openid-configuration ") { + let discovery = format!(r#"{{"token_endpoint":"http://{addr}/token"}}"#); + write_json_response(&mut stream, "200 OK", &discovery).await; + } else if request_line.starts_with("POST /token ") { + write_json_response( + &mut stream, + "400 Bad Request", + &format!(r#"{{"error":"{error}"}}"#), + ) + .await; + } else { + write_json_response(&mut stream, "404 Not Found", "{}").await; + } + } + }); + + (issuer_url, server) + } + async fn spawn_oauth_server() -> (String, Arc, JoinHandle<()>) { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); diff --git a/rust/lancedb/src/remote/token_cache.rs b/rust/lancedb/src/remote/token_cache.rs new file mode 100644 index 000000000..b67ced651 --- /dev/null +++ b/rust/lancedb/src/remote/token_cache.rs @@ -0,0 +1,1753 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Persistent OAuth token cache and session lifecycle APIs. +//! +//! By default, OAuth sessions are kept in process memory only (see +//! [`OAuthHeaderProvider`](crate::remote::OAuthHeaderProvider)). Short-lived +//! processes such as CLI tools, notebooks, or scripts would otherwise have to +//! run a full interactive browser or device flow on every start. Configuring +//! [`TokenCacheOptions`] on an [`OAuthConfig`](crate::remote::OAuthConfig) opts +//! in to an explicit, hardened, on-disk cache that stores only the refresh +//! token plus non-secret metadata, so a second process can silently refresh +//! instead of re-prompting. +//! +//! Security properties: +//! +//! - Opt-in only; callers that do not configure a cache stay memory-only. +//! - Only refresh tokens are persisted. Access tokens never touch disk, so +//! there are no local expiry decisions to get wrong when clocks move. +//! - No client secret is ever stored. +//! - The cache directory is private (`0700`) and each record is `0600`, +//! owner-checked, and symlink-rejected on Unix; records are replaced +//! atomically via `rename` so a crash can never leave a torn file. +//! - Cache filenames are SHA-256 hashes of the canonical issuer, client, +//! scope, flow, and client-auth identity. No secret appears in a filename. +//! - Refresh-token rotation is serialized across processes with a per-key +//! advisory file lock (`flock` on Unix, `LockFileEx` on Windows). The +//! operating system releases these locks when a process dies, so a crash +//! cannot leave a stale lock behind. +//! +//! Use [`OAuthSession`] to explicitly `login`, inspect `status`, or `logout` +//! without issuing a database request. +//! +//! # Example +//! +//! ``` +//! use lancedb::remote::{OAuthConfig, OAuthFlow, TokenCacheOptions}; +//! +//! # async fn example() -> Result<(), Box> { +//! let config = OAuthConfig { +//! issuer_url: "https://issuer.example.com".to_string(), +//! client_id: "my-app".to_string(), +//! client_secret: None, +//! scopes: vec!["openid".to_string()], +//! flow: OAuthFlow::DeviceCode, +//! refresh_buffer_secs: None, +//! token_cache: Some( +//! TokenCacheOptions::new().cache_dir("/tmp/my-app/oauth-cache"), +//! ), +//! }; +//! let session = lancedb::remote::OAuthSession::new(config)?; +//! session.login().await?; +//! let status = session.status().await?; +//! assert!(status.refreshable); +//! # Ok(()) +//! # } +//! ``` + +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use fs4::fs_std::FileExt; +use log::{debug, warn}; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; + +use crate::error::{Error, Result}; +use crate::remote::oauth::{OAuthConfig, OAuthFlow, RefreshResult, TokenResponse, TokenSource}; + +const CACHE_RECORD_VERSION: u32 = 1; +const DEFAULT_LOCK_TIMEOUT_SECS: u64 = 30; +const LOCK_POLL_INTERVAL: Duration = Duration::from_millis(100); + +fn now_unix_secs() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_secs()) + .unwrap_or(0) +} + +/// Options for the persistent OAuth token cache. +/// +/// The cache is opt-in: it is only used when set on +/// [`OAuthConfig::token_cache`](crate::remote::OAuthConfig::token_cache). See +/// the [module documentation](self) for the security properties. +#[derive(Clone, Debug, Default)] +pub struct TokenCacheOptions { + /// Directory that holds the 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, + + /// How long to wait for the cross-process refresh lock before failing. + /// + /// Defaults to 30 seconds. + pub lock_timeout_secs: Option, +} + +impl TokenCacheOptions { + /// Create cache options with all defaults. + pub fn new() -> Self { + Self::default() + } + + /// Set the directory that holds cached credentials. + pub fn cache_dir(mut self, cache_dir: impl Into) -> Self { + self.cache_dir = Some(cache_dir.into()); + self + } + + /// Set the cross-process refresh lock timeout in seconds. + pub fn lock_timeout_secs(mut self, secs: u64) -> Self { + self.lock_timeout_secs = Some(secs); + self + } + + fn resolved_dir(&self) -> Result { + if let Some(dir) = &self.cache_dir { + if dir.as_os_str().is_empty() { + return Err(Error::InvalidInput { + message: "OAuth token cache directory must not be empty".to_string(), + }); + } + return Ok(dir.clone()); + } + default_cache_dir().ok_or_else(|| Error::InvalidInput { + message: "Could not determine a default OAuth token cache directory; \ + set TokenCacheOptions::cache_dir or XDG_CACHE_HOME/HOME" + .to_string(), + }) + } + + fn lock_timeout(&self) -> Duration { + Duration::from_secs(self.lock_timeout_secs.unwrap_or(DEFAULT_LOCK_TIMEOUT_SECS)) + } +} + +#[cfg(unix)] +fn default_cache_dir() -> Option { + let base = std::env::var_os("XDG_CACHE_HOME") + .filter(|value| !value.is_empty()) + .map(PathBuf::from) + .or_else(|| { + std::env::var_os("HOME") + .filter(|value| !value.is_empty()) + .map(|home| { + let mut path = PathBuf::from(home); + path.push(".cache"); + path + }) + })?; + let mut dir = base; + dir.push("lancedb"); + dir.push("oauth"); + Some(dir) +} + +#[cfg(windows)] +fn default_cache_dir() -> Option { + let base = std::env::var_os("LOCALAPPDATA") + .filter(|value| !value.is_empty()) + .map(PathBuf::from)?; + let mut dir = base; + dir.push("lancedb"); + dir.push("oauth"); + Some(dir) +} + +#[cfg(not(any(unix, windows)))] +fn default_cache_dir() -> Option { + None +} + +/// A cached OAuth session record. +/// +/// Only the refresh token is persisted. The metadata mirrors the cache key so +/// `status` can report what a record belongs to without exposing secrets. +#[derive(Serialize, Deserialize)] +struct CachedTokenRecord { + version: u32, + issuer_url: String, + client_id: String, + scopes: Vec, + flow: String, + client_auth: String, + refresh_token: String, + obtained_at: u64, +} + +impl std::fmt::Debug for CachedTokenRecord { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CachedTokenRecord") + .field("version", &self.version) + .field("issuer_url", &self.issuer_url) + .field("client_id", &self.client_id) + .field("scopes", &self.scopes) + .field("flow", &self.flow) + .field("client_auth", &self.client_auth) + .field("refresh_token", &"") + .field("obtained_at", &self.obtained_at) + .finish() + } +} + +fn canonicalize_issuer(issuer_url: &str) -> String { + issuer_url.trim_end_matches('/').to_string() +} + +fn canonicalize_scopes(scopes: &[String]) -> Vec { + let mut canonical: Vec = scopes + .iter() + .map(|scope| scope.trim().to_string()) + .filter(|scope| !scope.is_empty()) + .collect(); + canonical.sort(); + canonical.dedup(); + canonical +} + +/// Returns the cache identity of a flow, or `None` for flows that never +/// persist: client credentials have no refresh token to store, and Azure +/// managed identity is machine identity that must not enter a user token +/// cache (rejected separately with an explicit error). +fn flow_key(flow: &OAuthFlow) -> Option<&'static str> { + match flow { + OAuthFlow::AuthorizationCode(_) => Some("authorization_code"), + OAuthFlow::DeviceCode => Some("device_code"), + OAuthFlow::ClientCredentials | OAuthFlow::AzureManagedIdentity { .. } => None, + } +} + +fn client_auth_key(client_secret: Option<&str>) -> &'static str { + if client_secret.is_some() { + "confidential" + } else { + "public" + } +} + +/// Identity of one cached session: canonical issuer, client, scopes, flow, +/// and client-auth mode, plus the hashed filename derived from it. +#[derive(Clone, Debug)] +struct CacheKey { + issuer_url: String, + client_id: String, + scopes: Vec, + flow: &'static str, + client_auth: &'static str, + file_stem: String, +} + +impl CacheKey { + fn new(config: &OAuthConfig) -> Result { + let flow = flow_key(&config.flow).ok_or_else(|| Error::InvalidInput { + message: format!( + "A persistent OAuth token cache is not supported for the {:?} flow; \ + remove TokenCacheOptions to keep tokens in memory", + config.flow + ), + })?; + let issuer_url = canonicalize_issuer(&config.issuer_url); + let scopes = canonicalize_scopes(&config.scopes); + let client_auth = client_auth_key(config.client_secret.as_deref()); + let identity = format!( + "v1\n{}\n{}\n{}\n{}\n{}", + issuer_url, + config.client_id, + scopes.join(" "), + flow, + client_auth + ); + let file_stem = hex_sha256(identity.as_bytes()); + Ok(Self { + issuer_url, + client_id: config.client_id.clone(), + scopes, + flow, + client_auth, + file_stem, + }) + } +} + +fn hex_sha256(bytes: &[u8]) -> String { + let digest = Sha256::digest(bytes); + let mut hex = String::with_capacity(digest.len() * 2); + for byte in digest { + use std::fmt::Write; + let _ = write!(hex, "{byte:02x}"); + } + hex +} + +/// Guard for the per-key cross-process refresh lock. +/// +/// The lock is an advisory exclusive lock on a per-key file. The operating +/// system releases it when the owning process exits, so crashes cannot strand +/// a stale lock. +struct LockGuard { + #[allow(dead_code)] + file: std::fs::File, +} + +/// The persistent token cache engine for one [`OAuthConfig`]. +pub(crate) struct TokenCache { + dir: PathBuf, + key: CacheKey, + lock_timeout: Duration, + #[cfg(unix)] + dir_owner: u32, +} + +impl std::fmt::Debug for TokenCache { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TokenCache") + .field("dir", &self.dir) + .field("flow", &self.key.flow) + .finish() + } +} + +impl TokenCache { + fn new(config: &OAuthConfig, options: &TokenCacheOptions) -> Result { + let key = CacheKey::new(config)?; + let dir = options.resolved_dir()?; + let lock_timeout = options.lock_timeout(); + prepare_cache_dir(&dir)?; + #[cfg(unix)] + let dir_owner = { + use std::os::unix::fs::MetadataExt; + let metadata = std::fs::metadata(&dir).map_err(|e| Error::Runtime { + message: format!( + "Failed to inspect OAuth token cache directory {}: {e}", + dir.display() + ), + })?; + metadata.uid() + }; + Ok(Self { + dir, + key, + lock_timeout, + #[cfg(unix)] + dir_owner, + }) + } + + fn record_path(&self) -> PathBuf { + let mut path = self.dir.clone(); + path.push(format!("{}.token.json", self.key.file_stem)); + path + } + + fn lock_path(&self) -> PathBuf { + let mut path = self.dir.clone(); + path.push(format!("{}.lock", self.key.file_stem)); + path + } + + /// Load the cached record, if one exists and passes hardening checks. + /// + /// Corrupt, truncated, unknown-version, or permission-invalid records + /// return an actionable error instead of being silently ignored or + /// deleted; the message names the file and how to recover. + async fn load(&self) -> Result> { + let path = self.record_path(); + let dir_owner = self.dir_owner_or_zero(); + tokio::task::spawn_blocking(move || read_record(&path, dir_owner)) + .await + .map_err(|e| Error::Runtime { + message: format!("Failed to join OAuth token cache read: {e}"), + })? + } + + #[cfg(unix)] + fn dir_owner_or_zero(&self) -> u32 { + self.dir_owner + } + + #[cfg(not(unix))] + fn dir_owner_or_zero(&self) -> u32 { + 0 + } + + /// Build a record from a token response, or `None` when the response + /// carries no refresh token (nothing may be persisted). + fn record_from_response(&self, response: &TokenResponse) -> Option { + let refresh_token = response.refresh_token.clone()?; + Some(CachedTokenRecord { + version: CACHE_RECORD_VERSION, + issuer_url: self.key.issuer_url.clone(), + client_id: self.key.client_id.clone(), + scopes: self.key.scopes.clone(), + flow: self.key.flow.to_string(), + client_auth: self.key.client_auth.to_string(), + refresh_token, + obtained_at: now_unix_secs(), + }) + } + + /// Atomically replace the cached record. + async fn store(&self, record: &CachedTokenRecord) -> Result<()> { + let path = self.record_path(); + let dir = self.dir.clone(); + let payload = serde_json::to_vec(record).map_err(|e| Error::Runtime { + message: format!("Failed to serialize OAuth token cache record: {e}"), + })?; + tokio::task::spawn_blocking(move || write_record(&dir, &path, &payload)) + .await + .map_err(|e| Error::Runtime { + message: format!("Failed to join OAuth token cache write: {e}"), + })? + } + + /// Delete the cached record. Returns whether a record was removed. + async fn delete(&self) -> Result { + let path = self.record_path(); + tokio::task::spawn_blocking(move || match std::fs::remove_file(&path) { + Ok(()) => Ok(true), + Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(false), + Err(e) => Err(Error::Runtime { + message: format!( + "Failed to remove OAuth token cache record {}: {e}", + path.display() + ), + }), + }) + .await + .map_err(|e| Error::Runtime { + message: format!("Failed to join OAuth token cache delete: {e}"), + })? + } + + /// Acquire the per-key cross-process lock, polling until the timeout. + async fn acquire_lock(&self) -> Result { + let path = self.lock_path(); + let timeout = self.lock_timeout; + tokio::task::spawn_blocking(move || { + use std::fs::OpenOptions; + let file = OpenOptions::new() + .create(true) + .read(true) + .write(true) + .truncate(false) + .open(&path) + .map_err(|e| Error::Runtime { + message: format!( + "Failed to open OAuth token cache lock file {}: {e}", + path.display() + ), + })?; + #[cfg(unix)] + set_owner_only_permissions(&file, &path); + let deadline = std::time::Instant::now() + timeout; + loop { + match file.try_lock_exclusive() { + Ok(true) => return Ok(LockGuard { file }), + Ok(false) => { + if std::time::Instant::now() >= deadline { + return Err(Error::Runtime { + message: format!( + "Timed out after {}s waiting for the OAuth token cache \ + lock at {}", + timeout.as_secs(), + path.display() + ), + }); + } + std::thread::sleep(LOCK_POLL_INTERVAL); + } + Err(error) => { + return Err(Error::Runtime { + message: format!( + "Failed to lock the OAuth token cache file {}: {error}", + path.display() + ), + }); + } + } + } + }) + .await + .map_err(|e| Error::Runtime { + message: format!("Failed to join OAuth token cache lock acquisition: {e}"), + })? + } + + /// Run the cross-process refresh critical section. + /// + /// Must be called with the in-process write lock held. Refresh grants are + /// serialized by the per-key cross-process lock; interactive flows (first + /// login or reauthentication) run outside it so a slow human-in-the-loop + /// flow never blocks refreshes in other processes. + pub(crate) async fn refresh_or_acquire( + &self, + source: &dyn TokenSource, + ) -> Result { + // Fast path under the lock: refresh from the durable record. + { + let _guard = self.acquire_lock().await?; + // Reread the record: another process may have rotated the refresh + // token since this process last looked. + if let Some(record) = self.load().await? { + match source.refresh_token(&record.refresh_token).await? { + RefreshResult::Refreshed(response) => { + self.store_if_refreshable(&response).await?; + return Ok(response); + } + RefreshResult::Reauthenticate => { + warn!( + "Cached OAuth refresh token was rejected; removing the cached \ + session before reauthenticating via {:?}", + source + ); + self.delete().await?; + } + RefreshResult::Unsupported => {} + } + } else { + debug!("No cached OAuth session; acquiring one via {:?}", source); + } + } + + // Interactive acquisition happens without the cross-process lock. + // Concurrent logins are independent sessions; the last store wins, + // which is the documented multi-session rule. + let response = source.fetch_token().await?; + let _guard = self.acquire_lock().await?; + self.store_if_refreshable(&response).await?; + Ok(response) + } + + /// Store a fresh login response (used by the eager `login` API). + /// + /// A successful login atomically replaces any prior session for this + /// cache identity: when the provider does not issue a refresh token, the + /// previous record is removed rather than left in place, so logging in + /// can never silently keep an earlier account's credential. + async fn store_login_response(&self, response: &TokenResponse) -> Result<()> { + let _guard = self.acquire_lock().await?; + match self.record_from_response(response) { + Some(record) => self.store(&record).await?, + None => { + self.delete().await?; + } + } + Ok(()) + } + + /// Replace the record when the response carries a refresh token. + /// + /// Called with the cross-process lock already held. Responses without a + /// refresh token are not persistable and are skipped. + async fn store_if_refreshable(&self, response: &TokenResponse) -> Result<()> { + if let Some(record) = self.record_from_response(response) { + self.store(&record).await?; + } else { + debug!( + "OAuth response did not include a refresh token; nothing to cache for {:?}", + self.key.flow + ); + } + Ok(()) + } +} + +fn prepare_cache_dir(dir: &Path) -> Result<()> { + if dir.is_dir() { + #[cfg(unix)] + { + use std::os::unix::fs::MetadataExt; + let metadata = std::fs::metadata(dir).map_err(|e| Error::Runtime { + message: format!( + "Failed to inspect OAuth token cache directory {}: {e}", + dir.display() + ), + })?; + let mode = metadata.mode(); + if mode & 0o077 != 0 { + return Err(Error::InvalidInput { + message: format!( + "OAuth token cache directory {} must not be accessible by group or \ + other users (mode {:o}); run `chmod 700` on it or choose a \ + private directory", + dir.display(), + mode & 0o777 + ), + }); + } + } + return Ok(()); + } + if dir.exists() { + return Err(Error::InvalidInput { + message: format!( + "OAuth token cache path {} exists and is not a directory", + dir.display() + ), + }); + } + let mut builder = std::fs::DirBuilder::new(); + builder.recursive(true); + #[cfg(unix)] + { + use std::os::unix::fs::DirBuilderExt; + builder.mode(0o700); + } + builder.create(dir).map_err(|e| Error::Runtime { + message: format!( + "Failed to create OAuth token cache directory {}: {e}", + dir.display() + ), + }) +} + +fn read_record(path: &Path, dir_owner: u32) -> Result> { + let metadata = match std::fs::symlink_metadata(path) { + Ok(metadata) => metadata, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(e) => { + return Err(Error::Runtime { + message: format!( + "Failed to read OAuth token cache record {}: {e}", + path.display() + ), + }); + } + }; + if !metadata.is_file() { + return Err(Error::InvalidInput { + message: format!( + "OAuth token cache record {} is not a regular file; refusing to use it. \ + Remove the file or call logout to clear it", + path.display() + ), + }); + } + #[cfg(unix)] + { + use std::os::unix::fs::MetadataExt; + let mode = metadata.mode(); + if mode & 0o077 != 0 { + return Err(Error::InvalidInput { + message: format!( + "OAuth token cache record {} must not be accessible by group or other \ + users (mode {:o}); run `chmod 600` on it or call logout to clear it", + path.display(), + mode & 0o777 + ), + }); + } + if dir_owner != 0 && metadata.uid() != dir_owner { + return Err(Error::InvalidInput { + message: format!( + "OAuth token cache record {} is owned by a different user; refusing to \ + use it. Remove the file or call logout to clear it", + path.display() + ), + }); + } + } + let payload = std::fs::read_to_string(path).map_err(|e| Error::Runtime { + message: format!( + "Failed to read OAuth token cache record {}: {e}", + path.display() + ), + })?; + let record: CachedTokenRecord = serde_json::from_str(&payload).map_err(|e| Error::Runtime { + message: format!( + "OAuth token cache record {} is corrupt ({e}); remove the file or call \ + logout to clear it", + path.display() + ), + })?; + if record.version != CACHE_RECORD_VERSION { + return Err(Error::Runtime { + message: format!( + "OAuth token cache record {} has unsupported version {}; expected {}. \ + Remove the file or call logout to clear it", + path.display(), + record.version, + CACHE_RECORD_VERSION + ), + }); + } + Ok(Some(record)) +} + +fn write_record(dir: &Path, path: &Path, payload: &[u8]) -> Result<()> { + let random_suffix: String = { + use rand::Rng; + let mut rng = rand::rng(); + (0..8) + .map(|_| format!("{:x}", rng.random_range(0..16u32))) + .collect() + }; + let mut temp_path = dir.to_path_buf(); + temp_path.push(format!( + "{}.tmp.{}.{}", + path.file_name() + .map(|name| name.to_string_lossy().to_string()) + .unwrap_or_default(), + std::process::id(), + random_suffix + )); + let write_attempt = || -> std::io::Result<()> { + let mut options = std::fs::OpenOptions::new(); + options.write(true).create_new(true); + let file = options.open(&temp_path)?; + #[cfg(unix)] + set_owner_only_permissions(&file, &temp_path); + use std::io::Write; + let mut writer = std::io::BufWriter::new(&file); + writer.write_all(payload)?; + writer.flush()?; + drop(writer); + file.sync_all()?; + std::fs::rename(&temp_path, path)?; + Ok(()) + }; + write_attempt().map_err(|e| { + let _ = std::fs::remove_file(&temp_path); + Error::Runtime { + message: format!( + "Failed to write OAuth token cache record {}: {e}", + path.display() + ), + } + }) +} + +#[cfg(unix)] +fn set_owner_only_permissions(file: &std::fs::File, path: &Path) { + use std::os::unix::fs::PermissionsExt; + if let Err(e) = file.set_permissions(std::fs::Permissions::from_mode(0o600)) { + debug!("Could not restrict permissions on {}: {e}", path.display()); + } +} + +/// Safe, non-secret view of a cached OAuth session, returned by +/// [`OAuthSession::status`] and [`OAuthSession::login`]. +#[derive(Clone, Debug, PartialEq)] +pub struct SessionStatus { + /// Whether a cached session exists that can obtain tokens without + /// interactive authentication. + /// + /// Because access tokens are not persisted, this is `true` exactly when a + /// refresh token is cached; the next connection refreshes with it rather + /// than opening a browser or device prompt. + 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) scope set of the cached session. + pub scopes: Vec, + + /// Flow that produced the cached session. + pub flow: String, + + /// When the cached session was obtained, as Unix seconds. + pub obtained_at: Option, +} + +/// Result of [`OAuthSession::logout`]. +#[derive(Clone, Debug, PartialEq)] +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: eager `login`, non-secret `status`, and +/// local `logout` for the persistent token cache. +/// +/// A session is built from the same [`OAuthConfig`](crate::remote::OAuthConfig) +/// used to connect (including its +/// [`token_cache`](crate::remote::OAuthConfig::token_cache) options). A +/// connection created with the same configuration shares the cache, so logging +/// in here prepares tokens for later processes without any database request. +/// +/// `login` always runs the configured interactive flow and replaces the +/// cached session (the most recent login wins; see the module documentation +/// about multiple accounts). `logout` removes only the local credential; it +/// does not revoke anything with the provider and does not terminate a +/// browser SSO session. +/// +/// # Example +/// +/// ``` +/// # use lancedb::remote::{OAuthConfig, OAuthFlow, OAuthSession, TokenCacheOptions}; +/// # fn example() -> lancedb::error::Result<()> { +/// let config = OAuthConfig { +/// issuer_url: "https://issuer.example.com".to_string(), +/// client_id: "my-app".to_string(), +/// client_secret: None, +/// scopes: vec!["openid".to_string()], +/// flow: OAuthFlow::DeviceCode, +/// refresh_buffer_secs: None, +/// token_cache: Some(TokenCacheOptions::new()), +/// }; +/// let session = OAuthSession::new(config)?; +/// # Ok(()) +/// # } +/// ``` +pub struct OAuthSession { + token_source: Box, + cache: TokenCache, +} + +impl std::fmt::Debug for OAuthSession { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("OAuthSession") + .field("cache", &self.cache) + .finish() + } +} + +impl OAuthSession { + /// Create a session manager for the given configuration. + /// + /// The configuration must enable [`TokenCacheOptions`] and use a flow + /// that supports persistent sessions (authorization code or device + /// authorization). + pub fn new(config: OAuthConfig) -> Result { + let cache_options = config + .token_cache + .clone() + .ok_or_else(|| Error::InvalidInput { + message: "OAuthSession requires OAuthConfig.token_cache to be set".to_string(), + })?; + if config.scopes.is_empty() { + return Err(Error::InvalidInput { + message: "At least one OAuth scope is required".to_string(), + }); + } + let token_source = crate::remote::oauth::build_token_source(&config)?; + let cache = TokenCache::new(&config, &cache_options)?; + Ok(Self { + token_source, + cache, + }) + } + + /// Eagerly run the configured authentication flow and store the session. + /// + /// Returns the resulting [`SessionStatus`]. 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`. + pub async fn login(&self) -> Result { + let response = self.token_source.fetch_token().await?; + self.cache.store_login_response(&response).await?; + self.status().await + } + + /// Report whether a matching cached session exists, with safe metadata. + /// + /// This never contacts the identity provider and never exposes token + /// values. + pub async fn status(&self) -> Result { + let record = self.cache.load().await?; + Ok(match record { + Some(record) => SessionStatus { + refreshable: true, + issuer_url: record.issuer_url, + client_id: record.client_id, + scopes: record.scopes, + flow: record.flow, + obtained_at: Some(record.obtained_at), + }, + None => SessionStatus { + refreshable: false, + issuer_url: self.cache.key.issuer_url.clone(), + client_id: self.cache.key.client_id.clone(), + scopes: self.cache.key.scopes.clone(), + flow: self.cache.key.flow.to_string(), + obtained_at: None, + }, + }) + } + + /// 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. + pub async fn logout(&self) -> Result { + let removed = self.cache.delete().await?; + Ok(SessionLogout { removed }) + } +} + +/// Resolve the persistent cache for a configuration, if one should exist. +/// +/// Returns `Ok(None)` for configurations without cache options and for the +/// client-credentials flow, which has no refresh token to persist (a debug +/// note is logged). The Azure managed-identity flow is rejected because +/// machine identity must not enter a user token cache. +pub(crate) fn token_cache_for_config(config: &OAuthConfig) -> Result>> { + let Some(options) = &config.token_cache else { + return Ok(None); + }; + if matches!(config.flow, OAuthFlow::AzureManagedIdentity { .. }) { + return Err(Error::InvalidInput { + message: "A persistent OAuth token cache cannot be used with the \ + AzureManagedIdentity flow; remove TokenCacheOptions to keep the \ + machine identity token in memory" + .to_string(), + }); + } + if matches!(config.flow, OAuthFlow::ClientCredentials) { + debug!( + "The client-credentials flow has no refresh token to persist; the OAuth \ + token cache is not used" + ); + return Ok(None); + } + Ok(Some(Arc::new(TokenCache::new(config, options)?))) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::{TcpListener, TcpStream}; + + use crate::remote::HeaderProvider; + use crate::remote::oauth::OAuthHeaderProvider; + use serial_test::serial; + + /// Temp directory that satisfies the cache hardening checks. CI runners + /// can create temp directories with group/other bits set, which the + /// private-directory validation correctly rejects. + fn cache_tempdir() -> tempfile::TempDir { + let dir = tempfile::tempdir().unwrap(); + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions(dir.path(), std::fs::Permissions::from_mode(0o700)).unwrap(); + } + dir + } + + fn device_config(cache_dir: &Path) -> OAuthConfig { + OAuthConfig { + issuer_url: "https://issuer.example.com".to_string(), + client_id: "client-id".to_string(), + client_secret: None, + scopes: vec!["openid".to_string()], + flow: OAuthFlow::DeviceCode, + refresh_buffer_secs: None, + token_cache: Some(TokenCacheOptions::new().cache_dir(cache_dir)), + } + } + + /// Stateful mock IdP covering discovery, device authorization, device + /// polling, client credentials, and refresh with strict rotation: a + /// refresh token that is not the currently issued one is rejected with + /// `invalid_grant`, which is exactly what real providers do on rotation. + struct MockIdp { + issuer_url: String, + device_authorizations: Arc, + refresh_attempts: Arc, + invalid_grant_rejections: Arc, + access_tokens_issued: Arc, + current_refresh: Arc>>, + fail_refreshes: Arc, + issue_refresh_tokens: Arc, + } + + impl MockIdp { + async fn start() -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let issuer_url = format!("http://{addr}"); + let server = Self { + issuer_url: issuer_url.clone(), + device_authorizations: Arc::new(AtomicUsize::new(0)), + refresh_attempts: Arc::new(AtomicUsize::new(0)), + invalid_grant_rejections: Arc::new(AtomicUsize::new(0)), + access_tokens_issued: Arc::new(AtomicUsize::new(0)), + current_refresh: Arc::new(std::sync::Mutex::new(None)), + fail_refreshes: Arc::new(AtomicBool::new(false)), + issue_refresh_tokens: Arc::new(AtomicBool::new(true)), + }; + let device_authorizations = Arc::clone(&server.device_authorizations); + let refresh_attempts = Arc::clone(&server.refresh_attempts); + let invalid_grant_rejections = Arc::clone(&server.invalid_grant_rejections); + let access_tokens_issued = Arc::clone(&server.access_tokens_issued); + let current_refresh = Arc::clone(&server.current_refresh); + let fail_refreshes = Arc::clone(&server.fail_refreshes); + let issue_refresh_tokens = Arc::clone(&server.issue_refresh_tokens); + + tokio::spawn(async move { + loop { + let Ok((mut stream, _)) = listener.accept().await else { + return; + }; + let device_authorizations = Arc::clone(&device_authorizations); + let refresh_attempts = Arc::clone(&refresh_attempts); + let invalid_grant_rejections = Arc::clone(&invalid_grant_rejections); + let access_tokens_issued = Arc::clone(&access_tokens_issued); + let current_refresh = Arc::clone(¤t_refresh); + let fail_refreshes = Arc::clone(&fail_refreshes); + let issue_refresh_tokens = Arc::clone(&issue_refresh_tokens); + tokio::spawn(async move { + let (request_line, body) = read_http_request(&mut stream).await; + if request_line.starts_with("GET /.well-known/openid-configuration ") { + let discovery = format!( + r#"{{"token_endpoint":"http://{addr}/token","device_authorization_endpoint":"http://{addr}/device"}}"# + ); + write_json_response(&mut stream, "200 OK", &discovery).await; + } else if request_line.starts_with("POST /device ") { + device_authorizations.fetch_add(1, Ordering::SeqCst); + let device = format!( + r#"{{"device_code":"device-code","user_code":"ABCD-EFGH","verification_uri":"http://{addr}/verify","expires_in":60,"interval":1}}"# + ); + write_json_response(&mut stream, "200 OK", &device).await; + } else if request_line.starts_with("POST /token ") { + if body.contains("grant_type=refresh_token") { + refresh_attempts.fetch_add(1, Ordering::SeqCst); + if fail_refreshes.load(Ordering::SeqCst) { + write_json_response( + &mut stream, + "503 Service Unavailable", + r#"{"error":"temporarily_unavailable"}"#, + ) + .await; + return; + } + let expected = current_refresh.lock().unwrap().clone(); + let matched = body + .split('&') + .find_map(|pair| pair.strip_prefix("refresh_token=")) + .map(|token| token.to_string()) + .zip(expected) + .is_some_and(|(offered, expected)| offered == expected); + if !matched { + invalid_grant_rejections.fetch_add(1, Ordering::SeqCst); + write_json_response( + &mut stream, + "400 Bad Request", + r#"{"error":"invalid_grant"}"#, + ) + .await; + return; + } + issue_access_token( + &mut stream, + &access_tokens_issued, + ¤t_refresh, + issue_refresh_tokens.load(Ordering::SeqCst), + ) + .await; + } else { + // Device polling or client credentials: issue + // a token and, for interactive grants, a fresh + // refresh token with strict rotation. + let grant_device = + body.contains("grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Adevice_code"); + if grant_device { + issue_access_token( + &mut stream, + &access_tokens_issued, + ¤t_refresh, + issue_refresh_tokens.load(Ordering::SeqCst), + ) + .await; + } else { + let token = format!( + r#"{{"access_token":"access-{}","expires_in":3600}}"#, + access_tokens_issued.fetch_add(1, Ordering::SeqCst) + 1 + ); + write_json_response(&mut stream, "200 OK", &token).await; + } + } + } else { + write_json_response(&mut stream, "404 Not Found", "{}").await; + } + }); + } + }); + server + } + + fn config(&self, cache_dir: &Path) -> OAuthConfig { + let mut config = device_config(cache_dir); + config.issuer_url = self.issuer_url.clone(); + config + } + } + + async fn issue_access_token( + stream: &mut TcpStream, + access_tokens_issued: &AtomicUsize, + current_refresh: &std::sync::Mutex>, + with_refresh: bool, + ) { + let number = access_tokens_issued.fetch_add(1, Ordering::SeqCst) + 1; + let token = if with_refresh { + *current_refresh.lock().unwrap() = Some(format!("refresh-{number}")); + format!( + r#"{{"access_token":"access-{number}","refresh_token":"refresh-{number}","expires_in":3600}}"# + ) + } else { + format!(r#"{{"access_token":"access-{number}","expires_in":3600}}"#) + }; + write_json_response(stream, "200 OK", &token).await; + } + + async fn read_http_request(stream: &mut TcpStream) -> (String, String) { + let mut buffer = Vec::new(); + let mut header_end = None; + while header_end.is_none() { + let mut chunk = [0; 1024]; + let read = stream.read(&mut chunk).await.unwrap(); + assert_ne!(read, 0, "connection closed before request headers"); + buffer.extend_from_slice(&chunk[..read]); + header_end = find_subsequence(&buffer, b"\r\n\r\n").map(|pos| pos + 4); + } + let header_end = header_end.unwrap(); + let headers = String::from_utf8_lossy(&buffer[..header_end]).to_string(); + let request_line = headers.lines().next().unwrap_or_default().to_string(); + let content_length = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().ok()) + .flatten() + }) + .unwrap_or(0); + while buffer.len() < header_end + content_length { + let mut chunk = [0; 1024]; + let read = stream.read(&mut chunk).await.unwrap(); + assert_ne!(read, 0, "connection closed before request body"); + buffer.extend_from_slice(&chunk[..read]); + } + let body = + String::from_utf8_lossy(&buffer[header_end..header_end + content_length]).to_string(); + (request_line, body) + } + + fn find_subsequence(haystack: &[u8], needle: &[u8]) -> Option { + haystack + .windows(needle.len()) + .position(|window| window == needle) + } + + async fn write_json_response(stream: &mut TcpStream, status: &str, body: &str) { + let response = format!( + "HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", + body.len() + ); + stream.write_all(response.as_bytes()).await.unwrap(); + } + + fn suppress_browser() { + // Point the OAuth browser helper at a no-op so device-flow tests never + // open a real browser window. + unsafe { std::env::set_var("LANCEDB_OAUTH_BROWSER", "/usr/bin/true") }; + } + + #[tokio::test] + #[serial] + async fn test_provider_reuses_cached_session_across_instances() { + suppress_browser(); + let dir = cache_tempdir(); + let idp = MockIdp::start().await; + + let first = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap(); + let headers = first.get_headers().await.unwrap(); + assert_eq!(headers.get("authorization").unwrap(), "Bearer access-1"); + + // A second, independent provider (simulating a second process) must + // refresh silently instead of starting another device flow. + let second = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap(); + let headers = second.get_headers().await.unwrap(); + assert_eq!(headers.get("authorization").unwrap(), "Bearer access-2"); + + assert_eq!(idp.device_authorizations.load(Ordering::SeqCst), 1); + assert_eq!(idp.refresh_attempts.load(Ordering::SeqCst), 1); + + let cache = crate::remote::token_cache::TokenCache::new( + &idp.config(dir.path()), + &TokenCacheOptions::new().cache_dir(dir.path()), + ) + .unwrap(); + let record = cache.load().await.unwrap().unwrap(); + assert_eq!(record.refresh_token, "refresh-2"); + } + + #[tokio::test] + #[serial] + async fn test_concurrent_providers_serialize_rotation() { + suppress_browser(); + let dir = cache_tempdir(); + let idp = MockIdp::start().await; + + let priming = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap(); + priming.get_headers().await.unwrap(); + assert_eq!(idp.device_authorizations.load(Ordering::SeqCst), 1); + + let provider_a = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap(); + let provider_b = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap(); + let (headers_a, headers_b) = + tokio::join!(provider_a.get_headers(), provider_b.get_headers()); + let token_a = headers_a.unwrap().remove("authorization").unwrap(); + let token_b = headers_b.unwrap().remove("authorization").unwrap(); + assert!( + { + let mut tokens = [token_a, token_b]; + tokens.sort(); + tokens + } == ["Bearer access-2".to_string(), "Bearer access-3".to_string()], + "each provider must observe its own refreshed token" + ); + + // Rotation raced would produce an invalid_grant and a second device + // flow; the lock prevents both. + assert_eq!(idp.invalid_grant_rejections.load(Ordering::SeqCst), 0); + assert_eq!(idp.device_authorizations.load(Ordering::SeqCst), 1); + assert_eq!(idp.refresh_attempts.load(Ordering::SeqCst), 2); + + let cache = crate::remote::token_cache::TokenCache::new( + &idp.config(dir.path()), + &TokenCacheOptions::new().cache_dir(dir.path()), + ) + .unwrap(); + let record = cache.load().await.unwrap().unwrap(); + assert_eq!(record.refresh_token, "refresh-3"); + } + + #[tokio::test] + #[serial] + async fn test_transient_refresh_failure_retains_record() { + suppress_browser(); + let dir = cache_tempdir(); + let idp = MockIdp::start().await; + + let priming = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap(); + priming.get_headers().await.unwrap(); + + idp.fail_refreshes.store(true, Ordering::SeqCst); + let second = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap(); + let error = second.get_headers().await.unwrap_err(); + assert!(error.to_string().contains("503")); + + let session = OAuthSession::new(idp.config(dir.path())).unwrap(); + let status = session.status().await.unwrap(); + assert!( + status.refreshable, + "transient failures must keep the record" + ); + } + + #[tokio::test] + #[serial] + async fn test_invalid_grant_deletes_record_and_reauthenticates() { + suppress_browser(); + let dir = cache_tempdir(); + let idp = MockIdp::start().await; + + let priming = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap(); + priming.get_headers().await.unwrap(); + + // Simulate a revoked refresh token by replacing the record with one + // the provider never issued. + let cache = crate::remote::token_cache::TokenCache::new( + &idp.config(dir.path()), + &TokenCacheOptions::new().cache_dir(dir.path()), + ) + .unwrap(); + let mut record = cache.load().await.unwrap().unwrap(); + record.refresh_token = "revoked-refresh".to_string(); + cache.store(&record).await.unwrap(); + + let second = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap(); + let headers = second.get_headers().await.unwrap(); + assert_eq!(headers.get("authorization").unwrap(), "Bearer access-2"); + + assert_eq!(idp.invalid_grant_rejections.load(Ordering::SeqCst), 1); + assert_eq!(idp.device_authorizations.load(Ordering::SeqCst), 2); + + let record = cache.load().await.unwrap().unwrap(); + assert_eq!(record.refresh_token, "refresh-2"); + } + + #[tokio::test] + #[serial] + async fn test_session_login_status_logout_lifecycle() { + suppress_browser(); + let dir = cache_tempdir(); + let idp = MockIdp::start().await; + + let session = OAuthSession::new(idp.config(dir.path())).unwrap(); + let status = session.status().await.unwrap(); + assert!(!status.refreshable); + + let status = session.login().await.unwrap(); + assert!(status.refreshable); + assert_eq!(status.issuer_url, idp.issuer_url); + assert_eq!(status.client_id, "client-id"); + assert_eq!(status.scopes, vec!["openid".to_string()]); + assert_eq!(status.flow, "device_code"); + assert!(status.obtained_at.is_some()); + + // An independent session manager sees the same cached login. + let other = OAuthSession::new(idp.config(dir.path())).unwrap(); + assert!(other.status().await.unwrap().refreshable); + + assert!(session.logout().await.unwrap().removed); + assert!(!session.logout().await.unwrap().removed); + assert!(!other.status().await.unwrap().refreshable); + assert_eq!(idp.device_authorizations.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + #[serial] + async fn test_login_without_refresh_token_clears_prior_record() { + suppress_browser(); + let dir = cache_tempdir(); + let idp = MockIdp::start().await; + + let session = OAuthSession::new(idp.config(dir.path())).unwrap(); + session.login().await.unwrap(); + assert!(session.status().await.unwrap().refreshable); + + // A provider that stops issuing refresh tokens (for example a login + // without offline_access) must not leave the earlier account behind. + idp.issue_refresh_tokens.store(false, Ordering::SeqCst); + let status = session.login().await.unwrap(); + assert!(!status.refreshable); + assert!(!session.status().await.unwrap().refreshable); + } + + #[tokio::test] + async fn test_client_credentials_with_cache_stays_memory_only() { + let dir = cache_tempdir(); + let idp = MockIdp::start().await; + + let mut config = idp.config(dir.path()); + config.flow = OAuthFlow::ClientCredentials; + config.client_secret = Some("secret".to_string()); + + let provider = OAuthHeaderProvider::new(config).unwrap(); + let headers = provider.get_headers().await.unwrap(); + assert_eq!(headers.get("authorization").unwrap(), "Bearer access-1"); + assert_eq!(dir.path().read_dir().unwrap().count(), 0); + } + + #[test] + fn test_managed_identity_with_cache_is_rejected() { + let dir = cache_tempdir(); + let mut config = device_config(dir.path()); + config.flow = OAuthFlow::AzureManagedIdentity { client_id: None }; + + let err = OAuthHeaderProvider::new(config).unwrap_err(); + assert!( + matches!(err, Error::InvalidInput { message } if message.contains("AzureManagedIdentity")) + ); + } + + #[tokio::test] + #[serial] + async fn test_provider_debug_and_status_reveal_no_secrets() { + suppress_browser(); + let dir = cache_tempdir(); + let idp = MockIdp::start().await; + + let provider = OAuthHeaderProvider::new(idp.config(dir.path())).unwrap(); + let headers = provider.get_headers().await.unwrap(); + assert_eq!(headers.get("authorization").unwrap(), "Bearer access-1"); + let debug = format!("{provider:?}"); + assert!(!debug.contains("access-1")); + assert!(!debug.contains("refresh-1")); + + let session = OAuthSession::new(idp.config(dir.path())).unwrap(); + let status = session.status().await.unwrap(); + assert!(!format!("{status:?}").contains("refresh-")); + } + + #[test] + fn test_cache_key_canonicalizes_scopes_and_issuer() { + let mut config = device_config(Path::new("/tmp/cache")); + config.issuer_url = "https://issuer.example.com/".to_string(); + config.scopes = vec!["b".to_string(), " a ".to_string(), "a".to_string()]; + let key = CacheKey::new(&config).unwrap(); + assert_eq!(key.issuer_url, "https://issuer.example.com"); + assert_eq!(key.scopes, vec!["a".to_string(), "b".to_string()]); + + config.issuer_url = "https://issuer.example.com".to_string(); + let canonical = CacheKey::new(&config).unwrap(); + assert_eq!(canonical.file_stem, key.file_stem); + } + + #[test] + fn test_cache_key_separates_identity_dimensions() { + let base = device_config(Path::new("/tmp/cache")); + let base_key = CacheKey::new(&base).unwrap(); + + let mut other = base.clone(); + other.client_id = "other-client".to_string(); + assert_ne!(CacheKey::new(&other).unwrap().file_stem, base_key.file_stem); + + let mut other = base.clone(); + other.issuer_url = "https://other.example.com".to_string(); + assert_ne!(CacheKey::new(&other).unwrap().file_stem, base_key.file_stem); + + let mut other = base.clone(); + other.scopes = vec!["profile".to_string()]; + assert_ne!(CacheKey::new(&other).unwrap().file_stem, base_key.file_stem); + + let mut other = base.clone(); + other.flow = OAuthFlow::AuthorizationCode(Default::default()); + assert_ne!(CacheKey::new(&other).unwrap().file_stem, base_key.file_stem); + + let mut other = base.clone(); + other.client_secret = Some("secret".to_string()); + assert_ne!(CacheKey::new(&other).unwrap().file_stem, base_key.file_stem); + } + + #[test] + fn test_cache_key_contains_no_secret_material() { + let mut config = device_config(Path::new("/tmp/cache")); + config.client_secret = Some("super-secret-value".to_string()); + let key = CacheKey::new(&config).unwrap(); + assert!(!key.file_stem.contains("super-secret-value")); + assert_eq!(key.file_stem.len(), 64); + } + + #[test] + fn test_flow_key_rejects_non_persistent_flows() { + let mut config = device_config(Path::new("/tmp/cache")); + config.flow = OAuthFlow::ClientCredentials; + let err = CacheKey::new(&config).unwrap_err(); + assert!( + matches!(err, Error::InvalidInput { message } if message.contains("not supported")) + ); + + let mut config = device_config(Path::new("/tmp/cache")); + config.flow = OAuthFlow::AzureManagedIdentity { client_id: None }; + assert!(CacheKey::new(&config).is_err()); + } + + #[test] + fn test_token_cache_options_defaults() { + let options = TokenCacheOptions::new(); + assert!(options.cache_dir.is_none()); + assert!(options.lock_timeout_secs.is_none()); + assert_eq!(options.lock_timeout(), Duration::from_secs(30)); + + let options = options.cache_dir("/tmp/x").lock_timeout_secs(5); + assert_eq!(options.cache_dir.as_deref(), Some(Path::new("/tmp/x"))); + assert_eq!(options.lock_timeout(), Duration::from_secs(5)); + } + + #[test] + fn test_token_cache_options_rejects_empty_dir() { + let options = TokenCacheOptions::new().cache_dir(""); + assert!(matches!( + options.resolved_dir(), + Err(Error::InvalidInput { message }) if message.contains("must not be empty") + )); + } + + #[tokio::test] + async fn test_session_lifecycle_without_cache_entry() { + let dir = cache_tempdir(); + let session = OAuthSession::new(device_config(dir.path())).unwrap(); + + let status = session.status().await.unwrap(); + assert!(!status.refreshable); + assert_eq!(status.issuer_url, "https://issuer.example.com"); + assert_eq!(status.client_id, "client-id"); + assert_eq!(status.scopes, vec!["openid".to_string()]); + assert_eq!(status.flow, "device_code"); + assert_eq!(status.obtained_at, None); + + let logout = session.logout().await.unwrap(); + assert!(!logout.removed); + } + + #[tokio::test] + async fn test_record_round_trip_and_redaction() { + let dir = cache_tempdir(); + let cache = TokenCache::new( + &device_config(dir.path()), + &TokenCacheOptions::new().cache_dir(dir.path()), + ) + .unwrap(); + let response = TokenResponse { + access_token: "access-token".to_string(), + refresh_token: Some("refresh-token".to_string()), + expires_in: Some(3600), + token_type: Some("Bearer".to_string()), + }; + let record = cache.record_from_response(&response).unwrap(); + let debug = format!("{record:?}"); + assert!(!debug.contains("refresh-token")); + assert!(debug.contains("")); + // No access-token material is persisted. + let json = serde_json::to_string(&record).unwrap(); + assert!(!json.contains("access-token")); + + cache.store(&record).await.unwrap(); + let loaded = cache.load().await.unwrap().unwrap(); + assert_eq!(loaded.refresh_token, "refresh-token"); + assert_eq!(loaded.version, CACHE_RECORD_VERSION); + + // The on-disk file must not be group/other readable and must not be a symlink target. + let path = cache.record_path(); + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + let mode = std::fs::metadata(&path).unwrap().permissions().mode(); + assert_eq!(mode & 0o077, 0); + } + assert!(std::fs::symlink_metadata(&path).unwrap().is_file()); + + assert!(cache.delete().await.unwrap()); + assert!(cache.load().await.unwrap().is_none()); + assert!(!cache.delete().await.unwrap()); + } + + #[cfg(unix)] + #[tokio::test] + async fn test_record_rejects_symlink() { + let dir = cache_tempdir(); + let cache = TokenCache::new( + &device_config(dir.path()), + &TokenCacheOptions::new().cache_dir(dir.path()), + ) + .unwrap(); + let record = cache + .record_from_response(&TokenResponse { + access_token: "a".to_string(), + refresh_token: Some("r".to_string()), + expires_in: None, + token_type: None, + }) + .unwrap(); + cache.store(&record).await.unwrap(); + let path = cache.record_path(); + let target = dir.path().join("evil.json"); + std::fs::write(&target, "{}").unwrap(); + std::fs::remove_file(&path).unwrap(); + std::os::unix::fs::symlink(&target, &path).unwrap(); + + let err = cache.load().await.unwrap_err(); + assert!( + matches!(err, Error::InvalidInput { message } if message.contains("not a regular file")) + ); + } + + #[cfg(unix)] + #[tokio::test] + async fn test_record_rejects_world_readable_file() { + let dir = cache_tempdir(); + let cache = TokenCache::new( + &device_config(dir.path()), + &TokenCacheOptions::new().cache_dir(dir.path()), + ) + .unwrap(); + let path = cache.record_path(); + std::fs::write(&path, "{}").unwrap(); + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o644)).unwrap(); + + let err = cache.load().await.unwrap_err(); + assert!( + matches!(err, Error::InvalidInput { message } if message.contains("group or other")) + ); + } + + /// Write a record file that passes the permission hardening checks. + fn write_record_file(path: &Path, contents: &str) { + std::fs::write(path, contents).unwrap(); + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600)).unwrap(); + } + } + + #[tokio::test] + async fn test_record_rejects_unknown_version_and_corruption() { + let dir = cache_tempdir(); + let cache = TokenCache::new( + &device_config(dir.path()), + &TokenCacheOptions::new().cache_dir(dir.path()), + ) + .unwrap(); + let path = cache.record_path(); + + // A complete, well-formed record with an unknown schema version. + let record = cache + .record_from_response(&TokenResponse { + access_token: "a".to_string(), + refresh_token: Some("r".to_string()), + expires_in: None, + token_type: None, + }) + .unwrap(); + let mut json = serde_json::to_value(&record).unwrap(); + json["version"] = serde_json::json!(99); + write_record_file(&path, &json.to_string()); + let err = cache.load().await.unwrap_err(); + assert!( + matches!(err, Error::Runtime { message } if message.contains("unsupported version")) + ); + + write_record_file(&path, r#"{"version":1,"issuer_url":"x""#); + let err = cache.load().await.unwrap_err(); + assert!(matches!(err, Error::Runtime { message } if message.contains("corrupt"))); + + write_record_file(&path, ""); + assert!( + cache + .load() + .await + .unwrap_err() + .to_string() + .contains("corrupt") + ); + } + + #[cfg(unix)] + #[tokio::test] + async fn test_cache_dir_rejects_open_permissions() { + let dir = cache_tempdir(); + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions(dir.path(), std::fs::Permissions::from_mode(0o755)).unwrap(); + let err = TokenCache::new( + &device_config(dir.path()), + &TokenCacheOptions::new().cache_dir(dir.path()), + ) + .unwrap_err(); + assert!(matches!(err, Error::InvalidInput { message } if message.contains("chmod 700"))); + } + + #[tokio::test] + async fn test_lock_serializes_and_releases() { + let dir = cache_tempdir(); + let cache = TokenCache::new( + &device_config(dir.path()), + &TokenCacheOptions::new().cache_dir(dir.path()), + ) + .unwrap(); + let guard = cache.acquire_lock().await.unwrap(); + + let contender = { + let dir = dir.path().to_path_buf(); + let cache2 = TokenCache::new( + &device_config(&dir), + &TokenCacheOptions::new() + .cache_dir(&dir) + .lock_timeout_secs(1), + ) + .unwrap(); + tokio::time::timeout(Duration::from_millis(300), cache2.acquire_lock()).await + }; + assert!(contender.is_err(), "second acquire must block while held"); + drop(guard); + + let cache3 = TokenCache::new( + &device_config(dir.path()), + &TokenCacheOptions::new().cache_dir(dir.path()), + ) + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), cache3.acquire_lock()) + .await + .expect("lock re-acquirable after release") + .unwrap(); + } + + #[test] + fn test_token_cache_for_config_gating() { + let dir = cache_tempdir(); + assert!( + token_cache_for_config(&device_config(dir.path())) + .unwrap() + .is_some() + ); + + let mut config = device_config(dir.path()); + config.token_cache = None; + assert!(token_cache_for_config(&config).unwrap().is_none()); + + let mut config = device_config(dir.path()); + config.flow = OAuthFlow::ClientCredentials; + assert!(token_cache_for_config(&config).unwrap().is_none()); + + let mut config = device_config(dir.path()); + config.flow = OAuthFlow::AzureManagedIdentity { client_id: None }; + let err = token_cache_for_config(&config).unwrap_err(); + assert!( + matches!(err, Error::InvalidInput { message } if message.contains("AzureManagedIdentity")) + ); + } + + #[test] + fn test_oauth_session_requires_cache_options() { + let mut config = device_config(Path::new("/tmp/cache")); + config.token_cache = None; + let err = OAuthSession::new(config).unwrap_err(); + assert!(matches!(err, Error::InvalidInput { message } if message.contains("token_cache"))); + } + + #[tokio::test] + async fn test_store_skips_responses_without_refresh_token() { + let dir = cache_tempdir(); + let cache = TokenCache::new( + &device_config(dir.path()), + &TokenCacheOptions::new().cache_dir(dir.path()), + ) + .unwrap(); + let response = TokenResponse { + access_token: "access-token".to_string(), + refresh_token: None, + expires_in: Some(3600), + token_type: None, + }; + assert!(cache.record_from_response(&response).is_none()); + cache.store_if_refreshable(&response).await.unwrap(); + assert!(cache.load().await.unwrap().is_none()); + } + + #[test] + fn test_session_status_debug_has_no_secrets() { + let status = SessionStatus { + refreshable: true, + issuer_url: "https://issuer.example.com".to_string(), + client_id: "client-id".to_string(), + scopes: vec!["openid".to_string()], + flow: "device_code".to_string(), + obtained_at: Some(100), + }; + let debug = format!("{status:?}"); + assert!(!debug.contains("refresh-token")); + } +}