Merge pull request #1206 from Kriechi/improve-http2

improve http2

1 files changed, 43 insertions, 15 deletions

diff --git a/mitmproxy/protocol/http2.py b/mitmproxy/protocol/http2.py
index 39512c8f..9247e657 100644
--- a/mitmproxy/protocol/http2.py
+++ b/mitmproxy/protocol/http2.py
@@ -73,7 +73,7 @@ class SafeH2Connection(connection.H2Connection):
                 frame_chunk = chunk[position:position + max_outbound_frame_size]
                 if self.local_flow_control_window(stream_id) < len(frame_chunk):
                     self.lock.release()
-                    time.sleep(0)
+                    time.sleep(0.1)
                     continue
                 self.send_data(stream_id, frame_chunk)
                 self.conn.send(self.data_to_send())
@@ -165,9 +165,21 @@ class Http2Layer(base.Layer):
             new_settings = dict([(id, cs.new_value) for (id, cs) in six.iteritems(event.changed_settings)])
             other_conn.h2.safe_update_settings(new_settings)
         elif isinstance(event, events.ConnectionTerminated):
-            # Do not immediately terminate the other connection.
-            # Some streams might be still sending data to the client.
-            return False
+            if event.error_code == h2.errors.NO_ERROR:
+                # Do not immediately terminate the other connection.
+                # Some streams might be still sending data to the client.
+                return False
+            else:
+                # Something terrible has happened - kill everything!
+                self.client_conn.h2.close_connection(
+                    error_code=event.error_code,
+                    last_stream_id=event.last_stream_id,
+                    additional_data=event.additional_data
+                )
+                self.client_conn.send(self.client_conn.h2.data_to_send())
+                self._kill_all_streams()
+                return False
+
         elif isinstance(event, events.PushedStreamReceived):
             # pushed stream ids should be unique and not dependent on race conditions
             # only the parent stream id must be looked up first
@@ -177,7 +189,6 @@ class Http2Layer(base.Layer):
                 self.client_conn.send(self.client_conn.h2.data_to_send())
 
             headers = netlib.http.Headers([[str(k), str(v)] for k, v in event.headers])
-            headers['x-mitmproxy-pushed'] = 'true'
             self.streams[event.pushed_stream_id] = Http2SingleStreamLayer(self, event.pushed_stream_id, headers)
             self.streams[event.pushed_stream_id].timestamp_start = time.time()
             self.streams[event.pushed_stream_id].pushed = True
@@ -352,8 +363,22 @@ class Http2SingleStreamLayer(http._HttpTransmissionLayer, threading.Thread):
         raise NotImplementedError()
 
     def send_request(self, message):
+        if not hasattr(self.server_conn, 'h2'):
+            raise exceptions.Http2ProtocolException("Zombie Stream")
+
+        while True:
+            self.server_conn.h2.lock.acquire()
+            max_streams = self.server_conn.h2.remote_settings.max_concurrent_streams
+            if self.server_conn.h2.open_outbound_streams + 1 >= max_streams:
+                # wait until we get a free slot for a new outgoing stream
+                self.server_conn.h2.lock.release()
+                time.sleep(0.1)
+            else:
+                break
+
         if self.pushed:
             # nothing to do here
+            self.server_conn.h2.lock.release()
             return
 
         with self.server_conn.h2.lock:
@@ -364,16 +389,19 @@ class Http2SingleStreamLayer(http._HttpTransmissionLayer, threading.Thread):
             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
 
-            headers = message.headers.copy()
-            headers.insert(0, ":path", message.path)
-            headers.insert(0, ":method", message.method)
-            headers.insert(0, ":scheme", message.scheme)
+        headers = message.headers.copy()
+        headers.insert(0, ":path", message.path)
+        headers.insert(0, ":method", message.method)
+        headers.insert(0, ":scheme", message.scheme)
+        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
+        self.server_conn.h2.safe_send_headers(
+            self.is_zombie,
+            self.server_stream_id,
+            headers,
+        )
+        self.server_conn.h2.lock.release()
 
-            self.server_conn.h2.safe_send_headers(
-                self.is_zombie,
-                self.server_stream_id,
-                headers
-            )
         self.server_conn.h2.safe_send_body(
             self.is_zombie,
             self.server_stream_id,
@@ -408,7 +436,7 @@ class Http2SingleStreamLayer(http._HttpTransmissionLayer, threading.Thread):
             if self.response_data_finished.is_set():
                 while self.response_data_queue.qsize() > 0:
                     yield self.response_data_queue.get()
-                return
+                break
             if self.zombie:  # pragma: no cover
                 raise exceptions.Http2ProtocolException("Zombie Stream")