Limit reads to block length.

2 files changed, 17 insertions, 2 deletions

diff --git a/netlib/tcp.py b/netlib/tcp.py
index 414c1237..f8f877de 100644
--- a/netlib/tcp.py
+++ b/netlib/tcp.py
@@ -106,13 +106,17 @@ class Writer(_FileLike):
 class Reader(_FileLike):
     def read(self, length):
         """
-            If length is None, we read until connection closes.
+            If length is -1, we read until connection closes.
         """
         result = ''
         start = time.time()
         while length == -1 or length > 0:
+            if length == -1 or length > self.BLOCKSIZE:
+                rlen = self.BLOCKSIZE
+            else:
+                rlen = length
             try:
-                data = self.o.read(self.BLOCKSIZE if length == -1 else length)
+                data = self.o.read(rlen)
             except SSL.ZeroReturnError:
                 break
             except SSL.WantReadError:
diff --git a/test/test_tcp.py b/test/test_tcp.py
index c833ce07..5a12da91 100644
--- a/test/test_tcp.py
+++ b/test/test_tcp.py
@@ -255,6 +255,17 @@ class TestTCPClient:
 
 
 class TestFileLike:
+    def test_blocksize(self):
+        s = cStringIO.StringIO("1234567890abcdefghijklmnopqrstuvwxyz")
+        s = tcp.Reader(s)
+        s.BLOCKSIZE = 2
+        assert s.read(1) == "1"
+        assert s.read(2) == "23"
+        assert s.read(3) == "456"
+        assert s.read(4) == "7890"
+        d = s.read(-1)
+        assert d.startswith("abc") and d.endswith("xyz")
+
     def test_wrap(self):
         s = cStringIO.StringIO("foobar\nfoobar")
         s.flush()