Skip to main content

os/mm/
memory_set.rs

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