From fb12642254563717c053b7ee4b59be3f183d14f9 Mon Sep 17 00:00:00 2001 From: "Hui, Chunyang" Date: Thu, 5 Jan 2023 16:47:44 +0800 Subject: [PATCH] Fix brk not reset memory --- src/libos/src/vm/process_vm.rs | 38 +++++-- src/libos/src/vm/vm_manager.rs | 42 +++++++ test/Makefile | 2 +- test/brk/Makefile | 5 + test/brk/main.c | 195 +++++++++++++++++++++++++++++++++ 5 files changed, 269 insertions(+), 13 deletions(-) create mode 100644 test/brk/Makefile create mode 100644 test/brk/main.c diff --git a/src/libos/src/vm/process_vm.rs b/src/libos/src/vm/process_vm.rs index 1a4d6766..912d4b41 100644 --- a/src/libos/src/vm/process_vm.rs +++ b/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, heap_range: VMRange, stack_range: VMRange, - brk: AtomicUsize, + brk: RwLock, // 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 { + pub fn brk(&self, brk: usize) -> Result { let heap_start = self.heap_range.start(); let heap_end = self.heap_range.end(); - if new_brk >= heap_start && new_brk <= heap_end { - self.brk - .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |old_brk| Some(new_brk)); - Ok(new_brk) + // Acquire lock first to avoid data-race. + let mut brk_guard = self.brk.write().unwrap(); + + 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) } } diff --git a/src/libos/src/vm/vm_manager.rs b/src/libos/src/vm/vm_manager.rs index 8b6df147..e7b3a653 100644 --- a/src/libos/src/vm/vm_manager.rs +++ b/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::>(); + + // 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 = { diff --git a/test/Makefile b/test/Makefile index fc664993..b0ec0f75 100644 --- a/test/Makefile +++ b/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 diff --git a/test/brk/Makefile b/test/brk/Makefile new file mode 100644 index 00000000..9e1b6dec --- /dev/null +++ b/test/brk/Makefile @@ -0,0 +1,5 @@ +include ../test_common.mk + +EXTRA_C_FLAGS := +EXTRA_LINK_FLAGS := +BIN_ARGS := diff --git a/test/brk/main.c b/test/brk/main.c new file mode 100644 index 00000000..7a4ea30e --- /dev/null +++ b/test/brk/main.c @@ -0,0 +1,195 @@ +#include +#include +#include +#include +#include +#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)); +}