Skip to main content

strat9_kernel/ostd/
mm.rs

1//! Memory management abstraction layer
2//!
3//! Provides safe abstractions for memory operations including:
4//! - Physical and virtual address types
5//! - Memory mapping abstractions (MappedPages)
6//! - Page table management
7//!
8//! Inspired by OSes Theseus MappedPages and Asterinas VM modules.
9
10#![allow(unsafe_code)]
11#![allow(unsafe_op_in_unsafe_fn)]
12
13extern crate alloc;
14
15use core::{marker::PhantomData, ops::Range};
16
17/// Physical address type
18#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
19#[repr(transparent)]
20pub struct PhysAddr(u64);
21
22impl PhysAddr {
23    /// Creates a new physical address
24    pub const fn new(addr: u64) -> Self {
25        Self(addr)
26    }
27
28    /// Creates a null physical address
29    pub const fn null() -> Self {
30        Self(0)
31    }
32
33    /// Returns the raw address value
34    pub const fn as_u64(&self) -> u64 {
35        self.0
36    }
37
38    /// Returns the raw address value as usize
39    pub const fn as_usize(&self) -> usize {
40        self.0 as usize
41    }
42
43    /// Checks if the address is null
44    pub const fn is_null(&self) -> bool {
45        self.0 == 0
46    }
47
48    /// Aligns the address up to the given alignment
49    pub const fn align_up(&self, align: u64) -> Self {
50        Self((self.0 + align - 1) & !(align - 1))
51    }
52
53    /// Aligns the address down to the given alignment
54    pub const fn align_down(&self, align: u64) -> Self {
55        Self(self.0 & !(align - 1))
56    }
57
58    /// Checks if the address is aligned to the given alignment
59    pub const fn is_aligned(&self, align: u64) -> bool {
60        self.0 & (align - 1) == 0
61    }
62
63    /// Adds an offset to the address
64    pub const fn add(&self, offset: u64) -> Self {
65        Self(self.0 + offset)
66    }
67
68    /// Subtracts an offset from the address
69    pub const fn sub(&self, offset: u64) -> Self {
70        Self(self.0 - offset)
71    }
72
73    /// Converts to a raw mutable pointer
74    pub fn as_mut_ptr<T>(&self) -> *mut T {
75        self.0 as *mut T
76    }
77
78    /// Converts to a raw pointer
79    pub fn as_ptr<T>(&self) -> *const T {
80        self.0 as *const T
81    }
82}
83
84impl From<u64> for PhysAddr {
85    /// Performs the from operation.
86    fn from(addr: u64) -> Self {
87        Self::new(addr)
88    }
89}
90
91impl From<PhysAddr> for u64 {
92    /// Performs the from operation.
93    fn from(addr: PhysAddr) -> u64 {
94        addr.as_u64()
95    }
96}
97
98/// Virtual address type
99#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
100#[repr(transparent)]
101pub struct VirtAddr(u64);
102
103impl VirtAddr {
104    /// Creates a new virtual address
105    pub const fn new(addr: u64) -> Self {
106        Self(addr)
107    }
108
109    /// Creates a null virtual address
110    pub const fn null() -> Self {
111        Self(0)
112    }
113
114    /// Returns the raw address value
115    pub const fn as_u64(&self) -> u64 {
116        self.0
117    }
118
119    /// Returns the raw address value as usize
120    pub const fn as_usize(&self) -> usize {
121        self.0 as usize
122    }
123
124    /// Checks if the address is null
125    pub const fn is_null(&self) -> bool {
126        self.0 == 0
127    }
128
129    /// Aligns the address up to the given alignment
130    pub const fn align_up(&self, align: u64) -> Self {
131        Self((self.0 + align - 1) & !(align - 1))
132    }
133
134    /// Aligns the address down to the given alignment
135    pub const fn align_down(&self, align: u64) -> Self {
136        Self(self.0 & !(align - 1))
137    }
138
139    /// Checks if the address is aligned to the given alignment
140    pub const fn is_aligned(&self, align: u64) -> bool {
141        self.0 & (align - 1) == 0
142    }
143
144    /// Adds an offset to the address
145    pub const fn add(&self, offset: u64) -> Self {
146        Self(self.0 + offset)
147    }
148
149    /// Subtracts an offset from the address
150    pub const fn sub(&self, offset: u64) -> Self {
151        Self(self.0 - offset)
152    }
153
154    /// Converts to a raw mutable pointer
155    pub fn as_mut_ptr<T>(&self) -> *mut T {
156        self.0 as *mut T
157    }
158
159    /// Converts to a raw pointer
160    pub fn as_ptr<T>(&self) -> *const T {
161        self.0 as *const T
162    }
163}
164
165impl core::ops::Add<u64> for VirtAddr {
166    type Output = VirtAddr;
167    fn add(self, rhs: u64) -> VirtAddr {
168        VirtAddr(self.0 + rhs)
169    }
170}
171
172impl core::ops::Sub<u64> for VirtAddr {
173    type Output = VirtAddr;
174    fn sub(self, rhs: u64) -> VirtAddr {
175        VirtAddr(self.0 - rhs)
176    }
177}
178
179impl core::ops::Add<u64> for PhysAddr {
180    type Output = PhysAddr;
181    fn add(self, rhs: u64) -> PhysAddr {
182        PhysAddr(self.0 + rhs)
183    }
184}
185
186impl core::ops::Sub<u64> for PhysAddr {
187    type Output = PhysAddr;
188    fn sub(self, rhs: u64) -> PhysAddr {
189        PhysAddr(self.0 - rhs)
190    }
191}
192
193impl From<u64> for VirtAddr {
194    /// Performs the from operation.
195    fn from(addr: u64) -> Self {
196        Self::new(addr)
197    }
198}
199
200impl From<VirtAddr> for u64 {
201    /// Performs the from operation.
202    fn from(addr: VirtAddr) -> u64 {
203        addr.as_u64()
204    }
205}
206
207/// Page size constant (4KB)
208pub const PAGE_SIZE: usize = 4096;
209
210/// Converts a physical address to a virtual address using HHDM offset
211#[inline]
212pub fn phys_to_virt(phys: PhysAddr) -> VirtAddr {
213    VirtAddr::new(crate::memory::phys_to_virt(phys.as_u64()))
214}
215
216/// Converts a virtual address to a physical address
217#[inline]
218pub fn virt_to_phys(virt: VirtAddr) -> PhysAddr {
219    PhysAddr::new(crate::memory::virt_to_phys(virt.as_u64()))
220}
221
222/// A safely mapped memory region
223///
224/// `MappedPages` represents a contiguous virtual memory mapping to physical frames.
225/// The mapping is automatically unmapped when the `MappedPages` is dropped.
226///
227/// This is inspired by Theseus's MappedPages abstraction.
228pub struct MappedPages {
229    /// Starting virtual address
230    start_vaddr: VirtAddr,
231    /// Size in bytes
232    size: usize,
233    /// Whether this mapping owns the underlying frames
234    owned: bool,
235    /// Marker to prevent Send/Sync (mapping is CPU-local)
236    _marker: PhantomData<*mut ()>,
237}
238
239// SAFETY: MappedPages can be sent between CPUs if explicitly transferred
240unsafe impl Send for MappedPages {}
241
242impl MappedPages {
243    /// Creates a new MappedPages from an existing mapping
244    ///
245    /// # Safety
246    ///
247    /// - The virtual address range must be a valid mapping
248    /// - The caller must ensure the mapping remains valid for the lifetime
249    /// - The size must match the actual mapping size
250    pub unsafe fn new(start_vaddr: VirtAddr, size: usize, owned: bool) -> Self {
251        Self {
252            start_vaddr,
253            size,
254            owned,
255            _marker: PhantomData,
256        }
257    }
258
259    /// Returns the starting virtual address
260    pub fn start_address(&self) -> VirtAddr {
261        self.start_vaddr
262    }
263
264    /// Returns the size in bytes
265    pub fn size(&self) -> usize {
266        self.size
267    }
268
269    /// Returns the ending virtual address (exclusive)
270    pub fn end_address(&self) -> VirtAddr {
271        self.start_vaddr.add(self.size as u64)
272    }
273
274    /// Returns the virtual address range
275    pub fn range(&self) -> Range<VirtAddr> {
276        self.start_vaddr..self.end_address()
277    }
278
279    /// Returns a pointer to the start of the mapping
280    pub fn as_ptr(&self) -> *const u8 {
281        self.start_vaddr.as_usize() as *const u8
282    }
283
284    /// Returns a mutable pointer to the start of the mapping
285    pub fn as_mut_ptr(&mut self) -> *mut u8 {
286        self.start_vaddr.as_usize() as *mut u8
287    }
288
289    /// Reads a value from the mapped memory at the given offset
290    ///
291    /// # Safety
292    ///
293    /// - The offset + size_of::<T>() must be within the mapping
294    /// - The memory must be properly initialized for type T
295    /// - Proper alignment must be ensured
296    pub unsafe fn read<T>(&self, offset: usize) -> Result<T, MapError> {
297        if offset + core::mem::size_of::<T>() > self.size {
298            return Err(MapError::OutOfBounds);
299        }
300        let ptr = self.start_vaddr.as_usize().wrapping_add(offset) as *const T;
301        // SAFETY: Caller guarantees the pointer is valid and properly aligned
302        Ok(ptr.read_volatile())
303    }
304
305    /// Writes a value to the mapped memory at the given offset
306    ///
307    /// # Safety
308    ///
309    /// - The offset + size_of::<T>() must be within the mapping
310    /// - The memory must be writable (not read-only)
311    /// - Proper alignment must be ensured
312    pub unsafe fn write<T>(&mut self, offset: usize, value: T) -> Result<(), MapError> {
313        if offset + core::mem::size_of::<T>() > self.size {
314            return Err(MapError::OutOfBounds);
315        }
316        let ptr = self.start_vaddr.as_usize().wrapping_add(offset) as *mut T;
317        // SAFETY: Caller guarantees the pointer is valid and writable
318        ptr.write_volatile(value);
319        Ok(())
320    }
321
322    /// Returns a slice reference to the mapped memory
323    ///
324    /// # Safety
325    ///
326    /// - The mapping must contain initialized data
327    /// - No other mutable references to this memory can exist
328    pub unsafe fn as_slice(&self, len: usize) -> Result<&[u8], MapError> {
329        if len > self.size {
330            return Err(MapError::OutOfBounds);
331        }
332        Ok(core::slice::from_raw_parts(self.as_ptr(), len))
333    }
334
335    /// Returns a mutable slice reference to the mapped memory
336    ///
337    /// # Safety
338    ///
339    /// - The mapping must be writable
340    /// - No other references to this memory can exist
341    pub unsafe fn as_mut_slice(&mut self, len: usize) -> Result<&mut [u8], MapError> {
342        if len > self.size {
343            return Err(MapError::OutOfBounds);
344        }
345        Ok(core::slice::from_raw_parts_mut(self.as_mut_ptr(), len))
346    }
347
348    /// Converts this MappedPages into an AllocatedPages, consuming the mapping
349    ///
350    /// This transfers ownership of the underlying frames.
351    pub fn into_allocated_pages(self) -> Result<AllocatedPages, MapError> {
352        if !self.owned {
353            return Err(MapError::NotOwner);
354        }
355        let pages = AllocatedPages {
356            start_vaddr: self.start_vaddr,
357            size: self.size,
358        };
359        // Prevent the Drop implementation from running
360        core::mem::forget(self);
361        Ok(pages)
362    }
363}
364
365impl Drop for MappedPages {
366    /// Performs the drop operation.
367    fn drop(&mut self) {
368        if self.owned {
369            // Calculate page count
370            let page_count = (self.size + PAGE_SIZE - 1) / PAGE_SIZE;
371            // We own the mapping and are responsible for unmapping
372            crate::memory::address_space::kernel_address_space()
373                .unmap_region(
374                    self.start_vaddr.as_u64(),
375                    page_count,
376                    crate::memory::address_space::VmaPageSize::Small,
377                )
378                .ok();
379        }
380    }
381}
382
383impl core::fmt::Debug for MappedPages {
384    /// Performs the fmt operation.
385    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
386        f.debug_struct("MappedPages")
387            .field("start", &self.start_vaddr)
388            .field("size", &self.size)
389            .field("owned", &self.owned)
390            .finish()
391    }
392}
393
394/// Allocated pages that can be mapped
395///
396/// Represents virtually allocated pages that own their underlying frames.
397pub struct AllocatedPages {
398    start_vaddr: VirtAddr,
399    size: usize,
400}
401
402impl AllocatedPages {
403    /// Returns the starting virtual address
404    pub fn start_address(&self) -> VirtAddr {
405        self.start_vaddr
406    }
407
408    /// Returns the size in bytes
409    pub fn size(&self) -> usize {
410        self.size
411    }
412
413    /// Returns the number of pages
414    pub fn page_count(&self) -> usize {
415        (self.size + PAGE_SIZE - 1) / PAGE_SIZE
416    }
417}
418
419impl Drop for AllocatedPages {
420    /// Performs the drop operation.
421    fn drop(&mut self) {
422        // Deallocate the frames using the buddy allocator
423        // SAFETY: we own these pages and are responsible for deallocation
424        let phys_addr = virt_to_phys(self.start_vaddr).as_u64();
425        // TODO: implement proper frame deallocation
426        // For now, we just leak the frames to avoid double-free issues
427        let _ = phys_addr;
428        let _ = self.size;
429        // crate::memory::frame::deallocate_frames(phys_addr, self.size);
430    }
431}
432
433/// Memory mapping flags
434#[derive(Debug, Clone, Copy)]
435pub struct MapFlags {
436    /// Page is present (mapped)
437    pub present: bool,
438    /// Page is writable
439    pub writable: bool,
440    /// Page is user-accessible
441    pub user: bool,
442    /// Write-through caching
443    pub write_through: bool,
444    /// Cache disabled
445    pub cache_disabled: bool,
446    /// No-execute (NX)
447    pub no_execute: bool,
448}
449
450impl MapFlags {
451    /// Creates flags for a read-only kernel mapping
452    pub const fn read_only() -> Self {
453        Self {
454            present: true,
455            writable: false,
456            user: false,
457            write_through: false,
458            cache_disabled: false,
459            no_execute: false,
460        }
461    }
462
463    /// Creates flags for a read-write kernel mapping
464    pub const fn read_write() -> Self {
465        Self {
466            present: true,
467            writable: true,
468            user: false,
469            write_through: false,
470            cache_disabled: false,
471            no_execute: false,
472        }
473    }
474
475    /// Creates flags for a user mapping
476    pub const fn user_read_write() -> Self {
477        Self {
478            present: true,
479            writable: true,
480            user: true,
481            write_through: false,
482            cache_disabled: false,
483            no_execute: false,
484        }
485    }
486
487    /// Creates flags for MMIO (device memory)
488    pub const fn mmio() -> Self {
489        Self {
490            present: true,
491            writable: true,
492            user: false,
493            write_through: false,
494            cache_disabled: true,
495            no_execute: true,
496        }
497    }
498}
499
500#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
501pub enum MapError {
502    #[error("address out of bounds")]
503    OutOfBounds,
504    #[error("not owner of mapping")]
505    NotOwner,
506    #[error("address already mapped")]
507    AlreadyMapped,
508    #[error("invalid address")]
509    InvalidAddress,
510    #[error("out of memory")]
511    OutOfMemory,
512    #[error("architecture error: {0}")]
513    ArchError(&'static str),
514}
515
516/// Virtual Memory Address Region (VMAR)
517///
518/// Manages a region of virtual address space, similar to Asterinas VMAR.
519/// Used for process address space management.
520pub struct Vmar {
521    /// Base virtual address
522    base: VirtAddr,
523    /// Size of the region
524    size: usize,
525    /// Child regions
526    children: spin::Mutex<alloc::vec::Vec<VmarChild>>,
527}
528
529struct VmarChild {
530    /// Offset from parent base
531    offset: usize,
532    /// Size of the child region
533    #[allow(dead_code)]
534    size: usize,
535    /// The actual mapping
536    #[allow(dead_code)]
537    mapping: Option<MappedPages>,
538}
539
540impl Vmar {
541    /// Creates a new VMAR
542    pub fn new(base: VirtAddr, size: usize) -> Self {
543        Self {
544            base,
545            size,
546            children: spin::Mutex::new(alloc::vec![]),
547        }
548    }
549
550    /// Returns the base virtual address
551    pub fn base(&self) -> VirtAddr {
552        self.base
553    }
554
555    /// Returns the size of the region
556    pub fn size(&self) -> usize {
557        self.size
558    }
559
560    /// Allocates a new region within this VMAR
561    pub fn alloc(&self, offset: usize, size: usize, flags: MapFlags) -> Result<VirtAddr, MapError> {
562        // TODO: implement proper allocation with conflict detection
563        let vaddr = self.base.add(offset as u64);
564
565        // TODO: map the region with the given flags
566        let _ = flags; // Suppress unused warning
567
568        let mut children = self.children.lock();
569        children.push(VmarChild {
570            offset,
571            size,
572            mapping: None,
573        });
574
575        Ok(vaddr)
576    }
577
578    /// Deallocates a region within this VMAR
579    pub fn dealloc(&self, offset: usize) -> Result<(), MapError> {
580        let mut children = self.children.lock();
581        if let Some(pos) = children.iter().position(|c| c.offset == offset) {
582            children.remove(pos);
583            Ok(())
584        } else {
585            Err(MapError::InvalidAddress)
586        }
587    }
588}
589
590impl core::fmt::Debug for Vmar {
591    /// Performs the fmt operation.
592    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
593        f.debug_struct("Vmar")
594            .field("base", &self.base)
595            .field("size", &self.size)
596            .finish()
597    }
598}
599
600/// TLB flush operation for SMP systems
601///
602/// Flushes TLB entries on all CPUs that may have cached the given virtual address.
603pub fn tlb_flush_virt_addr(vaddr: VirtAddr) {
604    // SAFETY: invlpg is a privileged instruction that invalidates a TLB entry.
605    // This is safe to call in kernel mode.
606    unsafe {
607        core::arch::asm!(
608            "invlpg [{}]",
609            in(reg) vaddr.as_u64(),
610            options(nostack, preserves_flags)
611        );
612    }
613}
614
615/// Flushes the entire TLB on the current CPU
616///
617/// This is more expensive than `tlb_flush_virt_addr` and should be used sparingly.
618pub fn tlb_flush_all() {
619    // SAFETY: writing to CR3 with the same value flushes the TLB (except global pages).
620    // This is safe to call in kernel mode.
621    unsafe {
622        let cr3: u64;
623        core::arch::asm!(
624            "mov {}, cr3",
625            out(reg) cr3,
626            options(nostack, preserves_flags)
627        );
628        core::arch::asm!(
629            "mov cr3, {}",
630            in(reg) cr3,
631            options(nostack)
632        );
633    }
634}