Harden upload handling and storage quota checks

apps/api/src/db/repository.ts

@@ -952,7 +952,9 @@ export async function totalMediaSize(input: {
     appId?: string;
 }): Promise<number> {
     const db = getDb();
-    const filters = [eq(media.userId, input.userId), eq(media.temp, false)];
+    // Unsealed uploads occupy storage too. Skipping them let an account stack
+    // uploads past its quota before sealing.
+    const filters = [eq(media.userId, input.userId)];
     if (input.appId) filters.push(eq(media.appId, input.appId));
     const [row] = await db
         .select({ total: sql<number>`coalesce(sum(${media.size}), 0)::float8` })

apps/api/src/mcp/tools/__tests__/upload.test.ts

@@ -0,0 +1,118 @@
+import { before, test } from "node:test";
+import { fileURLToPath } from "node:url";
+import { PGlite } from "@electric-sql/pglite";
+import { drizzle } from "drizzle-orm/pglite";
+import { migrate } from "drizzle-orm/pglite/migrator";
+import { setDb, type AppDb } from "../../../db/client.js";
+import * as schema from "../../../db/schema/index.js";
+import assert from "node:assert/strict";
+import { existsSync } from "node:fs";
+import os from "node:os";
+import path from "node:path";
+
+process.env.MEDIALIT_DEPLOYMENT_MODE = "oss";
+
+let handleUploadMediaTool: typeof import("../upload.js").handleUploadMediaTool;
+let mediaService: typeof import("../../../media/service.js").default;
+
+before(async () => {
+    ({ handleUploadMediaTool } = await import("../upload.js"));
+    ({ default: mediaService } = await import("../../../media/service.js"));
+    // The size/storage check reads totals from the database.
+    const db = drizzle(new PGlite(), { schema });
+    await migrate(db, {
+        migrationsFolder: path.resolve(
+            path.dirname(fileURLToPath(import.meta.url)),
+            "../../../../drizzle",
+        ),
+    });
+    setDb(db as unknown as AppDb);
+});
+
+const PNG_BASE64 = Buffer.from("not-really-a-png").toString("base64");
+
+function authExtra() {
+    return {
+        authInfo: {
+            token: "t",
+            clientId: "c",
+            scopes: ["data:write"],
+            extra: {
+                authKind: "apikey",
+                userId: "user-1",
+                user: { _id: "user-1" },
+                appId: "app-1",
+            },
+        },
+    };
+}
+
+async function uploadWithFileName(fileName: string) {
+    const seen: { tempFilePath: string; existed: boolean; name: string }[] = [];
+    const upload = mediaService.upload;
+    const getMediaDetails = mediaService.getMediaDetails;
+    (mediaService as any).upload = async ({ file }: any) => {
+        seen.push({
+            tempFilePath: file.tempFilePath,
+            existed: existsSync(file.tempFilePath),
+            name: file.name,
+        });
+        return "media-1";
+    };
+    (mediaService as any).getMediaDetails = async () => ({
+        mediaId: "media-1",
+    });
+    try {
+        const result = await handleUploadMediaTool(
+            {
+                fileBase64: PNG_BASE64,
+                fileName,
+                mimeType: "image/png",
+            },
+            authExtra(),
+        );
+        return { result, seen };
+    } finally {
+        (mediaService as any).upload = upload;
+        (mediaService as any).getMediaDetails = getMediaDetails;
+    }
+}
+
+test("upload_media keeps the temp file inside its own temp directory", async () => {
+    const marker = `mcp-traversal-${process.pid}-${Date.now()}`;
+    const escaped = path.join(os.tmpdir(), marker);
+
+    const { result, seen } = await uploadWithFileName(`../${marker}`);
+
+    assert.notEqual((result as any).isError, true);
+    assert.equal(seen.length, 1);
+    assert.equal(seen[0].existed, true);
+    const tempDir = path.dirname(seen[0].tempFilePath);
+    assert.ok(path.basename(tempDir).startsWith("mcp-upload-"));
+    assert.equal(path.dirname(tempDir), path.resolve(os.tmpdir()));
+    assert.equal(existsSync(escaped), false);
+});
+
+test("upload_media handles absolute and nested file names without writing elsewhere", async () => {
+    for (const fileName of [
+        "/tmp/mcp-abs-escape.txt",
+        "a/b/c.txt",
+        "..\\..\\win.txt",
+    ]) {
+        const { result, seen } = await uploadWithFileName(fileName);
+        assert.notEqual((result as any).isError, true, fileName);
+        assert.equal(
+            path.dirname(path.dirname(seen[0].tempFilePath)),
+            path.resolve(os.tmpdir()),
+            fileName,
+        );
+    }
+    assert.equal(existsSync("/tmp/mcp-abs-escape.txt"), false);
+});
+
+test("upload_media cleans up its temp directory", async () => {
+    const { seen } = await uploadWithFileName("photo.png");
+    // Cleanup is fire-and-forget.
+    await new Promise((r) => setTimeout(r, 100));
+    assert.equal(existsSync(path.dirname(seen[0].tempFilePath)), false);
+});

apps/api/src/mcp/tools/upload.ts

@@ -123,7 +123,8 @@ export async function handleUploadMediaTool(args: any, extra: any) {
     }
 
     const tempDir = await mkdtemp(path.join(os.tmpdir(), "mcp-upload-"));
-    const tempPath = path.join(tempDir, args.fileName);
+    // args.fileName is caller-controlled; never use it to build a path.
+    const tempPath = path.join(tempDir, "upload");
 
     try {
         await writeFile(tempPath, buffer);

apps/api/src/media/__tests__/quota.test.ts

@@ -0,0 +1,85 @@
+import { before, test } from "node:test";
+import assert from "node:assert/strict";
+import path from "node:path";
+import { fileURLToPath } from "node:url";
+import { PGlite } from "@electric-sql/pglite";
+import { drizzle } from "drizzle-orm/pglite";
+import { migrate } from "drizzle-orm/pglite/migrator";
+import { setDb, type AppDb } from "../../db/client.js";
+import * as schema from "../../db/schema/index.js";
+import {
+    createAccount,
+    createApp,
+    createMediaRecord,
+} from "../../db/repository.js";
+import { billingStateSource } from "../../billing/entitlements.js";
+import { maxStorageAllowedNotSubscribed } from "../../config/constants.js";
+import { validateUploadConstraints } from "../storage-middleware.js";
+import mediaService from "../service.js";
+
+// Quota is enforced in cloud mode only.
+process.env.MEDIALIT_DEPLOYMENT_MODE = "cloud";
+
+before(async () => {
+    const db = drizzle(new PGlite(), { schema });
+    await migrate(db, {
+        migrationsFolder: path.resolve(
+            path.dirname(fileURLToPath(import.meta.url)),
+            "../../../drizzle",
+        ),
+    });
+    setDb(db as unknown as AppDb);
+    // No paid plan: the account gets the free storage limit.
+    billingStateSource.read = async () => null;
+});
+
+async function accountWithTempUpload(email: string, size: number) {
+    const user = await createAccount({ email, name: email });
+    const app = await createApp({ userId: user.id, name: "app" });
+    await createMediaRecord({
+        fileName: "objects/a",
+        mediaId: `media-${email}`,
+        userId: user.id,
+        appId: app.id,
+        originalFileName: "a.bin",
+        mimeType: "application/octet-stream",
+        size,
+        thumbnailGenerated: false,
+        accessControl: "private",
+        temp: true,
+    });
+    return { user, app };
+}
+
+test("unsealed uploads count against the storage quota", async () => {
+    const limit = maxStorageAllowedNotSubscribed;
+    const { user } = await accountWithTempUpload("a@example.com", limit - 5);
+    const uploadUser = { ...user, _id: user.id } as any;
+
+    const tooBig = await validateUploadConstraints({
+        size: 10,
+        user: uploadUser,
+    });
+    assert.equal(tooBig.valid, false);
+    assert.equal((tooBig as any).reason, "not_enough_storage");
+
+    const fits = await validateUploadConstraints({ size: 5, user: uploadUser });
+    assert.equal(fits.valid, true);
+});
+
+test("sealing is refused when the account is over its quota", async () => {
+    const limit = maxStorageAllowedNotSubscribed;
+    const { user, app } = await accountWithTempUpload(
+        "b@example.com",
+        limit + 1,
+    );
+
+    await assert.rejects(
+        mediaService.sealMedia({
+            userId: user.id,
+            appId: app.id,
+            mediaId: "media-b@example.com",
+        }),
+        /Not enough storage/,
+    );
+});

apps/api/src/media/handlers.ts

@@ -1,5 +1,10 @@
 import Joi from "joi";
-import { FILE_IS_REQUIRED, NOT_FOUND, SUCCESS } from "../config/strings";
+import {
+    FILE_IS_REQUIRED,
+    NOT_ENOUGH_STORAGE,
+    NOT_FOUND,
+    SUCCESS,
+} from "../config/strings";
 import logger from "../services/log";
 import { Request } from "express";
 import mediaService from "./service";
@@ -160,6 +165,7 @@ export async function deleteMedia(req: any, res: any) {
 export function sealErrorStatus(err: { message?: string }): number {
     if (err.message === "Media not found") return 404;
     if (err.message === "Media is already sealed") return 409;
+    if (err.message === NOT_ENOUGH_STORAGE) return 403;
     return 500;
 }
 

apps/api/src/media/service.ts

@@ -27,6 +27,8 @@ import {
     UploadParams,
 } from "../services/s3";
 import logger from "../services/log";
+import { resolveAccountPlan, storageLimitFor } from "../billing/entitlements";
+import { NOT_ENOUGH_STORAGE } from "../config/strings";
 import generateKey from "./utils/generate-key";
 import { getMediaSettings } from "../media-settings/queries";
 import generateFileName from "./utils/generate-file-name";
@@ -37,6 +39,7 @@ import {
     getMediaCount,
     getPaginatedMedia,
     createMedia,
+    getTotalSpace,
 } from "./queries";
 import { listMediaRecordsByIds, sealMediaRecord } from "@/db";
 import * as presignedUrlService from "../signature/service";
@@ -451,6 +454,13 @@ async function sealMedia({
         return await getMediaResponse(media);
     }
 
+    // The total already includes this unsealed upload. Over the limit means
+    // uploads raced past the check at upload time.
+    const plan = await resolveAccountPlan({ id: String(userId) });
+    if ((await getTotalSpace({ userId })) > storageLimitFor(plan)) {
+        throw new Error(NOT_ENOUGH_STORAGE);
+    }
+
     const fileExtension = path.extname(media.fileName).replace(".", "");
 
     // Get tags from source object (in private bucket)

apps/api/src/media/utils/__tests__/generate-file-name.test.ts

@@ -0,0 +1,32 @@
+import { test } from "node:test";
+import assert from "node:assert/strict";
+import generateFileName from "../generate-file-name.js";
+
+test("keeps a plain alphanumeric extension", () => {
+    assert.equal(generateFileName("photo.PNG").ext, "PNG");
+    assert.equal(generateFileName("a.b.mp4").ext, "mp4");
+});
+
+test("falls back to bin for extensions that are not safe in paths or shells", () => {
+    for (const name of [
+        "x.mp4$(touch /tmp/pwn)",
+        'x.mp4"; touch /tmp/pwn; echo "',
+        "x.mp4`id`",
+        "x.mp4;id",
+        "x.mp4 && id",
+        "x.../../etc",
+        "x.",
+        "x.abcdefghijk",
+        "noextension",
+        "",
+    ]) {
+        assert.equal(generateFileName(name).ext, "bin", name);
+    }
+});
+
+test("always generates a fresh media id", () => {
+    assert.notEqual(
+        generateFileName("a.png").name,
+        generateFileName("a.png").name,
+    );
+});

apps/api/src/media/utils/__tests__/thumbnail-shell.test.ts

@@ -0,0 +1,26 @@
+import { test } from "node:test";
+import assert from "node:assert/strict";
+import { existsSync, mkdtempSync, rmSync } from "node:fs";
+import os from "node:os";
+import path from "node:path";
+import thumbnail from "@medialit/thumbnail";
+
+test("forVideo does not run shell syntax found in the source path", async () => {
+    const dir = mkdtempSync(path.join(os.tmpdir(), "thumb-shell-"));
+    const marker = path.join(dir, "PWNED");
+    try {
+        for (const source of [
+            path.join(dir, `main.mp4$(touch ${marker})`),
+            path.join(dir, `main.mp4"; touch ${marker}; echo "`),
+            path.join(dir, `main.mp4\`touch ${marker}\``),
+        ]) {
+            // The file does not exist, so ffmpeg fails; that is fine.
+            await assert.rejects(
+                thumbnail.forVideo(source, path.join(dir, "thumb.webp")),
+            );
+        }
+        assert.equal(existsSync(marker), false);
+    } finally {
+        rmSync(dir, { recursive: true, force: true });
+    }
+});

apps/api/src/media/utils/generate-file-name.ts

@@ -10,9 +10,14 @@ export default function generateFileName(
     filename: string,
 ): FileNameWithExtention {
     const extention = filename.split(".");
+    const ext = extention[extention.length - 1];
 
     return {
         name: getUniqueId(MEDIA_ID_LENGTH),
-        ext: extention[extention.length - 1],
+        // Extension ends up in file paths and S3 keys; never trust it raw.
+        ext:
+            extention.length > 1 && /^[A-Za-z0-9]{1,10}$/.test(ext)
+                ? ext
+                : "bin",
     };
 }

apps/api/src/tus/finalize.ts

@@ -86,6 +86,7 @@ export default async function finalizeUpload(
     }
 
     let fileExtension = path.extname(metadata.fileName).replace(".", "");
+    if (!/^[A-Za-z0-9]{1,10}$/.test(fileExtension)) fileExtension = "bin";
     let mimeType = metadata.mimeType;
     if (useWebP && imagePattern.test(mimeType)) {
         fileExtension = "webp";

packages/integration-tests/src/cli.ts

@@ -330,7 +330,8 @@ export async function runCliIntegrationTests(
         });
         pending.add(tempId);
         assert.deepEqual(await list(), [], "--temp uploads must stay hidden");
-        await checkTotals(0, 0);
+        // The draft takes up storage but is not counted as media.
+        await checkTotals(png.length, 0);
         log("PASS upload --temp leaves the upload temporary and hidden");
 
         const sealed = object(await json(["seal", tempId]), "sealed media");

packages/integration-tests/src/mcp.ts

@@ -518,7 +518,8 @@ export async function runMcpIntegrationTests(
             [],
             "Draft MCP uploads must stay hidden",
         );
-        await checkTotals(0, 0);
+        // The draft takes up storage but is not counted as media.
+        await checkTotals(png.length, 0);
         log("PASS upload_media and draft visibility");
 
         const sealed = await client.callTool("seal_media", { mediaId });

packages/integration-tests/src/rest.ts

@@ -348,7 +348,8 @@ export async function runIntegrationTests(
         );
         checkMedia(retrieved, expected);
         await download(retrieved);
-        await checkTotals(0, 0); // Drafts are not included in the public totals.
+        // Drafts take up storage but are not counted as media.
+        await checkTotals(png.length * uploads.length, 0);
         log(`PASS upload and get ${label} (${access})`);
     }
 
@@ -390,7 +391,7 @@ export async function runIntegrationTests(
                 ),
                 expected,
             );
-            await checkTotals(png.length * (index + 1), index + 1);
+            await checkTotals(png.length * uploads.length, index + 1);
             log(`PASS seal ${expected.filename} and storage decrement`);
         }
         const listed = await list({ group });