Merge pull request #46 from maralla/tranport_error

fix tranport error

tests/test_protocol_cybinary.py

@@ -6,9 +6,17 @@
 import pytest
 pytestmark = pytest.mark.skipif(PYPY, reason="cybin not enabled in pypy.")
 
+import multiprocessing
+import time
+
 from thriftpy._compat import u
 from thriftpy.thrift import TType, TPayload
-from thriftpy.transport import TMemoryBuffer, TCyBufferedTransport
+from thriftpy.transport import (
+    TMemoryBuffer,
+    TCyBufferedTransport,
+    TSocket,
+    TServerSocket
+)
 from thriftpy.utils import hexlify
 
 if not PYPY:
@@ -309,3 +317,27 @@ def test_skip_struct():
 
     proto.skip(b, TType.STRUCT)
     assert 123 == proto.read_val(b, TType.I32)
+
+
+def test_read_long_data():
+    val = 'z' * 97 * 1024
+
+    def serve():
+        server_sock = TServerSocket(host="127.0.0.1", port=9090)
+        server_sock.listen()
+        client = server_sock.accept()
+        t = TCyBufferedTransport(client)
+        proto.write_val(t, TType.STRING, val)
+        t.flush()
+
+    p = multiprocessing.Process(target=serve)
+    p.start()
+    time.sleep(0.1)
+
+    try:
+        sock = TSocket(host='127.0.0.1', port=9090)
+        b = TCyBufferedTransport(sock)
+        b.open()
+        assert val == proto.read_val(b, TType.STRING)
+    finally:
+        p.terminate()

thriftpy/transport/cytransport.pyx

@@ -1,6 +1,8 @@
 from libc.stdlib cimport malloc, free
 from libc.string cimport memcpy, memmove
 
+from ..transport import TTransportException
+
 DEF DEFAULT_BUFFER = 4096
 DEF MIN_BUFFER_SZIE = 1024
 
@@ -104,8 +106,14 @@ cdef class TCyBufferedTransport(object):
             new_data = self.trans.read(cap)
             new_data_len = len(new_data)
 
-            if new_data_len < sz - self.rbuf.data_size:
-                raise Exception("Transport error")
+            while new_data_len < sz - self.rbuf.data_size:
+                more = self.trans.read(cap - new_data_len)
+                if len(more) <= 0:
+                    raise TTransportException(
+                        TTransportException.END_OF_FILE,
+                        "End of file reading from transport")
+                new_data += more
+                new_data_len += len(more)
 
             if cap - self.rbuf.cur < new_data_len:
                 self.rbuf.move_to_start()