various fixes

11 files changed, 78 insertions, 26 deletions

diff --git a/libmproxy/filt.py b/libmproxy/filt.py
index bd17a807..25747bc6 100644
--- a/libmproxy/filt.py
+++ b/libmproxy/filt.py
@@ -246,14 +246,14 @@ class FSrc(_Rex):
     help = "Match source address"
 
     def __call__(self, f):
-        return f.client_conn and re.search(self.expr, repr(f.client_conn.address))
+        return f.client_conn.address and re.search(self.expr, repr(f.client_conn.address))
 
 class FDst(_Rex):
     code = "dst"
     help = "Match destination address"
 
     def __call__(self, f):
-        return f.server_conn and re.search(self.expr, repr(f.server_conn.address))
+        return f.server_conn.address and re.search(self.expr, repr(f.server_conn.address))
 
 class _Int(_Action):
     def __init__(self, num):
diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py
index 4c15c80d..4472cb2a 100644
--- a/libmproxy/protocol/http.py
+++ b/libmproxy/protocol/http.py
@@ -729,6 +729,7 @@ class RequestReplayThread(threading.Thread):
             if not self.flow.response:
                 # In all modes, we directly connect to the server displayed
                 if self.config.mode == "upstream":
+                    # FIXME
                     server_address = self.config.mode.get_upstream_server(
                         self.flow.client_conn
                     )[2:]
diff --git a/libmproxy/protocol2/http.py b/libmproxy/protocol2/http.py
index eadde3b3..53f40a72 100644
--- a/libmproxy/protocol2/http.py
+++ b/libmproxy/protocol2/http.py
@@ -10,13 +10,14 @@ from libmproxy.protocol import KILL
 from libmproxy.protocol.http import HTTPFlow
 from libmproxy.protocol.http_wrappers import HTTPResponse, HTTPRequest
 from netlib import tcp
-from netlib.http import status_codes, http1, HttpErrorConnClosed
+from netlib.http import status_codes, http1, HttpErrorConnClosed, HttpError
 from netlib.http.semantics import CONTENT_MISSING
 from netlib import odict
 from netlib.tcp import NetLibError, Address
 from netlib.http.http1 import HTTP1Protocol
 from netlib.http.http2 import HTTP2Protocol
 
+
 # TODO: The HTTP2 layer is missing multiplexing, which requires a major rewrite.
 
 
@@ -31,6 +32,7 @@ class Http1Layer(Layer):
         layer = HttpLayer(self, self.mode)
         for message in layer():
             yield message
+            self.server_protocol = HTTP1Protocol(self.server_conn)
 
 
 class Http2Layer(Layer):
@@ -41,10 +43,10 @@ class Http2Layer(Layer):
         self.server_protocol = HTTP2Protocol(self.server_conn, is_server=False)
 
     def __call__(self):
-        # FIXME: Handle Reconnect etc.
         layer = HttpLayer(self, self.mode)
         for message in layer():
             yield message
+            self.server_protocol = HTTP1Protocol(self.server_conn)
 
 
 def make_error_response(status_code, message, headers=None):
@@ -100,6 +102,7 @@ class ConnectServerConnection(object):
     """
     "Fake" ServerConnection to represent state after a CONNECT request to an upstream proxy.
     """
+
     def __init__(self, address, ctx):
         self.address = tcp.Address.wrap(address)
         self._ctx = ctx
@@ -124,6 +127,8 @@ class HttpLayer(Layer):
     def __call__(self):
         while True:
             try:
+                flow = HTTPFlow(self.client_conn, self.server_conn, live=True)
+
                 try:
                     request = HTTPRequest.from_protocol(
                         self.client_protocol,
@@ -148,7 +153,6 @@ class HttpLayer(Layer):
                 # Make sure that the incoming request matches our expectations
                 self.validate_request(request)
 
-                flow = HTTPFlow(self.client_conn, self.server_conn)
                 flow.request = request
                 for message in self.process_request_hook(flow):
                     yield message
@@ -164,15 +168,22 @@ class HttpLayer(Layer):
                 if self.check_close_connection(flow):
                     return
 
+                # TODO: Implement HTTP Upgrade
+
                 # 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
                     return
 
-            except (HttpErrorConnClosed, NetLibError) as e:
-                make_error_response(502, repr(e))
+            except (HttpErrorConnClosed, NetLibError, HttpError) as e:
+                self.send_to_client(make_error_response(
+                    getattr(e, "code", 502),
+                    repr(e)
+                ))
                 raise ProtocolException(repr(e), e)
+            finally:
+                flow.live = False
 
     def handle_regular_mode_connect(self, request):
         yield SetServer((request.host, request.port), False, None)
@@ -267,21 +278,43 @@ class HttpLayer(Layer):
 
             for chunk in chunks:
                 for part in chunk:
+                    # TODO: That's going to fail.
                     self.send_to_client(part)
                 self.client_conn.wfile.flush()
 
             flow.response.timestamp_end = utils.timestamp()
 
     def get_response_from_server(self, flow):
-        # TODO: Add second attempt.
-        self.send_to_server(flow.request)
-
-        flow.response = HTTPResponse.from_protocol(
-            self.server_protocol,
-            flow.request.method,
-            body_size_limit=self.config.body_size_limit,
-            include_body=False,
-        )
+        def get_response():
+            self.send_to_server(flow.request)
+            # Only get the headers at first...
+            flow.response = HTTPResponse.from_protocol(
+                self.server_protocol,
+                flow.request.method,
+                body_size_limit=self.config.body_size_limit,
+                include_body=False,
+            )
+
+        try:
+            get_response()
+        except (tcp.NetLibError, HttpErrorConnClosed) as v:
+            self.log(
+                "server communication error: %s" % repr(v),
+                level="debug"
+            )
+            # In any case, we try to reconnect at least once. This is
+            # necessary because it might be possible that we already
+            # initiated an upstream connection after clientconnect that
+            # has already been expired, e.g consider the following event
+            # log:
+            # > clientconnect (transparent mode destination known)
+            # > serverconnect (required for client tls handshake)
+            # > read n% of large request
+            # > server detects timeout, disconnects
+            # > read (100-n)% of large request
+            # > send large request upstream
+            yield Reconnect()
+            get_response()
 
         # call the appropriate script hook - this is an opportunity for an
         # inline script to set flow.stream = True
@@ -293,7 +326,6 @@ class HttpLayer(Layer):
             flow.response.content = CONTENT_MISSING
         else:
             flow.response.content = self.server_protocol.read_http_body(
-                self.server_conn,
                 flow.response.headers,
                 self.config.body_size_limit,
                 flow.request.method,
@@ -405,7 +437,7 @@ class HttpLayer(Layer):
                 self.send_to_client(make_error_response(
                     407,
                     "Proxy Authentication Required",
-                    self.config.authenticator.auth_challenge_headers()
+                    odict.ODictCaseless([[k,v] for k, v in self.config.authenticator.auth_challenge_headers().items()])
                 ))
                 raise InvalidCredentials("Proxy Authentication Required")
 
diff --git a/libmproxy/protocol2/http_proxy.py b/libmproxy/protocol2/http_proxy.py
index 19b5f7ef..652aa473 100644
--- a/libmproxy/protocol2/http_proxy.py
+++ b/libmproxy/protocol2/http_proxy.py
@@ -10,7 +10,8 @@ class HttpProxy(Layer, ServerConnectionMixin):
         for message in layer():
             if not self._handle_server_message(message):
                 yield message
-
+        if self.server_conn:
+            self._disconnect()
 
 class HttpUpstreamProxy(Layer, ServerConnectionMixin):
     def __init__(self, ctx, server_address):
@@ -21,3 +22,5 @@ class HttpUpstreamProxy(Layer, ServerConnectionMixin):
         for message in layer():
             if not self._handle_server_message(message):
                 yield message
+        if self.server_conn:
+            self._disconnect()
diff --git a/libmproxy/protocol2/layer.py b/libmproxy/protocol2/layer.py
index 67f3d549..eb41bab7 100644
--- a/libmproxy/protocol2/layer.py
+++ b/libmproxy/protocol2/layer.py
@@ -109,7 +109,9 @@ class ServerConnectionMixin(object):
 
     def _handle_server_message(self, message):
         if message == Reconnect:
+            address = self.server_conn.address
             self._disconnect()
+            self.server_conn.address = address
             self._connect()
             return True
         elif message == Connect:
diff --git a/libmproxy/protocol2/reverse_proxy.py b/libmproxy/protocol2/reverse_proxy.py
index 2ee3d9d8..767107ad 100644
--- a/libmproxy/protocol2/reverse_proxy.py
+++ b/libmproxy/protocol2/reverse_proxy.py
@@ -19,3 +19,5 @@ class ReverseProxy(Layer, ServerConnectionMixin):
         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
diff --git a/libmproxy/protocol2/root_context.py b/libmproxy/protocol2/root_context.py
index 6ba6ca9a..f8a645b0 100644
--- a/libmproxy/protocol2/root_context.py
+++ b/libmproxy/protocol2/root_context.py
@@ -1,4 +1,5 @@
 from __future__ import (absolute_import, print_function, division)
+import string
 
 from .messages import Kill
 from .rawtcp import RawTcpLayer
@@ -36,6 +37,8 @@ class RootContext(object):
             d[2] in ('\x00', '\x01', '\x02', '\x03')
         )
 
+        is_ascii = all(x in string.ascii_uppercase for x in d)
+
         # TODO: build is_http2_magic check here, maybe this is an easy way to detect h2c
 
         if not d:
@@ -43,10 +46,11 @@ class RootContext(object):
 
         if is_tls_client_hello:
             return TlsLayer(top_layer, True, True)
-        elif isinstance(top_layer, TlsLayer) and top_layer.client_conn.get_alpn_proto_negotiated() == 'h2':
+        elif isinstance(top_layer, TlsLayer) and is_ascii:
+            if top_layer.client_conn.get_alpn_proto_negotiated() == 'h2':
                 return Http2Layer(top_layer, 'transparent')
-        elif isinstance(top_layer, TlsLayer) and isinstance(top_layer.ctx, Http1Layer):
-            return Http1Layer(top_layer, "transparent")
+            else:
+                return Http1Layer(top_layer, "transparent")
         else:
             return RawTcpLayer(top_layer)
 
diff --git a/libmproxy/protocol2/socks_proxy.py b/libmproxy/protocol2/socks_proxy.py
index c6126a42..5bb8e5f8 100644
--- a/libmproxy/protocol2/socks_proxy.py
+++ b/libmproxy/protocol2/socks_proxy.py
@@ -20,3 +20,5 @@ class Socks5Proxy(ServerConnectionMixin, Layer):
         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
diff --git a/libmproxy/protocol2/transparent_proxy.py b/libmproxy/protocol2/transparent_proxy.py
index 4ed4c14b..28ad3726 100644
--- a/libmproxy/protocol2/transparent_proxy.py
+++ b/libmproxy/protocol2/transparent_proxy.py
@@ -21,3 +21,5 @@ class TransparentProxy(Layer, ServerConnectionMixin):
         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
diff --git a/libmproxy/proxy/connection.py b/libmproxy/proxy/connection.py
index f92b53aa..c9b57998 100644
--- a/libmproxy/proxy/connection.py
+++ b/libmproxy/proxy/connection.py
@@ -27,6 +27,9 @@ class ClientConnection(tcp.BaseHandler, stateobject.StateObject):
         self.timestamp_ssl_setup = None
         self.protocol = None
 
+    def __nonzero__(self):
+        return bool(self.connection) and not self.finished
+
     def __repr__(self):
         return "<ClientConnection: {ssl}{host}:{port}>".format(
             ssl="[ssl] " if self.ssl_established else "",
@@ -89,7 +92,7 @@ class ServerConnection(tcp.TCPClient, stateobject.StateObject):
     def __init__(self, address):
         tcp.TCPClient.__init__(self, address)
 
-        self.state = []  # a list containing (conntype, state) tuples
+        self.via = None
         self.timestamp_start = None
         self.timestamp_end = None
         self.timestamp_tcp_setup = None
@@ -97,7 +100,7 @@ class ServerConnection(tcp.TCPClient, stateobject.StateObject):
         self.protocol = None
 
     def __nonzero__(self):
-        return bool(self.connection)
+        return bool(self.connection) and not self.finished
 
     def __repr__(self):
         if self.ssl_established and self.sni:
@@ -117,7 +120,6 @@ class ServerConnection(tcp.TCPClient, stateobject.StateObject):
         return self.ssl_established
 
     _stateobject_attributes = dict(
-        state=list,
         timestamp_start=float,
         timestamp_end=float,
         timestamp_tcp_setup=float,
@@ -187,3 +189,5 @@ class ServerConnection(tcp.TCPClient, stateobject.StateObject):
     def finish(self):
         tcp.TCPClient.finish(self)
         self.timestamp_end = utils.timestamp()
+
+ServerConnection._stateobject_attributes["via"] = ServerConnection
\ No newline at end of file
diff --git a/test/test_protocol_http.py b/test/test_protocol_http.py
index 2da54093..c6a9159c 100644
--- a/test/test_protocol_http.py
+++ b/test/test_protocol_http.py
@@ -56,7 +56,7 @@ class TestInvalidRequests(tservers.HTTPProxTest):
         p = self.pathoc()
         r = p.request("connect:'%s:%s'" % ("127.0.0.1", self.server2.port))
         assert r.status_code == 400
-        assert "Must not CONNECT on already encrypted connection" in r.body
+        assert "Invalid HTTP request form" in r.body
 
     def test_relative_request(self):
         p = self.pathoc_raw()