1use crate::bitmap::BitmapSlice;
7use crate::volatile_memory::copy_slice_impl::{copy_from_volatile_slice, copy_to_volatile_slice};
8use crate::{VolatileMemoryError, VolatileSlice};
9use std::io::{Cursor, ErrorKind, Stdout};
10use std::os::fd::AsRawFd;
11
12pub trait ReadVolatile {
20 fn read_volatile<B: BitmapSlice>(
25 &mut self,
26 buf: &mut VolatileSlice<B>,
27 ) -> Result<usize, VolatileMemoryError>;
28
29 fn read_exact_volatile<B: BitmapSlice>(
34 &mut self,
35 buf: &mut VolatileSlice<B>,
36 ) -> Result<(), VolatileMemoryError> {
37 let mut partial_buf = buf.offset(0)?;
40
41 while !partial_buf.is_empty() {
42 match self.read_volatile(&mut partial_buf) {
43 Err(VolatileMemoryError::IOError(err)) if err.kind() == ErrorKind::Interrupted => {
44 continue
45 }
46 Ok(0) => {
47 return Err(VolatileMemoryError::IOError(std::io::Error::new(
48 ErrorKind::UnexpectedEof,
49 "failed to fill whole buffer",
50 )))
51 }
52 Ok(bytes_read) => partial_buf = partial_buf.offset(bytes_read)?,
53 Err(err) => return Err(err),
54 }
55 }
56
57 Ok(())
58 }
59}
60
61pub trait WriteVolatile {
69 fn write_volatile<B: BitmapSlice>(
74 &mut self,
75 buf: &VolatileSlice<B>,
76 ) -> Result<usize, VolatileMemoryError>;
77
78 fn write_all_volatile<B: BitmapSlice>(
83 &mut self,
84 buf: &VolatileSlice<B>,
85 ) -> Result<(), VolatileMemoryError> {
86 let mut partial_buf = buf.offset(0)?;
89
90 while !partial_buf.is_empty() {
91 match self.write_volatile(&partial_buf) {
92 Err(VolatileMemoryError::IOError(err)) if err.kind() == ErrorKind::Interrupted => {
93 continue
94 }
95 Ok(0) => {
96 return Err(VolatileMemoryError::IOError(std::io::Error::new(
97 ErrorKind::WriteZero,
98 "failed to write whole buffer",
99 )))
100 }
101 Ok(bytes_written) => partial_buf = partial_buf.offset(bytes_written)?,
102 Err(err) => return Err(err),
103 }
104 }
105
106 Ok(())
107 }
108}
109
110macro_rules! impl_read_write_volatile_for_raw_fd {
116 ($raw_fd_ty:ty) => {
117 impl ReadVolatile for $raw_fd_ty {
118 fn read_volatile<B: BitmapSlice>(
119 &mut self,
120 buf: &mut VolatileSlice<B>,
121 ) -> Result<usize, VolatileMemoryError> {
122 read_volatile_raw_fd(self, buf)
123 }
124 }
125
126 impl WriteVolatile for $raw_fd_ty {
127 fn write_volatile<B: BitmapSlice>(
128 &mut self,
129 buf: &VolatileSlice<B>,
130 ) -> Result<usize, VolatileMemoryError> {
131 write_volatile_raw_fd(self, buf)
132 }
133 }
134 };
135}
136
137impl WriteVolatile for Stdout {
138 fn write_volatile<B: BitmapSlice>(
139 &mut self,
140 buf: &VolatileSlice<B>,
141 ) -> Result<usize, VolatileMemoryError> {
142 write_volatile_raw_fd(self, buf)
143 }
144}
145
146impl_read_write_volatile_for_raw_fd!(std::fs::File);
147impl_read_write_volatile_for_raw_fd!(std::net::TcpStream);
148impl_read_write_volatile_for_raw_fd!(std::os::unix::net::UnixStream);
149impl_read_write_volatile_for_raw_fd!(std::os::fd::OwnedFd);
150impl_read_write_volatile_for_raw_fd!(std::os::fd::BorrowedFd<'_>);
151
152fn read_volatile_raw_fd<Fd: AsRawFd>(
157 raw_fd: &mut Fd,
158 buf: &mut VolatileSlice<impl BitmapSlice>,
159) -> Result<usize, VolatileMemoryError> {
160 let fd = raw_fd.as_raw_fd();
161 let guard = buf.ptr_guard_mut();
162
163 let dst = guard.as_ptr().cast::<libc::c_void>();
164
165 let bytes_read = unsafe { libc::read(fd, dst, buf.len()) };
169
170 if bytes_read < 0 {
171 buf.bitmap().mark_dirty(0, buf.len());
173
174 Err(VolatileMemoryError::IOError(std::io::Error::last_os_error()))
175 } else {
176 let bytes_read = bytes_read.try_into().unwrap();
177 buf.bitmap().mark_dirty(0, bytes_read);
178 Ok(bytes_read)
179 }
180}
181
182fn write_volatile_raw_fd<Fd: AsRawFd>(
187 raw_fd: &mut Fd,
188 buf: &VolatileSlice<impl BitmapSlice>,
189) -> Result<usize, VolatileMemoryError> {
190 let fd = raw_fd.as_raw_fd();
191 let guard = buf.ptr_guard();
192
193 let src = guard.as_ptr().cast::<libc::c_void>();
194
195 let bytes_written = unsafe { libc::write(fd, src, buf.len()) };
199
200 if bytes_written < 0 {
201 Err(VolatileMemoryError::IOError(std::io::Error::last_os_error()))
202 } else {
203 Ok(bytes_written.try_into().unwrap())
204 }
205}
206
207impl WriteVolatile for &mut [u8] {
208 fn write_volatile<B: BitmapSlice>(
209 &mut self,
210 buf: &VolatileSlice<B>,
211 ) -> Result<usize, VolatileMemoryError> {
212 let total = buf.len().min(self.len());
213 let src = buf.subslice(0, total)?;
214
215 let written = unsafe { copy_from_volatile_slice(self.as_mut_ptr(), &src, total) };
223
224 *self = std::mem::take(self).split_at_mut(written).1;
226
227 Ok(written)
228 }
229
230 fn write_all_volatile<B: BitmapSlice>(
231 &mut self,
232 buf: &VolatileSlice<B>,
233 ) -> Result<(), VolatileMemoryError> {
234 if self.write_volatile(buf)? == buf.len() {
236 Ok(())
237 } else {
238 Err(VolatileMemoryError::IOError(std::io::Error::new(
239 ErrorKind::WriteZero,
240 "failed to write whole buffer",
241 )))
242 }
243 }
244}
245
246impl ReadVolatile for &[u8] {
247 fn read_volatile<B: BitmapSlice>(
248 &mut self,
249 buf: &mut VolatileSlice<B>,
250 ) -> Result<usize, VolatileMemoryError> {
251 let total = buf.len().min(self.len());
252 let dst = buf.subslice(0, total)?;
253
254 let read = unsafe { copy_to_volatile_slice(&dst, self.as_ptr(), total) };
261
262 *self = self.split_at(read).1;
264
265 Ok(read)
266 }
267
268 fn read_exact_volatile<B: BitmapSlice>(
269 &mut self,
270 buf: &mut VolatileSlice<B>,
271 ) -> Result<(), VolatileMemoryError> {
272 if buf.len() > self.len() {
274 return Err(VolatileMemoryError::IOError(std::io::Error::new(
275 ErrorKind::UnexpectedEof,
276 "failed to fill whole buffer",
277 )));
278 }
279
280 self.read_volatile(buf).map(|_| ())
281 }
282}
283
284impl WriteVolatile for Vec<u8> {
287 fn write_volatile<B: BitmapSlice>(
288 &mut self,
289 buf: &VolatileSlice<B>,
290 ) -> Result<usize, VolatileMemoryError> {
291 let count = buf.len();
292 self.reserve(count);
293 let len = self.len();
294
295 unsafe {
306 let copied_len = copy_from_volatile_slice(self.as_mut_ptr().add(len), buf, count);
307
308 assert_eq!(copied_len, count);
309 self.set_len(len + count);
310 }
311 Ok(count)
312 }
313}
314
315impl<T> ReadVolatile for Cursor<T>
318where
319 T: AsRef<[u8]>,
320{
321 fn read_volatile<B: BitmapSlice>(
322 &mut self,
323 buf: &mut VolatileSlice<B>,
324 ) -> Result<usize, VolatileMemoryError> {
325 let inner = self.get_ref().as_ref();
326 let len = self.position().min(inner.len() as u64);
327 let n = ReadVolatile::read_volatile(&mut &inner[(len as usize)..], buf)?;
328 self.set_position(self.position() + n as u64);
329 Ok(n)
330 }
331
332 fn read_exact_volatile<B: BitmapSlice>(
333 &mut self,
334 buf: &mut VolatileSlice<B>,
335 ) -> Result<(), VolatileMemoryError> {
336 let inner = self.get_ref().as_ref();
337 let n = buf.len();
338 let len = self.position().min(inner.len() as u64);
339 ReadVolatile::read_exact_volatile(&mut &inner[(len as usize)..], buf)?;
340 self.set_position(self.position() + n as u64);
341 Ok(())
342 }
343}
344
345impl WriteVolatile for Cursor<&mut [u8]> {
346 fn write_volatile<B: BitmapSlice>(
347 &mut self,
348 buf: &VolatileSlice<B>,
349 ) -> Result<usize, VolatileMemoryError> {
350 let pos = self.position().min(self.get_ref().len() as u64);
351 let n = WriteVolatile::write_volatile(&mut &mut self.get_mut()[(pos as usize)..], buf)?;
352 self.set_position(self.position() + n as u64);
353 Ok(n)
354 }
355
356 }
359
360#[cfg(test)]
361mod tests {
362 use crate::io::{ReadVolatile, WriteVolatile};
363 use crate::{VolatileMemoryError, VolatileSlice};
364 use std::io::{Cursor, ErrorKind, Read, Seek, Write};
365 use vmm_sys_util::tempfile::TempFile;
366
367 fn read_4_bytes_to_5_byte_memory(source: Vec<u8>, expected_output: [u8; 5]) {
369 let mut memory = vec![0u8; 5];
371
372 assert_eq!(
373 (&source[..])
374 .read_volatile(&mut VolatileSlice::from(&mut memory[..4]))
375 .unwrap(),
376 source.len().min(4)
377 );
378 assert_eq!(&memory, &expected_output);
379
380 let mut memory = vec![0u8; 5];
382 let result = (&source[..]).read_exact_volatile(&mut VolatileSlice::from(&mut memory[..4]));
383
384 if source.len() < 4 {
387 match result.unwrap_err() {
388 VolatileMemoryError::IOError(ioe) => {
389 assert_eq!(ioe.kind(), ErrorKind::UnexpectedEof)
390 }
391 err => panic!("{:?}", err),
392 }
393 assert_eq!(memory, vec![0u8; 5]);
394 } else {
395 result.unwrap();
396 assert_eq!(&memory, &expected_output);
397 }
398 }
399
400 fn read_4_bytes_from_file(source: Vec<u8>, expected_output: [u8; 5]) {
402 let mut temp_file = TempFile::new().unwrap().into_file();
403 temp_file.write_all(source.as_ref()).unwrap();
404 temp_file.rewind().unwrap();
405
406 let mut memory = vec![0u8; 5];
408
409 assert_eq!(
410 temp_file
411 .read_volatile(&mut VolatileSlice::from(&mut memory[..4]))
412 .unwrap(),
413 source.len().min(4)
414 );
415 assert_eq!(&memory, &expected_output);
416
417 temp_file.rewind().unwrap();
418
419 let mut memory = vec![0u8; 5];
421
422 let read_exact_result =
423 temp_file.read_exact_volatile(&mut VolatileSlice::from(&mut memory[..4]));
424
425 if source.len() < 4 {
426 read_exact_result.unwrap_err();
427 } else {
428 read_exact_result.unwrap();
429 }
430 assert_eq!(&memory, &expected_output);
431 }
432
433 #[test]
434 fn test_read_volatile() {
435 let test_cases = [
436 (vec![1u8, 2], [1u8, 2, 0, 0, 0]),
437 (vec![1, 2, 3, 4], [1, 2, 3, 4, 0]),
438 (vec![5, 6, 7, 8, 9], [5, 6, 7, 8, 0]),
440 ];
441
442 for (input, output) in test_cases {
443 read_4_bytes_to_5_byte_memory(input.clone(), output);
444 read_4_bytes_from_file(input, output);
445 }
446 }
447
448 fn write_4_bytes_to_5_byte_vec(mut source: Vec<u8>, expected_result: [u8; 5]) {
450 let mut memory = vec![0u8; 5];
451
452 assert_eq!(
454 (&mut memory[..4])
455 .write_volatile(&VolatileSlice::from(source.as_mut_slice()))
456 .unwrap(),
457 source.len().min(4)
458 );
459 assert_eq!(&memory, &expected_result);
460
461 let mut memory = vec![0u8; 5];
463
464 let result =
465 (&mut memory[..4]).write_all_volatile(&VolatileSlice::from(source.as_mut_slice()));
466
467 if source.len() > 4 {
468 match result.unwrap_err() {
469 VolatileMemoryError::IOError(ioe) => {
470 assert_eq!(ioe.kind(), ErrorKind::WriteZero)
471 }
472 err => panic!("{:?}", err),
473 }
474 assert_eq!(&memory, &expected_result);
477 } else {
478 result.unwrap();
479 assert_eq!(&memory, &expected_result);
480 }
481 }
482
483 fn write_5_bytes_to_file(mut source: Vec<u8>) {
485 let mut temp_file = TempFile::new().unwrap().into_file();
487
488 temp_file
489 .write_volatile(&VolatileSlice::from(source.as_mut_slice()))
490 .unwrap();
491 temp_file.rewind().unwrap();
492
493 let mut written = vec![0u8; source.len()];
494 temp_file.read_exact(written.as_mut_slice()).unwrap();
495
496 assert_eq!(source, written);
497 assert_eq!(temp_file.read(&mut [0u8]).unwrap(), 0);
499
500 let mut temp_file = TempFile::new().unwrap().into_file();
502
503 temp_file
504 .write_all_volatile(&VolatileSlice::from(source.as_mut_slice()))
505 .unwrap();
506 temp_file.rewind().unwrap();
507
508 let mut written = vec![0u8; source.len()];
509 temp_file.read_exact(written.as_mut_slice()).unwrap();
510
511 assert_eq!(source, written);
512 assert_eq!(temp_file.read(&mut [0u8]).unwrap(), 0);
514 }
515
516 #[test]
517 fn test_write_volatile() {
518 let test_cases = [
519 (vec![1u8, 2], [1u8, 2, 0, 0, 0]),
520 (vec![1, 2, 3, 4], [1, 2, 3, 4, 0]),
521 (vec![5, 6, 7, 8, 9], [5, 6, 7, 8, 0]),
523 ];
524
525 for (input, output) in test_cases {
526 write_4_bytes_to_5_byte_vec(input.clone(), output);
527 write_5_bytes_to_file(input);
528 }
529 }
530
531 #[test]
532 fn test_read_volatile_for_cursor() {
533 let read_buffer = [1, 2, 3, 4, 5, 6, 7];
534 let mut output = vec![0u8; 5];
535
536 let mut cursor = Cursor::new(read_buffer);
537
538 assert_eq!(
540 cursor
541 .read_volatile(&mut VolatileSlice::from(&mut output[..4]))
542 .unwrap(),
543 4
544 );
545 assert_eq!(output, vec![1, 2, 3, 4, 0]);
546
547 assert_eq!(
549 cursor
550 .read_volatile(&mut VolatileSlice::from(&mut output[..4]))
551 .unwrap(),
552 3
553 );
554 assert_eq!(output, vec![5, 6, 7, 4, 0]);
555
556 cursor.set_position(0);
557 cursor
559 .read_exact_volatile(&mut VolatileSlice::from(&mut output[..4]))
560 .unwrap();
561 assert_eq!(output, vec![1, 2, 3, 4, 0]);
562
563 assert!(cursor
566 .read_exact_volatile(&mut VolatileSlice::from(&mut output[..4]))
567 .is_err());
568 assert_eq!(output, vec![1, 2, 3, 4, 0]);
569 }
570
571 #[test]
572 fn test_write_volatile_for_cursor() {
573 let mut write_buffer = vec![0u8; 7];
574 let mut input = [1, 2, 3, 4];
575
576 let mut cursor = Cursor::new(write_buffer.as_mut_slice());
577
578 assert_eq!(
580 cursor
581 .write_volatile(&VolatileSlice::from(input.as_mut_slice()))
582 .unwrap(),
583 4
584 );
585 assert_eq!(cursor.get_ref(), &[1, 2, 3, 4, 0, 0, 0]);
586
587 assert_eq!(
589 cursor
590 .write_volatile(&VolatileSlice::from(input.as_mut_slice()))
591 .unwrap(),
592 3
593 );
594 assert_eq!(cursor.get_ref(), &[1, 2, 3, 4, 1, 2, 3]);
595 }
596
597 #[test]
598 fn test_write_volatile_for_vec() {
599 let mut write_buffer = Vec::new();
600 let mut input = [1, 2, 3, 4];
601
602 assert_eq!(
603 write_buffer
604 .write_volatile(&VolatileSlice::from(input.as_mut_slice()))
605 .unwrap(),
606 4
607 );
608
609 assert_eq!(&write_buffer, &input);
610 }
611}