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/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/AmpersandContextProvider.tsx b/src/context/AmpersandContextProvider/AmpersandContextProvider.tsx index 6d90bf1d8..081960d53 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/ApiKeyContextProvider.tsx b/src/context/ApiKeyContextProvider.tsx index 902c76c59..887f6f8cc 100644 --- a/src/context/ApiKeyContextProvider.tsx +++ b/src/context/ApiKeyContextProvider.tsx @@ -6,10 +6,5 @@ export const ApiKeyProvider = ApiKeyContext.Provider; export const useApiKey = () => { const apiKey = useContext(ApiKeyContext); - - if (apiKey === null) { - throw new Error("useApiKey must be used within an ApiKeyProvider"); - } - return apiKey; }; diff --git a/src/context/JwtTokenContextProvider.tsx b/src/context/JwtTokenContextProvider.tsx new file mode 100644 index 000000000..f4f377797 --- /dev/null +++ b/src/context/JwtTokenContextProvider.tsx @@ -0,0 +1,213 @@ +import { + createContext, + useCallback, + useContext, + useEffect, + useState, +} from "react"; +import { jwtVerify } from "jose"; + +interface TokenCacheEntry { + token: string; + expiresAt: number; +} + +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 getSessionStorageKey = (cacheKey: string) => + `${SESSION_STORAGE_PREFIX}${cacheKey}`; + +interface JwtTokenContextValue { + getToken?: (consumerRef: string, groupRef: string) => Promise; +} + +const JwtTokenContext = createContext(null); + +interface JwtTokenProviderProps { + getTokenCallback: + | ((consumerRef: string, groupRef: string) => Promise) + | null; + children: React.ReactNode; +} + +/** + * Extract JWT token expiration time + */ +const getTokenExpirationTime = async ( + token: string, +): Promise => { + try { + const decoded = await jwtVerify(token, new Uint8Array(0), { + algorithms: [], // Skip signature verification + }); + const payload = decoded.payload; + 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; + } +}; + +/** + * Simplified JWT token provider with cleaner caching logic + */ +export function JwtTokenProvider({ + getTokenCallback, + children, +}: JwtTokenProviderProps) { + const [tokenCache, setTokenCache] = useState>( + new Map(), + ); + + // Load cached tokens from sessionStorage on mount + useEffect(() => { + try { + const newCache = new Map(); + const now = Date.now(); // current time in milliseconds + + 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 > now) { + newCache.set(cacheKey, cacheEntry); + } else { + sessionStorage.removeItem(key); + } + } catch { + sessionStorage.removeItem(key); + } + } + } + }); + + if (newCache.size > 0) { + setTokenCache(newCache); + } + } catch { + console.warn("Failed to load JWT tokens from sessionStorage"); + } + }, []); + + /** + * Get a cached token from the in-memory cache or sessionStorage + * @param consumerRef - The consumer reference + * @param groupRef - The group reference + * @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 => { + // Check all caches first + const cachedToken = getCachedToken(consumerRef, groupRef); + if (cachedToken) { + return cachedToken; + } + + // Fetch new token if no callback provided + if (!getTokenCallback) { + throw new Error("JWT token callback not provided"); + } + + try { + const token = await getTokenCallback(consumerRef, groupRef); + await setCachedToken(consumerRef, groupRef, token); + return token; + } catch (error) { + console.error("Failed to get JWT token:", error); + throw new Error("Failed to get JWT token"); + } + }, + [getTokenCallback, getCachedToken, setCachedToken], + ); + + const contextValue: JwtTokenContextValue = { + getToken: getTokenCallback ? getToken : undefined, + }; + + 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..ace77b28c 100644 --- a/src/services/api.ts +++ b/src/services/api.ts @@ -26,6 +26,8 @@ import { UpdateInstallationRequestInstallationConfig, } from "@generated/api/src"; import { useApiKey } from "src/context/ApiKeyContextProvider"; +import { useJwtToken } from "src/context/JwtTokenContextProvider"; +import { useInstallationProps } from "src/headless/InstallationProvider"; import { ApiService } from "./ApiService"; import { LIB_VERSION } from "./version"; @@ -110,6 +112,37 @@ export const setApi = (api: ApiService) => { */ export const api = () => apiValue; +// Authentication helper functions +const createApiKeyAuth = (apiKey: string) => ({ + header: "X-Api-Key", + value: apiKey, +}); + +const createJwtAuth = (token: string) => { + try { + 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, + }, + }); + +// TODO: remove this flag when we have a proper JWT auth flow +const ENABLE_JWT_AUTH_FF = false; + /** * hook to access the API service * @@ -118,25 +151,48 @@ export const api = () => apiValue; */ export function useAPI(): () => Promise { const apiKey = useApiKey(); + 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 */ 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 (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); } - 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( + "Unable to create API service without API key or JWT token.", + ); + }, [apiKey, getToken, consumerRef, groupRef]); 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"