1use rustc_hash::FxHashMap;
22
23use crate::tlb::TranslationCache;
24
25pub const PAGE_SIZE: u64 = 0x1000;
28const PAGE_MASK: u64 = PAGE_SIZE - 1;
29
30pub type Perm = u8;
32
33pub mod perm {
38 use super::Perm;
39
40 pub const NONE: Perm = 0;
41 pub const INIT: Perm = 1 << 0;
45 pub const READ: Perm = 1 << 1;
46 pub const WRITE: Perm = 1 << 2;
47 pub const EXEC: Perm = 1 << 3;
48 pub const MAP: Perm = 1 << 4;
50 pub const READ_WATCH: Perm = 1 << 5;
51 pub const WRITE_WATCH: Perm = 1 << 6;
52
53 pub const READ_WRITE: Perm = READ | WRITE;
54 pub const RW_INIT: Perm = MAP | READ | WRITE | INIT;
56 pub const RX_INIT: Perm = MAP | READ | EXEC | INIT;
58}
59
60#[derive(Debug, Clone, Copy, PartialEq, Eq)]
66pub enum FaultKind {
67 ReadUnmapped,
68 ReadPerm,
69 ReadUninit,
70 WriteUnmapped,
71 WritePerm,
72 ExecUnmapped,
73 ExecViolation,
74 ReadWatch,
75 WriteWatch,
76 AddressOverflow,
78}
79
80#[derive(Debug, Clone, Copy, PartialEq, Eq)]
84pub struct MemFault {
85 pub kind: FaultKind,
86 pub addr: u64,
87}
88
89impl MemFault {
90 fn new(kind: FaultKind, addr: u64) -> Self {
91 Self { kind, addr }
92 }
93
94 pub fn is_write(&self) -> bool {
100 matches!(
101 self.kind,
102 FaultKind::WriteUnmapped | FaultKind::WritePerm | FaultKind::WriteWatch
103 )
104 }
105}
106
107impl std::fmt::Display for MemFault {
108 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
109 let what = match self.kind {
110 FaultKind::ReadUnmapped => "read of unmapped memory",
111 FaultKind::ReadPerm => "read of unreadable memory",
112 FaultKind::ReadUninit => "read of uninitialized memory",
113 FaultKind::WriteUnmapped => "write to unmapped memory",
114 FaultKind::WritePerm => "write to unwritable memory",
115 FaultKind::ExecUnmapped => "execution of unmapped memory",
116 FaultKind::ExecViolation => "execution of non-executable memory",
117 FaultKind::ReadWatch => "watched read",
118 FaultKind::WriteWatch => "watched write",
119 FaultKind::AddressOverflow => "address space overflow",
120 };
121 write!(f, "{what} at {:#x}", self.addr)
122 }
123}
124
125impl std::error::Error for MemFault {}
126
127#[repr(C)]
135#[derive(Clone)]
136pub struct PageData {
137 pub data: [u8; PAGE_SIZE as usize],
138 pub perm: [Perm; PAGE_SIZE as usize],
139}
140
141pub const PAGE_PERM_OFFSET: usize = PAGE_SIZE as usize;
146
147#[derive(Clone)]
148struct Page {
149 inner: Box<PageData>,
150}
151
152impl Page {
153 fn unmapped() -> Self {
154 let inner = unsafe {
167 let layout = std::alloc::Layout::new::<PageData>();
168 let ptr = std::alloc::alloc_zeroed(layout).cast::<PageData>();
169 if ptr.is_null() {
170 std::alloc::handle_alloc_error(layout);
171 }
172 Box::from_raw(ptr)
173 };
174 Self { inner }
175 }
176
177 fn host_ptr(&mut self) -> *mut u8 {
179 std::ptr::from_mut(&mut *self.inner).cast()
180 }
181}
182
183impl std::ops::Deref for Page {
184 type Target = PageData;
185
186 fn deref(&self) -> &PageData {
187 &self.inner
188 }
189}
190
191impl std::ops::DerefMut for Page {
192 fn deref_mut(&mut self) -> &mut PageData {
193 &mut self.inner
194 }
195}
196
197#[derive(Default)]
202pub struct Mmu {
203 pages: FxHashMap<u64, Page>,
204 check_uninit: bool,
208 watchpoints_armed: bool,
212 tlb: TranslationCache,
217}
218
219impl Clone for Mmu {
225 fn clone(&self) -> Self {
226 Self {
227 pages: self.pages.clone(),
228 check_uninit: self.check_uninit,
229 watchpoints_armed: self.watchpoints_armed,
230 tlb: TranslationCache::default(),
231 }
232 }
233}
234
235fn page_chunks(addr: u64, len: u64) -> Result<impl Iterator<Item = (u64, usize, usize)>, MemFault> {
240 if len != 0 && addr.checked_add(len - 1).is_none() {
241 return Err(MemFault::new(FaultKind::AddressOverflow, addr));
242 }
243 let mut remaining = len;
244 let mut cursor = addr;
245 Ok(std::iter::from_fn(move || {
246 if remaining == 0 {
247 return None;
248 }
249 let offset = cursor & PAGE_MASK;
250 let take = (PAGE_SIZE - offset).min(remaining);
251 let chunk = (cursor >> 12, offset as usize, take as usize);
252 cursor = cursor.wrapping_add(take);
253 remaining -= take;
254 Some(chunk)
255 }))
256}
257
258impl Mmu {
259 pub fn new() -> Self {
260 Self::default()
261 }
262
263 pub fn resident_pages(&self) -> usize {
266 self.pages.len()
267 }
268
269 pub fn check_uninit(&self) -> bool {
270 self.check_uninit
271 }
272
273 pub fn set_check_uninit(&mut self, enabled: bool) {
275 self.check_uninit = enabled;
276 self.tlb.flush();
280 }
281
282 pub fn watchpoints_armed(&self) -> bool {
283 self.watchpoints_armed
284 }
285
286 pub fn set_watchpoints_armed(&mut self, armed: bool) {
288 self.watchpoints_armed = armed;
289 self.tlb.flush();
290 }
291
292 fn caching_allowed(&self) -> bool {
299 !self.check_uninit && !self.watchpoints_armed
300 }
301
302 pub fn tlb_ptr(&mut self) -> *mut u8 {
307 std::ptr::from_mut(&mut self.tlb).cast()
308 }
309
310 pub fn flush_tlb(&mut self) {
313 self.tlb.flush();
314 }
315
316 pub fn cache_translation(&mut self, addr: u64) -> bool {
323 if !self.caching_allowed() {
324 return false;
325 }
326 let Some(page) = self.pages.get_mut(&(addr >> 12)) else {
327 return false;
328 };
329 let host = page.host_ptr();
330 self.tlb.insert(addr, host);
331 true
332 }
333
334 pub fn map(&mut self, addr: u64, len: u64, permissions: Perm) -> Result<(), MemFault> {
340 self.tlb.flush();
341 for (index, offset, take) in page_chunks(addr, len)? {
342 let page = self.pages.entry(index).or_insert_with(Page::unmapped);
343 page.data[offset..offset + take].fill(0);
344 page.perm[offset..offset + take].fill(permissions | perm::MAP);
345 }
346 Ok(())
347 }
348
349 pub fn unmap(&mut self, addr: u64, len: u64) -> Result<(), MemFault> {
354 self.tlb.flush();
355 for (index, offset, take) in page_chunks(addr, len)? {
356 let Some(page) = self.pages.get_mut(&index) else {
357 continue;
358 };
359 page.data[offset..offset + take].fill(0);
360 page.perm[offset..offset + take].fill(perm::NONE);
361 if page.perm.iter().all(|&p| p == perm::NONE) {
362 self.pages.remove(&index);
363 }
364 }
365 Ok(())
366 }
367
368 pub fn protect(&mut self, addr: u64, len: u64, permissions: Perm) -> Result<(), MemFault> {
374 self.tlb.flush();
375 for (index, offset, take) in page_chunks(addr, len)? {
379 let base = index << 12;
380 match self.pages.get(&index) {
381 Some(page) => {
382 for byte in offset..offset + take {
383 if page.perm[byte] & perm::MAP == 0 {
384 let at = base + byte as u64;
385 return Err(MemFault::new(FaultKind::WriteUnmapped, at));
386 }
387 }
388 }
389 None => {
390 return Err(MemFault::new(
391 FaultKind::WriteUnmapped,
392 base + offset as u64,
393 ));
394 }
395 }
396 }
397
398 for (index, offset, take) in page_chunks(addr, len)? {
399 let page = self.pages.get_mut(&index).expect("checked above");
400 for byte in offset..offset + take {
401 let init = page.perm[byte] & perm::INIT;
402 page.perm[byte] = permissions | perm::MAP | init;
403 }
404 }
405 Ok(())
406 }
407
408 pub fn permissions(&self, addr: u64) -> Perm {
410 self.pages
411 .get(&(addr >> 12))
412 .map_or(perm::NONE, |page| page.perm[(addr & PAGE_MASK) as usize])
413 }
414
415 pub fn read(&self, addr: u64, out: &mut [u8]) -> Result<(), MemFault> {
417 self.read_with(addr, out, perm::READ)
418 }
419
420 pub fn read_code(&self, addr: u64, out: &mut [u8]) -> Result<(), MemFault> {
426 self.read_with(addr, out, perm::EXEC)
427 }
428
429 fn read_with(&self, addr: u64, out: &mut [u8], required: Perm) -> Result<(), MemFault> {
430 let executing = required & perm::EXEC != 0;
431 let mut written = 0;
432 for (index, offset, take) in page_chunks(addr, out.len() as u64)? {
433 let base = index << 12;
434 let Some(page) = self.pages.get(&index) else {
435 let kind = if executing {
436 FaultKind::ExecUnmapped
437 } else {
438 FaultKind::ReadUnmapped
439 };
440 return Err(MemFault::new(kind, base + offset as u64));
441 };
442 for byte in offset..offset + take {
443 let at = base + byte as u64;
444 let held = page.perm[byte];
445 if held & perm::MAP == 0 {
446 let kind = if executing {
447 FaultKind::ExecUnmapped
448 } else {
449 FaultKind::ReadUnmapped
450 };
451 return Err(MemFault::new(kind, at));
452 }
453 if held & required == 0 {
454 let kind = if executing {
455 FaultKind::ExecViolation
456 } else {
457 FaultKind::ReadPerm
458 };
459 return Err(MemFault::new(kind, at));
460 }
461 if self.check_uninit && held & perm::INIT == 0 {
462 return Err(MemFault::new(FaultKind::ReadUninit, at));
463 }
464 if self.watchpoints_armed && held & perm::READ_WATCH != 0 {
465 return Err(MemFault::new(FaultKind::ReadWatch, at));
466 }
467 }
468 out[written..written + take].copy_from_slice(&page.data[offset..offset + take]);
469 written += take;
470 }
471 Ok(())
472 }
473
474 pub fn write(&mut self, addr: u64, bytes: &[u8]) -> Result<(), MemFault> {
477 let mut read = 0;
481 for (index, offset, take) in page_chunks(addr, bytes.len() as u64)? {
482 let base = index << 12;
483 let Some(page) = self.pages.get(&index) else {
484 return Err(MemFault::new(
485 FaultKind::WriteUnmapped,
486 base + offset as u64,
487 ));
488 };
489 for byte in offset..offset + take {
490 let at = base + byte as u64;
491 let held = page.perm[byte];
492 if held & perm::MAP == 0 {
493 return Err(MemFault::new(FaultKind::WriteUnmapped, at));
494 }
495 if held & perm::WRITE == 0 {
496 return Err(MemFault::new(FaultKind::WritePerm, at));
497 }
498 if self.watchpoints_armed && held & perm::WRITE_WATCH != 0 {
499 return Err(MemFault::new(FaultKind::WriteWatch, at));
500 }
501 }
502 read += take;
503 }
504 debug_assert_eq!(read, bytes.len());
505
506 let mut written = 0;
507 for (index, offset, take) in page_chunks(addr, bytes.len() as u64)? {
508 let page = self.pages.get_mut(&index).expect("checked above");
509 page.data[offset..offset + take].copy_from_slice(&bytes[written..written + take]);
510 for byte in offset..offset + take {
511 page.perm[byte] |= perm::INIT;
512 }
513 written += take;
514 }
515 Ok(())
516 }
517
518 pub fn write_unchecked(&mut self, addr: u64, bytes: &[u8], permissions: Perm) {
524 self.tlb.flush();
525 let mut written = 0;
526 let chunks = page_chunks(addr, bytes.len() as u64)
527 .expect("write_unchecked range must fit the address space");
528 for (index, offset, take) in chunks {
529 let page = self.pages.entry(index).or_insert_with(Page::unmapped);
530 page.data[offset..offset + take].copy_from_slice(&bytes[written..written + take]);
531 page.perm[offset..offset + take].fill(permissions | perm::MAP | perm::INIT);
532 written += take;
533 }
534 }
535
536 pub fn snapshot(&self) -> MmuSnapshot {
542 MmuSnapshot {
543 pages: self.pages.clone(),
544 }
545 }
546
547 pub fn restore(&mut self, snapshot: &MmuSnapshot) {
549 self.pages.clone_from(&snapshot.pages);
550 self.tlb.flush();
553 }
554}
555
556#[derive(Clone)]
558pub struct MmuSnapshot {
559 pages: FxHashMap<u64, Page>,
560}
561
562#[cfg(test)]
563mod tests {
564 use super::*;
565
566 fn mapped() -> Mmu {
567 let mut mmu = Mmu::new();
568 mmu.map(0x1000, 0x2000, perm::RW_INIT).unwrap();
569 mmu
570 }
571
572 #[test]
573 fn the_permission_array_follows_the_data_array() {
574 let page = Page::unmapped();
577 let data = std::ptr::from_ref(&page.data) as usize;
578 let perm = std::ptr::from_ref(&page.perm) as usize;
579 assert_eq!(perm - data, PAGE_PERM_OFFSET);
580 assert_eq!(std::mem::size_of::<PageData>(), 2 * PAGE_SIZE as usize);
581 }
582
583 #[test]
584 fn a_translation_is_cached_only_for_a_resident_page() {
585 let mut mmu = mapped();
586 assert!(!mmu.cache_translation(0x9000), "unmapped page");
587 assert!(mmu.cache_translation(0x1abc));
588 mmu.write(0x1abc, &[0x5a]).unwrap();
590 let host = mmu.tlb.lookup(0x1abc).expect("just cached");
591 assert_eq!(unsafe { *host }, 0x5a);
592 assert_eq!(unsafe { *host.add(PAGE_PERM_OFFSET) }, perm::RW_INIT);
594 }
595
596 #[test]
597 fn a_dynamic_check_empties_the_cache_and_keeps_it_empty() {
598 let mut mmu = mapped();
599 assert!(mmu.cache_translation(0x1000));
600 mmu.set_check_uninit(true);
601 assert!(mmu.tlb.lookup(0x1000).is_none());
602 assert!(!mmu.cache_translation(0x1000));
605 mmu.set_check_uninit(false);
606 assert!(mmu.cache_translation(0x1000));
607 }
608
609 #[test]
610 fn unmapping_a_page_drops_its_cached_translation() {
611 let mut mmu = mapped();
612 assert!(mmu.cache_translation(0x1000));
613 mmu.unmap(0x1000, PAGE_SIZE).unwrap();
616 assert!(mmu.tlb.lookup(0x1000).is_none());
617 }
618
619 #[test]
620 fn a_cloned_mmu_starts_with_an_empty_cache() {
621 let mut mmu = mapped();
622 assert!(mmu.cache_translation(0x1000));
623 let clone = mmu.clone();
624 assert!(clone.tlb.lookup(0x1000).is_none());
625 }
626
627 #[test]
628 fn maps_reads_and_writes() {
629 let mut mmu = mapped();
630 mmu.write(0x1004, &[1, 2, 3, 4]).unwrap();
631 let mut out = [0; 4];
632 mmu.read(0x1004, &mut out).unwrap();
633 assert_eq!(out, [1, 2, 3, 4]);
634 }
635
636 #[test]
637 fn unmapped_access_faults_at_the_offending_byte() {
638 let mmu = mapped();
639 let mut out = [0; 4];
640 assert_eq!(
643 mmu.read(0x2ffe, &mut out).unwrap_err(),
644 MemFault::new(FaultKind::ReadUnmapped, 0x3000)
645 );
646 }
647
648 #[test]
649 fn access_spanning_pages_is_contiguous() {
650 let mut mmu = mapped();
651 let bytes: Vec<u8> = (0..16).collect();
652 mmu.write(0x1ff8, &bytes).unwrap();
653 let mut out = [0; 16];
654 mmu.read(0x1ff8, &mut out).unwrap();
655 assert_eq!(out.to_vec(), bytes);
656 }
657
658 #[test]
659 fn write_to_read_only_memory_faults_and_changes_nothing() {
660 let mut mmu = Mmu::new();
661 mmu.map(0x1000, PAGE_SIZE, perm::RX_INIT).unwrap();
662 assert_eq!(
663 mmu.write(0x1000, &[0xff]).unwrap_err(),
664 MemFault::new(FaultKind::WritePerm, 0x1000)
665 );
666 let mut out = [0xaa];
667 mmu.read(0x1000, &mut out).unwrap();
668 assert_eq!(out, [0]);
669 }
670
671 #[test]
672 fn partially_refused_write_is_not_applied() {
673 let mut mmu = Mmu::new();
674 mmu.map(0x1000, PAGE_SIZE, perm::RW_INIT).unwrap();
675 mmu.map(0x2000, PAGE_SIZE, perm::RX_INIT).unwrap();
676 assert_eq!(
678 mmu.write(0x1ffc, &[0xff; 8]).unwrap_err(),
679 MemFault::new(FaultKind::WritePerm, 0x2000)
680 );
681 let mut out = [0xaa; 4];
682 mmu.read(0x1ffc, &mut out).unwrap();
683 assert_eq!(out, [0; 4]);
684 }
685
686 #[test]
687 fn fetching_from_non_executable_memory_faults() {
688 let mut mmu = mapped();
689 let mut out = [0; 4];
690 assert_eq!(
691 mmu.read_code(0x1000, &mut out).unwrap_err(),
692 MemFault::new(FaultKind::ExecViolation, 0x1000)
693 );
694 mmu.protect(0x1000, PAGE_SIZE, perm::READ | perm::EXEC)
695 .unwrap();
696 mmu.read_code(0x1000, &mut out).unwrap();
697 }
698
699 #[test]
700 fn protect_preserves_contents_and_initializedness() {
701 let mut mmu = mapped();
702 mmu.write(0x1000, &[7; 4]).unwrap();
703 mmu.set_check_uninit(true);
704 mmu.protect(0x1000, PAGE_SIZE, perm::READ).unwrap();
705 let mut out = [0; 4];
706 mmu.read(0x1000, &mut out).unwrap();
707 assert_eq!(out, [7; 4]);
708 }
709
710 #[test]
711 fn protect_of_unmapped_memory_is_refused_entirely() {
712 let mut mmu = mapped();
713 assert!(mmu.protect(0x2000, 0x2000, perm::READ).is_err());
714 assert_eq!(mmu.permissions(0x2000), perm::RW_INIT);
716 }
717
718 #[test]
719 fn uninitialized_reads_fault_only_when_checked() {
720 let mut mmu = Mmu::new();
721 mmu.map(0x1000, PAGE_SIZE, perm::MAP | perm::READ_WRITE)
722 .unwrap();
723 let mut out = [0; 1];
724 mmu.read(0x1000, &mut out).unwrap();
725
726 mmu.set_check_uninit(true);
727 assert_eq!(
728 mmu.read(0x1000, &mut out).unwrap_err(),
729 MemFault::new(FaultKind::ReadUninit, 0x1000)
730 );
731 mmu.write(0x1000, &[1]).unwrap();
733 mmu.read(0x1000, &mut out).unwrap();
734 }
735
736 #[test]
737 fn watchpoints_fire_only_when_armed() {
738 let mut mmu = Mmu::new();
739 mmu.map(0x1000, PAGE_SIZE, perm::RW_INIT | perm::WRITE_WATCH)
740 .unwrap();
741 mmu.write(0x1000, &[1]).unwrap();
742
743 mmu.set_watchpoints_armed(true);
744 assert_eq!(
745 mmu.write(0x1000, &[2]).unwrap_err(),
746 MemFault::new(FaultKind::WriteWatch, 0x1000)
747 );
748 let mut out = [0; 1];
750 mmu.read(0x1000, &mut out).unwrap();
751 assert_eq!(out, [1]);
752 }
753
754 #[test]
755 fn unmap_releases_pages_and_faults_afterwards() {
756 let mut mmu = mapped();
757 assert_eq!(mmu.resident_pages(), 2);
758 mmu.unmap(0x1000, 0x2000).unwrap();
759 assert_eq!(mmu.resident_pages(), 0);
760 let mut out = [0; 1];
761 assert_eq!(
762 mmu.read(0x1000, &mut out).unwrap_err(),
763 MemFault::new(FaultKind::ReadUnmapped, 0x1000)
764 );
765 }
766
767 #[test]
768 fn partial_unmap_keeps_the_rest_of_the_page() {
769 let mut mmu = mapped();
770 mmu.write(0x1000, &[9; 8]).unwrap();
771 mmu.unmap(0x1000, 4).unwrap();
772 assert_eq!(mmu.resident_pages(), 2);
773 let mut out = [0; 4];
774 mmu.read(0x1004, &mut out).unwrap();
775 assert_eq!(out, [9; 4]);
776 }
777
778 #[test]
779 fn snapshot_and_restore_round_trips_contents_and_mappings() {
780 let mut mmu = mapped();
781 mmu.write(0x1000, &[1, 2, 3, 4]).unwrap();
782 let snapshot = mmu.snapshot();
783
784 mmu.write(0x1000, &[9, 9, 9, 9]).unwrap();
785 mmu.map(0x8000, PAGE_SIZE, perm::RW_INIT).unwrap();
786 mmu.unmap(0x2000, PAGE_SIZE).unwrap();
787
788 mmu.restore(&snapshot);
789 let mut out = [0; 4];
790 mmu.read(0x1000, &mut out).unwrap();
791 assert_eq!(out, [1, 2, 3, 4]);
792 assert_eq!(mmu.permissions(0x8000), perm::NONE);
794 mmu.read(0x2000, &mut out).unwrap();
795 }
796
797 #[test]
798 fn access_wrapping_the_address_space_faults() {
799 let mmu = mapped();
800 let mut out = [0; 8];
801 assert_eq!(
802 mmu.read(u64::MAX - 2, &mut out).unwrap_err(),
803 MemFault::new(FaultKind::AddressOverflow, u64::MAX - 2)
804 );
805 }
806
807 #[test]
808 fn write_unchecked_maps_and_ignores_permissions() {
809 let mut mmu = Mmu::new();
810 mmu.write_unchecked(0x1000, &[1, 2, 3, 4], perm::READ | perm::EXEC);
811 let mut out = [0; 4];
812 mmu.read(0x1000, &mut out).unwrap();
813 assert_eq!(out, [1, 2, 3, 4]);
814 assert_eq!(
815 mmu.write(0x1000, &[0]).unwrap_err(),
816 MemFault::new(FaultKind::WritePerm, 0x1000)
817 );
818 }
819}