feat: add persistent OAuth token cache and session APIs (#4182)

Stacked on #4173 (diff includes it until that merges; will rebase after). Addresses the token-cache part of [Colin's review](https://github.com/lancedb/lancedb/pull/4173#issuecomment-5674048100). Adds an explicit, opt-in persistent OAuth token cache shared by Rust, Python, and Node clients, plus `login` / `status` / `logout` session APIs, so short-lived processes (CLIs, scripts, notebooks) reuse one session instead of restarting a browser or device flow on every start. - **Opt-in and minimal**: existing callers stay memory-only and lazy. Only refresh tokens are persisted (never access tokens, never client secrets), so there are no local token-expiry decisions to get wrong when clocks move. Each process start performs one silent refresh grant. - **Hardened file backend**: private directory (`0700`), per-record files (`0600`), owner validation, symlink rejection, and atomic `rename` replacement. Corrupt, truncated, unknown-version, or permission-invalid records fail with actionable errors naming the file. Native keyring backends were evaluated (keyring crate routes Linux through D-Bus/zbus: heavy deps, headless/CI flakiness) and are deferred; the file store is the explicit opt-in, not a downgrade from a keyring. - **Cache key**: SHA-256 of the canonical identity (issuer, client ID, sorted/de-duplicated scopes, flow, public/confidential), so no secret appears in a filename and distinct identities never collide. Versioned record schema (`version: 1`). One record per identity: last login wins, documented. - **Cross-process rotation locking**: per-key `fs4` file lock (`flock` / `LockFileEx`) around the refresh critical section — acquire, reread the durable record, refresh exactly once, atomically store the rotated refresh token, release. The OS releases locks on process death, so crashes cannot strand stale locks. Only confirmed `invalid_grant`/`invalid_token` deletes a record and reauthenticates; transport, 5xx, 429, and parse failures retain it. - **Session APIs**: `OAuthSession::login/status/logout` in Rust, `lancedb.remote.OAuthSession` (async) in Python, `OAuthSession` class in Node. `status` returns non-secret metadata only. `logout` removes only the local credential — provider revocation (RFC 7009) is a deliberate follow-up, and local logout never terminates browser SSO. Azure managed identity is rejected for persistence (machine identity stays in memory); client credentials have nothing refreshable to persist and stay memory-only. - No CLI binary exists in this repo, so this ships library APIs plus doc examples in all three languages. Tests: Rust unit + mock-IdP integration (cache-key canonicalization/separation, record versioning/corruption/truncation/symlink/owner/perms, lock serialization + release, two concurrent providers proving no `invalid_grant` and correct rotation, transient-failure retention, `invalid_grant` delete + reauthenticate, login/status/logout lifecycle, client-credentials no-op, IMDS rejection, secret redaction); Python lifecycle + a true two-subprocess cross-process reuse test (second process refreshes once, never hits the device endpoint); Node lifecycle + device-flow login test. Local builds were skipped in development; CI validates all bindings. --------- Co-authored-by: Xuanwo <github@xuanwo.io>
2026-09-21 20:45:57 +00:00 · 2026-09-16 02:01:10 +08:00
parent 575286922b
commit 2f88b71c21
25 changed files with 5464 additions and 97 deletions
@@ -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<void> {
+    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<void>((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();
+  }
+}