From 0a222ea9aca4d34603f725ca0f3ff422df2196b5 Mon Sep 17 00:00:00 2001
From: Dirkjan Ochtman <dirkjan@ochtman.nl>
Date: Wed, 20 Sep 2023 16:32:30 +0200
Subject: [PATCH] More principled error handling for invalid frames

---
 quinn-proto/src/connection/mod.rs   | 26 ++++++++++++-------------
 quinn-proto/src/connection/stats.rs |  1 -
 quinn-proto/src/frame.rs            | 30 +++++++++++++++++++++--------
 3 files changed, 35 insertions(+), 22 deletions(-)

diff --git a/quinn-proto/src/connection/mod.rs b/quinn-proto/src/connection/mod.rs
index 0603bc976..1c695f698 100644
--- a/quinn-proto/src/connection/mod.rs
+++ b/quinn-proto/src/connection/mod.rs
@@ -2050,7 +2050,15 @@ impl Connection {
                 return Ok(());
             }
             State::Closed(_) => {
-                for frame in frame::Iter::new(packet.payload.freeze()) {
+                for result in frame::Iter::new(packet.payload.freeze()) {
+                    let frame = match result {
+                        Ok(frame) => frame,
+                        Err(err) => {
+                            debug!("frame decoding error: {err:?}");
+                            continue;
+                        }
+                    };
+
                     if let Frame::Padding = frame {
                         continue;
                     };
@@ -2291,7 +2299,8 @@ impl Connection {
         debug_assert_ne!(packet.header.space(), SpaceId::Data);
         let payload_len = packet.payload.len();
         let mut ack_eliciting = false;
-        for frame in frame::Iter::new(packet.payload.freeze()) {
+        for result in frame::Iter::new(packet.payload.freeze()) {
+            let frame = result?;
             let span = match frame {
                 Frame::Padding => continue,
                 _ => Some(trace_span!("frame", ty = %frame.ty())),
@@ -2316,11 +2325,6 @@ impl Connection {
                     self.state = State::Draining;
                     return Ok(());
                 }
-                Frame::Invalid { ty, reason } => {
-                    let mut err = TransportError::FRAME_ENCODING_ERROR(reason);
-                    err.frame = Some(ty);
-                    return Err(err);
-                }
                 _ => {
                     let mut err =
                         TransportError::PROTOCOL_VIOLATION("illegal frame type in handshake");
@@ -2349,7 +2353,8 @@ impl Connection {
         let mut close = None;
         let payload_len = payload.len();
         let mut ack_eliciting = false;
-        for frame in frame::Iter::new(payload) {
+        for result in frame::Iter::new(payload) {
+            let frame = result?;
             let span = match frame {
                 Frame::Padding => continue,
                 _ => Some(trace_span!("frame", ty = %frame.ty())),
@@ -2397,11 +2402,6 @@ impl Connection {
                 }
             }
             match frame {
-                Frame::Invalid { ty, reason } => {
-                    let mut err = TransportError::FRAME_ENCODING_ERROR(reason);
-                    err.frame = Some(ty);
-                    return Err(err);
-                }
                 Frame::Crypto(frame) => {
                     self.read_crypto(SpaceId::Data, &frame, payload_len)?;
                 }
diff --git a/quinn-proto/src/connection/stats.rs b/quinn-proto/src/connection/stats.rs
index dfae4e751..9270bef0e 100644
--- a/quinn-proto/src/connection/stats.rs
+++ b/quinn-proto/src/connection/stats.rs
@@ -82,7 +82,6 @@ impl FrameStats {
             Frame::PathResponse(_) => self.path_response += 1,
             Frame::Close(_) => self.connection_close += 1,
             Frame::HandshakeDone => self.handshake_done += 1,
-            Frame::Invalid { .. } => {}
         }
     }
 }
diff --git a/quinn-proto/src/frame.rs b/quinn-proto/src/frame.rs
index 6b7381b7d..209471cf3 100644
--- a/quinn-proto/src/frame.rs
+++ b/quinn-proto/src/frame.rs
@@ -157,7 +157,6 @@ pub enum Frame {
     PathResponse(u64),
     Close(Close),
     Datagram(Datagram),
-    Invalid { ty: Type, reason: &'static str },
     HandshakeDone,
 }
 
@@ -197,7 +196,6 @@ impl Frame {
             Crypto(_) => Type::CRYPTO,
             NewToken { .. } => Type::NEW_TOKEN,
             Datagram(_) => Type(*DATAGRAM_TYS.start()),
-            Invalid { ty, .. } => ty,
             HandshakeDone => Type::HANDSHAKE_DONE,
         }
     }
@@ -721,25 +719,39 @@ impl Iter {
 }
 
 impl Iterator for Iter {
-    type Item = Frame;
+    type Item = Result<Frame, InvalidFrame>;
     fn next(&mut self) -> Option<Self::Item> {
         if !self.bytes.has_remaining() {
             return None;
         }
         match self.try_next() {
-            Ok(x) => Some(x),
+            Ok(x) => Some(Ok(x)),
             Err(e) => {
                 // Corrupt frame, skip it and everything that follows
                 self.bytes = io::Cursor::new(Bytes::new());
-                Some(Frame::Invalid {
-                    ty: self.last_ty.unwrap(),
+                Some(Err(InvalidFrame {
+                    ty: self.last_ty,
                     reason: e.reason(),
-                })
+                }))
             }
         }
     }
 }
 
+#[derive(Debug)]
+pub(crate) struct InvalidFrame {
+    pub(crate) ty: Option<Type>,
+    pub(crate) reason: &'static str,
+}
+
+impl From<InvalidFrame> for TransportError {
+    fn from(err: InvalidFrame) -> Self {
+        let mut te = Self::FRAME_ENCODING_ERROR(err.reason);
+        te.frame = err.ty;
+        te
+    }
+}
+
 fn scan_ack_blocks(buf: &mut io::Cursor<Bytes>, largest: u64, n: usize) -> Result<(), IterErr> {
     let first_block = buf.get_var()?;
     let mut smallest = largest.checked_sub(first_block).ok_or(IterErr::Malformed)?;
@@ -890,7 +902,9 @@ mod test {
             ce: 12,
         };
         Ack::encode(42, &ranges, Some(&ECN), &mut buf);
-        let frames = Iter::new(Bytes::from(buf)).collect::<Vec<_>>();
+        let frames = Iter::new(Bytes::from(buf))
+            .collect::<Result<Vec<_>, _>>()
+            .unwrap();
         assert_eq!(frames.len(), 1);
         match frames[0] {
             Frame::Ack(ref ack) => {