Merge pull request #3103 from OpenNeuroOrg/support-gzip

fix(worker): Support gzip for git clients that require it

services/datalad/datalad_service/common/stream.py

@@ -1,5 +1,6 @@
 import os
 import tempfile
+import zlib
 
 CHUNK_SIZE_BYTES = 2048
 
@@ -23,9 +24,13 @@ def update_file(path, stream):
             raise
 
 
-def pipe_chunks(reader, writer):
-    while True:
-        chunk = reader.read(CHUNK_SIZE_BYTES)
-        if not chunk:
-            break
-        writer.write(chunk)
+def pipe_chunks(reader, writer, gzipped=False):
+    # If gzipped, we have to read the entire request and write once
+    if gzipped:
+        writer.write(zlib.decompress(reader.read(), zlib.MAX_WBITS|32))
+    else:
+        while True:
+            chunk = reader.read(CHUNK_SIZE_BYTES)
+            if not chunk:
+                break
+            writer.write(chunk)

services/datalad/datalad_service/handlers/git.py

@@ -92,7 +92,7 @@ def on_post(self, req, resp, worker, dataset):
                 os.symlink('/hooks/pre-receive', pre_receive_path)
             process = subprocess.Popen(
                 ['git-receive-pack', '--stateless-rpc', dataset_path], stdin=subprocess.PIPE, stdout=subprocess.PIPE)
-            pipe_chunks(reader=req.bounded_stream, writer=process.stdin)
+            pipe_chunks(reader=req.bounded_stream, writer=process.stdin, gzipped=req.get_header('content-encoding') == 'gzip')
             process.stdin.flush()
             resp.status = falcon.HTTP_OK
             resp.stream = process.stdout
@@ -118,7 +118,7 @@ def on_post(self, req, resp, worker, dataset):
             dataset_path = self.store.get_dataset_path(dataset)
             process = subprocess.Popen(
                 ['git-upload-pack', '--stateless-rpc', dataset_path], stdin=subprocess.PIPE, stdout=subprocess.PIPE)
-            pipe_chunks(reader=req.bounded_stream, writer=process.stdin)
+            pipe_chunks(reader=req.bounded_stream, writer=process.stdin, gzipped=req.get_header('content-encoding') == 'gzip')
             process.stdin.flush()
             resp.status = falcon.HTTP_OK
             resp.stream = process.stdout

services/datalad/tests/test_git.py

@@ -1,4 +1,5 @@
 import os
+import zlib
 
 import falcon
 from falcon import testing
@@ -7,6 +8,10 @@
 from datalad_service.common import git
 from datalad.api import Dataset
 
+
+gzip_compress = zlib.compressobj(9, zlib.DEFLATED, zlib.MAX_WBITS | 16)
+
+
 test_auth = "Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJmZDQ0ZjVjNS1iMjFiLTQyMGItOTU1NS1hZjg1NmVmYzk0NTIiLCJlbWFpbCI6Im5lbGxAc3F1aXNoeW1lZGlhLmNvbSIsInByb3ZpZGVyIjoiZ29vZ2xlIiwibmFtZSI6Ik5lbGwgSGFyZGNhc3RsZSIsImFkbWluIjp0cnVlLCJzY29wZXMiOlsiZGF0YXNldDpnaXQiXSwiZGF0YXNldCI6ImRzMDAwMDAxIiwiaWF0IjoxNjA4NDEwNjEyLCJleHAiOjIxNDc0ODM2NDd9.0aA9cZWMieYr9zbmVrTeFEhpATqmT_X4tVX1VR1uabA"
 
 
@@ -111,6 +116,24 @@ def test_git_upload_resource(client):
     assert response.content[0:12] == b'0008NAK\nPACK'
 
 
+def test_git_upload_resource_gzip(client):
+    ds_id = 'ds000001'
+    get_response = client.simulate_get(
+        f'/git/0/{ds_id}/info/refs?service=git-upload-pack', headers={"authorization": test_auth})
+    lines = get_response.content.decode().split('\n')
+    # Grab two refs to ask for
+    annex = lines[2][4:44]
+    head = lines[3][4:44]
+    upload_pack_input = "0032want {}\n0032want {}\n00000009done\n""".format(
+        head, annex)
+    gzipped_input = gzip_compress.compress(upload_pack_input.encode()) + gzip_compress.flush()
+    # Ask for them
+    response = client.simulate_post(
+        f'/git/0/{ds_id}/git-upload-pack', headers={"authorization": test_auth, "content-encoding": "gzip"}, body=gzipped_input)
+    assert response.status == falcon.HTTP_OK
+    # Just look for the start of a pack stream
+    assert response.content[0:12] == b'0008NAK\nPACK'
+
 def test_git_receive_resource(client):
     ds_id = 'ds000001'
     get_response = client.simulate_get(