yield -> callbacks

1 files changed, 23 insertions, 31 deletions

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()