Use u32 for Vsock related buffer sizes

src/vmm/src/devices/virtio/vsock/csm/connection.rs

@@ -237,8 +237,7 @@ where
                         // length of the read data.
                         // Safe to unwrap because read_cnt is no more than max_len, which is bounded
                         // by self.peer_avail_credit(), a u32 internally.
-                        pkt.set_op(uapi::VSOCK_OP_RW)
-                            .set_len(u32::try_from(read_cnt).unwrap());
+                        pkt.set_op(uapi::VSOCK_OP_RW).set_len(read_cnt);
                         METRICS.rx_bytes_count.add(read_cnt as u64);
                     }
                     self.rx_cnt += Wrapping(pkt.len());
@@ -605,7 +604,7 @@ where
     /// Raw data can either be sent straight to the host stream, or to our TX buffer, if the
     /// former fails.
     fn send_bytes(&mut self, pkt: &VsockPacket) -> Result<(), VsockError> {
-        let len = pkt.len() as usize;
+        let len = pkt.len();
 
         // If there is data in the TX buffer, that means we're already registered for EPOLLOUT
         // events on the underlying stream. Therefore, there's no point in attempting a write
@@ -635,7 +634,7 @@ where
         };
         // Move the "forwarded bytes" counter ahead by how much we were able to send out.
         // Safe to unwrap because the maximum value is pkt.len(), which is a u32.
-        self.fwd_cnt += wrap_usize_to_u32(written);
+        self.fwd_cnt += written;
         METRICS.tx_bytes_count.add(written as u64);
 
         // If we couldn't write the whole slice, we'll need to push the remaining data to our
@@ -662,8 +661,8 @@ where
 
     /// Get the maximum number of bytes that we can send to our peer, without overflowing its
     /// buffer.
-    fn peer_avail_credit(&self) -> usize {
-        (Wrapping(self.peer_buf_alloc) - (self.rx_cnt - self.peer_fwd_cnt)).0 as usize
+    fn peer_avail_credit(&self) -> u32 {
+        (Wrapping(self.peer_buf_alloc) - (self.rx_cnt - self.peer_fwd_cnt)).0
     }
 
     /// Prepare a packet header for transmission to our peer.
@@ -916,7 +915,7 @@ mod tests {
             assert!(credit < self.conn.peer_buf_alloc);
             self.conn.peer_fwd_cnt = Wrapping(0);
             self.conn.rx_cnt = Wrapping(self.conn.peer_buf_alloc - credit);
-            assert_eq!(self.conn.peer_avail_credit(), credit as usize);
+            assert_eq!(self.conn.peer_avail_credit(), credit);
         }
 
         fn send(&mut self) {
@@ -941,11 +940,13 @@ mod tests {
         }
 
         fn init_data_tx_pkt(&mut self, mut data: &[u8]) -> &VsockPacket {
-            assert!(data.len() <= self.tx_pkt.buf_size());
+            assert!(data.len() <= self.tx_pkt.buf_size() as usize);
             self.init_tx_pkt(uapi::VSOCK_OP_RW, u32::try_from(data.len()).unwrap());
 
             let len = data.len();
-            self.rx_pkt.read_at_offset_from(&mut data, 0, len).unwrap();
+            self.rx_pkt
+                .read_at_offset_from(&mut data, 0, len.try_into().unwrap())
+                .unwrap();
             &self.tx_pkt
         }
     }
@@ -1282,7 +1283,7 @@ mod tests {
         ctx.set_stream(stream);
 
         // Fill up the TX buffer.
-        let data = vec![0u8; ctx.tx_pkt.buf_size()];
+        let data = vec![0u8; ctx.tx_pkt.buf_size() as usize];
         ctx.init_data_tx_pkt(data.as_slice());
         for _i in 0..(csm_defs::CONN_TX_BUF_SIZE as usize / data.len()) {
             ctx.send();

src/vmm/src/devices/virtio/vsock/packet.rs

@@ -219,57 +219,61 @@ impl VsockPacket {
     ///
     /// Return value will equal the total length of the underlying descriptor chain's buffers,
     /// minus the length of the vsock header.
-    pub fn buf_size(&self) -> usize {
+    pub fn buf_size(&self) -> u32 {
         let chain_length = match self.buffer {
             VsockPacketBuffer::Tx(ref iovec_buf) => iovec_buf.len(),
             VsockPacketBuffer::Rx(ref iovec_buf) => iovec_buf.len(),
         };
-        (chain_length - VSOCK_PKT_HDR_SIZE) as usize
+        chain_length - VSOCK_PKT_HDR_SIZE
     }
 
     pub fn read_at_offset_from<T: ReadVolatile + Debug>(
         &mut self,
         src: &mut T,
-        offset: usize,
-        count: usize,
-    ) -> Result<usize, VsockError> {
+        offset: u32,
+        count: u32,
+    ) -> Result<u32, VsockError> {
         match self.buffer {
             VsockPacketBuffer::Tx(_) => Err(VsockError::UnwritableDescriptor),
             VsockPacketBuffer::Rx(ref mut buffer) => {
                 if count
-                    > (buffer.len() as usize)
-                        .saturating_sub(VSOCK_PKT_HDR_SIZE as usize)
+                    > buffer
+                        .len()
+                        .saturating_sub(VSOCK_PKT_HDR_SIZE)
                         .saturating_sub(offset)
                 {
                     return Err(VsockError::GuestMemoryBounds);
                 }
 
                 buffer
-                    .write_volatile_at(src, offset + VSOCK_PKT_HDR_SIZE as usize, count)
+                    .write_volatile_at(src, (offset + VSOCK_PKT_HDR_SIZE) as usize, count as usize)
                     .map_err(|err| VsockError::GuestMemoryMmap(GuestMemoryError::from(err)))
+                    .and_then(|read| read.try_into().map_err(|_| VsockError::DescChainOverflow))
             }
         }
     }
 
     pub fn write_from_offset_to<T: WriteVolatile + Debug>(
         &self,
         dst: &mut T,
-        offset: usize,
-        count: usize,
-    ) -> Result<usize, VsockError> {
+        offset: u32,
+        count: u32,
+    ) -> Result<u32, VsockError> {
         match self.buffer {
             VsockPacketBuffer::Tx(ref buffer) => {
                 if count
-                    > (buffer.len() as usize)
-                        .saturating_sub(VSOCK_PKT_HDR_SIZE as usize)
+                    > buffer
+                        .len()
+                        .saturating_sub(VSOCK_PKT_HDR_SIZE)
                         .saturating_sub(offset)
                 {
                     return Err(VsockError::GuestMemoryBounds);
                 }
 
                 buffer
-                    .read_volatile_at(dst, offset + VSOCK_PKT_HDR_SIZE as usize, count)
+                    .read_volatile_at(dst, (offset + VSOCK_PKT_HDR_SIZE) as usize, count as usize)
                     .map_err(|err| VsockError::GuestMemoryMmap(GuestMemoryError::from(err)))
+                    .and_then(|read| read.try_into().map_err(|_| VsockError::DescChainOverflow))
             }
             VsockPacketBuffer::Rx(_) => Err(VsockError::UnreadableDescriptor),
         }
@@ -537,10 +541,7 @@ mod tests {
                 handler_ctx.device.queues[RXQ_INDEX].pop().unwrap(),
             )
             .unwrap();
-            assert_eq!(
-                pkt.buf_size(),
-                handler_ctx.guest_rxvq.dtable[1].len.get() as usize
-            );
+            assert_eq!(pkt.buf_size(), handler_ctx.guest_rxvq.dtable[1].len.get());
         }
 
         // Test case: read-only RX packet header.
@@ -646,35 +647,38 @@ mod tests {
         .unwrap();
 
         let buf_desc = &mut handler_ctx.guest_rxvq.dtable[1];
-        assert_eq!(pkt.buf_size(), buf_desc.len.get() as usize);
-        let zeros = vec![0_u8; pkt.buf_size()];
+        assert_eq!(pkt.buf_size(), buf_desc.len.get());
+        let zeros = vec![0_u8; pkt.buf_size() as usize];
         let data: Vec<u8> = (0..pkt.buf_size())
             .map(|i| ((i as u64) & 0xff) as u8)
             .collect();
         for offset in 0..pkt.buf_size() {
             buf_desc.set_data(&zeros);
 
-            let mut expected_data = zeros[..offset].to_vec();
-            expected_data.extend_from_slice(&data[..pkt.buf_size() - offset]);
+            let mut expected_data = zeros[..offset as usize].to_vec();
+            expected_data.extend_from_slice(&data[..(pkt.buf_size() - offset) as usize]);
 
             pkt.read_at_offset_from(&mut data.as_slice(), offset, pkt.buf_size() - offset)
                 .unwrap();
 
             buf_desc.check_data(&expected_data);
 
-            let mut buf = vec![0; pkt.buf_size()];
+            let mut buf = vec![0; pkt.buf_size() as usize];
             pkt2.write_from_offset_to(&mut buf.as_mut_slice(), offset, pkt.buf_size() - offset)
                 .unwrap();
-            assert_eq!(&buf[..pkt.buf_size() - offset], &expected_data[offset..]);
+            assert_eq!(
+                &buf[..(pkt.buf_size() - offset) as usize],
+                &expected_data[offset as usize..]
+            );
         }
 
         let oob_cases = vec![
             (1, pkt.buf_size()),
             (pkt.buf_size(), 1),
-            (usize::MAX, 1),
-            (1, usize::MAX),
+            (u32::MAX, 1),
+            (1, u32::MAX),
         ];
-        let mut buf = vec![0; pkt.buf_size()];
+        let mut buf = vec![0; pkt.buf_size() as usize];
         for (offset, count) in oob_cases {
             let res = pkt.read_at_offset_from(&mut data.as_slice(), offset, count);
             assert!(matches!(res, Err(VsockError::GuestMemoryBounds)));

src/vmm/src/devices/virtio/vsock/test_utils.rs

@@ -69,7 +69,7 @@ impl VsockChannel for TestBackend {
                 let buf_size = pkt.buf_size();
                 if buf_size > 0 {
                     let buf: Vec<u8> = (0..buf_size)
-                        .map(|i| cool_buf[i % cool_buf.len()])
+                        .map(|i| cool_buf[i as usize % cool_buf.len()])
                         .collect();
                     pkt.read_at_offset_from(&mut buf.as_slice(), 0, buf_size)
                         .unwrap();
@@ -206,8 +206,8 @@ impl<'a> EventHandlerContext<'a> {
 }
 
 #[cfg(test)]
-pub fn read_packet_data(pkt: &VsockPacket, how_much: usize) -> Vec<u8> {
-    let mut buf = vec![0; how_much];
+pub fn read_packet_data(pkt: &VsockPacket, how_much: u32) -> Vec<u8> {
+    let mut buf = vec![0; how_much as usize];
     pkt.write_from_offset_to(&mut buf.as_mut_slice(), 0, how_much)
         .unwrap();
     buf

src/vmm/src/devices/virtio/vsock/unix/muxer.rs

@@ -869,11 +869,11 @@ mod tests {
             peer_port: u32,
             mut data: &[u8],
         ) -> &mut VsockPacket {
-            assert!(data.len() <= self.tx_pkt.buf_size());
+            assert!(data.len() <= self.tx_pkt.buf_size() as usize);
             self.init_tx_pkt(local_port, peer_port, uapi::VSOCK_OP_RW)
                 .set_len(u32::try_from(data.len()).unwrap());
 
-            let data_len = data.len(); // store in tmp var to make borrow checker happy.
+            let data_len = data.len().try_into().unwrap(); // store in tmp var to make borrow checker happy.
             self.rx_pkt
                 .read_at_offset_from(&mut data, 0, data_len)
                 .unwrap();