Limit amount of _closed_streams

Author: hellysmile

h2/connection.py

@@ -33,7 +33,7 @@
 from .frame_buffer import FrameBuffer
 from .settings import Settings, SettingCodes
 from .stream import H2Stream, StreamClosedBy
-from .utilities import guard_increment_window
+from .utilities import SizeLimitDict, guard_increment_window
 from .windows import WindowManager
 
 
@@ -281,6 +281,9 @@ class H2Connection(object):
     # The initial default value of SETTINGS_MAX_HEADER_LIST_SIZE.
     DEFAULT_MAX_HEADER_LIST_SIZE = 2**16
 
+    # Keep in memory limited amount of results for streams closes
+    MAX_CLOSED_STREAMS = 2**31
+
     def __init__(self, config=None):
         self.state_machine = H2ConnectionStateMachine()
         self.streams = {}
@@ -355,7 +358,9 @@ def __init__(self, config=None):
         # Also used to determine whether we should consider a frame received
         # while a stream is closed as either a stream error or a connection
         # error.
-        self._closed_streams = {}
+        self._closed_streams = SizeLimitDict(
+            size_limit=self.MAX_CLOSED_STREAMS
+        )
 
         # The flow control window manager for the connection.
         self._inbound_flow_control_window_manager = WindowManager(

h2/utilities.py

@@ -617,3 +617,22 @@ def validate_outbound_headers(headers, hdr_validation_flags):
     headers = _check_path_header(headers, hdr_validation_flags)
 
     return headers
+
+
+class SizeLimitDict(collections.OrderedDict):
+
+    def __init__(self, *args, **kwargs):
+        self._size_limit = kwargs.pop("size_limit", None)
+        super(SizeLimitDict, self).__init__(*args, **kwargs)
+
+        self._check_size_limit()
+
+    def __setitem__(self, key, value):
+        super(SizeLimitDict, self).__setitem__(key, value)
+
+        self._check_size_limit()
+
+    def _check_size_limit(self):
+        if self._size_limit is not None:
+            while len(self) > self._size_limit:
+                self.popitem(last=False)