diff --git a/testbench/common.py b/testbench/common.py index 34233d69..f418d207 100644 --- a/testbench/common.py +++ b/testbench/common.py @@ -16,6 +16,8 @@ import base64 from functools import wraps +import gzip +import io import json import random import re @@ -801,6 +803,19 @@ def wrapper(*args, **kwargs): return retry_test +def handle_gzip_request(request): + """ + Handle gzip compressed JSON payloads when Content-Encoding: gzip is present on metadata requests. + No decompressions for media uploads when object's metadata includes Content-Encoding: gzip. + """ + if ( + request.headers.get("Content-Encoding", None) == "gzip" + and request.args.get("contentEncoding", None) != "gzip" + ): + request.data = gzip.decompress(request.data) + request.environ["wsgi.input"] = io.BytesIO(request.data) + + def rest_crc32c_to_proto(crc32c): """Convert from the REST representation of crc32c checksums to the proto representation. diff --git a/testbench/rest_server.py b/testbench/rest_server.py index 8a9e6ca0..90a91004 100644 --- a/testbench/rest_server.py +++ b/testbench/rest_server.py @@ -182,6 +182,11 @@ def start_grpc(): gcs.register_error_handler(Exception, testbench.error.RestException.handler) +@gcs.before_request +def handle_gzip_compressed_request(): + return testbench.common.handle_gzip_request(flask.request) + + # === BUCKET === # @@ -888,6 +893,11 @@ def download_object_get(bucket_name, object_name): upload.register_error_handler(Exception, testbench.error.RestException.handler) +@upload.before_request +def handle_gzip_compressed_request(): + return testbench.common.handle_gzip_request(flask.request) + + @upload.route("/b//o", methods=["POST"]) @retry_test(method="storage.objects.insert") def object_insert(bucket_name): diff --git a/testbench/servers/iam_rest_server.py b/testbench/servers/iam_rest_server.py index 6082ecc0..572e4503 100644 --- a/testbench/servers/iam_rest_server.py +++ b/testbench/servers/iam_rest_server.py @@ -28,6 +28,11 @@ _iam.register_error_handler(Exception, testbench.error.RestException.handler) +@_iam.before_request +def handle_gzip_compressed_request(): + return testbench.common.handle_gzip_request(flask.request) + + @_iam.route("/projects/-/serviceAccounts/:signBlob", methods=["POST"]) def sign_blob(service_account): """Implement the `projects.serviceAccounts.signBlob` API.""" diff --git a/testbench/servers/projects_rest_server.py b/testbench/servers/projects_rest_server.py index a3a1586f..193b9a52 100644 --- a/testbench/servers/projects_rest_server.py +++ b/testbench/servers/projects_rest_server.py @@ -30,6 +30,10 @@ def get_projects_app(db): projects.debug = False projects.register_error_handler(Exception, testbench.error.RestException.handler) + @projects.before_request + def handle_gzip_compressed_request(): + return testbench.common.handle_gzip_request(flask.request) + @projects.route("//serviceAccount") @retry_test("storage.serviceaccount.get") def projects_get(project_id): diff --git a/tests/test_common.py b/tests/test_common.py index 68083c2b..eb371c64 100644 --- a/tests/test_common.py +++ b/tests/test_common.py @@ -17,6 +17,7 @@ """Unit test for utils""" import base64 +import gzip import hashlib import json import types @@ -1205,6 +1206,46 @@ def test_bucket_to_from_proto(self): ), ) + def test_handle_gzip_request(self): + # Test gzip decompresses request payload when Content-Encoding: gzip is present. + payload = b'{"name": "bucket-name"}' + compressed = gzip.compress(payload) + request = testbench.common.FakeRequest( + headers={"Content-Encoding": "gzip"}, + args={}, + data=compressed, + environ={}, + ) + testbench.common.handle_gzip_request(request) + self.assertEqual(request.data, payload) + + # Test no decompression when object's metadata includes Content-Encoding: gzip, + # so that data can be stored in compressed form if applicable (PR#322). + media = "How vexingly quick daft zebras jump!" + compressed = gzip.compress(media.encode("utf-8")) + request = testbench.common.FakeRequest( + headers={"Content-Encoding": "gzip"}, + args={ + "name": "test_blob", + "uploadType": "media", + "contentEncoding": "gzip", + }, + data=compressed, + environ={}, + ) + testbench.common.handle_gzip_request(request) + self.assertEqual(request.data, compressed) + + # Test no decompression when Content-Encoding: gzip is not present. + request = testbench.common.FakeRequest( + headers={}, + args={}, + data=payload, + environ={}, + ) + testbench.common.handle_gzip_request(request) + self.assertEqual(request.data, payload) + if __name__ == "__main__": unittest.main()