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