feat: resource restriction in presign token (#2150)

apps/remix/app/routes/embed+/v1+/authoring+/document.edit.$id.tsx

@@ -41,7 +41,9 @@ export const loader = async ({ request, params }: Route.LoaderArgs) => {
   const token = url.searchParams.get('token') || '';
 
   // We also know that the token is valid, but we need the userId + teamId
-  const result = await verifyEmbeddingPresignToken({ token }).catch(() => null);
+  const result = await verifyEmbeddingPresignToken({ token, scope: `documentId:${id}` }).catch(
+    () => null,
+  );
 
   if (!result) {
     throw new Error('Invalid token');

apps/remix/app/routes/embed+/v1+/authoring+/template.edit.$id.tsx

@@ -41,7 +41,9 @@ export const loader = async ({ request, params }: Route.LoaderArgs) => {
   const token = url.searchParams.get('token') || '';
 
   // We also know that the token is valid, but we need the userId + teamId
-  const result = await verifyEmbeddingPresignToken({ token }).catch(() => null);
+  const result = await verifyEmbeddingPresignToken({ token, scope: `templateId:${id}` }).catch(
+    () => null,
+  );
 
   if (!result) {
     throw new Error('Invalid token');

packages/app-tests/e2e/api/v2/embedding-presign.spec.ts

@@ -1,6 +1,9 @@
 import { expect, test } from '@playwright/test';
+import type { APIRequestContext } from 'playwright-core';
 
 import { NEXT_PUBLIC_WEBAPP_URL } from '@documenso/lib/constants/app';
+import type { CreateEmbeddingPresignTokenOptions } from '@documenso/lib/server-only/embedding-presign/create-embedding-presign-token';
+import type { VerifyEmbeddingPresignTokenOptions } from '@documenso/lib/server-only/embedding-presign/verify-embedding-presign-token';
 import { createApiToken } from '@documenso/lib/server-only/public-api/create-api-token';
 import { seedUser } from '@documenso/prisma/seed/users';
 
@@ -17,18 +20,7 @@ test.describe('Embedding Presign API', () => {
       expiresIn: null,
     });
 
-    const response = await request.post(
+    const response = await createPresignToken(request, token);
-      `${NEXT_PUBLIC_WEBAPP_URL()}/api/v2-beta/embedding/create-presign-token`,
-      {
-        headers: {
-          Authorization: `Bearer ${token}`,
-          'Content-Type': 'application/json',
-        },
-        data: {
-          apiToken: token,
-        },
-      },
-    );
 
     const responseData = await response.json();
 
@@ -54,19 +46,9 @@ test.describe('Embedding Presign API', () => {
       expiresIn: null,
     });
 
-    const response = await request.post(
+    const response = await createPresignToken(request, token, {
-      `${NEXT_PUBLIC_WEBAPP_URL()}/api/v2-beta/embedding/create-presign-token`,
+      expiresIn: 120, // 2 hours
-      {
+    });
-        headers: {
-          Authorization: `Bearer ${token}`,
-          'Content-Type': 'application/json',
-        },
-        data: {
-          apiToken: token,
-          expiresIn: 120, // 2 hours
-        },
-      },
-    );
 
     const responseData = await response.json();
 
@@ -92,19 +74,9 @@ test.describe('Embedding Presign API', () => {
       expiresIn: null,
     });
 
-    const response = await request.post(
+    const response = await createPresignToken(request, token, {
-      `${NEXT_PUBLIC_WEBAPP_URL()}/api/v2-beta/embedding/create-presign-token`,
+      expiresIn: 0, // Immediate expiration
-      {
+    });
-        headers: {
-          Authorization: `Bearer ${token}`,
-          'Content-Type': 'application/json',
-        },
-        data: {
-          apiToken: token,
-          expiresIn: 0, // Immediate expiration
-        },
-      },
-    );
 
     expect(response.ok()).toBeTruthy();
     expect(response.status()).toBe(200);
@@ -129,18 +101,7 @@ test.describe('Embedding Presign API', () => {
     });
 
     // First create a token
-    const createResponse = await request.post(
+    const createResponse = await createPresignToken(request, token);
-      `${NEXT_PUBLIC_WEBAPP_URL()}/api/v2-beta/embedding/create-presign-token`,
-      {
-        headers: {
-          Authorization: `Bearer ${token}`,
-          'Content-Type': 'application/json',
-        },
-        data: {
-          apiToken: token,
-        },
-      },
-    );
 
     expect(createResponse.ok()).toBeTruthy();
     const createResponseData = await createResponse.json();
@@ -150,18 +111,9 @@ test.describe('Embedding Presign API', () => {
     const presignToken = createResponseData.token;
 
     // Then verify it
-    const verifyResponse = await request.post(
+    const verifyResponse = await verifyPresignToken(request, token, {
-      `${NEXT_PUBLIC_WEBAPP_URL()}/api/v2-beta/embedding/verify-presign-token`,
+      token: presignToken,
-      {
+    });
-        headers: {
-          Authorization: `Bearer ${token}`,
-          'Content-Type': 'application/json',
-        },
-        data: {
-          token: presignToken,
-        },
-      },
-    );
 
     expect(verifyResponse.ok()).toBeTruthy();
     expect(verifyResponse.status()).toBe(200);
@@ -183,18 +135,87 @@ test.describe('Embedding Presign API', () => {
       expiresIn: null,
     });
 
-    const response = await request.post(
+    const response = await verifyPresignToken(request, token, {
-      `${NEXT_PUBLIC_WEBAPP_URL()}/api/v2-beta/embedding/verify-presign-token`,
+      token: 'invalid-token',
-      {
+    });
-        headers: {
+
-          Authorization: `Bearer ${token}`,
+    const responseData = await response.json();
-          'Content-Type': 'application/json',
+
-        },
+    console.log('Invalid token response:', responseData);
-        data: {
+
-          token: 'invalid-token',
+    expect(response.ok()).toBeTruthy();
-        },
+    expect(response.status()).toBe(200);
-      },
+
-    );
+    expect(responseData.success).toBe(false);
+  });
+
+  test('verifyEmbeddingPresignToken: should verify a valid scoped token', async ({ request }) => {
+    const { user, team } = await seedUser();
+
+    const { token } = await createApiToken({
+      userId: user.id,
+      teamId: team.id,
+      tokenName: 'test',
+      expiresIn: null,
+    });
+
+    // First create a token
+    const createResponse = await createPresignToken(request, token, {
+      scope: 'documentId:1',
+    });
+
+    expect(createResponse.ok()).toBeTruthy();
+    const createResponseData = await createResponse.json();
+
+    console.log('Create response:', createResponseData);
+
+    const presignToken = createResponseData.token;
+
+    // Then verify it
+    const verifyResponse = await verifyPresignToken(request, token, {
+      token: presignToken,
+      scope: 'documentId:1',
+    });
+
+    expect(verifyResponse.ok()).toBeTruthy();
+    expect(verifyResponse.status()).toBe(200);
+
+    const verifyResponseData = await verifyResponse.json();
+
+    console.log('Verify response:', verifyResponseData);
+
+    expect(verifyResponseData.success).toBe(true);
+  });
+
+  test('verifyEmbeddingPresignToken: should reject a scope mismatched token', async ({
+    request,
+  }) => {
+    const { user, team } = await seedUser();
+
+    const { token } = await createApiToken({
+      userId: user.id,
+      teamId: team.id,
+      tokenName: 'test',
+      expiresIn: null,
+    });
+
+    // First create a token
+    const createResponse = await createPresignToken(request, token, {
+      scope: 'documentId:1',
+    });
+
+    expect(createResponse.ok()).toBeTruthy();
+    const createResponseData = await createResponse.json();
+
+    console.log('Create response:', createResponseData);
+
+    const presignToken = createResponseData.token;
+
+    // Then verify it
+    const response = await verifyPresignToken(request, token, {
+      token: presignToken,
+      scope: 'documentId:2',
+    });
 
     const responseData = await response.json();
 
@@ -206,3 +227,40 @@ test.describe('Embedding Presign API', () => {
     expect(responseData.success).toBe(false);
   });
 });
+
+const createPresignToken = async (
+  request: APIRequestContext,
+  apiToken: string,
+  data?: Partial<CreateEmbeddingPresignTokenOptions>,
+) => {
+  return await request.post(
+    `${NEXT_PUBLIC_WEBAPP_URL()}/api/v2-beta/embedding/create-presign-token`,
+    {
+      headers: {
+        Authorization: `Bearer ${apiToken}`,
+        'Content-Type': 'application/json',
+      },
+      data: {
+        apiToken,
+        ...data,
+      },
+    },
+  );
+};
+
+const verifyPresignToken = async (
+  request: APIRequestContext,
+  apiToken: string,
+  data: VerifyEmbeddingPresignTokenOptions,
+) => {
+  return await request.post(
+    `${NEXT_PUBLIC_WEBAPP_URL()}/api/v2-beta/embedding/verify-presign-token`,
+    {
+      headers: {
+        Authorization: `Bearer ${apiToken}`,
+        'Content-Type': 'application/json',
+      },
+      data,
+    },
+  );
+};

packages/lib/server-only/embedding-presign/create-embedding-presign-token.ts

@@ -12,11 +12,13 @@ export type CreateEmbeddingPresignTokenOptions = {
    * In development mode, can be set to 0 to create a token that expires immediately (for testing)
    */
   expiresIn?: number;
+  scope?: string;
 };
 
 export const createEmbeddingPresignToken = async ({
   apiToken,
   expiresIn,
+  scope,
 }: CreateEmbeddingPresignTokenOptions) => {
   try {
     // Validate the API token
@@ -40,6 +42,7 @@ export const createEmbeddingPresignToken = async ({
     const token = await new SignJWT({
       aud: String(validatedToken.teamId ?? validatedToken.userId),
       sub: String(validatedToken.id),
+      scope,
     })
       .setProtectedHeader({ alg: 'HS256' })
       .setIssuedAt(now.toJSDate())

packages/lib/server-only/embedding-presign/verify-embedding-presign-token.ts

@@ -7,10 +7,12 @@ import { AppError, AppErrorCode } from '../../errors/app-error';
 
 export type VerifyEmbeddingPresignTokenOptions = {
   token: string;
+  scope?: string;
 };
 
 export const verifyEmbeddingPresignToken = async ({
   token,
+  scope,
 }: VerifyEmbeddingPresignTokenOptions) => {
   // First decode the JWT to get the claims without verification
   let decodedToken: JWTPayload;
@@ -81,6 +83,12 @@ export const verifyEmbeddingPresignToken = async ({
     });
   }
 
+  if (decodedToken.scope && scope && decodedToken.scope !== scope) {
+    throw new AppError(AppErrorCode.UNAUTHORIZED, {
+      message: 'Presign token scope not matched',
+    });
+  }
+
   // Now verify the token with the actual secret
   const secret = new TextEncoder().encode(apiToken.token);
 

packages/trpc/server/embedding-router/create-embedding-presign-token.ts

@@ -29,7 +29,7 @@ export const createEmbeddingPresignTokenRoute = procedure
         });
       }
 
-      const { expiresIn } = input;
+      const { expiresIn, scope } = input;
 
       if (IS_BILLING_ENABLED()) {
         const token = await getApiTokenByToken({ token: apiToken });
@@ -54,6 +54,7 @@ export const createEmbeddingPresignTokenRoute = procedure
       const presignToken = await createEmbeddingPresignToken({
         apiToken,
         expiresIn,
+        scope,
       });
 
       return { ...presignToken };

packages/trpc/server/embedding-router/create-embedding-presign-token.types.ts

@@ -21,6 +21,10 @@ export const ZCreateEmbeddingPresignTokenRequestSchema = z.object({
     .optional()
     .default(60)
     .describe('Expiration time in minutes (default: 60, max: 10,080)'),
+  scope: z
+    .string()
+    .optional()
+    .describe('Resource restriction. Example: documentId:1, templateId:2'),
 });
 
 export const ZCreateEmbeddingPresignTokenResponseSchema = z.object({

packages/trpc/server/embedding-router/update-embedding-document.ts

@@ -34,7 +34,10 @@ export const updateEmbeddingDocumentRoute = procedure
         });
       }
 
-      const apiToken = await verifyEmbeddingPresignToken({ token: presignToken });
+      const apiToken = await verifyEmbeddingPresignToken({
+        token: presignToken,
+        scope: `documentId:${input.documentId}`,
+      });
 
       const { documentId, title, externalId, recipients, meta } = input;
 

packages/trpc/server/embedding-router/update-embedding-template.ts

@@ -33,7 +33,10 @@ export const updateEmbeddingTemplateRoute = procedure
         });
       }
 
-      const apiToken = await verifyEmbeddingPresignToken({ token: presignToken });
+      const apiToken = await verifyEmbeddingPresignToken({
+        token: presignToken,
+        scope: `templateId:${input.templateId}`,
+      });
 
       const { templateId, title, externalId, recipients, meta } = input;
 

packages/trpc/server/embedding-router/verify-embedding-presign-token.ts

@@ -17,10 +17,11 @@ export const verifyEmbeddingPresignTokenRoute = procedure
   .output(ZVerifyEmbeddingPresignTokenResponseSchema)
   .mutation(async ({ input }) => {
     try {
-      const { token } = input;
+      const { token, scope } = input;
 
       const apiToken = await verifyEmbeddingPresignToken({
         token,
+        scope,
       }).catch(() => null);
 
       return { success: !!apiToken };

packages/trpc/server/embedding-router/verify-embedding-presign-token.types.ts

@@ -18,6 +18,7 @@ export const ZVerifyEmbeddingPresignTokenRequestSchema = z.object({
     .string()
     .min(1, { message: 'Token is required' })
     .describe('The presign token to verify'),
+  scope: z.string().optional().describe('The scope to verify'),
 });
 
 export const ZVerifyEmbeddingPresignTokenResponseSchema = z.object({