From 50bbe05334fc6a375e4091b338acdb46cc326f24 Mon Sep 17 00:00:00 2001
From: Tom Moor <tom@getoutline.com>
Date: Tue, 28 May 2024 08:12:25 -0400
Subject: [PATCH] Add range header support to `files.get` (#6950)

---
 plugins/storage/server/api/files.ts  | 41 ++++++++++++++++++++++++++--
 server/storage/files/BaseStorage.ts  |  3 +-
 server/storage/files/LocalStorage.ts |  8 ++++--
 server/storage/files/S3Storage.ts    |  6 +++-
 4 files changed, 52 insertions(+), 6 deletions(-)

diff --git a/plugins/storage/server/api/files.ts b/plugins/storage/server/api/files.ts
index 41f66de9e..6cb0b450f 100644
--- a/plugins/storage/server/api/files.ts
+++ b/plugins/storage/server/api/files.ts
@@ -15,6 +15,7 @@ import { Attachment } from "@server/models";
 import AttachmentHelper from "@server/models/helpers/AttachmentHelper";
 import { authorize } from "@server/policies";
 import FileStorage from "@server/storage/files";
+import LocalStorage from "@server/storage/files/LocalStorage";
 import { APIContext } from "@server/types";
 import { RateLimiterStrategy } from "@server/utils/RateLimiter";
 import { getJWTPayload } from "@server/utils/jwt";
@@ -84,12 +85,12 @@ router.get(
           "application/octet-stream"
       );
       ctx.attachment(fileName);
-      ctx.body = await FileStorage.getFileStream(key);
     } else {
       const attachment = await Attachment.findOne({
         where: { key },
         rejectOnEmpty: true,
       });
+
       authorize(actor, "read", attachment);
 
       ctx.set("Cache-Control", cacheHeader);
@@ -97,11 +98,47 @@ router.get(
       ctx.attachment(attachment.name, {
         type: FileStorage.getContentDisposition(attachment.contentType),
       });
-      ctx.body = attachment.stream;
     }
+
+    // Handle byte range requests
+    // https://developer.mozilla.org/en-US/docs/Web/HTTP/Range_requests
+    const stats = await (FileStorage as LocalStorage).stat(key);
+    const range = getByteRange(ctx, stats.size);
+
+    if (range) {
+      ctx.set("Content-Length", String(range.end - range.start + 1));
+      ctx.set(
+        "Content-Range",
+        `bytes ${range.start}-${range.end}/${stats.size}`
+      );
+    } else {
+      ctx.set("Content-Length", String(stats.size));
+    }
+
+    ctx.body = await FileStorage.getFileStream(key, range);
   }
 );
 
+function getByteRange(
+  ctx: APIContext<T.FilesGetReq>,
+  size: number
+): { start: number; end: number } | undefined {
+  const { range } = ctx.headers;
+  if (!range) {
+    return;
+  }
+
+  const match = range.match(/bytes=(\d+)-(\d+)?/);
+  if (!match) {
+    return;
+  }
+
+  const start = parseInt(match[1], 10);
+  const end = parseInt(match[2], 10) || size - 1;
+
+  return { start, end };
+}
+
 function getKeyFromContext(ctx: APIContext<T.FilesGetReq>): string {
   const { key, sig } = ctx.input.query;
   if (sig) {
diff --git a/server/storage/files/BaseStorage.ts b/server/storage/files/BaseStorage.ts
index 11afa44c3..304d1f199 100644
--- a/server/storage/files/BaseStorage.ts
+++ b/server/storage/files/BaseStorage.ts
@@ -32,7 +32,8 @@ export default abstract class BaseStorage {
    * @param key The path to the file
    */
   public abstract getFileStream(
-    key: string
+    key: string,
+    range?: { start?: number; end?: number }
   ): Promise<NodeJS.ReadableStream | null>;
 
   /**
diff --git a/server/storage/files/LocalStorage.ts b/server/storage/files/LocalStorage.ts
index 65d0313be..a2b63c853 100644
--- a/server/storage/files/LocalStorage.ts
+++ b/server/storage/files/LocalStorage.ts
@@ -131,8 +131,12 @@ export default class LocalStorage extends BaseStorage {
     };
   }
 
-  public getFileStream(key: string) {
-    return Promise.resolve(fs.createReadStream(this.getFilePath(key)));
+  public getFileStream(key: string, range?: { start: number; end: number }) {
+    return Promise.resolve(fs.createReadStream(this.getFilePath(key), range));
+  }
+
+  public stat(key: string) {
+    return fs.stat(this.getFilePath(key));
   }
 
   private getFilePath(key: string) {
diff --git a/server/storage/files/S3Storage.ts b/server/storage/files/S3Storage.ts
index 47a9ca7ff..0d14b397b 100644
--- a/server/storage/files/S3Storage.ts
+++ b/server/storage/files/S3Storage.ts
@@ -201,12 +201,16 @@ export default class S3Storage extends BaseStorage {
     });
   }
 
-  public getFileStream(key: string): Promise<NodeJS.ReadableStream | null> {
+  public getFileStream(
+    key: string,
+    range?: { start: number; end: number }
+  ): Promise<NodeJS.ReadableStream | null> {
     return this.client
       .send(
         new GetObjectCommand({
           Bucket: this.getBucket(),
           Key: key,
+          Range: range ? `bytes=${range.start}-${range.end}` : undefined,
         })
       )
       .then((item) => item.Body as NodeJS.ReadableStream)