chore: rework handling of submission attachment under 100 bytes (#1171)

Author: craigzour

lambda-code/reliability/src/lib/s3FileInput.ts

@@ -8,6 +8,13 @@ import {
 } from "@aws-sdk/client-s3";
 import { NodeJsClient } from "@smithy/types";
 
+export class FileSizeUnder100BytesException extends Error {
+  constructor() {
+    super("FileSizeUnder100BytesException");
+    Object.setPrototypeOf(this, FileSizeUnder100BytesException.prototype);
+  }
+}
+
 const s3Client = new S3Client({
   region: process.env.REGION ?? "ca-central-1",
   forcePathStyle: true,
@@ -124,22 +131,6 @@ export const getFileTags = async (filePath: string) => {
   }
 };
 
-export const getFileSize = async (filePath: string) => {
-  const response = await s3Client
-    .send(
-      new HeadObjectCommand({
-        Bucket: reliabilityBucketName,
-        Key: filePath,
-      })
-    )
-    .catch((error) => {
-      console.error(error);
-      throw new Error(`Failed to retrieve size for file: ${filePath}`);
-    });
-
-  return Number(response.ContentLength ?? "0");
-};
-
 export const getFileMetaData = async (filePath: string) => {
   const response = await s3Client
     .send(
@@ -165,19 +156,28 @@ export const getFileMetaData = async (filePath: string) => {
 export async function getObjectFirst100BytesInReliabilityBucket(
   objectKey: string
 ): Promise<Uint8Array<ArrayBufferLike>> {
-  const response = await s3Client.send(
-    new GetObjectCommand({
-      Bucket: reliabilityBucketName,
-      Key: objectKey,
-      Range: "bytes=0-99",
-    })
-  );
-
-  const bytes = await response.Body?.transformToByteArray();
-
-  if (bytes === undefined) {
-    throw new Error("Object first 100 bytes failed to be retrieved");
-  }
+  try {
+    const response = await s3Client.send(
+      new GetObjectCommand({
+        Bucket: reliabilityBucketName,
+        Key: objectKey,
+        Range: "bytes=0-99",
+      })
+    );
 
-  return bytes;
+    const bytes = await response.Body?.transformToByteArray();
+
+    if (bytes === undefined || bytes.length < 100) {
+      throw new FileSizeUnder100BytesException();
+    }
+
+    return bytes;
+  } catch (error) {
+    // AWS S3 will throw the exception only if the file size is 0 bytes.
+    if ((error as Error).name === "InvalidRange") {
+      throw new FileSizeUnder100BytesException();
+    }
+
+    throw error;
+  }
 }

lambda-code/reliability/src/lib/vaultProcessing.ts

@@ -4,7 +4,7 @@ import {} from "./file_scanning.js";
 import { isFileValid } from "./fileValidation.js";
 import {
   copyFilesFromReliabilityToVaultStorage,
-  getFileSize,
+  FileSizeUnder100BytesException,
   getObjectFirst100BytesInReliabilityBucket,
   removeFilesFromReliabilityStorage,
 } from "./s3FileInput.js";
@@ -93,29 +93,39 @@ async function verifyAndFlagMaliciousSubmissionAttachments(
 ): Promise<SubmissionAttachmentInformation[]> {
   return Promise.all(
     submissionAttachmentsWithScanStatuses.map(async (item) => {
-      const fileSize = await getFileSize(item.attachmentPath);
+      try {
+        const attachmentFirst100Bytes = await getObjectFirst100BytesInReliabilityBucket(
+          item.attachmentPath
+        );
+
+        const isFileValidResult = isFileValid(item.attachmentPath, attachmentFirst100Bytes);
+
+        // If we flagged the file as invalid we return the same scan status value AWS Guard Duty offers to avoid breaking Data Retrieval API service integration
+        const scanStatus = isFileValidResult ? item.scanStatus : "THREATS_FOUND";
 
-      if (fileSize < 100) {
-        // File size is too small to test and contains no content that can be of value to the end user
         return {
           ...item,
-          scanStatus: "THREATS_FOUND",
+          scanStatus: scanStatus,
         };
+      } catch (error) {
+        switch ((error as Error).constructor) {
+          case FileSizeUnder100BytesException:
+            console.warn(
+              JSON.stringify({
+                level: "warn",
+                msg: `Will not try to validate file as its size is under 100 bytes and probably contains no content that can be of value to the end user. Flagging ${item.attachmentPath} as potentially malicious.`,
+              })
+            );
+
+            return {
+              ...item,
+              scanStatus: "THREATS_FOUND",
+            };
+          default: {
+            throw new Error(`Failed to verify submission attachment ${item.attachmentPath}`);
+          }
+        }
       }
-
-      const attachmentFirst100Bytes = await getObjectFirst100BytesInReliabilityBucket(
-        item.attachmentPath
-      );
-
-      const isFileValidResult = isFileValid(item.attachmentPath, attachmentFirst100Bytes);
-
-      // If we flagged the file as invalid we return the same scan status value AWS Guard Duty offers to avoid breaking Data Retrieval API service integration
-      const scanStatus = isFileValidResult ? item.scanStatus : "THREATS_FOUND";
-
-      return {
-        ...item,
-        scanStatus: scanStatus,
-      };
     })
   );
 }