Add a sepparate InternalError for every step of encoding a packet.

neqo-transport/src/addr_valid.rs

@@ -16,7 +16,7 @@ use crate::cid::ConnectionId;
 use crate::packet::PacketBuilder;
 use crate::recovery::RecoveryToken;
 use crate::stats::FrameStats;
-use crate::Res;
+use crate::{Error, Res};
 
 use smallvec::SmallVec;
 use std::convert::TryFrom;
@@ -355,10 +355,11 @@ impl NewTokenState {
         builder: &mut PacketBuilder,
         tokens: &mut Vec<RecoveryToken>,
         stats: &mut FrameStats,
-    ) {
+    ) -> Res<()> {
         if let Self::Server(ref mut sender) = self {
-            sender.write_frames(builder, tokens, stats);
+            sender.write_frames(builder, tokens, stats)?;
         }
+        Ok(())
     }
 
     /// If this a server, buffer a NEW_TOKEN for sending.
@@ -429,18 +430,22 @@ impl NewTokenSender {
         builder: &mut PacketBuilder,
         tokens: &mut Vec<RecoveryToken>,
         stats: &mut FrameStats,
-    ) {
+    ) -> Res<()> {
         for t in self.tokens.iter_mut() {
             if t.needs_sending && t.len() <= builder.remaining() {
                 t.needs_sending = false;
 
                 builder.encode_varint(crate::frame::FRAME_TYPE_NEW_TOKEN);
                 builder.encode_vvec(&t.token);
+                if builder.len() > builder.limit() {
+                    return Err(Error::InternalError(7));
+                }
 
                 tokens.push(RecoveryToken::NewToken(t.seqno));
                 stats.new_token += 1;
             }
         }
+        Ok(())
     }
 
     pub fn lost(&mut self, seqno: usize) {

neqo-transport/src/cid.rs

@@ -497,35 +497,38 @@ impl ConnectionIdManager {
         entry: &ConnectionIdEntry<[u8; 16]>,
         builder: &mut PacketBuilder,
         stats: &mut FrameStats,
-    ) -> bool {
+    ) -> Res<bool> {
         let len = 1 + Encoder::varint_len(entry.seqno) + 1 + 1 + entry.cid.len() + 16;
         if builder.remaining() < len {
-            return false;
+            return Ok(false);
         }
 
         builder.encode_varint(FRAME_TYPE_NEW_CONNECTION_ID);
         builder.encode_varint(entry.seqno);
         builder.encode_varint(0u64);
         builder.encode_vec(1, &entry.cid);
         builder.encode(&entry.srt);
+        if builder.len() > builder.limit() {
+            return Err(Error::InternalError(8));
+        }
 
         stats.new_connection_id += 1;
-        true
+        Ok(true)
     }
 
     pub fn write_frames(
         &mut self,
         builder: &mut PacketBuilder,
         tokens: &mut Vec<RecoveryToken>,
         stats: &mut FrameStats,
-    ) {
+    ) -> Res<()> {
         if self.generator.deref().borrow().generates_empty_cids() {
             debug_assert_eq!(self.generator.borrow_mut().generate_cid().unwrap().len(), 0);
-            return;
+            return Ok(());
         }
 
         while let Some(entry) = self.lost_new_connection_id.pop() {
-            if self.write_entry(&entry, builder, stats) {
+            if self.write_entry(&entry, builder, stats)? {
                 tokens.push(RecoveryToken::NewConnectionId(entry));
             } else {
                 // This shouldn't happen often.
@@ -550,10 +553,11 @@ impl ConnectionIdManager {
                     .add_local(ConnectionIdEntry::new(seqno, cid.clone(), ()));
 
                 let entry = ConnectionIdEntry::new(seqno, cid, srt);
-                debug_assert!(self.write_entry(&entry, builder, stats));
+                debug_assert!(self.write_entry(&entry, builder, stats)?);
                 tokens.push(RecoveryToken::NewConnectionId(entry));
             }
         }
+        Ok(())
     }
 
     pub fn lost(&mut self, entry: &ConnectionIdEntry<[u8; 16]>) {

neqo-transport/src/connection/mod.rs

@@ -1633,7 +1633,7 @@ impl Connection {
         address_validation: &AddressValidationInfo,
         quic_version: QuicVersion,
         grease_quic_bit: bool,
-    ) -> (PacketType, PacketBuilder) {
+    ) -> Res<(PacketType, PacketBuilder)> {
         let pt = PacketType::from(cspace);
         let mut builder = if pt == PacketType::Short {
             qdebug!("Building Short dcid {}", path.remote_cid());
@@ -1656,17 +1656,17 @@ impl Connection {
         };
         builder.scramble(grease_quic_bit);
         if pt == PacketType::Initial {
-            builder.initial_token(address_validation.token());
+            builder.initial_token(address_validation.token())?;
         }
 
-        (pt, builder)
+        Ok((pt, builder))
     }
 
     fn add_packet_number(
         builder: &mut PacketBuilder,
         tx: &CryptoDxState,
         largest_acknowledged: Option<PacketNumber>,
-    ) -> PacketNumber {
+    ) -> Res<PacketNumber> {
         // Get the packet number and work out how long it is.
         let pn = tx.next_pn();
         let unacked_range = if let Some(la) = largest_acknowledged {
@@ -1680,8 +1680,8 @@ impl Connection {
             - usize::try_from(unacked_range.leading_zeros() / 8).unwrap();
         // pn_len can't be zero (unacked_range is > 0)
         // TODO(mt) also use `4*path CWND/path MTU` to set a minimum length.
-        builder.pn(pn, pn_len);
-        pn
+        builder.pn(pn, pn_len)?;
+        Ok(pn)
     }
 
     fn can_grease_quic_bit(&self) -> bool {
@@ -1715,13 +1715,19 @@ impl Connection {
                 &AddressValidationInfo::None,
                 version,
                 grease_quic_bit,
-            );
+            )?;
             builder.set_limit(min(path.amplification_limit(), path.mtu()) - tx.expansion());
+            if builder.limit() > 2048 {
+                return Err(Error::InternalError(9));
+            }
+            if builder.len() > builder.limit() {
+                return Err(Error::InternalError(25));
+            }
             let _ = Self::add_packet_number(
                 &mut builder,
                 tx,
                 self.loss_recovery.largest_acknowledged_pn(*space),
-            );
+            )?;
 
             // ConnectionError::Application is only allowed at 1RTT.
             let sanitized = if *space == PNSpace::ApplicationData {
@@ -1733,6 +1739,9 @@ impl Connection {
                 .as_ref()
                 .unwrap_or(&close)
                 .write_frame(&mut builder);
+            if builder.len() > builder.limit() {
+                return Err(Error::InternalError(10));
+            }
             encoder = builder.build(tx)?;
         }
 
@@ -1750,14 +1759,14 @@ impl Connection {
         builder: &mut PacketBuilder,
         mut pad: bool,
         now: Instant,
-    ) -> (Vec<RecoveryToken>, bool, bool) {
+    ) -> Res<(Vec<RecoveryToken>, bool, bool)> {
         let mut tokens = Vec::new();
         let stats = &mut self.stats.borrow_mut().frame_tx;
         let primary = path.borrow().is_primary();
         let mut ack_eliciting = false;
 
         let ack_token = if primary {
-            self.acks.write_frame(space, now, builder, stats)
+            self.acks.write_frame(space, now, builder, stats)?
         } else {
             None
         };
@@ -1770,7 +1779,7 @@ impl Connection {
             // The probing code needs to know so it can track that.
             if path
                 .borrow_mut()
-                .write_frames(builder, stats, full_mtu, now)
+                .write_frames(builder, stats, full_mtu, now)?
             {
                 pad = true;
                 ack_eliciting = true;
@@ -1782,31 +1791,32 @@ impl Connection {
             if let Some(t) = ack_token {
                 tokens.push(t);
             }
-            return (tokens, false, false);
+            return Ok((tokens, false, false));
         }
 
         if primary {
             if space == PNSpace::ApplicationData && self.role == Role::Server {
-                if let Some(t) = self.state_signaling.write_done(builder) {
+                if let Some(t) = self.state_signaling.write_done(builder)? {
                     tokens.push(t);
                     stats.handshake_done += 1;
                 }
             }
 
-            if let Some(t) = self.crypto.streams.write_frame(space, builder) {
+            if let Some(t) = self.crypto.streams.write_frame(space, builder)? {
                 tokens.push(t);
                 stats.crypto += 1;
             }
 
             if space == PNSpace::ApplicationData {
                 self.flow_mgr
                     .borrow_mut()
-                    .write_frames(builder, &mut tokens, stats);
+                    .write_frames(builder, &mut tokens, stats)?;
 
-                self.send_streams.write_frames(builder, &mut tokens, stats);
-                self.new_token.write_frames(builder, &mut tokens, stats);
-                self.cid_manager.write_frames(builder, &mut tokens, stats);
-                self.paths.write_frames(builder, &mut tokens, stats);
+                self.send_streams
+                    .write_frames(builder, &mut tokens, stats)?;
+                self.new_token.write_frames(builder, &mut tokens, stats)?;
+                self.cid_manager.write_frames(builder, &mut tokens, stats)?;
+                self.paths.write_frames(builder, &mut tokens, stats)?;
             }
         }
 
@@ -1816,6 +1826,9 @@ impl Connection {
             // Nothing ack-eliciting and we need to probe; send PING.
             debug_assert_ne!(builder.remaining(), 0);
             builder.encode_varint(crate::frame::FRAME_TYPE_PING);
+            if builder.len() > builder.limit() {
+                return Err(Error::InternalError(11));
+            }
             stats.ping += 1;
             stats.all += 1;
             ack_eliciting = true;
@@ -1829,7 +1842,7 @@ impl Connection {
         // And avoid padding if we don't have a full MTU available.
         pad &= ack_eliciting && space == PNSpace::ApplicationData && full_mtu;
         if pad {
-            builder.pad();
+            builder.pad()?;
             stats.padding += 1;
             stats.all += 1;
         }
@@ -1838,7 +1851,7 @@ impl Connection {
             tokens.push(t);
         }
         stats.all += tokens.len();
-        (tokens, ack_eliciting, pad)
+        Ok((tokens, ack_eliciting, pad))
     }
 
     /// Build a datagram, possibly from multiple packets (for different PN
@@ -1877,12 +1890,12 @@ impl Connection {
                 &self.address_validation,
                 version,
                 grease_quic_bit,
-            );
+            )?;
             let pn = Self::add_packet_number(
                 &mut builder,
                 tx,
                 self.loss_recovery.largest_acknowledged_pn(*space),
-            );
+            )?;
             let payload_start = builder.len();
 
             // Work out if we have space left.
@@ -1894,10 +1907,16 @@ impl Connection {
             }
             let limit = profile.limit() - aead_expansion;
             builder.set_limit(limit);
+            if builder.limit() > 2048 {
+                return Err(Error::InternalError(12));
+            }
+            if builder.len() > builder.limit() {
+                return Err(Error::InternalError(13));
+            }
 
             // Add frames to the packet.
             let (tokens, ack_eliciting, padded) =
-                self.write_frames(path, *space, &profile, &mut builder, needs_padding, now);
+                self.write_frames(path, *space, &profile, &mut builder, needs_padding, now)?;
 
             if builder.packet_empty() {
                 // Nothing to include in this packet.

neqo-transport/src/connection/state.rs

@@ -17,7 +17,7 @@ use crate::frame::{
 use crate::packet::PacketBuilder;
 use crate::path::PathRef;
 use crate::recovery::RecoveryToken;
-use crate::{ConnectionError, Error};
+use crate::{ConnectionError, Error, Res};
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 /// The state of the Connection.
@@ -185,13 +185,16 @@ impl StateSignaling {
         *self = Self::HandshakeDone
     }
 
-    pub fn write_done(&mut self, builder: &mut PacketBuilder) -> Option<RecoveryToken> {
+    pub fn write_done(&mut self, builder: &mut PacketBuilder) -> Res<Option<RecoveryToken>> {
         if matches!(self, Self::HandshakeDone) && builder.remaining() >= 1 {
             *self = Self::Idle;
             builder.encode_varint(FRAME_TYPE_HANDSHAKE_DONE);
-            Some(RecoveryToken::HandshakeDone)
+            if builder.len() > builder.limit() {
+                return Err(Error::InternalError(14));
+            }
+            Ok(Some(RecoveryToken::HandshakeDone))
         } else {
-            None
+            Ok(None)
         }
     }
 

neqo-transport/src/connection/tests/idle.rs

@@ -253,12 +253,14 @@ fn idle_caching() {
     let crypto = server
         .crypto
         .streams
-        .write_frame(PNSpace::Initial, &mut builder);
+        .write_frame(PNSpace::Initial, &mut builder)
+        .unwrap();
     assert!(crypto.is_some());
     let crypto = server
         .crypto
         .streams
-        .write_frame(PNSpace::Initial, &mut builder);
+        .write_frame(PNSpace::Initial, &mut builder)
+        .unwrap();
     assert!(crypto.is_none());
     let dgram = server.process_output(middle).dgram();
 

neqo-transport/src/crypto.rs

@@ -1241,14 +1241,14 @@ impl CryptoStreams {
         &mut self,
         space: PNSpace,
         builder: &mut PacketBuilder,
-    ) -> Option<RecoveryToken> {
+    ) -> Res<Option<RecoveryToken>> {
         let cs = self.get_mut(space).unwrap();
         if let Some((offset, data)) = cs.tx.next_bytes() {
             let mut header_len = 1 + Encoder::varint_len(offset) + 1;
 
             // Don't bother if there isn't room for the header and some data.
             if builder.remaining() < header_len + 1 {
-                return None;
+                return Ok(None);
             }
             // Calculate length of data based on the minimum of:
             // - available data
@@ -1261,16 +1261,20 @@ impl CryptoStreams {
             builder.encode_varint(crate::frame::FRAME_TYPE_CRYPTO);
             builder.encode_varint(offset);
             builder.encode_vvec(&data[..length]);
+            if builder.len() > builder.limit() {
+                return Err(Error::InternalError(15));
+            }
+
             cs.tx.mark_as_sent(offset, length);
 
             qdebug!("CRYPTO for {} offset={}, len={}", space, offset, length);
-            Some(RecoveryToken::Crypto(CryptoRecoveryToken {
+            Ok(Some(RecoveryToken::Crypto(CryptoRecoveryToken {
                 space,
                 offset,
                 length,
-            }))
+            })))
         } else {
-            None
+            Ok(None)
         }
     }
 }