From 3c481d12976ef5e922e800dba3ae84066c6f0365 Mon Sep 17 00:00:00 2001
From: "Hui, Chunyang" <sanqian.hcy@antfin.com>
Date: Fri, 22 Sep 2023 09:04:57 +0000
Subject: [PATCH] Add support for AMX

---
 etc/template/Occlum.json               |  3 ++-
 src/libos/src/process/do_arch_prctl.rs | 25 +++++++++++++++++++++++++
 src/libos/src/process/syscalls.rs      |  1 -
 tools/gen_internal_conf/src/main.rs    |  4 ++++
 4 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/etc/template/Occlum.json b/etc/template/Occlum.json
index be56e8b4..64b4b4de 100644
--- a/etc/template/Occlum.json
+++ b/etc/template/Occlum.json
@@ -34,7 +34,8 @@
             "high": "0x0",
             "low": "0x0"
         },
-        "pkru": 0
+        "pkru": 0,
+        "amx": 0
     },
     "mount": [
         {
diff --git a/src/libos/src/process/do_arch_prctl.rs b/src/libos/src/process/do_arch_prctl.rs
index 1a3ce496..9e54e066 100644
--- a/src/libos/src/process/do_arch_prctl.rs
+++ b/src/libos/src/process/do_arch_prctl.rs
@@ -1,17 +1,40 @@
 use crate::prelude::*;
+use crate::util::mem_util::from_user::check_mut_ptr;
 
 pub fn do_arch_prctl(code: ArchPrctlCode, addr: *mut usize) -> Result<()> {
     debug!("do_arch_prctl: code: {:?}, addr: {:?}", code, addr);
     match code {
         ArchPrctlCode::ARCH_SET_FS => {
+            check_mut_ptr(addr)?;
             current!().task().set_user_fs(addr as usize);
         }
         ArchPrctlCode::ARCH_GET_FS => unsafe {
+            check_mut_ptr(addr)?;
             *addr = current!().task().user_fs();
         },
         ArchPrctlCode::ARCH_SET_GS | ArchPrctlCode::ARCH_GET_GS => {
+            check_mut_ptr(addr)?;
             return_errno!(EINVAL, "GS cannot be accessed from the user space");
         }
+        ArchPrctlCode::ARCH_REQ_XCOMP_PERM => {
+            // Allows to request permission for a dynamically enabled feature or a feature set
+            // Currently only used to enable AMX
+            use crate::util::sgx::get_self_target;
+            const XFEATURE_XTILEDATA: u64 = 18;
+
+            let features = addr as u64;
+            if features == XFEATURE_XTILEDATA {
+                // Check if AMX is enabled for current Enclave
+                let target_info = get_self_target()?;
+                if target_info.attributes.xfrm & SGX_XFRM_AMX != SGX_XFRM_AMX {
+                    return_errno!(EINVAL, "AMX is not enabled for this enclave");
+                } else {
+                    info!("AMX is enabled for this enclave");
+                }
+            } else {
+                return_errno!(ENOSYS, "feature not supported");
+            }
+        }
     }
     Ok(())
 }
@@ -23,6 +46,7 @@ pub enum ArchPrctlCode {
     ARCH_SET_FS = 0x1002,
     ARCH_GET_FS = 0x1003,
     ARCH_GET_GS = 0x1004,
+    ARCH_REQ_XCOMP_PERM = 0x1023,
 }
 
 impl ArchPrctlCode {
@@ -32,6 +56,7 @@ impl ArchPrctlCode {
             0x1002 => Ok(ArchPrctlCode::ARCH_SET_FS),
             0x1003 => Ok(ArchPrctlCode::ARCH_GET_FS),
             0x1004 => Ok(ArchPrctlCode::ARCH_GET_GS),
+            0x1023 => Ok(ArchPrctlCode::ARCH_REQ_XCOMP_PERM),
             _ => return_errno!(EINVAL, "Unknown code for arch_prctl"),
         }
     }
diff --git a/src/libos/src/process/syscalls.rs b/src/libos/src/process/syscalls.rs
index 75f6c639..563c8a69 100644
--- a/src/libos/src/process/syscalls.rs
+++ b/src/libos/src/process/syscalls.rs
@@ -345,7 +345,6 @@ pub fn do_prctl(option: i32, arg2: u64, arg3: u64, arg4: u64, arg5: u64) -> Resu
 
 pub fn do_arch_prctl(code: u32, addr: *mut usize) -> Result<isize> {
     let code = ArchPrctlCode::from_u32(code)?;
-    check_mut_ptr(addr)?;
     super::do_arch_prctl::do_arch_prctl(code, addr).map(|_| 0)
 }
 
diff --git a/tools/gen_internal_conf/src/main.rs b/tools/gen_internal_conf/src/main.rs
index e6d75ef3..31caffe1 100644
--- a/tools/gen_internal_conf/src/main.rs
+++ b/tools/gen_internal_conf/src/main.rs
@@ -461,6 +461,7 @@ fn main() {
             ISVFAMILYID_H: kss_tuple.3,
             ISVFAMILYID_L: kss_tuple.4,
             PKRU: occlum_config.metadata.pkru,
+            AMX: occlum_config.metadata.amx,
         };
         let enclave_config = serde_xml_rs::to_string(&sgx_enclave_configuration).unwrap();
         debug!("The enclave config:{:?}", enclave_config);
@@ -724,6 +725,8 @@ struct OcclumMetadata {
     ext_prod_id: OcclumMetaID,
     #[serde(default)]
     pkru: u32,
+    #[serde(default)]
+    amx: u32,
 }
 
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
@@ -792,6 +795,7 @@ struct EnclaveConfiguration {
     ISVFAMILYID_H: u64,
     ISVFAMILYID_L: u64,
     PKRU: u32,
+    AMX: u32,
 }
 
 #[derive(Debug, PartialEq, Clone, Serialize)]