http2: more coverage

2 files changed, 231 insertions, 43 deletions

diff --git a/mitmproxy/protocol/http2.py b/mitmproxy/protocol/http2.py
index 4b57174b..27c2a664 100644
--- a/mitmproxy/protocol/http2.py
+++ b/mitmproxy/protocol/http2.py
@@ -78,7 +78,7 @@ class SafeH2Connection(connection.H2Connection):
                 self.send_data(stream_id, frame_chunk)
                 try:
                     self.conn.send(self.data_to_send())
-                except Exception as e:
+                except Exception as e:  # pragma: no cover
                     raise e
                 finally:
                     self.lock.release()
@@ -142,9 +142,9 @@ class Http2Layer(base.Layer):
             self.streams[eid].timestamp_start = time.time()
             self.streams[eid].no_body = (event.stream_ended is not None)
             if event.priority_updated is not None:
-                self.streams[eid].priority_weight = event.priority_updated.weight
-                self.streams[eid].priority_depends_on = event.priority_updated.depends_on
                 self.streams[eid].priority_exclusive = event.priority_updated.exclusive
+                self.streams[eid].priority_depends_on = event.priority_updated.depends_on
+                self.streams[eid].priority_weight = event.priority_updated.weight
                 self.streams[eid].handled_priority_event = event.priority_updated
             self.streams[eid].start()
         elif isinstance(event, events.ResponseReceived):
@@ -155,10 +155,13 @@ class Http2Layer(base.Layer):
             self.streams[eid].response_arrived.set()
         elif isinstance(event, events.DataReceived):
             if self.config.body_size_limit and self.streams[eid].queued_data_length > self.config.body_size_limit:
-                raise netlib.exceptions.HttpException("HTTP body too large. Limit is {}.".format(self.config.body_size_limit))
-            self.streams[eid].data_queue.put(event.data)
-            self.streams[eid].queued_data_length += len(event.data)
-            source_conn.h2.safe_increment_flow_control(event.stream_id, event.flow_controlled_length)
+                self.streams[eid].zombie = time.time()
+                source_conn.h2.safe_reset_stream(event.stream_id, 0x7)
+                self.log("HTTP body too large. Limit is {}.".format(self.config.body_size_limit), "info")
+            else:
+                self.streams[eid].data_queue.put(event.data)
+                self.streams[eid].queued_data_length += len(event.data)
+                source_conn.h2.safe_increment_flow_control(event.stream_id, event.flow_controlled_length)
         elif isinstance(event, events.StreamEnded):
             self.streams[eid].timestamp_end = time.time()
             self.streams[eid].data_finished.set()
@@ -206,6 +209,11 @@ class Http2Layer(base.Layer):
             self.streams[event.pushed_stream_id].request_data_finished.set()
             self.streams[event.pushed_stream_id].start()
         elif isinstance(event, events.PriorityUpdated):
+            if eid in self.streams and self.streams[eid].handled_priority_event is event:
+                # this event was already handled during stream creation
+                # HeadersFrame + Priority information as RequestReceived
+                return True
+
             mapped_stream_id = event.stream_id
             if mapped_stream_id in self.streams and self.streams[mapped_stream_id].server_stream_id:
                 # if the stream is already up and running and was sent to the server
@@ -213,13 +221,9 @@ class Http2Layer(base.Layer):
                 mapped_stream_id = self.streams[mapped_stream_id].server_stream_id
 
             if eid in self.streams:
-                if self.streams[eid].handled_priority_event is event:
-                    # this event was already handled during stream creation
-                    # HeadersFrame + Priority information as RequestReceived
-                    return True
-                self.streams[eid].priority_weight = event.weight
-                self.streams[eid].priority_depends_on = event.depends_on
                 self.streams[eid].priority_exclusive = event.exclusive
+                self.streams[eid].priority_depends_on = event.depends_on
+                self.streams[eid].priority_weight = event.weight
 
             with self.server_conn.h2.lock:
                 self.server_conn.h2.prioritize(
@@ -248,10 +252,12 @@ class Http2Layer(base.Layer):
 
     def _cleanup_streams(self):
         death_time = time.time() - 10
-        for stream_id in self.streams.keys():
-            zombie = self.streams[stream_id].zombie
-            if zombie and zombie <= death_time:
-                self.streams.pop(stream_id, None)
+
+        zombie_streams = [(stream_id, stream) for stream_id, stream in list(self.streams.items()) if stream.zombie]
+        outdated_streams = [stream_id for stream_id, stream in zombie_streams if stream.zombie <= death_time]
+
+        for stream_id in outdated_streams:  # pragma: no cover
+            self.streams.pop(stream_id, None)
 
     def _kill_all_streams(self):
         for stream in self.streams.values():
@@ -296,7 +302,7 @@ class Http2Layer(base.Layer):
                                 return
 
                     self._cleanup_streams()
-        except Exception as e:
+        except Exception as e:  # pragma: no cover
             self.log(repr(e), "info")
             self.log(traceback.format_exc(), "debug")
             self._kill_all_streams()
@@ -326,9 +332,9 @@ class Http2SingleStreamLayer(http._HttpTransmissionLayer, basethread.BaseThread)
 
         self.no_body = False
 
-        self.priority_weight = None
-        self.priority_depends_on = None
         self.priority_exclusive = None
+        self.priority_depends_on = None
+        self.priority_weight = None
         self.handled_priority_event = None
 
     @property
@@ -428,11 +434,11 @@ class Http2SingleStreamLayer(http._HttpTransmissionLayer, basethread.BaseThread)
                 self.server_stream_id,
                 headers,
                 end_stream=self.no_body,
-                priority_weight=self.priority_weight,
-                priority_depends_on=self._map_depends_on_stream_id(self.server_stream_id, self.priority_depends_on),
                 priority_exclusive=self.priority_exclusive,
+                priority_depends_on=self._map_depends_on_stream_id(self.server_stream_id, self.priority_depends_on),
+                priority_weight=self.priority_weight,
             )
-        except Exception as e:
+        except Exception as e:  # pragma: no cover
             raise e
         finally:
             self.server_conn.h2.lock.release()
@@ -523,7 +529,7 @@ class Http2SingleStreamLayer(http._HttpTransmissionLayer, basethread.BaseThread)
 
         try:
             layer()
-        except exceptions.ProtocolException as e:
+        except exceptions.ProtocolException as e:  # pragma: no cover
             self.log(repr(e), "info")
             self.log(traceback.format_exc(), "debug")
 
diff --git a/test/mitmproxy/test_protocol_http2.py b/test/mitmproxy/test_protocol_http2.py
index a4f6b574..cb7cebca 100644
--- a/test/mitmproxy/test_protocol_http2.py
+++ b/test/mitmproxy/test_protocol_http2.py
@@ -60,7 +60,10 @@ class _Http2ServerBase(netlib_tservers.ServerTestBase):
                 except HttpException:
                     print(traceback.format_exc())
                     assert False
+                except netlib.exceptions.TcpDisconnect:
+                    break
                 except:
+                    print(traceback.format_exc())
                     break
                 self.wfile.write(h2_conn.data_to_send())
                 self.wfile.flush()
@@ -70,8 +73,11 @@ class _Http2ServerBase(netlib_tservers.ServerTestBase):
                         if not self.server.handle_server_event(event, h2_conn, self.rfile, self.wfile):
                             done = True
                             break
+                    except netlib.exceptions.TcpDisconnect:
+                        done = True
                     except:
                         done = True
+                        print(traceback.format_exc())
                         break
 
     def handle_server_event(self, h2_conn, rfile, wfile):
@@ -138,11 +144,22 @@ class _Http2TestBase(object):
 
         return client, h2_conn
 
-    def _send_request(self, wfile, h2_conn, stream_id=1, headers=[], body=b''):
+    def _send_request(self,
+                      wfile,
+                      h2_conn,
+                      stream_id=1,
+                      headers=[],
+                      body=b'',
+                      priority_exclusive=None,
+                      priority_depends_on=None,
+                      priority_weight=None):
         h2_conn.send_headers(
             stream_id=stream_id,
             headers=headers,
             end_stream=(len(body) == 0),
+            priority_exclusive=priority_exclusive,
+            priority_depends_on=priority_depends_on,
+            priority_weight=priority_weight,
         )
         if body:
             h2_conn.send_data(stream_id, body)
@@ -166,6 +183,7 @@ class _Http2Test(_Http2TestBase, _Http2ServerBase):
 
 @requires_alpn
 class TestSimple(_Http2Test):
+    request_body_buffer = b''
 
     @classmethod
     def handle_server_event(self, event, h2_conn, rfile, wfile):
@@ -190,13 +208,16 @@ class TestSimple(_Http2Test):
                     ('föo', 'bär'),
                     ('X-Stream-ID', str(event.stream_id)),
                 ])
-            h2_conn.send_data(event.stream_id, b'foobar')
+            h2_conn.send_data(event.stream_id, b'response body')
             h2_conn.end_stream(event.stream_id)
             wfile.write(h2_conn.data_to_send())
             wfile.flush()
+        elif isinstance(event, h2.events.DataReceived):
+            self.request_body_buffer += event.data
         return True
 
     def test_simple(self):
+        response_body_buffer = b''
         client, h2_conn = self._setup_connection()
 
         self._send_request(
@@ -210,7 +231,7 @@ class TestSimple(_Http2Test):
                 ('ClIeNt-FoO', 'client-bar-1'),
                 ('ClIeNt-FoO', 'client-bar-2'),
             ],
-            body=b'my request body echoed back to me')
+            body=b'request body')
 
         done = False
         while not done:
@@ -225,7 +246,9 @@ class TestSimple(_Http2Test):
             client.wfile.flush()
 
             for event in events:
-                if isinstance(event, h2.events.StreamEnded):
+                if isinstance(event, h2.events.DataReceived):
+                    response_body_buffer += event.data
+                elif isinstance(event, h2.events.StreamEnded):
                     done = True
 
         h2_conn.close_connection()
@@ -236,31 +259,41 @@ class TestSimple(_Http2Test):
         assert self.master.state.flows[0].response.status_code == 200
         assert self.master.state.flows[0].response.headers['server-foo'] == 'server-bar'
         assert self.master.state.flows[0].response.headers['föo'] == 'bär'
-        assert self.master.state.flows[0].response.body == b'foobar'
+        assert self.master.state.flows[0].response.body == b'response body'
+        assert self.request_body_buffer == b'request body'
+        assert response_body_buffer == b'response body'
 
 
 @requires_alpn
-class TestWithBodies(_Http2Test):
-    tmp_data_buffer_foobar = b''
+class TestRequestWithPriority(_Http2Test):
 
     @classmethod
     def handle_server_event(self, event, h2_conn, rfile, wfile):
         if isinstance(event, h2.events.ConnectionTerminated):
             return False
-        if isinstance(event, h2.events.DataReceived):
-            self.tmp_data_buffer_foobar += event.data
-        elif isinstance(event, h2.events.StreamEnded):
-            h2_conn.send_headers(1, [
-                (':status', '200'),
-            ])
-            h2_conn.send_data(1, self.tmp_data_buffer_foobar)
-            h2_conn.end_stream(1)
+        elif isinstance(event, h2.events.RequestReceived):
+            import warnings
+            with warnings.catch_warnings():
+                # Ignore UnicodeWarning:
+                # h2/utilities.py:64: UnicodeWarning: Unicode equal comparison
+                # failed to convert both arguments to Unicode - interpreting
+                # them as being unequal.
+                #     elif header[0] in (b'cookie', u'cookie') and len(header[1]) < 20:
+
+                warnings.simplefilter("ignore")
+
+                headers = [(':status', '200')]
+                if event.priority_updated:
+                    headers.append(('priority_exclusive', event.priority_updated.exclusive))
+                    headers.append(('priority_depends_on', event.priority_updated.depends_on))
+                    headers.append(('priority_weight', event.priority_updated.weight))
+                h2_conn.send_headers(event.stream_id, headers)
+            h2_conn.end_stream(event.stream_id)
             wfile.write(h2_conn.data_to_send())
             wfile.flush()
-
         return True
 
-    def test_with_bodies(self):
+    def test_request_with_priority(self):
         client, h2_conn = self._setup_connection()
 
         self._send_request(
@@ -272,7 +305,9 @@ class TestWithBodies(_Http2Test):
                 (':scheme', 'https'),
                 (':path', '/'),
             ],
-            body=b'foobar with request body',
+            priority_exclusive = True,
+            priority_depends_on = 42424242,
+            priority_weight = 42,
         )
 
         done = False
@@ -295,7 +330,149 @@ class TestWithBodies(_Http2Test):
         client.wfile.write(h2_conn.data_to_send())
         client.wfile.flush()
 
-        assert self.master.state.flows[0].response.body == b'foobar with request body'
+        assert len(self.master.state.flows) == 1
+        assert self.master.state.flows[0].response.headers['priority_exclusive'] == 'True'
+        assert self.master.state.flows[0].response.headers['priority_depends_on'] == '42424242'
+        assert self.master.state.flows[0].response.headers['priority_weight'] == '42'
+
+    def test_request_without_priority(self):
+        client, h2_conn = self._setup_connection()
+
+        self._send_request(
+            client.wfile,
+            h2_conn,
+            headers=[
+                (':authority', "127.0.0.1:%s" % self.server.server.address.port),
+                (':method', 'GET'),
+                (':scheme', 'https'),
+                (':path', '/'),
+            ],
+        )
+
+        done = False
+        while not done:
+            try:
+                raw = b''.join(framereader.http2_read_raw_frame(client.rfile))
+                events = h2_conn.receive_data(raw)
+            except HttpException:
+                print(traceback.format_exc())
+                assert False
+
+            client.wfile.write(h2_conn.data_to_send())
+            client.wfile.flush()
+
+            for event in events:
+                if isinstance(event, h2.events.StreamEnded):
+                    done = True
+
+        h2_conn.close_connection()
+        client.wfile.write(h2_conn.data_to_send())
+        client.wfile.flush()
+
+        assert len(self.master.state.flows) == 1
+        assert 'priority_exclusive' not in self.master.state.flows[0].response.headers
+        assert 'priority_depends_on' not in self.master.state.flows[0].response.headers
+        assert 'priority_weight' not in self.master.state.flows[0].response.headers
+
+
+@requires_alpn
+class TestStreamResetFromServer(_Http2Test):
+
+    @classmethod
+    def handle_server_event(self, event, h2_conn, rfile, wfile):
+        if isinstance(event, h2.events.ConnectionTerminated):
+            return False
+        elif isinstance(event, h2.events.RequestReceived):
+            h2_conn.reset_stream(event.stream_id, 0x8)
+            wfile.write(h2_conn.data_to_send())
+            wfile.flush()
+        return True
+
+    def test_request_with_priority(self):
+        client, h2_conn = self._setup_connection()
+
+        self._send_request(
+            client.wfile,
+            h2_conn,
+            headers=[
+                (':authority', "127.0.0.1:%s" % self.server.server.address.port),
+                (':method', 'GET'),
+                (':scheme', 'https'),
+                (':path', '/'),
+            ],
+        )
+
+        done = False
+        while not done:
+            try:
+                raw = b''.join(framereader.http2_read_raw_frame(client.rfile))
+                events = h2_conn.receive_data(raw)
+            except HttpException:
+                print(traceback.format_exc())
+                assert False
+
+            client.wfile.write(h2_conn.data_to_send())
+            client.wfile.flush()
+
+            for event in events:
+                if isinstance(event, h2.events.StreamReset):
+                    done = True
+
+        h2_conn.close_connection()
+        client.wfile.write(h2_conn.data_to_send())
+        client.wfile.flush()
+
+        assert len(self.master.state.flows) == 1
+        assert self.master.state.flows[0].response is None
+
+
+@requires_alpn
+class TestBodySizeLimit(_Http2Test):
+
+    @classmethod
+    def handle_server_event(self, event, h2_conn, rfile, wfile):
+        if isinstance(event, h2.events.ConnectionTerminated):
+            return False
+        return True
+
+    def test_body_size_limit(self):
+        self.config.body_size_limit = 20
+
+        client, h2_conn = self._setup_connection()
+
+        self._send_request(
+            client.wfile,
+            h2_conn,
+            headers=[
+                (':authority', "127.0.0.1:%s" % self.server.server.address.port),
+                (':method', 'GET'),
+                (':scheme', 'https'),
+                (':path', '/'),
+            ],
+            body=b'very long body over 20 characters long',
+        )
+
+        done = False
+        while not done:
+            try:
+                raw = b''.join(framereader.http2_read_raw_frame(client.rfile))
+                events = h2_conn.receive_data(raw)
+            except HttpException:
+                print(traceback.format_exc())
+                assert False
+
+            client.wfile.write(h2_conn.data_to_send())
+            client.wfile.flush()
+
+            for event in events:
+                if isinstance(event, h2.events.StreamReset):
+                    done = True
+
+        h2_conn.close_connection()
+        client.wfile.write(h2_conn.data_to_send())
+        client.wfile.flush()
+
+        assert len(self.master.state.flows) == 0
 
 
 @requires_alpn
@@ -497,6 +674,11 @@ class TestConnectionLost(_Http2Test):
 class TestMaxConcurrentStreams(_Http2Test):
 
     @classmethod
+    def setup_class(self):
+        _Http2TestBase.setup_class()
+        _Http2ServerBase.setup_class(h2_server_settings={h2.settings.MAX_CONCURRENT_STREAMS: 2})
+
+    @classmethod
     def handle_server_event(self, event, h2_conn, rfile, wfile):
         if isinstance(event, h2.events.ConnectionTerminated):
             return False