Repository: firecracker-microvm/firecracker
File: src/vmm/src/devices/virtio/net/device.rs (Rust)

Code before:
        // We only want to read from TAP (or mmds) if we have at least 64K of available capacity as
        // this is the max size of 1 packet.
        // SAFETY:
        // * MAX_BUFFER_SIZE is constant and fits into u32
        #[allow(clippy::cast_possible_truncation)]
        if self.rx_buffer.capacity() < MAX_BUFFER_SIZE as u32 {
            self.parse_rx_descriptors();

            // If after parsing the RX queue we still don't have enough capacity, stop processing RX
            // frames.
            if self.rx_buffer.capacity() < MAX_BUFFER_SIZE as u32 {
                return Ok(None);
            }
        }

        if let Some(ns) = self.mmds_ns.as_mut() {
            if let Some(len) =
                ns.write_next_frame(frame_bytes_from_buf_mut(&mut self.rx_frame_buf)?)
            {
                let len = len.get();
                METRICS.mmds.tx_frames.inc();
                METRICS.mmds.tx_bytes.add(len as u64);
                init_vnet_hdr(&mut self.rx_frame_buf);
                self.rx_buffer
                    .iovec
                    .write_all_volatile_at(&self.rx_frame_buf[..vnet_hdr_len() + len], 0)?;

                // SAFETY:
                // * len is never 0
                // * Maximum size of the packet is 65562 which fits into u32
                #[allow(clippy::cast_possible_truncation)]
                let len = unsafe { NonZeroU32::new_unchecked((vnet_hdr_len() + len) as u32) };

                // SAFETY: This is safe:
                // * We checked that `rx_buffer` includes at least one `DescriptorChain`
                // * `rx_frame_buf` has size of `MAX_BUFFER_SIZE` and all `DescriptorChain` objects
                //   are at least that big.
                unsafe {
                    self.rx_buffer
                        .mark_used(len.get(), &mut self.queues[RX_INDEX]);
                }
                return Ok(Some(len));
            }
        }

        // SAFETY: this is safe because we ensured that `self.rx_buffer` has at least one
        // DescriptorChain parsed in it.
        let len = unsafe { self.read_tap().map_err(NetError::IO) }?;

        // SAFETY:
        // * len is never 0

Review comment (suggestion):
I don't really think using `NonZeroU32` over `u32` adds anything here. We just `.get()` the nonzero at every possible occasion, so I think it's fine to just return `u32` and have `0` mean "out of buffers" (or turn that into a new `NetError` variant, as it's not immediately obvious to me why readv cannot return 0), like it was before @bchalios' PR.

Code after:
                return Ok(None);
            }
        }

        if let Some(ns) = self.mmds_ns.as_mut() {
            if let Some(len) =
                ns.write_next_frame(frame_bytes_from_buf_mut(&mut self.rx_frame_buf)?)
            {
                let len = len.get();
                METRICS.mmds.tx_frames.inc();
                METRICS.mmds.tx_bytes.add(len as u64);
                init_vnet_hdr(&mut self.rx_frame_buf);
                self.rx_buffer
                    .iovec
                    .write_all_volatile_at(&self.rx_frame_buf[..vnet_hdr_len() + len], 0)?;
                // SAFETY:
                // * len will never be bigger that u32::MAX because mmds is bound
                // by the size of `self.rx_frame_buf` which is MAX_BUFFER_SIZE size.
                let len: u32 = (vnet_hdr_len() + len).try_into().unwrap();

                // SAFETY:
                // * We checked that `rx_buffer` includes at least one `DescriptorChain`
                // * `rx_frame_buf` has size of `MAX_BUFFER_SIZE` and all `DescriptorChain` objects
                //   are at least that big.
                unsafe {
                    self.rx_buffer.mark_used(len, &mut self.queues[RX_INDEX]);
                }
                return Ok(Some(len));
            }
        }

        // SAFETY:
        // * We ensured that `self.rx_buffer` has at least one DescriptorChain parsed in it.
        let len = unsafe { self.read_tap().map_err(NetError::IO) }?;
        // SAFETY:
        // * len will never be bigger that u32::MAX
        let len: u32 = len.try_into().unwrap();

        // SAFETY:
        // * `rx_buffer` has at least one `DescriptorChain`
        // * `read_tap` passes the first `DescriptorChain` to `readv` so we can't have read more
        //   bytes than its capacity.
        unsafe {
            self.rx_buffer.mark_used(len, &mut self.queues[RX_INDEX]);
        }
        Ok(Some(len))
    }

    /// Read as many frames as possible.
    fn process_rx(&mut self) -> Result<(), DeviceError> {
        loop {