http2: more coverage

1 files changed, 201 insertions, 19 deletions

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