fix bugs, fix tests

5 files changed, 58 insertions, 56 deletions

diff --git a/libmproxy/protocol2/http.py b/libmproxy/protocol2/http.py
index f093f7c5..973f169c 100644
--- a/libmproxy/protocol2/http.py
+++ b/libmproxy/protocol2/http.py
@@ -266,12 +266,15 @@ class HttpLayer(Layer):
                     self.handle_upstream_mode_connect(flow.request.copy())
                     return
 
-            except (HttpErrorConnClosed, NetLibError, HttpError) as e:
+            except (HttpErrorConnClosed, NetLibError, HttpError, ProtocolException) as e:
                 self.send_to_client(make_error_response(
                     getattr(e, "code", 502),
                     repr(e)
                 ))
-                raise ProtocolException(repr(e), e)
+                if isinstance(e, ProtocolException):
+                    raise e
+                else:
+                    raise ProtocolException(repr(e), e)
             finally:
                 flow.live = False
 
@@ -468,7 +471,7 @@ class HttpLayer(Layer):
 
     def validate_request(self, request):
         if request.form_in == "absolute" and request.scheme != "http":
-            self.send_response(make_error_response(400, "Invalid request scheme: %s" % request.scheme))
+            self.send_to_client(make_error_response(400, "Invalid request scheme: %s" % request.scheme))
             raise HttpException("Invalid request scheme: %s" % request.scheme)
 
         expected_request_forms = {
diff --git a/libmproxy/protocol2/root_context.py b/libmproxy/protocol2/root_context.py
index 9b18f0aa..d0c62be4 100644
--- a/libmproxy/protocol2/root_context.py
+++ b/libmproxy/protocol2/root_context.py
@@ -41,7 +41,7 @@ class RootContext(object):
         d = top_layer.client_conn.rfile.peek(3)
         is_ascii = (
             len(d) == 3 and
-            all(x in string.ascii_uppercase for x in d)
+            all(x in string.ascii_letters for x in d) # better be safe here and don't expect uppercase...
         )
 
         d = top_layer.client_conn.rfile.peek(len(HTTP2Protocol.CLIENT_CONNECTION_PREFACE))
diff --git a/libmproxy/protocol2/tls.py b/libmproxy/protocol2/tls.py
index 28480388..98c5d603 100644
--- a/libmproxy/protocol2/tls.py
+++ b/libmproxy/protocol2/tls.py
@@ -17,6 +17,7 @@ class TlsLayer(Layer):
         self.client_sni = None
         self._sni_from_server_change = None
         self.client_alpn_protos = None
+        self.__server_tls_exception = None
 
         # foo alpn protos = [netlib.http.http1.HTTP1Protocol.ALPN_PROTO_HTTP1, netlib.http.http2.HTTP2Protocol.ALPN_PROTO_H2],  # TODO: read this from client_conn first
 
@@ -107,49 +108,48 @@ class TlsLayer(Layer):
         This callback gets called during the TLS handshake with the client.
         The client has just sent the Sever Name Indication (SNI).
         """
-        try:
-            old_upstream_sni = self.sni_for_upstream_connection
-
-            sn = connection.get_servername()
-            if not sn:
-                return
-            self.client_sni = sn.decode("utf8").encode("idna")
-
-            if old_upstream_sni != self.sni_for_upstream_connection:
-                # Perform reconnect
-                if self.server_conn and self._server_tls:
-                    self.reconnect()
-
-            if self.client_sni:
-                # Now, change client context to reflect possibly changed certificate:
-                cert, key, chain_file = self._find_cert()
-                new_context = self.client_conn.create_ssl_context(
-                    cert, key,
-                    method=self.config.openssl_method_client,
-                    options=self.config.openssl_options_client,
-                    cipher_list=self.config.ciphers_client,
-                    dhparams=self.config.certstore.dhparams,
-                    chain_file=chain_file,
-                    alpn_select_callback=self.__handle_alpn_select,
-                )
-                connection.set_context(new_context)
-        # An unhandled exception in this method will core dump PyOpenSSL, so
-        # make dang sure it doesn't happen.
-        except:  # pragma: no cover
-            self.log("Error in handle_sni:\r\n" + traceback.format_exc(), "error")
+        old_upstream_sni = self.sni_for_upstream_connection
+
+        sn = connection.get_servername()
+        if not sn:
+            return
+
+        self.client_sni = sn.decode("utf8").encode("idna")
+
+        server_sni_changed = (old_upstream_sni != self.sni_for_upstream_connection)
+        server_conn_with_tls_exists = (self.server_conn and self._server_tls)
+        if server_sni_changed and server_conn_with_tls_exists:
+            try:
+                self.reconnect()
+            except Exception as e:
+                self.__server_tls_exception = e
+
+        # Now, change client context to reflect possibly changed certificate:
+        cert, key, chain_file = self._find_cert()
+        new_context = self.client_conn.create_ssl_context(
+            cert, key,
+            method=self.config.openssl_method_client,
+            options=self.config.openssl_options_client,
+            cipher_list=self.config.ciphers_client,
+            dhparams=self.config.certstore.dhparams,
+            chain_file=chain_file,
+            alpn_select_callback=self.__handle_alpn_select,
+        )
+        connection.set_context(new_context)
 
     def __handle_alpn_select(self, conn_, options):
         # TODO: change to something meaningful?
-        alpn_preference = netlib.http.http1.HTTP1Protocol.ALPN_PROTO_HTTP1
+        # alpn_preference = netlib.http.http1.HTTP1Protocol.ALPN_PROTO_HTTP1
         alpn_preference = netlib.http.http2.HTTP2Protocol.ALPN_PROTO_H2
-        ###
 
-        # TODO: Not
-        if self.client_alpn_protos != options:
-            # Perform reconnect
-            # TODO: Avoid double reconnect.
-            if self.server_conn and self._server_tls:
+        # TODO: Don't reconnect twice?
+        upstream_alpn_changed = (self.client_alpn_protos != options)
+        server_conn_with_tls_exists = (self.server_conn and self._server_tls)
+        if upstream_alpn_changed and server_conn_with_tls_exists:
+            try:
                 self.reconnect()
+            except Exception as e:
+                self.__server_tls_exception = e
 
         self.client_alpn_protos = options
 
@@ -177,6 +177,11 @@ class TlsLayer(Layer):
             print("alpn: %s" % self.client_alpn_protos)
             raise ProtocolException(repr(e), e)
 
+        # Do not raise server tls exceptions immediately.
+        # We want to try to finish the client handshake so that other layers can send error messages over it.
+        if self.__server_tls_exception:
+            raise self.__server_tls_exception
+
     def _establish_tls_with_server(self):
         self.log("Establish TLS with server", "debug")
         try:
diff --git a/test/test_proxy.py b/test/test_proxy.py
index 6ab19e02..b9eec53b 100644
--- a/test/test_proxy.py
+++ b/test/test_proxy.py
@@ -36,7 +36,7 @@ class TestServerConnection:
         sc.send(protocol.assemble(f.request))
 
         protocol = http.http1.HTTP1Protocol(rfile=sc.rfile)
-        assert protocol.read_response(f.request.method, 1000)
+        assert protocol.read_response(f.request, 1000)
         assert self.d.last_log()
 
         sc.finish()
diff --git a/test/test_server.py b/test/test_server.py
index 77ba4576..529024f5 100644
--- a/test/test_server.py
+++ b/test/test_server.py
@@ -319,17 +319,6 @@ class TestHTTPAuth(tservers.HTTPProxTest):
         assert ret.status_code == 202
 
 
-class TestHTTPConnectSSLError(tservers.HTTPProxTest):
-    certfile = True
-
-    def test_go(self):
-        self.config.ssl_ports.append(self.proxy.port)
-        p = self.pathoc_raw()
-        dst = ("localhost", self.proxy.port)
-        p.connect(connect_to=dst)
-        tutils.raises("502 - Bad Gateway", p.http_connect, dst)
-
-
 class TestHTTPS(tservers.HTTPProxTest, CommonMixin, TcpMixin):
     ssl = True
     ssloptions = pathod.SSLOptions(request_client_cert=True)
@@ -390,26 +379,31 @@ class TestHTTPSUpstreamServerVerificationWBadCert(tservers.HTTPProxTest):
             ("untrusted-cert", tutils.test_data.path("data/untrusted-server.crt"))
         ])
 
+    def _request(self):
+        p = self.pathoc()
+        # We need to make an actual request because the upstream connection is lazy-loaded.
+        return p.request("get:/p/242")
+
     def test_default_verification_w_bad_cert(self):
         """Should use no verification."""
         self.config.openssl_trusted_ca_server = tutils.test_data.path(
             "data/trusted-cadir/trusted-ca.pem")
 
-        self.pathoc()
+        assert self._request().status_code == 242
 
     def test_no_verification_w_bad_cert(self):
         self.config.openssl_verification_mode_server = SSL.VERIFY_NONE
         self.config.openssl_trusted_ca_server = tutils.test_data.path(
             "data/trusted-cadir/trusted-ca.pem")
 
-        self.pathoc()
+        assert self._request().status_code == 242
 
     def test_verification_w_bad_cert(self):
         self.config.openssl_verification_mode_server = SSL.VERIFY_PEER
         self.config.openssl_trusted_ca_server = tutils.test_data.path(
             "data/trusted-cadir/trusted-ca.pem")
 
-        tutils.raises("SSL handshake error", self.pathoc)
+        assert self._request().status_code == 502
 
 
 class TestHTTPSNoCommonName(tservers.HTTPProxTest):