separate code jwt input and output schemas

Author: jasonhedman

packages/app/control/src/app/(auth)/(oauth)/(authorize)/_lib/authorize.ts

@@ -3,13 +3,10 @@ import { z } from 'zod';
 import { nanoid } from 'nanoid';
 
 import { auth } from '@/auth';
-import { authCodeJwtBodySchema, createAuthCodeJwt } from './code';
+import { authCodeJwtInputSchema, createAuthCodeJwt } from './code';
 
-export const authorizeParamsSchema = authCodeJwtBodySchema
+export const authorizeParamsSchema = authCodeJwtInputSchema
   .omit({
-    code: true,
-    exp: true,
-    iat: true,
     user_id: true,
   })
   .extend({

packages/app/control/src/app/(auth)/(oauth)/(authorize)/_lib/code.ts

@@ -10,7 +10,7 @@ import { addSeconds, getUnixTime } from 'date-fns';
 const AUTH_CODE_TTL = 300; // 5 minutes
 const JWT_SECRET = new TextEncoder().encode(env.OAUTH_CODE_SIGNING_JWT_SECRET);
 
-export const authCodeJwtBodySchema = z.object({
+export const authCodeJwtInputSchema = z.object({
   client_id: z.uuid('client_id must be a valid UUID'),
   user_id: z.uuid('user_id must be a valid UUID'),
   redirect_uri: z.url('redirect_uri must be a valid URL'),
@@ -25,20 +25,26 @@ export const authCodeJwtBodySchema = z.object({
     error: 'Only S256 code challenge method is supported',
   }),
   scope: z.string('missing scope').default('llm:invoke offline_access'),
-  code: z.string().default(nanoid(32)),
-  exp: z
-    .number()
-    .default(() => getUnixTime(addSeconds(new Date(), AUTH_CODE_TTL))),
-  iat: z.number().default(() => getUnixTime(new Date())),
 });
 
 export const createAuthCodeJwt = async (
-  payloadInput: z.input<typeof authCodeJwtBodySchema>
+  payloadInput: z.input<typeof authCodeJwtInputSchema>
 ) => {
-  const payload = authCodeJwtBodySchema.parse(payloadInput);
-  const authCodeJwt = await new SignJWT(payload)
+  const input = authCodeJwtInputSchema.parse(payloadInput);
+  const authCodeJwt = await new SignJWT({
+    ...input,
+    code: nanoid(32),
+  })
     .setProtectedHeader({ alg: 'HS256' })
+    .setExpirationTime(getUnixTime(addSeconds(new Date(), AUTH_CODE_TTL)))
+    .setIssuedAt(getUnixTime(new Date()))
     .sign(JWT_SECRET);
 
   return authCodeJwt;
 };
+
+export const authCodeJwtPayloadSchema = authCodeJwtInputSchema.extend({
+  code: z.string(),
+  exp: z.number(),
+  iat: z.number(),
+});

packages/app/control/src/app/(auth)/(oauth)/api/oauth/token/_lib/create-jwt.ts

@@ -8,7 +8,6 @@ import { addSeconds, getUnixTime } from 'date-fns';
 
 import { env } from '@/env';
 import { logger } from '@/logger';
-import { Prisma } from '@/generated/prisma';
 
 const API_ECHO_ACCESS_JWT_SECRET = new TextEncoder().encode(
   env.API_ECHO_ACCESS_JWT_SECRET

packages/app/control/src/app/(auth)/(oauth)/api/oauth/token/_lib/issue.ts

@@ -3,7 +3,7 @@ import z from 'zod';
 import { createHash } from 'crypto';
 import { jwtVerify } from 'jose';
 
-import { authCodeJwtBodySchema } from '@/app/(auth)/(oauth)/(authorize)/_lib/code';
+import { authCodeJwtPayloadSchema } from '@/app/(auth)/(oauth)/(authorize)/_lib/code';
 import { isValidRedirectUri } from '@/app/(auth)/(oauth)/_lib/redirect-uri';
 
 import { db } from '@/lib/db';
@@ -57,7 +57,7 @@ export async function handleIssueToken(
     code_challenge,
     user_id,
     scope,
-  } = authCodeJwtBodySchema.parse(payload);
+  } = authCodeJwtPayloadSchema.parse(payload);
 
   /* 4️⃣ Validate the authorization code data */
   if (codeClientId !== client_id || codeRedirectUri !== redirect_uri) {