Add peek_request and peek_response

These handler methods get a first pass at socket data before the
consuming recv. These can be used to, for example, hook into libraries
that read from sockets themselves or otherwise consume the data in the
handler.

This is the start of moving all the SSL MiTM code out of Connection and
into SSL specific handlers as well as adding handlers dynamically.

This change is pretty straightfoward except for the work to be done to
support peek on mitm'd connections. pyOpenSSL does not support peek so
we need to read into a buffer and read from that when peeking. This
requires some pretty hacky code to keep select working correct on a
connection where there is data remaning in the peek buffer as the
underlying socket is no longer ready to be selected for reading (as it
would be with real MSG_PEEK) so in that case we use a different fd for
select that is always ready for reading.

nogotofail/mitm/connection/connection.py

@@ -28,35 +28,81 @@
 import os
 
 class ConnectionWrapper(object):
-    """Wrapper around OpenSSL's Connection object to make recv act like socket.recv()
+    """Wrapper around OpenSSL's Connection object to make it act like a real socket.
     """
 
     def __init__(self, connection):
         self._connection = connection
+        self.buffer = ""
+        self._is_short_send = False
 
     def __getattr__(self, name):
         return getattr(self._connection, name)
 
-    def recv(self, size):
+    def recv(self, size, flags=0):
         """Wrapper around pyOpenSSL's Connection.recv
         PyOpenSSL doesn't return "" on error like socket.recv does,
         instead it throws a SSL.ZeroReturnError or (-1, "Unexpected EOF") erorrs.
 
         Wrap recv so we don't have to deal with that noise.
         """
-        buf = ""
+        if flags & socket.MSG_PEEK == 0:
+            return self._recv(size)
+        if len(self.buffer) >= size:
+            return self.buffer[:size]
         try:
-            buf = self._connection.recv(size)
+            self.buffer += self._recv(size - len(self.buffer))
+        except SSL.WantReadError:
+            pass
+        return self.buffer[:size]
+
+    def _recv(self, size):
+        if size <= len(self.buffer):
+            out = self.buffer[:size]
+            self.buffer = self.buffer[size:]
+            return out
+        buf = self.buffer
+        size -= len(buf)
+        try:
+            buf += self._connection.recv(size)
         except SSL.SysCallError as e:
             if e.args != (-1, "Unexpected EOF"):
                 raise e
         except SSL.ZeroReturnError:
             pass
+        except SSL.WantReadError as e:
+            # Rethrow the WantRead if we really have no data
+            if not buf:
+                raise e
         except SSL.Error as e:
             if e.args != (-1, "Unexpected EOF"):
                 raise e
+        self.buffer = ""
         return buf
 
+    def send(self, string):
+        sent = self._connection.send(string)
+        # Track short send state for our awful fileno hacks
+        self._is_short_send = sent != len(string)
+        return sent
+
+    _always_read_fd = None
+    def always_read_fd(self):
+        """Return an fd that is always ready for read when passed to select.select. See fileno for why this is needed."""
+        if ConnectionWrapper._always_read_fd:
+            return ConnectionWrapper._always_read_fd
+        ConnectionWrapper._always_read_fd = open("/dev/zero")
+        return ConnectionWrapper._always_read_fd
+
+    def fileno(self):
+        # _AWFUL_ HACK to support MSG_PEEK without breaking select.select.
+        # If we read data with a peeking recv then return a fd that is always selectable on read to make sure the connection keeps flowing.
+        # Note that if the conneciton is handling a short send then we're only waiting for write not read, so use the underlying connection.
+        # Once the backlog is sent the connection will start trying to read again and we'll return the always_read_fd.
+        if self.buffer and not self._is_short_send:
+            return self.always_read_fd().fileno()
+
+        return self._connection.fileno()
 
 def stub_verify(conn, cert, errno, errdepth, code):
     """We don't verify the server when we attempt a MiTM.
@@ -411,17 +457,20 @@ def _handle_hello(self, client_hello):
 
     def _bridge_client(self):
         try:
-            # Check for a TLS client hello we might need to intercept
-            if not self.ssl:
+            try:
                 client_request = self.client_socket.recv(65536, socket.MSG_PEEK)
-                if not client_request:
-                    return False
-                # If a MiTM was attempted discard client_request, we used it
-                # for establishing a MiTM with the client.
-                if self._check_for_ssl(client_request):
+                handled = self.handler.peek_request(client_request)
+                if handled:
                     return not self.closed
-
-            try:
+                for handler in self.data_handlers:
+                    if handler.peek_request(client_request):
+                        return not self.closed
+                # Check for a TLS client hello we might need to intercept
+                if not self.ssl:
+                    # If a MiTM was attempted discard client_request, we used it
+                    # for establishing a MiTM with the client.
+                    if self._check_for_ssl(client_request):
+                        return not self.closed
                 client_request = self.client_socket.recv(65536)
             except (socket.error, SSL.WantReadError):
                 # recv can still time out even if select returned this socket
@@ -456,6 +505,13 @@ def _bridge_client(self):
     def _bridge_server(self):
         try:
             try:
+                server_response = self.server_socket.recv(65536, socket.MSG_PEEK)
+                handled = self.handler.peek_response(server_response)
+                if handled:
+                    return not self.closed
+                for handler in self.data_handlers:
+                    if handler.peek_response(server_response):
+                        return not self.closed
                 server_response = self.server_socket.recv(65536)
             except (socket.error, SSL.WantReadError):
                 # recv can still time out even if select returned this socket

nogotofail/mitm/connection/handlers/base.py

@@ -101,6 +101,24 @@ def on_ssl(self, client_hello):
         """
         pass
 
+    def peek_request(self, request):
+        """Called with the data from a request _before_ it has been read from the socket.
+
+        This can be used to prempt the socket recv and handle data yourself.
+
+        Returns if the request should be considered handled and recv should not be called on the underlying socket
+        """
+        return False
+
+    def peek_response(self, response):
+        """Called with the data from a response _before_ it has been read from the socket.
+
+        This can be used to prempt the socket recv and handle data yourself.
+
+        Returns if the response should be considered handled and recv should not be called on the underlying socket
+        """
+        return False
+
 
 class BaseConnectionHandler(BaseHandler):