1use rand::{RngExt, SeedableRng};
19use rand_chacha::ChaCha8Rng;
20use std::io;
21
22pub const SECTOR_SIZE: usize = 512;
27
28const MAX_OVERLAYS: usize = 2;
34
35#[derive(Debug, Clone)]
40pub struct SectorBitSet {
41 bits: Vec<u64>,
42 len: usize,
43}
44
45impl SectorBitSet {
46 #[must_use]
50 pub fn new(num_sectors: usize) -> Self {
51 let num_words = num_sectors.div_ceil(64);
52 Self {
53 bits: vec![0; num_words],
54 len: num_sectors,
55 }
56 }
57
58 fn indices(&self, sector: usize) -> (usize, usize) {
64 assert!(sector < self.len, "sector index out of bounds");
65 (sector / 64, sector % 64)
66 }
67
68 pub fn set(&mut self, sector: usize) {
74 let (word, bit) = self.indices(sector);
75 self.bits[word] |= 1 << bit;
76 }
77
78 pub fn clear(&mut self, sector: usize) {
84 let (word, bit) = self.indices(sector);
85 self.bits[word] &= !(1 << bit);
86 }
87
88 #[must_use]
94 pub fn is_set(&self, sector: usize) -> bool {
95 let (word, bit) = self.indices(sector);
96 (self.bits[word] & (1 << bit)) != 0
97 }
98
99 #[must_use]
101 pub fn len(&self) -> usize {
102 self.len
103 }
104
105 #[must_use]
107 pub fn is_empty(&self) -> bool {
108 self.len == 0
109 }
110
111 #[must_use]
115 pub fn resize_copy(other: &Self, new_len: usize) -> Self {
116 let mut new_bitset = Self::new(new_len);
117 let copy_len = other.len.min(new_len);
118 for sector in 0..copy_len {
119 if other.is_set(sector) {
120 new_bitset.set(sector);
121 }
122 }
123 new_bitset
124 }
125}
126
127#[derive(Debug, Clone)]
132struct WriteOverlay {
133 offset: u64,
135 size: u32,
137 data: Vec<u8>,
139 active: bool,
141}
142
143#[derive(Debug, Clone)]
148struct PendingWrite {
149 offset: u64,
151 data: Vec<u8>,
153 is_phantom: bool,
155}
156
157#[derive(Debug)]
183pub struct InMemoryStorage {
184 data: Vec<u8>,
186 written: SectorBitSet,
188 faults: SectorBitSet,
190 overlays: [Option<WriteOverlay>; MAX_OVERLAYS],
192 pending_writes: Vec<PendingWrite>,
194 size: u64,
196 seed: u64,
198}
199
200impl InMemoryStorage {
201 #[must_use]
212 pub fn new(size: u64, seed: u64) -> Self {
213 let size_usize = usize::try_from(size).expect("storage size fits in usize");
214 let num_sectors = size_usize.div_ceil(SECTOR_SIZE);
215 Self {
216 data: vec![0; size_usize],
217 written: SectorBitSet::new(num_sectors),
218 faults: SectorBitSet::new(num_sectors),
219 overlays: [const { None }; MAX_OVERLAYS],
220 pending_writes: Vec::new(),
221 size,
222 seed,
223 }
224 }
225
226 #[must_use]
228 pub fn size(&self) -> u64 {
229 self.size
230 }
231
232 pub fn resize(&mut self, new_size: u64) {
241 let old_size = self.size;
242 self.size = new_size;
243
244 let new_size_usize = usize::try_from(new_size).expect("storage size fits in usize");
245
246 self.data.resize(new_size_usize, 0);
248
249 let new_num_sectors = new_size_usize.div_ceil(SECTOR_SIZE);
251 let old_num_sectors = self.written.len();
252
253 if new_num_sectors != old_num_sectors {
254 self.written = SectorBitSet::resize_copy(&self.written, new_num_sectors);
255 self.faults = SectorBitSet::resize_copy(&self.faults, new_num_sectors);
256 }
257
258 tracing::trace!(
260 "InMemoryStorage resized from {} to {} bytes ({} to {} sectors)",
261 old_size,
262 new_size,
263 old_num_sectors,
264 new_num_sectors
265 );
266 }
267
268 pub fn read(&self, offset: u64, buf: &mut [u8]) -> io::Result<()> {
289 let end = offset
291 .checked_add(buf.len() as u64)
292 .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "offset overflow"))?;
293
294 if end > self.size {
295 return Err(io::Error::new(
296 io::ErrorKind::InvalidInput,
297 format!(
298 "read past end of storage: offset={}, len={}, size={}",
299 offset,
300 buf.len(),
301 self.size
302 ),
303 ));
304 }
305
306 let offset_usize = usize::try_from(offset).expect("offset fits in usize");
308 buf.copy_from_slice(&self.data[offset_usize..offset_usize + buf.len()]);
309
310 let start_sector = offset_usize / SECTOR_SIZE;
312 let end_sector = (offset_usize + buf.len()).div_ceil(SECTOR_SIZE);
313
314 for sector in start_sector..end_sector {
315 if sector >= self.written.len() {
316 break;
317 }
318
319 let sector_start = sector * SECTOR_SIZE;
321 let sector_end = sector_start + SECTOR_SIZE;
322
323 let buf_start = sector_start.saturating_sub(offset_usize);
324 let buf_end = (sector_end.saturating_sub(offset_usize)).min(buf.len());
325
326 if buf_start >= buf_end {
327 continue;
328 }
329
330 let sector_buf = &mut buf[buf_start..buf_end];
331
332 if !self.written.is_set(sector) {
334 self.fill_unwritten_sector(sector, sector_buf, sector_start, offset_usize);
335 }
336
337 if self.faults.is_set(sector) {
339 self.apply_corruption(sector, sector_buf);
340 }
341 }
342
343 self.apply_overlays(offset, buf);
345
346 Ok(())
347 }
348
349 fn fill_unwritten_sector(
351 &self,
352 sector: usize,
353 buf: &mut [u8],
354 sector_start: usize,
355 read_offset: usize,
356 ) {
357 let mut rng = ChaCha8Rng::seed_from_u64(self.seed.wrapping_add(sector as u64));
359
360 let mut sector_data = [0u8; SECTOR_SIZE];
362 rng.fill(&mut sector_data);
363
364 let offset_in_sector = read_offset.saturating_sub(sector_start);
366 let copy_start = offset_in_sector.min(SECTOR_SIZE);
367 let copy_len = buf.len().min(SECTOR_SIZE - copy_start);
368
369 buf[..copy_len].copy_from_slice(§or_data[copy_start..copy_start + copy_len]);
370 }
371
372 fn apply_corruption(&self, sector: usize, buf: &mut [u8]) {
379 if buf.is_empty() {
380 return;
381 }
382
383 let sector_start = sector * SECTOR_SIZE;
385 let mut seed_bytes = [0u8; 8];
386 if sector_start + 8 <= self.data.len() {
387 seed_bytes.copy_from_slice(&self.data[sector_start..sector_start + 8]);
388 }
389 let seed = u64::from_le_bytes(seed_bytes);
390
391 let mut rng = ChaCha8Rng::seed_from_u64(seed);
392 let byte_idx = rng.random_range(0..buf.len());
393 let bit_idx = rng.random_range(0..8u8);
394 buf[byte_idx] ^= 1 << bit_idx;
395 }
396
397 fn apply_overlays(&self, offset: u64, buf: &mut [u8]) {
399 for overlay in self.overlays.iter().flatten() {
400 if !overlay.active {
401 continue;
402 }
403
404 let overlay_end = overlay.offset + u64::from(overlay.size);
406 let read_end = offset + buf.len() as u64;
407
408 if overlay.offset >= read_end || overlay_end <= offset {
409 continue;
410 }
411
412 let intersect_start = overlay.offset.max(offset);
414 let intersect_end = overlay_end.min(read_end);
415
416 let buf_offset =
417 usize::try_from(intersect_start - offset).expect("offset fits in usize");
418 let overlay_offset =
419 usize::try_from(intersect_start - overlay.offset).expect("offset fits in usize");
420 let copy_len = usize::try_from(intersect_end - intersect_start)
421 .expect("overlay length fits in usize");
422
423 buf[buf_offset..buf_offset + copy_len]
424 .copy_from_slice(&overlay.data[overlay_offset..overlay_offset + copy_len]);
425 }
426 }
427
428 pub fn write(&mut self, offset: u64, data: &[u8], is_synced: bool) -> io::Result<()> {
447 let end = offset
449 .checked_add(data.len() as u64)
450 .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "offset overflow"))?;
451
452 if end > self.size {
454 self.resize(end);
455 }
456
457 let offset_usize = usize::try_from(offset).expect("offset fits in usize");
458
459 let start_sector = offset_usize / SECTOR_SIZE;
461 let end_sector = (offset_usize + data.len()).div_ceil(SECTOR_SIZE);
462
463 for sector in start_sector..end_sector {
464 if sector < self.written.len() {
465 self.written.set(sector);
466 self.faults.clear(sector);
467 }
468 }
469
470 self.data[offset_usize..offset_usize + data.len()].copy_from_slice(data);
472
473 if !is_synced {
475 self.pending_writes.push(PendingWrite {
476 offset,
477 data: data.to_vec(),
478 is_phantom: false,
479 });
480 }
481
482 Ok(())
483 }
484
485 pub fn sync(&mut self) {
489 self.pending_writes.clear();
490 }
491
492 pub fn apply_misdirected_write(
514 &mut self,
515 intended_offset: u64,
516 mistaken_offset: u64,
517 data: &[u8],
518 ) -> io::Result<()> {
519 let intended_end = intended_offset
521 .checked_add(data.len() as u64)
522 .ok_or_else(|| {
523 io::Error::new(io::ErrorKind::InvalidInput, "intended offset overflow")
524 })?;
525
526 let mistaken_end = mistaken_offset
527 .checked_add(data.len() as u64)
528 .ok_or_else(|| {
529 io::Error::new(io::ErrorKind::InvalidInput, "mistaken offset overflow")
530 })?;
531
532 let required_size = intended_end.max(mistaken_end);
534 if required_size > self.size {
535 self.resize(required_size);
536 }
537
538 let intended_usize = usize::try_from(intended_offset).expect("offset fits in usize");
540 let old_data = self.data[intended_usize..intended_usize + data.len()].to_vec();
541 let overlay_size = u32::try_from(data.len()).expect("data length fits in u32");
542
543 self.overlays[0] = Some(WriteOverlay {
545 offset: intended_offset,
546 size: overlay_size,
547 data: old_data,
548 active: true,
549 });
550
551 self.overlays[1] = Some(WriteOverlay {
553 offset: mistaken_offset,
554 size: overlay_size,
555 data: data.to_vec(),
556 active: true,
557 });
558
559 self.data[intended_usize..intended_usize + data.len()].copy_from_slice(data);
561
562 let start_sector = intended_usize / SECTOR_SIZE;
564 let end_sector = (intended_usize + data.len()).div_ceil(SECTOR_SIZE);
565 for sector in start_sector..end_sector {
566 if sector < self.written.len() {
567 self.written.set(sector);
568 }
569 }
570
571 Ok(())
572 }
573
574 pub fn record_phantom_write(&mut self, offset: u64, data: &[u8]) {
585 self.pending_writes.push(PendingWrite {
586 offset,
587 data: data.to_vec(),
588 is_phantom: true,
589 });
590 }
592
593 pub fn apply_crash(&mut self, crash_fault_probability: f64) {
608 let mut rng = ChaCha8Rng::seed_from_u64(self.seed);
609
610 for pending in &self.pending_writes {
611 if pending.is_phantom {
612 continue;
614 }
615
616 if rng.random::<f64>() >= crash_fault_probability {
618 continue;
619 }
620
621 let offset_usize =
623 usize::try_from(pending.offset).expect("pending offset fits in usize");
624 let start_sector = offset_usize / SECTOR_SIZE;
625 let end_sector = (offset_usize + pending.data.len()).div_ceil(SECTOR_SIZE);
626
627 if start_sector < end_sector && end_sector <= self.faults.len() {
628 let faulted_sector = rng.random_range(start_sector..end_sector);
629 self.faults.set(faulted_sector);
630 }
631 }
632
633 self.pending_writes.clear();
635 }
636
637 pub fn set_fault(&mut self, sector: usize) {
641 if sector < self.faults.len() {
642 self.faults.set(sector);
643 }
644 }
645
646 #[must_use]
648 pub fn has_fault(&self, sector: usize) -> bool {
649 sector < self.faults.len() && self.faults.is_set(sector)
650 }
651}
652
653#[cfg(test)]
654mod tests {
655 use super::*;
656
657 #[test]
658 fn test_basic_write_read() {
659 let mut storage = InMemoryStorage::new(4096, 42);
660
661 let data = b"Hello, World!";
663 storage.write(0, data, true).expect("write failed");
664
665 let mut buf = vec![0u8; data.len()];
667 storage.read(0, &mut buf).expect("read failed");
668
669 assert_eq!(&buf, data);
670 }
671
672 #[test]
673 fn test_unwritten_sector_deterministic() {
674 let storage1 = InMemoryStorage::new(4096, 42);
675 let storage2 = InMemoryStorage::new(4096, 42);
676
677 let mut buf1 = vec![0u8; SECTOR_SIZE];
679 let mut buf2 = vec![0u8; SECTOR_SIZE];
680
681 storage1.read(0, &mut buf1).expect("read1 failed");
682 storage2.read(0, &mut buf2).expect("read2 failed");
683
684 assert_eq!(buf1, buf2);
686
687 let storage3 = InMemoryStorage::new(4096, 99);
689 let mut buf3 = vec![0u8; SECTOR_SIZE];
690 storage3.read(0, &mut buf3).expect("read3 failed");
691
692 assert_ne!(buf1, buf3);
693 }
694
695 #[test]
696 fn test_fault_corruption() {
697 let mut storage = InMemoryStorage::new(4096, 42);
698
699 let data = vec![0xAA; SECTOR_SIZE];
701 storage.write(0, &data, true).expect("write failed");
702
703 let mut buf_clean = vec![0u8; SECTOR_SIZE];
705 storage.read(0, &mut buf_clean).expect("read failed");
706 assert_eq!(buf_clean, data);
707
708 storage.set_fault(0);
710 let mut buf_faulted = vec![0u8; SECTOR_SIZE];
711 storage.read(0, &mut buf_faulted).expect("read failed");
712
713 assert_ne!(buf_faulted, data);
715
716 let bit_diffs: u32 = buf_clean
718 .iter()
719 .zip(buf_faulted.iter())
720 .map(|(a, b)| (*a ^ *b).count_ones())
721 .sum();
722 assert_eq!(bit_diffs, 1, "Expected exactly one bit flip");
723 }
724
725 #[test]
726 fn test_corruption_determinism() {
727 let mut storage = InMemoryStorage::new(4096, 42);
728
729 let data = vec![0xAA; SECTOR_SIZE];
731 storage.write(0, &data, true).expect("write failed");
732 storage.set_fault(0);
733
734 let mut buf1 = vec![0u8; SECTOR_SIZE];
736 let mut buf2 = vec![0u8; SECTOR_SIZE];
737 storage.read(0, &mut buf1).expect("read1 failed");
738 storage.read(0, &mut buf2).expect("read2 failed");
739
740 assert_eq!(buf1, buf2);
742 }
743
744 #[test]
745 fn test_misdirected_write() {
746 let mut storage = InMemoryStorage::new(4096, 42);
747
748 let original_intended = vec![0x11; SECTOR_SIZE];
750 let original_mistaken = vec![0x22; SECTOR_SIZE];
751 storage
752 .write(0, &original_intended, true)
753 .expect("write1 failed");
754 storage
755 .write(SECTOR_SIZE as u64, &original_mistaken, true)
756 .expect("write2 failed");
757
758 let new_data = vec![0xFF; SECTOR_SIZE];
760 storage
761 .apply_misdirected_write(0, SECTOR_SIZE as u64, &new_data)
762 .expect("misdirect failed");
763
764 let mut buf_intended = vec![0u8; SECTOR_SIZE];
766 storage.read(0, &mut buf_intended).expect("read failed");
767 assert_eq!(buf_intended, original_intended);
768
769 let mut buf_mistaken = vec![0u8; SECTOR_SIZE];
771 storage
772 .read(SECTOR_SIZE as u64, &mut buf_mistaken)
773 .expect("read failed");
774 assert_eq!(buf_mistaken, new_data);
775 }
776
777 #[test]
778 fn test_phantom_write_lost_on_crash() {
779 let mut storage = InMemoryStorage::new(4096, 42);
780
781 let real_data = vec![0x11; SECTOR_SIZE];
783 storage.write(0, &real_data, true).expect("write failed");
784
785 let phantom_data = vec![0xFF; SECTOR_SIZE];
787 storage.record_phantom_write(0, &phantom_data);
788
789 let mut buf = vec![0u8; SECTOR_SIZE];
791 storage.read(0, &mut buf).expect("read failed");
792 assert_eq!(buf, real_data);
793
794 storage.apply_crash(1.0);
796
797 storage.read(0, &mut buf).expect("read failed");
799 assert_eq!(buf, real_data);
800 }
801
802 #[test]
803 fn test_crash_faults_pending_writes() {
804 let mut storage = InMemoryStorage::new(4096, 42);
805
806 let data = vec![0xAA; SECTOR_SIZE];
808 storage.write(0, &data, false).expect("write failed");
809
810 storage.apply_crash(1.0);
812
813 assert!(storage.has_fault(0));
815
816 let mut buf = vec![0u8; SECTOR_SIZE];
818 storage.read(0, &mut buf).expect("read failed");
819 assert_ne!(buf, data);
820 }
821
822 #[test]
823 fn test_sync_clears_pending() {
824 let mut storage = InMemoryStorage::new(4096, 42);
825
826 let data = vec![0xAA; SECTOR_SIZE];
828 storage.write(0, &data, false).expect("write failed");
829
830 storage.sync();
832
833 storage.apply_crash(1.0);
835
836 assert!(!storage.has_fault(0));
838
839 let mut buf = vec![0u8; SECTOR_SIZE];
841 storage.read(0, &mut buf).expect("read failed");
842 assert_eq!(buf, data);
843 }
844
845 #[test]
846 fn test_sector_bitset() {
847 let mut bitset = SectorBitSet::new(100);
848
849 assert!(!bitset.is_set(0));
850 assert!(!bitset.is_set(50));
851 assert!(!bitset.is_set(99));
852
853 bitset.set(0);
854 bitset.set(50);
855 bitset.set(99);
856
857 assert!(bitset.is_set(0));
858 assert!(bitset.is_set(50));
859 assert!(bitset.is_set(99));
860 assert!(!bitset.is_set(1));
861
862 bitset.clear(50);
863 assert!(!bitset.is_set(50));
864
865 assert_eq!(bitset.len(), 100);
866 }
867
868 #[test]
869 fn test_read_past_end() {
870 let storage = InMemoryStorage::new(1024, 42);
871
872 let mut buf = vec![0u8; 100];
873 let result = storage.read(1000, &mut buf);
874
875 assert!(result.is_err());
876 }
877
878 #[test]
879 fn test_write_past_end_auto_extends() {
880 let mut storage = InMemoryStorage::new(1024, 42);
881
882 let data = vec![0xAB; 100];
884 let result = storage.write(1000, &data, true);
885
886 assert!(result.is_ok());
887 assert_eq!(storage.size(), 1100); let mut read_buf = vec![0u8; 100];
891 storage.read(1000, &mut read_buf).expect("read failed");
892 assert_eq!(read_buf, data);
893 }
894
895 #[test]
896 fn test_partial_sector_read() {
897 let mut storage = InMemoryStorage::new(4096, 42);
898
899 let data: Vec<u8> = (0..SECTOR_SIZE)
901 .map(|i| u8::try_from(i % 256).expect("modulo 256 fits in u8"))
902 .collect();
903 storage.write(0, &data, true).expect("write failed");
904
905 let mut buf = vec![0u8; 100];
907 storage.read(50, &mut buf).expect("read failed");
908
909 assert_eq!(buf, &data[50..150]);
910 }
911
912 #[test]
913 fn test_multi_sector_read() {
914 let mut storage = InMemoryStorage::new(4096, 42);
915
916 let data = vec![0xAB; SECTOR_SIZE * 3];
918 storage.write(0, &data, true).expect("write failed");
919
920 let mut buf = vec![0u8; SECTOR_SIZE * 3];
922 storage.read(0, &mut buf).expect("read failed");
923
924 assert_eq!(buf, data);
925 }
926}