improved zombie detection

libmproxy/protocol/http.py

@@ -157,7 +157,7 @@ class SafeH2Connection(H2Connection):
         with self.lock:
             try:
                 self.reset_stream(stream_id, error_code)
-            except StreamClosedError:
+            except h2.exceptions.ProtocolError:
                 # stream is already closed - good
                 pass
             self.conn.send(self.data_to_send())
@@ -172,30 +172,33 @@ class SafeH2Connection(H2Connection):
             self.update_settings(new_settings)
             self.conn.send(self.data_to_send())
 
-    def safe_send_headers(self, stream_id, headers):
+    def safe_send_headers(self, is_zombie, stream_id, headers):
         with self.lock:
+            if is_zombie(self, stream_id):
+                return
             self.send_headers(stream_id, headers)
             self.conn.send(self.data_to_send())
 
-    def safe_send_body(self, stream_id, chunks):
+    def safe_send_body(self, is_zombie, stream_id, chunks):
-        # TODO: this assumes the MAX_FRAME_SIZE does not change in the middle
-        # of a transfer - it could though. Then we need to re-chunk everything.
         for chunk in chunks:
-            max_outbound_frame_size = self.max_outbound_frame_size
+            position = 0
-            for i in xrange(0, len(chunk), max_outbound_frame_size):
+            while position < len(chunk):
-                frame_chunk = chunk[i:i+max_outbound_frame_size]
-
                 self.lock.acquire()
-                while True:
+                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):
                     self.lock.release()
                     time.sleep(0)
-                    else:
+                    continue
-                        break
+                if is_zombie(self, stream_id):
+                    return
                 self.send_data(stream_id, frame_chunk)
                 self.conn.send(self.data_to_send())
                 self.lock.release()
+                position += max_outbound_frame_size
         with self.lock:
+            if is_zombie(self, stream_id):
+                return
             self.end_stream(stream_id)
             self.conn.send(self.data_to_send())
 
@@ -314,6 +317,18 @@ class Http2SingleStreamLayer(_HttpLayer, threading.Thread):
         self.response_arrived = threading.Event()
         self.data_finished = threading.Event()
 
+    def is_zombie(self, h2_conn, stream_id):
+        if self.zombie:
+            return True
+
+        try:
+            h2_conn._get_stream_by_id(stream_id)
+        except Exception as e:
+            if isinstance(e, h2.exceptions.StreamClosedError):
+                return true
+
+        return False
+
     def read_request(self):
         self.data_finished.wait()
         self.data_finished.clear()
@@ -364,18 +379,17 @@ class Http2SingleStreamLayer(_HttpLayer, threading.Thread):
         )
 
     def send_request(self, message):
-        if self.zombie:
-            return
-
         with self.server_conn.h2.lock:
             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,
                 message.headers
             )
         self.server_conn.h2.safe_send_body(
+            self.is_zombie,
             self.server_stream_id,
             message.body
         )
@@ -409,19 +423,15 @@ class Http2SingleStreamLayer(_HttpLayer, threading.Thread):
                 return
 
     def send_response_headers(self, response):
-        if self.zombie:
-            return
-
         self.client_conn.h2.safe_send_headers(
+            self.is_zombie,
             self.client_stream_id,
             response.headers
         )
 
     def send_response_body(self, _response, chunks):
-        if self.zombie:
-            return
-
         self.client_conn.h2.safe_send_body(
+            self.is_zombie,
             self.client_stream_id,
             chunks
         )