secGear/0079-Optimize-the-registration-shared-memory.patch

From 2e22f45d9c20941823761fa858e1faa9ce050a2c Mon Sep 17 00:00:00 2001
From: zhengxiaoxiao <zhengxiaoxiao2@huawei.com>
Date: Tue, 20 Aug 2024 23:41:45 +0800
Subject: [PATCH] Optimize the registration shared memory

Reference:https://gitee.com/openeuler/secGear/commit/2e22f45d9c20941823761fa858e1faa9ce050a2c
Conflict:NA
---
 inc/host_inc/enclave_internal.h      |   2 +-
 src/host_src/gp/gp_enclave.c         | 112 +++++++++++++++++++++++----
 src/host_src/gp/gp_shared_memory.c   |  22 +++++-
 src/host_src/gp/gp_shared_memory.h   |   2 +-
 src/host_src/secgear_shared_memory.c |  24 +++---
 src/host_src/sgx/sgx_shared_memory.c | 3 ++-
 src/host_src/sgx/sgx_shared_memory.h | 2 +-
 7 files changed, 134 insertions(+), 33 deletions(-)

diff --git a/inc/host_inc/enclave_internal.h b/inc/host_inc/enclave_internal.h
index fa0cbf4..ac88f46 100644
--- a/inc/host_inc/enclave_internal.h
+++ b/inc/host_inc/enclave_internal.h
@@ -74,7 +74,7 @@ struct cc_enclave_ops {
     cc_enclave_result_t (*cc_sl_async_ecall_get_result)(cc_enclave_t *enclave, int task_id, void *retval);
 
     /* shared memory */
-    void *(*cc_malloc_shared_memory)(cc_enclave_t *enclave, size_t size, bool is_control_buf);
+    void *(*cc_malloc_shared_memory)(cc_enclave_t *enclave, size_t size, bool is_control_buf, int try_cnt);
     cc_enclave_result_t (*cc_free_shared_memory)(cc_enclave_t *enclave, void *ptr);
     cc_enclave_result_t (*cc_register_shared_memory)(cc_enclave_t *enclave, void *ptr);
     cc_enclave_result_t (*cc_unregister_shared_memory)(cc_enclave_t *enclave, void *ptr);
diff --git a/src/host_src/gp/gp_enclave.c b/src/host_src/gp/gp_enclave.c
index ad07c30..a2ff9f4 100644
--- a/src/host_src/gp/gp_enclave.c
+++ b/src/host_src/gp/gp_enclave.c
@@ -377,27 +377,34 @@ cc_enclave_result_t init_uswitchless(cc_enclave_t *enclave, const enclave_featur
     uswitchless_adjust_config(&cfg);
 
     size_t pool_buf_len = sl_get_pool_buf_len_by_config(&cfg);
-    void *pool_buf = gp_malloc_shared_memory(enclave, pool_buf_len, true);
-    if (pool_buf == NULL) {
-        return CC_ERROR_OUT_OF_MEMORY;
-    }
-    (void)memset(pool_buf, 0, pool_buf_len);
+    cc_enclave_result_t ret;
+    sl_task_pool_t *pool;
+    for (int i = 0; i < 2; i++) {
+        void *pool_buf = gp_malloc_shared_memory(enclave, pool_buf_len, true, i);
+        if (pool_buf == NULL) {
+            return CC_ERROR_OUT_OF_MEMORY;
+        }
+        (void)memset(pool_buf, 0, pool_buf_len);
 
-    // Fill config
-    (void)memcpy(pool_buf, &cfg, sizeof(cc_sl_config_t));
+        // Fill config
+        (void)memcpy(pool_buf, &cfg, sizeof(cc_sl_config_t));
 
-    // Layout task pool
-    sl_task_pool_t *pool = uswitchless_create_task_pool(pool_buf, &cfg);
-    if (pool == NULL) {
-        (void)gp_free_shared_memory(enclave, pool_buf);
-        return CC_ERROR_OUT_OF_MEMORY;
-    }
+        // Layout task pool
+        pool = uswitchless_create_task_pool(pool_buf, &cfg);
+        if (pool == NULL) {
+            (void)gp_free_shared_memory(enclave, pool_buf);
+            return CC_ERROR_OUT_OF_MEMORY;
+        }
 
-    // Registering a task pool
-    cc_enclave_result_t ret = gp_register_shared_memory(enclave, pool_buf);
-    if (ret != CC_SUCCESS) {
+        // Registering a task pool
+        ret = gp_register_shared_memory(enclave, pool_buf);
+        if (ret == CC_SUCCESS) {
+            break;
+        }
         free(pool);
         (void)gp_free_shared_memory(enclave, pool_buf);
+    }
+    if (ret != CC_SUCCESS) {
         return ret;
     }
 
diff --git a/src/host_src/gp/gp_shared_memory.c b/src/host_src/gp/gp_shared_memory.c
index cd1a4c5..232edbf 100644
--- a/src/host_src/gp/gp_shared_memory.c
+++ b/src/host_src/gp/gp_shared_memory.c
@@ -47,7 +47,7 @@ static void gp_add_shared_mem_to_list(gp_shared_memory_t *shared_mem)
     CC_RWLOCK_UNLOCK(&g_shared_mem_list_lock);
 }
 
-void *gp_malloc_shared_memory(cc_enclave_t *context, size_t size, bool is_control_buf)
+void *gp_malloc_shared_memory(cc_enclave_t *context, size_t size, bool is_control_buf, int try_cnt)
 {
     gp_context_t *gp_context = (gp_context_t *)context->private_data;
     gp_shared_memory_t gp_shared_mem = {
@@ -63,7 +63,7 @@ void *gp_malloc_shared_memory(cc_enclave_t *context, size_t size, bool is_contro
     }
     TEEC_SharedMemory *teec_shared_mem = (TEEC_SharedMemory *)(&gp_shared_mem.shared_mem);
     teec_shared_mem->size = size + sizeof(gp_shared_memory_t);
-    teec_shared_mem->flags = TEEC_MEM_REGISTER_INOUT;
+    teec_shared_mem->flags = try_cnt == 0 ? TEEC_MEM_REGISTER_INOUT : TEEC_MEM_SHARED_INOUT;
 
     TEEC_Result result = TEEC_AllocateSharedMemory(&gp_context->ctx, teec_shared_mem);
     if (result == TEEC_ERROR_BAD_PARAMETERS) {
diff --git a/src/host_src/gp/gp_shared_memory.h b/src/host_src/gp/gp_shared_memory.h
index 6914193..4659b4a 100644
--- a/src/host_src/gp/gp_shared_memory.h
+++ b/src/host_src/gp/gp_shared_memory.h
@@ -31,7 +31,7 @@ extern "C" {
  *     is_control_buf: whether it is a control area buffer
  * Return: A pointer to the allocated memory. On error, return NULL.
  */
-void *gp_malloc_shared_memory(cc_enclave_t *context, size_t size, bool is_control_buf);
+void *gp_malloc_shared_memory(cc_enclave_t *context, size_t size, bool is_control_buf, int try_cnt);
 
 /*
  * Summary: Frees the memory space pointed to by ptr, which must have been returned by gp_malloc_shared_memory.
diff --git a/src/host_src/secgear_shared_memory.c b/src/host_src/secgear_shared_memory.c
index d7e8d35..258f329 100644
--- a/src/host_src/secgear_shared_memory.c
+++ b/src/host_src/secgear_shared_memory.c
@@ -40,21 +40,27 @@ void *cc_malloc_shared_memory(cc_enclave_t *enclave, size_t size)
         return NULL;
     }
 
-    void *ptr = FUNC_CREATE_SHARED_MEM(enclave)(enclave, size, false);
-    if (ptr == NULL) {
-        CC_RWLOCK_UNLOCK(&enclave->rwlock);
-        return NULL;
+    cc_enclave_result_t ret;
+    void *ptr;
+    for (int i = 0; i < 2; i++) {
+        ptr = FUNC_CREATE_SHARED_MEM(enclave)(enclave, size, false, i);
+        if (ptr == NULL) {
+            CC_RWLOCK_UNLOCK(&enclave->rwlock);
+            return NULL;
+        }
+
+        ret = FUNC_REGISTER_SHARED_MEM(enclave)(enclave, ptr);
+        if (ret == CC_SUCCESS) {
+            break;
+        }
+        CC_IGNORE(FUNC_FREE_SHARED_MEM(enclave)(enclave, ptr));
     }
 
-    cc_enclave_result_t ret = FUNC_REGISTER_SHARED_MEM(enclave)(enclave, ptr);
+    CC_RWLOCK_UNLOCK(&enclave->rwlock);
     if (ret != CC_SUCCESS) {
-        CC_IGNORE(FUNC_FREE_SHARED_MEM(enclave)(enclave, ptr));
-        CC_RWLOCK_UNLOCK(&enclave->rwlock);
         return NULL;
     }
 
-    CC_RWLOCK_UNLOCK(&enclave->rwlock);
-
     return ptr;
 }
 
diff --git a/src/host_src/sgx/sgx_shared_memory.c b/src/host_src/sgx/sgx_shared_memory.c
index b9ecf9a..2699580 100644
--- a/src/host_src/sgx/sgx_shared_memory.c
+++ b/src/host_src/sgx/sgx_shared_memory.c
@@ -15,10 +15,11 @@
 #include <stdlib.h>
 #include "secgear_defs.h"
 
-void *sgx_malloc_shared_memory(cc_enclave_t *enclave, size_t size, bool is_control_buf)
+void *sgx_malloc_shared_memory(cc_enclave_t *enclave, size_t size, bool is_control_buf, int try_cnt)
 {
     CC_IGNORE(enclave);
     CC_IGNORE(is_control_buf);
+    CC_IGNORE(try_cnt);
 
     return malloc(size);
 }
diff --git a/src/host_src/sgx/sgx_shared_memory.h b/src/host_src/sgx/sgx_shared_memory.h
index 861cea7..b7f886a 100644
--- a/src/host_src/sgx/sgx_shared_memory.h
+++ b/src/host_src/sgx/sgx_shared_memory.h
@@ -27,7 +27,7 @@
  *     is_control_buf: whether it is a control area buffer
  * Return: A pointer to the allocated memory. On error, return NULL.
  */
-void *sgx_malloc_shared_memory(cc_enclave_t *context, size_t size, bool is_control_buf);
+void *sgx_malloc_shared_memory(cc_enclave_t *context, size_t size, bool is_control_buf, int try_cnt);
 
 /*
  * Summary: Frees the memory space pointed to by ptr, which must have been returned by sgx_malloc_shared_memory.
diff --git a/src/enclave_src/gp/gp.c b/src/enclave_src/gp/gp.c
index 3f30a16..4f3c3b2 100644
--- a/src/enclave_src/gp/gp.c
+++ b/src/enclave_src/gp/gp.c
@@ -69,6 +69,13 @@ TEE_Result TA_OpenSessionEntryPoint(uint32_t paramTypes,
     TEE_Result ret = TEE_SUCCESS;
     SLogTrace("---- TA_OpenSessionEntryPoint -------- ");
 
+    uint32_t param_in = 0;
+    uint32_t param_shared_mem = 1;
+    if (TEE_PARAM_TYPE_GET(paramTypes, param_shared_mem) == TEE_PARAM_TYPE_MEMREF_REGISTER_INOUT) {
+        ret = register_shared_memory_by_session(params[param_in].memref.buffer,
+            params[param_shared_mem].memref.buffer, sessionContext);
+        tlogi("[secGear]TA_OpenSessionEntryPoint register shared memory ret:%d, shared_mem:%p", ret, *sessionContext);
+    }
     return ret;
 }
 
@@ -83,6 +90,12 @@ void TA_CloseSessionEntryPoint(void *sessionContext)
 {
     (void)sessionContext;  /* -Wunused-parameter */
     SLogTrace("---- TA_CloseSessionEntryPoint ----- ");
+
+    // find shared mem block by session, and destroy
+    if (sessionContext != NULL) {
+        tlogi("[secGear]TA_CloseSessionEntryPoint unregister shared_mem:%p", sessionContext);
+        open_session_unregister_shared_memory(sessionContext);
+    }
 }
 
 /**
-- 
2.27.0