Re-enable PKU support

Also add a testcase for PKU support

.github/workflows/hw_mode_test.yml

@@ -166,7 +166,7 @@ jobs:
     runs-on: ${{ matrix.self_runner }}
     strategy:
       matrix:
-        self_runner: [[self-hosted, SGX2-HW]]
+        self_runner: [[self-hosted, SGX2-HW, PKU]]
 
     steps:
     - name: Clean before running
@@ -205,6 +205,9 @@ jobs:
     - name: Run processBuilder
       run: docker exec ${{ env.CONTAINER_NAME }} bash -c "cd /root/occlum/demos/java && ./run_java_on_occlum.sh processBuilder"
 
+    - name: Run hello PKU
+      run: docker exec ${{ env.CONTAINER_NAME }} bash -c "cd /root/occlum/demos/java && ./run_java_on_occlum.sh hello_pku"
+
     - name: Clean the environment
       if: ${{ always() }}
       run: docker stop ${{ env.CONTAINER_NAME }}

demos/java/run_java_on_occlum.sh

@@ -6,7 +6,7 @@ NC='\033[0m'
 
 show_usage() {
     echo "Error: invalid arguments"
-    echo "Usage: $0 web_app/hello/processBuilder"
+    echo "Usage: $0 web_app/hello/processBuilder/hello_pku"
     exit 1
 }
 
@@ -33,6 +33,10 @@ init_instance() {
     echo "${new_json}" > Occlum.json
 }
 
+update_pku_config() {
+    new_json="$(jq '.metadata.pkru = 1' Occlum.json)" && echo "${new_json}" > Occlum.json
+}
+
 build_web() {
     # Copy JVM and JAR file into Occlum instance and build
     rm -rf image
@@ -66,6 +70,16 @@ run_hello() {
     occlum run /usr/lib/jvm/java-11-alibaba-dragonwell/jre/bin/java -Xmx512m -XX:-UseCompressedOops -XX:MaxMetaspaceSize=64m -Dos.name=Linux Main
 }
 
+run_hello_pku() {
+    hello=./hello_world/Main.class
+    check_file_exist ${hello}
+    init_instance
+    update_pku_config
+    build_hello
+    echo -e "${BLUE}occlum run JVM hello with PKU enabled${NC}"
+    occlum run /usr/lib/jvm/java-11-alibaba-dragonwell/jre/bin/java -Xmx512m -XX:-UseCompressedOops -XX:MaxMetaspaceSize=64m -Dos.name=Linux Main
+}
+
 build_processBuilder() {
     # Copy JVM and class file into Occlum instance and build
     rm -rf image
@@ -97,6 +111,9 @@ case "$arg" in
     processBuilder)
         run_processBuilder
         ;;
+    hello_pku)
+        run_hello_pku
+        ;;
     *)
         show_usage
 esac

src/libos/src/util/pku_util.rs

@@ -1,5 +1,6 @@
 use super::*;
 
+use crate::vm::{VMPerms, VMRange};
 use std::sync::atomic::{AtomicBool, Ordering};
 
 /// Status variable accessed by assembly code
@@ -31,55 +32,91 @@ 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) {
+pub fn pkey_mprotect_userspace_mem(
+    user_space_range: &VMRange,
+    gap_range: Option<&VMRange>,
+    perm: VMPerms,
+) {
     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_space_range.start(),
-        user_mem_base + user_mem_len,
+        user_space_range.end(),
-        user_mem_len,
+        user_space_range.size(),
         PKEY_USER
     );
-    let sgx_status = unsafe {
+
-        occlum_ocall_pkey_mprotect(
+    pkey_mprotect_user_space(user_space_range, gap_range, perm.bits() as i32, PKEY_USER);
-            &mut retval,
+}
-            user_mem_base as *const c_void,
+
-            user_mem_len,
+pub fn clear_pku_when_libos_exit(
-            perm,
+    user_space_range: &VMRange,
-            PKEY_USER,
+    gap_range: Option<&VMRange>,
-        )
+    perm: VMPerms,
-    };
+) {
+    if !self::check_pku_enabled() {
+        return;
+    }
+    debug!(
+        "re-associate memory region  0x{:x} -> 0x{:x}, size: 0x{:x} with pkey for libos: {:?}",
+        user_space_range.start(),
+        user_space_range.end(),
+        user_space_range.size(),
+        PKEY_LIBOS
+    );
+    pkey_mprotect_user_space(user_space_range, gap_range, perm.bits() as i32, PKEY_LIBOS);
+
+    debug!("free pkey: {:?}", PKEY_USER);
+    let mut retval = -1;
+    let sgx_status = unsafe { occlum_ocall_pkey_free(&mut retval, 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) {
+fn pkey_mprotect_user_space(
-    if !self::check_pku_enabled() {
+    user_space_range: &VMRange,
-        return;
+    gap_range: Option<&VMRange>,
-    }
+    perm: i32,
+    pkey: i32,
+) {
     let mut retval = -1;
-    debug!(
+
-        "re-associate memory region  0x{:x} -> 0x{:x}, size: 0x{:x} with pkey for libos: {:?}",
+    if let Some(gap_range) = gap_range {
-        user_mem_base,
+        // user_space_left
-        user_mem_base + user_mem_len,
+        let user_space_left = VMRange::new(user_space_range.start(), gap_range.start()).unwrap();
-        user_mem_len,
+        let user_space_right = VMRange::new(gap_range.end(), user_space_range.end()).unwrap();
-        PKEY_LIBOS
+        let sgx_status = unsafe {
-    );
+            occlum_ocall_pkey_mprotect(
-    let sgx_status = unsafe {
+                &mut retval,
-        occlum_ocall_pkey_mprotect(
+                user_space_left.start() as *const c_void,
-            &mut retval,
+                user_space_left.size(),
-            user_mem_base as *const c_void,
+                perm,
-            user_mem_len,
+                pkey,
-            perm,
+            )
-            PKEY_LIBOS,
+        };
-        )
+        assert!(sgx_status == sgx_status_t::SGX_SUCCESS && retval == 0);
-    };
+        let sgx_status = unsafe {
-    assert!(sgx_status == sgx_status_t::SGX_SUCCESS && retval == 0);
+            occlum_ocall_pkey_mprotect(
-    debug!("free pkey: {:?}", PKEY_USER);
+                &mut retval,
-    let sgx_status = unsafe { occlum_ocall_pkey_free(&mut retval, PKEY_USER) };
+                user_space_right.start() as *const c_void,
-    assert!(sgx_status == sgx_status_t::SGX_SUCCESS && retval == 0);
+                user_space_right.size(),
+                perm,
+                pkey,
+            )
+        };
+    } else {
+        let sgx_status = unsafe {
+            occlum_ocall_pkey_mprotect(
+                &mut retval,
+                user_space_range.start() as *const c_void,
+                user_space_range.size(),
+                perm,
+                pkey,
+            )
+        };
+        assert!(sgx_status == sgx_status_t::SGX_SUCCESS && retval == 0);
+    }
 }
 
 extern "C" {