Change packet lengths to sized types

IPv4 packet header sizes fit in a u8 value, and all IP/Ethernet-related
packet total sizes fit in a u16 value. This changes those sizes to use
smaller types that don't ever need to be converted to smaller types.

Signed-off-by: Jonathan Browne <[email protected]>

Author: JBYoshi

src/vmm/src/dumbo/pdu/ipv4.rs

@@ -25,7 +25,7 @@ const PROTOCOL_OFFSET: usize = 9;
 const HEADER_CHECKSUM_OFFSET: usize = 10;
 const SOURCE_ADDRESS_OFFSET: usize = 12;
 const DESTINATION_ADDRESS_OFFSET: usize = 16;
-const OPTIONS_OFFSET: usize = 20;
+const OPTIONS_OFFSET: u8 = 20;
 
 /// Indicates version 4 of the IP protocol
 pub const IPV4_VERSION: u8 = 0x04;
@@ -82,7 +82,7 @@ impl<'a, T: NetworkBytes + Debug> IPv4Packet<'a, T> {
     pub fn from_bytes(bytes: T, verify_checksum: bool) -> Result<Self, Error> {
         let bytes_len = bytes.len();
 
-        if bytes_len < OPTIONS_OFFSET {
+        if bytes_len < usize::from(OPTIONS_OFFSET) {
             return Err(Error::SliceTooShort);
         }
 
@@ -96,7 +96,7 @@ impl<'a, T: NetworkBytes + Debug> IPv4Packet<'a, T> {
 
         let total_len = packet.total_len() as usize;
 
-        if total_len < header_len {
+        if total_len < header_len.into() {
             return Err(Error::InvalidTotalLen);
         }
 
@@ -111,7 +111,7 @@ impl<'a, T: NetworkBytes + Debug> IPv4Packet<'a, T> {
         // We ignore the TTL field since only routers should care about it. An end host has no
         // reason really to discard an otherwise valid packet.
 
-        if verify_checksum && packet.compute_checksum_unchecked(header_len) != 0 {
+        if verify_checksum && packet.compute_checksum_unchecked(header_len.into()) != 0 {
             return Err(Error::Checksum);
         }
 
@@ -123,16 +123,16 @@ impl<'a, T: NetworkBytes + Debug> IPv4Packet<'a, T> {
     /// This method returns the actual length (in bytes) of the header, and not the value of the
     /// `ihl` header field).
     #[inline]
-    pub fn version_and_header_len(&self) -> (u8, usize) {
+    pub fn version_and_header_len(&self) -> (u8, u8) {
         let x = self.bytes[VERSION_AND_IHL_OFFSET];
         let ihl = x & 0x0f;
-        let header_len = (ihl << 2) as usize;
+        let header_len = ihl << 2;
         (x >> 4, header_len)
     }
 
     /// Returns the packet header length (in bytes).
     #[inline]
-    pub fn header_len(&self) -> usize {
+    pub fn header_len(&self) -> u8 {
         let (_, header_len) = self.version_and_header_len();
         header_len
     }
@@ -207,7 +207,7 @@ impl<'a, T: NetworkBytes + Debug> IPv4Packet<'a, T> {
     /// Returns a byte slice that contains the payload of the packet.
     #[inline]
     pub fn payload(&self) -> &[u8] {
-        self.payload_unchecked(self.header_len())
+        self.payload_unchecked(self.header_len().into())
     }
 
     /// Returns the length of the inner byte sequence.
@@ -245,7 +245,7 @@ impl<'a, T: NetworkBytes + Debug> IPv4Packet<'a, T> {
     /// Computes and returns the packet header checksum.
     #[inline]
     pub fn compute_checksum(&self) -> u16 {
-        self.compute_checksum_unchecked(self.header_len())
+        self.compute_checksum_unchecked(self.header_len().into())
     }
 }
 
@@ -263,7 +263,7 @@ impl<'a, T: NetworkBytesMut + Debug> IPv4Packet<'a, T> {
         src_addr: Ipv4Addr,
         dst_addr: Ipv4Addr,
     ) -> Result<Incomplete<Self>, Error> {
-        if buf.len() < OPTIONS_OFFSET {
+        if buf.len() < usize::from(OPTIONS_OFFSET) {
             return Err(Error::SliceTooShort);
         }
         let mut packet = IPv4Packet::from_bytes_unchecked(buf);
@@ -283,9 +283,9 @@ impl<'a, T: NetworkBytesMut + Debug> IPv4Packet<'a, T> {
     /// Sets the values of the `version` and `ihl` header fields (the latter is computed from the
     /// value of `header_len`).
     #[inline]
-    pub fn set_version_and_header_len(&mut self, version: u8, header_len: usize) -> &mut Self {
+    pub fn set_version_and_header_len(&mut self, version: u8, header_len: u8) -> &mut Self {
         let version = version << 4;
-        let ihl = ((header_len as u8) >> 2) & 0xf;
+        let ihl = (header_len >> 2) & 0xf;
         self.bytes[VERSION_AND_IHL_OFFSET] = version | ihl;
         self
     }
@@ -374,7 +374,7 @@ impl<'a, T: NetworkBytesMut + Debug> IPv4Packet<'a, T> {
         // Can't use self.header_len() as a fn parameter on the following line, because
         // the borrow checker complains. This may change when it becomes smarter.
         let header_len = self.header_len();
-        self.payload_mut_unchecked(header_len)
+        self.payload_mut_unchecked(header_len.into())
     }
 }
 
@@ -394,24 +394,24 @@ impl<'a, T: NetworkBytesMut + Debug> Incomplete<IPv4Packet<'a, T>> {
     #[inline]
     pub fn with_header_and_payload_len_unchecked(
         mut self,
-        header_len: usize,
-        payload_len: usize,
+        header_len: u8,
+        payload_len: u16,
         compute_checksum: bool,
     ) -> IPv4Packet<'a, T> {
-        let total_len = header_len + payload_len;
+        let total_len = u16::from(header_len) + payload_len;
         {
             let packet = &mut self.inner;
 
             // This unchecked is fine as long as total_len is smaller than the length of the
             // original slice, which should be the case if our code is not wrong.
-            packet.bytes.shrink_unchecked(total_len);
+            packet.bytes.shrink_unchecked(total_len.into());
             // Set the total_len.
-            packet.set_total_len(total_len as u16);
+            packet.set_total_len(total_len);
             if compute_checksum {
                 // Ensure this is set to 0 first.
                 packet.set_header_checksum(0);
                 // Now compute the actual checksum.
-                let checksum = packet.compute_checksum_unchecked(header_len);
+                let checksum = packet.compute_checksum_unchecked(header_len.into());
                 packet.set_header_checksum(checksum);
             }
         }
@@ -427,8 +427,8 @@ impl<'a, T: NetworkBytesMut + Debug> Incomplete<IPv4Packet<'a, T>> {
     #[inline]
     pub fn with_options_and_payload_len_unchecked(
         self,
-        options_len: usize,
-        payload_len: usize,
+        options_len: u8,
+        payload_len: u16,
         compute_checksum: bool,
     ) -> IPv4Packet<'a, T> {
         let header_len = OPTIONS_OFFSET + options_len;
@@ -444,7 +444,7 @@ impl<'a, T: NetworkBytesMut + Debug> Incomplete<IPv4Packet<'a, T>> {
     #[inline]
     pub fn with_payload_len_unchecked(
         self,
-        payload_len: usize,
+        payload_len: u16,
         compute_checksum: bool,
     ) -> IPv4Packet<'a, T> {
         let header_len = self.inner().header_len();
@@ -457,7 +457,7 @@ impl<'a, T: NetworkBytesMut + Debug> Incomplete<IPv4Packet<'a, T>> {
 #[inline]
 pub fn test_speculative_dst_addr(buf: &[u8], addr: Ipv4Addr) -> bool {
     // The unchecked methods are safe because we actually check the buffer length beforehand.
-    if buf.len() >= ethernet::PAYLOAD_OFFSET + OPTIONS_OFFSET {
+    if buf.len() >= ethernet::PAYLOAD_OFFSET + usize::from(OPTIONS_OFFSET) {
         let bytes = &buf[ethernet::PAYLOAD_OFFSET..];
         if IPv4Packet::from_bytes_unchecked(bytes).destination_address() == addr {
             return true;

src/vmm/src/dumbo/pdu/mod.rs

@@ -79,28 +79,26 @@ fn compute_checksum<T: NetworkBytes + Debug>(
     dst_addr: Ipv4Addr,
     protocol: ChecksumProto,
 ) -> u16 {
-    // TODO: Is u32 enough to prevent overflow for the code in this function? I think so, but it
-    // would be nice to double-check.
-    let mut sum = 0u32;
+    let mut sum = 0usize;
 
-    let a = u32::from(src_addr);
+    let a = u32::from(src_addr) as usize;
     sum += a & 0xffff;
     sum += a >> 16;
 
-    let b = u32::from(dst_addr);
+    let b = u32::from(dst_addr) as usize;
     sum += b & 0xffff;
     sum += b >> 16;
 
     let len = bytes.len();
-    sum += protocol as u32;
-    sum += len as u32;
+    sum += protocol as usize;
+    sum += len;
 
     for i in 0..len / 2 {
-        sum += u32::from(bytes.ntohs_unchecked(i * 2));
+        sum += usize::from(bytes.ntohs_unchecked(i * 2));
     }
 
     if len % 2 != 0 {
-        sum += u32::from(bytes[len - 1]) << 8;
+        sum += usize::from(bytes[len - 1]) << 8;
     }
 
     while sum >> 16 != 0 {

src/vmm/src/dumbo/pdu/tcp.rs

@@ -30,9 +30,9 @@ const WINDOW_SIZE_OFFSET: usize = 14;
 const CHECKSUM_OFFSET: usize = 16;
 const URG_POINTER_OFFSET: usize = 18;
 
-const OPTIONS_OFFSET: usize = 20;
+const OPTIONS_OFFSET: u8 = 20;
 
-const MAX_HEADER_LEN: usize = 60;
+const MAX_HEADER_LEN: u8 = 60;
 
 const OPTION_KIND_EOL: u8 = 0x00;
 const OPTION_KIND_NOP: u8 = 0x01;
@@ -127,18 +127,18 @@ impl<'a, T: NetworkBytes + Debug> TcpSegment<'a, T> {
     /// Returns the header length, the value of the reserved bits, and whether the `NS` flag
     /// is set or not.
     #[inline]
-    pub fn header_len_rsvd_ns(&self) -> (usize, u8, bool) {
+    pub fn header_len_rsvd_ns(&self) -> (u8, u8, bool) {
         let value = self.bytes[DATAOFF_RSVD_NS_OFFSET];
         let data_offset = value >> 4;
-        let header_len = data_offset as usize * 4;
+        let header_len = data_offset * 4;
         let rsvd = value & 0x0e;
         let ns = (value & 1) != 0;
         (header_len, rsvd, ns)
     }
 
     /// Returns the length of the header.
     #[inline]
-    pub fn header_len(&self) -> usize {
+    pub fn header_len(&self) -> u8 {
         self.header_len_rsvd_ns().0
     }
 
@@ -174,7 +174,7 @@ impl<'a, T: NetworkBytes + Debug> TcpSegment<'a, T> {
     /// This method may panic if the value of `header_len` is invalid.
     #[inline]
     pub fn options_unchecked(&self, header_len: usize) -> &[u8] {
-        &self.bytes[OPTIONS_OFFSET..header_len]
+        &self.bytes[usize::from(OPTIONS_OFFSET)..header_len]
     }
 
     /// Returns a slice which contains the payload of the segment. May panic if the value of
@@ -190,20 +190,23 @@ impl<'a, T: NetworkBytes + Debug> TcpSegment<'a, T> {
 
     /// Returns the length of the segment.
     #[inline]
-    pub fn len(&self) -> usize {
-        self.bytes.len()
+    pub fn len(&self) -> u16 {
+        // NOTE: This appears to be a safe conversion in all current cases.
+        // Packets are always set up in the context of an Ipv4Packet, which is
+        // capped at a u16 size. However, I'd rather be safe here.
+        u16::try_from(self.bytes.len()).unwrap_or(u16::MAX)
     }
 
     /// Returns a slice which contains the payload of the segment.
     #[inline]
     pub fn payload(&self) -> &[u8] {
-        self.payload_unchecked(self.header_len())
+        self.payload_unchecked(self.header_len().into())
     }
 
     /// Returns the length of the payload.
     #[inline]
-    pub fn payload_len(&self) -> usize {
-        self.len() - self.header_len()
+    pub fn payload_len(&self) -> u16 {
+        self.len() - u16::from(self.header_len())
     }
 
     /// Computes the TCP checksum of the segment. More details about TCP checksum computation can
@@ -285,7 +288,7 @@ impl<'a, T: NetworkBytes + Debug> TcpSegment<'a, T> {
         bytes: T,
         verify_checksum: Option<(Ipv4Addr, Ipv4Addr)>,
     ) -> Result<Self, Error> {
-        if bytes.len() < OPTIONS_OFFSET {
+        if bytes.len() < usize::from(OPTIONS_OFFSET) {
             return Err(Error::SliceTooShort);
         }
 
@@ -295,7 +298,9 @@ impl<'a, T: NetworkBytes + Debug> TcpSegment<'a, T> {
 
         let header_len = segment.header_len();
 
-        if header_len < OPTIONS_OFFSET || header_len > min(MAX_HEADER_LEN, segment.len()) {
+        if header_len < OPTIONS_OFFSET
+            || u16::from(header_len) > min(u16::from(MAX_HEADER_LEN), segment.len())
+        {
             return Err(Error::HeaderLen);
         }
 
@@ -342,8 +347,8 @@ impl<'a, T: NetworkBytesMut + Debug> TcpSegment<'a, T> {
     /// of 4), clears the reserved bits, and sets the `NS` flag according to the last parameter.
     // TODO: Check that header_len | 0b11 == 0 and the resulting data_offset is valid?
     #[inline]
-    pub fn set_header_len_rsvd_ns(&mut self, header_len: usize, ns: bool) -> &mut Self {
-        let mut value = (header_len as u8) << 2;
+    pub fn set_header_len_rsvd_ns(&mut self, header_len: u8, ns: bool) -> &mut Self {
+        let mut value = header_len << 2;
         if ns {
             value |= 1;
         }
@@ -393,7 +398,7 @@ impl<'a, T: NetworkBytesMut + Debug> TcpSegment<'a, T> {
     #[inline]
     pub fn payload_mut(&mut self) -> &mut [u8] {
         let header_len = self.header_len();
-        self.payload_mut_unchecked(header_len)
+        self.payload_mut_unchecked(header_len.into())
     }
 
     /// Writes a complete TCP segment.
@@ -479,24 +484,24 @@ impl<'a, T: NetworkBytesMut + Debug> TcpSegment<'a, T> {
         mss_remaining: u16,
         payload: Option<(&R, usize)>,
     ) -> Result<Incomplete<Self>, Error> {
-        let mut mss_left = mss_remaining as usize;
+        let mut mss_left = mss_remaining;
 
         // We're going to need at least this many bytes.
-        let mut segment_len = OPTIONS_OFFSET;
+        let mut segment_len = u16::from(OPTIONS_OFFSET);
 
         // The TCP options will require this much more bytes.
         let options_len = if mss_option.is_some() {
             mss_left = mss_left
                 .checked_sub(OPTION_LEN_MSS.into())
                 .ok_or(Error::MssRemaining)?;
-            usize::from(OPTION_LEN_MSS)
+            OPTION_LEN_MSS
         } else {
             0
         };
 
-        segment_len += options_len;
+        segment_len += u16::from(options_len);
 
-        if buf.len() < segment_len {
+        if buf.len() < usize::from(segment_len) {
             return Err(Error::SliceTooShort);
         }
 
@@ -513,9 +518,11 @@ impl<'a, T: NetworkBytesMut + Debug> TcpSegment<'a, T> {
 
         // Let's write the MSS option if we have to.
         if let Some(value) = mss_option {
-            segment.bytes[OPTIONS_OFFSET] = OPTION_KIND_MSS;
-            segment.bytes[OPTIONS_OFFSET + 1] = OPTION_LEN_MSS;
-            segment.bytes.htons_unchecked(OPTIONS_OFFSET + 2, value);
+            segment.bytes[usize::from(OPTIONS_OFFSET)] = OPTION_KIND_MSS;
+            segment.bytes[usize::from(OPTIONS_OFFSET) + 1] = OPTION_LEN_MSS;
+            segment
+                .bytes
+                .htons_unchecked(usize::from(OPTIONS_OFFSET) + 2, value);
         }
 
         let payload_bytes_count = if let Some((payload_buf, max_payload_bytes)) = payload {
@@ -524,7 +531,9 @@ impl<'a, T: NetworkBytesMut + Debug> TcpSegment<'a, T> {
             // The subtraction makes sense because we previously checked that
             // buf.len() >= segment_len.
             let mut room_for_payload = min(segment.len() - segment_len, mss_left);
-            room_for_payload = min(room_for_payload, left_to_read);
+            // The unwrap is safe because room_for_payload is a u16.
+            room_for_payload =
+                u16::try_from(min(usize::from(room_for_payload), left_to_read)).unwrap();
 
             if room_for_payload == 0 {
                 return Err(Error::EmptyPayload);
@@ -535,7 +544,8 @@ impl<'a, T: NetworkBytesMut + Debug> TcpSegment<'a, T> {
             // `offset + room_for_payload <= payload_buf.len()`.
             payload_buf.read_to_slice(
                 0,
-                &mut segment.bytes[segment_len..segment_len + room_for_payload],
+                &mut segment.bytes
+                    [usize::from(segment_len)..usize::from(segment_len + room_for_payload)],
             );
             room_for_payload
         } else {
@@ -544,7 +554,7 @@ impl<'a, T: NetworkBytesMut + Debug> TcpSegment<'a, T> {
         segment_len += payload_bytes_count;
 
         // This is ok because segment_len <= buf.len().
-        segment.bytes.shrink_unchecked(segment_len);
+        segment.bytes.shrink_unchecked(segment_len.into());
 
         // Shrink the resulting segment to a slice of exact size, so using self.len() makes sense.
         Ok(Incomplete::new(segment))