refactor: change local reset counter to use type system more

Author: seanmonstar

src/proto/connection.rs

@@ -428,24 +428,7 @@ where
             // error. This is handled by setting a GOAWAY frame followed by
             // terminating the connection.
             Err(Error::GoAway(debug_data, reason, initiator)) => {
-                let e = Error::GoAway(debug_data.clone(), reason, initiator);
-                tracing::debug!(error = ?e, "Connection::poll; connection error");
-
-                // We may have already sent a GOAWAY for this error,
-                // if so, don't send another, just flush and close up.
-                if self
-                    .go_away
-                    .going_away()
-                    .map_or(false, |frame| frame.reason() == reason)
-                {
-                    tracing::trace!("    -> already going away");
-                    *self.state = State::Closing(reason, initiator);
-                    return Ok(());
-                }
-
-                // Reset all active streams
-                self.streams.handle_error(e);
-                self.go_away_now_data(reason, debug_data);
+                self.handle_go_away(reason, debug_data, initiator);
                 Ok(())
             }
             // Attempting to read a frame resulted in a stream level error.
@@ -454,7 +437,12 @@ where
             Err(Error::Reset(id, reason, initiator)) => {
                 debug_assert_eq!(initiator, Initiator::Library);
                 tracing::trace!(?id, ?reason, "stream error");
-                self.streams.send_reset(id, reason);
+                match self.streams.send_reset(id, reason) {
+                    Ok(()) => (),
+                    Err(crate::proto::error::GoAway { debug_data, reason }) => {
+                        self.handle_go_away(reason, debug_data, Initiator::Library);
+                    }
+                }
                 Ok(())
             }
             // Attempting to read a frame resulted in an I/O error. All
@@ -474,6 +462,27 @@ where
         }
     }
 
+    fn handle_go_away(&mut self, reason: Reason, debug_data: Bytes, initiator: Initiator) {
+        let e = Error::GoAway(debug_data.clone(), reason, initiator);
+        tracing::debug!(error = ?e, "Connection::poll; connection error");
+
+        // We may have already sent a GOAWAY for this error,
+        // if so, don't send another, just flush and close up.
+        if self
+            .go_away
+            .going_away()
+            .map_or(false, |frame| frame.reason() == reason)
+        {
+            tracing::trace!("    -> already going away");
+            *self.state = State::Closing(reason, initiator);
+            return;
+        }
+
+        // Reset all active streams
+        self.streams.handle_error(e);
+        self.go_away_now_data(reason, debug_data);
+    }
+
     fn recv_frame(&mut self, frame: Option<Frame>) -> Result<ReceivedFrame, Error> {
         use crate::frame::Frame::*;
         match frame {

src/proto/error.rs

@@ -13,6 +13,11 @@ pub enum Error {
     Io(io::ErrorKind, Option<String>),
 }
 
+pub struct GoAway {
+    pub debug_data: Bytes,
+    pub reason: Reason,
+}
+
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum Initiator {
     User,
@@ -60,6 +65,10 @@ impl Initiator {
             Self::Remote => false,
         }
     }
+
+    pub(crate) fn is_library(&self) -> bool {
+        matches!(self, Self::Library)
+    }
 }
 
 impl fmt::Display for Error {

src/proto/streams/streams.rs

@@ -366,7 +366,11 @@ impl<B> DynStreams<'_, B> {
         me.recv_eof(self.send_buffer, clear_pending_accept)
     }
 
-    pub fn send_reset(&mut self, id: StreamId, reason: Reason) {
+    pub fn send_reset(
+        &mut self,
+        id: StreamId,
+        reason: Reason,
+    ) -> Result<(), crate::proto::error::GoAway> {
         let mut me = self.inner.lock().unwrap();
         me.send_reset(self.send_buffer, id, reason)
     }
@@ -637,15 +641,23 @@ impl Inner {
             // The remote may send window updates for streams that the local now
             // considers closed. It's ok...
             if let Some(mut stream) = self.store.find_mut(&id) {
-                // This result is ignored as there is nothing to do when there
-                // is an error. The stream is reset by the function on error and
-                // the error is informational.
-                let _ = self.actions.send.recv_stream_window_update(
-                    frame.size_increment(),
+                let res = self
+                    .actions
+                    .send
+                    .recv_stream_window_update(
+                        frame.size_increment(),
+                        send_buffer,
+                        &mut stream,
+                        &mut self.counts,
+                        &mut self.actions.task,
+                    )
+                    .map_err(|reason| Error::library_reset(id, reason));
+
+                return self.actions.reset_on_recv_stream_err(
                     send_buffer,
                     &mut stream,
                     &mut self.counts,
-                    &mut self.actions.task,
+                    res,
                 );
             } else {
                 self.actions
@@ -882,7 +894,12 @@ impl Inner {
         Poll::Ready(Ok(()))
     }
 
-    fn send_reset<B>(&mut self, send_buffer: &SendBuffer<B>, id: StreamId, reason: Reason) {
+    fn send_reset<B>(
+        &mut self,
+        send_buffer: &SendBuffer<B>,
+        id: StreamId,
+        reason: Reason,
+    ) -> Result<(), crate::proto::error::GoAway> {
         let key = match self.store.find_entry(id) {
             Entry::Occupied(e) => e.key(),
             Entry::Vacant(e) => {
@@ -923,7 +940,7 @@ impl Inner {
             Initiator::Library,
             &mut self.counts,
             send_buffer,
-        );
+        )
     }
 }
 
@@ -1095,8 +1112,20 @@ impl<B> StreamRef<B> {
         let mut send_buffer = self.send_buffer.inner.lock().unwrap();
         let send_buffer = &mut *send_buffer;
 
-        me.actions
-            .send_reset(stream, reason, Initiator::User, &mut me.counts, send_buffer);
+        match me
+            .actions
+            .send_reset(stream, reason, Initiator::User, &mut me.counts, send_buffer)
+        {
+            Ok(()) => (),
+            Err(crate::proto::error::GoAway { .. }) => {
+                // this should never happen, because Initiator::User resets do
+                // not count toward the local limit.
+                // we could perhaps make this state impossible, if we made the
+                // initiator argument a generic, and so this could return
+                // Infallible instead of an impossible GoAway, but oh well.
+                unreachable!("Initiator::User should not error sending reset");
+            }
+        }
     }
 
     pub fn send_response(
@@ -1517,8 +1546,23 @@ impl Actions {
         initiator: Initiator,
         counts: &mut Counts,
         send_buffer: &mut Buffer<Frame<B>>,
-    ) {
+    ) -> Result<(), crate::proto::error::GoAway> {
         counts.transition(stream, |counts, stream| {
+            if initiator.is_library() {
+                if counts.can_inc_num_local_error_resets() {
+                    counts.inc_num_local_error_resets();
+                } else {
+                    tracing::warn!(
+                        "locally-reset streams reached limit ({:?})",
+                        counts.max_local_error_resets().unwrap(),
+                    );
+                    return Err(crate::proto::error::GoAway {
+                        reason: Reason::ENHANCE_YOUR_CALM,
+                        debug_data: "too_many_internal_resets".into(),
+                    });
+                }
+            }
+
             self.send.send_reset(
                 reason,
                 initiator,
@@ -1530,7 +1574,9 @@ impl Actions {
             self.recv.enqueue_reset_expiration(stream, counts);
             // if a RecvStream is parked, ensure it's notified
             stream.notify_recv();
-        });
+
+            Ok(())
+        })
     }
 
     fn reset_on_recv_stream_err<B>(