Skip to content

Commit

Permalink
Use u32 for Vsock related buffer sizes
Browse files Browse the repository at this point in the history 
Use u32 for all Vsock related buffers instead of usize,
the Virtio specification states that lengths of virtio-buffers
fit into a u32.

Signed-off-by: River Phillips <[email protected]>
  • Loading branch information
@RiverPhillips @roypat
RiverPhillips authored and roypat committed Sep 24, 2024
1 parent 946cf20 commit d772c0f
Show file tree
Hide file tree
Showing 4 changed files with 47 additions and 42 deletions.
21 changes: 11 additions & 10 deletions src/vmm/src/devices/virtio/vsock/csm/connection.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -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) {
Expand All @@ -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
}
}
Expand Down Expand Up @@ -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();
Expand Down
58 changes: 31 additions & 27 deletions src/vmm/src/devices/virtio/vsock/packet.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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),
}
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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)));
Expand Down
6 changes: 3 additions & 3 deletions src/vmm/src/devices/virtio/vsock/test_utils.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions src/vmm/src/devices/virtio/vsock/unix/muxer.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down

0 comments on commit d772c0f

Please sign in to comment.