fix presigned urls for custom endpoint

yunojuno · Mar 25, 2022 · 99773e8 · 99773e8
1 parent 97e5aea
commit 99773e8
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 8 deletions.
diff --git a/s3upload/utils.py b/s3upload/utils.py
@@ -142,10 +142,12 @@ def get_signed_download_url(
     ttl: int = 60,
 ) -> str:
     bucket_name = bucket_name or settings.AWS_STORAGE_BUCKET_NAME
+    bucket_endpoint = get_bucket_endpoint_url(bucket_name)
     s3 = boto3.client(
         "s3",
         aws_access_key_id=settings.AWS_ACCESS_KEY_ID,
         aws_secret_access_key=settings.AWS_SECRET_ACCESS_KEY,
+        endpoint_url=bucket_endpoint,
     )
     download_url = s3.generate_presigned_url(
         "get_object", Params={"Bucket": bucket_name, "Key": key}, ExpiresIn=ttl

diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -1,6 +1,9 @@
+from urllib.parse import urlparse
+
 import pytest
+from django.test import override_settings
 
-from s3upload.utils import get_s3_path_from_url, get_bucket_endpoint_url
+from s3upload.utils import get_bucket_endpoint_url, get_s3_path_from_url
 
 TEST_BUCKET = "test-bucket-name"
 TEST_KEY = "folder1/folder2/file1.json"
@@ -26,13 +29,35 @@ def test_get_s3_path_from_url(url: str) -> None:
     "url,expected",
     [
         (None, f"https://{TEST_BUCKET}.s3.{TEST_REGION}.amazonaws.com"),
-        ("https://{bucket}.s3.{region}.amazonaws.com", f"https://{TEST_BUCKET}.s3.{TEST_REGION}.amazonaws.com"),
-        ("https://s3.{region}.amazonaws.com", f"https://s3.{TEST_REGION}.amazonaws.com"),
-        ("https://{bucket}.s3.amazonaws.com", f"https://{TEST_BUCKET}.s3.amazonaws.com"),
-    ]
+        (
+            "https://{bucket}.s3.{region}.amazonaws.com",
+            f"https://{TEST_BUCKET}.s3.{TEST_REGION}.amazonaws.com",
+        ),
+        (
+            "https://s3.{region}.amazonaws.com",
+            f"https://s3.{TEST_REGION}.amazonaws.com",
+        ),
+        (
+            "https://{bucket}.s3.amazonaws.com",
+            f"https://{TEST_BUCKET}.s3.amazonaws.com",
+        ),
+    ],
 )
 def test_get_bucket_endpoint_url(url: str, expected: str) -> None:
     kwargs = {}
     if url:
         kwargs["default"] = url
-    assert get_bucket_endpoint_url(bucket_name=TEST_BUCKET, region=TEST_REGION, **kwargs) == expected
+    assert (
+        get_bucket_endpoint_url(bucket_name=TEST_BUCKET, region=TEST_REGION, **kwargs)
+        == expected
+    )
+
+
+@override_settings(
+    S3UPLOAD_REGION=TEST_REGION,
+)
+def test_get_signed_download_url() -> None:
+    presigned_url = get_signed_download_url(key=TEST_KEY, bucket_name=TEST_BUCKET)
+    parsed_url = urlparse(presigned_url)
+    assert parsed_url.scheme == "https"
+    assert parsed_url.netloc == f"{TEST_BUCKET}.s3.{TEST_REGION}.amazonaws.com"
diff --git a/tests/test_widgets.py b/tests/test_widgets.py
@@ -181,7 +181,8 @@ def test_check_policy_conditions(self):
                 "acl": "private",
                 "bucket": "test-bucket",
             }
-        }
+        },
+        S3UPLOAD_REGION="us",
     )
     def test_check_signed_url(self):
         data = {"dest": "misc", "name": "image.jpg", "type": "image/jpeg"}
@@ -190,7 +191,7 @@ def test_check_signed_url(self):
         parsed_url = urlparse(response_dict["private_access_url"])
         parsed_qs = parse_qs(parsed_url.query)
         self.assertEqual(parsed_url.scheme, "https")
-        self.assertEqual(parsed_url.netloc, "test-bucket.s3.amazonaws.com")
+        self.assertEqual(parsed_url.netloc, "test-bucket.s3.us.amazonaws.com")
         self.assertTrue("Signature" in parsed_qs)
         self.assertTrue("Expires" in parsed_qs)