launchpad-reviewers team mailing list archive

[Colin Watson <cjwatson@xxxxxxxxxxxxx>]

Colin Watson has proposed merging ~cjwatson/turnip:netstring-recv-split-data into turnip:master.

Commit message:
Fix netstring_recv to handle large responses

If the response is large, then netstring_recv may need to make multiple
sock.recv calls to receive all the data.

Requested reviews:
  Launchpad code reviewers (launchpad-reviewers)

For more details, see:
https://code.launchpad.net/~cjwatson/turnip/+git/turnip/+merge/359740
-- 
Your team Launchpad code reviewers is requested to review the proposed merge of ~cjwatson/turnip:netstring-recv-split-data into turnip:master.

diff --git a/turnip/pack/hooks/hook.py b/turnip/pack/hooks/hook.py
index 1594a4b..bb9c5c7 100755
--- a/turnip/pack/hooks/hook.py
+++ b/turnip/pack/hooks/hook.py
@@ -104,9 +104,11 @@ def netstring_recv(sock):
         lengthstr += c
         c = sock.recv(1)
     length = int(lengthstr)
-    s = sock.recv(length)
+    s = bytearray()
+    while len(s) < length:
+        s += sock.recv(length - len(s))
     assert sock.recv(1) == b','
-    return s
+    return bytes(s)
 
 
 def rpc_invoke(sock, method, args):
diff --git a/turnip/pack/tests/test_hooks.py b/turnip/pack/tests/test_hooks.py
index ebf9c80..5127349 100644
--- a/turnip/pack/tests/test_hooks.py
+++ b/turnip/pack/tests/test_hooks.py
@@ -79,6 +79,51 @@ class MockRepo(object):
         return MockRef(self.ancestor)
 
 
+class MockSocket(object):
+
+    def __init__(self, blocks=None):
+        self.blocks = blocks or []
+
+    def recv(self, bufsize, flags=None):
+        if not self.blocks:
+            return b''
+        elif bufsize < len(self.blocks[0]):
+            block = self.blocks[0][:bufsize]
+            self.blocks[0] = self.blocks[0][bufsize:]
+            return block
+        else:
+            return self.blocks.pop(0)
+
+
+class TestNetstringRecv(TestCase):
+    """Tests for netstring_recv."""
+
+    def test_nondigit(self):
+        sock = MockSocket([b'zzz:abc,'])
+        self.assertRaises(AssertionError, hook.netstring_recv, sock)
+
+    def test_short(self):
+        sock = MockSocket([b'4:abc,'])
+        self.assertRaises(AssertionError, hook.netstring_recv, sock)
+
+    def test_unterminated(self):
+        sock = MockSocket([b'4:abcd'])
+        self.assertRaises(AssertionError, hook.netstring_recv, sock)
+
+    def test_split_data(self):
+        sock = MockSocket([b'12:abcd', b'efgh', b'ijkl,'])
+        self.assertEqual(b'abcdefghijkl', hook.netstring_recv(sock))
+        self.assertEqual([], sock.blocks)
+
+    def test_valid(self):
+        sock = MockSocket(
+            [b'11:\x00\x01\x02\x03\x04,\x05\x06\x07\x08\x09,remaining'])
+        self.assertEqual(
+            b'\x00\x01\x02\x03\x04,\x05\x06\x07\x08\x09',
+            hook.netstring_recv(sock))
+        self.assertEqual([b'remaining'], sock.blocks)
+
+
 class HookTestMixin(object):
     run_tests_with = AsynchronousDeferredRunTest.make_factory(timeout=5)