Forbid too many streams being added to priority tree.

Author: Lukasa

src/priority/__init__.py

@@ -4,5 +4,5 @@
 """
 from .priority import (  # noqa
     Stream, PriorityTree, DeadlockError, PriorityLoop, DuplicateStreamError,
-    MissingStreamError
+    MissingStreamError, TooManyStreamsError
 )

src/priority/priority.py

@@ -43,6 +43,16 @@ class MissingStreamError(KeyError, Exception):
     pass
 
 
+class TooManyStreamsError(Exception):
+    """
+    An attempt was made to insert a dangerous number of streams into the
+    priority tree at the same time.
+
+    .. versionadded:: 1.2.0
+    """
+    pass
+
+
 class Stream(object):
     """
     Priority information for a given stream.
@@ -199,13 +209,33 @@ class PriorityTree(object):
     A HTTP/2 Priority Tree.
 
     This tree stores HTTP/2 streams according to their HTTP/2 priorities.
+
+    .. versionchanged:: 1.2.0
+       Added ``maximum_streams`` keyword argument.
+
+    :param maximum_streams: The maximum number of streams that may be active in
+        the priority tree at any one time. If this number is exceeded, the
+        priority tree will raise a :class:`TooManyStreamsError
+        <priority.TooManyStreamsError>` and will refuse to insert the stream.
+
+        This parameter exists to defend against the possibility of DoS attack
+        by attempting to overfill the priority tree. If any endpoint is
+        attempting to manage the priority of this many streams at once it is
+        probably trying to screw with you, so it is sensible to simply refuse
+        to play ball at that point.
+
+        While we allow the user to configure this, we don't really *expect*
+        them too, unless they want to be even more conservative than we are by
+        default.
+    :type maximum_streams: ``int``
     """
-    def __init__(self):
+    def __init__(self, maximum_streams=1000):
         # This flat array keeps hold of all the streams that are logically
         # dependent on stream 0.
         self._root_stream = Stream(stream_id=0, weight=1)
         self._root_stream.active = False
         self._streams = {0: self._root_stream}
+        self._maximum_streams = maximum_streams
 
     def _exclusive_insert(self, parent_stream, inserted_stream):
         """
@@ -233,6 +263,13 @@ def insert_stream(self,
         if stream_id in self._streams:
             raise DuplicateStreamError("Stream %d already in tree" % stream_id)
 
+        if (len(self._streams) + 1) > self._maximum_streams:
+            raise TooManyStreamsError(
+                "Refusing to insert %d streams into priority tree at once" % (
+                    self._maximum_streams + 1
+                )
+            )
+
         stream = Stream(stream_id, weight)
 
         if exclusive:

test/test_priority.py

@@ -292,6 +292,21 @@ def test_priority_raises_good_errors_for_missing_streams(self):
         with pytest.raises(priority.MissingStreamError):
             p.remove_stream(3)
 
+    @pytest.mark.parametrize('count', range(2, 10000, 100))
+    def test_priority_refuses_to_allow_too_many_streams_in_tree(self, count):
+        """
+        Attempting to insert more streams than maximum_streams into the tree
+        fails.
+        """
+        p = priority.PriorityTree(maximum_streams=count)
+
+        # This isn't an off-by-one error: stream 0 is in the tree by default.
+        for x in range(1, count):
+            p.insert_stream(x)
+
+        with pytest.raises(priority.TooManyStreamsError):
+            p.insert_stream(x + 1)
+
 
 class TestPriorityTreeOutput(object):
     """