[libos] Add PKU support

src/Enclave.edl

@@ -173,6 +173,10 @@ enclave {
 
         int occlum_ocall_mprotect([user_check] void* addr, size_t len, int prot);
 
+        int occlum_ocall_pkey_alloc(unsigned int flags, unsigned int access_rights);
+        int occlum_ocall_pkey_mprotect([user_check] void* addr, size_t len, int prot, int pkey);
+        int occlum_ocall_pkey_free(int pkey);
+
         int occlum_ocall_get_numa_topology(
             [out, count=ncpus] uint32_t *numa_buf,
             size_t ncpus

src/libos/Cargo.toml

@@ -43,6 +43,7 @@ sgx1_exception_sim = [] # Simulate #PF and #GP exceptions on SGX 1
 dcap = []               # DCAP support. The compilation relies on DCAP package.
 cov = ["sgx_cov"]       # Enable coverage colletcion.
 hyper_mode = []         # For running in hyper mode.
+pku = []                # PKU Support
 
 [target.'cfg(not(target_env = "sgx"))'.dependencies]
 sgx_types = { path = "../../deps/rust-sgx-sdk/sgx_types" }

src/libos/Makefile

@@ -45,9 +45,12 @@ LIBOS_LOG ?= error
 
 LIBOS_SONAME := libocclum-libos.so.$(MAJOR_VER_NUM)
 
+LIBOS_FEATURES :=
+
 ifeq ($(SGX_MODE), HW)
 	LIBOS_CORE_LIB_NAME := occlum-libos-core
 	LIBOS_SO_REAL := $(BUILD_DIR)/lib/libocclum-libos.so.$(VERSION_NUM)
+	LIBOS_FEATURES += pku
 else ifeq ($(SGX_MODE), HYPER)
 	LIBOS_CORE_LIB_NAME := occlum-libos-core_hyper
 	LIBOS_SO_REAL := $(BUILD_DIR)/lib/libocclum-libos_hyper.so.$(VERSION_NUM)
@@ -122,8 +125,6 @@ else
     RUSTC_WRAPPER :=
 endif
 
-LIBOS_FEATURES :=
-
 ifndef OCCLUM_DISABLE_DCAP
     LIBOS_FEATURES += dcap
 endif

src/libos/include/task.h

@@ -37,6 +37,9 @@ void do_exit_task(void);
 
 #else  /* __ASSEMBLY__ */
 
+#define PKRU_LIBOS          (0x0)
+#define PKRU_USER           (0x55555551)
+
 /* See /<path-to-linux-sgx>/common/inc/internal/thread_data.h */
 #define TD_STACKGUARD_OFFSET        (8 * 5)
 /* Override the field for stack guard */

src/libos/src/entry.rs

@@ -70,11 +70,19 @@ pub extern "C" fn occlum_ecall_init(
         // Init the log infrastructure first so that log messages will be printed afterwards
         util::log::init(log_level);
 
-        // Init MPX for SFI if MPX is available
         let report = rsgx_self_report();
+        // Init MPX for SFI if MPX is available
         if (report.body.attributes.xfrm & SGX_XFRM_MPX != 0) {
             util::mpx_util::mpx_enable();
         }
+        // Init PKU for isolating LibOS form user-apps if PKU is available
+        // Occlum only turns on `pku` feature in HW mode
+        #[cfg(feature = "pku")]
+        {
+            if (report.body.attributes.xfrm & SGX_XFRM_PKRU != 0) {
+                crate::util::pku_util::try_set_pku_enabled();
+            }
+        }
 
         // Register exception handlers (support cpuid & rdtsc for now)
         register_exception_handlers();

src/libos/src/process/do_spawn/mod.rs

@@ -14,6 +14,7 @@ use crate::fs::{
 };
 use crate::prelude::*;
 use crate::process::pgrp::{get_spawn_attribute_pgrp, update_pgrp_for_new_process};
+use crate::util::pku_util;
 use crate::vm::ProcessVM;
 
 mod aux_vec;
@@ -447,7 +448,11 @@ fn init_auxvec(process_vm: &ProcessVM, exec_elf_file: &ElfFile) -> Result<AuxVec
     let ldso_elf_base = process_vm.get_elf_ranges()[1].start() as u64;
     auxvec.set(AuxKey::AT_BASE, ldso_elf_base)?;
 
-    let syscall_addr = __occlum_syscall_linux_abi as *const () as u64;
+    let syscall_addr = if pku_util::check_pku_enabled() {
+        __occlum_syscall_linux_pku_abi
+    } else {
+        __occlum_syscall_linux_abi
+    } as *const () as u64;
     auxvec.set(AuxKey::AT_OCCLUM_ENTRY, syscall_addr)?;
     // TODO: init AT_EXECFN
     // auxvec.set_val(AuxKey::AT_EXECFN, "program_name")?;
@@ -457,5 +462,6 @@ fn init_auxvec(process_vm: &ProcessVM, exec_elf_file: &ElfFile) -> Result<AuxVec
 
 extern "C" {
     fn __occlum_syscall_linux_abi() -> i64;
+    fn __occlum_syscall_linux_pku_abi() -> i64;
     fn occlum_gdb_hook_load_elf(elf_base: u64, elf_path: *const u8, elf_path_len: u64);
 }

src/libos/src/process/task/task_x86-64.S

@@ -68,6 +68,20 @@ __exec_task:
 
     // Use user stack
     movq TASK_USER_RSP(%rdi), %rsp
-    // Run user code
+    // Get user code address
     movq TASK_USER_ENTRY_ADDR(%rdi), %r11
+
+    // Whether to switch PKRU value
+    mov pku_enabled(%rip), %r10
+    cmp $1, %r10
+    je update_pkru_in_exec_task
+
+    // Run user code
+    jmp *%r11
+
+update_pkru_in_exec_task:
+    xor %ecx, %ecx
+    xor %edx, %edx
+    mov $PKRU_USER, %eax
+    wrpkru
     jmp *%r11

src/libos/src/syscall/syscall_entry_x86-64.S

@@ -3,6 +3,22 @@
 
 
     .file "syscall_entry_x86-64.S"
+    .global __occlum_syscall_linux_pku_abi
+    .type __occlum_syscall_linux_pku_abi, @function
+__occlum_syscall_linux_pku_abi:
+    pushq %rcx
+    pushq %rdx
+    pushq %rax
+
+    xor %ecx, %ecx
+    xor %edx, %edx
+    mov $PKRU_LIBOS, %eax
+    wrpkru
+
+    popq %rax
+    popq %rdx
+    popq %rcx
+
     .global __occlum_syscall_linux_abi
     .type __occlum_syscall_linux_abi, @function
 __occlum_syscall_linux_abi:
@@ -128,7 +144,42 @@ __occlum_sysret:
     pop %rax
     pop %rcx
     pop %rsp
+
+    // Store RFLAGS since `cmp` operation may overwrite it
+    pushfq
+    push %rax
+
+    mov pku_enabled(%rip), %rax
+    cmp $1, %rax
+    je update_pkru_in_sysret
+
+    pop %rax
+    popfq
+
     jmp *%gs:(TD_SYSCALL_RET_ADDR_OFFSET)
+    // This should never happen
+    ud2
+
+update_pkru_in_sysret:
+    pop %rax
+    popfq
+
+    sub $0x20, %rsp
+    mov %rax, (%rsp)
+    mov %rdx, 0x8(%rsp)
+    mov %rcx, 0x10(%rsp)
+    mov %gs:(TD_SYSCALL_RET_ADDR_OFFSET), %rcx
+    mov %rcx, 0x18(%rsp)
+
+    xor %ecx, %ecx
+    xor %edx, %edx
+    mov $PKRU_USER, %eax
+    wrpkru
+
+    pop %rax
+    pop %rdx
+    pop %rcx
+    ret
 
     .global __occlum_syscall_c_abi
     .type __occlum_syscall_c_abi, @function

src/libos/src/util/mod.rs

@@ -6,5 +6,6 @@ pub mod hosts_parser_util;
 pub mod log;
 pub mod mem_util;
 pub mod mpx_util;
+pub mod pku_util;
 pub mod sgx;
 pub mod sync;

src/libos/src/util/pku_util.rs

@@ -0,0 +1,101 @@
+use super::*;
+
+use std::sync::atomic::{AtomicBool, Ordering};
+
+/// Status variable accessed by assembly code
+#[no_mangle]
+pub static mut pku_enabled: u64 = 0;
+
+lazy_static! {
+    pub static ref PKU_ENABLED: AtomicBool = AtomicBool::new(false);
+}
+
+const PKEY_LIBOS: i32 = 0;
+const PKEY_USER: i32 = 1;
+
+/// Try enable PKU features in Occlum.
+pub fn try_set_pku_enabled() {
+    // Alloc pkey
+    let mut pkey = -1;
+    let sgx_status = unsafe { occlum_ocall_pkey_alloc(&mut pkey, 0, 0) };
+    assert!(sgx_status == sgx_status_t::SGX_SUCCESS && pkey == PKEY_USER);
+
+    unsafe {
+        pku_enabled = 1;
+    }
+    assert!(PKU_ENABLED.load(Ordering::Relaxed) == false);
+    PKU_ENABLED.store(true, Ordering::Release);
+}
+
+pub fn check_pku_enabled() -> bool {
+    PKU_ENABLED.load(Ordering::Acquire)
+}
+
+pub fn pkey_mprotect_userspace_mem(user_mem_base: usize, user_mem_len: usize, perm: i32) {
+    if !self::check_pku_enabled() {
+        return;
+    }
+    let mut retval = -1;
+    debug!(
+        "associate memory region: 0x{:x} -> 0x{:x}, size: 0x{:x} with pkey for userspace: {:?}",
+        user_mem_base,
+        user_mem_base + user_mem_len,
+        user_mem_len,
+        PKEY_USER
+    );
+    let sgx_status = unsafe {
+        occlum_ocall_pkey_mprotect(
+            &mut retval,
+            user_mem_base as *const c_void,
+            user_mem_len,
+            perm,
+            PKEY_USER,
+        )
+    };
+    assert!(sgx_status == sgx_status_t::SGX_SUCCESS && retval == 0);
+}
+
+pub fn clear_pku_when_libos_exit(user_mem_base: usize, user_mem_len: usize, perm: i32) {
+    if !self::check_pku_enabled() {
+        return;
+    }
+    let mut retval = -1;
+    debug!(
+        "re-associate memory region  0x{:x} -> 0x{:x}, size: 0x{:x} with pkey for libos: {:?}",
+        user_mem_base,
+        user_mem_base + user_mem_len,
+        user_mem_len,
+        PKEY_LIBOS
+    );
+    let sgx_status = unsafe {
+        occlum_ocall_pkey_mprotect(
+            &mut retval,
+            user_mem_base as *const c_void,
+            user_mem_len,
+            perm,
+            PKEY_LIBOS,
+        )
+    };
+    assert!(sgx_status == sgx_status_t::SGX_SUCCESS && retval == 0);
+    debug!("free pkey: {:?}", PKEY_USER);
+    let sgx_status = unsafe { occlum_ocall_pkey_free(&mut retval, PKEY_USER) };
+    assert!(sgx_status == sgx_status_t::SGX_SUCCESS && retval == 0);
+}
+
+extern "C" {
+    pub fn occlum_ocall_pkey_alloc(
+        retval: *mut i32,
+        flags: u32,
+        access_rights: u32,
+    ) -> sgx_status_t;
+
+    pub fn occlum_ocall_pkey_mprotect(
+        retval: *mut i32,
+        addr: *const c_void,
+        len: usize,
+        prot: i32,
+        pkey: i32,
+    ) -> sgx_status_t;
+
+    pub fn occlum_ocall_pkey_free(retval: *mut i32, pkey: i32) -> sgx_status_t;
+}

src/libos/src/vm/user_space_vm.rs

@@ -1,10 +1,14 @@
 use super::ipc::SHM_MANAGER;
 use super::*;
 use crate::ctor::dtor;
+use crate::util::pku_util;
 use config::LIBOS_CONFIG;
 use std::ops::{Deref, DerefMut};
 use vm_manager::VMManager;
 
+const RSRV_MEM_PERM: MemPerm =
+    MemPerm::from_bits_truncate(MemPerm::READ.bits() | MemPerm::WRITE.bits());
+
 /// The virtual memory manager for the entire user space
 pub struct UserSpaceVMManager(VMManager);
 
@@ -16,13 +20,13 @@ impl UserSpaceVMManager {
             // a lot of time. When EDMM is supported, there is no need to commit all the pages at the initialization stage. A function
             // which reserves memory but not commit pages should be provided then.
             let ptr = sgx_alloc_rsrv_mem(rsrv_mem_size);
-            let perm = MemPerm::READ | MemPerm::WRITE;
             if ptr.is_null() {
                 return_errno!(ENOMEM, "run out of reserved memory");
             }
             // Change the page permission to RW (default)
             assert!(
-                sgx_tprotect_rsrv_mem(ptr, rsrv_mem_size, perm.bits()) == sgx_status_t::SGX_SUCCESS
+                sgx_tprotect_rsrv_mem(ptr, rsrv_mem_size, RSRV_MEM_PERM.bits())
+                    == sgx_status_t::SGX_SUCCESS
             );
 
             let addr = ptr as usize;
@@ -30,6 +34,7 @@ impl UserSpaceVMManager {
                 "allocated rsrv addr is 0x{:x}, len is 0x{:x}",
                 addr, rsrv_mem_size
             );
+            pku_util::pkey_mprotect_userspace_mem(addr, rsrv_mem_size, RSRV_MEM_PERM.bits());
             VMRange::new(addr, addr + rsrv_mem_size)?
         };
 
@@ -50,10 +55,11 @@ fn free_user_space() {
     SHM_MANAGER.clean_when_libos_exit();
     let range = USER_SPACE_VM_MANAGER.range();
     assert!(USER_SPACE_VM_MANAGER.verified_clean_when_exit());
-    let addr = range.start() as *const c_void;
+    let addr = range.start();
     let size = range.size();
     info!("free user space VM: {:?}", range);
-    assert!(unsafe { sgx_free_rsrv_mem(addr, size) == 0 });
+    pku_util::clear_pku_when_libos_exit(addr, size, RSRV_MEM_PERM.bits());
+    assert!(unsafe { sgx_free_rsrv_mem(addr as *const c_void, size) == 0 });
 }
 
 impl Deref for UserSpaceVMManager {

src/pal/src/ocalls/mem.c

@@ -1,3 +1,4 @@
+#define _GNU_SOURCE
 #include <stdlib.h>
 #include <sys/mman.h>
 #include "ocalls.h"
@@ -30,3 +31,15 @@ void occlum_ocall_free(void *ptr) {
 int occlum_ocall_mprotect(void *addr, size_t len, int prot) {
     return mprotect(addr, len, prot);
 }
+
+int occlum_ocall_pkey_alloc(unsigned int flags, unsigned int access_rights) {
+    return pkey_alloc(flags, access_rights);
+}
+
+int occlum_ocall_pkey_mprotect(void *addr, size_t len, int prot, int pkey) {
+    return pkey_mprotect(addr, len, prot, pkey);
+}
+
+int occlum_ocall_pkey_free(int pkey) {
+    return pkey_free(pkey);
+}