From 3f834150176bc2d65a16d0c97794ddfcd57a1362 Mon Sep 17 00:00:00 2001 From: Amey Kulkarni Date: Sun, 1 May 2022 22:16:56 +0530 Subject: [PATCH 1/4] s3fs: Adds support for SSE client keys Fixes #6141 --- dvc/config_schema.py | 2 ++ dvc/fs/s3.py | 25 +++++++++++++++++++++++-- tests/unit/fs/test_s3.py | 30 ++++++++++++++++++++++++++++++ 3 files changed, 55 insertions(+), 2 deletions(-) diff --git a/dvc/config_schema.py b/dvc/config_schema.py index aa495a890a..f86c3fdc25 100644 --- a/dvc/config_schema.py +++ b/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, diff --git a/dvc/fs/s3.py b/dvc/fs/s3.py index 7eb9632dec..bb6bd407a6 100644 --- a/dvc/fs/s3.py +++ b/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"]: diff --git a/tests/unit/fs/test_s3.py b/tests/unit/fs/test_s3.py index d2a8bda782..730d7e3d99 100644 --- a/tests/unit/fs/test_s3.py +++ b/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) From 96963dbb1636e6d01d1c85dc90674bc61a662e28 Mon Sep 17 00:00:00 2001 From: Amey Kulkarni Date: Mon, 2 May 2022 12:23:13 +0530 Subject: [PATCH 2/4] s3fs: modified logic so that sse_customer_algorithm will default to AES256 if sse_customer_key is specified and no algorithm is specified --- dvc/fs/s3.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dvc/fs/s3.py b/dvc/fs/s3.py index bb6bd407a6..8454ef417f 100644 --- a/dvc/fs/s3.py +++ b/dvc/fs/s3.py @@ -116,7 +116,7 @@ def _prepare_credentials(self, **config): sse_customer_key = base64.b64decode(config.get("sse_customer_key")) sse_customer_algorithm = config.get("sse_customer_algorithm") - if not sse_customer_algorithm: + if not sse_customer_algorithm and sse_customer_key: sse_customer_algorithm = "AES256" sse_params = SSEParams( server_side_encryption=config.get("sse"), From cce9eff090cafedd83104b65de71a4438b19e23d Mon Sep 17 00:00:00 2001 From: Amey Kulkarni Date: Tue, 21 Jun 2022 19:14:33 +0530 Subject: [PATCH 3/4] Removed tests for SSE-C in test_s3.py --- tests/unit/fs/test_s3.py | 28 ---------------------------- 1 file changed, 28 deletions(-) diff --git a/tests/unit/fs/test_s3.py b/tests/unit/fs/test_s3.py index 22871e2255..95a442e0ef 100644 --- a/tests/unit/fs/test_s3.py +++ b/tests/unit/fs/test_s3.py @@ -137,31 +137,3 @@ 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) From 6bcebc197eb7db968b77cac13e599e5f0138951c Mon Sep 17 00:00:00 2001 From: Amey Kulkarni Date: Tue, 21 Jun 2022 19:16:33 +0530 Subject: [PATCH 4/4] Removed unused variables in test_s3.py --- tests/unit/fs/test_s3.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/unit/fs/test_s3.py b/tests/unit/fs/test_s3.py index 95a442e0ef..cc24fee1f4 100644 --- a/tests/unit/fs/test_s3.py +++ b/tests/unit/fs/test_s3.py @@ -10,8 +10,6 @@ key_id = "key-id" key_secret = "key-secret" session_token = "session-token" -sse_customer_key = "3v2Q08ia8HOMGi5N21vsXI1fyfGAM6aVkW5pMu0A5dE=" -sse_customer_algorithm = "AES256" @pytest.fixture(autouse=True)