feat: generate signed URLs for blobs/buckets using virtual hostname (#58)

tseaver · web-flow · commit 23df542d0669 · 2020-02-13T15:31:56.000-05:00
* Add 'virtual_hosted_style' arg to 'Blob.generate_signed_url'
* Add 'virtual_hosted_style arg to 'Bucket.generate_signed_url'
diff --git a/google/cloud/storage/blob.py b/google/cloud/storage/blob.py
@@ -361,6 +361,7 @@ def generate_signed_url(
         version=None,
         service_account_email=None,
         access_token=None,
+        virtual_hosted_style=False,
     ):
         """Generates a signed URL for this blob.
 
@@ -454,6 +455,11 @@ def generate_signed_url(
         :type access_token: str
         :param access_token: (Optional) Access token for a service account.
 
+        :type virtual_hosted_style: bool
+        :param virtual_hosted_style:
+            (Optional) If true, then construct the URL relative the bucket's
+            virtual hostname, e.g., '<bucket-name>.storage.googleapis.com'.
+
         :raises: :exc:`ValueError` when version is invalid.
         :raises: :exc:`TypeError` when expiration is not a valid type.
         :raises: :exc:`AttributeError` if credentials is not an instance
@@ -469,9 +475,16 @@ def generate_signed_url(
             raise ValueError("'version' must be either 'v2' or 'v4'")
 
         quoted_name = _quote(self.name, safe=b"/~")
-        resource = "/{bucket_name}/{quoted_name}".format(
-            bucket_name=self.bucket.name, quoted_name=quoted_name
-        )
+
+        if virtual_hosted_style:
+            api_access_endpoint = "https://{bucket_name}.storage.googleapis.com".format(
+                bucket_name=self.bucket.name
+            )
+            resource = "/{quoted_name}".format(quoted_name=quoted_name)
+        else:
+            resource = "/{bucket_name}/{quoted_name}".format(
+                bucket_name=self.bucket.name, quoted_name=quoted_name
+            )
 
         if credentials is None:
             client = self._require_client(client)
diff --git a/google/cloud/storage/bucket.py b/google/cloud/storage/bucket.py
@@ -2354,6 +2354,7 @@ def generate_signed_url(
         client=None,
         credentials=None,
         version=None,
+        virtual_hosted_style=False,
     ):
         """Generates a signed URL for this bucket.
 
@@ -2416,6 +2417,11 @@ def generate_signed_url(
         :param version: (Optional) The version of signed credential to create.
                         Must be one of 'v2' | 'v4'.
 
+        :type virtual_hosted_style: bool
+        :param virtual_hosted_style:
+            (Optional) If true, then construct the URL relative the bucket's
+            virtual hostname, e.g., '<bucket-name>.storage.googleapis.com'.
+
         :raises: :exc:`ValueError` when version is invalid.
         :raises: :exc:`TypeError` when expiration is not a valid type.
         :raises: :exc:`AttributeError` if credentials is not an instance
@@ -2430,7 +2436,13 @@ def generate_signed_url(
         elif version not in ("v2", "v4"):
             raise ValueError("'version' must be either 'v2' or 'v4'")
 
-        resource = "/{bucket_name}".format(bucket_name=self.name)
+        if virtual_hosted_style:
+            api_access_endpoint = "https://{bucket_name}.storage.googleapis.com".format(
+                bucket_name=self.name
+            )
+            resource = "/"
+        else:
+            resource = "/{bucket_name}".format(bucket_name=self.name)
 
         if credentials is None:
             client = self._require_client(client)
diff --git a/tests/unit/test_blob.py b/tests/unit/test_blob.py
@@ -399,6 +399,7 @@ def _generate_signed_url_helper(
         encryption_key=None,
         access_token=None,
         service_account_email=None,
+        virtual_hosted_style=False,
     ):
         from six.moves.urllib import parse
         from google.cloud._helpers import UTC
@@ -442,6 +443,7 @@ def _generate_signed_url_helper(
                 version=version,
                 access_token=access_token,
                 service_account_email=service_account_email,
+                virtual_hosted_style=virtual_hosted_style,
             )
 
         self.assertEqual(signed_uri, signer.return_value)
@@ -452,7 +454,17 @@ def _generate_signed_url_helper(
             expected_creds = credentials
 
         encoded_name = blob_name.encode("utf-8")
-        expected_resource = "/name/{}".format(parse.quote(encoded_name, safe=b"/~"))
+        quoted_name = parse.quote(encoded_name, safe=b"/~")
+
+        if virtual_hosted_style:
+            expected_api_access_endpoint = "https://{}.storage.googleapis.com".format(
+                bucket.name
+            )
+            expected_resource = "/{}".format(quoted_name)
+        else:
+            expected_api_access_endpoint = api_access_endpoint
+            expected_resource = "/{}/{}".format(bucket.name, quoted_name)
+
         if encryption_key is not None:
             expected_headers = headers or {}
             if effective_version == "v2":
@@ -465,7 +477,7 @@ def _generate_signed_url_helper(
         expected_kwargs = {
             "resource": expected_resource,
             "expiration": expiration,
-            "api_access_endpoint": api_access_endpoint,
+            "api_access_endpoint": expected_api_access_endpoint,
             "method": method.upper(),
             "content_md5": content_md5,
             "content_type": content_type,
@@ -604,6 +616,9 @@ def test_generate_signed_url_v4_w_csek_and_headers(self):
             encryption_key=os.urandom(32), headers={"x-goog-foo": "bar"}
         )
 
+    def test_generate_signed_url_v4_w_virtual_hostname(self):
+        self._generate_signed_url_v4_helper(virtual_hosted_style=True)
+
     def test_generate_signed_url_v4_w_credentials(self):
         credentials = object()
         self._generate_signed_url_v4_helper(credentials=credentials)
diff --git a/tests/unit/test_bucket.py b/tests/unit/test_bucket.py
@@ -2739,6 +2739,7 @@ def _generate_signed_url_helper(
         query_parameters=None,
         credentials=None,
         expiration=None,
+        virtual_hosted_style=False,
     ):
         from six.moves.urllib import parse
         from google.cloud._helpers import UTC
@@ -2773,6 +2774,7 @@ def _generate_signed_url_helper(
                 headers=headers,
                 query_parameters=query_parameters,
                 version=version,
+                virtual_hosted_style=virtual_hosted_style,
             )
 
         self.assertEqual(signed_uri, signer.return_value)
@@ -2782,12 +2784,19 @@ def _generate_signed_url_helper(
         else:
             expected_creds = credentials
 
-        encoded_name = bucket_name.encode("utf-8")
-        expected_resource = "/{}".format(parse.quote(encoded_name))
+        if virtual_hosted_style:
+            expected_api_access_endpoint = "https://{}.storage.googleapis.com".format(
+                bucket_name
+            )
+            expected_resource = "/"
+        else:
+            expected_api_access_endpoint = api_access_endpoint
+            expected_resource = "/{}".format(parse.quote(bucket_name))
+
         expected_kwargs = {
             "resource": expected_resource,
             "expiration": expiration,
-            "api_access_endpoint": api_access_endpoint,
+            "api_access_endpoint": expected_api_access_endpoint,
             "method": method.upper(),
             "headers": headers,
             "query_parameters": query_parameters,
@@ -2916,6 +2925,9 @@ def test_generate_signed_url_v4_w_credentials(self):
         credentials = object()
         self._generate_signed_url_v4_helper(credentials=credentials)
 
+    def test_generate_signed_url_v4_w_virtual_hostname(self):
+        self._generate_signed_url_v4_helper(virtual_hosted_style=True)
+
 
 class _Connection(object):
     _delete_bucket = False
