Merge commit from fork

Author: diced

src/lib/api/upload.ts

@@ -36,11 +36,12 @@ export async function checkQuota(
   user: Pick<User, 'id' | 'quota'> | null,
   newSize: number,
   fileCount: number,
+  db: Pick<typeof prisma, 'file'> = prisma,
 ): Promise<true | string> {
   if (!user?.quota) return true;
 
   if (user.quota.filesQuota === 'BY_BYTES') {
-    const stats = await prisma.file.aggregate({
+    const stats = await db.file.aggregate({
       where: { userId: user.id },
       _sum: { size: true },
     });
@@ -51,7 +52,7 @@ export async function checkQuota(
     return true;
   }
 
-  const count = await prisma.file.count({ where: { userId: user.id } });
+  const count = await db.file.count({ where: { userId: user.id } });
   if (count + fileCount > user.quota.maxFiles!)
     return `uploading will exceed your file count quota of ${user.quota.maxFiles} files`;
 

src/offload/partial.ts

@@ -228,6 +228,8 @@ async function runComplete(id: string, size: number) {
 async function failPartial(config: Config, incompleteFile: IncompleteFile) {
   logger.error('failing incomplete file', { id: incompleteFile.id });
 
+  await dbProxy('file.delete', { where: { id: file.id } });
+
   const partials = await readdir(config.core.tempDirectory).then((files) =>
     files.filter((file) => file.startsWith(`zipline_partial_${options.partial!.identifier}`)),
   );

src/server/routes/api/auth/register.ts

@@ -56,46 +56,56 @@ export default typedPlugin(
         });
         if (oUser) throw new ApiError(1039);
 
-        if (code) {
-          const invite = await prisma.invite.findFirst({
-            where: {
-              OR: [{ id: code }, { code }],
+        const hashedPassword = await hashPassword(password);
+        const token = createToken();
+
+        const createUser = (db: Pick<typeof prisma, 'user'>) =>
+          db.user.create({
+            data: {
+              username,
+              password: hashedPassword,
+              role: 'USER',
+              token,
+            },
+            select: {
+              ...userSelect,
+              password: true,
+              token: true,
             },
           });
 
-          if (!invite) throw new ApiError(1035);
-          if (invite.expiresAt && new Date(invite.expiresAt) < new Date()) throw new ApiError(1035);
-          if (invite.maxUses && invite.uses >= invite.maxUses) throw new ApiError(1035);
+        let user;
+        if (code) {
+          const result = await prisma.$transaction(async (tx) => {
+            const [invite] = await tx.invite.updateManyAndReturn({
+              where: {
+                AND: [
+                  { OR: [{ id: code }, { code }] },
+                  { OR: [{ expiresAt: null }, { expiresAt: { gt: new Date() } }] },
+                  {
+                    OR: [{ maxUses: null }, { uses: { lt: tx.invite.fields.maxUses } }],
+                  },
+                ],
+              },
+              data: { uses: { increment: 1 } },
+              select: { id: true },
+            });
 
-          await prisma.invite.update({
-            where: {
-              id: invite.id,
-            },
-            data: {
-              uses: invite.uses + 1,
-            },
+            if (!invite) throw new ApiError(1035);
+
+            return { inviteId: invite.id, user: await createUser(tx) };
           });
 
+          user = result.user;
+
           logger.info('invite used', {
             user: username,
-            invite: invite.id,
+            invite: result.inviteId,
           });
+        } else {
+          user = await createUser(prisma);
         }
 
-        const user = await prisma.user.create({
-          data: {
-            username,
-            password: await hashPassword(password),
-            role: 'USER',
-            token: createToken(),
-          },
-          select: {
-            ...userSelect,
-            password: true,
-            token: true,
-          },
-        });
-
         await saveSession(session, <User>user);
 
         delete (user as any).password;

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

@@ -232,7 +232,8 @@ export default typedPlugin(
           named.push({ ...item, fileName: nameResult.fileName });
         }
 
-        response.files = await mapConcurrent(named, 4, async (item, i) => {
+        const password = options.password ? await hashPassword(options.password) : undefined;
+        const uploads = named.map((item, i) => {
           const { file, fileName, extension, mimetype, size, compressed, removedGps } = item;
 
           const data: Prisma.FileCreateInput = {
@@ -245,7 +246,7 @@ export default typedPlugin(
           if (!req.user && folder) data.anonymous = true;
 
           if (options.maxViews) data.maxViews = options.maxViews;
-          if (options.password) data.password = await hashPassword(options.password);
+          if (password) data.password = password;
           if (folder) data.Folder = { connect: { id: folder.id } };
           if (options.addOriginalName) {
             const sanitizedOG = sanitizeFilename(file.filename);
@@ -256,10 +257,39 @@ export default typedPlugin(
 
           data.deletesAt = options.deletesAt && options.deletesAt !== 'never' ? options.deletesAt : null;
 
-          const fileUpload = await prisma.file.create({
-            data,
-            select: fileSelect,
-          });
+          return { compressed, data, extension, file, removedGps, size };
+        });
+
+        const fileUploads = await prisma.$transaction(async (tx) => {
+          if (quotaUser?.quota) {
+            await tx.$queryRaw`SELECT "id" FROM "User" WHERE "id" = ${quotaUser.id} FOR UPDATE`;
+
+            const quotaCheck = await checkQuota(
+              quotaUser,
+              uploads.reduce((total, upload) => total + upload.size, 0),
+              uploads.length,
+              tx,
+            );
+            if (quotaCheck !== true)
+              throw new ApiError(5002, typeof quotaCheck === 'string' ? quotaCheck : undefined);
+          }
+
+          const created = [];
+          for (const upload of uploads) {
+            created.push(
+              await tx.file.create({
+                data: upload.data,
+                select: fileSelect,
+              }),
+            );
+          }
+
+          return created;
+        });
+
+        response.files = await mapConcurrent(uploads, 4, async (upload, uploadIndex) => {
+          const { compressed, extension, file, removedGps } = upload;
+          const fileUpload = fileUploads[uploadIndex];
 
           const storageData = compressed?.buffer ?? file.filepath;
           await datasource.put(fileUpload.name, storageData, {

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

@@ -70,16 +70,25 @@ function createPartial(options: UploadOptions, actorKey: string, quotaUserId: st
 }
 
 function activePartials(actorKey: string) {
-  return [...partialsCache.values()].filter((partial) => partial.actorKey === actorKey).length;
+  let count = 0;
+  for (const partial of partialsCache.values()) {
+    if (partial.actorKey === actorKey && ++count >= MAX_PARTIALS) return count;
+  }
+
+  return count;
 }
 
 function quotaReservations(quotaUserId: string) {
-  const reservations = [...partialsCache.values()].filter((partial) => partial.quotaUserId === quotaUserId);
+  let size = 0;
+  let files = 0;
+  for (const partial of partialsCache.values()) {
+    if (partial.quotaUserId !== quotaUserId || partial.finalized) continue;
+
+    size += partial.total;
+    files++;
+  }
 
-  return {
-    size: reservations.reduce((total, partial) => total + partial.total, 0),
-    files: reservations.filter((partial) => !partial.finalized).length,
-  };
+  return { size, files };
 }
 
 async function deletePartial(identifier: string, deleteFiles = true) {
@@ -98,12 +107,16 @@ async function deletePartial(identifier: string, deleteFiles = true) {
 }
 
 async function deleteOrphanedPartialFiles() {
-  const activePrefixes = [...partialsCache.values()].map((partial) => partial.prefix);
   const tempFiles = await readdir(config.core.tempDirectory);
-  const orphaned = tempFiles.filter(
-    (file) =>
-      file.startsWith('zipline_partial_') && !activePrefixes.some((prefix) => file.startsWith(prefix)),
-  );
+  const orphaned = tempFiles.filter((file) => {
+    if (!file.startsWith('zipline_partial_')) return false;
+
+    for (const partial of partialsCache.values()) {
+      if (file.startsWith(partial.prefix)) return false;
+    }
+
+    return true;
+  });
 
   await Promise.all(orphaned.map((file) => rm(join(config.core.tempDirectory, file), { force: true })));
 
@@ -307,7 +320,7 @@ export default typedPlugin(
 
           const data: Prisma.FileCreateInput = {
             name: `${fileName}${extension}`,
-            size: 0,
+            size: total,
             type: mimetype,
             User: {
               connect: {
@@ -327,9 +340,23 @@ export default typedPlugin(
           }
           if (!req.user && folder) data.anonymous = true;
 
-          const fileUpload = await prisma.file.create({
-            data,
-          });
+          let fileUpload;
+          try {
+            fileUpload = await prisma.$transaction(async (tx) => {
+              if (quotaUser?.quota) {
+                await tx.$queryRaw`SELECT "id" FROM "User" WHERE "id" = ${quotaUser.id} FOR UPDATE`;
+
+                const quotaCheck = await checkQuota(quotaUser, total, 1, tx);
+                if (quotaCheck !== true)
+                  throw new ApiError(5002, typeof quotaCheck === 'string' ? quotaCheck : undefined);
+              }
+
+              return tx.file.create({ data });
+            });
+          } catch (error) {
+            await deletePartial(options.partial.identifier);
+            throw error;
+          }
 
           const urlPath =
             options.extensionless && config.files.extensionlessUrls
@@ -378,6 +405,9 @@ export default typedPlugin(
                   result = await prisma.file.update(msg.data);
                   await deletePartial(partialIdentifier, false);
                   break;
+                case 'file.delete':
+                  result = await prisma.file.delete(msg.data);
+                  break;
                 case 'user.findUnique':
                   result = await prisma.user.findUnique(msg.data);
                   break;

src/server/routes/api/user/urls/index.ts

@@ -76,17 +76,6 @@ export default typedPlugin(
         const { vanity, destination, enabled } = req.body;
         const noJson = req.headers['x-zipline-no-json'];
 
-        const countUrls = await prisma.url.count({
-          where: {
-            userId: req.user.id,
-          },
-        });
-        if (req.user.quota && req.user.quota.maxUrls && countUrls + 1 > req.user.quota.maxUrls)
-          throw new ApiError(
-            3012,
-            `Shortening this URL would exceed your quota of ${req.user.quota.maxUrls} URLs.`,
-          );
-
         let returnDomain;
         const headerDomain = req.headers['x-zipline-domain'];
         if (headerDomain) {
@@ -116,19 +105,32 @@ export default typedPlugin(
           existingCode = await prisma.url.findFirst({ where: { code } });
         } while (existingCode);
 
-        const url = await prisma.url.create({
-          data: {
-            userId: req.user.id,
-            destination: destination,
-            code,
-            ...(vanity && { vanity: vanity }),
-            ...(maxViews && { maxViews: maxViews }),
-            ...(password && { password: password }),
-            ...(enabled !== undefined && { enabled: enabled }),
-          },
-          omit: {
-            password: true,
-          },
+        const url = await prisma.$transaction(async (tx) => {
+          await tx.$queryRaw`SELECT "id" FROM "User" WHERE "id" = ${req.user.id} FOR UPDATE`;
+
+          const countUrls = await tx.url.count({
+            where: { userId: req.user.id },
+          });
+          if (req.user.quota?.maxUrls && countUrls + 1 > req.user.quota.maxUrls)
+            throw new ApiError(
+              3012,
+              `Shortening this URL would exceed your quota of ${req.user.quota.maxUrls} URLs.`,
+            );
+
+          return tx.url.create({
+            data: {
+              userId: req.user.id,
+              destination: destination,
+              code,
+              ...(vanity && { vanity: vanity }),
+              ...(maxViews && { maxViews: maxViews }),
+              ...(password && { password: password }),
+              ...(enabled !== undefined && { enabled: enabled }),
+            },
+            omit: {
+              password: true,
+            },
+          });
         });
 
         let domain;