fix(storage): enable CSEK w/ V4 signed URLs

storage/google/cloud/storage/blob.py

@@ -468,6 +468,17 @@ def generate_signed_url(
         else:
             helper = generate_signed_url_v4
 
+        if self._encryption_key is not None:
+            encryption_headers = _get_encryption_headers(self._encryption_key)
+            if headers is None:
+                headers = {}
+            if version == "v2":
+                # See: https://cloud.google.com/storage/docs/access-control/signed-urls-v2#about-canonical-extension-headers
+                v2_copy_only = "X-Goog-Encryption-Algorithm"
+                headers[v2_copy_only] = encryption_headers[v2_copy_only]
+            else:
+                headers.update(encryption_headers)
+
         return helper(
             credentials,
             resource=resource,

storage/tests/system.py

@@ -12,7 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import base64
 import datetime
+import hashlib
 import os
 import re
 import tempfile
@@ -860,11 +862,12 @@ def _create_signed_read_url_helper(
         version="v2",
         payload=None,
         expiration=None,
+        encryption_key=None,
     ):
         expiration = self._morph_expiration(version, expiration)
 
         if payload is not None:
-            blob = self.bucket.blob(blob_name)
+            blob = self.bucket.blob(blob_name, encryption_key=encryption_key)
             blob.upload_from_string(payload)
         else:
             blob = self.blob
@@ -873,7 +876,17 @@ def _create_signed_read_url_helper(
             expiration=expiration, method=method, client=Config.CLIENT, version=version
         )
 
-        response = requests.get(signed_url)
+        headers = {}
+
+        if encryption_key is not None:
+            headers["x-goog-encryption-algorithm"] = "AES256"
+            encoded_key = base64.b64encode(encryption_key).decode("utf-8")
+            headers["x-goog-encryption-key"] = encoded_key
+            key_hash = hashlib.sha256(encryption_key).digest()
+            key_hash = base64.b64encode(key_hash).decode("utf-8")
+            headers["x-goog-encryption-key-sha256"] = key_hash
+
+        response = requests.get(signed_url, headers=headers)
         self.assertEqual(response.status_code, 200)
         if payload is not None:
             self.assertEqual(response.content, payload)
@@ -916,6 +929,23 @@ def test_create_signed_read_url_v4_w_non_ascii_name(self):
             version="v4",
         )
 
+    def test_create_signed_read_url_v2_w_csek(self):
+        encryption_key = os.urandom(32)
+        self._create_signed_read_url_helper(
+            blob_name="v2-w-csek.txt",
+            payload=b"Test signed URL for blob w/ CSEK",
+            encryption_key=encryption_key,
+        )
+
+    def test_create_signed_read_url_v4_w_csek(self):
+        encryption_key = os.urandom(32)
+        self._create_signed_read_url_helper(
+            blob_name="v2-w-csek.txt",
+            payload=b"Test signed URL for blob w/ CSEK",
+            encryption_key=encryption_key,
+            version="v4",
+        )
+
     def _create_signed_delete_url_helper(self, version="v2", expiration=None):
         expiration = self._morph_expiration(version, expiration)
 

storage/tests/unit/test_blob.py

@@ -391,10 +391,12 @@ def _generate_signed_url_helper(
         query_parameters=None,
         credentials=None,
         expiration=None,
+        encryption_key=None,
     ):
         from six.moves.urllib import parse
         from google.cloud._helpers import UTC
         from google.cloud.storage.blob import _API_ACCESS_ENDPOINT
+        from google.cloud.storage.blob import _get_encryption_headers
 
         api_access_endpoint = api_access_endpoint or _API_ACCESS_ENDPOINT
 
@@ -406,7 +408,7 @@ def _generate_signed_url_helper(
         connection = _Connection()
         client = _Client(connection)
         bucket = _Bucket(client)
-        blob = self._make_one(blob_name, bucket=bucket)
+        blob = self._make_one(blob_name, bucket=bucket, encryption_key=encryption_key)
 
         if version is None:
             effective_version = "v2"
@@ -442,6 +444,15 @@ def _generate_signed_url_helper(
 
         encoded_name = blob_name.encode("utf-8")
         expected_resource = "/name/{}".format(parse.quote(encoded_name, safe=b"/~"))
+        if encryption_key is not None:
+            expected_headers = headers or {}
+            if effective_version == "v2":
+                expected_headers["X-Goog-Encryption-Algorithm"] = "AES256"
+            else:
+                expected_headers.update(_get_encryption_headers(encryption_key))
+        else:
+            expected_headers = headers
+
         expected_kwargs = {
             "resource": expected_resource,
             "expiration": expiration,
@@ -452,7 +463,7 @@ def _generate_signed_url_helper(
             "response_type": response_type,
             "response_disposition": response_disposition,
             "generation": generation,
-            "headers": headers,
+            "headers": expected_headers,
             "query_parameters": query_parameters,
         }
         signer.assert_called_once_with(expected_creds, **expected_kwargs)
@@ -514,6 +525,14 @@ def test_generate_signed_url_v2_w_generation(self):
     def test_generate_signed_url_v2_w_headers(self):
         self._generate_signed_url_v2_helper(headers={"x-goog-foo": "bar"})
 
+    def test_generate_signed_url_v2_w_csek(self):
+        self._generate_signed_url_v2_helper(encryption_key=os.urandom(32))
+
+    def test_generate_signed_url_v2_w_csek_and_headers(self):
+        self._generate_signed_url_v2_helper(
+            encryption_key=os.urandom(32), headers={"x-goog-foo": "bar"}
+        )
+
     def test_generate_signed_url_v2_w_credentials(self):
         credentials = object()
         self._generate_signed_url_v2_helper(credentials=credentials)
@@ -566,6 +585,14 @@ def test_generate_signed_url_v4_w_generation(self):
     def test_generate_signed_url_v4_w_headers(self):
         self._generate_signed_url_v4_helper(headers={"x-goog-foo": "bar"})
 
+    def test_generate_signed_url_v4_w_csek(self):
+        self._generate_signed_url_v4_helper(encryption_key=os.urandom(32))
+
+    def test_generate_signed_url_v4_w_csek_and_headers(self):
+        self._generate_signed_url_v4_helper(
+            encryption_key=os.urandom(32), headers={"x-goog-foo": "bar"}
+        )
+
     def test_generate_signed_url_v4_w_credentials(self):
         credentials = object()
         self._generate_signed_url_v4_helper(credentials=credentials)