From 5cd10ce51d4d0c718ae7cdcc450a691d6f443f47 Mon Sep 17 00:00:00 2001
From: Alex Vandiver <alexmv@zulip.com>
Date: Tue, 18 Jun 2024 17:19:17 +0000
Subject: [PATCH] s3: Allow setting a CloudFront URL prefix for avatar and
 emoji images.

---
 zerver/lib/upload/s3.py            |  9 +++++++++
 zerver/tests/test_upload_s3.py     | 19 +++++++++++++++++++
 zproject/default_settings.py       |  1 +
 zproject/prod_settings_template.py |  1 +
 4 files changed, 30 insertions(+)

diff --git a/zerver/lib/upload/s3.py b/zerver/lib/upload/s3.py
index 8108c729cc..b065ae5f8b 100644
--- a/zerver/lib/upload/s3.py
+++ b/zerver/lib/upload/s3.py
@@ -154,6 +154,15 @@ class S3UploadBackend(ZulipUploadBackend):
         # hundreds of avatar URLs in single `GET /messages` request,
         # we instead back-compute the URL pattern here.
 
+        # The S3_AVATAR_PUBLIC_URL_PREFIX setting is used to override
+        # this prefix, for instance if a CloudFront distribution is
+        # used.
+        if settings.S3_AVATAR_PUBLIC_URL_PREFIX is not None:
+            prefix = settings.S3_AVATAR_PUBLIC_URL_PREFIX
+            if not prefix.endswith("/"):
+                prefix += "/"
+            return prefix
+
         DUMMY_KEY = "dummy_key_ignored"
 
         # We do not access self.avatar_bucket.meta.client directly,
diff --git a/zerver/tests/test_upload_s3.py b/zerver/tests/test_upload_s3.py
index 238f44439c..309958acea 100644
--- a/zerver/tests/test_upload_s3.py
+++ b/zerver/tests/test_upload_s3.py
@@ -247,6 +247,25 @@ class S3Test(ZulipTestCase):
         body = f"First message ...[zulip.txt](http://{hamlet.realm.host}" + url + ")"
         self.send_stream_message(hamlet, "Denmark", body, "test")
 
+    @use_s3_backend
+    def test_user_avatars_base(self) -> None:
+        backend = zerver.lib.upload.upload_backend
+        assert isinstance(backend, S3UploadBackend)
+        self.assertEqual(
+            backend.construct_public_upload_url_base(),
+            f"https://{settings.S3_AVATAR_BUCKET}.s3.amazonaws.com/",
+        )
+
+        with self.settings(S3_AVATAR_PUBLIC_URL_PREFIX="https://avatars.example.com"):
+            self.assertEqual(
+                backend.construct_public_upload_url_base(), "https://avatars.example.com/"
+            )
+
+        with self.settings(S3_AVATAR_PUBLIC_URL_PREFIX="https://avatars.example.com/"):
+            self.assertEqual(
+                backend.construct_public_upload_url_base(), "https://avatars.example.com/"
+            )
+
     @use_s3_backend
     def test_user_avatars_redirect(self) -> None:
         create_s3_buckets(settings.S3_AVATAR_BUCKET)[0]
diff --git a/zproject/default_settings.py b/zproject/default_settings.py
index 45f8195dc3..eb2d33de14 100644
--- a/zproject/default_settings.py
+++ b/zproject/default_settings.py
@@ -161,6 +161,7 @@ S3_UPLOADS_STORAGE_CLASS: Literal[
     "STANDARD",
     "STANDARD_IA",
 ] = "STANDARD"
+S3_AVATAR_PUBLIC_URL_PREFIX: Optional[str] = None
 LOCAL_UPLOADS_DIR: Optional[str] = None
 LOCAL_AVATARS_DIR: Optional[str] = None
 LOCAL_FILES_DIR: Optional[str] = None
diff --git a/zproject/prod_settings_template.py b/zproject/prod_settings_template.py
index 63d3c0408b..911364d44d 100644
--- a/zproject/prod_settings_template.py
+++ b/zproject/prod_settings_template.py
@@ -790,6 +790,7 @@ LOCAL_UPLOADS_DIR = "/home/zulip/uploads"
 # S3_AVATAR_BUCKET = ""
 # S3_REGION = None
 # S3_ENDPOINT_URL = None
+# S3_AVATAR_PUBLIC_URL_PREFIX = None
 # S3_ADDRESSING_STYLE = "auto"
 # S3_SKIP_PROXY = True
 # S3_UPLOADS_STORAGE_CLASS = "STANDARD"