diff --git a/src/devices/src/virtio/file_traits.rs b/src/devices/src/virtio/file_traits.rs index e496896d..3d75c44a 100644 --- a/src/devices/src/virtio/file_traits.rs +++ b/src/devices/src/virtio/file_traits.rs @@ -7,14 +7,14 @@ use std::io::{Error, ErrorKind, Result}; use std::os::unix::io::AsRawFd; #[cfg(feature = "blk")] -use imago::io_buffers::{IoVector, IoVectorMut}; +use imago::io_buffers::{IoBuffer, IoVector, IoVectorMut}; use vm_memory::VolatileSlice; use libc::{c_int, c_void, read, readv, size_t, write, writev}; use super::bindings::{off64_t, pread64, preadv64, pwrite64, pwritev64}; #[cfg(feature = "blk")] -use super::block::device::DiskProperties; +use super::block::{device::DiskProperties, disk}; /// A trait for setting the size of a file. /// This is equivalent to File's `set_len` method, but @@ -440,9 +440,12 @@ impl FileReadWriteAtVolatile for DiskProperties { } fn write_at_volatile(&self, slice: VolatileSlice, offset: u64) -> Result { - let slices = &[&slice]; - let iovec = IoVector::from_volatile_slice(slices); - self.file().writev(iovec.0, offset)?; + let slices = &[slice]; + + let mut _copied_head: Option = None; + let iovec = create_iovec(self, slices, offset, &mut _copied_head)?; + + self.file().writev(iovec, offset)?; Ok(slice.len()) } @@ -451,7 +454,9 @@ impl FileReadWriteAtVolatile for DiskProperties { return Ok(0); } - let (iovec, _guard) = IoVector::from_volatile_slice(bufs); + let mut _copied_head: Option = None; + let iovec = create_iovec(self, bufs, offset, &mut _copied_head)?; + let full_length = iovec .len() .try_into() @@ -460,3 +465,54 @@ impl FileReadWriteAtVolatile for DiskProperties { Ok(full_length) } } + +#[cfg(feature = "blk")] +fn create_iovec<'a>( + disk: &DiskProperties, + bufs: &'a [VolatileSlice], + offset: u64, + head_buffer: &'a mut Option, +) -> Result> { + match disk.image_format() { + disk::ImageType::Raw => { + let (iovec, _guard) = IoVector::from_volatile_slice(bufs); + guard_header_bytes(disk.file(), iovec, offset, head_buffer) + } + _ => Ok(IoVector::from_volatile_slice(bufs).0), + } +} + +#[cfg(feature = "blk")] +fn guard_header_bytes<'a>( + image: &imago::SyncFormatAccess, + iovec: IoVector<'a>, + offset: u64, + head_buffer: &'a mut Option, +) -> Result> { + if (1..4).contains(&offset) || (offset + iovec.len() < 4) { + // Make sure the guest writes to the first four bytes as a whole. This is done to prevent a + // Raw disk image byte-by-byte, and in any order, writing the header of a formatted disk + // image into the first sector of the disk image. + return Err(Error::new( + ErrorKind::InvalidInput, + "partial write to first four bytes of disk image", + )); + } + + if offset > 0 { + return Ok(iovec); + } + + let (head, tail) = iovec.split_at(std::cmp::max(image.mem_align(), 4) as u64); + *head_buffer = Some(head.try_into_owned(image.mem_align())?); + let head_slice = head_buffer.as_ref().unwrap().as_ref().into_slice(); + + if head_slice.get(0..4) == Some(&disk::QCOW_MAGIC.to_be_bytes()) { + return Err(Error::new( + ErrorKind::InvalidData, + "writing QCOW2 header to disk image", + )); + } + + Ok(tail.with_inserted(0, head_slice)) +}