send reset signal when read Error occurs

Author: DougAnderson444

misc/webrtc-utils/src/stream.rs

@@ -30,7 +30,7 @@ use std::{
 
 use crate::proto::{Flag, Message};
 use crate::{
-    stream::drop_listener::GracefullyClosed,
+    stream::drop_listener::DropMessage,
     stream::framed_dc::FramedDc,
     stream::state::{Closing, State},
 };
@@ -62,7 +62,7 @@ pub struct Stream<T> {
     state: State,
     read_buffer: Bytes,
     /// Dropping this will close the oneshot and notify the receiver by emitting `Canceled`.
-    drop_notifier: Option<oneshot::Sender<GracefullyClosed>>,
+    drop_notifier: Option<oneshot::Sender<DropMessage>>,
 }
 
 impl<T> Stream<T>
@@ -139,21 +139,29 @@ where
                 ..
             } = &mut *self;
 
-            match ready!(io_poll_next(io, cx))? {
-                Some((flag, message)) => {
+            match ready!(io_poll_next(io, cx)) {
+                Ok(Some((flag, message))) => {
                     if let Some(flag) = flag {
                         state.handle_inbound_flag(flag, read_buffer);
                     }
 
-                    debug_assert!(read_buffer.is_empty());
                     if let Some(message) = message {
                         *read_buffer = message.into();
                     }
                 }
-                None => {
+                Ok(None) => {
                     state.handle_inbound_flag(Flag::FIN, read_buffer);
                     return Poll::Ready(Ok(0));
                 }
+                Err(e) => {
+                    log::error!("Error while reading next message: {:?}", e);
+                    let _ = self
+                        .drop_notifier
+                        .take()
+                        .expect("should be able to take drop_notifier value")
+                        .send(DropMessage::SendReset);
+                    return Poll::Ready(Err(e));
+                }
             }
         }
     }
@@ -231,7 +239,7 @@ where
                         .drop_notifier
                         .take()
                         .expect("to not close twice")
-                        .send(GracefullyClosed {});
+                        .send(DropMessage::GracefullyClosed);
 
                     return Poll::Ready(Ok(()));
                 }

misc/webrtc-utils/src/stream/drop_listener.rs

@@ -36,7 +36,7 @@ pub struct DropListener<T> {
 }
 
 impl<T> DropListener<T> {
-    pub fn new(stream: FramedDc<T>, receiver: oneshot::Receiver<GracefullyClosed>) -> Self {
+    pub fn new(stream: FramedDc<T>, receiver: oneshot::Receiver<DropMessage>) -> Self {
         Self {
             state: State::Idle { stream, receiver },
         }
@@ -47,7 +47,7 @@ enum State<T> {
     /// The [`DropListener`] is idle and waiting to be activated.
     Idle {
         stream: FramedDc<T>,
-        receiver: oneshot::Receiver<GracefullyClosed>,
+        receiver: oneshot::Receiver<DropMessage>,
     },
     /// The stream got dropped and we are sending a reset flag.
     SendingReset {
@@ -75,9 +75,14 @@ where
                     stream,
                     mut receiver,
                 } => match receiver.poll_unpin(cx) {
-                    Poll::Ready(Ok(GracefullyClosed {})) => {
+                    Poll::Ready(Ok(DropMessage::GracefullyClosed)) => {
                         return Poll::Ready(Ok(()));
                     }
+                    Poll::Ready(Ok(DropMessage::SendReset)) => {
+                        log::info!("Stream gracefully errored, sending Reset");
+                        *state = State::SendingReset { stream };
+                        continue;
+                    }
                     Poll::Ready(Err(Canceled)) => {
                         log::info!("Stream dropped without graceful close, sending Reset");
                         *state = State::SendingReset { stream };
@@ -117,5 +122,10 @@ where
     }
 }
 
-/// Indicates that our stream got gracefully closed.
-pub struct GracefullyClosed {}
+/// The reason we are dropping the Stream.
+pub enum DropMessage {
+    /// The stream was closed gracefully.
+    GracefullyClosed,
+    /// The stream errored (such as receiving too much data) and we are sending a reset signal
+    SendReset,
+}