Fix brk not reset memory

src/libos/src/vm/process_vm.rs

@@ -11,7 +11,6 @@ use super::vm_util::{
     FileBacked, VMInitializer, VMMapAddr, VMMapOptions, VMMapOptionsBuilder, VMRemapOptions,
 };
 use std::collections::HashSet;
-use std::sync::atomic::{AtomicUsize, Ordering};
 
 // Used for heap and stack start address randomization.
 const RANGE_FOR_RANDOMIZATION: usize = 256 * 4096; // 1M
@@ -162,7 +161,7 @@ impl<'a, 'b> ProcessVMBuilder<'a, 'b> {
         })?;
         debug_assert!(heap_range.start() % heap_layout.align() == 0);
         trace!("heap range = {:?}", heap_range);
-        let brk = AtomicUsize::new(heap_range.start());
+        let brk = RwLock::new(heap_range.start());
         chunks.insert(chunk_ref);
 
         // Init the stack memory in the process
@@ -274,7 +273,7 @@ pub struct ProcessVM {
     elf_ranges: Vec<VMRange>,
     heap_range: VMRange,
     stack_range: VMRange,
-    brk: AtomicUsize,
+    brk: RwLock<usize>,
     // Memory safety notes: the mem_chunks field must be the last one.
     //
     // Rust drops fields in the same order as they are declared. So by making
@@ -441,25 +440,40 @@ impl ProcessVM {
     }
 
     pub fn get_brk(&self) -> usize {
-        self.brk.load(Ordering::SeqCst)
+        *self.brk.read().unwrap()
     }
 
-    pub fn brk(&self, new_brk: usize) -> Result<usize> {
+    pub fn brk(&self, brk: usize) -> Result<usize> {
         let heap_start = self.heap_range.start();
         let heap_end = self.heap_range.end();
 
-        if new_brk >= heap_start && new_brk <= heap_end {
+        // Acquire lock first to avoid data-race.
-            self.brk
+        let mut brk_guard = self.brk.write().unwrap();
-                .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |old_brk| Some(new_brk));
+
-            Ok(new_brk)
+        if brk >= heap_start && brk <= heap_end {
+            // Get page-aligned brk address.
+            let new_brk = align_up(brk, PAGE_SIZE);
+            // Get page-aligned old brk address.
+            let old_brk = align_up(*brk_guard, PAGE_SIZE);
+
+            // Reset the memory when brk shrinks.
+            if new_brk < old_brk {
+                let shrink_brk_range =
+                    VMRange::new(new_brk, old_brk).expect("shrink brk range must be valid");
+                USER_SPACE_VM_MANAGER.reset_memory(shrink_brk_range)?;
+            }
+
+            // Return the user-specified brk address without page aligned. This is same as Linux.
+            *brk_guard = brk;
+            Ok(brk)
         } else {
-            if new_brk < heap_start {
+            if brk < heap_start {
                 error!("New brk address is too low");
-            } else if new_brk > heap_end {
+            } else if brk > heap_end {
                 error!("New brk address is too high");
             }
 
-            Ok(self.get_brk())
+            Ok(*brk_guard)
         }
     }
 

src/libos/src/vm/vm_manager.rs

@@ -362,6 +362,48 @@ impl VMManager {
         }
     }
 
+    // Reset memory permission to default (R/W) and reset the memory contents to zero. Currently only used by brk.
+    pub fn reset_memory(&self, reset_range: VMRange) -> Result<()> {
+        let intersect_chunks = {
+            let chunks = self
+                .internal()
+                .chunks
+                .iter()
+                .filter(|&chunk| chunk.range().intersect(&reset_range).is_some())
+                .map(|chunk| chunk.clone())
+                .collect::<Vec<_>>();
+
+            // In the heap area, there shouldn't be any default chunks or chunks owned by other process.
+            if chunks
+                .iter()
+                .any(|chunk| !chunk.is_owned_by_current_process() || !chunk.is_single_vma())
+            {
+                return_errno!(EINVAL, "There is something wrong with the intersect chunks");
+            }
+            chunks
+        };
+
+        intersect_chunks.iter().for_each(|chunk| {
+            if let ChunkType::SingleVMA(vma) = chunk.internal() {
+                if let Some(intersection_range) = chunk.range().intersect(&reset_range) {
+                    let mut internal_manager = self.internal();
+                    internal_manager.mprotect_single_vma_chunk(
+                        &chunk,
+                        intersection_range,
+                        VMPerms::DEFAULT,
+                    );
+
+                    unsafe {
+                        let buf = intersection_range.as_slice_mut();
+                        buf.iter_mut().for_each(|b| *b = 0)
+                    }
+                }
+            }
+        });
+
+        Ok(())
+    }
+
     pub fn msync(&self, addr: usize, size: usize) -> Result<()> {
         let sync_range = VMRange::new_with_size(addr, size)?;
         let chunk = {

test/Makefile

@@ -22,7 +22,7 @@ TESTS ?= env empty hello_world malloc mmap file fs_perms getpid spawn sched pipe
 	truncate readdir mkdir open stat link symlink chmod chown tls pthread system_info rlimit \
 	server server_epoll unix_socket cout hostfs cpuid rdtsc device sleep exit_group posix_flock \
 	ioctl fcntl eventfd emulate_syscall access signal sysinfo prctl rename procfs wait \
-	spawn_attribute exec statfs random umask pgrp vfork mount flock utimes shm epoll
+	spawn_attribute exec statfs random umask pgrp vfork mount flock utimes shm epoll brk
 # Benchmarks: need to be compiled and run by bench-% target
 BENCHES := spawn_and_exit_latency pipe_throughput unix_socket_throughput
 

test/brk/Makefile

@@ -0,0 +1,5 @@
+include ../test_common.mk
+
+EXTRA_C_FLAGS :=
+EXTRA_LINK_FLAGS :=
+BIN_ARGS :=

test/brk/main.c

@@ -0,0 +1,195 @@
+#include <unistd.h>
+#include <stdlib.h>
+#include <sys/syscall.h>
+#include <stdint.h>
+#include <sys/mman.h>
+#include "test.h"
+
+// ============================================================================
+// Helper function
+// ============================================================================
+
+#define PAGE_SIZE 4096
+
+const static uint8_t magic_num_01 = 0xFF;
+typedef struct syscall_args {
+    int num;
+    unsigned long arg0;
+} syscall_args_t;
+
+static inline uint64_t native_syscall(syscall_args_t *p) {
+    uint64_t ret;
+    register int num asm ("rax") = p->num;
+    register unsigned long arg0 asm ("rdi") = p->arg0;
+
+    asm volatile("syscall"
+                 : "=a" (ret)
+                 : "r" (num), "r" (arg0));
+    return ret;
+}
+
+static uint64_t brk_syscall(uint64_t brk) {
+    syscall_args_t brk_arg = {
+        .num = __NR_brk,
+        .arg0 = brk,
+    };
+
+    return native_syscall(&brk_arg);
+}
+
+// ============================================================================
+// Test cases for access
+// ============================================================================
+
+static int test_brk_shrinks() {
+    char *zero_buf = malloc(PAGE_SIZE * 2);
+    if (zero_buf == NULL) {
+        THROW_ERROR("malloc failed");
+    }
+    memset(zero_buf, 0, PAGE_SIZE * 2);
+
+    uint64_t original_brk = brk_syscall(0);
+    if (original_brk == 0) {
+        THROW_ERROR("sbrk failed");
+    }
+    printf("original brk = %lx\n", original_brk);
+
+    // increase brk
+    printf("increase brk\n");
+    uint64_t ret = brk_syscall(original_brk + PAGE_SIZE * 4);
+    if (ret == 0) {
+        THROW_ERROR("extend brk failed");
+    }
+
+    // set some values to the brk memory
+    uint64_t test_range_start = original_brk + PAGE_SIZE * 2;
+    for (int i = 0; i < PAGE_SIZE; i++) {
+        *(int *)test_range_start = magic_num_01;
+    }
+
+    // decrease brk
+    printf("decrease brk\n");
+    ret = brk_syscall(original_brk + PAGE_SIZE * 2);
+    if (ret != test_range_start) {
+        THROW_ERROR("shrink brk failed");
+    }
+    printf("test range start = %lx\n", test_range_start);
+
+    // increase brk
+    uint64_t test_range_end = brk_syscall(original_brk + PAGE_SIZE * 4);
+    if (test_range_end != original_brk + PAGE_SIZE * 4) {
+        THROW_ERROR("extend brk failed");
+    }
+
+    if ( memcmp((const void *)test_range_start, zero_buf, PAGE_SIZE * 2) != 0) {
+        THROW_ERROR("sbrk not reset memory");
+    }
+
+    free(zero_buf);
+
+    return 0;
+}
+
+#ifdef SGX_MODE_HW
+// This test case will fail in simulation mode. Because the raw syscall interface are not handled by Occlum
+// in simulation mode.
+//
+// Use brk to allocate 4 pages and test brk and mprotect
+//    original brk
+//       | page 00          page 02
+//       |         page 01          page 03
+// ...---|-------|-------|-------|-------|
+static int test_brk_shrinks_spans_multiple_chunks() {
+    const static uint8_t magic_num_02 = 0xFE;
+    char *zero_buf = malloc(PAGE_SIZE * 4);
+    if (zero_buf == NULL) {
+        THROW_ERROR("malloc failed");
+    }
+    memset(zero_buf, 0, PAGE_SIZE * 4);
+
+    size_t original_brk = brk_syscall(0);
+    if (original_brk == 0) {
+        THROW_ERROR("brk failed");
+    }
+    printf("original brk = %lx\n", original_brk);
+
+    // increase brk to the end of page 03
+    size_t ret = brk_syscall(original_brk + PAGE_SIZE * 4);
+    if (ret != original_brk + PAGE_SIZE * 4) {
+        THROW_ERROR("extend brk failed");
+    }
+
+    // set some values to the brk memory page 02
+    size_t test_range_start = original_brk + PAGE_SIZE * 2;
+    for (int i = 0; i < PAGE_SIZE; i++) {
+        *(int *)test_range_start = magic_num_01;
+    }
+
+    // mprotect page 01 - 03 to PROT_NONE and decrease brk to the end of page 00
+    int rc = mprotect((void *)(original_brk + PAGE_SIZE * 1), PAGE_SIZE * 3, PROT_NONE);
+    if (rc < 0) {
+        THROW_ERROR("mprotect failure");
+    }
+    ret = brk_syscall(original_brk + PAGE_SIZE * 1);
+    if (ret != original_brk + PAGE_SIZE * 1) {
+        THROW_ERROR("shrink brk failed");
+    }
+
+    // increase brk to the end of page 02
+    ret = brk_syscall(original_brk + PAGE_SIZE * 3);
+    if (ret != original_brk + PAGE_SIZE * 3) {
+        THROW_ERROR("extend brk failed");
+    }
+
+    // set some values to the brk memory page 01
+    test_range_start = original_brk + PAGE_SIZE * 1;
+    for (int i = 0; i < PAGE_SIZE; i++) {
+        *(int *)test_range_start = magic_num_02;
+    }
+
+    // decrease brk again to the end of page 00
+    rc = mprotect((void *)(original_brk + PAGE_SIZE * 1), PAGE_SIZE * 2, PROT_NONE);
+    if (rc < 0) {
+        THROW_ERROR("mprotect failure");
+    }
+    ret = brk_syscall(original_brk + PAGE_SIZE * 1);
+    if (ret != original_brk + PAGE_SIZE * 1) {
+        THROW_ERROR("shrink brk failed");
+    }
+
+    // increase brk to the end of page 03
+    ret = brk_syscall(original_brk + PAGE_SIZE * 4);
+    if (ret != original_brk + PAGE_SIZE * 4) {
+        THROW_ERROR("extend brk failed");
+    }
+
+    if ( memcmp((const void *)original_brk, zero_buf, PAGE_SIZE * 4) != 0) {
+        THROW_ERROR("brk not reset memory");
+    }
+
+    // decrease brk to the original brk
+    ret = brk_syscall(original_brk);
+    if (ret != original_brk) {
+        THROW_ERROR("shrink brk failed");
+    }
+
+    free(zero_buf);
+
+    return 0;
+}
+#endif
+
+// ============================================================================
+// Test suite main
+// ============================================================================
+
+static test_case_t test_cases[] = {
+    TEST_CASE(test_brk_shrinks),
+#ifdef SGX_MODE_HW
+    TEST_CASE(test_brk_shrinks_spans_multiple_chunks),
+#endif
+};
+
+int main(int argc, const char *argv[]) {
+    return test_suite_run(test_cases, ARRAY_SIZE(test_cases));
+}