s3fs: Adds support for SSE client keys

dvc/config_schema.py

@@ -152,6 +152,8 @@ class RelPath(str):
                     "ssl_verify": Any(Bool, str),
                     "sse": str,
                     "sse_kms_key_id": str,
+                    "sse_customer_algorithm": str,
+                    "sse_customer_key": str,
                     "acl": str,
                     "grant_read": str,
                     "grant_read_acp": str,

dvc/fs/s3.py

@@ -76,6 +76,8 @@ def _load_aws_config_file(self, profile):
         return self._split_s3_config(s3_config)
 
     def _prepare_credentials(self, **config):
+        from s3fs.utils import SSEParams
+
         from dvc.config import ConfigError
         from dvc.utils.flatten import flatten, unflatten
 
@@ -103,9 +105,28 @@ def _prepare_credentials(self, **config):
 
         # encryptions
         additional = login_info["s3_additional_kwargs"]
-        additional["ServerSideEncryption"] = config.get("sse")
-        additional["SSEKMSKeyId"] = config.get("sse_kms_key_id")
+        sse_customer_key = None
+        if config.get("sse_customer_key"):
+            if config.get("sse_kms_key_id"):
+                raise ConfigError(
+                    "`sse_kms_key_id` and `sse_customer_key` AWS S3 config "
+                    "options are mutually exclusive"
+                )
+            import base64
+
+            sse_customer_key = base64.b64decode(config.get("sse_customer_key"))
+        sse_customer_algorithm = config.get("sse_customer_algorithm")
+        if not sse_customer_algorithm:
+            sse_customer_algorithm = "AES256"
+        sse_params = SSEParams(
+            server_side_encryption=config.get("sse"),
+            sse_customer_algorithm=sse_customer_algorithm,
+            sse_customer_key=sse_customer_key,
+            sse_kms_key_id=config.get("sse_kms_key_id"),
+        )
+        additional.update(sse_params.to_kwargs())
         additional["ACL"] = config.get("acl")
+
         for grant_option, grant_key in self._GRANTS.items():
             if config.get(grant_option):
                 if additional["ACL"]:

tests/unit/fs/test_s3.py

@@ -11,6 +11,8 @@
 key_id = "key-id"
 key_secret = "key-secret"
 session_token = "session-token"
+sse_customer_key = "3v2Q08ia8HOMGi5N21vsXI1fyfGAM6aVkW5pMu0A5dE="
+sse_customer_algorithm = "AES256"
 
 
 @pytest.fixture(autouse=True)
@@ -136,3 +138,31 @@ def test_key_id_and_secret(dvc):
     assert fs.fs_args["key"] == key_id
     assert fs.fs_args["secret"] == key_secret
     assert fs.fs_args["token"] == session_token
+
+
+def test_sse_customer_key(dvc):
+    fs = S3FileSystem(
+        url=url,
+        sse_customer_key=sse_customer_key,
+        sse_customer_algorithm=sse_customer_algorithm,
+    )
+
+    import base64
+
+    sse_key = fs.fs_args["s3_additional_kwargs"]["SSECustomerKey"]
+    assert base64.b64encode(sse_key).decode() == sse_customer_key
+    assert (
+        fs.fs_args["s3_additional_kwargs"]["SSECustomerAlgorithm"]
+        == sse_customer_algorithm
+    )
+
+
+def test_sse_customer_key_and_sse_kms_key_id_mutually_exclusive(dvc):
+    config = {
+        "url": url,
+        "sse_customer_key": sse_customer_key,
+        "sse_kms_key_id": "key",
+    }
+
+    with pytest.raises(ConfigError):
+        S3FileSystem(**config)