Fix: Limit the number in the slab

Author: p00512853

src/proto/streams/counts.rs

@@ -50,6 +50,11 @@ impl Counts {
         self.num_send_streams != 0 || self.num_recv_streams != 0
     }
 
+    /// Returns the maximum number of streams.
+    pub(crate) fn max_streams(&self) -> usize {
+        self.max_recv_streams + self.max_reset_streams
+    }
+
     /// Returns true if the receive stream concurrency can be incremented
     pub fn can_inc_num_recv_streams(&self) -> bool {
         self.max_recv_streams > self.num_recv_streams

src/proto/streams/store.rs

@@ -212,7 +212,7 @@ impl Store {
         self.ids.len()
     }
 
-    #[cfg(feature = "unstable")]
+    /// Returns the number of streams that are held in slab.
     pub fn num_wired_streams(&self) -> usize {
         self.slab.len()
     }

src/proto/streams/streams.rs

@@ -311,7 +311,16 @@ where
 impl<B> DynStreams<'_, B> {
     pub fn recv_headers(&mut self, frame: frame::Headers) -> Result<(), Error> {
         let mut me = self.inner.lock().unwrap();
-
+        if me.store.num_wired_streams() > me.counts.max_streams() {
+            tracing::error!(
+                "HEADERS: number of streams exceeds the upper limit ({})",
+                me.counts.max_streams()
+            );
+            return Err(Error::remote_go_away(
+                Bytes::from("out of stream limit"),
+                Reason::PROTOCOL_ERROR,
+            ));
+        }
         me.recv_headers(self.peer, &self.send_buffer, frame)
     }
 
@@ -348,6 +357,17 @@ impl<B> DynStreams<'_, B> {
 
     pub fn recv_push_promise(&mut self, frame: frame::PushPromise) -> Result<(), Error> {
         let mut me = self.inner.lock().unwrap();
+
+        if me.store.num_wired_streams() > me.counts.max_streams() {
+            tracing::error!(
+                "PUSH_PROMISE: number of streams exceeds the upper limit ({})",
+                me.counts.max_streams()
+            );
+            return Err(Error::remote_go_away(
+                Bytes::from("out of stream limit"),
+                Reason::PROTOCOL_ERROR,
+            ));
+        }
         me.recv_push_promise(&self.send_buffer, frame)
     }