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, TRAP_CONTEXT, USER_STACK_SIZE};
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 and TrapContext and user stack,
180    /// also returns user_sp 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        // map user stack with U flags
217        let max_end_va: VirtAddr = max_end_vpn.into();
218        let mut user_stack_bottom: usize = max_end_va.into();
219        // guard page
220        user_stack_bottom += PAGE_SIZE;
221        let user_stack_top = user_stack_bottom + USER_STACK_SIZE;
222        memory_set.push(
223            MapArea::new(
224                user_stack_bottom.into(),
225                user_stack_top.into(),
226                MapType::Framed,
227                MapPermission::R | MapPermission::W | MapPermission::U,
228            ),
229            None,
230        );
231        // map TrapContext
232        memory_set.push(
233            MapArea::new(
234                TRAP_CONTEXT.into(),
235                TRAMPOLINE.into(),
236                MapType::Framed,
237                MapPermission::R | MapPermission::W,
238            ),
239            None,
240        );
241        (
242            memory_set,
243            user_stack_top,
244            elf.header.pt2.entry_point() as usize,
245        )
246    }
247    pub fn from_existed_user(user_space: &MemorySet) -> MemorySet {
248        let mut memory_set = Self::new_bare();
249        // map trampoline
250        memory_set.map_trampoline();
251        // copy data sections/trap_context/user_stack
252        for area in user_space.areas.iter() {
253            let new_area = MapArea::from_another(area);
254            memory_set.push(new_area, None);
255            // copy data from another space
256            for vpn in area.vpn_range {
257                let src_ppn = user_space.translate(vpn).unwrap().ppn();
258                let dst_ppn = memory_set.translate(vpn).unwrap().ppn();
259                dst_ppn
260                    .get_bytes_array()
261                    .copy_from_slice(src_ppn.get_bytes_array());
262            }
263        }
264        memory_set
265    }
266    pub fn activate(&self) {
267        let satp = self.page_table.token();
268        unsafe {
269            satp::write(satp::Satp::from_bits(satp));
270            asm!("sfence.vma");
271        }
272    }
273    pub fn translate(&self, vpn: VirtPageNum) -> Option<PageTableEntry> {
274        self.page_table.translate(vpn)
275    }
276    pub fn recycle_data_pages(&mut self) {
277        //*self = Self::new_bare();
278        self.areas.clear();
279    }
280}
281
282pub struct MapArea {
283    vpn_range: VPNRange,
284    data_frames: BTreeMap<VirtPageNum, FrameTracker>,
285    map_type: MapType,
286    map_perm: MapPermission,
287}
288
289impl MapArea {
290    pub fn new(
291        start_va: VirtAddr,
292        end_va: VirtAddr,
293        map_type: MapType,
294        map_perm: MapPermission,
295    ) -> Self {
296        let start_vpn: VirtPageNum = start_va.floor();
297        let end_vpn: VirtPageNum = end_va.ceil();
298        Self {
299            vpn_range: VPNRange::new(start_vpn, end_vpn),
300            data_frames: BTreeMap::new(),
301            map_type,
302            map_perm,
303        }
304    }
305    pub fn from_another(another: &MapArea) -> Self {
306        Self {
307            vpn_range: VPNRange::new(another.vpn_range.get_start(), another.vpn_range.get_end()),
308            data_frames: BTreeMap::new(),
309            map_type: another.map_type,
310            map_perm: another.map_perm,
311        }
312    }
313    pub fn map_one(&mut self, page_table: &mut PageTable, vpn: VirtPageNum) {
314        let ppn: PhysPageNum;
315        match self.map_type {
316            MapType::Identical => {
317                ppn = PhysPageNum(vpn.0);
318            }
319            MapType::Framed => {
320                let frame = frame_alloc().unwrap();
321                ppn = frame.ppn;
322                self.data_frames.insert(vpn, frame);
323            }
324        }
325        let pte_flags = PTEFlags::from_bits(self.map_perm.bits).unwrap();
326        page_table.map(vpn, ppn, pte_flags);
327    }
328    pub fn unmap_one(&mut self, page_table: &mut PageTable, vpn: VirtPageNum) {
329        if self.map_type == MapType::Framed {
330            self.data_frames.remove(&vpn);
331        }
332        page_table.unmap(vpn);
333    }
334    pub fn map(&mut self, page_table: &mut PageTable) {
335        for vpn in self.vpn_range {
336            self.map_one(page_table, vpn);
337        }
338    }
339    pub fn unmap(&mut self, page_table: &mut PageTable) {
340        for vpn in self.vpn_range {
341            self.unmap_one(page_table, vpn);
342        }
343    }
344    /// data: start-aligned but maybe with shorter length
345    /// assume that all frames were cleared before
346    pub fn copy_data(&mut self, page_table: &PageTable, data: &[u8]) {
347        assert_eq!(self.map_type, MapType::Framed);
348        let mut start: usize = 0;
349        let mut current_vpn = self.vpn_range.get_start();
350        let len = data.len();
351        loop {
352            let src = &data[start..len.min(start + PAGE_SIZE)];
353            let dst = &mut page_table
354                .translate(current_vpn)
355                .unwrap()
356                .ppn()
357                .get_bytes_array()[..src.len()];
358            dst.copy_from_slice(src);
359            start += PAGE_SIZE;
360            if start >= len {
361                break;
362            }
363            current_vpn.step();
364        }
365    }
366}
367
368#[derive(Copy, Clone, PartialEq, Debug)]
369pub enum MapType {
370    Identical,
371    Framed,
372}
373
374bitflags! {
375    pub struct MapPermission: u8 {
376        const R = 1 << 1;
377        const W = 1 << 2;
378        const X = 1 << 3;
379        const U = 1 << 4;
380    }
381}
382
383#[allow(unused)]
384pub fn remap_test() {
385    let mut kernel_space = KERNEL_SPACE.exclusive_access();
386    let mid_text: VirtAddr = ((linker_symbol_addr!(stext) + linker_symbol_addr!(etext)) / 2).into();
387    let mid_rodata: VirtAddr =
388        ((linker_symbol_addr!(srodata) + linker_symbol_addr!(erodata)) / 2).into();
389    let mid_data: VirtAddr = ((linker_symbol_addr!(sdata) + linker_symbol_addr!(edata)) / 2).into();
390    assert!(
391        !kernel_space
392            .page_table
393            .translate(mid_text.floor())
394            .unwrap()
395            .writable(),
396    );
397    assert!(
398        !kernel_space
399            .page_table
400            .translate(mid_rodata.floor())
401            .unwrap()
402            .writable(),
403    );
404    assert!(
405        !kernel_space
406            .page_table
407            .translate(mid_data.floor())
408            .unwrap()
409            .executable(),
410    );
411    println!("remap_test passed!");
412}