Stop using Response errors when validating API Keys, instead introduce a new "Result" type that has success and failure conditions. Adding in a way to progressively adopt because this touches everything.

ericallam · ericallam · commit a39790287435 · 2024-11-25T15:10:02.000Z
diff --git a/apps/webapp/app/services/apiAuth.server.ts b/apps/webapp/app/services/apiAuth.server.ts
@@ -29,30 +29,67 @@ export type AuthenticatedEnvironment = Optional<
   "orgMember"
 >;
 
-export type ApiAuthenticationResult = {
+export type ApiAuthenticationResult =
+  | ApiAuthenticationResultSuccess
+  | ApiAuthenticationResultFailure;
+
+export type ApiAuthenticationResultSuccess = {
+  ok: true;
   apiKey: string;
   type: "PUBLIC" | "PRIVATE" | "PUBLIC_JWT";
   environment: AuthenticatedEnvironment;
   scopes?: string[];
 };
 
+export type ApiAuthenticationResultFailure = {
+  ok: false;
+  error: string;
+};
+
+/**
+ * @deprecated Use `authenticateApiRequestWithFailure` instead.
+ */
 export async function authenticateApiRequest(
   request: Request,
   options: { allowPublicKey?: boolean; allowJWT?: boolean } = {}
+): Promise<ApiAuthenticationResultSuccess | undefined> {
+  const apiKey = getApiKeyFromRequest(request);
+
+  if (!apiKey) {
+    return;
+  }
+
+  const authentication = await authenticateApiKey(apiKey, options);
+
+  return authentication;
+}
+
+/**
+ * This method is the same as `authenticateApiRequest` but it returns a failure result instead of undefined.
+ * It should be used from now on to ensure that the API key is always validated and provide a failure result.
+ */
+export async function authenticateApiRequestWithFailure(
+  request: Request,
+  options: { allowPublicKey?: boolean; allowJWT?: boolean } = {}
 ): Promise<ApiAuthenticationResult | undefined> {
   const apiKey = getApiKeyFromRequest(request);
 
   if (!apiKey) {
     return;
   }
 
-  return authenticateApiKey(apiKey, options);
+  const authentication = await authenticateApiKeyWithFailure(apiKey, options);
+
+  return authentication;
 }
 
+/**
+ * @deprecated Use `authenticateApiKeyWithFailure` instead.
+ */
 export async function authenticateApiKey(
   apiKey: string,
   options: { allowPublicKey?: boolean; allowJWT?: boolean } = {}
-): Promise<ApiAuthenticationResult | undefined> {
+): Promise<ApiAuthenticationResultSuccess | undefined> {
   const result = getApiKeyResult(apiKey);
 
   if (!result) {
@@ -70,30 +107,120 @@ export async function authenticateApiKey(
   switch (result.type) {
     case "PUBLIC": {
       const environment = await findEnvironmentByPublicApiKey(result.apiKey);
-      if (!environment) return;
+      if (!environment) {
+        return;
+      }
+
       return {
+        ok: true,
         ...result,
         environment,
       };
     }
     case "PRIVATE": {
       const environment = await findEnvironmentByApiKey(result.apiKey);
-      if (!environment) return;
+      if (!environment) {
+        return;
+      }
+
       return {
+        ok: true,
         ...result,
         environment,
       };
     }
     case "PUBLIC_JWT": {
       const validationResults = await validatePublicJwtKey(result.apiKey);
 
-      if (!validationResults) {
+      if (!validationResults.ok) {
         return;
       }
 
       const parsedClaims = ClaimsSchema.safeParse(validationResults.claims);
 
       return {
+        ok: true,
+        ...result,
+        environment: validationResults.environment,
+        scopes: parsedClaims.success ? parsedClaims.data.scopes : [],
+      };
+    }
+  }
+}
+
+/**
+ * This method is the same as `authenticateApiKey` but it returns a failure result instead of undefined.
+ * It should be used from now on to ensure that the API key is always validated and provide a failure result.
+ */
+export async function authenticateApiKeyWithFailure(
+  apiKey: string,
+  options: { allowPublicKey?: boolean; allowJWT?: boolean } = {}
+): Promise<ApiAuthenticationResult> {
+  const result = getApiKeyResult(apiKey);
+
+  if (!result) {
+    return {
+      ok: false,
+      error: "Invalid API Key",
+    };
+  }
+
+  if (!options.allowPublicKey && result.type === "PUBLIC") {
+    return {
+      ok: false,
+      error: "Public API keys are not allowed for this request",
+    };
+  }
+
+  if (!options.allowJWT && result.type === "PUBLIC_JWT") {
+    return {
+      ok: false,
+      error: "Public JWT API keys are not allowed for this request",
+    };
+  }
+
+  switch (result.type) {
+    case "PUBLIC": {
+      const environment = await findEnvironmentByPublicApiKey(result.apiKey);
+      if (!environment) {
+        return {
+          ok: false,
+          error: "Invalid API Key",
+        };
+      }
+
+      return {
+        ok: true,
+        ...result,
+        environment,
+      };
+    }
+    case "PRIVATE": {
+      const environment = await findEnvironmentByApiKey(result.apiKey);
+      if (!environment) {
+        return {
+          ok: false,
+          error: "Invalid API Key",
+        };
+      }
+
+      return {
+        ok: true,
+        ...result,
+        environment,
+      };
+    }
+    case "PUBLIC_JWT": {
+      const validationResults = await validatePublicJwtKey(result.apiKey);
+
+      if (!validationResults.ok) {
+        return validationResults;
+      }
+
+      const parsedClaims = ClaimsSchema.safeParse(validationResults.claims);
+
+      return {
+        ok: true,
         ...result,
         environment: validationResults.environment,
         scopes: parsedClaims.success ? parsedClaims.data.scopes : [],
@@ -207,6 +334,10 @@ export async function authenticatedEnvironmentForAuthentication(
 
   switch (auth.type) {
     case "apiKey": {
+      if (!auth.result.ok) {
+        throw json({ error: auth.result.error }, { status: 401 });
+      }
+
       if (auth.result.environment.project.externalRef !== projectRef) {
         throw json(
           {
@@ -337,6 +468,14 @@ export async function validateJWTTokenAndRenew<T extends z.ZodTypeAny>(
         return;
       }
 
+      if (!authenticatedEnv.ok) {
+        logger.error("Failed to renew JWT token, invalid API key", {
+          error: error.message,
+        });
+
+        return;
+      }
+
       const payload = payloadSchema.safeParse(error.payload);
 
       if (!payload.success) {
diff --git a/apps/webapp/app/services/apiRateLimit.server.ts b/apps/webapp/app/services/apiRateLimit.server.ts
@@ -29,7 +29,7 @@ export const apiRateLimiter = authorizationRateLimitMiddleware({
       allowJWT: true,
     });
 
-    if (!authenticatedEnv) {
+    if (!authenticatedEnv || !authenticatedEnv.ok) {
       return;
     }
 
diff --git a/apps/webapp/app/services/realtime/jwtAuth.server.ts b/apps/webapp/app/services/realtime/jwtAuth.server.ts
@@ -1,58 +1,67 @@
 import { json } from "@remix-run/server-runtime";
 import { validateJWT } from "@trigger.dev/core/v3/jwt";
 import { findEnvironmentById } from "~/models/runtimeEnvironment.server";
+import { AuthenticatedEnvironment } from "../apiAuth.server";
 
-export async function validatePublicJwtKey(token: string) {
+export type ValidatePublicJwtKeySuccess = {
+  ok: true;
+  environment: AuthenticatedEnvironment;
+  claims: Record<string, unknown>;
+};
+
+export type ValidatePublicJwtKeyError = {
+  ok: false;
+  error: string;
+};
+
+export type ValidatePublicJwtKeyResult = ValidatePublicJwtKeySuccess | ValidatePublicJwtKeyError;
+
+export async function validatePublicJwtKey(token: string): Promise<ValidatePublicJwtKeyResult> {
   // Get the sub claim from the token
   // Use the sub claim to find the environment
   // Validate the token against the environment.apiKey
   // Once that's done, return the environment and the claims
   const sub = extractJWTSub(token);
 
   if (!sub) {
-    throw json({ error: "Invalid Public Access Token, missing subject." }, { status: 401 });
+    return { ok: false, error: "Invalid Public Access Token, missing subject." };
   }
 
   const environment = await findEnvironmentById(sub);
 
   if (!environment) {
-    throw json({ error: "Invalid Public Access Token, environment not found." }, { status: 401 });
+    return { ok: false, error: "Invalid Public Access Token, environment not found." };
   }
 
   const result = await validateJWT(token, environment.apiKey);
 
   if (!result.ok) {
     switch (result.code) {
       case "ERR_JWT_EXPIRED": {
-        throw json(
-          {
-            error:
-              "Public Access Token has expired. See https://trigger.dev/docs/frontend/overview#authentication for more information.",
-          },
-          { status: 401 }
-        );
+        return {
+          ok: false,
+          error:
+            "Public Access Token has expired. See https://trigger.dev/docs/frontend/overview#authentication for more information.",
+        };
       }
       case "ERR_JWT_CLAIM_INVALID": {
-        throw json(
-          {
-            error: `Public Access Token is invalid: ${result.error}. See https://trigger.dev/docs/frontend/overview#authentication for more information.`,
-          },
-          { status: 401 }
-        );
+        return {
+          ok: false,
+          error: `Public Access Token is invalid: ${result.error}. See https://trigger.dev/docs/frontend/overview#authentication for more information.`,
+        };
       }
       default: {
-        throw json(
-          {
-            error:
-              "Public Access Token is invalid. See https://trigger.dev/docs/frontend/overview#authentication for more information.",
-          },
-          { status: 401 }
-        );
+        return {
+          ok: false,
+          error:
+            "Public Access Token is invalid. See https://trigger.dev/docs/frontend/overview#authentication for more information.",
+        };
       }
     }
   }
 
   return {
+    ok: true,
     environment,
     claims: result.payload,
   };
diff --git a/apps/webapp/app/services/routeBuilders/apiBuilder.server.ts b/apps/webapp/app/services/routeBuilders/apiBuilder.server.ts
@@ -1,5 +1,8 @@
 import { z } from "zod";
-import { ApiAuthenticationResult, authenticateApiRequest } from "../apiAuth.server";
+import {
+  ApiAuthenticationResultSuccess,
+  authenticateApiRequestWithFailure,
+} from "../apiAuth.server";
 import { ActionFunctionArgs, json, LoaderFunctionArgs } from "@remix-run/server-runtime";
 import { fromZodError } from "zod-validation-error";
 import { apiCors } from "~/utils/apiCors";
@@ -48,7 +51,7 @@ type ApiKeyHandlerFunction<
     ? z.infer<TSearchParamsSchema>
     : undefined;
   headers: THeadersSchema extends z.AnyZodObject ? z.infer<THeadersSchema> : undefined;
-  authentication: ApiAuthenticationResult;
+  authentication: ApiAuthenticationResultSuccess;
   request: Request;
 }) => Promise<Response>;
 
@@ -75,7 +78,7 @@ export function createLoaderApiRoute<
     }
 
     try {
-      const authenticationResult = await authenticateApiRequest(request, { allowJWT });
+      const authenticationResult = await authenticateApiRequestWithFailure(request, { allowJWT });
 
       if (!authenticationResult) {
         return await wrapResponse(
@@ -85,6 +88,14 @@ export function createLoaderApiRoute<
         );
       }
 
+      if (!authenticationResult.ok) {
+        return await wrapResponse(
+          request,
+          json({ error: authenticationResult.error }, { status: 401 }),
+          corsStrategy !== "none"
+        );
+      }
+
       let parsedParams: any = undefined;
       if (paramsSchema) {
         const parsed = paramsSchema.safeParse(params);
@@ -352,7 +363,7 @@ type ApiKeyActionHandlerFunction<
     : undefined;
   headers: THeadersSchema extends z.AnyZodObject ? z.infer<THeadersSchema> : undefined;
   body: TBodySchema extends z.AnyZodObject ? z.infer<TBodySchema> : undefined;
-  authentication: ApiAuthenticationResult;
+  authentication: ApiAuthenticationResultSuccess;
   request: Request;
 }) => Promise<Response>;
 
@@ -396,7 +407,7 @@ export function createActionApiRoute<
 
   async function action({ request, params }: ActionFunctionArgs) {
     try {
-      const authenticationResult = await authenticateApiRequest(request, { allowJWT });
+      const authenticationResult = await authenticateApiRequestWithFailure(request, { allowJWT });
 
       if (!authenticationResult) {
         return await wrapResponse(
@@ -406,6 +417,14 @@ export function createActionApiRoute<
         );
       }
 
+      if (!authenticationResult.ok) {
+        return await wrapResponse(
+          request,
+          json({ error: authenticationResult.error }, { status: 401 }),
+          corsStrategy !== "none"
+        );
+      }
+
       if (maxContentLength) {
         const contentLength = request.headers.get("content-length");
 
diff --git a/apps/webapp/app/v3/handleWebsockets.server.ts b/apps/webapp/app/v3/handleWebsockets.server.ts

Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,7 @@ export const apiRateLimiter = authorizationRateLimitMiddleware({`
`29`	`29`	`allowJWT: true,`
`30`	`30`	`});`
`31`	`31`
`32`		`- if (!authenticatedEnv) {`
	`32`	`+ if (!authenticatedEnv \|\| !authenticatedEnv.ok) {`
`33`	`33`	`return;`
`34`	`34`	`}`
`35`	`35`