yield -> callbacks

12 files changed, 156 insertions, 205 deletions

diff --git a/libmproxy/protocol2/__init__.py b/libmproxy/protocol2/__init__.py
index cf6032da..d5dafaae 100644
--- a/libmproxy/protocol2/__init__.py
+++ b/libmproxy/protocol2/__init__.py
@@ -4,8 +4,7 @@ from .socks_proxy import Socks5Proxy
 from .reverse_proxy import ReverseProxy
 from .http_proxy import HttpProxy, HttpUpstreamProxy
 from .rawtcp import RawTcpLayer
-from . import messages
 
 __all__ = [
-    "Socks5Proxy", "RawTcpLayer", "RootContext", "ReverseProxy", "HttpProxy", "HttpUpstreamProxy", "messages"
+    "Socks5Proxy", "RawTcpLayer", "RootContext", "ReverseProxy", "HttpProxy", "HttpUpstreamProxy"
 ]
diff --git a/libmproxy/protocol2/http.py b/libmproxy/protocol2/http.py
index 53f40a72..db5aabaf 100644
--- a/libmproxy/protocol2/http.py
+++ b/libmproxy/protocol2/http.py
@@ -4,7 +4,7 @@ from .. import version
 from ..exceptions import InvalidCredentials, HttpException, ProtocolException
 from .layer import Layer
 from libmproxy import utils
-from .messages import SetServer, Connect, Reconnect, Kill
+from libmproxy.protocol2.layer import Kill
 from libmproxy.protocol import KILL
 
 from libmproxy.protocol.http import HTTPFlow
@@ -28,12 +28,21 @@ class Http1Layer(Layer):
         self.client_protocol = HTTP1Protocol(self.client_conn)
         self.server_protocol = HTTP1Protocol(self.server_conn)
 
+    def connect(self):
+        self.ctx.connect()
+        self.server_protocol = HTTP1Protocol(self.server_conn)
+
+    def reconnect(self):
+        self.ctx.reconnect()
+        self.server_protocol = HTTP1Protocol(self.server_conn)
+
+    def set_server(self, *args, **kwargs):
+        self.ctx.set_server(*args, **kwargs)
+        self.server_protocol = HTTP1Protocol(self.server_conn)
+
     def __call__(self):
         layer = HttpLayer(self, self.mode)
-        for message in layer():
-            yield message
-            self.server_protocol = HTTP1Protocol(self.server_conn)
-
+        layer()
 
 class Http2Layer(Layer):
     def __init__(self, ctx, mode):
@@ -42,11 +51,21 @@ class Http2Layer(Layer):
         self.client_protocol = HTTP2Protocol(self.client_conn, is_server=True)
         self.server_protocol = HTTP2Protocol(self.server_conn, is_server=False)
 
+    def connect(self):
+        self.ctx.connect()
+        self.server_protocol = HTTP2Protocol(self.server_conn)
+
+    def reconnect(self):
+        self.ctx.reconnect()
+        self.server_protocol = HTTP2Protocol(self.server_conn)
+
+    def set_server(self, *args, **kwargs):
+        self.ctx.set_server(*args, **kwargs)
+        self.server_protocol = HTTP2Protocol(self.server_conn)
+
     def __call__(self):
         layer = HttpLayer(self, self.mode)
-        for message in layer():
-            yield message
-            self.server_protocol = HTTP1Protocol(self.server_conn)
+        layer()
 
 
 def make_error_response(status_code, message, headers=None):
@@ -115,6 +134,37 @@ class ConnectServerConnection(object):
         return getattr(self.via, item)
 
 
+class UpstreamConnectLayer(Layer):
+    def __init__(self, ctx, connect_request):
+        super(UpstreamConnectLayer, self).__init__(ctx)
+        self.connect_request = connect_request
+        self.server_conn = ConnectServerConnection((connect_request.host, connect_request.port), self.ctx)
+
+    def __call__(self):
+        layer = self.ctx.next_layer(self)
+        layer()
+
+    def connect(self):
+        if not self.server_conn:
+            self.ctx.connect()
+            self.send_to_server(self.connect_request)
+        else:
+            pass  # swallow the message
+
+    def reconnect(self):
+        self.ctx.reconnect()
+        self.send_to_server(self.connect_request)
+
+    def set_server(self, address, server_tls, sni, depth=1):
+        if depth == 1:
+            if self.ctx.server_conn:
+                self.ctx.reconnect()
+            self.connect_request.host = address.host
+            self.connect_request.port = address.port
+            self.server_conn.address = address
+        else:
+            self.ctx.set_server(address, server_tls, sni, depth-1)
+
 class HttpLayer(Layer):
     """
     HTTP 1 Layer
@@ -146,22 +196,18 @@ class HttpLayer(Layer):
 
                 # Regular Proxy Mode: Handle CONNECT
                 if self.mode == "regular" and request.form_in == "authority":
-                    for message in self.handle_regular_mode_connect(request):
-                        yield message
+                    self.handle_regular_mode_connect(request)
                     return
 
                 # Make sure that the incoming request matches our expectations
                 self.validate_request(request)
 
                 flow.request = request
-                for message in self.process_request_hook(flow):
-                    yield message
+                self.process_request_hook(flow)
 
                 if not flow.response:
-                    for message in self.establish_server_connection(flow):
-                        yield message
-                    for message in self.get_response_from_server(flow):
-                        yield message
+                    self.establish_server_connection(flow)
+                    self.get_response_from_server(flow)
 
                 self.send_response_to_client(flow)
 
@@ -172,8 +218,7 @@ class HttpLayer(Layer):
 
                 # Upstream Proxy Mode: Handle CONNECT
                 if flow.request.form_in == "authority" and flow.response.code == 200:
-                    for message in self.handle_upstream_mode_connect(flow.request.copy()):
-                        yield message
+                    self.handle_upstream_mode_connect(flow.request.copy())
                     return
 
             except (HttpErrorConnClosed, NetLibError, HttpError) as e:
@@ -186,38 +231,14 @@ class HttpLayer(Layer):
                 flow.live = False
 
     def handle_regular_mode_connect(self, request):
-        yield SetServer((request.host, request.port), False, None)
+        self.set_server((request.host, request.port), False, None)
         self.send_to_client(make_connect_response(request.httpversion))
         layer = self.ctx.next_layer(self)
-        for message in layer():
-            yield message
+        layer()
 
     def handle_upstream_mode_connect(self, connect_request):
-        layer = self.ctx.next_layer(self)
-        self.server_conn = ConnectServerConnection((connect_request.host, connect_request.port), self.ctx)
-
-        for message in layer():
-            if message == Connect:
-                if not self.server_conn:
-                    yield message
-                    self.send_to_server(connect_request)
-                else:
-                    pass  # swallow the message
-            elif message == Reconnect:
-                yield message
-                self.send_to_server(connect_request)
-            elif message == SetServer:
-                if message.depth == 1:
-                    if self.ctx.server_conn:
-                        yield Reconnect()
-                    connect_request.host = message.address.host
-                    connect_request.port = message.address.port
-                    self.server_conn.address = message.address
-                else:
-                    message.depth -= 1
-                    yield message
-            else:
-                yield message
+        layer = UpstreamConnectLayer(self, connect_request)
+        layer()
 
     def check_close_connection(self, flow):
         """
@@ -313,14 +334,14 @@ class HttpLayer(Layer):
             # > server detects timeout, disconnects
             # > read (100-n)% of large request
             # > send large request upstream
-            yield Reconnect()
+            self.reconnect()
             get_response()
 
         # call the appropriate script hook - this is an opportunity for an
         # inline script to set flow.stream = True
         flow = self.channel.ask("responseheaders", flow)
         if flow is None or flow == KILL:
-            yield Kill()
+            raise Kill()
 
         if flow.response.stream and isinstance(self.server_protocol, http1.HTTP1Protocol):
             flow.response.content = CONTENT_MISSING
@@ -345,7 +366,7 @@ class HttpLayer(Layer):
         )
         response_reply = self.channel.ask("response", flow)
         if response_reply is None or response_reply == KILL:
-            yield Kill()
+            raise Kill()
 
     def process_request_hook(self, flow):
         # Determine .scheme, .host and .port attributes for inline scripts.
@@ -363,10 +384,10 @@ class HttpLayer(Layer):
             flow.request.port = self.ctx.server_conn.address.port
             flow.request.scheme = "https" if self.server_conn.tls_established else "http"
 
-        # TODO: Expose SetServer functionality to inline scripts somehow? (yield_from_callback?)
+        # TODO: Expose .set_server functionality to inline scripts
         request_reply = self.channel.ask("request", flow)
         if request_reply is None or request_reply == KILL:
-            yield Kill()
+            raise Kill()
         if isinstance(request_reply, HTTPResponse):
             flow.response = request_reply
             return
@@ -378,10 +399,10 @@ class HttpLayer(Layer):
         if self.mode == "regular" or self.mode == "transparent":
             # If there's an existing connection that doesn't match our expectations, kill it.
             if address != self.server_conn.address or tls != self.server_conn.ssl_established:
-                yield SetServer(address, tls, address.host)
+                self.set_server(address, tls, address.host)
             # Establish connection is neccessary.
             if not self.server_conn:
-                yield Connect()
+                self.connect()
 
             # SetServer is not guaranteed to work with TLS:
             # If there's not TlsLayer below which could catch the exception,
@@ -391,16 +412,16 @@ class HttpLayer(Layer):
 
         else:
             if not self.server_conn:
-                yield Connect()
+                self.connect()
             if tls:
                 raise HttpException("Cannot change scheme in upstream proxy mode.")
             """
             # This is a very ugly (untested) workaround to solve a very ugly problem.
             if self.server_conn and self.server_conn.tls_established and not ssl:
-                yield Reconnect()
+                self.reconnect()
             elif ssl and not hasattr(self, "connected_to") or self.connected_to != address:
                 if self.server_conn.tls_established:
-                    yield Reconnect()
+                    self.reconnect()
 
                 self.send_to_server(make_connect_request(address))
                 tls_layer = TlsLayer(self, False, True)
diff --git a/libmproxy/protocol2/http_proxy.py b/libmproxy/protocol2/http_proxy.py
index 652aa473..c24af6cf 100644
--- a/libmproxy/protocol2/http_proxy.py
+++ b/libmproxy/protocol2/http_proxy.py
@@ -7,11 +7,11 @@ from .http import Http1Layer
 class HttpProxy(Layer, ServerConnectionMixin):
     def __call__(self):
         layer = Http1Layer(self, "regular")
-        for message in layer():
-            if not self._handle_server_message(message):
-                yield message
-        if self.server_conn:
-            self._disconnect()
+        try:
+            layer()
+        finally:
+            if self.server_conn:
+                self._disconnect()
 
 class HttpUpstreamProxy(Layer, ServerConnectionMixin):
     def __init__(self, ctx, server_address):
@@ -19,8 +19,8 @@ class HttpUpstreamProxy(Layer, ServerConnectionMixin):
 
     def __call__(self):
         layer = Http1Layer(self, "upstream")
-        for message in layer():
-            if not self._handle_server_message(message):
-                yield message
-        if self.server_conn:
-            self._disconnect()
+        try:
+            layer()
+        finally:
+            if self.server_conn:
+                self._disconnect()
\ No newline at end of file
diff --git a/libmproxy/protocol2/layer.py b/libmproxy/protocol2/layer.py
index eb41bab7..7cb76591 100644
--- a/libmproxy/protocol2/layer.py
+++ b/libmproxy/protocol2/layer.py
@@ -21,7 +21,7 @@ Automated protocol detection by peeking into the buffer:
 
 Communication between layers is done as follows:
     - lower layers provide context information to higher layers
-    - higher layers can "yield" commands to lower layers,
+    - higher layers can call functions provided by lower layers,
       which are propagated until they reach a suitable layer.
 
 Further goals:
@@ -35,7 +35,6 @@ import threading
 from netlib import tcp
 from ..proxy import Log
 from ..proxy.connection import ServerConnection
-from .messages import Connect, Reconnect, SetServer, Kill
 from ..exceptions import ProtocolException
 
 
@@ -69,7 +68,7 @@ class Layer(_LayerCodeCompletion):
         """
         Logic of the layer.
         Raises:
-            ProxyError2 in case of protocol exceptions.
+            ProtocolException in case of protocol exceptions.
         """
         raise NotImplementedError
 
@@ -107,29 +106,20 @@ class ServerConnectionMixin(object):
         super(ServerConnectionMixin, self).__init__()
         self.server_conn = ServerConnection(server_address)
 
-    def _handle_server_message(self, message):
-        if message == Reconnect:
-            address = self.server_conn.address
-            self._disconnect()
+    def reconnect(self):
+        address = self.server_conn.address
+        self._disconnect()
+        self.server_conn.address = address
+        self.connect()
+
+    def set_server(self, address, server_tls, sni, depth=1):
+        if depth == 1:
+            if self.server_conn:
+                self._disconnect()
+            self.log("Set new server address: " + repr(address), "debug")
             self.server_conn.address = address
-            self._connect()
-            return True
-        elif message == Connect:
-            self._connect()
-            return True
-        elif message == SetServer:
-            if message.depth == 1:
-                if self.server_conn:
-                    self._disconnect()
-                self.log("Set new server address: " + repr(message.address), "debug")
-                self.server_conn.address = message.address
-                return True
-            else:
-                message.depth -= 1
-        elif message == Kill:
-            self._disconnect()
-
-        return False
+        else:
+            self.ctx.set_server(address, server_tls, sni, depth-1)
 
     def _disconnect(self):
         """
@@ -141,7 +131,7 @@ class ServerConnectionMixin(object):
         # self.channel.tell("serverdisconnect", self)
         self.server_conn = ServerConnection(None)
 
-    def _connect(self):
+    def connect(self):
         if not self.server_conn.address:
             raise ProtocolException("Cannot connect to server, no server address given.")
         self.log("serverconnect", "debug", [repr(self.server_conn.address)])
@@ -151,55 +141,7 @@ class ServerConnectionMixin(object):
             raise ProtocolException("Server connection to '%s' failed: %s" % (self.server_conn.address, e), e)
 
 
-def yield_from_callback(fun):
+class Kill(Exception):
     """
-    Decorator which makes it possible to yield from callbacks in the original thread.
-    As a use case, take the pyOpenSSL handle_sni callback: If we receive a new SNI from the client,
-    we need to reconnect to the server with the new SNI. Reconnecting would normally be done using "yield Reconnect()",
-    but we're in a pyOpenSSL callback here, outside of the main program flow. With this decorator, it looks as follows:
-
-    def handle_sni(self):
-        # ...
-        self.yield_from_callback(Reconnect())
-
-    @yield_from_callback
-    def establish_ssl_with_client():
-        self.client_conn.convert_to_ssl(...)
-
-    for message in self.establish_ssl_with_client():  # will yield Reconnect at some point
-        yield message
-
-
-    Limitations:
-        - You cannot yield True.
-    """
-    yield_queue = Queue.Queue()
-
-    def do_yield(msg):
-        yield_queue.put(msg)
-        yield_queue.get()
-
-    def wrapper(self, *args, **kwargs):
-        self.yield_from_callback = do_yield
-
-        def run():
-            try:
-                fun(self, *args, **kwargs)
-                yield_queue.put(True)
-            except Exception as e:
-                yield_queue.put(e)
-
-        threading.Thread(target=run, name="YieldFromCallbackThread").start()
-        while True:
-            msg = yield_queue.get()
-            if msg is True:
-                break
-            elif isinstance(msg, Exception):
-                raise ProtocolException("Error in %s: %s" % (fun.__name__, repr(msg)), msg)
-            else:
-                yield msg
-                yield_queue.put(None)
-
-        self.yield_from_callback = None
-
-    return wrapper
+    Kill a connection.
+    """
\ No newline at end of file
diff --git a/libmproxy/protocol2/messages.py b/libmproxy/protocol2/messages.py
index f5907537..de049486 100644
--- a/libmproxy/protocol2/messages.py
+++ b/libmproxy/protocol2/messages.py
@@ -44,7 +44,3 @@ class SetServer(_Message):
         self.depth = depth
 
 
-class Kill(_Message):
-    """
-    Kill a connection.
-    """
\ No newline at end of file
diff --git a/libmproxy/protocol2/rawtcp.py b/libmproxy/protocol2/rawtcp.py
index 167c8c79..6819ad6e 100644
--- a/libmproxy/protocol2/rawtcp.py
+++ b/libmproxy/protocol2/rawtcp.py
@@ -4,12 +4,11 @@ import OpenSSL
 from ..exceptions import ProtocolException
 from ..protocol.tcp import TCPHandler
 from .layer import Layer
-from .messages import Connect
 
 
 class RawTcpLayer(Layer):
     def __call__(self):
-        yield Connect()
+        self.connect()
         tcp_handler = TCPHandler(self)
         try:
             tcp_handler.handle_messages()
diff --git a/libmproxy/protocol2/reverse_proxy.py b/libmproxy/protocol2/reverse_proxy.py
index 767107ad..9d5a4beb 100644
--- a/libmproxy/protocol2/reverse_proxy.py
+++ b/libmproxy/protocol2/reverse_proxy.py
@@ -16,8 +16,9 @@ class ReverseProxy(Layer, ServerConnectionMixin):
             layer = TlsLayer(self, self._client_tls, self._server_tls)
         else:
             layer = self.ctx.next_layer(self)
-        for message in layer():
-            if not self._handle_server_message(message):
-                yield message
-        if self.server_conn:
-            self._disconnect()
\ No newline at end of file
+
+        try:
+            layer()
+        finally:
+            if self.server_conn:
+                self._disconnect()
\ No newline at end of file
diff --git a/libmproxy/protocol2/root_context.py b/libmproxy/protocol2/root_context.py
index f8a645b0..f0e5b9a7 100644
--- a/libmproxy/protocol2/root_context.py
+++ b/libmproxy/protocol2/root_context.py
@@ -1,7 +1,7 @@
 from __future__ import (absolute_import, print_function, division)
 import string
 
-from .messages import Kill
+from libmproxy.protocol2.layer import Kill
 from .rawtcp import RawTcpLayer
 from .tls import TlsLayer
 from .http import Http1Layer, Http2Layer, HttpLayer
diff --git a/libmproxy/protocol2/socks_proxy.py b/libmproxy/protocol2/socks_proxy.py
index 5bb8e5f8..18b363d5 100644
--- a/libmproxy/protocol2/socks_proxy.py
+++ b/libmproxy/protocol2/socks_proxy.py
@@ -5,7 +5,7 @@ from ..proxy import ProxyError, Socks5ProxyMode
 from .layer import Layer, ServerConnectionMixin
 
 
-class Socks5Proxy(ServerConnectionMixin, Layer):
+class Socks5Proxy(Layer, ServerConnectionMixin):
     def __call__(self):
         try:
             s5mode = Socks5ProxyMode(self.config.ssl_ports)
@@ -16,9 +16,12 @@ class Socks5Proxy(ServerConnectionMixin, Layer):
 
         self.server_conn.address = address
 
+        # TODO: Kill event
+
         layer = self.ctx.next_layer(self)
-        for message in layer():
-            if not self._handle_server_message(message):
-                yield message
-        if self.server_conn:
-            self._disconnect()
\ No newline at end of file
+
+        try:
+            layer()
+        finally:
+            if self.server_conn:
+                self._disconnect()
\ No newline at end of file
diff --git a/libmproxy/protocol2/tls.py b/libmproxy/protocol2/tls.py
index 970abe62..28480388 100644
--- a/libmproxy/protocol2/tls.py
+++ b/libmproxy/protocol2/tls.py
@@ -6,8 +6,7 @@ from netlib import tcp
 import netlib.http.http2
 
 from ..exceptions import ProtocolException
-from .layer import Layer, yield_from_callback
-from .messages import Connect, Reconnect, SetServer
+from .layer import Layer
 
 
 class TlsLayer(Layer):
@@ -50,36 +49,35 @@ class TlsLayer(Layer):
         )
 
         if client_tls_requires_server_cert:
-            for m in self._establish_tls_with_client_and_server():
-                yield m
+            self._establish_tls_with_client_and_server()
         elif self._client_tls:
-            for m in self._establish_tls_with_client():
-                yield m
+            self._establish_tls_with_client()
 
         layer = self.ctx.next_layer(self)
+        layer()
 
-        for message in layer():
-            self.log("TlsLayer: %s" % message,"debug")
-            if not (message == Connect and self._connected):
-                yield message
+    def connect(self):
+        if not self.server_conn:
+            self.ctx.connect()
+        if self._server_tls and not self._server_tls_established:
+            self._establish_tls_with_server()
+
+    def reconnect(self):
+        self.ctx.reconnect()
+        if self._server_tls and not self._server_tls_established:
+            self._establish_tls_with_server()
 
-            if message == Connect or message == Reconnect:
-                if self._server_tls and not self._server_tls_established:
-                    self._establish_tls_with_server()
-            if message == SetServer and message.depth == 1:
-                if message.server_tls is not None:
-                    self._sni_from_server_change = message.sni
-                    self._server_tls = message.server_tls
+    def set_server(self, address, server_tls, sni, depth=1):
+        self.ctx.set_server(address, server_tls, sni, depth)
+        if server_tls is not None:
+            self._sni_from_server_change = sni
+            self._server_tls = server_tls
 
     @property
     def _server_tls_established(self):
         return self.server_conn and self.server_conn.tls_established
 
     @property
-    def _connected(self):
-        return bool(self.server_conn)
-
-    @property
     def sni_for_upstream_connection(self):
         if self._sni_from_server_change is False:
             return None
@@ -92,19 +90,14 @@ class TlsLayer(Layer):
         """
 
         # First, try to connect to the server.
-        yield Connect()
+        self.ctx.connect()
         server_err = None
         try:
             self._establish_tls_with_server()
         except ProtocolException as e:
             server_err = e
 
-        for message in self._establish_tls_with_client():
-            if message == Reconnect:
-                yield message
-                self._establish_tls_with_server()
-            else:
-                raise RuntimeError("Unexpected Message: %s" % message)
+        self._establish_tls_with_client()
 
         if server_err and not self.client_sni:
             raise server_err
@@ -125,7 +118,7 @@ class TlsLayer(Layer):
             if old_upstream_sni != self.sni_for_upstream_connection:
                 # Perform reconnect
                 if self.server_conn and self._server_tls:
-                    self.yield_from_callback(Reconnect())
+                    self.reconnect()
 
             if self.client_sni:
                 # Now, change client context to reflect possibly changed certificate:
@@ -156,7 +149,7 @@ class TlsLayer(Layer):
             # Perform reconnect
             # TODO: Avoid double reconnect.
             if self.server_conn and self._server_tls:
-                self.yield_from_callback(Reconnect())
+                self.reconnect()
 
         self.client_alpn_protos = options
 
@@ -165,7 +158,6 @@ class TlsLayer(Layer):
         else:  # pragma no cover
             return options[0]
 
-    @yield_from_callback
     def _establish_tls_with_client(self):
         self.log("Establish TLS with client", "debug")
         cert, key, chain_file = self._find_cert()
diff --git a/libmproxy/protocol2/transparent_proxy.py b/libmproxy/protocol2/transparent_proxy.py
index 28ad3726..9263dbde 100644
--- a/libmproxy/protocol2/transparent_proxy.py
+++ b/libmproxy/protocol2/transparent_proxy.py
@@ -18,8 +18,8 @@ class TransparentProxy(Layer, ServerConnectionMixin):
             raise ProtocolException("Transparent mode failure: %s" % repr(e), e)
 
         layer = self.ctx.next_layer(self)
-        for message in layer():
-            if not self._handle_server_message(message):
-                yield message
-        if self.server_conn:
-            self._disconnect()
\ No newline at end of file
+        try:
+            layer()
+        finally:
+            if self.server_conn:
+                self._disconnect()
\ No newline at end of file
diff --git a/libmproxy/proxy/server.py b/libmproxy/proxy/server.py
index e23a7d72..9957caa0 100644
--- a/libmproxy/proxy/server.py
+++ b/libmproxy/proxy/server.py
@@ -3,6 +3,7 @@ from __future__ import absolute_import, print_function
 import traceback
 import sys
 import socket
+from libmproxy.protocol2.layer import Kill
 from netlib import tcp
 
 from ..protocol.handle import protocol_handler
@@ -88,12 +89,9 @@ class ConnectionHandler2:
             root_layer = protocol2.HttpProxy(root_context)
 
         try:
-            for message in root_layer():
-                if message == protocol2.messages.Kill:
-                    self.log("Connection killed", "info")
-                    break
-
-                print("Root layer receveived: %s" % message)
+            root_layer()
+        except Kill as e:
+            self.log("Connection killed", "info")
         except ProtocolException as e:
             self.log(e, "info")
         except Exception: