From 9bc162e5b25dff34f812d2986f1f1d96e5f17f4d Mon Sep 17 00:00:00 2001
From: Tetsuya Morimoto <tetsuya.morimoto@gmail.com>
Date: Wed, 28 Jan 2015 00:01:55 +0900
Subject: [PATCH] Fixed to read out all data from socket when Data Frame
 payload is quite large like SETTINGS_MAX_FRAME_SIZE

---
 hyper/http20/connection.py | 26 +++++++++++++++++++++++---
 test/test_hyper.py         | 26 ++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/hyper/http20/connection.py b/hyper/http20/connection.py
index 9d8e0bac..144e4736 100644
--- a/hyper/http20/connection.py
+++ b/hyper/http20/connection.py
@@ -483,11 +483,31 @@ def _consume_single_frame(self):
         frame, length = Frame.parse_frame_header(header)
 
         # Read the remaining data from the socket.
-        if length:
+        data = self._recv_payload(length)
+        self._consume_frame_payload(frame, data)
+
+    def _recv_payload(self, length):
+        if not length:
+            return memoryview(b'')
+
+        buffer = bytearray(length)
+        buffer_view = memoryview(buffer)
+        index = 0
+        # _sock.recv(length) might not read out all data if the given length is
+        # very large. So it should be to retrieve from socket repeatedly.
+        while True:
             data = self._sock.recv(length)
-        else:
-            data = memoryview(b'')
+            data_length = len(data)
+            end = index + data_length
+            buffer_view[index:end] = data[:]
+            length -= data_length
+            if length == 0 or data_length == 0:
+                break
+            index = end
+
+        return buffer_view[:end]
 
+    def _consume_frame_payload(self, frame, data):
         frame.parse_body(data)
 
         log.info(
diff --git a/test/test_hyper.py b/test/test_hyper.py
index 675b63ae..083d58f0 100644
--- a/test/test_hyper.py
+++ b/test/test_hyper.py
@@ -1072,6 +1072,23 @@ def test_we_can_read_from_the_socket(self):
         s = c.recent_stream
         assert s.data == [b'testdata']
 
+    def test_we_can_read_fitfully_from_the_socket(self):
+        sock = DummyFitfullySocket()
+        sock.buffer = BytesIO(
+            b'\x00\x00\x18\x00\x01\x00\x00\x00\x01'
+            b'testdata'
+            b'+payload'
+        )
+
+        c = HTTP20Connection('www.google.com')
+        c._sock = sock
+        c.putrequest('GET', '/')
+        c.endheaders()
+        c._recv_cb()
+
+        s = c.recent_stream
+        assert s.data == [b'testdata+payload']
+
     def test_putrequest_sends_data(self):
         sock = DummySocket()
 
@@ -2045,6 +2062,15 @@ def recv(self, l):
     def close(self):
         pass
 
+
+class DummyFitfullySocket(DummySocket):
+    def recv(self, l):
+        length = l
+        if l != 9 and l >= 4:
+            length = int(round(l / 2))
+        return memoryview(self.buffer.read(length))
+
+
 class DummyStream(object):
     def __init__(self, data, trailers=None):
         self.data = data