1use std::cmp::min;
24use std::mem::size_of;
25use std::ptr::copy;
26use std::ptr::read_volatile;
27use std::ptr::write_bytes;
28use std::ptr::write_volatile;
29use std::result;
30use std::slice;
31
32use remain::sorted;
33use thiserror::Error;
34use zerocopy::FromBytes;
35use zerocopy::IntoBytes;
36
37use crate::IoBufMut;
38
39#[sorted]
40#[derive(Error, Eq, PartialEq, Debug)]
41pub enum VolatileMemoryError {
42 #[error("address 0x{addr:x} is out of bounds")]
44 OutOfBounds { addr: usize },
45 #[error("address 0x{base:x} offset by 0x{offset:x} would overflow")]
47 Overflow { base: usize, offset: usize },
48}
49
50pub type VolatileMemoryResult<T> = result::Result<T, VolatileMemoryError>;
51
52use crate::VolatileMemoryError as Error;
53type Result<T> = VolatileMemoryResult<T>;
54
55pub trait VolatileMemory {
57 fn get_slice(&self, offset: usize, count: usize) -> Result<VolatileSlice>;
60}
61
62#[derive(Copy, Clone, Debug)]
66#[repr(transparent)]
67pub struct VolatileSlice<'a>(IoBufMut<'a>);
68
69impl<'a> VolatileSlice<'a> {
70 pub fn new(buf: &mut [u8]) -> VolatileSlice {
72 VolatileSlice(IoBufMut::new(buf))
73 }
74
75 pub unsafe fn from_raw_parts(addr: *mut u8, len: usize) -> VolatileSlice<'a> {
82 VolatileSlice(IoBufMut::from_raw_parts(addr, len))
83 }
84
85 pub fn as_ptr(&self) -> *const u8 {
87 self.0.as_ptr()
88 }
89
90 pub fn as_mut_ptr(&self) -> *mut u8 {
92 self.0.as_mut_ptr()
93 }
94
95 pub fn size(&self) -> usize {
97 self.0.len()
98 }
99
100 pub fn advance(&mut self, count: usize) {
104 self.0.advance(count)
105 }
106
107 pub fn truncate(&mut self, len: usize) {
111 self.0.truncate(len)
112 }
113
114 pub fn as_iobuf(&self) -> &IoBufMut {
116 &self.0
117 }
118
119 #[allow(clippy::wrong_self_convention)]
121 pub fn as_iobufs<'mem, 'slice>(
122 iovs: &'slice [VolatileSlice<'mem>],
123 ) -> &'slice [IoBufMut<'mem>] {
124 unsafe { slice::from_raw_parts(iovs.as_ptr() as *const IoBufMut, iovs.len()) }
127 }
128
129 #[inline]
131 pub fn as_iobufs_mut<'mem, 'slice>(
132 iovs: &'slice mut [VolatileSlice<'mem>],
133 ) -> &'slice mut [IoBufMut<'mem>] {
134 unsafe { slice::from_raw_parts_mut(iovs.as_mut_ptr() as *mut IoBufMut, iovs.len()) }
137 }
138
139 pub fn offset(self, count: usize) -> Result<VolatileSlice<'a>> {
142 let new_addr = (self.as_mut_ptr() as usize).checked_add(count).ok_or(
143 VolatileMemoryError::Overflow {
144 base: self.as_mut_ptr() as usize,
145 offset: count,
146 },
147 )?;
148 let new_size = self
149 .size()
150 .checked_sub(count)
151 .ok_or(VolatileMemoryError::OutOfBounds { addr: new_addr })?;
152
153 unsafe { Ok(VolatileSlice::from_raw_parts(new_addr as *mut u8, new_size)) }
157 }
158
159 pub fn sub_slice(self, offset: usize, count: usize) -> Result<VolatileSlice<'a>> {
163 let mem_end = offset
164 .checked_add(count)
165 .ok_or(VolatileMemoryError::Overflow {
166 base: offset,
167 offset: count,
168 })?;
169 if mem_end > self.size() {
170 return Err(Error::OutOfBounds { addr: mem_end });
171 }
172 let new_addr = (self.as_mut_ptr() as usize).checked_add(offset).ok_or(
173 VolatileMemoryError::Overflow {
174 base: self.as_mut_ptr() as usize,
175 offset,
176 },
177 )?;
178
179 Ok(unsafe { VolatileSlice::from_raw_parts(new_addr as *mut u8, count) })
182 }
183
184 pub fn write_bytes(&self, value: u8) {
202 unsafe {
205 write_bytes(self.as_mut_ptr(), value, self.size());
206 }
207 }
208
209 pub fn copy_to<T>(&self, buf: &mut [T])
232 where
233 T: FromBytes + IntoBytes + Copy,
234 {
235 let mut addr = self.as_mut_ptr() as *const u8;
236 for v in buf.iter_mut().take(self.size() / size_of::<T>()) {
237 unsafe {
239 *v = read_volatile(addr as *const T);
240 addr = addr.add(size_of::<T>());
241 }
242 }
243 }
244
245 pub fn copy_to_volatile_slice(&self, slice: VolatileSlice) {
261 unsafe {
263 copy(
264 self.as_mut_ptr() as *const u8,
265 slice.as_mut_ptr(),
266 min(self.size(), slice.size()),
267 );
268 }
269 }
270
271 pub fn copy_from<T>(&self, buf: &[T])
297 where
298 T: IntoBytes + Copy,
299 {
300 let mut addr = self.as_mut_ptr();
301 for v in buf.iter().take(self.size() / size_of::<T>()) {
302 unsafe {
304 write_volatile(addr as *mut T, *v);
305 addr = addr.add(size_of::<T>());
306 }
307 }
308 }
309
310 pub fn is_all_zero(&self) -> bool {
316 const MASK_4BIT: usize = 0x0f;
317 let head_addr = self.as_ptr() as usize;
318 let aligned_head_addr = (head_addr + MASK_4BIT) & !MASK_4BIT;
320 let tail_addr = head_addr + self.size();
321 let aligned_tail_addr = tail_addr & !MASK_4BIT;
323
324 if aligned_head_addr > aligned_tail_addr {
325 return unsafe { is_all_zero_naive(head_addr, tail_addr) };
329 }
330
331 if (aligned_head_addr..aligned_tail_addr).step_by(16).any(
333 |aligned_addr|
334 unsafe { *(aligned_addr as *const u128) } != 0,
336 ) {
337 return false;
338 }
339
340 if head_addr == aligned_head_addr && tail_addr == aligned_tail_addr {
341 true
344 } else {
345 unsafe {
349 is_all_zero_naive(head_addr, aligned_head_addr)
350 && is_all_zero_naive(aligned_tail_addr, tail_addr)
351 }
352 }
353 }
354}
355
356unsafe fn is_all_zero_naive(head_addr: usize, tail_addr: usize) -> bool {
365 (head_addr..tail_addr).all(|addr| *(addr as *const u8) == 0)
366}
367
368impl VolatileMemory for VolatileSlice<'_> {
369 fn get_slice(&self, offset: usize, count: usize) -> Result<VolatileSlice> {
370 self.sub_slice(offset, count)
371 }
372}
373
374impl PartialEq<VolatileSlice<'_>> for VolatileSlice<'_> {
375 fn eq(&self, other: &VolatileSlice) -> bool {
376 let size = self.size();
377 if size != other.size() {
378 return false;
379 }
380
381 let cmp = unsafe { libc::memcmp(self.as_ptr() as _, other.as_ptr() as _, size) };
383
384 cmp == 0
385 }
386}
387
388impl Eq for VolatileSlice<'_> {}
390
391impl std::io::Write for VolatileSlice<'_> {
392 fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
393 let len = buf.len().min(self.size());
394 self.copy_from(&buf[..len]);
395 self.advance(len);
396 Ok(len)
397 }
398
399 fn flush(&mut self) -> std::io::Result<()> {
400 Ok(())
401 }
402}
403
404#[cfg(test)]
405mod tests {
406 use std::io::Write;
407 use std::sync::Arc;
408 use std::sync::Barrier;
409 use std::thread::spawn;
410
411 use super::*;
412
413 #[derive(Clone)]
414 struct VecMem {
415 mem: Arc<Vec<u8>>,
416 }
417
418 impl VecMem {
419 fn new(size: usize) -> VecMem {
420 VecMem {
421 mem: Arc::new(vec![0u8; size]),
422 }
423 }
424 }
425
426 impl VolatileMemory for VecMem {
427 fn get_slice(&self, offset: usize, count: usize) -> Result<VolatileSlice> {
428 let mem_end = offset
429 .checked_add(count)
430 .ok_or(VolatileMemoryError::Overflow {
431 base: offset,
432 offset: count,
433 })?;
434 if mem_end > self.mem.len() {
435 return Err(Error::OutOfBounds { addr: mem_end });
436 }
437
438 let new_addr = (self.mem.as_ptr() as usize).checked_add(offset).ok_or(
439 VolatileMemoryError::Overflow {
440 base: self.mem.as_ptr() as usize,
441 offset,
442 },
443 )?;
444
445 Ok(
446 unsafe { VolatileSlice::from_raw_parts(new_addr as *mut u8, count) },
448 )
449 }
450 }
451
452 #[test]
453 fn observe_mutate() {
454 let a = VecMem::new(1);
455 let a_clone = a.clone();
456 a.get_slice(0, 1).unwrap().write_bytes(99);
457
458 let start_barrier = Arc::new(Barrier::new(2));
459 let thread_start_barrier = start_barrier.clone();
460 let end_barrier = Arc::new(Barrier::new(2));
461 let thread_end_barrier = end_barrier.clone();
462 spawn(move || {
463 thread_start_barrier.wait();
464 a_clone.get_slice(0, 1).unwrap().write_bytes(0);
465 thread_end_barrier.wait();
466 });
467
468 let mut byte = [0u8; 1];
469 a.get_slice(0, 1).unwrap().copy_to(&mut byte);
470 assert_eq!(byte[0], 99);
471
472 start_barrier.wait();
473 end_barrier.wait();
474
475 a.get_slice(0, 1).unwrap().copy_to(&mut byte);
476 assert_eq!(byte[0], 0);
477 }
478
479 #[test]
480 fn slice_size() {
481 let a = VecMem::new(100);
482 let s = a.get_slice(0, 27).unwrap();
483 assert_eq!(s.size(), 27);
484
485 let s = a.get_slice(34, 27).unwrap();
486 assert_eq!(s.size(), 27);
487
488 let s = s.get_slice(20, 5).unwrap();
489 assert_eq!(s.size(), 5);
490 }
491
492 #[test]
493 fn slice_overflow_error() {
494 let a = VecMem::new(1);
495 let res = a.get_slice(usize::MAX, 1).unwrap_err();
496 assert_eq!(
497 res,
498 Error::Overflow {
499 base: usize::MAX,
500 offset: 1,
501 }
502 );
503 }
504
505 #[test]
506 fn slice_oob_error() {
507 let a = VecMem::new(100);
508 a.get_slice(50, 50).unwrap();
509 let res = a.get_slice(55, 50).unwrap_err();
510 assert_eq!(res, Error::OutOfBounds { addr: 105 });
511 }
512
513 #[test]
514 fn is_all_zero_16bytes_aligned() {
515 let a = VecMem::new(1024);
516 let slice = a.get_slice(0, 1024).unwrap();
517
518 assert!(slice.is_all_zero());
519 a.get_slice(129, 1).unwrap().write_bytes(1);
520 assert!(!slice.is_all_zero());
521 }
522
523 #[test]
524 fn is_all_zero_head_not_aligned() {
525 let a = VecMem::new(1024);
526 let slice = a.get_slice(1, 1023).unwrap();
527
528 assert!(slice.is_all_zero());
529 a.get_slice(0, 1).unwrap().write_bytes(1);
530 assert!(slice.is_all_zero());
531 a.get_slice(1, 1).unwrap().write_bytes(1);
532 assert!(!slice.is_all_zero());
533 a.get_slice(1, 1).unwrap().write_bytes(0);
534 a.get_slice(129, 1).unwrap().write_bytes(1);
535 assert!(!slice.is_all_zero());
536 }
537
538 #[test]
539 fn is_all_zero_tail_not_aligned() {
540 let a = VecMem::new(1024);
541 let slice = a.get_slice(0, 1023).unwrap();
542
543 assert!(slice.is_all_zero());
544 a.get_slice(1023, 1).unwrap().write_bytes(1);
545 assert!(slice.is_all_zero());
546 a.get_slice(1022, 1).unwrap().write_bytes(1);
547 assert!(!slice.is_all_zero());
548 a.get_slice(1022, 1).unwrap().write_bytes(0);
549 a.get_slice(0, 1).unwrap().write_bytes(1);
550 assert!(!slice.is_all_zero());
551 }
552
553 #[test]
554 fn is_all_zero_no_aligned_16bytes() {
555 let a = VecMem::new(1024);
556 let slice = a.get_slice(1, 16).unwrap();
557
558 assert!(slice.is_all_zero());
559 a.get_slice(0, 1).unwrap().write_bytes(1);
560 assert!(slice.is_all_zero());
561 for i in 1..17 {
562 a.get_slice(i, 1).unwrap().write_bytes(1);
563 assert!(!slice.is_all_zero());
564 a.get_slice(i, 1).unwrap().write_bytes(0);
565 }
566 a.get_slice(17, 1).unwrap().write_bytes(1);
567 assert!(slice.is_all_zero());
568 }
569
570 #[test]
571 fn write_partial() {
572 let mem = VecMem::new(1024);
573 let mut slice = mem.get_slice(1, 16).unwrap();
574 slice.write_bytes(0xCC);
575
576 let write_len = slice.write(&[1, 2, 3, 4]).unwrap();
578 assert_eq!(write_len, 4);
579 assert_eq!(slice.size(), 16 - 4);
580
581 assert_eq!(mem.mem[1..=4], [1, 2, 3, 4]);
583
584 assert_eq!(mem.mem[5], 0xCC);
586 }
587
588 #[test]
589 fn write_multiple() {
590 let mem = VecMem::new(1024);
591 let mut slice = mem.get_slice(1, 16).unwrap();
592 slice.write_bytes(0xCC);
593
594 let write_len = slice.write(&[1, 2, 3, 4]).unwrap();
596 assert_eq!(write_len, 4);
597 assert_eq!(slice.size(), 16 - 4);
598
599 assert_eq!(mem.mem[5], 0xCC);
601
602 let write2_len = slice.write(&[5, 6, 7, 8]).unwrap();
604 assert_eq!(write2_len, 4);
605 assert_eq!(slice.size(), 16 - 4 - 4);
606
607 assert_eq!(mem.mem[1..=8], [1, 2, 3, 4, 5, 6, 7, 8]);
609
610 assert_eq!(mem.mem[9], 0xCC);
612 }
613
614 #[test]
615 fn write_exact_slice_size() {
616 let mem = VecMem::new(1024);
617 let mut slice = mem.get_slice(1, 4).unwrap();
618 slice.write_bytes(0xCC);
619
620 let write_len = slice.write(&[1, 2, 3, 4]).unwrap();
622 assert_eq!(write_len, 4);
623 assert_eq!(slice.size(), 0);
624
625 assert_eq!(mem.mem[1..=4], [1, 2, 3, 4]);
627
628 assert_eq!(mem.mem[5], 0);
630 }
631
632 #[test]
633 fn write_more_than_slice_size() {
634 let mem = VecMem::new(1024);
635 let mut slice = mem.get_slice(1, 4).unwrap();
636 slice.write_bytes(0xCC);
637
638 let write_len = slice.write(&[1, 2, 3, 4, 5]).unwrap();
640 assert_eq!(write_len, 4);
641 assert_eq!(slice.size(), 0);
642
643 assert_eq!(mem.mem[1..=4], [1, 2, 3, 4]);
645
646 assert_eq!(mem.mem[5], 0);
648 }
649
650 #[test]
651 fn write_empty_slice() {
652 let mem = VecMem::new(1024);
653 let mut slice = mem.get_slice(1, 0).unwrap();
654
655 assert_eq!(slice.write(&[1, 2, 3, 4]).unwrap(), 0);
657 assert_eq!(slice.write(&[5, 6, 7, 8]).unwrap(), 0);
658 assert_eq!(slice.write(&[]).unwrap(), 0);
659 }
660
661 #[test]
662 fn is_all_zero_sub_16bytes_unaligned() {
663 let a = VecMem::new(64);
664 a.get_slice(0, 64).unwrap().write_bytes(0xFF);
666 let slice = a.get_slice(3, 4).unwrap();
667 slice.write_bytes(0);
668 assert!(slice.is_all_zero());
669
670 let empty_slice = a.get_slice(3, 0).unwrap();
672 assert!(empty_slice.is_all_zero());
673 }
674}
675
676#[cfg(kani)]
677mod kani_proofs {
678 use super::*;
679
680 #[repr(align(16))]
681 struct AlignedBuf([u8; 36]);
682
683 #[kani::proof]
688 #[kani::unwind(38)]
689 fn proof_volatile_slice_is_all_zero() {
690 let mut aligned = AlignedBuf(kani::any());
691 let offset: usize = kani::any();
692 let len: usize = kani::any();
693 kani::assume(offset <= 36);
694 kani::assume(len <= 36 - offset);
695
696 let expected = aligned.0[offset..offset + len].iter().all(|&b| b == 0);
697 let slice = VolatileSlice::new(&mut aligned.0[offset..offset + len]);
698 assert_eq!(slice.is_all_zero(), expected);
699 }
700
701 #[kani::proof]
704 #[kani::unwind(16)]
705 fn proof_volatile_slice_small_allocation_memory_safety() {
706 let mut small_buf: [u8; 4] = kani::any();
707 let expected = small_buf.iter().all(|&b| b == 0);
708 let slice = VolatileSlice::new(&mut small_buf);
709 assert_eq!(slice.is_all_zero(), expected);
710 }
711
712 #[kani::proof]
715 fn proof_volatile_slice_sub_slice_and_offset() {
716 let mut buf = [0u8; 64];
717 let total_len: usize = kani::any();
718 kani::assume(total_len <= 64);
719
720 let base_slice = VolatileSlice::new(&mut buf[..total_len]);
721 let offset: usize = kani::any();
722 let count: usize = kani::any();
723
724 match base_slice.sub_slice(offset, count) {
725 Ok(sub) => {
726 assert_eq!(sub.size(), count);
727 assert!(offset
728 .checked_add(count)
729 .is_some_and(|end| end <= total_len));
730 assert_eq!(
731 sub.as_ptr() as usize,
732 (base_slice.as_ptr() as usize) + offset
733 );
734 }
735 Err(VolatileMemoryError::Overflow { .. }) => {
736 assert!(offset.checked_add(count).is_none());
737 }
738 Err(VolatileMemoryError::OutOfBounds { .. }) => {
739 assert!(offset.checked_add(count).is_some_and(|end| end > total_len));
740 }
741 }
742
743 match base_slice.offset(offset) {
744 Ok(off_slice) => {
745 assert!(offset <= total_len);
746 assert_eq!(off_slice.size(), total_len - offset);
747 assert_eq!(
748 off_slice.as_ptr() as usize,
749 (base_slice.as_ptr() as usize) + offset
750 );
751 }
752 Err(_) => {
753 assert!(
754 offset > total_len
755 || (base_slice.as_ptr() as usize).checked_add(offset).is_none()
756 );
757 }
758 }
759 }
760}