Skip to content

Commit

Permalink
fix: implement refresh token flow (#1136)
Browse files Browse the repository at this point in the history
  • Loading branch information
abvthecity authored Jul 12, 2024
1 parent fafd0e8 commit 0086714
Show file tree
Hide file tree
Showing 19 changed files with 390 additions and 132 deletions.
52 changes: 25 additions & 27 deletions packages/ui/app/src/api-playground/PlaygroundAuthorizationForm.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,8 @@ import { isEmpty } from "lodash-es";
import { useRouter } from "next/router";
import { Dispatch, FC, ReactElement, SetStateAction, useCallback, useEffect, useState } from "react";
import { Key } from "react-feather";
import { useApiKey } from "../atoms/auth";
import { useApiKeyInjectionEnabled } from "../services/useApiKeyInjectionEnabled";
import { Callout } from "../mdx/components/Callout";
import { useApiKeyInjectionConfig } from "../services/useApiKeyInjectionConfig";
import { PasswordInputGroup } from "./PasswordInputGroup";
import { PlaygroundSecretsModal, SecretBearer } from "./PlaygroundSecretsModal";
import { PlaygroundRequestFormAuth } from "./types";
Expand Down Expand Up @@ -263,25 +263,28 @@ export function PlaygroundAuthorizationFormCard({
disabled,
}: PlaygroundAuthorizationFormCardProps): ReactElement | null {
const isOpen = useBooleanState(false);
const apiKeyInjectionUrl = useApiKeyInjectionEnabled();
const apiKeyInjection = useApiKeyInjectionConfig();
const router = useRouter();
const apiKey = useApiKey();
const apiKey = apiKeyInjection.enabled && apiKeyInjection.authenticated ? apiKeyInjection.access_token : null;
const [loginError, setLoginError] = useState<string | null>(null);

const redirectOrOpenAuthForm = () => {
if (apiKeyInjectionUrl != null) {
if (apiKeyInjection.enabled && !apiKeyInjection.authenticated) {
// const redirect_uri = encodeURIComponent(
// urlJoin(window.location.origin, basePath ?? "", "/api/fern-docs/auth/login"),
// );
const url = new URL(apiKeyInjectionUrl);
url.searchParams.set("state", encodeURIComponent(window.location.href));
const url = new URL(apiKeyInjection.url);
const state = new URL(window.location.href);
if (state.searchParams.has("loginError")) {
state.searchParams.delete("loginError");
}
url.searchParams.set("state", encodeURIComponent(state.toString()));
window.location.replace(url);
} else {
isOpen.toggleValue();
}
};
const hasApiInjectionConfig = apiKeyInjectionUrl !== undefined;
const authButtonCopy = hasApiInjectionConfig
const authButtonCopy = apiKeyInjection.enabled
? "Login to send a real request"
: "Authenticate with your API key to send a real request";

Expand All @@ -297,26 +300,21 @@ export function PlaygroundAuthorizationFormCard({
}, [router.query, router.isReady]);

// TODO change this login
if (apiKey && authState && authState.type === "bearerAuth") {
if (authState.token === "") {
setAuthorization({ type: "bearerAuth", token: apiKey });
useEffect(() => {
if (apiKey && authState && authState.type === "bearerAuth") {
if (authState.token === "") {
setAuthorization({ type: "bearerAuth", token: apiKey });
}
}
}
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [apiKey]);

return (
<div>
{hasApiInjectionConfig && !apiKey && (
{apiKeyInjection.enabled && !apiKey && (
<>
<FernCard className="rounded-xl p-4 shadow-sm mb-2">
{loginError && (
<FernButton
className="w-full text-left pointer-events-none mb-2"
size="large"
intent="danger"
variant="outlined"
text={loginError}
active={true}
/>
)}
{loginError && <Callout intent="error">{loginError}</Callout>}

<h5 className="t-muted m-0">Login to send a real request</h5>
<div className="flex justify-center my-5 gap-2">
Expand All @@ -340,7 +338,7 @@ export function PlaygroundAuthorizationFormCard({
</>
)}

{hasApiInjectionConfig && apiKey && (
{apiKeyInjection.enabled && apiKey && (
<>
<FernCard className="rounded-xl p-4 shadow-sm mb-3" title="Login to send a real request">
<FernButton
Expand Down Expand Up @@ -378,7 +376,7 @@ export function PlaygroundAuthorizationFormCard({
</>
)}

{!hasApiInjectionConfig && (isAuthed(auth, authState) || apiKey) && (
{!apiKeyInjection.enabled && (isAuthed(auth, authState) || apiKey) && (
<FernButton
className="w-full text-left"
size="large"
Expand All @@ -395,7 +393,7 @@ export function PlaygroundAuthorizationFormCard({
active={isOpen.value}
/>
)}
{!hasApiInjectionConfig && !(isAuthed(auth, authState) || apiKey) && (
{!apiKeyInjection.enabled && !(isAuthed(auth, authState) || apiKey) && (
<FernButton
className="w-full text-left"
size="large"
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,12 +15,12 @@ import { useAtom, useAtomValue } from "jotai";
import { atomWithStorage } from "jotai/utils";
import { isEmpty, round } from "lodash-es";
import { Dispatch, FC, SetStateAction, useEffect, useRef, useState } from "react";
import { useApiKey } from "../atoms/auth";
import { useFeatureFlags } from "../atoms/flags";
import { useDomain } from "../atoms/navigation";
import { IS_MOBILE_SCREEN_ATOM } from "../atoms/viewport";
import { FernErrorTag } from "../components/FernErrorBoundary";
import { ResolvedEndpointDefinition, ResolvedTypeDefinition } from "../resolver/types";
import { useApiKeyInjectionConfig } from "../services/useApiKeyInjectionConfig";
import { PlaygroundAuthorizationFormCard } from "./PlaygroundAuthorizationForm";
import { PlaygroundEndpointForm } from "./PlaygroundEndpointForm";
import { PlaygroundEndpointFormButtons } from "./PlaygroundEndpointFormButtons";
Expand Down Expand Up @@ -64,7 +64,8 @@ export const PlaygroundEndpointContent: FC<PlaygroundEndpointContentProps> = ({

const isMobileScreen = useAtomValue(IS_MOBILE_SCREEN_ATOM);

const apiKey = useApiKey();
const config = useApiKeyInjectionConfig();
const apiKey = config.enabled && config.authenticated ? config.access_token : null;

if (apiKey && formState.auth == null) {
formState.auth = {
Expand Down
6 changes: 0 additions & 6 deletions packages/ui/app/src/atoms/auth.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,3 @@ export const FERN_USER_ATOM = atom<FernUser | undefined>(undefined);
export function useFernUser(): FernUser | undefined {
return useAtomValue(FERN_USER_ATOM);
}

export const API_KEY_ATOM = atom<string | undefined>(undefined);

export function useApiKey(): string | undefined {
return useAtomValue(API_KEY_ATOM);
}
142 changes: 142 additions & 0 deletions packages/ui/app/src/auth/OAuth2Client.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,142 @@
import { JWTPayload, createRemoteJWKSet, decodeJwt, jwtVerify } from "jose";
import { NextApiRequestCookies } from "next/dist/server/api-utils";
import type { NextRequest } from "next/server";
import urlJoin from "url-join";
import { AuthEdgeConfigOAuth2, OAuthTokenResponse, OAuthTokenResponseSchema } from "./types";

interface TokenInfo {
access_token: string;
exp?: number;
refresh_token?: string;
}

export class OAuth2Client {
private readonly clientId: string;
private readonly clientSecret: string;
private readonly environment: string;
private readonly scope: string | undefined;
private readonly jwks: string | undefined;

constructor(
config: AuthEdgeConfigOAuth2,
private readonly redirect_uri?: string,
) {
this.clientId = config.clientId;
this.clientSecret = config.clientSecret;
this.environment = config.environment;
this.scope = config.scope;
this.jwks = config.jwks;
}

public async getToken(code: string): Promise<OAuthTokenResponse> {
const form = new FormData();
form.append("code", code);
form.append("client_secret", this.clientSecret);
form.append("grant_type", "authorization_code");
form.append("client_id", this.clientId);
if (this.redirect_uri != null) {
form.append("redirect_uri", this.redirect_uri);
}

const response = await fetch(urlJoin(this.environment, "/token"), {
method: "POST",
body: form,
});

if (response.ok) {
return OAuthTokenResponseSchema.parse(await response.json());
}
throw new Error(`Failed to get OAuth token: ${response.status} ${await response.text()}`);
}

public async refreshToken(refresh_token: string): Promise<OAuthTokenResponse> {
const form = new FormData();
form.append("refresh_token", refresh_token);
form.append("client_secret", this.clientSecret);
form.append("grant_type", "refresh_token");
form.append("client_id", this.clientId);
if (this.redirect_uri != null) {
form.append("redirect_uri", this.redirect_uri);
}

const response = await fetch(urlJoin(this.environment, "/token"), {
method: "POST",
body: form,
});

if (response.ok) {
return OAuthTokenResponseSchema.parse(await response.json());
}
throw new Error(`Failed to refresh OAuth token: ${response.status} ${await response.text()}`);
}

public getRedirectUrl(state?: string): string {
const url = new URL(urlJoin(this.environment, "/auth"));
url.searchParams.set("response_type", "code");
url.searchParams.set("client_id", this.clientId);
// if (this.redirect_uri != null) {
// url.searchParams.set("redirect_uri", encodeURIComponent(this.redirect_uri));
// }
if (state != null) {
url.searchParams.set("state", state);
}
if (this.scope != null) {
url.searchParams.set("scope", encodeURIComponent(this.scope).replaceAll(/%20/g, "+"));
}
return url.toString();
}

public async decode(access_token: string): Promise<JWTPayload> {
if (this.jwks == null) {
return decodeJwt(access_token);
}
const JWKS = createRemoteJWKSet(new URL(this.jwks));
const { payload } = await jwtVerify(access_token, JWKS);
return payload;
}

public async safeDecode(access_token: string): Promise<JWTPayload | null> {
try {
return await this.decode(access_token);
} catch (e) {
// eslint-disable-next-line no-console
console.error(e);
return null;
}
}

public async getOrRefreshAccessTokenEdge(cookies: NextRequest["cookies"]): Promise<TokenInfo | undefined> {
const access_token = cookies.get("access_token")?.value;
const refresh_token = cookies.get("refresh_token")?.value;
return this.getOrRefreshAccessToken(access_token, refresh_token);
}

public async getOrRefreshAccessTokenNode(cookies: NextApiRequestCookies): Promise<TokenInfo | undefined> {
const access_token = cookies.access_token;
const refresh_token = cookies.refresh_token;
return this.getOrRefreshAccessToken(access_token, refresh_token);
}

private async getOrRefreshAccessToken(
access_token: string | undefined,
refresh_token: string | undefined,
): Promise<TokenInfo | undefined> {
if (access_token != null) {
let payload = await this.safeDecode(access_token);

if (payload == null && refresh_token != null) {
const refreshed = await this.refreshToken(refresh_token);
access_token = refreshed.access_token;
refresh_token = refreshed.refresh_token;
payload = await this.safeDecode(access_token);
}

if (payload == null) {
return undefined;
}

return { access_token, exp: payload.exp, refresh_token };
}
return undefined;
}
}
92 changes: 92 additions & 0 deletions packages/ui/app/src/auth/getApiKeyInjectionConfig.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,92 @@
import type { NextApiRequestCookies } from "next/dist/server/api-utils";
import type { NextRequest } from "next/server";
import { OAuth2Client } from "./OAuth2Client";
import { getAuthEdgeConfig } from "./getAuthEdgeConfig";

interface APIKeyInjectionConfigDisabled {
enabled: false;
}
interface APIKeyInjectionConfigEnabledUnauthorized {
enabled: true;
authenticated: false;
url: string;
}
interface APIKeyInjectionConfigEnabledAuthorized {
enabled: true;
authenticated: true;
access_token: string;
refresh_token?: string;
exp?: number;
}

export type APIKeyInjectionConfig =
| APIKeyInjectionConfigDisabled
| APIKeyInjectionConfigEnabledUnauthorized
| APIKeyInjectionConfigEnabledAuthorized;

export async function getAPIKeyInjectionConfig(
domain: string,
cookies?: NextRequest["cookies"],
state?: string,
): Promise<APIKeyInjectionConfig> {
const config = await getAuthEdgeConfig(domain);
if (config?.type === "oauth2" && config["api-key-injection-enabled"]) {
const client = new OAuth2Client(config, `https://${domain}/api/auth/callback`);
const tokens = cookies != null ? await client.getOrRefreshAccessTokenEdge(cookies) : undefined;

if (tokens != null) {
return {
enabled: true,
authenticated: true,
access_token: tokens.access_token,
refresh_token: tokens.refresh_token,
exp: tokens.exp,
};
}

const url = client.getRedirectUrl(state);
if (url != null) {
return {
enabled: true,
authenticated: false,
url,
};
}
}
return {
enabled: false,
};
}
export async function getAPIKeyInjectionConfigNode(
domain: string,
cookies?: NextApiRequestCookies,
state?: string,
): Promise<APIKeyInjectionConfig> {
const config = await getAuthEdgeConfig(domain);
if (config?.type === "oauth2" && config["api-key-injection-enabled"]) {
const client = new OAuth2Client(config, `https://${domain}/api/auth/callback`);
const tokens = cookies != null ? await client.getOrRefreshAccessTokenNode(cookies) : undefined;

if (tokens != null) {
return {
enabled: true,
authenticated: true,
access_token: tokens.access_token,
refresh_token: tokens.refresh_token,
exp: tokens.exp,
};
}

const url = client.getRedirectUrl(state);
if (url != null) {
return {
enabled: true,
authenticated: false,
url,
};
}
}
return {
enabled: false,
};
}
Loading

0 comments on commit 0086714

Please sign in to comment.