WC-4185 Support variable asset count based on claims from upstream

.changeset/two-comics-win.md

@@ -0,0 +1,5 @@
+---
+"wrangler": minor
+---
+
+Support more files in an upload for non-free users

packages/wrangler/src/__tests__/pages/project-upload.test.ts

@@ -2,6 +2,7 @@
 import { mkdirSync, writeFileSync } from "node:fs";
 import { http, HttpResponse } from "msw";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { maxFileCountAllowedFromClaims } from "../../pages/upload";
 import { endEventLoop } from "../helpers/end-event-loop";
 import { mockAccountId, mockApiToken } from "../helpers/mock-account-id";
 import { mockConsoleMethods } from "../helpers/mock-console";
@@ -665,3 +666,50 @@ describe("pages project upload", () => {
 		expect(std.err).toMatchInlineSnapshot(`""`);
 	});
 });
+
+describe("maxFileCountAllowedFromClaims", () => {
+	it("should return the value from max_file_count_allowed claim when present", () => {
+		// JWT payload: {"max_file_count_allowed": 100000}
+		const jwt =
+			"header." +
+			Buffer.from(JSON.stringify({ max_file_count_allowed: 100000 })).toString(
+				"base64"
+			) +
+			".signature";
+		expect(maxFileCountAllowedFromClaims(jwt)).toBe(100000);
+	});
+
+	it("should return default value when max_file_count_allowed is not a number", () => {
+		// JWT payload: {"max_file_count_allowed": "invalid"}
+		const jwt =
+			"header." +
+			Buffer.from(
+				JSON.stringify({ max_file_count_allowed: "invalid" })
+			).toString("base64") +
+			".signature";
+		expect(maxFileCountAllowedFromClaims(jwt)).toBe(20000);
+	});
+
+	it("should return default value when JWT does not have max_file_count_allowed claim", () => {
+		// JWT payload: {"sub": "user"}
+		const jwt =
+			"header." +
+			Buffer.from(JSON.stringify({ sub: "user" })).toString("base64") +
+			".signature";
+		expect(maxFileCountAllowedFromClaims(jwt)).toBe(20000);
+	});
+
+	it("should return default value for test tokens without parsing", () => {
+		expect(maxFileCountAllowedFromClaims("<<funfetti-auth-jwt>>")).toBe(20000);
+		expect(maxFileCountAllowedFromClaims("<<funfetti-auth-jwt2>>")).toBe(20000);
+		expect(maxFileCountAllowedFromClaims("<<aus-completion-token>>")).toBe(
+			20000
+		);
+	});
+
+	it("should throw error for invalid JWT format", () => {
+		expect(() => maxFileCountAllowedFromClaims("invalid-jwt")).toThrow(
+			"Invalid token:"
+		);
+	});
+});

packages/wrangler/src/__tests__/pages/project-validate.test.ts

@@ -1,6 +1,7 @@
 // /* eslint-disable no-shadow */
 import { writeFileSync } from "node:fs";
 import { afterEach, describe, expect, it, vi } from "vitest";
+import { validate } from "../../pages/validate";
 import { endEventLoop } from "../helpers/end-event-loop";
 import { mockConsoleMethods } from "../helpers/mock-console";
 import { runInTempDir } from "../helpers/run-in-tmp";
@@ -10,7 +11,7 @@ vi.mock("../../pages/constants", async (importActual) => ({
 	// eslint-disable-next-line @typescript-eslint/consistent-type-imports
 	...(await importActual<typeof import("../../pages/constants")>()),
 	MAX_ASSET_SIZE: 1 * 1024 * 1024,
-	MAX_ASSET_COUNT: 10,
+	MAX_ASSET_COUNT_DEFAULT: 10,
 }));
 
 describe("pages project validate", () => {
@@ -54,7 +55,32 @@ describe("pages project validate", () => {
 		await expect(() =>
 			runWrangler("pages project validate .")
 		).rejects.toThrowErrorMatchingInlineSnapshot(
-			`[Error: Error: Pages only supports up to 10 files in a deployment. Ensure you have specified your build output directory correctly.]`
+			`[Error: Error: Pages only supports up to 10 files in a deployment for your current plan. Ensure you have specified your build output directory correctly.]`
+		);
+	});
+
+	it("should succeed with custom fileCountLimit even when exceeding default limit", async () => {
+		// Create 11 files, which exceeds the mocked MAX_ASSET_COUNT_DEFAULT of 10
+		for (let i = 0; i < 11; i++) {
+			writeFileSync(`logo${i}.png`, Buffer.alloc(1));
+		}
+
+		// Should succeed when passing a custom fileCountLimit of 20
+		const fileMap = await validate({ directory: ".", fileCountLimit: 20 });
+		expect(fileMap.size).toBe(11);
+	});
+
+	it("should error with custom fileCountLimit when exceeding custom limit", async () => {
+		// Create 6 files
+		for (let i = 0; i < 6; i++) {
+			writeFileSync(`logo${i}.png`, Buffer.alloc(1));
+		}
+
+		// Should fail when passing a custom fileCountLimit of 5
+		await expect(() =>
+			validate({ directory: ".", fileCountLimit: 5 })
+		).rejects.toThrowError(
+			"Error: Pages only supports up to 5 files in a deployment for your current plan. Ensure you have specified your build output directory correctly."
 		);
 	});
 });

packages/wrangler/src/pages/constants.ts

@@ -2,7 +2,7 @@ import { version as wranglerVersion } from "../../package.json";
 
 const isWindows = process.platform === "win32";
 
-export const MAX_ASSET_COUNT = 20_000;
+export const MAX_ASSET_COUNT_DEFAULT = 20_000;
 export const MAX_ASSET_SIZE = 25 * 1024 * 1024;
 export const PAGES_CONFIG_CACHE_FILENAME = "pages.json";
 export const MAX_BUCKET_SIZE = 40 * 1024 * 1024;

packages/wrangler/src/pages/upload.ts

@@ -13,6 +13,7 @@ import isInteractive from "../is-interactive";
 import { logger } from "../logger";
 import {
 	BULK_UPLOAD_CONCURRENCY,
+	MAX_ASSET_COUNT_DEFAULT,
 	MAX_BUCKET_FILE_COUNT,
 	MAX_BUCKET_SIZE,
 	MAX_CHECK_MISSING_ATTEMPTS,
@@ -59,7 +60,12 @@ export const pagesProjectUploadCommand = createCommand({
 			throw new FatalError("No JWT given.", 1);
 		}
 
-		const fileMap = await validate({ directory });
+		const fileMap = await validate({
+			directory,
+			fileCountLimit: maxFileCountAllowedFromClaims(
+				process.env.CF_PAGES_UPLOAD_JWT
+			),
+		});
 
 		const manifest = await upload({
 			fileMap,
@@ -400,6 +406,38 @@ export const isJwtExpired = (token: string): boolean | undefined => {
 	}
 };
 
+export const maxFileCountAllowedFromClaims = (token: string): number => {
+	// During testing we don't use valid JWTs, so don't try and parse them
+	if (
+		typeof vitest !== "undefined" &&
+		(token === "<<funfetti-auth-jwt>>" ||
+			token === "<<funfetti-auth-jwt2>>" ||
+			token === "<<aus-completion-token>>")
+	) {
+		return MAX_ASSET_COUNT_DEFAULT;
+	}
+	try {
+		// Not validating the JWT here, which ordinarily would be a big red flag.
+		// However, if the JWT is invalid, no uploads (calls to /pages/assets/upload)
+		// will succeed.
+		const decodedJwt = JSON.parse(
+			Buffer.from(token.split(".")[1], "base64").toString()
+		);
+
+		const maxFileCountAllowed = decodedJwt["max_file_count_allowed"];
+		if (typeof maxFileCountAllowed == "number") {
+			return maxFileCountAllowed;
+		}
+
+		return MAX_ASSET_COUNT_DEFAULT;
+	} catch (e) {
+		if (e instanceof Error) {
+			throw new Error(`Invalid token: ${e.message}`);
+		}
+		return MAX_ASSET_COUNT_DEFAULT;
+	}
+};
+
 function formatTime(duration: number) {
 	return `(${(duration / 1000).toFixed(2)} sec)`;
 }

packages/wrangler/src/pages/validate.ts

@@ -5,7 +5,7 @@ import { getType } from "mime";
 import { Minimatch } from "minimatch";
 import prettyBytes from "pretty-bytes";
 import { createCommand } from "../core/create-command";
-import { MAX_ASSET_COUNT, MAX_ASSET_SIZE } from "./constants";
+import { MAX_ASSET_COUNT_DEFAULT, MAX_ASSET_SIZE } from "./constants";
 import { hashFile } from "./hash";
 
 export const pagesProjectValidateCommand = createCommand({
@@ -46,6 +46,7 @@ export type FileContainer = {
 
 export const validate = async (args: {
 	directory: string;
+	fileCountLimit?: number;
 }): Promise<Map<string, FileContainer>> => {
 	const IGNORE_LIST = [
 		"_worker.js",
@@ -68,6 +69,8 @@ export const validate = async (args: {
 	// 	maxMemory = (parsed['max-old-space-size'] ? parsed['max-old-space-size'] : parsed['max_old_space_size']) * 1000 * 1000; // Turn MB into bytes
 	// }
 
+	const fileCountLimit = args.fileCountLimit ?? MAX_ASSET_COUNT_DEFAULT;
+
 	const walk = async (
 		dir: string,
 		fileMap: Map<string, FileContainer> = new Map(),
@@ -124,9 +127,9 @@ export const validate = async (args: {
 
 	const fileMap = await walk(directory);
 
-	if (fileMap.size > MAX_ASSET_COUNT) {
+	if (fileMap.size > fileCountLimit) {
 		throw new FatalError(
-			`Error: Pages only supports up to ${MAX_ASSET_COUNT.toLocaleString()} files in a deployment. Ensure you have specified your build output directory correctly.`,
+			`Error: Pages only supports up to ${fileCountLimit.toLocaleString()} files in a deployment for your current plan. Ensure you have specified your build output directory correctly.`,
 			1
 		);
 	}