@@ -24,13 +24,15 @@ pub enum IoVecError {
2424 GuestMemory ( #[ from] GuestMemoryError ) ,
2525}
2626
27+ const SMALL_VEC_SIZE : usize = 4 ;
28+
2729// Using SmallVec in the kani proofs causes kani to use unbounded amounts of memory
2830// during post-processing, and then crash.
2931// TODO: remove new-type once kani performance regression are resolved
3032#[ cfg( kani) ]
3133type IoVecVec = Vec < iovec > ;
3234#[ cfg( not( kani) ) ]
33- type IoVecVec = SmallVec < [ iovec ; 4 ] > ;
35+ type IoVecVec = SmallVec < [ iovec ; SMALL_VEC_SIZE ] > ;
3436
3537/// This is essentially a wrapper of a `Vec<libc::iovec>` which can be passed to `libc::writev`.
3638///
@@ -45,6 +47,15 @@ pub struct IoVecBuffer {
4547 len : u32 ,
4648}
4749
50+ impl IoVecBuffer {
51+ pub fn with_capacity ( cap : usize ) -> IoVecBuffer {
52+ IoVecBuffer {
53+ vecs : IoVecVec :: with_capacity ( cap) ,
54+ len : 0 ,
55+ }
56+ }
57+ }
58+
4859// SAFETY: `IoVecBuffer` doesn't allow for interior mutability and no shared ownership is possible
4960// as it doesn't implement clone
5061unsafe impl Send for IoVecBuffer { }
@@ -57,12 +68,11 @@ impl IoVecBuffer {
5768 /// The descriptor chain cannot be referencing the same memory location as another chain
5869 pub unsafe fn load_descriptor_chain (
5970 & mut self ,
60- head : DescriptorChain ,
71+ mut desc : DescriptorChain ,
6172 ) -> Result < ( ) , IoVecError > {
6273 self . clear ( ) ;
6374
64- let mut next_descriptor = Some ( head) ;
65- while let Some ( desc) = next_descriptor {
75+ loop {
6676 if desc. is_write_only ( ) {
6777 return Err ( IoVecError :: WriteOnlyDescriptor ) ;
6878 }
@@ -85,7 +95,9 @@ impl IoVecBuffer {
8595 . checked_add ( desc. len )
8696 . ok_or ( IoVecError :: OverflowedDescriptor ) ?;
8797
88- next_descriptor = desc. next_descriptor ( ) ;
98+ if desc. load_next_descriptor ( ) . is_none ( ) {
99+ break ;
100+ }
89101 }
90102
91103 Ok ( ( ) )
@@ -217,21 +229,38 @@ impl IoVecBuffer {
217229/// It describes a write-only buffer passed to us by the guest that is scattered across multiple
218230/// memory regions. Additionally, this wrapper provides methods that allow reading arbitrary ranges
219231/// of data from that buffer.
220- #[ derive( Debug ) ]
232+ #[ derive( Debug , Default ) ]
221233pub struct IoVecBufferMut {
222234 // container of the memory regions included in this IO vector
223- vecs : IoVecVec ,
235+ pub ( crate ) vecs : IoVecVec ,
224236 // Total length of the IoVecBufferMut
225- len : u32 ,
237+ pub ( crate ) len : u32 ,
238+ }
239+
240+ impl IoVecBufferMut {
241+ pub fn with_capacity ( cap : usize ) -> IoVecBufferMut {
242+ IoVecBufferMut {
243+ vecs : IoVecVec :: with_capacity ( cap) ,
244+ len : 0 ,
245+ }
246+ }
226247}
227248
249+ // SAFETY: iovec pointers are safe to send across threads.
250+ unsafe impl Send for IoVecBufferMut { }
251+
228252impl IoVecBufferMut {
229253 /// Create an `IoVecBufferMut` from a `DescriptorChain`
230- pub fn from_descriptor_chain ( head : DescriptorChain ) -> Result < Self , IoVecError > {
231- let mut vecs = IoVecVec :: new ( ) ;
232- let mut len = 0u32 ;
254+ /// # Safety
255+ /// The descriptor chain cannot be referencing the same memory location as another chain.
256+ pub unsafe fn load_descriptor_chain (
257+ & mut self ,
258+ mut desc : DescriptorChain ,
259+ max_size : Option < u32 > ,
260+ ) -> Result < ( ) , IoVecError > {
261+ self . clear ( ) ;
233262
234- for desc in head {
263+ loop {
235264 if !desc. is_write_only ( ) {
236265 return Err ( IoVecError :: ReadOnlyDescriptor ) ;
237266 }
@@ -247,16 +276,34 @@ impl IoVecBufferMut {
247276 slice. bitmap ( ) . mark_dirty ( 0 , desc. len as usize ) ;
248277
249278 let iov_base = slice. ptr_guard_mut ( ) . as_ptr ( ) . cast :: < c_void > ( ) ;
250- vecs. push ( iovec {
279+ self . vecs . push ( iovec {
251280 iov_base,
252281 iov_len : desc. len as size_t ,
253282 } ) ;
254- len = len
283+ self . len = self
284+ . len
255285 . checked_add ( desc. len )
256286 . ok_or ( IoVecError :: OverflowedDescriptor ) ?;
287+ if matches ! ( max_size, Some ( max) if self . len >= max) {
288+ break ;
289+ }
290+ if desc. load_next_descriptor ( ) . is_none ( ) {
291+ break ;
292+ }
257293 }
258294
259- Ok ( Self { vecs, len } )
295+ Ok ( ( ) )
296+ }
297+
298+ /// Create an `IoVecBufferMut` from a `DescriptorChain`
299+ /// # Safety
300+ /// The descriptor chain cannot be referencing the same memory location as another chain.
301+ pub unsafe fn from_descriptor_chain ( head : DescriptorChain ) -> Result < Self , IoVecError > {
302+ let mut new_buffer = Self :: default ( ) ;
303+
304+ Self :: load_descriptor_chain ( & mut new_buffer, head, None ) ?;
305+
306+ Ok ( new_buffer)
260307 }
261308
262309 /// Get the total length of the memory regions covered by this `IoVecBuffer`
@@ -298,6 +345,22 @@ impl IoVecBufferMut {
298345 }
299346 }
300347
348+ /// Returns a pointer to the memory keeping the `iovec` structs
349+ pub fn as_iovec_ptr ( & self ) -> * const iovec {
350+ self . vecs . as_ptr ( )
351+ }
352+
353+ /// Returns the length of the `iovec` array.
354+ pub fn iovec_count ( & self ) -> usize {
355+ self . vecs . len ( )
356+ }
357+
358+ /// Clears the `iovec` array
359+ pub fn clear ( & mut self ) {
360+ self . vecs . clear ( ) ;
361+ self . len = 0u32 ;
362+ }
363+
301364 /// Writes up to `len` bytes into the `IoVecBuffer` starting at the given offset.
302365 ///
303366 /// This will try to write to the given [`WriteVolatile`].
@@ -468,11 +531,13 @@ mod tests {
468531
469532 let ( mut q, _) = read_only_chain ( & mem) ;
470533 let head = q. pop ( & mem) . unwrap ( ) ;
471- IoVecBufferMut :: from_descriptor_chain ( head) . unwrap_err ( ) ;
534+ // SAFETY: This descriptor chain is only loaded into one buffer.
535+ unsafe { IoVecBufferMut :: from_descriptor_chain ( head) . unwrap_err ( ) } ;
472536
473537 let ( mut q, _) = write_only_chain ( & mem) ;
474538 let head = q. pop ( & mem) . unwrap ( ) ;
475- IoVecBufferMut :: from_descriptor_chain ( head) . unwrap ( ) ;
539+ // SAFETY: This descriptor chain is only loaded into one buffer.
540+ unsafe { IoVecBufferMut :: from_descriptor_chain ( head) . unwrap ( ) } ;
476541 }
477542
478543 #[ test]
@@ -493,7 +558,7 @@ mod tests {
493558 let head = q. pop ( & mem) . unwrap ( ) ;
494559
495560 // SAFETY: This descriptor chain is only loaded once in this test
496- let iovec = IoVecBufferMut :: from_descriptor_chain ( head) . unwrap ( ) ;
561+ let iovec = unsafe { IoVecBufferMut :: from_descriptor_chain ( head) . unwrap ( ) } ;
497562 assert_eq ! ( iovec. len( ) , 4 * 64 ) ;
498563 }
499564
@@ -558,7 +623,8 @@ mod tests {
558623 // This is a descriptor chain with 4 elements 64 bytes long each.
559624 let head = q. pop ( & mem) . unwrap ( ) ;
560625
561- let mut iovec = IoVecBufferMut :: from_descriptor_chain ( head) . unwrap ( ) ;
626+ // SAFETY: This descriptor chain is only loaded into one buffer.
627+ let mut iovec = unsafe { IoVecBufferMut :: from_descriptor_chain ( head) . unwrap ( ) } ;
562628 let buf = vec ! [ 0u8 , 1 , 2 , 3 , 4 ] ;
563629
564630 // One test vector for each part of the chain
0 commit comments