From 093ee74a9033bee0c164e31c3281de518417e663 Mon Sep 17 00:00:00 2001
From: Tom Moor <tom.moor@gmail.com>
Date: Sat, 2 Sep 2023 22:11:53 -0400
Subject: [PATCH] fix: Protect against exports larger than memory/max

---
 server/env.ts                     |  9 +++++++++
 server/models/Attachment.ts       | 25 ++++++++++++++++++++++++
 server/queues/tasks/ExportTask.ts | 32 ++++++++++++++++++++++++++++++-
 3 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/server/env.ts b/server/env.ts
index fa228a4ae..35459e035 100644
--- a/server/env.ts
+++ b/server/env.ts
@@ -5,6 +5,7 @@ require("dotenv").config({
   silent: true,
 });
 
+import os from "os";
 import {
   validate,
   IsNotEmpty,
@@ -622,6 +623,14 @@ export class Environment {
     this.AWS_S3_UPLOAD_MAX_SIZE
   );
 
+  /**
+   * Limit on export size in bytes. Defaults to the total memory available to
+   * the container.
+   */
+  @IsNumber()
+  public MAXIMUM_EXPORT_SIZE =
+    this.toOptionalNumber(process.env.MAXIMUM_EXPORT_SIZE) ?? os.totalmem();
+
   /**
    * Iframely url
    */
diff --git a/server/models/Attachment.ts b/server/models/Attachment.ts
index fe1518a95..0fe297e13 100644
--- a/server/models/Attachment.ts
+++ b/server/models/Attachment.ts
@@ -1,4 +1,5 @@
 import path from "path";
+import { QueryTypes } from "sequelize";
 import {
   BeforeDestroy,
   BelongsTo,
@@ -116,6 +117,30 @@ class Attachment extends IdModel {
     await FileStorage.deleteFile(model.key);
   }
 
+  // static methods
+
+  /**
+   * Get the total size of all attachments for a given team.
+   *
+   * @param teamId - The ID of the team to get the total size for.
+   * @returns A promise resolving to the total size of all attachments for the given team in bytes.
+   */
+  static async getTotalSizeForTeam(teamId: string): Promise<number> {
+    const result = await this.sequelize!.query<{ total: string }>(
+      `
+      SELECT SUM(size) as total
+      FROM attachments
+      WHERE "teamId" = :teamId
+    `,
+      {
+        replacements: { teamId },
+        type: QueryTypes.SELECT,
+      }
+    );
+
+    return parseInt(result?.[0]?.total ?? "0", 10);
+  }
+
   // associations
 
   @BelongsTo(() => Team, "teamId")
diff --git a/server/queues/tasks/ExportTask.ts b/server/queues/tasks/ExportTask.ts
index 8aca07483..65ee136ca 100644
--- a/server/queues/tasks/ExportTask.ts
+++ b/server/queues/tasks/ExportTask.ts
@@ -1,10 +1,20 @@
 import fs from "fs";
 import truncate from "lodash/truncate";
 import { FileOperationState, NotificationEventType } from "@shared/types";
+import { bytesToHumanReadable } from "@shared/utils/files";
 import ExportFailureEmail from "@server/emails/templates/ExportFailureEmail";
 import ExportSuccessEmail from "@server/emails/templates/ExportSuccessEmail";
+import env from "@server/env";
+import { ValidationError } from "@server/errors";
 import Logger from "@server/logging/Logger";
-import { Collection, Event, FileOperation, Team, User } from "@server/models";
+import {
+  Attachment,
+  Collection,
+  Event,
+  FileOperation,
+  Team,
+  User,
+} from "@server/models";
 import fileOperationPresenter from "@server/presenters/fileOperation";
 import FileStorage from "@server/storage/files";
 import BaseTask, { TaskPriority } from "./BaseTask";
@@ -43,6 +53,26 @@ export default abstract class ExportTask extends BaseTask<Props> {
     let filePath: string | undefined;
 
     try {
+      if (!fileOperation.collectionId) {
+        const totalAttachmentsSize = await Attachment.getTotalSizeForTeam(
+          user.teamId
+        );
+
+        if (
+          fileOperation.includeAttachments &&
+          env.MAXIMUM_EXPORT_SIZE &&
+          totalAttachmentsSize > env.MAXIMUM_EXPORT_SIZE
+        ) {
+          throw ValidationError(
+            `${bytesToHumanReadable(
+              totalAttachmentsSize
+            )} of attachments in workspace is larger than maximum export size of ${bytesToHumanReadable(
+              env.MAXIMUM_EXPORT_SIZE
+            )}.`
+          );
+        }
+      }
+
       Logger.info("task", `ExportTask processing data for ${fileOperationId}`, {
         includeAttachments: fileOperation.includeAttachments,
       });