[s3] fix support for retrieving SSE-C files

jschneier · Sep 4, 2023 · 49b5207 · 49b5207
1 parent ad36a9f
commit 49b5207
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 2 deletions.
diff --git a/storages/backends/s3boto3.py b/storages/backends/s3boto3.py
@@ -34,6 +34,7 @@
 
 try:
     import boto3.session
+    import s3transfer.constants
     from boto3.s3.transfer import TransferConfig
     from botocore.client import Config
     from botocore.exceptions import ClientError
@@ -88,6 +89,13 @@ def _cloud_front_signer_from_pem(key_id, pem):
             'Supported backends are packages: cryptography and rsa.')
 
 
+def _filter_download_params(params):
+    return {
+        key: value for (key, value) in params.items()
+        if key in s3transfer.constants.ALLOWED_DOWNLOAD_ARGS
+    }
+
+
 @deconstructible
 class S3Boto3StorageFile(CompressedFileMixin, File):
     """
@@ -117,7 +125,8 @@ def __init__(self, name, mode, storage, buffer_size=None):
         self.obj = storage.bucket.Object(name)
         if 'w' not in mode:
             # Force early RAII-style exception if object does not exist
-            self.obj.load()
+            params = _filter_download_params(self._storage.get_object_parameters(self.name))
+            self.obj.load(**params)
         self._is_dirty = False
         self._raw_bytes_written = 0
         self._file = None
@@ -147,7 +156,12 @@ def _get_file(self):
             )
             if 'r' in self._mode:
                 self._is_dirty = False
-                self.obj.download_fileobj(self._file, Config=self._storage.transfer_config)
+                params = _filter_download_params(self._storage.get_object_parameters(self.name))
+                self.obj.download_fileobj(
+                    self._file,
+                    ExtraArgs=params,
+                    Config=self._storage.transfer_config
+                )
                 self._file.seek(0)
                 if self._storage.gzip and self.obj.content_encoding == 'gzip':
                     self._file = self._decompress_file(mode=self._mode, file=self._file)

diff --git a/tests/test_s3boto3.py b/tests/test_s3boto3.py
@@ -854,6 +854,19 @@ def setUp(self) -> None:
         self.storage = s3boto3.S3Boto3Storage()
         self.storage._connections.connection = mock.MagicMock()
 
+    def test_loading_ssec(self):
+        params = {"SSECustomerKey": "xyz", "CacheControl": "never"}
+        self.storage.get_object_parameters = lambda name: params
+
+        filtered = {"SSECustomerKey": "xyz"}
+        f = s3boto3.S3Boto3StorageFile('test', 'r', self.storage)
+        f.obj.load.assert_called_once_with(**filtered)
+
+        f.file
+        f.obj.download_fileobj.assert_called_once_with(
+            mock.ANY, ExtraArgs=filtered, Config=self.storage.transfer_config
+        )
+
     def test_closed(self):
         f = s3boto3.S3Boto3StorageFile('test', 'wb', self.storage)