Add sbrk syscall and heap management from ch4 address space design

os/src/mm/memory_set.rs

@@ -245,6 +245,32 @@ impl MemorySet {
         //*self = Self::new_bare();
         self.areas.clear();
     }
+    #[allow(unused)]
+    pub fn shrink_to(&mut self, start: VirtAddr, new_end: VirtAddr) -> bool {
+        if let Some(area) = self
+            .areas
+            .iter_mut()
+            .find(|area| area.vpn_range.get_start() == start.floor())
+        {
+            area.shrink_to(&mut self.page_table, new_end.ceil());
+            true
+        } else {
+            false
+        }
+    }
+    #[allow(unused)]
+    pub fn append_to(&mut self, start: VirtAddr, new_end: VirtAddr) -> bool {
+        if let Some(area) = self
+            .areas
+            .iter_mut()
+            .find(|area| area.vpn_range.get_start() == start.floor())
+        {
+            area.append_to(&mut self.page_table, new_end.ceil());
+            true
+        } else {
+            false
+        }
+    }
 }
 
 pub struct MapArea {
@@ -336,6 +362,20 @@ impl MapArea {
             current_vpn.step();
         }
     }
+    #[allow(unused)]
+    pub fn shrink_to(&mut self, page_table: &mut PageTable, new_end: VirtPageNum) {
+        for vpn in VPNRange::new(new_end, self.vpn_range.get_end()) {
+            self.unmap_one(page_table, vpn)
+        }
+        self.vpn_range = VPNRange::new(self.vpn_range.get_start(), new_end);
+    }
+    #[allow(unused)]
+    pub fn append_to(&mut self, page_table: &mut PageTable, new_end: VirtPageNum) {
+        for vpn in VPNRange::new(self.vpn_range.get_end(), new_end) {
+            self.map_one(page_table, vpn)
+        }
+        self.vpn_range = VPNRange::new(self.vpn_range.get_start(), new_end);
+    }
 }
 
 #[derive(Copy, Clone, PartialEq, Debug)]

os/src/syscall/mod.rs

@@ -13,6 +13,7 @@ const SYSCALL_YIELD: usize = 124;
 const SYSCALL_KILL: usize = 129;
 const SYSCALL_GET_TIME: usize = 169;
 const SYSCALL_GETPID: usize = 172;
+const SYSCALL_SBRK: usize = 214;
 const SYSCALL_FORK: usize = 220;
 const SYSCALL_EXEC: usize = 221;
 const SYSCALL_WAITPID: usize = 260;
@@ -66,6 +67,7 @@ pub fn syscall(syscall_id: usize, args: [usize; 3]) -> isize {
         SYSCALL_KILL => sys_kill(args[0], args[1] as u32),
         SYSCALL_GET_TIME => sys_get_time(),
         SYSCALL_GETPID => sys_getpid(),
+        SYSCALL_SBRK => sys_sbrk(args[0] as i32),
         SYSCALL_FORK => sys_fork(),
         SYSCALL_EXEC => sys_exec(args[0] as *const u8, args[1] as *const usize),
         SYSCALL_WAITPID => sys_waitpid(args[0] as isize, args[1] as *mut i32),

os/src/syscall/process.rs

@@ -1,8 +1,8 @@
 use crate::fs::{OpenFlags, open_file};
 use crate::mm::{translated_ref, translated_refmut, translated_str};
 use crate::task::{
-    SignalFlags, current_process, current_task, current_user_token, exit_current_and_run_next,
-    pid2process, suspend_current_and_run_next,
+    SignalFlags, change_program_brk, current_process, current_task, current_user_token,
+    exit_current_and_run_next, pid2process, suspend_current_and_run_next,
 };
 use crate::timer::get_time_ms;
 use alloc::string::String;
@@ -27,6 +27,15 @@ pub fn sys_getpid() -> isize {
     current_task().unwrap().process.upgrade().unwrap().getpid() as isize
 }
 
+/// change data segment size
+pub fn sys_sbrk(size: i32) -> isize {
+    if let Some(old_brk) = change_program_brk(size) {
+        old_brk as isize
+    } else {
+        -1
+    }
+}
+
 pub fn sys_fork() -> isize {
     let current_process = current_process();
     let new_process = current_process.fork();

os/src/task/mod.rs

@@ -164,3 +164,11 @@ pub fn current_add_signal(signal: SignalFlags) {
     let mut process_inner = process.inner_exclusive_access();
     process_inner.signals |= signal;
 }
+
+/// Change the current process's program break.
+/// Returns the old break value on success, or None on failure.
+pub fn change_program_brk(size: i32) -> Option<usize> {
+    let process = current_process();
+    let mut process_inner = process.inner_exclusive_access();
+    process_inner.change_program_brk(size)
+}

os/src/task/process.rs

@@ -3,8 +3,9 @@ use super::id::RecycleAllocator;
 use super::manager::insert_into_pid2process;
 use super::{PidHandle, pid_alloc};
 use super::{SignalFlags, add_task};
+use crate::config::USER_STACK_SIZE;
 use crate::fs::{File, Stdin, Stdout};
-use crate::mm::{KERNEL_SPACE, MemorySet, translated_refmut};
+use crate::mm::{KERNEL_SPACE, MapArea, MapPermission, MapType, MemorySet, VirtAddr, translated_refmut};
 use crate::sync::{Condvar, Mutex, Semaphore, UPIntrFreeCell, UPIntrRefMut};
 use crate::trap::{TrapContext, trap_handler};
 use alloc::string::String;
@@ -32,6 +33,8 @@ pub struct ProcessControlBlockInner {
     pub mutex_list: Vec<Option<Arc<dyn Mutex>>>,
     pub semaphore_list: Vec<Option<Arc<Semaphore>>>,
     pub condvar_list: Vec<Option<Arc<Condvar>>>,
+    pub heap_bottom: usize,
+    pub program_brk: usize,
 }
 
 impl ProcessControlBlockInner {
@@ -64,6 +67,29 @@ impl ProcessControlBlockInner {
     pub fn get_task(&self, tid: usize) -> Arc<TaskControlBlock> {
         self.tasks[tid].as_ref().unwrap().clone()
     }
+
+    /// Change the location of the program break. Return None if failed.
+    pub fn change_program_brk(&mut self, size: i32) -> Option<usize> {
+        let old_break = self.program_brk;
+        let new_brk = self.program_brk as isize + size as isize;
+        if new_brk < self.heap_bottom as isize {
+            return None;
+        }
+        let heap_bottom = self.heap_bottom;
+        let result = if size < 0 {
+            self.memory_set
+                .shrink_to(VirtAddr(heap_bottom), VirtAddr(new_brk as usize))
+        } else {
+            self.memory_set
+                .append_to(VirtAddr(heap_bottom), VirtAddr(new_brk as usize))
+        };
+        if result {
+            self.program_brk = new_brk as usize;
+            Some(old_break)
+        } else {
+            None
+        }
+    }
 }
 
 impl ProcessControlBlock {
@@ -73,7 +99,19 @@ impl ProcessControlBlock {
 
     pub fn new(elf_data: &[u8]) -> Arc<Self> {
         // memory_set with elf program headers/trampoline/trap context/user stack
-        let (memory_set, ustack_base, entry_point) = MemorySet::from_elf(elf_data);
+        let (mut memory_set, ustack_base, entry_point) = MemorySet::from_elf(elf_data);
+        // heap starts right after the main thread's user stack
+        let heap_bottom = ustack_base + USER_STACK_SIZE;
+        // map an initial empty heap area (grows via sbrk)
+        memory_set.push(
+            MapArea::new(
+                heap_bottom.into(),
+                heap_bottom.into(),
+                MapType::Framed,
+                MapPermission::R | MapPermission::W | MapPermission::U,
+            ),
+            None,
+        );
         // allocate a pid
         let pid_handle = pid_alloc();
         let process = Arc::new(Self {
@@ -99,6 +137,8 @@ impl ProcessControlBlock {
                     mutex_list: Vec::new(),
                     semaphore_list: Vec::new(),
                     condvar_list: Vec::new(),
+                    heap_bottom,
+                    program_brk: heap_bottom,
                 })
             },
         });
@@ -135,10 +175,27 @@ impl ProcessControlBlock {
     pub fn exec(self: &Arc<Self>, elf_data: &[u8], args: Vec<String>) {
         assert_eq!(self.inner_exclusive_access().thread_count(), 1);
         // memory_set with elf program headers/trampoline/trap context/user stack
-        let (memory_set, ustack_base, entry_point) = MemorySet::from_elf(elf_data);
+        let (mut memory_set, ustack_base, entry_point) = MemorySet::from_elf(elf_data);
+        // heap starts right after the main thread's user stack
+        let heap_bottom = ustack_base + USER_STACK_SIZE;
+        // map an initial empty heap area (grows via sbrk)
+        memory_set.push(
+            MapArea::new(
+                heap_bottom.into(),
+                heap_bottom.into(),
+                MapType::Framed,
+                MapPermission::R | MapPermission::W | MapPermission::U,
+            ),
+            None,
+        );
         let new_token = memory_set.token();
         // substitute memory_set
-        self.inner_exclusive_access().memory_set = memory_set;
+        let mut process_inner = self.inner_exclusive_access();
+        process_inner.memory_set = memory_set;
+        // initialize heap info
+        process_inner.heap_bottom = heap_bottom;
+        process_inner.program_brk = heap_bottom;
+        drop(process_inner);
         // then we alloc user resource for main thread again
         // since memory_set has been changed
         let task = self.inner_exclusive_access().get_task(0);
@@ -218,6 +275,8 @@ impl ProcessControlBlock {
                     mutex_list: Vec::new(),
                     semaphore_list: Vec::new(),
                     condvar_list: Vec::new(),
+                    heap_bottom: parent.heap_bottom,
+                    program_brk: parent.program_brk,
                 })
             },
         });

user/src/bin/sbrk_test.rs

@@ -0,0 +1,48 @@
+#![no_std]
+#![no_main]
+
+#[macro_use]
+extern crate user_lib;
+
+use core::ptr::slice_from_raw_parts_mut;
+use user_lib::sbrk;
+
+#[unsafe(no_mangle)]
+fn main() -> i32 {
+    println!("Test sbrk start.");
+    const PAGE_SIZE: usize = 0x1000;
+    let origin_brk = sbrk(0);
+    println!("origin break point = {:x}", origin_brk);
+    let brk = sbrk(PAGE_SIZE as i32);
+    if brk != origin_brk {
+        return -1;
+    }
+    let brk = sbrk(0);
+    println!("one page allocated,  break point = {:x}", brk);
+    println!("try write to allocated page");
+    let new_page = unsafe {
+        &mut *slice_from_raw_parts_mut(origin_brk as usize as *const u8 as *mut u8, PAGE_SIZE)
+    };
+    for pos in 0..PAGE_SIZE {
+        new_page[pos] = 1;
+    }
+    println!("write ok");
+    sbrk(PAGE_SIZE as i32 * 10);
+    let brk = sbrk(0);
+    println!("10 page allocated,  break point = {:x}", brk);
+    sbrk(PAGE_SIZE as i32 * -11);
+    let brk = sbrk(0);
+    println!("11 page DEALLOCATED,  break point = {:x}", brk);
+    println!("try DEALLOCATED more one page, should be failed.");
+    let ret = sbrk(PAGE_SIZE as i32 * -1);
+    if ret != -1 {
+        println!("Test sbrk failed!");
+        return -1;
+    }
+    println!("Test sbrk almost OK!");
+    println!("now write to deallocated page, should cause page fault.");
+    for pos in 0..PAGE_SIZE {
+        new_page[pos] = 2;
+    }
+    0
+}

user/src/syscall.rs

@@ -13,6 +13,7 @@ const SYSCALL_YIELD: usize = 124;
 const SYSCALL_KILL: usize = 129;
 const SYSCALL_GET_TIME: usize = 169;
 const SYSCALL_GETPID: usize = 172;
+const SYSCALL_SBRK: usize = 214;
 const SYSCALL_FORK: usize = 220;
 const SYSCALL_EXEC: usize = 221;
 const SYSCALL_WAITPID: usize = 260;
@@ -115,6 +116,10 @@ pub fn sys_getpid() -> isize {
     syscall(SYSCALL_GETPID, [0, 0, 0])
 }
 
+pub fn sys_sbrk(size: i32) -> isize {
+    syscall(SYSCALL_SBRK, [size as usize, 0, 0])
+}
+
 pub fn sys_fork() -> isize {
     syscall(SYSCALL_FORK, [0, 0, 0])
 }

user/src/task.rs

@@ -81,3 +81,7 @@ pub fn waittid(tid: usize) -> isize {
         }
     }
 }
+
+pub fn sbrk(size: i32) -> isize {
+    sys_sbrk(size)
+}