Merge pull request #54 from yusufipk/fix/bunny-upload-reservation

fix(uploads): count a Bunny upload from the moment it is admitted

tests/api/bunny-upload-reservation.test.ts

@@ -0,0 +1,263 @@
+// The Bunny init routes used to ask the quota about zero bytes, which meant two
+// things at once: an upload that could never fit was only discovered after it had
+// been sent, and nothing an upload was about to consume was visible to the next
+// request. Bunny reports its own storage on a delay and the figure is cached for
+// two minutes on top, so every init inside that window read the same stale total
+// and every one of them passed.
+//
+// These tests pin the two halves of the fix: the declared size is checked before
+// a byte moves, and it is held as a reservation the next init has to see.
+
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+import { db } from '@/lib/db';
+import {
+  DELETE as cancelProjectBunnyUpload,
+  POST as initProjectBunnyUpload,
+} from '@/app/api/projects/[projectId]/videos/bunny-init/route';
+import { POST as createAsset } from '@/app/api/videos/[videoId]/assets/route';
+import { UPLOAD_RESERVATION_PURPOSES } from '@/lib/storage-quota';
+import { apiRequest, callRoute, readData, readError } from '../helpers/request';
+import { signedInAs } from '../helpers/session';
+import { createVideo, seedProject } from '../factories';
+
+const GIB = BigInt(1024) * BigInt(1024) * BigInt(1024);
+
+/**
+ * Bunny answers every call the same way, with a fresh video id each time so two
+ * inits in one test are distinguishable. The cancel path talks to Bunny too, so
+ * the stub has to cover it rather than just the creation call.
+ */
+function stubBunnyApi(): void {
+  let created = 0;
+  vi.stubGlobal(
+    'fetch',
+    vi.fn(async () => {
+      created += 1;
+      return new Response(JSON.stringify({ guid: `bunnyvideo-${created}-abcdefgh` }), {
+        status: 200,
+        headers: { 'Content-Type': 'application/json' },
+      });
+    })
+  );
+}
+
+beforeEach(() => {
+  vi.stubEnv('OPENFRAME_ENABLE_STRIPE', 'true');
+  vi.stubEnv('OPENFRAME_ENABLE_BUNNY_UPLOADS', 'true');
+  vi.stubEnv('OPENFRAME_ENABLE_S3_VIDEO_UPLOADS', 'false');
+  vi.stubEnv('BUNNY_STREAM_API_KEY', 'test-bunny-key');
+  vi.stubEnv('BUNNY_STREAM_LIBRARY_ID', '999999');
+  vi.stubEnv('BUNNY_UPLOAD_TOKEN_SECRET', 'test-bunny-upload-token-secret');
+  stubBunnyApi();
+});
+
+afterEach(() => {
+  vi.unstubAllGlobals();
+});
+
+function initRequest(projectId: string, body: Record<string, unknown>) {
+  return apiRequest(`/api/projects/${projectId}/videos/bunny-init`, { body });
+}
+
+async function initUpload(projectId: string, sizeBytes: bigint) {
+  return callRoute(
+    initProjectBunnyUpload,
+    initRequest(projectId, { title: 'A clip', sizeBytes: sizeBytes.toString() }),
+    { projectId }
+  );
+}
+
+describe('POST /api/projects/[projectId]/videos/bunny-init', () => {
+  it('refuses an init that does not say how big the upload is', async () => {
+    const scenario = await seedProject();
+    signedInAs(scenario.owner);
+
+    const response = await callRoute(
+      initProjectBunnyUpload,
+      initRequest(scenario.project.id, { title: 'A clip' }),
+      { projectId: scenario.project.id }
+    );
+
+    expect(response.status).toBe(400);
+    expect(await readError(response)).toContain('sizeBytes');
+    expect(await db.uploadReservation.count()).toBe(0);
+  });
+
+  it('refuses a size beyond the host per-file ceiling', async () => {
+    vi.stubEnv('OPENFRAME_MAX_VIDEO_UPLOAD_BYTES', '1024');
+    const scenario = await seedProject();
+    signedInAs(scenario.owner);
+
+    const response = await initUpload(scenario.project.id, BigInt(2048));
+
+    expect(response.status).toBe(400);
+    expect(await readError(response)).toContain('maximum allowed upload size');
+  });
+
+  // The trial ceiling is 3 GiB, so this is refused on the way in rather than
+  // after four gigabytes have been pushed to Bunny.
+  it('refuses an upload the remaining quota cannot hold', async () => {
+    const scenario = await seedProject();
+    signedInAs(scenario.owner);
+
+    const response = await initUpload(scenario.project.id, BigInt(4) * GIB);
+
+    expect(response.status).toBe(507);
+    expect(await db.uploadReservation.count()).toBe(0);
+  });
+
+  it('holds the declared size as a reservation for the workspace owner', async () => {
+    const scenario = await seedProject();
+    signedInAs(scenario.owner);
+
+    const response = await initUpload(scenario.project.id, BigInt(2) * GIB);
+
+    expect(response.status).toBe(200);
+    const reservations = await db.uploadReservation.findMany();
+    expect(reservations).toHaveLength(1);
+    expect(reservations[0].billedUserId).toBe(scenario.owner.id);
+    expect(reservations[0].sizeBytes).toBe(BigInt(2) * GIB);
+  });
+
+  // The regression this whole change exists for. Both of these used to be
+  // granted, because neither could see what the other was about to upload.
+  it('refuses a second upload that no longer fits beside the first', async () => {
+    const scenario = await seedProject();
+    signedInAs(scenario.owner);
+
+    const first = await initUpload(scenario.project.id, BigInt(2) * GIB);
+    const second = await initUpload(scenario.project.id, BigInt(2) * GIB);
+
+    expect(first.status).toBe(200);
+    expect(second.status).toBe(507);
+    expect(await db.uploadReservation.count()).toBe(1);
+  });
+});
+
+describe('DELETE /api/projects/[projectId]/videos/bunny-init', () => {
+  it('gives the quota back when a pending upload is cancelled', async () => {
+    const scenario = await seedProject();
+    signedInAs(scenario.owner);
+
+    const init = await initUpload(scenario.project.id, BigInt(2) * GIB);
+    const { videoId, uploadToken } = await readData(init);
+
+    const response = await callRoute(
+      cancelProjectBunnyUpload,
+      initRequest(scenario.project.id, { videoId, uploadToken }),
+      { projectId: scenario.project.id }
+    );
+
+    expect(response.status).toBe(200);
+    expect(await db.uploadReservation.count()).toBe(0);
+  });
+
+  // Why the reservation id travels inside the signed token instead of being
+  // handed to the client as a field of its own: a caller who could name a
+  // reservation could start two uploads, cancel the cheap one while quoting the
+  // expensive one's reservation, and keep uploading against quota it no longer
+  // holds.
+  it('will not let one upload cancel release another upload reservation', async () => {
+    const scenario = await seedProject();
+    signedInAs(scenario.owner);
+
+    const first = await initUpload(scenario.project.id, BigInt(1) * GIB);
+    const second = await initUpload(scenario.project.id, BigInt(1) * GIB);
+    const firstData = await readData(first);
+    const secondData = await readData(second);
+
+    const response = await callRoute(
+      cancelProjectBunnyUpload,
+      initRequest(scenario.project.id, {
+        videoId: secondData.videoId,
+        uploadToken: firstData.uploadToken,
+      }),
+      { projectId: scenario.project.id }
+    );
+
+    expect(response.status).toBe(403);
+    expect(await db.uploadReservation.count()).toBe(2);
+  });
+});
+
+// The reservation id is not a secret and was never going to be one. The upload
+// token is `base64url(payload).signature`, so the client can read every claim in
+// it, and the two R2 upload routes hand their reservation ids to the client
+// outright. What keeps a hold from being dropped by whoever can name it is that
+// a reservation records what it was opened for, and every finalize route matches
+// on that as well as on the id.
+describe('a Bunny hold cannot be consumed by another flow', () => {
+  /** The claims the client can read out of an upload token without our help. */
+  function claimsOf(uploadToken: string): Record<string, unknown> {
+    return JSON.parse(Buffer.from(uploadToken.split('.')[0], 'base64url').toString('utf8'));
+  }
+
+  it('puts the reservation id somewhere the client can read it', async () => {
+    const scenario = await seedProject();
+    signedInAs(scenario.owner);
+
+    const init = await initUpload(scenario.project.id, BigInt(1) * GIB);
+    const { uploadToken } = await readData(init);
+
+    const reservation = (await db.uploadReservation.findMany())[0];
+    expect(claimsOf(uploadToken).rid).toBe(reservation.id);
+    expect(reservation.purpose).toBe(UPLOAD_RESERVATION_PURPOSES.BUNNY);
+  });
+
+  // The attack the purpose column closes. Creating a YouTube asset costs nothing
+  // and consumes no storage, so quoting a Bunny reservation on one was a way to
+  // hand back the quota of an upload that was still running and then start
+  // another. Repeat and a trial worth three gigabytes uploads as much as it
+  // likes for as long as Bunny takes to report.
+  it('ignores a Bunny reservation id quoted on a YouTube asset create', async () => {
+    const scenario = await seedProject();
+    const video = await createVideo({ projectId: scenario.project.id });
+    signedInAs(scenario.owner);
+
+    const init = await initUpload(scenario.project.id, BigInt(2) * GIB);
+    const { uploadToken } = await readData(init);
+    const reservationId = claimsOf(uploadToken).rid as string;
+
+    const response = await callRoute(
+      createAsset,
+      apiRequest(`/api/videos/${video.id}/assets`, {
+        body: {
+          provider: 'YOUTUBE',
+          sourceUrl: 'https://www.youtube.com/watch?v=dQw4w9WgXcQ',
+          reservationId,
+        },
+      }),
+      { videoId: video.id }
+    );
+
+    expect(response.status).toBe(201);
+    const reservations = await db.uploadReservation.findMany();
+    expect(reservations).toHaveLength(1);
+    expect(reservations[0].id).toBe(reservationId);
+  });
+
+  // And with the hold still standing, the next init has to see it.
+  it('still refuses the next upload after the quoted release attempt', async () => {
+    const scenario = await seedProject();
+    const video = await createVideo({ projectId: scenario.project.id });
+    signedInAs(scenario.owner);
+
+    const init = await initUpload(scenario.project.id, BigInt(2) * GIB);
+    const { uploadToken } = await readData(init);
+
+    await callRoute(
+      createAsset,
+      apiRequest(`/api/videos/${video.id}/assets`, {
+        body: {
+          provider: 'YOUTUBE',
+          sourceUrl: 'https://www.youtube.com/watch?v=dQw4w9WgXcQ',
+          reservationId: claimsOf(uploadToken).rid,
+        },
+      }),
+      { videoId: video.id }
+    );
+
+    const second = await initUpload(scenario.project.id, BigInt(2) * GIB);
+    expect(second.status).toBe(507);
+  });
+});

tests/api/lib-r2-upload-session.test.ts

@@ -11,6 +11,7 @@
 import { describe, expect, it } from 'vitest';
 import { randomUUID } from 'crypto';
 import { db } from '@/lib/db';
+import { UPLOAD_RESERVATION_PURPOSES } from '@/lib/storage-quota';
 import { cancelR2UploadSession, createR2UploadSession } from '@/lib/r2-upload-session';
 import { seedProject } from '../factories';
 
@@ -101,6 +102,7 @@ describe('createR2UploadSession', () => {
       data: {
         billedUserId: scenario.owner.id,
         sizeBytes: BigInt(4096),
+        purpose: UPLOAD_RESERVATION_PURPOSES.R2_VIDEO,
         expiresAt: new Date(Date.now() + HOUR_MS),
       },
     });

tests/api/lib-r2-video-finalize.test.ts

@@ -16,6 +16,7 @@
 import { beforeEach, describe, expect, it, vi } from 'vitest';
 import { randomUUID } from 'crypto';
 import { db } from '@/lib/db';
+import { UPLOAD_RESERVATION_PURPOSES } from '@/lib/storage-quota';
 import { deleteR2Object, deleteVideoObject, headVideoObject, readVideoObjectBytes } from '@/lib/r2';
 import { createR2UploadToken } from '@/lib/r2-upload-token';
 import { createR2UploadSession } from '@/lib/r2-upload-session';
@@ -515,6 +516,7 @@ describe('finalizeR2VideoUpload success', () => {
       data: {
         billedUserId: scenario.owner.id,
         sizeBytes: BigInt(4096),
+        purpose: UPLOAD_RESERVATION_PURPOSES.R2_VIDEO,
         expiresAt: new Date(Date.now() + 30 * 60 * 1000),
       },
     });

tests/api/storage-quota.test.ts

@@ -15,6 +15,7 @@ import { db } from '@/lib/db';
 import { getCachedUserBunnyStorage } from '@/lib/admin-stats';
 import {
   PLAN_STORAGE_LIMIT_BYTES,
+  UPLOAD_RESERVATION_PURPOSES,
   enforceStorageQuota,
   getUserStorageInfo,
   getUserTotalStorageBytes,
@@ -83,7 +84,11 @@ describe('the trial ceiling', () => {
     const user = await createUser();
     await createUploadReservation({ billedUserId: user.id, sizeBytes: BigInt(2) * GIB });
 
-    const result = await reserveStorageQuota(user.id, BigInt(2) * GIB);
+    const result = await reserveStorageQuota(
+      user.id,
+      BigInt(2) * GIB,
+      UPLOAD_RESERVATION_PURPOSES.R2_VIDEO
+    );
 
     expect('error' in result).toBe(true);
     expect((result as { error: Response }).error.status).toBe(507);
@@ -92,7 +97,11 @@ describe('the trial ceiling', () => {
   it('still lets a trial account upload inside its own ceiling', async () => {
     const user = await createUser();
 
-    const result = await reserveStorageQuota(user.id, BigInt(1) * GIB);
+    const result = await reserveStorageQuota(
+      user.id,
+      BigInt(1) * GIB,
+      UPLOAD_RESERVATION_PURPOSES.R2_VIDEO
+    );
 
     expect('reservationId' in result).toBe(true);
   });
@@ -277,7 +286,11 @@ describe('reserveStorageQuota', () => {
   it('writes a reservation row billed to the user with the requested size', async () => {
     const user = await createSubscribedUser();
 
-    const result = await reserveStorageQuota(user.id, BigInt(4096));
+    const result = await reserveStorageQuota(
+      user.id,
+      BigInt(4096),
+      UPLOAD_RESERVATION_PURPOSES.R2_VIDEO
+    );
 
     expect('reservationId' in result).toBe(true);
     const reservation = await db.uploadReservation.findFirstOrThrow();
@@ -291,7 +304,11 @@ describe('reserveStorageQuota', () => {
     vi.stubEnv('OPENFRAME_ENABLE_STRIPE', 'false');
     const user = await createSubscribedUser();
 
-    const result = await reserveStorageQuota(user.id, BigInt(4096));
+    const result = await reserveStorageQuota(
+      user.id,
+      BigInt(4096),
+      UPLOAD_RESERVATION_PURPOSES.R2_VIDEO
+    );
 
     expect(result).toEqual({ reservationId: null });
     expect(await db.uploadReservation.count()).toBe(0);
@@ -304,7 +321,11 @@ describe('reserveStorageQuota', () => {
       sizeBytes: PLAN_STORAGE_LIMIT_BYTES - BigInt(1024),
     });
 
-    const result = await reserveStorageQuota(user.id, BigInt(2048));
+    const result = await reserveStorageQuota(
+      user.id,
+      BigInt(2048),
+      UPLOAD_RESERVATION_PURPOSES.R2_VIDEO
+    );
 
     expect('error' in result).toBe(true);
     expect('error' in result && result.error.status).toBe(507);
@@ -321,7 +342,11 @@ describe('reserveStorageQuota', () => {
       sizeBytes: PLAN_STORAGE_LIMIT_BYTES - BigInt(100),
     });
 
-    const result = await reserveStorageQuota(scenario.owner.id, BigInt(200));
+    const result = await reserveStorageQuota(
+      scenario.owner.id,
+      BigInt(200),
+      UPLOAD_RESERVATION_PURPOSES.R2_VIDEO
+    );
 
     expect('error' in result).toBe(true);
     expect(await db.uploadReservation.count()).toBe(0);
@@ -331,7 +356,11 @@ describe('reserveStorageQuota', () => {
     const user = await createSubscribedUser();
     bunnyStorage({ [user.id]: Number(PLAN_STORAGE_LIMIT_BYTES - BigInt(1024)) });
 
-    const result = await reserveStorageQuota(user.id, BigInt(2048));
+    const result = await reserveStorageQuota(
+      user.id,
+      BigInt(2048),
+      UPLOAD_RESERVATION_PURPOSES.R2_VIDEO
+    );
 
     expect('error' in result).toBe(true);
     expect(await db.uploadReservation.count()).toBe(0);
@@ -345,7 +374,11 @@ describe('reserveStorageQuota', () => {
       expiresInMs: -60_000,
     });
 
-    const result = await reserveStorageQuota(user.id, BigInt(2048));
+    const result = await reserveStorageQuota(
+      user.id,
+      BigInt(2048),
+      UPLOAD_RESERVATION_PURPOSES.R2_VIDEO
+    );
 
     expect('reservationId' in result).toBe(true);
   });
@@ -358,7 +391,11 @@ describe('reserveStorageQuota', () => {
       sizeBytes: PLAN_STORAGE_LIMIT_BYTES - BigInt(1024),
     });
 
-    const result = await reserveStorageQuota(light.id, BigInt(10) * GIB);
+    const result = await reserveStorageQuota(
+      light.id,
+      BigInt(10) * GIB,
+      UPLOAD_RESERVATION_PURPOSES.R2_VIDEO
+    );
 
     expect('reservationId' in result).toBe(true);
   });
@@ -375,8 +412,8 @@ describe('reserveStorageQuota', () => {
     expect(headroom(used)).toBe(BigInt(30) * GIB);
 
     const [first, second] = await Promise.all([
-      reserveStorageQuota(user.id, request),
-      reserveStorageQuota(user.id, request),
+      reserveStorageQuota(user.id, request, UPLOAD_RESERVATION_PURPOSES.R2_VIDEO),
+      reserveStorageQuota(user.id, request, UPLOAD_RESERVATION_PURPOSES.R2_VIDEO),
     ]);
 
     const granted = [first, second].filter((result) => 'reservationId' in result);
@@ -401,8 +438,8 @@ describe('reserveStorageQuota', () => {
     const request = BigInt(20) * GIB;
 
     const results = await Promise.all([
-      reserveStorageQuota(user.id, request),
-      reserveStorageQuota(user.id, request),
+      reserveStorageQuota(user.id, request, UPLOAD_RESERVATION_PURPOSES.R2_VIDEO),
+      reserveStorageQuota(user.id, request, UPLOAD_RESERVATION_PURPOSES.R2_VIDEO),
     ]);
 
     expect(results.every((result) => 'reservationId' in result)).toBe(true);
@@ -417,7 +454,9 @@ describe('reserveStorageQuota', () => {
     const request = BigInt(20) * GIB;
 
     const results = await Promise.all(
-      Array.from({ length: 5 }, () => reserveStorageQuota(user.id, request))
+      Array.from({ length: 5 }, () =>
+        reserveStorageQuota(user.id, request, UPLOAD_RESERVATION_PURPOSES.R2_VIDEO)
+      )
     );
 
     const granted = results.filter((result) => 'reservationId' in result);
@@ -438,8 +477,8 @@ describe('reserveStorageQuota', () => {
     const request = BigInt(150) * GIB;
 
     const results = await Promise.all([
-      reserveStorageQuota(first.id, request),
-      reserveStorageQuota(second.id, request),
+      reserveStorageQuota(first.id, request, UPLOAD_RESERVATION_PURPOSES.R2_VIDEO),
+      reserveStorageQuota(second.id, request, UPLOAD_RESERVATION_PURPOSES.R2_VIDEO),
     ]);
 
     expect(results.every((result) => 'reservationId' in result)).toBe(true);
@@ -450,7 +489,11 @@ describe('reserveStorageQuota', () => {
 describe('releaseStorageReservation', () => {
   it('deletes the reservation and frees the headroom', async () => {
     const user = await createSubscribedUser();
-    const result = await reserveStorageQuota(user.id, BigInt(10) * GIB);
+    const result = await reserveStorageQuota(
+      user.id,
+      BigInt(10) * GIB,
+      UPLOAD_RESERVATION_PURPOSES.R2_VIDEO
+    );
     const reservationId = 'reservationId' in result ? result.reservationId : null;
     expect(reservationId).toBeTruthy();
     expect(await getUserTotalStorageBytes(user.id)).toBe(BigInt(10) * GIB);
@@ -484,6 +527,25 @@ describe('releaseStorageReservation', () => {
 
     expect(await db.uploadReservation.count()).toBe(1);
   });
+
+  // The same scoping in the other direction. Every hold an account owns is
+  // billed to the same user, so `billedUserId` alone does not separate them: an
+  // image being attached would release a Bunny upload that was still in flight
+  // if it could name it, and the ids are readable by the client.
+  it('refuses to delete a reservation opened for a different flow', async () => {
+    const owner = await createSubscribedUser();
+    const reservation = await createUploadReservation({
+      billedUserId: owner.id,
+      sizeBytes: BigInt(4096),
+      purpose: UPLOAD_RESERVATION_PURPOSES.BUNNY,
+    });
+
+    await releaseStorageReservation(reservation.id, owner.id, UPLOAD_RESERVATION_PURPOSES.IMAGE);
+    expect(await db.uploadReservation.count()).toBe(1);
+
+    await releaseStorageReservation(reservation.id, owner.id, UPLOAD_RESERVATION_PURPOSES.BUNNY);
+    expect(await db.uploadReservation.count()).toBe(0);
+  });
 });
 
 describe('GET /api/settings/storage', () => {