Skip to main content

os/mm/
memory_set.rs

1//! Implementation of [`MapArea`] and [`MemorySet`].
2
3use super::{FrameTracker, frame_alloc};
4use super::{PTEFlags, PageTable, PageTableEntry};
5use super::{PhysAddr, PhysPageNum, VirtAddr, VirtPageNum};
6use super::{StepByOne, VPNRange};
7use crate::config::{MEMORY_END, MMIO, PAGE_SIZE, TRAMPOLINE, TRAP_CONTEXT, USER_STACK_SIZE};
8use crate::sync::UPSafeCell;
9use alloc::collections::BTreeMap;
10use alloc::sync::Arc;
11use alloc::vec::Vec;
12use core::arch::asm;
13use lazy_static::*;
14use riscv::register::satp;
15
16unsafe extern "C" {
17    safe fn stext();
18    safe fn etext();
19    safe fn srodata();
20    safe fn erodata();
21    safe fn sdata();
22    safe fn edata();
23    safe fn sbss_with_stack();
24    safe fn ebss();
25    safe fn ekernel();
26    safe fn strampoline();
27}
28
29lazy_static! {
30    /// a memory set instance through lazy_static! managing kernel space
31    pub static ref KERNEL_SPACE: Arc<UPSafeCell<MemorySet>> =
32        Arc::new(unsafe { UPSafeCell::new(MemorySet::new_kernel()) });
33}
34
35/// memory set structure, controls virtual-memory space
36pub struct MemorySet {
37    page_table: PageTable,
38    areas: Vec<MapArea>,
39}
40
41impl MemorySet {
42    pub fn new_bare() -> Self {
43        Self {
44            page_table: PageTable::new(),
45            areas: Vec::new(),
46        }
47    }
48    pub fn token(&self) -> usize {
49        self.page_table.token()
50    }
51    /// Assume that no conflicts.
52    pub fn insert_framed_area(
53        &mut self,
54        start_va: VirtAddr,
55        end_va: VirtAddr,
56        permission: MapPermission,
57    ) {
58        self.push(
59            MapArea::new(start_va, end_va, MapType::Framed, permission),
60            None,
61        );
62    }
63    fn push(&mut self, mut map_area: MapArea, data: Option<&[u8]>) {
64        map_area.map(&mut self.page_table);
65        if let Some(data) = data {
66            map_area.copy_data(&self.page_table, data);
67        }
68        self.areas.push(map_area);
69    }
70    /// Mention that trampoline is not collected by areas.
71    fn map_trampoline(&mut self) {
72        self.page_table.map(
73            VirtAddr::from(TRAMPOLINE).into(),
74            PhysAddr::from(linker_symbol_addr!(strampoline)).into(),
75            PTEFlags::R | PTEFlags::X,
76        );
77    }
78    /// Without kernel stacks.
79    pub fn new_kernel() -> Self {
80        let mut memory_set = Self::new_bare();
81        // map trampoline
82        memory_set.map_trampoline();
83        // map kernel sections
84        println!(
85            ".text [{:#x}, {:#x})",
86            linker_symbol_addr!(stext),
87            linker_symbol_addr!(etext)
88        );
89        println!(
90            ".rodata [{:#x}, {:#x})",
91            linker_symbol_addr!(srodata),
92            linker_symbol_addr!(erodata)
93        );
94        println!(
95            ".data [{:#x}, {:#x})",
96            linker_symbol_addr!(sdata),
97            linker_symbol_addr!(edata)
98        );
99        println!(
100            ".bss [{:#x}, {:#x})",
101            linker_symbol_addr!(sbss_with_stack),
102            linker_symbol_addr!(ebss)
103        );
104        println!("mapping .text section");
105        memory_set.push(
106            MapArea::new(
107                (linker_symbol_addr!(stext)).into(),
108                (linker_symbol_addr!(etext)).into(),
109                MapType::Identical,
110                MapPermission::R | MapPermission::X,
111            ),
112            None,
113        );
114        println!("mapping .rodata section");
115        memory_set.push(
116            MapArea::new(
117                (linker_symbol_addr!(srodata)).into(),
118                (linker_symbol_addr!(erodata)).into(),
119                MapType::Identical,
120                MapPermission::R,
121            ),
122            None,
123        );
124        println!("mapping .data section");
125        memory_set.push(
126            MapArea::new(
127                (linker_symbol_addr!(sdata)).into(),
128                (linker_symbol_addr!(edata)).into(),
129                MapType::Identical,
130                MapPermission::R | MapPermission::W,
131            ),
132            None,
133        );
134        println!("mapping .bss section");
135        memory_set.push(
136            MapArea::new(
137                (linker_symbol_addr!(sbss_with_stack)).into(),
138                (linker_symbol_addr!(ebss)).into(),
139                MapType::Identical,
140                MapPermission::R | MapPermission::W,
141            ),
142            None,
143        );
144        println!("mapping physical memory");
145        memory_set.push(
146            MapArea::new(
147                (linker_symbol_addr!(ekernel)).into(),
148                MEMORY_END.into(),
149                MapType::Identical,
150                MapPermission::R | MapPermission::W,
151            ),
152            None,
153        );
154        println!("mapping memory-mapped registers");
155        for pair in MMIO {
156            memory_set.push(
157                MapArea::new(
158                    (*pair).0.into(),
159                    ((*pair).0 + (*pair).1).into(),
160                    MapType::Identical,
161                    MapPermission::R | MapPermission::W,
162                ),
163                None,
164            );
165        }
166        memory_set
167    }
168    /// Include sections in elf and trampoline and TrapContext and user stack,
169    /// also returns user_sp and entry point.
170    pub fn from_elf(elf_data: &[u8]) -> (Self, usize, usize) {
171        let mut memory_set = Self::new_bare();
172        // map trampoline
173        memory_set.map_trampoline();
174        // map program headers of elf, with U flag
175        let elf = xmas_elf::ElfFile::new(elf_data).unwrap();
176        let elf_header = elf.header;
177        let magic = elf_header.pt1.magic;
178        assert_eq!(magic, [0x7f, 0x45, 0x4c, 0x46], "invalid elf!");
179        let ph_count = elf_header.pt2.ph_count();
180        let mut max_end_vpn = VirtPageNum(0);
181        for i in 0..ph_count {
182            let ph = elf.program_header(i).unwrap();
183            if ph.get_type().unwrap() == xmas_elf::program::Type::Load {
184                let start_va: VirtAddr = (ph.virtual_addr() as usize).into();
185                let end_va: VirtAddr = ((ph.virtual_addr() + ph.mem_size()) as usize).into();
186                let mut map_perm = MapPermission::U;
187                let ph_flags = ph.flags();
188                if ph_flags.is_read() {
189                    map_perm |= MapPermission::R;
190                }
191                if ph_flags.is_write() {
192                    map_perm |= MapPermission::W;
193                }
194                if ph_flags.is_execute() {
195                    map_perm |= MapPermission::X;
196                }
197                let map_area = MapArea::new(start_va, end_va, MapType::Framed, map_perm);
198                max_end_vpn = map_area.vpn_range.get_end();
199                memory_set.push(
200                    map_area,
201                    Some(&elf.input[ph.offset() as usize..(ph.offset() + ph.file_size()) as usize]),
202                );
203            }
204        }
205        // map user stack with U flags
206        let max_end_va: VirtAddr = max_end_vpn.into();
207        let mut user_stack_bottom: usize = max_end_va.into();
208        // guard page
209        user_stack_bottom += PAGE_SIZE;
210        let user_stack_top = user_stack_bottom + USER_STACK_SIZE;
211        memory_set.push(
212            MapArea::new(
213                user_stack_bottom.into(),
214                user_stack_top.into(),
215                MapType::Framed,
216                MapPermission::R | MapPermission::W | MapPermission::U,
217            ),
218            None,
219        );
220        // used in sbrk
221        memory_set.push(
222            MapArea::new(
223                user_stack_top.into(),
224                user_stack_top.into(),
225                MapType::Framed,
226                MapPermission::R | MapPermission::W | MapPermission::U,
227            ),
228            None,
229        );
230        // map TrapContext
231        memory_set.push(
232            MapArea::new(
233                TRAP_CONTEXT.into(),
234                TRAMPOLINE.into(),
235                MapType::Framed,
236                MapPermission::R | MapPermission::W,
237            ),
238            None,
239        );
240        (
241            memory_set,
242            user_stack_top,
243            elf.header.pt2.entry_point() as usize,
244        )
245    }
246    pub fn activate(&self) {
247        let satp = self.page_table.token();
248        unsafe {
249            satp::write(satp::Satp::from_bits(satp));
250            asm!("sfence.vma");
251        }
252    }
253    pub fn translate(&self, vpn: VirtPageNum) -> Option<PageTableEntry> {
254        self.page_table.translate(vpn)
255    }
256    #[allow(unused)]
257    pub fn shrink_to(&mut self, start: VirtAddr, new_end: VirtAddr) -> bool {
258        if let Some(area) = self
259            .areas
260            .iter_mut()
261            .find(|area| area.vpn_range.get_start() == start.floor())
262        {
263            area.shrink_to(&mut self.page_table, new_end.ceil());
264            true
265        } else {
266            false
267        }
268    }
269    #[allow(unused)]
270    pub fn append_to(&mut self, start: VirtAddr, new_end: VirtAddr) -> bool {
271        if let Some(area) = self
272            .areas
273            .iter_mut()
274            .find(|area| area.vpn_range.get_start() == start.floor())
275        {
276            area.append_to(&mut self.page_table, new_end.ceil());
277            true
278        } else {
279            false
280        }
281    }
282}
283
284/// map area structure, controls a contiguous piece of virtual memory
285pub struct MapArea {
286    vpn_range: VPNRange,
287    data_frames: BTreeMap<VirtPageNum, FrameTracker>,
288    map_type: MapType,
289    map_perm: MapPermission,
290}
291
292impl MapArea {
293    pub fn new(
294        start_va: VirtAddr,
295        end_va: VirtAddr,
296        map_type: MapType,
297        map_perm: MapPermission,
298    ) -> Self {
299        let start_vpn: VirtPageNum = start_va.floor();
300        let end_vpn: VirtPageNum = end_va.ceil();
301        Self {
302            vpn_range: VPNRange::new(start_vpn, end_vpn),
303            data_frames: BTreeMap::new(),
304            map_type,
305            map_perm,
306        }
307    }
308    pub fn map_one(&mut self, page_table: &mut PageTable, vpn: VirtPageNum) {
309        let ppn: PhysPageNum;
310        match self.map_type {
311            MapType::Identical => {
312                ppn = PhysPageNum(vpn.0);
313            }
314            MapType::Framed => {
315                let frame = frame_alloc().unwrap();
316                ppn = frame.ppn;
317                self.data_frames.insert(vpn, frame);
318            }
319        }
320        let pte_flags = PTEFlags::from_bits(self.map_perm.bits).unwrap();
321        page_table.map(vpn, ppn, pte_flags);
322    }
323    #[allow(unused)]
324    pub fn unmap_one(&mut self, page_table: &mut PageTable, vpn: VirtPageNum) {
325        if self.map_type == MapType::Framed {
326            self.data_frames.remove(&vpn);
327        }
328        page_table.unmap(vpn);
329    }
330    pub fn map(&mut self, page_table: &mut PageTable) {
331        for vpn in self.vpn_range {
332            self.map_one(page_table, vpn);
333        }
334    }
335    #[allow(unused)]
336    pub fn unmap(&mut self, page_table: &mut PageTable) {
337        for vpn in self.vpn_range {
338            self.unmap_one(page_table, vpn);
339        }
340    }
341    #[allow(unused)]
342    pub fn shrink_to(&mut self, page_table: &mut PageTable, new_end: VirtPageNum) {
343        for vpn in VPNRange::new(new_end, self.vpn_range.get_end()) {
344            self.unmap_one(page_table, vpn)
345        }
346        self.vpn_range = VPNRange::new(self.vpn_range.get_start(), new_end);
347    }
348    #[allow(unused)]
349    pub fn append_to(&mut self, page_table: &mut PageTable, new_end: VirtPageNum) {
350        for vpn in VPNRange::new(self.vpn_range.get_end(), new_end) {
351            self.map_one(page_table, vpn)
352        }
353        self.vpn_range = VPNRange::new(self.vpn_range.get_start(), new_end);
354    }
355    /// data: start-aligned but maybe with shorter length
356    /// assume that all frames were cleared before
357    pub fn copy_data(&mut self, page_table: &PageTable, data: &[u8]) {
358        assert_eq!(self.map_type, MapType::Framed);
359        let mut start: usize = 0;
360        let mut current_vpn = self.vpn_range.get_start();
361        let len = data.len();
362        loop {
363            let src = &data[start..len.min(start + PAGE_SIZE)];
364            let dst = &mut page_table
365                .translate(current_vpn)
366                .unwrap()
367                .ppn()
368                .get_bytes_array()[..src.len()];
369            dst.copy_from_slice(src);
370            start += PAGE_SIZE;
371            if start >= len {
372                break;
373            }
374            current_vpn.step();
375        }
376    }
377}
378
379#[derive(Copy, Clone, PartialEq, Debug)]
380/// map type for memory set: identical or framed
381pub enum MapType {
382    Identical,
383    Framed,
384}
385
386bitflags! {
387    /// map permission corresponding to that in pte: `R W X U`
388    pub struct MapPermission: u8 {
389        const R = 1 << 1;
390        const W = 1 << 2;
391        const X = 1 << 3;
392        const U = 1 << 4;
393    }
394}
395
396#[allow(unused)]
397pub fn remap_test() {
398    let mut kernel_space = KERNEL_SPACE.exclusive_access();
399    let mid_text: VirtAddr = ((linker_symbol_addr!(stext) + linker_symbol_addr!(etext)) / 2).into();
400    let mid_rodata: VirtAddr =
401        ((linker_symbol_addr!(srodata) + linker_symbol_addr!(erodata)) / 2).into();
402    let mid_data: VirtAddr = ((linker_symbol_addr!(sdata) + linker_symbol_addr!(edata)) / 2).into();
403    assert!(
404        !kernel_space
405            .page_table
406            .translate(mid_text.floor())
407            .unwrap()
408            .writable(),
409    );
410    assert!(
411        !kernel_space
412            .page_table
413            .translate(mid_rodata.floor())
414            .unwrap()
415            .writable(),
416    );
417    assert!(
418        !kernel_space
419            .page_table
420            .translate(mid_data.floor())
421            .unwrap()
422            .executable(),
423    );
424    println!("remap_test passed!");
425}