refactor: upload/partial logic + more sanitzation

src/lib/api/upload.ts

@@ -0,0 +1,116 @@
+import { extname } from 'path';
+import { User } from '@/lib/db/models/user';
+import { prisma } from '@/lib/db';
+import { bytes } from '@/lib/bytes';
+import { config } from '../config';
+import { Config } from '../config/validate';
+import { sanitizeFilename } from '../fs';
+import { formatFileName } from '../uploader/formatFileName';
+import { guess } from '../mimes';
+
+const commonDoubleExts = [
+  '.tar.gz',
+  '.tar.xz',
+  '.tar.bz2',
+  '.tar.lz',
+  '.tar.lzma',
+  '.tar.Z',
+  '.tar.7z',
+  '.zip.gz',
+  '.zip.xz',
+  '.rar.gz',
+  '.log.gz',
+  '.csv.gz',
+  '.pdf.gz',
+  // feel free to PR more
+];
+
+export function getExtension(filename: string, override?: string) {
+  return override ?? commonDoubleExts.find((ext) => filename.endsWith(ext)) ?? extname(filename);
+}
+
+export async function checkQuota(
+  user: User | null,
+  newSize: number,
+  fileCount: number,
+): Promise<true | string> {
+  if (!user?.quota) return true;
+
+  const stats = await prisma.file.aggregate({
+    where: {
+      userId: user.id,
+    },
+    _sum: {
+      size: true,
+    },
+    _count: {
+      _all: true,
+    },
+  });
+
+  const aggSize = stats?._sum?.size ? stats._sum.size : 0n;
+
+  if (user.quota.filesQuota === 'BY_BYTES' && Number(aggSize) + newSize > bytes(user.quota.maxBytes!))
+    return `uploading will exceed your storage quota of ${user.quota.maxFiles} files`;
+
+  if (user.quota.filesQuota === 'BY_FILES' && stats?._count?._all + fileCount > user.quota.maxFiles!)
+    return `uploading will exceed your file count quota of ${user.quota.maxFiles} files`;
+
+  return true;
+}
+
+export function getDomain(
+  overrideDomain?: string | null,
+  defaultDomain?: string | null,
+  hostDomain?: string,
+) {
+  const base = `${config.core.returnHttpsUrls ? 'https' : 'http'}://`;
+
+  if (overrideDomain) return base + overrideDomain;
+  if (defaultDomain) return base + defaultDomain;
+
+  // using localhost as a fallback in the 1% chance theres no host header
+  return base + (hostDomain ?? 'localhost');
+}
+
+export async function getFilename(
+  format: Config['files']['defaultFormat'],
+  originalName: string,
+  extension: string,
+  override?: string,
+): Promise<{ error: string } | { fileName: string }> {
+  let fileName = override ? sanitizeFilename(override) : formatFileName(format, originalName);
+  if (!fileName) return { error: 'invalid file name' };
+
+  let fullFileName = `${fileName}${extension}`;
+  let existing = await prisma.file.findFirst({ where: { name: fullFileName } });
+
+  if (existing && (override || format === 'name')) {
+    return { error: 'file with the same name already exists' };
+  }
+
+  while (existing && format === 'random') {
+    fileName = formatFileName(format, originalName);
+    if (!fileName) return { error: 'invalid file name' };
+
+    fullFileName = `${fileName}${extension}`;
+    existing = await prisma.file.findFirst({ where: { name: fullFileName } });
+  }
+
+  return { fileName };
+}
+
+export async function getMimetype(
+  originalMimetype: string,
+  extension: string,
+): Promise<{ mimetype: string; assumed: boolean }> {
+  const mimetype = originalMimetype;
+
+  if (config.files.assumeMimetypes) {
+    const mime = await guess(extension.substring(1));
+
+    if (mime) return { mimetype: mime, assumed: true };
+  }
+
+  return { mimetype, assumed: false };
+}

src/lib/uploader/formatFileName.ts

@@ -15,8 +15,10 @@ export function formatFileName(nameFormat: Config['files']['defaultFormat'], ori
     case 'uuid':
       return randomUUID({ disableEntropyCache: true });
     case 'name':
-      const { name } = parse(originalName!);
+      const sanitized = originalName ? parse(originalName).name : null;
+      if (!sanitized) return null;
 
+      const { name } = parse(sanitized);
       return name;
     case 'random-words':
     case 'gfycat':

src/server/routes/api/upload/index.ts

@@ -1,3 +1,4 @@
+import { checkQuota, getDomain, getExtension, getFilename, getMimetype } from '@/lib/api/upload';
 import { bytes } from '@/lib/bytes';
 import { compressFile, CompressResult } from '@/lib/compress';
 import { config } from '@/lib/config';
@@ -8,36 +9,12 @@ import { fileSelect } from '@/lib/db/models/file';
 import { sanitizeFilename } from '@/lib/fs';
 import { removeGps } from '@/lib/gps';
 import { log } from '@/lib/logger';
-import { guess } from '@/lib/mimes';
-import { formatFileName } from '@/lib/uploader/formatFileName';
 import { parseHeaders, UploadHeaders } from '@/lib/uploader/parseHeaders';
 import { onUpload } from '@/lib/webhooks';
 import { Prisma } from '@/prisma/client';
 import { userMiddleware } from '@/server/middleware/user';
 import typedPlugin from '@/server/typedPlugin';
 import { stat } from 'fs/promises';
-import { extname } from 'path';
-
-const commonDoubleExts = [
-  '.tar.gz',
-  '.tar.xz',
-  '.tar.bz2',
-  '.tar.lz',
-  '.tar.lzma',
-  '.tar.Z',
-  '.tar.7z',
-  '.zip.gz',
-  '.zip.xz',
-  '.rar.gz',
-  '.log.gz',
-  '.csv.gz',
-  '.pdf.gz',
-  // feel free to PR more
-];
-
-export const getExtension = (filename: string, override?: string): string => {
-  return override ?? commonDoubleExts.find((ext) => filename.endsWith(ext)) ?? extname(filename);
-};
 
 export type ApiUploadResponse = {
   files: {
@@ -84,40 +61,9 @@ export default typedPlugin(
 
       const files = await req.saveRequestFiles({ tmpdir: config.core.tempDirectory });
 
-      if (req.user?.quota) {
-        const totalFileSize = files.reduce((acc, x) => acc + x.file.bytesRead, 0);
-
-        const userAggregateStats = await prisma.file.aggregate({
-          where: {
-            userId: req.user.id,
-          },
-          _sum: {
-            size: true,
-          },
-          _count: {
-            _all: true,
-          },
-        });
-        const aggSize: bigint =
-          userAggregateStats!._sum?.size === null
-            ? 0n
-            : (userAggregateStats!._sum?.size as unknown as bigint);
-        if (
-          req.user.quota.filesQuota === 'BY_BYTES' &&
-          Number(aggSize) + totalFileSize > bytes(req.user.quota.maxBytes!)
-        )
-          return res.payloadTooLarge(
-            `uploading will exceed your storage quota of ${bytes(req.user.quota.maxBytes!)} bytes`,
-          );
-
-        if (
-          req.user.quota.filesQuota === 'BY_FILES' &&
-          userAggregateStats!._count?._all + req.files.length > req.user.quota.maxFiles!
-        )
-          return res.payloadTooLarge(
-            `uploading will exceed your file count quota of ${req.user.quota.maxFiles} files`,
-          );
-      }
+      const totalFileSize = files.reduce((acc, x) => acc + x.file.bytesRead, 0);
+      const quotaCheck = await checkQuota(req.user, totalFileSize, files.length);
+      if (quotaCheck !== true) return res.payloadTooLarge(quotaCheck);
 
       const response: ApiUploadResponse = {
         files: [],
@@ -127,14 +73,7 @@ export default typedPlugin(
         ...(config.files.assumeMimetypes && { assumedMimetypes: Array(req.files.length) }),
       };
 
-      let domain;
-      if (options.overrides?.returnDomain) {
-        domain = `${config.core.returnHttpsUrls ? 'https' : 'http'}://${options.overrides.returnDomain}`;
-      } else if (config.core.defaultDomain) {
-        domain = `${config.core.returnHttpsUrls ? 'https' : 'http'}://${config.core.defaultDomain}`;
-      } else {
-        domain = `${config.core.returnHttpsUrls ? 'https' : 'http'}://${req.headers.host}`;
-      }
+      const domain = getDomain(options.overrides?.returnDomain, config.core.defaultDomain, req.headers.host);
 
       logger.debug('uploading files', { files: files.map((x) => x.filename) });
 
@@ -151,36 +90,21 @@ export default typedPlugin(
 
         // determine filename
         const format = options.format || config.files.defaultFormat;
-        let fileName = formatFileName(format, file.filename);
-        if (options.overrides?.filename || format === 'name') {
-          if (options.overrides?.filename) {
-            const sanitized = sanitizeFilename(options.overrides.filename!);
-            if (!sanitized) return res.badRequest(`file[${i}]: Invalid characters in filename override`);
+        const nameResult = await getFilename(format, file.filename, extension, options.overrides?.filename);
+        if ('error' in nameResult) return res.badRequest(`file[${i}]: ${nameResult.error}`);
 
-            fileName = sanitized;
-          }
-
-          const fullFileName = `${fileName}${extension}`;
-          const existing = await prisma.file.findFirst({ where: { name: fullFileName } });
-          if (existing)
-            return res.badRequest(`file[${i}]: A file with the name "${fullFileName}" already exists`);
-        } else if (format === 'random') {
-          let fullFileName = `${fileName}${extension}`;
-          let existing = await prisma.file.findFirst({ where: { name: fullFileName } });
-          while (existing) {
-            fileName = formatFileName(format, file.filename);
-            fullFileName = `${fileName}${extension}`;
-            existing = await prisma.file.findFirst({ where: { name: fullFileName } });
-          }
-        }
+        const { fileName } = nameResult;
 
         // determine mimetype
-        let mimetype = file.mimetype;
-        if (mimetype === 'application/octet-stream' && config.files.assumeMimetypes) {
-          const mime = await guess(extension.substring(1));
+        const { mimetype, assumed } = await getMimetype(file.mimetype, extension);
+        if (!assumed && config.files.assumeMimetypes) {
+          logger.warn(
+            `file[${i}]: mimetype ${file.mimetype} was not recognized, to ignore this warning, turn off assume mimetypes.`,
+          );
 
-          response.assumedMimetypes![i] = !!mime;
-          if (mime) mimetype = mime;
+          return res.badRequest(
+            `file[${i}]: mimetype ${file.mimetype} was not recognized, supply a valid mimetype`,
+          );
         }
 
         // compress the image if requested
@@ -220,7 +144,13 @@ export default typedPlugin(
         if (options.maxViews) data.maxViews = options.maxViews;
         if (options.password) data.password = await hashPassword(options.password);
         if (folder) data.Folder = { connect: { id: folder.id } };
-        if (options.addOriginalName) data.originalName = file.filename;
+        if (options.addOriginalName) {
+          const sanitizedOG = sanitizeFilename(file.filename);
+          if (!sanitizedOG) return res.badRequest(`file[${i}]: Invalid characters in original filename`);
+
+          data.originalName = sanitizedOG;
+        }
+
         data.deletesAt = options.deletesAt && options.deletesAt !== 'never' ? options.deletesAt : null;
 
         const fileUpload = await prisma.file.create({

src/server/routes/api/upload/partial.ts

@@ -1,3 +1,4 @@
+import { checkQuota, getDomain, getExtension, getFilename } from '@/lib/api/upload';
 import { bytes } from '@/lib/bytes';
 import { config } from '@/lib/config';
 import { hashPassword } from '@/lib/crypto';
@@ -6,7 +7,6 @@ import { sanitizeFilename } from '@/lib/fs';
 import { log } from '@/lib/logger';
 import { guess } from '@/lib/mimes';
 import { randomCharacters } from '@/lib/random';
-import { formatFileName } from '@/lib/uploader/formatFileName';
 import { UploadHeaders, UploadOptions, parseHeaders } from '@/lib/uploader/parseHeaders';
 import { Prisma } from '@/prisma/client';
 import { userMiddleware } from '@/server/middleware/user';
@@ -14,11 +14,34 @@ import typedPlugin from '@/server/typedPlugin';
 import { readdir, rename, rm } from 'fs/promises';
 import { join } from 'path';
 import { Worker } from 'worker_threads';
-import { ApiUploadResponse, getExtension } from '.';
+import { ApiUploadResponse } from '.';
 
 const logger = log('api').c('upload').c('partial');
 
-const partialsCache = new Map<string, { length: number; options: UploadOptions }>();
+const partialsCache = new Map<string, { length: number; options: UploadOptions; prefix: string }>();
+
+function createPartial(length: number, options: UploadOptions) {
+  const identifier = randomCharacters(8);
+
+  const prefix = `zipline_partial_${identifier}_`;
+
+  partialsCache.set(identifier, { length, options, prefix });
+  return identifier;
+}
+
+async function deletePartial(identifier: string, deleteFiles = true) {
+  const cache = partialsCache.get(identifier);
+  if (!cache) return;
+
+  partialsCache.delete(identifier);
+
+  if (deleteFiles) {
+    const tempFiles = await readdir(config.core.tempDirectory);
+    await Promise.all(
+      tempFiles.filter((f) => f.startsWith(cache.prefix)).map((f) => rm(join(config.core.tempDirectory, f))),
+    );
+  }
+}
 
 export type ApiUploadPartialResponse = ApiUploadResponse & {
   partialSuccess?: boolean;
@@ -54,41 +77,6 @@ export default typedPlugin(
 
       const files = await req.saveRequestFiles({ tmpdir: config.core.tempDirectory });
 
-      if (req.user?.quota) {
-        const totalFileSize = files.reduce((acc, x) => acc + x.file.bytesRead, 0);
-
-        const userAggregateStats = await prisma.file.aggregate({
-          where: {
-            userId: req.user.id,
-          },
-          _sum: {
-            size: true,
-          },
-          _count: {
-            _all: true,
-          },
-        });
-        const aggSize: bigint =
-          userAggregateStats!._sum?.size === null
-            ? 0n
-            : (userAggregateStats!._sum?.size as unknown as bigint);
-        if (
-          req.user.quota.filesQuota === 'BY_BYTES' &&
-          Number(aggSize) + totalFileSize > bytes(req.user.quota.maxBytes!)
-        )
-          return res.payloadTooLarge(
-            `uploading will exceed your storage quota of ${bytes(req.user.quota.maxBytes!)} bytes`,
-          );
-
-        if (
-          req.user.quota.filesQuota === 'BY_FILES' &&
-          userAggregateStats!._count?._all + req.files.length > req.user.quota.maxFiles!
-        )
-          return res.payloadTooLarge(
-            `uploading will exceed your file count quota of ${req.user.quota.maxFiles} files`,
-          );
-      }
-
       const response: ApiUploadPartialResponse = {
         files: [],
         ...(options.deletesAt && {
@@ -97,14 +85,7 @@ export default typedPlugin(
         ...(config.files.assumeMimetypes && { assumedMimetypes: Array(req.files.length) }),
       };
 
-      let domain;
-      if (options.overrides?.returnDomain) {
-        domain = `${config.core.returnHttpsUrls ? 'https' : 'http'}://${options.overrides.returnDomain}`;
-      } else if (config.core.defaultDomain) {
-        domain = `${config.core.returnHttpsUrls ? 'https' : 'http'}://${config.core.defaultDomain}`;
-      } else {
-        domain = `${config.core.returnHttpsUrls ? 'https' : 'http'}://${req.headers.host}`;
-      }
+      const domain = getDomain(options.overrides?.returnDomain, config.core.defaultDomain, req.headers.host);
 
       logger.debug('saving partial files', { partial: options.partial, files: files.map((x) => x.filename) });
 
@@ -114,27 +95,26 @@ export default typedPlugin(
 
       // caching for partial uploads server side checks and performance
       if (options.partial.range[0] === 0) {
-        const identifier = randomCharacters(8);
-        partialsCache.set(identifier, { length: fileSize, options });
-        options.partial.identifier = identifier;
+        options.partial.identifier = createPartial(fileSize, options);
       } else {
         if (!options.partial.identifier || !partialsCache.has(options.partial.identifier))
-          return res.badRequest('No partial upload identifier provided');
+          return res.badRequest('No/Invalid partial upload identifier provided');
       }
 
       const cache = partialsCache.get(options.partial.identifier);
-      if (!cache) throw 'No partial upload cache found';
+      if (!cache) return res.badRequest('No/Invalid partial upload identifier provided');
 
-      const prefix = `zipline_partial_${options.partial.identifier}_`;
+      // check quota, using the current added length, and only just adding one file
+      const quotaCheck = await checkQuota(req.user, cache.length + fileSize, 1);
+      if (quotaCheck !== true) {
+        await deletePartial(options.partial.identifier);
+
+        return res.payloadTooLarge(quotaCheck);
+      }
 
       // file is too large so we delete everything
       if (cache.length + fileSize > bytes(config.files.maxFileSize)) {
-        partialsCache.delete(options.partial.identifier);
-
-        const tempFiles = await readdir(config.core.tempDirectory);
-        await Promise.all(
-          tempFiles.filter((f) => f.startsWith(prefix)).map((f) => rm(join(config.core.tempDirectory, f))),
-        );
+        await deletePartial(options.partial.identifier!);
 
         return res.payloadTooLarge('File is too large');
       }
@@ -142,10 +122,12 @@ export default typedPlugin(
       cache.length += fileSize;
 
       // handle partial stuff
-      const tempFile = join(
-        config.core.tempDirectory,
-        `${prefix}${options.partial.range[0]}_${options.partial.range[1]}`,
+      const sanitized = sanitizeFilename(
+        `${cache.prefix}${options.partial.range[0]}_${options.partial.range[1]}`,
       );
+      if (!sanitized) return res.badRequest('Invalid characters in filename');
+
+      const tempFile = join(config.core.tempDirectory, sanitized);
       await rename(file.filepath, tempFile);
 
       if (options.partial.lastchunk) {
@@ -155,32 +137,15 @@ export default typedPlugin(
 
         // determine filename
         const format = options.format || config.files.defaultFormat;
-        let fileName = formatFileName(format, decodeURIComponent(options.partial.filename));
+        const nameResult = await getFilename(
+          format,
+          options.partial.filename,
+          extension,
+          options.overrides?.filename,
+        );
+        if ('error' in nameResult) return res.badRequest(nameResult.error);
 
-        if (options.overrides?.filename || format === 'name') {
-          if (options.overrides?.filename) {
-            const sanitized = sanitizeFilename(options.overrides!.filename!);
-            if (!sanitized) return res.badRequest('Invalid characters in filename override');
-
-            fileName = sanitized;
-          }
-          const fullFileName = `${fileName}${extension}`;
-
-          const existing = await prisma.file.findFirst({
-            where: {
-              name: fullFileName,
-            },
-          });
-          if (existing) return res.badRequest(`A file with the name "${fullFileName}" already exists`);
-        } else if (format === 'random') {
-          let fullFileName = `${fileName}${extension}`;
-          let existing = await prisma.file.findFirst({ where: { name: fullFileName } });
-          while (existing) {
-            fileName = formatFileName(format, decodeURIComponent(options.partial.filename));
-            fullFileName = `${fileName}${extension}`;
-            existing = await prisma.file.findFirst({ where: { name: fullFileName } });
-          }
-        }
+        const { fileName } = nameResult;
 
         // determine mimetype
         let mimetype = options.partial.contentType;
@@ -208,10 +173,12 @@ export default typedPlugin(
         if (options.password) data.password = await hashPassword(options.password);
         if (options.maxViews) data.maxViews = options.maxViews;
         if (folder) data.Folder = { connect: { id: folder.id } };
-        if (options.addOriginalName)
-          data.originalName = options.partial.filename
-            ? decodeURIComponent(options.partial.filename)
-            : file.filename; // this will prolly be "blob" but should hopefully never happen
+        if (options.addOriginalName) {
+          const sanitizedOG = sanitizeFilename(options.partial.filename);
+          if (!sanitizedOG) return res.badRequest('Invalid characters in original filename');
+
+          data.originalName = sanitizedOG || file.filename; // this will prolly be "blob" but should hopefully never happen
+        }
 
         const fileUpload = await prisma.file.create({
           data,
@@ -276,7 +243,7 @@ export default typedPlugin(
           pending: true,
         });
 
-        partialsCache.delete(options.partial.identifier);
+        await deletePartial(options.partial.identifier, false);
       }
 
       response.partialSuccess = true;