simplify raw tcp protocol

libmproxy/protocol2/rawtcp.py

@@ -1,21 +1,67 @@
 from __future__ import (absolute_import, print_function, division)
+import socket
+import select
+
+from OpenSSL import SSL
 
-import OpenSSL
 from ..exceptions import ProtocolException
+from netlib.tcp import NetLibError
+from netlib.utils import cleanBin
 from ..protocol.tcp import TCPHandler
 from .layer import Layer
 
 
 class RawTcpLayer(Layer):
+    chunk_size = 4096
+
+    def __init__(self, ctx, logging=True):
+        self.logging = logging
+        super(RawTcpLayer, self).__init__(ctx)
+
     def __call__(self):
         self.connect()
-        tcp_handler = TCPHandler(self)
+
+        buf = memoryview(bytearray(self.chunk_size))
+
+        client = self.client_conn.connection
+        server = self.server_conn.connection
+        conns = [client, server]
+
         try:
-            tcp_handler.handle_messages()
-        except OpenSSL.SSL.Error as e:
-            raise ProtocolException("SSL error: %s" % repr(e), e)
+            while True:
+                r, _, _ = select.select(conns, [], [], 10)
+                for conn in r:
 
+                    size = conn.recv_into(buf, self.chunk_size)
+                    if not size:
+                        conns.remove(conn)
+                        # Shutdown connection to the other peer
+                        if isinstance(conn, SSL.Connection):
+                            # We can't half-close a connection, so we just close everything here.
+                            # Sockets will be cleaned up on a higher level.
+                            return
+                        else:
+                            conn.shutdown(socket.SHUT_WR)
 
-    def establish_server_connection(self):
-        pass
-        # FIXME: Remove method, currently just here to mock TCPHandler's call to it.
+                        if len(conns) == 0:
+                            return
+                        continue
+
+                    dst = server if conn == client else client
+                    dst.sendall(buf[:size])
+
+                    if self.logging:
+                        # log messages are prepended with the client address,
+                        # hence the "weird" direction string.
+                        if dst == server:
+                            direction = "-> tcp -> {!r}".format(self.server_conn.address)
+                        else:
+                            direction = "<- tcp <- {!r}".format(self.server_conn.address)
+                        data = cleanBin(buf[:size].tobytes())
+                        self.log(
+                            "{}\r\n{}".format(direction, data),
+                            "info"
+                        )
+
+        except (socket.error, NetLibError, SSL.Error) as e:
+            raise ProtocolException("TCP connection closed unexpectedly: {}".format(repr(e)), e)