diff --git a/web-admin/src/features/edit-session/updateEditSessionDevJWT.spec.ts b/web-admin/src/features/edit-session/updateEditSessionDevJWT.spec.ts new file mode 100644 index 000000000000..d4f060a909ea --- /dev/null +++ b/web-admin/src/features/edit-session/updateEditSessionDevJWT.spec.ts @@ -0,0 +1,184 @@ +import { adminServiceGetDeployment } from "@rilldata/web-admin/client"; +import { + createMockUserJwtRefresher, + createUpdateEditSessionDevJWT, + editSessionRuntimeAuth, +} from "@rilldata/web-admin/features/edit-session/updateEditSessionDevJWT"; +import { + selectedMockUserJWT, + selectedMockUserStore, +} from "@rilldata/web-common/features/dashboards/granular-access-policies/stores"; +import { createQueryClient } from "@rilldata/web-common/lib/svelte-query/globalQueryClient"; +import { RUNTIME_ACCESS_TOKEN_DEFAULT_TTL } from "@rilldata/web-common/runtime-client/constants"; +import { + RuntimeClient, + type JwtRefresher, +} from "@rilldata/web-common/runtime-client/v2"; +import { get } from "svelte/store"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +vi.mock("@rilldata/web-admin/client", async (importOriginal) => ({ + ...(await importOriginal()), + adminServiceGetDeployment: vi.fn(), +})); + +describe("edit session View As", () => { + const fetchMock = vi.fn(); + const getDeployment = vi.mocked(adminServiceGetDeployment); + const mockUser = { email: "viewer@example.com" }; + + beforeEach(() => { + vi.useFakeTimers(); + fetchMock.mockReset().mockImplementation(() => + Promise.resolve( + new Response(new Uint8Array(0), { + status: 200, + headers: { "content-type": "application/proto" }, + }), + ), + ); + vi.stubGlobal("fetch", fetchMock); + getDeployment.mockReset(); + selectedMockUserStore.set(null); + selectedMockUserJWT.set(null); + }); + + afterEach(() => { + vi.useRealTimers(); + vi.unstubAllGlobals(); + }); + + function createClient() { + return new RuntimeClient({ + host: "http://localhost:9009", + instanceId: "test-instance", + jwt: "edit-session-jwt", + }); + } + + /** Sends a request; ConnectRPC may reject the minimal mock response, which does not matter here. */ + async function ping(client: RuntimeClient) { + try { + await client.runtimeService.ping({}); + } catch { + // no-op + } + } + + function sentAuthHeader(index: number) { + const headers = fetchMock.mock.calls[index]?.[1]?.headers as + | Headers + | undefined; + return headers?.get("Authorization") ?? null; + } + + it("installs the mock user's JWT and restores the user's when View As ends", async () => { + getDeployment.mockResolvedValueOnce({ accessToken: "mock-token" }); + const client = createClient(); + const queryClient = createQueryClient(); + const updateDevJWT = createUpdateEditSessionDevJWT( + "deployment-1", + "edit-session-jwt", + ); + + await updateDevJWT(queryClient, client, mockUser); + expect(client.getJwt()).toBe("mock-token"); + expect(client.authContext).toBe("mock"); + expect(get(selectedMockUserStore)).toBe(mockUser); + expect(get(selectedMockUserJWT)).toBe("mock-token"); + + await updateDevJWT(queryClient, client, null); + expect(client.getJwt()).toBe("edit-session-jwt"); + expect(client.authContext).toBe("user"); + expect(get(selectedMockUserStore)).toBeNull(); + expect(get(selectedMockUserJWT)).toBeNull(); + }); + + it("reissues the mock user's JWT with the same attributes", async () => { + getDeployment.mockResolvedValueOnce({ accessToken: "mock-token-2" }); + const abort = new AbortController(); + + const refresh = createMockUserJwtRefresher("deployment-1", mockUser); + await expect(refresh(abort.signal)).resolves.toBe("mock-token-2"); + + expect(getDeployment).toHaveBeenCalledTimes(1); + const [deploymentId, body, signal] = getDeployment.mock.calls[0]; + expect(deploymentId).toBe("deployment-1"); + expect(body.attributes).toMatchObject({ + email: "viewer@example.com", + domain: "example.com", + }); + expect(body.accessTokenTtlSeconds).toBe( + RUNTIME_ACCESS_TOKEN_DEFAULT_TTL / 1000, + ); + expect(signal).toBe(abort.signal); + expect(get(selectedMockUserJWT)).toBe("mock-token-2"); + }); + + it("keeps the mock credentials when a user JWT refresh lands after View As started", async () => { + getDeployment + .mockResolvedValueOnce({ accessToken: "mock-token" }) + .mockResolvedValueOnce({ accessToken: "mock-token-2" }); + const client = createClient(); + const queryClient = createQueryClient(); + let resolveUserRefresh!: (jwt: string) => void; + const userRefreshJwt: JwtRefresher = vi.fn( + () => + new Promise((resolve) => { + resolveUserRefresh = resolve; + }), + ); + const mockRefreshJwt = createMockUserJwtRefresher("deployment-1", mockUser); + // Plays the edit layout and RuntimeProvider: the client gets whatever the layout + // derives from the GetProject JWT and the View As selection + function forwardToProvider(userJwt: string) { + const auth = editSessionRuntimeAuth({ + userJwt, + userRefreshJwt, + mockUser: get(selectedMockUserStore), + mockJwt: get(selectedMockUserJWT), + mockRefreshJwt, + }); + client.updateJwt(auth.jwt, auth.authContext); + client.setJwtRefresher(auth.refreshJwt); + } + forwardToProvider("edit-session-jwt"); + + // The user's JWT expires and a request starts an on-demand refresh + vi.advanceTimersByTime(RUNTIME_ACCESS_TOKEN_DEFAULT_TTL); + const request = ping(client); + await vi.advanceTimersByTimeAsync(1); + expect(userRefreshJwt).toHaveBeenCalledTimes(1); + + // View As starts while that refresh is still in flight + const updateDevJWT = createUpdateEditSessionDevJWT( + "deployment-1", + "edit-session-jwt", + ); + await updateDevJWT(queryClient, client, mockUser); + forwardToProvider("edit-session-jwt"); + await request; + expect(sentAuthHeader(0)).toBe("Bearer mock-token"); + + // The user refresh completes and the layout forwards the new user JWT; the mock credentials stay + resolveUserRefresh("edit-session-jwt-2"); + await vi.advanceTimersByTimeAsync(1); + forwardToProvider("edit-session-jwt-2"); + expect(client.getJwt()).toBe("mock-token"); + expect(client.authContext).toBe("mock"); + + // When the mock token expires, it is the mock issuer that is asked, not the user's + vi.advanceTimersByTime(RUNTIME_ACCESS_TOKEN_DEFAULT_TTL); + await ping(client); + expect(getDeployment).toHaveBeenCalledTimes(2); + expect(userRefreshJwt).toHaveBeenCalledTimes(1); + expect(client.getJwt()).toBe("mock-token-2"); + expect(sentAuthHeader(1)).toBe("Bearer mock-token-2"); + + // View As ends: the latest user JWT applies + await updateDevJWT(queryClient, client, null); + forwardToProvider("edit-session-jwt-2"); + expect(client.getJwt()).toBe("edit-session-jwt-2"); + expect(client.authContext).toBe("user"); + }); +}); diff --git a/web-admin/src/features/edit-session/updateEditSessionDevJWT.ts b/web-admin/src/features/edit-session/updateEditSessionDevJWT.ts index c94b182d044a..2aea541daef1 100644 --- a/web-admin/src/features/edit-session/updateEditSessionDevJWT.ts +++ b/web-admin/src/features/edit-session/updateEditSessionDevJWT.ts @@ -8,7 +8,11 @@ import { invalidateCanvasQueries, removeInactiveResourceQueries, } from "@rilldata/web-common/runtime-client/invalidation"; -import type { RuntimeClient } from "@rilldata/web-common/runtime-client/v2"; +import type { + AuthContext, + JwtRefresher, + RuntimeClient, +} from "@rilldata/web-common/runtime-client/v2"; import type { QueryClient } from "@tanstack/svelte-query"; import { adminServiceGetDeployment } from "@rilldata/web-admin/client"; import { getDomain } from "@rilldata/web-admin/features/projects/user-management/selectors.ts"; @@ -29,28 +33,14 @@ export function createUpdateEditSessionDevJWT( const prevMockUser = get(selectedMockUserStore); selectedMockUserStore.set(mockUser); + // The edit layout derives the provider's credentials from the stores set here (see editSessionRuntimeAuth); + // the client is updated directly as well so the invalidations below already run with the new token if (mockUser === null) { selectedMockUserJWT.set(null); client.updateJwt(editSessionJwt, "user"); } else { try { - const { name, email, groups, admin, ...customAttributes } = mockUser; - - const { accessToken } = await adminServiceGetDeployment(deploymentId, { - attributes: { - email, - ...(email ? { domain: getDomain(email) } : {}), - name: name || "Mock User", - admin: !!admin, - groups: groups || [], - ...customAttributes, - }, - // Make sure to match TTL, GetDeployment defaults to 24-hours - accessTokenTtlSeconds: RUNTIME_ACCESS_TOKEN_DEFAULT_TTL / 1000, // Seconds param vs milliseconds constant - }); - - if (!accessToken) throw new Error("No JWT returned"); - + const accessToken = await issueMockUserJwt(deploymentId, mockUser); selectedMockUserJWT.set(accessToken); client.updateJwt(accessToken, "mock"); } catch { @@ -69,3 +59,78 @@ export function createUpdateEditSessionDevJWT( return invalidateCanvasQueries(queryClient, client.instanceId); }; } + +/** + * Reissues the mock user's JWT when a request finds it expired, e.g. after a laptop sleep, + * so the preview does not wait for a token that never comes. The mock token has the same TTL as the user's. + */ +export function createMockUserJwtRefresher( + deploymentId: string, + mockUser: MockUser, +): JwtRefresher { + return async (signal) => { + const accessToken = await issueMockUserJwt(deploymentId, mockUser, signal); + selectedMockUserJWT.set(accessToken); + return accessToken; + }; +} + +/** + * Picks the credentials the edit layout's RuntimeProvider runs with: the mock user's JWT and refresher while + * View As is active, the user's otherwise. The provider must follow the View As selection: if it kept forwarding + * the user's JWT, the next GetProject update (the 15-minute poll, a focus refetch, or an on-demand refresh that was + * still in flight when View As started) would push the user's token back onto the client with the "user" context, + * and the preview would show the user's data while View As stays selected. + */ +export function editSessionRuntimeAuth(opts: { + userJwt: string | undefined; + userRefreshJwt: JwtRefresher | undefined; + mockUser: MockUser | null; + mockJwt: string | null; + mockRefreshJwt: JwtRefresher | undefined; +}): { + jwt: string | undefined; + authContext: AuthContext; + refreshJwt: JwtRefresher | undefined; +} { + if (opts.mockUser && opts.mockJwt) { + return { + jwt: opts.mockJwt, + authContext: "mock", + refreshJwt: opts.mockRefreshJwt, + }; + } + return { + jwt: opts.userJwt, + authContext: "user", + refreshJwt: opts.userRefreshJwt, + }; +} + +async function issueMockUserJwt( + deploymentId: string, + mockUser: MockUser, + signal?: AbortSignal, +): Promise { + const { name, email, groups, admin, ...customAttributes } = mockUser; + + const { accessToken } = await adminServiceGetDeployment( + deploymentId, + { + attributes: { + email, + ...(email ? { domain: getDomain(email) } : {}), + name: name || "Mock User", + admin: !!admin, + groups: groups || [], + ...customAttributes, + }, + // Make sure to match TTL, GetDeployment defaults to 24-hours + accessTokenTtlSeconds: RUNTIME_ACCESS_TOKEN_DEFAULT_TTL / 1000, // Seconds param vs milliseconds constant + }, + signal, + ); + + if (!accessToken) throw new Error("No JWT returned"); + return accessToken; +} diff --git a/web-admin/src/features/projects/project-query-options.spec.ts b/web-admin/src/features/projects/project-query-options.spec.ts new file mode 100644 index 000000000000..4b800d9dd5e8 --- /dev/null +++ b/web-admin/src/features/projects/project-query-options.spec.ts @@ -0,0 +1,140 @@ +import { V1DeploymentStatus, type RpcStatus } from "@rilldata/web-admin/client"; +import { + baseGetProjectQueryOptions, + createRuntimeJwtRefresher, +} from "@rilldata/web-admin/features/projects/project-query-options"; +import { queryClient } from "@rilldata/web-common/lib/svelte-query/globalQueryClient"; +import { afterEach, describe, expect, it, vi } from "vitest"; + +const FifteenMinutes = 15 * 60 * 1000; +const runningDeployment = { + deployment: { status: V1DeploymentStatus.DEPLOYMENT_STATUS_RUNNING }, +}; + +function httpError(status: number): RpcStatus { + return { response: { status, data: { message: "nope" } } } as RpcStatus; +} + +function refetchInterval(state: { data?: unknown; error?: unknown }) { + const compute = baseGetProjectQueryOptions.refetchInterval; + if (typeof compute !== "function") { + throw new Error("expected refetchInterval to be a function"); + } + return compute({ state } as unknown as Parameters[0]); +} + +describe("baseGetProjectQueryOptions.refetchInterval", () => { + it("polls a running deployment at half the JWT TTL", () => { + expect(refetchInterval({ data: runningDeployment })).toBe(FifteenMinutes); + }); + + it("keeps polling after a transient failure so the JWT still gets refreshed", () => { + expect( + refetchInterval({ data: runningDeployment, error: httpError(503) }), + ).toBe(FifteenMinutes); + }); + + it("stops polling once the session or the project access is gone", () => { + for (const status of [401, 403]) { + expect( + refetchInterval({ data: runningDeployment, error: httpError(status) }), + ).toBe(false); + } + }); +}); + +describe("createRuntimeJwtRefresher", () => { + const neverAborted = new AbortController().signal; + + afterEach(() => { + vi.restoreAllMocks(); + }); + + function mockFetchQuery(data: unknown) { + return vi + .spyOn(queryClient, "fetchQuery") + .mockImplementation(() => Promise.resolve(data) as never); + } + + it("refetches GetProject for the signed-in user and returns its JWT", async () => { + const fetchQuery = mockFetchQuery({ jwt: "fresh" }); + const refresh = createRuntimeJwtRefresher("user", { + organization: "org", + project: "proj", + branch: "dev", + }); + + await expect(refresh?.(neverAborted)).resolves.toBe("fresh"); + expect(fetchQuery).toHaveBeenCalledWith( + expect.objectContaining({ + queryKey: ["/v1/orgs/org/projects/proj", { branch: "dev" }], + }), + ); + }); + + it("refetches GetProject with the magic token on public URLs", async () => { + const fetchQuery = mockFetchQuery({ jwt: "fresh" }); + const refresh = createRuntimeJwtRefresher("magic", { + organization: "org", + project: "proj", + token: "magic", + }); + + await expect(refresh?.(neverAborted)).resolves.toBe("fresh"); + expect(fetchQuery).toHaveBeenCalledWith( + expect.objectContaining({ + queryKey: ["/v1/orgs/org/projects/proj", "token/magic"], + }), + ); + }); + + it("refetches the impersonated user's deployment credentials for View As", async () => { + const fetchQuery = mockFetchQuery({ accessToken: "fresh" }); + const refresh = createRuntimeJwtRefresher("mock", { + organization: "org", + project: "proj", + branch: "dev", + mockedUserId: "u1", + }); + + await expect(refresh?.(neverAborted)).resolves.toBe("fresh"); + expect(fetchQuery).toHaveBeenCalledWith( + expect.objectContaining({ + queryKey: [ + "/v1/orgs/org/projects/proj/credentials", + { userId: "u1", branch: "dev" }, + ], + }), + ); + }); + + it("cancels the admin request when the client abandons the refresh", () => { + vi.spyOn(queryClient, "fetchQuery").mockImplementation( + () => new Promise(() => {}) as never, + ); + const cancelQueries = vi + .spyOn(queryClient, "cancelQueries") + .mockResolvedValue(); + const abort = new AbortController(); + const refresh = createRuntimeJwtRefresher("user", { + organization: "org", + project: "proj", + }); + + void refresh?.(abort.signal); + abort.abort(); + + expect(cancelQueries).toHaveBeenCalledWith({ + queryKey: ["/v1/orgs/org/projects/proj"], + }); + }); + + it("has nothing to refetch for embeds", () => { + expect( + createRuntimeJwtRefresher("embed", { + organization: "org", + project: "proj", + }), + ).toBeUndefined(); + }); +}); diff --git a/web-admin/src/features/projects/project-query-options.ts b/web-admin/src/features/projects/project-query-options.ts index 012c09240c9d..401de92386ef 100644 --- a/web-admin/src/features/projects/project-query-options.ts +++ b/web-admin/src/features/projects/project-query-options.ts @@ -1,20 +1,44 @@ import { + adminServiceGetDeploymentCredentials, + adminServiceGetProject, + getAdminServiceGetDeploymentCredentialsQueryKey, + getAdminServiceGetProjectQueryKey, V1DeploymentStatus, type RpcStatus, + type V1GetDeploymentCredentialsResponse, type V1GetProjectResponse, } from "@rilldata/web-admin/client"; +import { + adminServiceGetProjectWithBearerToken, + getAdminServiceGetProjectWithBearerTokenQueryKey, +} from "@rilldata/web-admin/features/public-urls/get-project-with-bearer-token"; +import { extractErrorStatusCode } from "@rilldata/web-common/lib/errors"; +import { queryClient } from "@rilldata/web-common/lib/svelte-query/globalQueryClient"; import { RUNTIME_ACCESS_TOKEN_DEFAULT_TTL } from "@rilldata/web-common/runtime-client/constants"; -import type { CreateQueryOptions } from "@tanstack/svelte-query"; +import type { + AuthContext, + JwtRefresher, +} from "@rilldata/web-common/runtime-client/v2"; +import type { + CreateQueryOptions, + FetchQueryOptions, +} from "@tanstack/svelte-query"; const PollTimeWhenProjectDeploymentPending = 1000; const PollTimeWhenProjectDeploymentError = 5000; const PollTimeWhenProjectDeploymentOk = RUNTIME_ACCESS_TOKEN_DEFAULT_TTL / 2; // Proactively refetch the JWT before it expires +// A refetch that completed this recently already carries a fresh JWT; hand it over instead of fetching again. +const FreshJwtStaleTime = 5000; export const baseGetProjectQueryOptions: Partial< CreateQueryOptions > = { gcTime: Math.min(RUNTIME_ACCESS_TOKEN_DEFAULT_TTL, 1000 * 60 * 5), // Make sure we don't keep a stale JWT in the cache refetchInterval: (query) => { + // The session or the project access is gone: polling would repeat the same 401/403 until the tab is closed. + const errorStatus = extractErrorStatusCode(query.state.error); + if (errorStatus === 401 || errorStatus === 403) return false; + const status = query.state.data?.deployment?.status; switch (status) { case V1DeploymentStatus.DEPLOYMENT_STATUS_PENDING: @@ -38,3 +62,107 @@ export const baseGetProjectQueryOptions: Partial< // refetch. Focus/reconnect refetch still kicks in once data exceeds this window. staleTime: PollTimeWhenProjectDeploymentOk, }; + +/** + * Creates the callback `RuntimeProvider` uses to fetch a fresh JWT when a request finds the current one + * expired, which happens when a laptop sleep or a network drop skipped the scheduled refetch above. + * It re-runs the admin query that issued the JWT for the active auth mode, keyed exactly like the + * layout's own query so a refetch here and a scheduled one dedupe. + */ +export function createRuntimeJwtRefresher( + authContext: AuthContext, + opts: { + organization: string; + project: string; + branch?: string; + /** Magic auth token of the public URL being viewed. */ + token?: string; + /** User impersonated with View As. */ + mockedUserId?: string; + }, +): JwtRefresher | undefined { + const { organization, project, branch, token, mockedUserId } = opts; + switch (authContext) { + case "mock": { + if (!mockedUserId) return undefined; + const body = { userId: mockedUserId, ...(branch ? { branch } : {}) }; + return (signal) => + fetchFreshJwt( + { + queryKey: getAdminServiceGetDeploymentCredentialsQueryKey( + organization, + project, + body, + ), + queryFn: ({ signal }) => + adminServiceGetDeploymentCredentials( + organization, + project, + body, + signal, + ), + }, + (data) => data.accessToken, + signal, + ); + } + case "magic": { + if (!token) return undefined; + return (signal) => + fetchFreshJwt( + { + queryKey: getAdminServiceGetProjectWithBearerTokenQueryKey( + organization, + project, + token, + ), + queryFn: ({ signal }) => + adminServiceGetProjectWithBearerToken( + organization, + project, + token, + undefined, + signal, + ), + }, + (data) => data.jwt, + signal, + ); + } + case "user": { + const params = branch ? { branch } : undefined; + return (signal) => + fetchFreshJwt( + { + queryKey: getAdminServiceGetProjectQueryKey( + organization, + project, + params, + ), + queryFn: ({ signal }) => + adminServiceGetProject(organization, project, params, signal), + }, + (data) => data.jwt, + signal, + ); + } + case "embed": + // Embed tokens come from the host page and last 24 hours; there is nothing to refetch. + return undefined; + } +} + +function fetchFreshJwt( + options: FetchQueryOptions, + selectJwt: (data: TData) => string | undefined, + signal: AbortSignal, +): Promise { + // The client abandons a call that stalls; cancel the request so the retry is not deduped into it + const cancel = () => + void queryClient.cancelQueries({ queryKey: options.queryKey }); + signal.addEventListener("abort", cancel, { once: true }); + return queryClient + .fetchQuery({ ...options, staleTime: FreshJwtStaleTime }) + .then(selectJwt) + .finally(() => signal.removeEventListener("abort", cancel)); +} diff --git a/web-admin/src/routes/[organization]/[project]/+layout.svelte b/web-admin/src/routes/[organization]/[project]/+layout.svelte index aa7818fa991d..d57598bcc2ae 100644 --- a/web-admin/src/routes/[organization]/[project]/+layout.svelte +++ b/web-admin/src/routes/[organization]/[project]/+layout.svelte @@ -45,7 +45,10 @@ import ProjectHeader from "../../../features/projects/header/ProjectHeader.svelte"; import ProjectTabs from "@rilldata/web-admin/features/projects/ProjectTabs.svelte"; import { resolveProjectError } from "@rilldata/web-admin/features/projects/project-error"; - import { baseGetProjectQueryOptions } from "@rilldata/web-admin/features/projects/project-query-options"; + import { + baseGetProjectQueryOptions, + createRuntimeJwtRefresher, + } from "@rilldata/web-admin/features/projects/project-query-options"; import { resolveRuntimeConnection } from "@rilldata/web-admin/features/projects/project-runtime"; import RedeployProjectCta from "@rilldata/web-admin/features/projects/RedeployProjectCTA.svelte"; import SlimProjectHeader from "@rilldata/web-admin/features/projects/SlimProjectHeader.svelte"; @@ -232,6 +235,18 @@ resolveRuntimeConnection(projectData, mockUser, onPublicURLPage), ); + // Re-requests the runtime JWT from the query that issued it when a request finds it expired, e.g. after a + // laptop sleep or a network drop skipped the scheduled refetch, instead of waiting for the next poll. + let refreshJwt = $derived( + createRuntimeJwtRefresher(runtime.authContext, { + organization, + project, + branch: activeBranch, + token, + mockedUserId, + }), + ); + // --- Side effects (cache invalidation and telemetry) --- // Track previous status so we invalidate only on *transitions*, not on every @@ -353,6 +368,7 @@ instanceId={runtime.instanceId} jwt={runtime.jwt} authContext={runtime.authContext} + {refreshJwt} > {#if !onWelcomePage} {:else if isReady && deployment?.id && instanceId && runtimeHost && jwt} {#key `${runtimeHost}::${instanceId}::${hasPrimaryDeployment}`} - + {#if !inProjectWelcomePage} import { useQueryClient } from "@tanstack/svelte-query"; - import { onDestroy, setContext } from "svelte"; + import { onDestroy, onMount, setContext } from "svelte"; import { featureFlags } from "../../features/feature-flags"; import { invalidateRuntimeQueries } from "../invalidation"; import { @@ -8,7 +8,7 @@ evictRuntimeClient, RUNTIME_CONTEXT_KEY, } from "./context"; - import type { AuthContext } from "./runtime-client"; + import type { AuthContext, JwtRefresher } from "./runtime-client"; import { createEnvFileStore } from "@rilldata/web-common/features/env-management/env-file-store.ts"; const queryClient = useQueryClient(); @@ -17,6 +17,9 @@ export let instanceId: string; export let jwt: string | undefined = undefined; export let authContext: AuthContext = "user"; + // Fetches a fresh JWT when a request finds the current one expired, e.g. after a laptop sleep + // skipped the scheduled refetch. Without it, requests wait for the parent to pass a new `jwt`. + export let refreshJwt: JwtRefresher | undefined = undefined; // Returns a cached instance if a load function already created one for this host+instanceId. // If host/instanceId change, the parent's {#key} re-mounts us. @@ -29,10 +32,28 @@ // Handle JWT-only changes (15-min refresh, View As with same host) $: { const authContextChanged = client.updateJwt(jwt, authContext); + client.setJwtRefresher(refreshJwt); if (authContextChanged) void invalidateRuntimeQueries(queryClient, instanceId); } + // A tab that was asleep, hidden or offline has likely missed the scheduled JWT refetch. + // Refresh as soon as it is back so the first requests don't have to wait for it. + onMount(() => { + const refreshIfStale = () => void client.refreshJwtIfStale(); + const refreshWhenVisible = () => { + if (document.visibilityState !== "hidden") refreshIfStale(); + }; + window.addEventListener("online", refreshIfStale); + window.addEventListener("focus", refreshWhenVisible); + document.addEventListener("visibilitychange", refreshWhenVisible); + return () => { + window.removeEventListener("online", refreshIfStale); + window.removeEventListener("focus", refreshWhenVisible); + document.removeEventListener("visibilitychange", refreshWhenVisible); + }; + }); + onDestroy(() => { featureFlags.clearRuntimeClient(); evictRuntimeClient(client); diff --git a/web-common/src/runtime-client/v2/index.ts b/web-common/src/runtime-client/v2/index.ts index 921cb40acbe9..e6c05b42c678 100644 --- a/web-common/src/runtime-client/v2/index.ts +++ b/web-common/src/runtime-client/v2/index.ts @@ -1,4 +1,8 @@ -export { RuntimeClient, type AuthContext } from "./runtime-client"; +export { + RuntimeClient, + type AuthContext, + type JwtRefresher, +} from "./runtime-client"; export { getRuntimeClient, evictRuntimeClient, diff --git a/web-common/src/runtime-client/v2/runtime-client.spec.ts b/web-common/src/runtime-client/v2/runtime-client.spec.ts index 682069036d8a..87d5f2492990 100644 --- a/web-common/src/runtime-client/v2/runtime-client.spec.ts +++ b/web-common/src/runtime-client/v2/runtime-client.spec.ts @@ -25,12 +25,41 @@ function createFetchMock() { ); } +/** Extract the Authorization header from the given fetch call. */ +function getAuthHeader( + fetchMock: ReturnType, + index: number, +): string | null { + const call = fetchMock.mock.calls[index] as + | [RequestInfo, RequestInit | undefined] + | undefined; + const headers = call?.[1]?.headers as Headers | undefined; + return headers?.get?.("Authorization") ?? null; +} + /** Extract the Authorization header from the most recent fetch call. */ function getLastAuthHeader(fetchMock: ReturnType): string | null { - if (fetchMock.mock.calls.length === 0) return null; - const lastCall = fetchMock.mock.calls[fetchMock.mock.calls.length - 1]; - const headers: Headers | undefined = lastCall[1]?.headers; - return headers?.get?.("Authorization") ?? null; + return getAuthHeader(fetchMock, fetchMock.mock.calls.length - 1); +} + +/** Builds an unsigned JWT whose payload carries the given claims (times in seconds since the epoch). */ +function makeJwt(claims: Record): string { + const payload = btoa(JSON.stringify(claims)) + .replace(/\+/g, "-") + .replace(/\//g, "_") + .replace(/=+$/, ""); + return `eyJhbGciOiJub25lIn0.${payload}.signature`; +} + +/** A unary Connect error response as the runtime returns it for an expired token. */ +function expiredJwtResponse() { + return new Response( + JSON.stringify({ + code: "unauthenticated", + message: "rpc error: code = Unauthenticated desc = jwt is expired", + }), + { status: 401, headers: { "content-type": "application/json" } }, + ); } describe("RuntimeClient JWT refresh", () => { @@ -319,4 +348,424 @@ describe("RuntimeClient JWT refresh", () => { expect(client.authContext).toBe("user"); }); }); + + // ── Active refresh through the registered refresher ──────────────── + + describe("refresher", () => { + it("asks the refresher for a new JWT instead of waiting for updateJwt", async () => { + const client = createClient({ jwt: "expiring-token" }); + const refreshJwt = vi.fn().mockResolvedValue("fresh-token"); + client.setJwtRefresher(refreshJwt); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + await makeRequest(client); + + expect(refreshJwt).toHaveBeenCalledTimes(1); + expect(getLastAuthHeader(fetchMock)).toBe("Bearer fresh-token"); + }); + + it("shares one in-flight refresh between concurrent requests", async () => { + const client = createClient({ jwt: "expiring-token" }); + let resolveRefresh!: (jwt: string) => void; + const refreshJwt = vi.fn( + () => + new Promise((resolve) => { + resolveRefresh = resolve; + }), + ); + client.setJwtRefresher(refreshJwt); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + const requests = Promise.all([ + makeRequest(client), + makeRequest(client), + makeRequest(client), + ]); + await vi.advanceTimersByTimeAsync(1); + expect(refreshJwt).toHaveBeenCalledTimes(1); + expect(fetchMock).not.toHaveBeenCalled(); + + resolveRefresh("fresh-token"); + await requests; + + expect(fetchMock).toHaveBeenCalledTimes(3); + for (let i = 0; i < 3; i++) { + expect(getAuthHeader(fetchMock, i)).toBe("Bearer fresh-token"); + } + }); + + it("retries a failed refresh with backoff until it succeeds", async () => { + const client = createClient({ jwt: "expiring-token" }); + const refreshJwt = vi + .fn() + .mockRejectedValueOnce(new Error("Network Error")) + .mockResolvedValueOnce("fresh-token"); + client.setJwtRefresher(refreshJwt); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + let resolved = false; + const promise = makeRequest(client).then(() => { + resolved = true; + }); + + // The first attempt fails; the retry waits for the 1s backoff + await vi.advanceTimersByTimeAsync(500); + expect(refreshJwt).toHaveBeenCalledTimes(1); + expect(resolved).toBe(false); + + await vi.advanceTimersByTimeAsync(1000); + await promise; + + expect(refreshJwt).toHaveBeenCalledTimes(2); + expect(getLastAuthHeader(fetchMock)).toBe("Bearer fresh-token"); + }); + + it("fails at once when the issuer refuses because the session or access is gone", async () => { + const client = createClient({ jwt: "expiring-token" }); + const refreshJwt = vi.fn().mockRejectedValue({ + response: { + status: 403, + data: { message: "does not have permission to read project" }, + }, + }); + client.setJwtRefresher(refreshJwt); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + const error = await client.runtimeService.ping({}).catch((e: Error) => e); + + expect(error).toBeInstanceOf(Error); + expect((error as Error).message).toContain( + "Failed to refresh the runtime JWT: does not have permission to read project", + ); + expect(refreshJwt).toHaveBeenCalledTimes(1); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it("gives up after the timeout and reports why the refresh kept failing", async () => { + const client = createClient({ jwt: "expiring-token" }); + const refreshJwt = vi.fn().mockRejectedValue(new Error("Network Error")); + client.setJwtRefresher(refreshJwt); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + const promise = client.runtimeService.ping({}).catch((e: Error) => e); + // 60s deadline plus the longest backoff + await vi.advanceTimersByTimeAsync(75_000); + const error = await promise; + + expect((error as Error).message).toContain( + "Timed out waiting for a fresh JWT: Network Error", + ); + // Several attempts, but spaced out by the backoff rather than a tight loop + expect(refreshJwt.mock.calls.length).toBeGreaterThan(1); + expect(refreshJwt.mock.calls.length).toBeLessThan(20); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it("backs off when the issuer hands back an already expired JWT", async () => { + const client = createClient({ jwt: "expiring-token" }); + const now = Date.now() / 1000; + const refreshJwt = vi + .fn() + .mockResolvedValue(makeJwt({ iat: now - 3600, exp: now - 60 })); + client.setJwtRefresher(refreshJwt); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + const promise = client.runtimeService.ping({}).catch((e: Error) => e); + await vi.advanceTimersByTimeAsync(75_000); + const error = await promise; + + expect((error as Error).message).toContain( + "Timed out waiting for a fresh JWT: the issuer returned an expired JWT", + ); + expect(refreshJwt.mock.calls.length).toBeLessThan(20); + }); + + it("abandons a refresh that stalls, cancels it and tries again", async () => { + const client = createClient({ jwt: "expiring-token" }); + const signals: AbortSignal[] = []; + const refreshJwt = vi.fn((signal: AbortSignal) => { + signals.push(signal); + return new Promise(() => {}); + }); + client.setJwtRefresher(refreshJwt); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + const promise = client.runtimeService.ping({}).catch((e: Error) => e); + + // An attempt is given 20s, then the retry waits out the 1s backoff + await vi.advanceTimersByTimeAsync(21_001); + expect(signals[0].aborted).toBe(true); + expect(refreshJwt).toHaveBeenCalledTimes(2); + + // The request itself still fails at the 60s deadline + await vi.advanceTimersByTimeAsync(40_000); + const error = await promise; + expect((error as Error).message).toContain( + "Timed out waiting for a fresh JWT: the issuer did not respond within 20s", + ); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it("releases requests waiting on a stalled refresh as soon as a JWT arrives from elsewhere", async () => { + // e.g. the scheduled GetProject refetch completes while the refresh hangs + const client = createClient({ jwt: "expiring-token" }); + client.setJwtRefresher(vi.fn(() => new Promise(() => {}))); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + let resolved = false; + const promise = makeRequest(client).then(() => { + resolved = true; + }); + await vi.advanceTimersByTimeAsync(5_000); + expect(resolved).toBe(false); + + client.updateJwt("fresh-token"); + await promise; + + expect(resolved).toBe(true); + expect(getLastAuthHeader(fetchMock)).toBe("Bearer fresh-token"); + }); + + it("does not apply a user token refreshed during the backoff after View As switched the context", async () => { + const client = createClient({ jwt: "user-token", authContext: "user" }); + const refreshJwt = vi + .fn() + .mockRejectedValueOnce(new Error("Network Error")) + .mockResolvedValue("fresh-user-token"); + client.setJwtRefresher(refreshJwt); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + const promise = makeRequest(client); + // The first attempt failed; the retry is sleeping through the backoff + await vi.advanceTimersByTimeAsync(500); + expect(refreshJwt).toHaveBeenCalledTimes(1); + + // View As installs a mock token while the retry sleeps + client.updateJwt("mock-token", "mock"); + await promise; + await vi.advanceTimersByTimeAsync(2_000); + + expect(refreshJwt).toHaveBeenCalledTimes(1); + expect(client.getJwt()).toBe("mock-token"); + expect(client.authContext).toBe("mock"); + expect(getLastAuthHeader(fetchMock)).toBe("Bearer mock-token"); + }); + + it("keeps a refresher per auth context so View As can register its own", async () => { + const client = createClient({ jwt: "user-token", authContext: "user" }); + const refreshUserJwt = vi.fn().mockResolvedValue("fresh-user-token"); + client.setJwtRefresher(refreshUserJwt); + + // View As in the editor installs a mock token and a refresher that reissues it + client.updateJwt("mock-token", "mock"); + const refreshMockJwt = vi.fn().mockResolvedValue("fresh-mock-token"); + client.setJwtRefresher(refreshMockJwt); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + await makeRequest(client); + expect(refreshMockJwt).toHaveBeenCalledTimes(1); + expect(refreshUserJwt).not.toHaveBeenCalled(); + expect(getLastAuthHeader(fetchMock)).toBe("Bearer fresh-mock-token"); + + // View As ends: the user's refresher applies again without being registered anew + client.updateJwt("user-token", "user"); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + await makeRequest(client); + expect(refreshUserJwt).toHaveBeenCalledTimes(1); + expect(getLastAuthHeader(fetchMock)).toBe("Bearer fresh-user-token"); + }); + + it("removes the current context's refresher when given undefined", async () => { + const client = createClient({ jwt: "token" }); + const refreshJwt = vi.fn().mockResolvedValue("fresh-token"); + client.setJwtRefresher(refreshJwt); + client.setJwtRefresher(undefined); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + const promise = client.runtimeService.ping({}).catch((e: Error) => e); + await vi.advanceTimersByTimeAsync(61_000); + const error = await promise; + + expect(refreshJwt).not.toHaveBeenCalled(); + expect((error as Error).message).toContain( + "Timed out waiting for a fresh JWT", + ); + }); + + it("ignores a refresher registered for another auth context", async () => { + // View As in Rill Developer sets a mock JWT directly on the client; the provider's + // refresher for the "user" context must not replace it with the user's token + const client = createClient({ jwt: "user-token", authContext: "user" }); + const refreshJwt = vi.fn().mockResolvedValue("fresh-user-token"); + client.setJwtRefresher(refreshJwt); + client.updateJwt("mock-token", "mock"); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + const promise = client.runtimeService.ping({}).catch((e: Error) => e); + await vi.advanceTimersByTimeAsync(61_000); + const error = await promise; + + expect(refreshJwt).not.toHaveBeenCalled(); + expect((error as Error).message).toContain( + "Timed out waiting for a fresh JWT", + ); + + // Back in the provider's context, the refresher applies again + client.updateJwt("user-token", "user"); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + await makeRequest(client); + + expect(refreshJwt).toHaveBeenCalledTimes(1); + expect(getLastAuthHeader(fetchMock)).toBe("Bearer fresh-user-token"); + }); + + it("refreshes once and retries when the runtime rejects the JWT as expired", async () => { + // The client's clock can disagree with the runtime's; the runtime wins + const client = createClient({ jwt: "token-runtime-rejects" }); + const refreshJwt = vi.fn().mockResolvedValue("fresh-token"); + client.setJwtRefresher(refreshJwt); + // The retry reuses the request's Headers object, so capture the header as each call is made + const sentAuthHeaders: (string | null)[] = []; + fetchMock.mockImplementation((_, init) => { + sentAuthHeaders.push( + (init?.headers as Headers | undefined)?.get("Authorization") ?? null, + ); + return Promise.resolve( + sentAuthHeaders.length === 1 + ? expiredJwtResponse() + : new Response(new Uint8Array(0), { + status: 200, + headers: { "content-type": "application/proto" }, + }), + ); + }); + + await makeRequest(client); + + expect(refreshJwt).toHaveBeenCalledTimes(1); + expect(sentAuthHeaders).toEqual([ + "Bearer token-runtime-rejects", + "Bearer fresh-token", + ]); + }); + + it("does not retry on other authentication errors", async () => { + const client = createClient({ jwt: "bad-token" }); + const refreshJwt = vi.fn().mockResolvedValue("fresh-token"); + client.setJwtRefresher(refreshJwt); + fetchMock.mockImplementationOnce(() => + Promise.resolve( + new Response( + JSON.stringify({ code: "unauthenticated", message: "bad token" }), + { status: 401, headers: { "content-type": "application/json" } }, + ), + ), + ); + + const error = await client.runtimeService.ping({}).catch((e: Error) => e); + + expect(error).toBeInstanceOf(Error); + expect(refreshJwt).not.toHaveBeenCalled(); + expect(fetchMock).toHaveBeenCalledTimes(1); + }); + + it("does not retry without a refresher, e.g. for embeds", async () => { + const client = createClient({ jwt: "embed-token", authContext: "embed" }); + fetchMock.mockImplementationOnce(() => + Promise.resolve(expiredJwtResponse()), + ); + + const error = await client.runtimeService.ping({}).catch((e: Error) => e); + + expect(error).toBeInstanceOf(Error); + expect(fetchMock).toHaveBeenCalledTimes(1); + }); + }); + + // ── Proactive refresh (tab visible again, network back) ──────────── + + describe("refreshJwtIfStale", () => { + it("refreshes a stale JWT and leaves a fresh one alone", async () => { + const client = createClient({ jwt: "token" }); + const refreshJwt = vi.fn().mockResolvedValue("fresh-token"); + client.setJwtRefresher(refreshJwt); + + await client.refreshJwtIfStale(); + expect(refreshJwt).not.toHaveBeenCalled(); + + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + await client.refreshJwtIfStale(); + + expect(refreshJwt).toHaveBeenCalledTimes(1); + expect(client.getJwt()).toBe("fresh-token"); + }); + + it("leaves refresh failures to the next request", async () => { + const client = createClient({ jwt: "token" }); + const refreshJwt = vi.fn().mockRejectedValue(new Error("Network Error")); + client.setJwtRefresher(refreshJwt); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + await expect(client.refreshJwtIfStale()).resolves.toBeUndefined(); + expect(client.getJwt()).toBe("token"); + }); + }); + + // ── Expiry derived from the token's claims ───────────────────────── + + describe("expiry from the token's claims", () => { + it("treats a token handed over from a cache as expired once its exp claim has passed", async () => { + // The GetProject cache can hand a new RuntimeProvider a token issued long ago. Crediting it + // with a full TTL from receipt made the runtime reject requests with "jwt is expired". + const now = Date.now() / 1000; + const client = createClient({ + jwt: makeJwt({ iat: now - 40 * 60, exp: now - 10 * 60 }), + }); + const refreshJwt = vi.fn().mockResolvedValue("fresh-token"); + client.setJwtRefresher(refreshJwt); + + await makeRequest(client); + + expect(refreshJwt).toHaveBeenCalledTimes(1); + expect(getLastAuthHeader(fetchMock)).toBe("Bearer fresh-token"); + }); + + it("trusts an exp claim beyond the default TTL", async () => { + const now = Date.now() / 1000; + const oneHourToken = makeJwt({ iat: now, exp: now + 60 * 60 }); + const client = createClient({ jwt: oneHourToken }); + vi.advanceTimersByTime(RUNTIME_ACCESS_TOKEN_DEFAULT_TTL + 60_000); + + await makeRequest(client); + + expect(fetchMock).toHaveBeenCalledTimes(1); + expect(getLastAuthHeader(fetchMock)).toBe(`Bearer ${oneHourToken}`); + }); + + it("caps a fresh token at its lifetime from receipt when the client clock runs behind the issuer", async () => { + // The issuer's clock is 5 minutes ahead: exp looks 35 minutes away, but the token only lives 30 + const now = Date.now() / 1000; + const skew = 5 * 60; + const client = createClient({ + jwt: makeJwt({ iat: now + skew, exp: now + skew + 30 * 60 }), + }); + const refreshJwt = vi.fn().mockResolvedValue("fresh-token"); + client.setJwtRefresher(refreshJwt); + vi.advanceTimersByTime(PAST_WARNING_WINDOW); + + await makeRequest(client); + + expect(refreshJwt).toHaveBeenCalledTimes(1); + expect(getLastAuthHeader(fetchMock)).toBe("Bearer fresh-token"); + }); + + it("falls back to the default TTL for tokens without claims", async () => { + const client = createClient({ jwt: makeJwt({ sub: "dev-user" }) }); + vi.advanceTimersByTime(RUNTIME_ACCESS_TOKEN_DEFAULT_TTL / 2); + + await makeRequest(client); + + expect(fetchMock).toHaveBeenCalledTimes(1); + }); + }); }); diff --git a/web-common/src/runtime-client/v2/runtime-client.ts b/web-common/src/runtime-client/v2/runtime-client.ts index e46e4f7e3212..e6f44d9b93e8 100644 --- a/web-common/src/runtime-client/v2/runtime-client.ts +++ b/web-common/src/runtime-client/v2/runtime-client.ts @@ -1,17 +1,40 @@ -import { createClient, type Client, type Transport } from "@connectrpc/connect"; +import { + Code, + ConnectError, + createClient, + type Client, + type Transport, +} from "@connectrpc/connect"; import { createConnectTransport } from "@connectrpc/connect-web"; import { QueryService } from "../../proto/gen/rill/runtime/v1/queries_connect"; import { RuntimeService } from "../../proto/gen/rill/runtime/v1/api_connect"; import { ConnectorService } from "../../proto/gen/rill/runtime/v1/connectors_connect"; +import { extractErrorMessage, extractErrorStatusCode } from "../../lib/errors"; import { RUNTIME_ACCESS_TOKEN_DEFAULT_TTL, JWT_EXPIRY_WARNING_WINDOW, - CHECK_RUNTIME_STORE_FOR_JWT_INTERVAL, } from "../constants"; import { RequestQueue, createQueueInterceptor } from "./request-queue"; export type AuthContext = "user" | "mock" | "magic" | "embed"; +/** + * Fetches a fresh JWT from whoever issued the current one, e.g. by refetching GetProject in Rill Cloud. + * Resolves with the new token; rejects when the issuer could not be reached or refused. + * The signal aborts when the client gives up on the call, so the issuer request can be cancelled. + */ +export type JwtRefresher = (signal: AbortSignal) => Promise; + +// How long a request waits for a fresh JWT before it fails. +const JWT_REFRESH_TIMEOUT = 60_000; +// How long one call to the refresher may take before it is abandoned and retried. +const JWT_REFRESH_ATTEMPT_TIMEOUT = 20_000; +// Delay before retrying a failed refresh: doubles from the base up to the max while the issuer keeps failing. +const JWT_REFRESH_BACKOFF_BASE = 1_000; +const JWT_REFRESH_BACKOFF_MAX = 10_000; + +const TimedOut = Symbol("timed out"); + export class RuntimeClient { readonly host: string; readonly instanceId: string; @@ -20,8 +43,18 @@ export class RuntimeClient { // JWT state (mutable; read by the transport interceptor) private currentJwt: string | undefined; - private jwtReceivedAt: number; + private jwtExpiresAt = 0; private _authContext: AuthContext; + // Resolved and replaced whenever the JWT, the auth context or a refresher changes, or the client + // is disposed, so that requests waiting for a fresh JWT re-check their situation right away + private jwtChanged = createDeferred(); + // Refreshers per auth context: RuntimeProvider registers one for its context and View As one for "mock" + private readonly jwtRefreshers = new Map(); + // In-flight refresh shared by every request waiting on it, and the backoff after failures + private jwtRefresh: Promise | undefined; + private jwtRefreshAbort: AbortController | undefined; + private jwtRefreshBackoff = 0; + private jwtRefreshNotBefore = 0; private disposed = false; // Cached service clients (created once per RuntimeClient) @@ -50,7 +83,7 @@ export class RuntimeClient { this.host = opts.host; this.instanceId = opts.instanceId; this.currentJwt = opts.jwt; - this.jwtReceivedAt = opts.jwt ? Date.now() : 0; + this.jwtExpiresAt = opts.jwt ? jwtExpiry(opts.jwt, Date.now()) : 0; this._authContext = opts.authContext ?? "user"; this.requestQueue = new RequestQueue(); @@ -61,11 +94,32 @@ export class RuntimeClient { createQueueInterceptor(this.requestQueue), // JWT interceptor adds auth header just before the network call (next) => async (req) => { - if (this.currentJwt) { + await this.waitForFreshJwt(); + const jwt = this.currentJwt; + if (jwt) req.header.set("Authorization", `Bearer ${jwt}`); + try { + return await next(req); + } catch (err) { + // The runtime's clock is the source of truth: when it rejects a token this client still + // considered valid, refresh once and retry. The runtime rejects expired tokens before the + // handler runs, so the retry is safe; only unary requests can be re-sent as is. + if ( + !jwt || + req.stream || + !this.canRefreshJwt() || + !isExpiredJwtError(err) + ) { + throw err; + } + if (jwt === this.currentJwt) this.jwtExpiresAt = 0; await this.waitForFreshJwt(); - req.header.set("Authorization", `Bearer ${this.currentJwt}`); + if (this.currentJwt) { + req.header.set("Authorization", `Bearer ${this.currentJwt}`); + } else { + req.header.delete("Authorization"); + } + return next(req); } - return next(req); }, ], }); @@ -79,15 +133,51 @@ export class RuntimeClient { const authContextChanged = !!this._authContext && !!authContext && authContext !== this._authContext; - if (jwt !== this.currentJwt) { + const jwtChanged = jwt !== this.currentJwt; + if (jwtChanged) { this.currentJwt = jwt; - this.jwtReceivedAt = Date.now(); + this.jwtExpiresAt = jwt ? jwtExpiry(jwt, Date.now()) : 0; } if (authContext) this._authContext = authContext; + if (jwtChanged || authContextChanged) this.wakeWaitingRequests(); return authContextChanged; } + /** + * Registers the callback that fetches a fresh JWT for the current auth context, replacing any earlier + * one for that context; undefined removes it. Refreshers are kept per context: RuntimeProvider's stays + * registered for its context while View As installs a mock token with its own refresher for "mock", + * and applies again once View As ends and the context switches back. + */ + setJwtRefresher(refreshJwt: JwtRefresher | undefined): void { + const authContext = this._authContext; + if (refreshJwt === this.jwtRefreshers.get(authContext)) return; + if (refreshJwt) { + this.jwtRefreshers.set(authContext, refreshJwt); + } else { + this.jwtRefreshers.delete(authContext); + } + // A new issuer starts without the backoff the previous one earned + this.jwtRefreshBackoff = 0; + this.jwtRefreshNotBefore = 0; + this.wakeWaitingRequests(); + } + + /** + * Refreshes the JWT if it is expired or about to expire, e.g. when the tab becomes visible again + * after a laptop sleep skipped the scheduled refresh. Failures are left to the next request, + * which retries with backoff and surfaces the error. + */ + async refreshJwtIfStale(): Promise { + if (this.isJwtFresh() || !this.canRefreshJwt()) return; + try { + await this.runJwtRefresh(); + } catch { + // Reported by the next request that needs the token + } + } + /** Returns the current JWT (used by SSE clients and other non-query consumers). */ getJwt(): string | undefined { return this.currentJwt; @@ -117,33 +207,207 @@ export class RuntimeClient { } /** - * If the JWT is close to expiring, wait in a polling loop for - * RuntimeProvider to call updateJwt() with a fresh token. - * Ported from http-client.ts:maybeWaitForFreshJWT. + * Blocks until the JWT is fresh. With a refresher registered for the current auth context it asks + * the issuer for a new token right away, retrying with backoff while the issuer is unreachable; + * without one it waits for an updateJwt() call. Gives up after JWT_REFRESH_TIMEOUT, or at once + * when the issuer refuses because the session or the project access is gone. */ private async waitForFreshJwt(): Promise { - // Embeds have a 24h backend-issued TTL; skip client-side expiry check + // Embeds have a 24h backend-issued TTL and no refresher; skip client-side expiry check if (this._authContext === "embed") return; - const deadline = Date.now() + 60_000; - let expiresAt = this.jwtReceivedAt + RUNTIME_ACCESS_TOKEN_DEFAULT_TTL; - while (Date.now() + JWT_EXPIRY_WARNING_WINDOW > expiresAt) { + const deadline = Date.now() + JWT_REFRESH_TIMEOUT; + let lastError: unknown; + while (!this.isJwtFresh()) { if (this.disposed) { throw new DOMException("Request cancelled", "AbortError"); } - if (Date.now() > deadline) { - throw new Error("Timed out waiting for a fresh JWT"); + const remaining = deadline - Date.now(); + if (remaining <= 0) { + throw new Error( + "Timed out waiting for a fresh JWT" + + (lastError ? `: ${extractErrorMessage(lastError)}` : ""), + ); + } + // Whatever happens below, a change to the JWT or the client wakes us to re-check + const changed = this.jwtChanged.promise; + if (!this.canRefreshJwt()) { + // Wait for whoever owns the token (e.g. View As in Rill Developer) to call updateJwt() + await waitAtMost(changed, remaining); + continue; + } + try { + await waitAtMost( + Promise.race([this.runJwtRefresh(), changed]), + remaining, + ); + } catch (err) { + const status = extractErrorStatusCode(err); + if (status === 401 || status === 403) { + throw new Error( + `Failed to refresh the runtime JWT: ${extractErrorMessage(err)}`, + ); + } + lastError = err; + } + } + } + + private isJwtFresh(): boolean { + return ( + !this.currentJwt || + Date.now() + JWT_EXPIRY_WARNING_WINDOW <= this.jwtExpiresAt + ); + } + + private canRefreshJwt(): boolean { + return this.jwtRefreshers.has(this._authContext); + } + + private wakeWaitingRequests(): void { + const { resolve } = this.jwtChanged; + this.jwtChanged = createDeferred(); + resolve(); + } + + /** Runs one refresh, shared by every caller that joins while it is in flight. */ + private runJwtRefresh(): Promise { + this.jwtRefresh ??= this.refreshJwtOnce().finally(() => { + this.jwtRefresh = undefined; + }); + return this.jwtRefresh; + } + + private async refreshJwtOnce(): Promise { + // Pin the refresher and the context it serves: both can change while this attempt waits or runs + // (e.g. View As switching to a mock token), and a token must only be applied to the context it + // was issued for + const authContext = this._authContext; + const refreshJwt = this.jwtRefreshers.get(authContext); + if (!refreshJwt) return; + const delay = this.jwtRefreshNotBefore - Date.now(); + if (delay > 0) await sleep(delay); + if (!this.servesCurrentContext(refreshJwt, authContext)) return; + + // A call that stalls is abandoned and cancelled so the retry is not stuck behind it + const abort = new AbortController(); + this.jwtRefreshAbort = abort; + try { + const jwt = await waitAtMost( + refreshJwt(abort.signal), + JWT_REFRESH_ATTEMPT_TIMEOUT, + ); + if (jwt === TimedOut) { + abort.abort(); + throw new Error( + `the issuer did not respond within ${JWT_REFRESH_ATTEMPT_TIMEOUT / 1000}s`, + ); + } + if (!this.servesCurrentContext(refreshJwt, authContext)) return; + if (!jwt) throw new Error("the issuer returned no JWT"); + this.updateJwt(jwt); + if (!this.isJwtFresh()) { + throw new Error("the issuer returned an expired JWT"); } - await new Promise((r) => - setTimeout(r, CHECK_RUNTIME_STORE_FOR_JWT_INTERVAL), + this.jwtRefreshBackoff = 0; + this.jwtRefreshNotBefore = 0; + } catch (err) { + this.jwtRefreshBackoff = Math.min( + Math.max(this.jwtRefreshBackoff * 2, JWT_REFRESH_BACKOFF_BASE), + JWT_REFRESH_BACKOFF_MAX, ); - // Re-check: provider may have called updateJwt() while we waited - expiresAt = this.jwtReceivedAt + RUNTIME_ACCESS_TOKEN_DEFAULT_TTL; + this.jwtRefreshNotBefore = Date.now() + this.jwtRefreshBackoff; + throw err; + } finally { + if (this.jwtRefreshAbort === abort) this.jwtRefreshAbort = undefined; } } + /** Whether the refresher an attempt pinned is still the one registered for the client's current context. */ + private servesCurrentContext( + refreshJwt: JwtRefresher, + authContext: AuthContext, + ): boolean { + return ( + !this.disposed && + this._authContext === authContext && + this.jwtRefreshers.get(authContext) === refreshJwt + ); + } + dispose(): void { this.disposed = true; + this.jwtRefreshAbort?.abort(); + this.wakeWaitingRequests(); this.requestQueue.clear(); } } + +/** + * Matches the runtime's rejection of an expired token. It is raised by the auth interceptor + * before the handler runs, so a request that gets it can be retried with a fresh token. + */ +function isExpiredJwtError(err: unknown): boolean { + return ( + err instanceof ConnectError && + err.code === Code.Unauthenticated && + err.rawMessage.includes("jwt is expired") + ); +} + +/** + * Returns when the JWT expires, in milliseconds since the epoch. + * The exp claim is the source of truth, so a token handed over from a cache is not credited with a + * full TTL again. A fresh token is also capped at receivedAt plus its own lifetime (exp minus iat), + * which keeps a client clock that runs behind the issuer's from presenting tokens the runtime + * already rejects. Tokens without readable claims fall back to the default TTL from receipt. + */ +function jwtExpiry(jwt: string, receivedAt: number): number { + const claims = decodeJwtClaims(jwt); + if (typeof claims?.exp !== "number") { + return receivedAt + RUNTIME_ACCESS_TOKEN_DEFAULT_TTL; + } + const expiresAt = claims.exp * 1000; + if (typeof claims.iat !== "number") return expiresAt; + return Math.min(expiresAt, receivedAt + (claims.exp - claims.iat) * 1000); +} + +/** Decodes the payload of a JWT without verifying it. Returns undefined for anything that is not a JWT. */ +function decodeJwtClaims(jwt: string): Record | undefined { + const parts = jwt.split("."); + if (parts.length !== 3) return undefined; + try { + const base64 = parts[1].replace(/-/g, "+").replace(/_/g, "/"); + const bytes = Uint8Array.from(atob(base64), (c) => c.charCodeAt(0)); + const claims: unknown = JSON.parse(new TextDecoder().decode(bytes)); + return typeof claims === "object" && claims !== null + ? (claims as Record) + : undefined; + } catch { + return undefined; + } +} + +/** Settles like the promise, or resolves with TimedOut once `ms` elapse first; the timer that lost is cleared. */ +function waitAtMost( + promise: Promise, + ms: number, +): Promise { + let timer: ReturnType | undefined; + const timeout = new Promise((resolve) => { + timer = setTimeout(() => resolve(TimedOut), ms); + }); + return Promise.race([promise, timeout]).finally(() => clearTimeout(timer)); +} + +function createDeferred(): { promise: Promise; resolve: () => void } { + let resolve!: () => void; + const promise = new Promise((r) => { + resolve = r; + }); + return { promise, resolve }; +} + +function sleep(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +}