Merge pull request #1126 from Kriechi/safeguard

Safeguard

4 files changed, 85 insertions, 13 deletions

diff --git a/mitmproxy/exceptions.py b/mitmproxy/exceptions.py
index 52fd36d1..8f989063 100644
--- a/mitmproxy/exceptions.py
+++ b/mitmproxy/exceptions.py
@@ -50,6 +50,10 @@ class HttpProtocolException(ProtocolException):
     pass
 
 
+class Http2ProtocolException(ProtocolException):
+    pass
+
+
 class ServerException(ProxyException):
     pass
 
@@ -83,4 +87,4 @@ class ScriptException(ProxyException):
 
 
 class FlowReadException(ProxyException):
-    pass
\ No newline at end of file
+    pass
diff --git a/mitmproxy/protocol/http.py b/mitmproxy/protocol/http.py
index fc181b5d..9cb35176 100644
--- a/mitmproxy/protocol/http.py
+++ b/mitmproxy/protocol/http.py
@@ -11,7 +11,7 @@ from netlib.http import Headers
 from h2.exceptions import H2Error
 
 from .. import utils
-from ..exceptions import HttpProtocolException, ProtocolException
+from ..exceptions import HttpProtocolException, Http2ProtocolException, ProtocolException
 from ..models import (
     HTTPFlow,
     HTTPResponse,
@@ -243,7 +243,7 @@ class HttpLayer(Layer):
         try:
             response = make_error_response(code, message)
             self.send_response(response)
-        except (NetlibException, H2Error):
+        except (NetlibException, H2Error, Http2ProtocolException):
             self.log(traceback.format_exc(), "debug")
 
     def change_upstream_proxy_server(self, address):
@@ -288,9 +288,9 @@ class HttpLayer(Layer):
 
         try:
             get_response()
-        except NetlibException as v:
+        except NetlibException as e:
             self.log(
-                "server communication error: %s" % repr(v),
+                "server communication error: %s" % repr(e),
                 level="debug"
             )
             # In any case, we try to reconnect at least once. This is
@@ -304,6 +304,11 @@ class HttpLayer(Layer):
             # > server detects timeout, disconnects
             # > read (100-n)% of large request
             # > send large request upstream
+
+            if isinstance(e, Http2ProtocolException):
+                # do not try to reconnect for HTTP2
+                raise ProtocolException("First and only attempt to get response via HTTP2 failed.")
+
             self.disconnect()
             self.connect()
             get_response()
diff --git a/mitmproxy/protocol/http2.py b/mitmproxy/protocol/http2.py
index 140a94d4..f2f31b17 100644
--- a/mitmproxy/protocol/http2.py
+++ b/mitmproxy/protocol/http2.py
@@ -4,6 +4,7 @@ import threading
 import time
 from six.moves import queue
 
+import traceback
 import h2
 import six
 from h2.connection import H2Connection
@@ -15,6 +16,7 @@ from netlib.utils import http2_read_raw_frame
 
 from .base import Layer
 from .http import _HttpTransmissionLayer, HttpLayer
+from ..exceptions import ProtocolException, Http2ProtocolException
 from .. import utils
 from ..models import HTTPRequest, HTTPResponse
 
@@ -55,7 +57,7 @@ class SafeH2Connection(H2Connection):
     def safe_send_headers(self, is_zombie, stream_id, headers):
         with self.lock:
             if is_zombie():  # pragma: no cover
-                return
+                raise Http2ProtocolException("Zombie Stream")
             self.send_headers(stream_id, headers.fields)
             self.conn.send(self.data_to_send())
 
@@ -66,7 +68,7 @@ class SafeH2Connection(H2Connection):
                 self.lock.acquire()
                 if is_zombie():  # pragma: no cover
                     self.lock.release()
-                    return
+                    raise Http2ProtocolException("Zombie Stream")
                 max_outbound_frame_size = self.max_outbound_frame_size
                 frame_chunk = chunk[position:position + max_outbound_frame_size]
                 if self.local_flow_control_window(stream_id) < len(frame_chunk):
@@ -79,7 +81,7 @@ class SafeH2Connection(H2Connection):
                 position += max_outbound_frame_size
         with self.lock:
             if is_zombie():  # pragma: no cover
-                return
+                raise Http2ProtocolException("Zombie Stream")
             self.end_stream(stream_id)
             self.conn.send(self.data_to_send())
 
@@ -104,7 +106,7 @@ class Http2Layer(Layer):
         self.active_conns.append(self.server_conn.connection)
 
     def connect(self):  # pragma: no cover
-        raise ValueError("CONNECT inside an HTTP2 stream is not supported.")
+        raise Http2ProtocolException("HTTP2 layer should already have a connection.")
 
     def set_server(self):  # pragma: no cover
         raise NotImplementedError("Cannot change server for HTTP2 connections.")
@@ -215,6 +217,7 @@ class Http2Layer(Layer):
                         raw_frame = b''.join(http2_read_raw_frame(source_conn.rfile))
                     except:
                         # read frame failed: connection closed
+                        # kill all streams
                         for stream in self.streams.values():
                             stream.zombie = time.time()
                         return
@@ -232,7 +235,7 @@ class Http2Layer(Layer):
 class Http2SingleStreamLayer(_HttpTransmissionLayer, threading.Thread):
 
     def __init__(self, ctx, stream_id, request_headers):
-        super(Http2SingleStreamLayer, self).__init__(ctx)
+        super(Http2SingleStreamLayer, self).__init__(ctx, name="Thread-Http2SingleStreamLayer-{}".format(stream_id))
         self.zombie = None
         self.client_stream_id = stream_id
         self.server_stream_id = None
@@ -335,7 +338,7 @@ class Http2SingleStreamLayer(_HttpTransmissionLayer, threading.Thread):
         with self.server_conn.h2.lock:
             # We must not assign a stream id if we are already a zombie.
             if self.zombie:  # pragma: no cover
-                return
+                raise Http2ProtocolException("Zombie Stream")
 
             self.server_stream_id = self.server_conn.h2.get_next_available_stream_id()
             self.server_to_client_stream_ids[self.server_stream_id] = self.client_stream_id
@@ -350,6 +353,8 @@ class Http2SingleStreamLayer(_HttpTransmissionLayer, threading.Thread):
             self.server_stream_id,
             message.body
         )
+        if self.zombie:  # pragma: no cover
+            raise Http2ProtocolException("Zombie Stream")
 
     def read_response_headers(self):
         self.response_arrived.wait()
@@ -377,7 +382,7 @@ class Http2SingleStreamLayer(_HttpTransmissionLayer, threading.Thread):
                     yield self.response_data_queue.get()
                 return
             if self.zombie:  # pragma: no cover
-                return
+                raise Http2ProtocolException("Zombie Stream")
 
     def send_response_headers(self, response):
         self.client_conn.h2.safe_send_headers(
@@ -385,6 +390,8 @@ class Http2SingleStreamLayer(_HttpTransmissionLayer, threading.Thread):
             self.client_stream_id,
             response.headers
         )
+        if self.zombie:  # pragma: no cover
+            raise Http2ProtocolException("Zombie Stream")
 
     def send_response_body(self, _response, chunks):
         self.client_conn.h2.safe_send_body(
@@ -392,6 +399,8 @@ class Http2SingleStreamLayer(_HttpTransmissionLayer, threading.Thread):
             self.client_stream_id,
             chunks
         )
+        if self.zombie:  # pragma: no cover
+            raise Http2ProtocolException("Zombie Stream")
 
     def check_close_connection(self, flow):
         # This layer only handles a single stream.
@@ -404,5 +413,11 @@ class Http2SingleStreamLayer(_HttpTransmissionLayer, threading.Thread):
 
     def run(self):
         layer = HttpLayer(self, self.mode)
-        layer()
+
+        try:
+            layer()
+        except ProtocolException as e:
+            self.log(repr(e), "info")
+            self.log(traceback.format_exc(), "debug")
+
         self.zombie = time.time()
diff --git a/test/mitmproxy/test_protocol_http2.py b/test/mitmproxy/test_protocol_http2.py
index dcb6ff3c..c3950975 100644
--- a/test/mitmproxy/test_protocol_http2.py
+++ b/test/mitmproxy/test_protocol_http2.py
@@ -434,3 +434,51 @@ class TestPushPromise(_Http2TestBase, _Http2ServerBase):
         assert len(bodies) >= 1
         assert b'regular_stream' in bodies
         # the other two bodies might not be transmitted before the reset
+
+@requires_alpn
+class TestConnectionLost(_Http2TestBase, _Http2ServerBase):
+
+    @classmethod
+    def setup_class(self):
+        _Http2TestBase.setup_class()
+        _Http2ServerBase.setup_class()
+
+    @classmethod
+    def teardown_class(self):
+        _Http2TestBase.teardown_class()
+        _Http2ServerBase.teardown_class()
+
+    @classmethod
+    def handle_server_event(self, event, h2_conn, rfile, wfile):
+        if isinstance(event, h2.events.RequestReceived):
+            h2_conn.send_headers(1, [(':status', '200')])
+            wfile.write(h2_conn.data_to_send())
+            wfile.flush()
+            return False
+
+    def test_connection_lost(self):
+        client, h2_conn = self._setup_connection()
+
+        self._send_request(client.wfile, h2_conn, stream_id=1, headers=[
+            (':authority', "127.0.0.1:%s" % self.server.server.address.port),
+            (':method', 'GET'),
+            (':scheme', 'https'),
+            (':path', '/'),
+            ('foo', 'bar')
+        ])
+
+        done = False
+        ended_streams = 0
+        pushed_streams = 0
+        responses = 0
+        while not done:
+            try:
+                raw = b''.join(http2_read_raw_frame(client.rfile))
+                events = h2_conn.receive_data(raw)
+            except:
+                break
+            client.wfile.write(h2_conn.data_to_send())
+            client.wfile.flush()
+
+        if len(self.master.state.flows) == 1:
+            assert self.master.state.flows[0].response is None