From 8ae084206e243ac2e0b70cd50c78c22894979205 Mon Sep 17 00:00:00 2001 From: Dion Low Date: Thu, 17 Jul 2025 19:02:37 -0500 Subject: [PATCH 1/4] feat: add jwt token provider to ampersand context - adds jose library to handle jwt token expiration / decoding - adds jwt token provider to ampersand context - refactor api service to use jwt token provider --- package.json | 1 + .../AmpersandContextProvider.tsx | 26 ++- .../AmpersandContextProvider/README.md | 93 ++++++++ src/context/ApiKeyContextProvider.tsx | 9 +- src/context/JwtTokenContextProvider.tsx | 205 ++++++++++++++++++ src/services/api.ts | 59 +++-- yarn.lock | 5 + 7 files changed, 377 insertions(+), 21 deletions(-) create mode 100644 src/context/AmpersandContextProvider/README.md create mode 100644 src/context/JwtTokenContextProvider.tsx diff --git a/package.json b/package.json index 4fe962eaf..740a81aee 100644 --- a/package.json +++ b/package.json @@ -99,6 +99,7 @@ "eslint-plugin-react-refresh": "^0.4.20", "globals": "^16.0.0", "immer": "^10.0.3", + "jose": "^6.0.12", "lodash.isequal": "^4.5.0", "react-tooltip": "^5.28.0" }, diff --git a/src/context/AmpersandContextProvider/AmpersandContextProvider.tsx b/src/context/AmpersandContextProvider/AmpersandContextProvider.tsx index 6d90bf1d8..97e461007 100644 --- a/src/context/AmpersandContextProvider/AmpersandContextProvider.tsx +++ b/src/context/AmpersandContextProvider/AmpersandContextProvider.tsx @@ -10,10 +10,11 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { ApiKeyProvider } from "../ApiKeyContextProvider"; import { ErrorStateProvider } from "../ErrorContextProvider"; +import { JwtTokenProvider } from "../JwtTokenContextProvider"; interface AmpersandProviderProps { options: { - apiKey: string; + apiKey?: string; /** * Use `project` instead of `projectId`. * @deprecated @@ -24,6 +25,11 @@ interface AmpersandProviderProps { */ project?: string; styles?: object; + /** + * Callback function to get a JWT token for authorization. + * This function should return a Promise that resolves to a JWT token string. + */ + getToken?: (consumerRef: string, groupRef: string) => Promise; }; children: React.ReactNode; } @@ -52,7 +58,7 @@ const queryClient = new QueryClient(); export function AmpersandProvider(props: AmpersandProviderProps) { const { - options: { apiKey, projectId, project }, + options: { apiKey, projectId, project, getToken }, children, } = props; const projectIdOrName = project || projectId; @@ -67,8 +73,16 @@ export function AmpersandProvider(props: AmpersandProviderProps) { ); } - if (!apiKey) { - throw new Error("Cannot use AmpersandProvider without an apiKey."); + if (!apiKey && !getToken) { + throw new Error( + "Cannot use AmpersandProvider without an apiKey or getToken.", + ); + } + + if (apiKey && getToken) { + throw new Error( + "Cannot use AmpersandProvider with both apiKey and getToken.", + ); } const contextValue: AmpersandContextValue = { @@ -80,7 +94,9 @@ export function AmpersandProvider(props: AmpersandProviderProps) { - {children} + + {children} + diff --git a/src/context/AmpersandContextProvider/README.md b/src/context/AmpersandContextProvider/README.md new file mode 100644 index 000000000..8b0bb1e12 --- /dev/null +++ b/src/context/AmpersandContextProvider/README.md @@ -0,0 +1,93 @@ +# AmpersandProvider with JWT Token Support + +The `AmpersandProvider` now supports both API key and JWT token authentication methods. + +## Usage + +### API Key Authentication (Existing) + +```tsx +import { AmpersandProvider } from './AmpersandContextProvider'; + +function App() { + return ( + + {/* Your app components */} + + ); +} +``` + +### JWT Token Authentication (New) + +```tsx +import { AmpersandProvider } from './AmpersandContextProvider'; + +function App() { + const getToken = async (consumerRef: string, groupRef: string): Promise => { + // Your custom token retrieval logic here + // This could involve calling your auth service, checking localStorage, etc. + const response = await fetch('/api/auth/token', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ consumerRef, groupRef }), + }); + + const data = await response.json(); + return data.token; + }; + + return ( + + {/* Your app components */} + + ); +} +``` + +## Features + +### JWT Token Caching + +The JWT token provider automatically caches tokens in both memory and localStorage: + +- **Memory Cache**: Fast access for the current session +- **localStorage**: Persistence across browser sessions +- **Automatic Expiration**: Tokens are automatically refreshed when they expire +- **Cache Key**: Tokens are cached using the pattern `{consumerRef}:{groupRef}` + +### Authentication Methods + +The provider supports two mutually exclusive authentication methods: + +1. **API Key**: Traditional API key authentication +2. **JWT Token**: Dynamic token retrieval with caching + +You cannot use both methods simultaneously - the provider will throw an error if both `apiKey` and `getToken` are provided. + +### Error Handling + +- If neither `apiKey` nor `getToken` is provided, the provider will throw an error +- If both are provided, the provider will throw an error +- JWT token retrieval failures are properly handled and logged + +## API Service Integration + +The `useAPI` hook automatically detects which authentication method is available and configures the API service accordingly: + +- **API Key**: Uses `X-Api-Key` header +- **JWT Token**: Uses `Authorization: Bearer {token}` header + +The API service will automatically retrieve fresh tokens when needed, handling the caching and refresh logic transparently. \ No newline at end of file diff --git a/src/context/ApiKeyContextProvider.tsx b/src/context/ApiKeyContextProvider.tsx index 902c76c59..a7511c0c0 100644 --- a/src/context/ApiKeyContextProvider.tsx +++ b/src/context/ApiKeyContextProvider.tsx @@ -1,14 +1,19 @@ import { createContext, useContext } from "react"; +import { JwtTokenContext } from "./JwtTokenContextProvider"; + export const ApiKeyContext = createContext(null); export const ApiKeyProvider = ApiKeyContext.Provider; export const useApiKey = () => { const apiKey = useContext(ApiKeyContext); + const jwtToken = useContext(JwtTokenContext); - if (apiKey === null) { - throw new Error("useApiKey must be used within an ApiKeyProvider"); + if (apiKey === null && jwtToken === null) { + console.error( + "useApiKey must be used within an ApiKeyProvider, or there is no JWT token callback", + ); } return apiKey; diff --git a/src/context/JwtTokenContextProvider.tsx b/src/context/JwtTokenContextProvider.tsx new file mode 100644 index 000000000..866568ff5 --- /dev/null +++ b/src/context/JwtTokenContextProvider.tsx @@ -0,0 +1,205 @@ +import { + createContext, + useCallback, + useContext, + useEffect, + useState, +} from "react"; +import { jwtVerify } from "jose"; + +// Token cache types +interface TokenCacheEntry { + token: string; + expiresAt: number; +} + +const createCacheKey = (consumerRef: string, groupRef: string) => + `${consumerRef}:${groupRef}`; + +type TokenCache = Map; + +interface JwtTokenContextValue { + getToken: (consumerRef: string, groupRef: string) => Promise; + clearToken: () => void; +} + +export const JwtTokenContext = createContext(null); + +interface JwtTokenProviderProps { + getTokenCallback: + | ((consumerRef: string, groupRef: string) => Promise) + | null; + children: React.ReactNode; +} + +// JWT token expiration extraction using jose library +const getTokenExpirationTime = async ( + token: string, +): Promise => { + try { + // Use jose library to decode and verify the JWT token + // We're not verifying the signature since we just want to extract the payload + const decoded = await jwtVerify(token, new Uint8Array(0), { + algorithms: [], // Skip signature verification + }); + + const payload = decoded.payload; + if (payload.exp && typeof payload.exp === "number") { + // JWT exp is in seconds, convert to milliseconds + return payload.exp * 1000; + } + + return null; + } catch (error) { + console.warn("Failed to decode JWT token:", error); + return null; + } +}; + +export function JwtTokenProvider({ + getTokenCallback, + children, +}: JwtTokenProviderProps) { + const [tokenCache, setTokenCache] = useState(new Map()); + + const getCachedToken = useCallback( + (consumerRef: string, groupRef: string): string | null => { + const cacheKey = createCacheKey(consumerRef, groupRef); + const cached = tokenCache.get(cacheKey); + + if (cached && cached.expiresAt > Date.now()) { + return cached.token; + } + + // Remove expired token from cache + if (cached && cached.expiresAt < Date.now()) { + setTokenCache((prev) => { + const newCache = new Map(prev); + newCache.delete(cacheKey); + return newCache; + }); + } + + return null; + }, + [tokenCache], + ); + + const setCachedToken = useCallback( + async (consumerRef: string, groupRef: string, token: string) => { + const cacheKey = createCacheKey(consumerRef, groupRef); + + // Extract actual expiration time from JWT token using jose library + const tokenExpiration = await getTokenExpirationTime(token); + const expiresAt = tokenExpiration || Date.now() + 3600 * 1000; // fallback to 1 hour + + const cacheEntry: TokenCacheEntry = { token, expiresAt }; + + setTokenCache((prev) => new Map(prev).set(cacheKey, cacheEntry)); + + // Also store in localStorage for persistence + try { + localStorage.setItem( + `amp-labs_jwt_${cacheKey}`, + JSON.stringify(cacheEntry), + ); + } catch { + console.warn("Failed to store JWT token in localStorage"); + } + }, + [], + ); + + const clearToken = useCallback(() => { + setTokenCache(new Map()); + // Clear from localStorage + try { + const keys = Object.keys(localStorage); + keys.forEach((key) => { + if (key.startsWith("amp-labs_jwt_")) { + localStorage.removeItem(key); + } + }); + } catch { + console.warn("Failed to clear JWT tokens from localStorage"); + } + }, []); + + // Load cached tokens from localStorage on mount + useEffect(() => { + try { + const newCache: TokenCache = new Map(); + const keys = Object.keys(localStorage); + + keys.forEach((key) => { + if (key.startsWith("amp-labs_jwt_")) { + const cacheKey = key.replace("amp-labs_jwt_", ""); + const stored = localStorage.getItem(key); + if (stored) { + try { + const cacheEntry: TokenCacheEntry = JSON.parse(stored); + if (cacheEntry.expiresAt > Date.now()) { + newCache.set(cacheKey, cacheEntry); + } else { + localStorage.removeItem(key); + } + } catch { + localStorage.removeItem(key); + } + } + } + }); + + if (newCache.size > 0) { + setTokenCache(newCache); + } + } catch { + console.warn("Failed to load JWT tokens from localStorage"); + } + }, []); + + const getToken = useCallback( + async (consumerRef: string, groupRef: string): Promise => { + // First try to get from cache + const cachedToken = getCachedToken(consumerRef, groupRef); + if (cachedToken) return cachedToken; + + if (!getTokenCallback) { + console.error("getTokenCallback is not set"); + throw new Error("getTokenCallback is not set"); + } + + // If not cached, fetch new token + try { + const token = await getTokenCallback(consumerRef, groupRef); + await setCachedToken(consumerRef, groupRef, token); + return token; + } catch { + console.error("Failed to get JWT token"); + throw new Error("Failed to get JWT token"); + } + }, + [getTokenCallback, getCachedToken, setCachedToken], + ); + + const contextValue: JwtTokenContextValue = { + getToken, + clearToken, + }; + + return ( + + {children} + + ); +} + +export const useJwtToken = () => { + const context = useContext(JwtTokenContext); + + if (!context) { + throw new Error("useJwtToken must be used within a JwtTokenProvider"); + } + + return context; +}; diff --git a/src/services/api.ts b/src/services/api.ts index 2dfef39d3..4c16f96c7 100644 --- a/src/services/api.ts +++ b/src/services/api.ts @@ -1,5 +1,5 @@ // currently not using a bundler to support alias imports -import { useCallback } from "react"; +import { useCallback, useContext } from "react"; import { BackfillConfig, BaseWriteConfigObject, @@ -26,6 +26,7 @@ import { UpdateInstallationRequestInstallationConfig, } from "@generated/api/src"; import { useApiKey } from "src/context/ApiKeyContextProvider"; +import { JwtTokenContext } from "src/context/JwtTokenContextProvider"; import { ApiService } from "./ApiService"; import { LIB_VERSION } from "./version"; @@ -110,6 +111,35 @@ export const setApi = (api: ApiService) => { */ export const api = () => apiValue; +// Authentication helper functions +const createApiKeyAuth = (apiKey: string) => ({ + header: "X-Api-Key", + value: apiKey, +}); + +const createJwtAuth = async (jwtToken: any) => { + try { + const token = await jwtToken.getToken("default", "default"); + return { + header: "Authorization", + value: `Bearer ${token}`, + }; + } catch (error) { + console.error("Failed to get JWT token for API authentication:", error); + throw new Error("Failed to authenticate with JWT token"); + } +}; + +const createAuthConfig = (authHeader: string, authValue: string) => + new Configuration({ + basePath: AMP_API_ROOT, + headers: { + "X-Amp-Client": "react", + "X-Amp-Client-Version": LIB_VERSION, + [authHeader]: authValue, + }, + }); + /** * hook to access the API service * @@ -118,25 +148,26 @@ export const api = () => apiValue; */ export function useAPI(): () => Promise { const apiKey = useApiKey(); + const jwtToken = useContext(JwtTokenContext); /** Even though it doesn't need to be be async right now, we want to be able to support other ways * to authenticating to the API in the future which may require async operations */ const getAPI = useCallback(async () => { - if (!apiKey) { - console.error("Unable to create API service without API key."); + if (apiKey) { + const auth = createApiKeyAuth(apiKey); + const configWithAuth = createAuthConfig(auth.header, auth.value); + return new ApiService(configWithAuth); + } + + if (jwtToken) { + const auth = await createJwtAuth(jwtToken); + const configWithAuth = createAuthConfig(auth.header, auth.value); + return new ApiService(configWithAuth); } - const configWithApiKey = new Configuration({ - basePath: AMP_API_ROOT, - headers: { - "X-Amp-Client": "react", - "X-Amp-Client-Version": LIB_VERSION, - "X-Api-Key": apiKey, - }, - }); - - return new ApiService(configWithApiKey); - }, [apiKey]); + console.error("Unable to create API service without API key or JWT token."); + throw new Error("No authentication method provided"); + }, [apiKey, jwtToken]); return getAPI; } diff --git a/yarn.lock b/yarn.lock index 8e532a520..8f8c0bca9 100644 --- a/yarn.lock +++ b/yarn.lock @@ -5761,6 +5761,11 @@ jju@~1.4.0: resolved "https://registry.npmjs.org/jju/-/jju-1.4.0.tgz" integrity sha512-8wb9Yw966OSxApiCt0K3yNJL8pnNeIv+OEq2YMidz4FKP6nonSRoOXc80iXY4JaN2FC11B9qsNmDsm+ZOfMROA== +jose@^6.0.12: + version "6.0.12" + resolved "https://registry.yarnpkg.com/jose/-/jose-6.0.12.tgz#56253d94d46bd784addc4bde3691c323552fe7e4" + integrity sha512-T8xypXs8CpmiIi78k0E+Lk7T2zlK4zDyg+o1CZ4AkOHgDg98ogdP2BeZ61lTFKFyoEwJ9RgAgN+SdM3iPgNonQ== + "js-tokens@^3.0.0 || ^4.0.0", js-tokens@^4.0.0: version "4.0.0" resolved "https://registry.npmjs.org/js-tokens/-/js-tokens-4.0.0.tgz" From 9e574686dd28982d6039ca1c0aade6feb0f6e6cb Mon Sep 17 00:00:00 2001 From: Dion Low Date: Wed, 23 Jul 2025 12:09:09 -0700 Subject: [PATCH 2/4] refactor: use session storage for jwt token cache add installation provider to connect provider add feature flag for jwt auth --- src/components/Connect/ConnectProvider.tsx | 50 ++++++---- .../AmpersandContextProvider/README.md | 93 ------------------- src/context/ApiKeyContextProvider.tsx | 6 +- src/context/JwtTokenContextProvider.tsx | 80 +++++++++------- src/services/api.ts | 43 +++++++-- 5 files changed, 114 insertions(+), 158 deletions(-) delete mode 100644 src/context/AmpersandContextProvider/README.md diff --git a/src/components/Connect/ConnectProvider.tsx b/src/components/Connect/ConnectProvider.tsx index d9eba17e6..9f205f422 100644 --- a/src/components/Connect/ConnectProvider.tsx +++ b/src/components/Connect/ConnectProvider.tsx @@ -1,5 +1,6 @@ import { useCallback } from "react"; import { ConnectionsProvider } from "context/ConnectionsContextProvider"; +import { InstallationProvider } from "src/headless"; import { useForceUpdate } from "src/hooks/useForceUpdate"; import { Connection } from "src/services/api"; @@ -60,26 +61,35 @@ export function ConnectProvider({ return (
- - - - - - - + {/* InstallationProvider is nested in ConnectionsProvider and API service JWT auth */} + + + + + + + + +
); } diff --git a/src/context/AmpersandContextProvider/README.md b/src/context/AmpersandContextProvider/README.md deleted file mode 100644 index 8b0bb1e12..000000000 --- a/src/context/AmpersandContextProvider/README.md +++ /dev/null @@ -1,93 +0,0 @@ -# AmpersandProvider with JWT Token Support - -The `AmpersandProvider` now supports both API key and JWT token authentication methods. - -## Usage - -### API Key Authentication (Existing) - -```tsx -import { AmpersandProvider } from './AmpersandContextProvider'; - -function App() { - return ( - - {/* Your app components */} - - ); -} -``` - -### JWT Token Authentication (New) - -```tsx -import { AmpersandProvider } from './AmpersandContextProvider'; - -function App() { - const getToken = async (consumerRef: string, groupRef: string): Promise => { - // Your custom token retrieval logic here - // This could involve calling your auth service, checking localStorage, etc. - const response = await fetch('/api/auth/token', { - method: 'POST', - headers: { - 'Content-Type': 'application/json', - }, - body: JSON.stringify({ consumerRef, groupRef }), - }); - - const data = await response.json(); - return data.token; - }; - - return ( - - {/* Your app components */} - - ); -} -``` - -## Features - -### JWT Token Caching - -The JWT token provider automatically caches tokens in both memory and localStorage: - -- **Memory Cache**: Fast access for the current session -- **localStorage**: Persistence across browser sessions -- **Automatic Expiration**: Tokens are automatically refreshed when they expire -- **Cache Key**: Tokens are cached using the pattern `{consumerRef}:{groupRef}` - -### Authentication Methods - -The provider supports two mutually exclusive authentication methods: - -1. **API Key**: Traditional API key authentication -2. **JWT Token**: Dynamic token retrieval with caching - -You cannot use both methods simultaneously - the provider will throw an error if both `apiKey` and `getToken` are provided. - -### Error Handling - -- If neither `apiKey` nor `getToken` is provided, the provider will throw an error -- If both are provided, the provider will throw an error -- JWT token retrieval failures are properly handled and logged - -## API Service Integration - -The `useAPI` hook automatically detects which authentication method is available and configures the API service accordingly: - -- **API Key**: Uses `X-Api-Key` header -- **JWT Token**: Uses `Authorization: Bearer {token}` header - -The API service will automatically retrieve fresh tokens when needed, handling the caching and refresh logic transparently. \ No newline at end of file diff --git a/src/context/ApiKeyContextProvider.tsx b/src/context/ApiKeyContextProvider.tsx index a7511c0c0..0a9c98f79 100644 --- a/src/context/ApiKeyContextProvider.tsx +++ b/src/context/ApiKeyContextProvider.tsx @@ -1,6 +1,6 @@ import { createContext, useContext } from "react"; -import { JwtTokenContext } from "./JwtTokenContextProvider"; +import { useJwtToken } from "./JwtTokenContextProvider"; export const ApiKeyContext = createContext(null); @@ -8,9 +8,9 @@ export const ApiKeyProvider = ApiKeyContext.Provider; export const useApiKey = () => { const apiKey = useContext(ApiKeyContext); - const jwtToken = useContext(JwtTokenContext); + const { getToken } = useJwtToken(); - if (apiKey === null && jwtToken === null) { + if (apiKey === null && getToken == null) { console.error( "useApiKey must be used within an ApiKeyProvider, or there is no JWT token callback", ); diff --git a/src/context/JwtTokenContextProvider.tsx b/src/context/JwtTokenContextProvider.tsx index 866568ff5..4027b6a87 100644 --- a/src/context/JwtTokenContextProvider.tsx +++ b/src/context/JwtTokenContextProvider.tsx @@ -13,17 +13,24 @@ interface TokenCacheEntry { expiresAt: number; } +/** + * Create a cache key for the JWT token + * @param consumerRef - The consumer reference + * @param groupRef - The group reference + * @returns The cache key + */ const createCacheKey = (consumerRef: string, groupRef: string) => `${consumerRef}:${groupRef}`; +const SESSION_STORAGE_PREFIX = "amp-labs_jwt_"; + type TokenCache = Map; interface JwtTokenContextValue { - getToken: (consumerRef: string, groupRef: string) => Promise; - clearToken: () => void; + getToken?: (consumerRef: string, groupRef: string) => Promise; } -export const JwtTokenContext = createContext(null); +const JwtTokenContext = createContext(null); interface JwtTokenProviderProps { getTokenCallback: @@ -32,7 +39,11 @@ interface JwtTokenProviderProps { children: React.ReactNode; } -// JWT token expiration extraction using jose library +/** + * JWT token expiration extraction using jose library + * @param token - The JWT token + * @returns The expiration time in milliseconds or null if the token is invalid + */ const getTokenExpirationTime = async ( token: string, ): Promise => { @@ -56,6 +67,12 @@ const getTokenExpirationTime = async ( } }; +/** + * JwtTokenProvider is a context provider for the JWT token + * @param getTokenCallback - The callback function to get the JWT token + * @param children - The children components + * @returns The JwtTokenProvider component + */ export function JwtTokenProvider({ getTokenCallback, children, @@ -85,6 +102,12 @@ export function JwtTokenProvider({ [tokenCache], ); + /** + * Set the token in the cache and sessionStorage + * @param consumerRef - The consumer reference + * @param groupRef - The group reference + * @param token - The JWT token + */ const setCachedToken = useCallback( async (consumerRef: string, groupRef: string, token: string) => { const cacheKey = createCacheKey(consumerRef, groupRef); @@ -97,54 +120,39 @@ export function JwtTokenProvider({ setTokenCache((prev) => new Map(prev).set(cacheKey, cacheEntry)); - // Also store in localStorage for persistence + // Also store in sessionStorage for persistence try { - localStorage.setItem( - `amp-labs_jwt_${cacheKey}`, + sessionStorage.setItem( + `${SESSION_STORAGE_PREFIX}${cacheKey}`, JSON.stringify(cacheEntry), ); } catch { - console.warn("Failed to store JWT token in localStorage"); + console.warn("Failed to store JWT token in sessionStorage"); } }, [], ); - const clearToken = useCallback(() => { - setTokenCache(new Map()); - // Clear from localStorage - try { - const keys = Object.keys(localStorage); - keys.forEach((key) => { - if (key.startsWith("amp-labs_jwt_")) { - localStorage.removeItem(key); - } - }); - } catch { - console.warn("Failed to clear JWT tokens from localStorage"); - } - }, []); - - // Load cached tokens from localStorage on mount + // Load cached tokens from sessionStorage on mount useEffect(() => { try { const newCache: TokenCache = new Map(); - const keys = Object.keys(localStorage); + const keys = Object.keys(sessionStorage); keys.forEach((key) => { - if (key.startsWith("amp-labs_jwt_")) { - const cacheKey = key.replace("amp-labs_jwt_", ""); - const stored = localStorage.getItem(key); + if (key.startsWith(SESSION_STORAGE_PREFIX)) { + const cacheKey = key.replace(SESSION_STORAGE_PREFIX, ""); + const stored = sessionStorage.getItem(key); if (stored) { try { const cacheEntry: TokenCacheEntry = JSON.parse(stored); if (cacheEntry.expiresAt > Date.now()) { newCache.set(cacheKey, cacheEntry); } else { - localStorage.removeItem(key); + sessionStorage.removeItem(key); } } catch { - localStorage.removeItem(key); + sessionStorage.removeItem(key); } } } @@ -154,10 +162,16 @@ export function JwtTokenProvider({ setTokenCache(newCache); } } catch { - console.warn("Failed to load JWT tokens from localStorage"); + console.warn("Failed to load JWT tokens from sessionStorage"); } }, []); + /** + * Get the token from the cache or fetch a new one + * @param consumerRef - The consumer reference + * @param groupRef - The group reference + * @returns The JWT token + */ const getToken = useCallback( async (consumerRef: string, groupRef: string): Promise => { // First try to get from cache @@ -183,8 +197,8 @@ export function JwtTokenProvider({ ); const contextValue: JwtTokenContextValue = { - getToken, - clearToken, + // If getTokenCallback is set, use it to get the token, otherwise return undefined + getToken: getTokenCallback ? getToken : undefined, }; return ( diff --git a/src/services/api.ts b/src/services/api.ts index 4c16f96c7..ace77b28c 100644 --- a/src/services/api.ts +++ b/src/services/api.ts @@ -1,5 +1,5 @@ // currently not using a bundler to support alias imports -import { useCallback, useContext } from "react"; +import { useCallback } from "react"; import { BackfillConfig, BaseWriteConfigObject, @@ -26,7 +26,8 @@ import { UpdateInstallationRequestInstallationConfig, } from "@generated/api/src"; import { useApiKey } from "src/context/ApiKeyContextProvider"; -import { JwtTokenContext } from "src/context/JwtTokenContextProvider"; +import { useJwtToken } from "src/context/JwtTokenContextProvider"; +import { useInstallationProps } from "src/headless/InstallationProvider"; import { ApiService } from "./ApiService"; import { LIB_VERSION } from "./version"; @@ -117,9 +118,8 @@ const createApiKeyAuth = (apiKey: string) => ({ value: apiKey, }); -const createJwtAuth = async (jwtToken: any) => { +const createJwtAuth = (token: string) => { try { - const token = await jwtToken.getToken("default", "default"); return { header: "Authorization", value: `Bearer ${token}`, @@ -140,6 +140,9 @@ const createAuthConfig = (authHeader: string, authValue: string) => }, }); +// TODO: remove this flag when we have a proper JWT auth flow +const ENABLE_JWT_AUTH_FF = false; + /** * hook to access the API service * @@ -148,7 +151,8 @@ const createAuthConfig = (authHeader: string, authValue: string) => */ export function useAPI(): () => Promise { const apiKey = useApiKey(); - const jwtToken = useContext(JwtTokenContext); + const { getToken } = useJwtToken(); + const { consumerRef, groupRef } = useInstallationProps(); /** Even though it doesn't need to be be async right now, we want to be able to support other ways * to authenticating to the API in the future which may require async operations */ @@ -159,15 +163,36 @@ export function useAPI(): () => Promise { return new ApiService(configWithAuth); } - if (jwtToken) { - const auth = await createJwtAuth(jwtToken); + if (getToken) { + if (!ENABLE_JWT_AUTH_FF) { + console.warn( + "JWT authentication is disabled. Please use API key authentication.", + ); + throw new Error( + "JWT authentication is disabled. Please use API key authentication.", + ); + } + + if (!consumerRef || !groupRef) { + console.error( + "Unable to create JWT API service without consumerRef or groupRef.", + { consumerRef, groupRef }, + ); + throw new Error( + "Unable to create JWT API service without consumerRef or groupRef.", + ); + } + const token = await getToken(consumerRef, groupRef); + const auth = createJwtAuth(token); const configWithAuth = createAuthConfig(auth.header, auth.value); return new ApiService(configWithAuth); } console.error("Unable to create API service without API key or JWT token."); - throw new Error("No authentication method provided"); - }, [apiKey, jwtToken]); + throw new Error( + "Unable to create API service without API key or JWT token.", + ); + }, [apiKey, getToken, consumerRef, groupRef]); return getAPI; } From aebb79b8b413b507c8f60ff17f295699191e4e7c Mon Sep 17 00:00:00 2001 From: Dion Low Date: Thu, 31 Jul 2025 12:30:16 -0700 Subject: [PATCH 3/4] fix: api flow --- .../AmpersandContextProvider/AmpersandContextProvider.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/context/AmpersandContextProvider/AmpersandContextProvider.tsx b/src/context/AmpersandContextProvider/AmpersandContextProvider.tsx index 97e461007..081960d53 100644 --- a/src/context/AmpersandContextProvider/AmpersandContextProvider.tsx +++ b/src/context/AmpersandContextProvider/AmpersandContextProvider.tsx @@ -95,7 +95,7 @@ export function AmpersandProvider(props: AmpersandProviderProps) { - {children} + {children} From 50ca8ef327db1ded03a0c56a344981c2a25ca6cb Mon Sep 17 00:00:00 2001 From: Dion Low Date: Thu, 31 Jul 2025 15:36:14 -0700 Subject: [PATCH 4/4] refactor: jwt logic - sessionStorage is now used to cache the JWT token - default token expiration time is 10 minutes - jwt token is now cached in sessionStorage - jwt token is now cached in the in-memory cache - jwt token is now cached in the sessionStorage - jwt token is now cached in the in-memory cache --- src/context/ApiKeyContextProvider.tsx | 10 -- src/context/JwtTokenContextProvider.tsx | 186 ++++++++++++------------ 2 files changed, 90 insertions(+), 106 deletions(-) diff --git a/src/context/ApiKeyContextProvider.tsx b/src/context/ApiKeyContextProvider.tsx index 0a9c98f79..887f6f8cc 100644 --- a/src/context/ApiKeyContextProvider.tsx +++ b/src/context/ApiKeyContextProvider.tsx @@ -1,20 +1,10 @@ import { createContext, useContext } from "react"; -import { useJwtToken } from "./JwtTokenContextProvider"; - export const ApiKeyContext = createContext(null); export const ApiKeyProvider = ApiKeyContext.Provider; export const useApiKey = () => { const apiKey = useContext(ApiKeyContext); - const { getToken } = useJwtToken(); - - if (apiKey === null && getToken == null) { - console.error( - "useApiKey must be used within an ApiKeyProvider, or there is no JWT token callback", - ); - } - return apiKey; }; diff --git a/src/context/JwtTokenContextProvider.tsx b/src/context/JwtTokenContextProvider.tsx index 4027b6a87..f4f377797 100644 --- a/src/context/JwtTokenContextProvider.tsx +++ b/src/context/JwtTokenContextProvider.tsx @@ -7,24 +7,20 @@ import { } from "react"; import { jwtVerify } from "jose"; -// Token cache types interface TokenCacheEntry { token: string; expiresAt: number; } -/** - * Create a cache key for the JWT token - * @param consumerRef - The consumer reference - * @param groupRef - The group reference - * @returns The cache key - */ +const SESSION_STORAGE_PREFIX = "amp-labs_jwt_"; + +const DEFAULT_TOKEN_EXPIRATION_TIME = 10 * 60 * 1000; // 10 minutes + const createCacheKey = (consumerRef: string, groupRef: string) => `${consumerRef}:${groupRef}`; -const SESSION_STORAGE_PREFIX = "amp-labs_jwt_"; - -type TokenCache = Map; +const getSessionStorageKey = (cacheKey: string) => + `${SESSION_STORAGE_PREFIX}${cacheKey}`; interface JwtTokenContextValue { getToken?: (consumerRef: string, groupRef: string) => Promise; @@ -40,27 +36,19 @@ interface JwtTokenProviderProps { } /** - * JWT token expiration extraction using jose library - * @param token - The JWT token - * @returns The expiration time in milliseconds or null if the token is invalid + * Extract JWT token expiration time */ const getTokenExpirationTime = async ( token: string, ): Promise => { try { - // Use jose library to decode and verify the JWT token - // We're not verifying the signature since we just want to extract the payload const decoded = await jwtVerify(token, new Uint8Array(0), { algorithms: [], // Skip signature verification }); - const payload = decoded.payload; - if (payload.exp && typeof payload.exp === "number") { - // JWT exp is in seconds, convert to milliseconds - return payload.exp * 1000; - } - - return null; + return payload.exp && typeof payload.exp === "number" + ? payload.exp * 1000 // jwt expiration is in seconds, convert to milliseconds + : null; } catch (error) { console.warn("Failed to decode JWT token:", error); return null; @@ -68,85 +56,31 @@ const getTokenExpirationTime = async ( }; /** - * JwtTokenProvider is a context provider for the JWT token - * @param getTokenCallback - The callback function to get the JWT token - * @param children - The children components - * @returns The JwtTokenProvider component + * Simplified JWT token provider with cleaner caching logic */ export function JwtTokenProvider({ getTokenCallback, children, }: JwtTokenProviderProps) { - const [tokenCache, setTokenCache] = useState(new Map()); - - const getCachedToken = useCallback( - (consumerRef: string, groupRef: string): string | null => { - const cacheKey = createCacheKey(consumerRef, groupRef); - const cached = tokenCache.get(cacheKey); - - if (cached && cached.expiresAt > Date.now()) { - return cached.token; - } - - // Remove expired token from cache - if (cached && cached.expiresAt < Date.now()) { - setTokenCache((prev) => { - const newCache = new Map(prev); - newCache.delete(cacheKey); - return newCache; - }); - } - - return null; - }, - [tokenCache], - ); - - /** - * Set the token in the cache and sessionStorage - * @param consumerRef - The consumer reference - * @param groupRef - The group reference - * @param token - The JWT token - */ - const setCachedToken = useCallback( - async (consumerRef: string, groupRef: string, token: string) => { - const cacheKey = createCacheKey(consumerRef, groupRef); - - // Extract actual expiration time from JWT token using jose library - const tokenExpiration = await getTokenExpirationTime(token); - const expiresAt = tokenExpiration || Date.now() + 3600 * 1000; // fallback to 1 hour - - const cacheEntry: TokenCacheEntry = { token, expiresAt }; - - setTokenCache((prev) => new Map(prev).set(cacheKey, cacheEntry)); - - // Also store in sessionStorage for persistence - try { - sessionStorage.setItem( - `${SESSION_STORAGE_PREFIX}${cacheKey}`, - JSON.stringify(cacheEntry), - ); - } catch { - console.warn("Failed to store JWT token in sessionStorage"); - } - }, - [], + const [tokenCache, setTokenCache] = useState>( + new Map(), ); // Load cached tokens from sessionStorage on mount useEffect(() => { try { - const newCache: TokenCache = new Map(); - const keys = Object.keys(sessionStorage); + const newCache = new Map(); + const now = Date.now(); // current time in milliseconds - keys.forEach((key) => { + Object.keys(sessionStorage).forEach((key) => { if (key.startsWith(SESSION_STORAGE_PREFIX)) { const cacheKey = key.replace(SESSION_STORAGE_PREFIX, ""); const stored = sessionStorage.getItem(key); + if (stored) { try { const cacheEntry: TokenCacheEntry = JSON.parse(stored); - if (cacheEntry.expiresAt > Date.now()) { + if (cacheEntry.expiresAt > now) { newCache.set(cacheKey, cacheEntry); } else { sessionStorage.removeItem(key); @@ -167,29 +101,92 @@ export function JwtTokenProvider({ }, []); /** - * Get the token from the cache or fetch a new one + * Get a cached token from the in-memory cache or sessionStorage * @param consumerRef - The consumer reference * @param groupRef - The group reference - * @returns The JWT token + * @returns The cached token or null if not found */ + const getCachedToken = useCallback( + (consumerRef: string, groupRef: string): string | null => { + const cacheKey = createCacheKey(consumerRef, groupRef); + const now = Date.now(); + + // Check in-memory cache first + const cached = tokenCache.get(cacheKey); + if (cached && cached.expiresAt > now) { + return cached.token; + } else { + tokenCache.delete(cacheKey); + } + + // Check sessionStorage + const sessionKey = getSessionStorageKey(cacheKey); + const stored = sessionStorage.getItem(sessionKey); + + if (stored) { + try { + const cacheEntry: TokenCacheEntry = JSON.parse(stored); + if (cacheEntry.expiresAt > now) { + // Update in-memory cache + setTokenCache((prev) => new Map(prev).set(cacheKey, cacheEntry)); + return cacheEntry.token; + } else { + sessionStorage.removeItem(sessionKey); + } + } catch { + sessionStorage.removeItem(sessionKey); + } + } + + return null; + }, + [tokenCache], + ); + + const setCachedToken = useCallback( + async (consumerRef: string, groupRef: string, token: string) => { + const cacheKey = createCacheKey(consumerRef, groupRef); + const tokenExpiration = await getTokenExpirationTime(token); + const expiresAt = + tokenExpiration || Date.now() + DEFAULT_TOKEN_EXPIRATION_TIME; + + const cacheEntry: TokenCacheEntry = { token, expiresAt }; + + // Update both in-memory cache + setTokenCache((prev) => new Map(prev).set(cacheKey, cacheEntry)); + + // Update sessionStorage + try { + sessionStorage.setItem( + getSessionStorageKey(cacheKey), + JSON.stringify(cacheEntry), + ); + } catch { + console.warn("Failed to store JWT token in sessionStorage"); + } + }, + [], + ); + const getToken = useCallback( async (consumerRef: string, groupRef: string): Promise => { - // First try to get from cache + // Check all caches first const cachedToken = getCachedToken(consumerRef, groupRef); - if (cachedToken) return cachedToken; + if (cachedToken) { + return cachedToken; + } + // Fetch new token if no callback provided if (!getTokenCallback) { - console.error("getTokenCallback is not set"); - throw new Error("getTokenCallback is not set"); + throw new Error("JWT token callback not provided"); } - // If not cached, fetch new token try { const token = await getTokenCallback(consumerRef, groupRef); await setCachedToken(consumerRef, groupRef, token); return token; - } catch { - console.error("Failed to get JWT token"); + } catch (error) { + console.error("Failed to get JWT token:", error); throw new Error("Failed to get JWT token"); } }, @@ -197,7 +194,6 @@ export function JwtTokenProvider({ ); const contextValue: JwtTokenContextValue = { - // If getTokenCallback is set, use it to get the token, otherwise return undefined getToken: getTokenCallback ? getToken : undefined, }; @@ -210,10 +206,8 @@ export function JwtTokenProvider({ export const useJwtToken = () => { const context = useContext(JwtTokenContext); - if (!context) { throw new Error("useJwtToken must be used within a JwtTokenProvider"); } - return context; };