From 232ba565206caf01e7f514c0c5735a8e8d3ae06a Mon Sep 17 00:00:00 2001
From: modric <wangyu283@huawei.com>
Date: Wed, 9 Nov 2022 15:17:28 +0800
Subject: [PATCH 3/4] rollback to common invoking when async invoking fails

---
 inc/host_inc/secgear_uswitchless.h |  7 ++-
 inc/host_inc/status.h              |  1 +
 src/host_src/gp/gp_enclave.c       |  5 ++
 src/host_src/gp/gp_uswitchless.c   |  5 ++
 src/host_src/gp/gp_uswitchless.h   | 10 ++++
 tools/codegener/Genheader.ml       | 13 ++++-
 tools/codegener/Gentrust.ml        | 10 ++--
 tools/codegener/Genuntrust.ml      | 92 +++++++++++++++++++++++++++++-
 8 files changed, 133 insertions(+), 10 deletions(-)

diff --git a/inc/host_inc/secgear_uswitchless.h b/inc/host_inc/secgear_uswitchless.h
index 8e21fd9..2ea4691 100644
--- a/inc/host_inc/secgear_uswitchless.h
+++ b/inc/host_inc/secgear_uswitchless.h
@@ -81,10 +81,13 @@ typedef struct {
     uint32_t retries_before_sleep;
 
     /* Worker thread scheduling policy, refer to cc_workers_policy_t, only for GP */
-    uint64_t workers_policy;
+    uint32_t workers_policy;
+
+    /* Indicates whether to roll back to common invoking when asynchronous switchless invoking fails, only for GP */
+    uint32_t rollback_to_common;
 } cc_sl_config_t;
 
-#define CC_USWITCHLESS_CONFIG_INITIALIZER   {1, 1, 1, 16, 0, 0, WORKERS_POLICY_BUSY}
+#define CC_USWITCHLESS_CONFIG_INITIALIZER   {1, 1, 1, 16, 0, 0, WORKERS_POLICY_BUSY, 0}
 
 #ifdef __cplusplus
 }
diff --git a/inc/host_inc/status.h b/inc/host_inc/status.h
index 4f982f8..7f8daaa 100644
--- a/inc/host_inc/status.h
+++ b/inc/host_inc/status.h
@@ -168,6 +168,7 @@ typedef enum _enclave_result_t
     CC_ERROR_ADDRESS_UNACCESSABLE,                      /* Memory address is not within enclave */
     CC_ERROR_SWITCHLESS_INVALID_TASK_ID,                /* Invalid invoking task ID */
     CC_ERROR_SWITCHLESS_ASYNC_TASK_UNFINISHED,          /* The asynchronous invoking task is not completed */
+    CC_ERROR_SWITCHLESS_ROLLBACK2COMMON,                /* rollback to common invoking when async invoking fails */
     CC_MAXIMUM_ERROR,
 } cc_enclave_result_t;
 
diff --git a/src/host_src/gp/gp_enclave.c b/src/host_src/gp/gp_enclave.c
index 5345973..521a850 100644
--- a/src/host_src/gp/gp_enclave.c
+++ b/src/host_src/gp/gp_enclave.c
@@ -857,6 +857,11 @@ cc_enclave_result_t cc_sl_async_ecall(cc_enclave_t *enclave, int *task_id, sl_ec
 
     int task_index = uswitchless_get_idle_task_index(enclave);
     if (task_index < 0) {
+        /* Need roll back to common invoking when asynchronous invoking fails. */
+        if (uswitchless_need_rollback_to_common(enclave)) {
+            return CC_ERROR_SWITCHLESS_ROLLBACK2COMMON;
+        }
+
         return CC_ERROR_SWITCHLESS_TASK_POOL_FULL;
     }
 
diff --git a/src/host_src/gp/gp_uswitchless.c b/src/host_src/gp/gp_uswitchless.c
index 2a315ea..53ecc55 100644
--- a/src/host_src/gp/gp_uswitchless.c
+++ b/src/host_src/gp/gp_uswitchless.c
@@ -112,6 +112,11 @@ bool uswitchless_is_valid_task_index(cc_enclave_t *enclave, int task_index)
     return !((*(pool->free_bit_buf + i)) & (1UL << j));
 }
 
+bool uswitchless_need_rollback_to_common(cc_enclave_t *enclave)
+{
+    return USWITCHLESS_TASK_POOL(enclave)->pool_cfg.rollback_to_common > 0;
+}
+
 int uswitchless_get_idle_task_index(cc_enclave_t *enclave)
 {
     sl_task_pool_t *pool = USWITCHLESS_TASK_POOL(enclave);
diff --git a/src/host_src/gp/gp_uswitchless.h b/src/host_src/gp/gp_uswitchless.h
index 13ac14a..a0ea117 100644
--- a/src/host_src/gp/gp_uswitchless.h
+++ b/src/host_src/gp/gp_uswitchless.h
@@ -134,6 +134,16 @@ bool uswitchless_is_valid_param_num(cc_enclave_t *enclave, uint32_t argc);
  */
 bool uswitchless_is_valid_task_index(cc_enclave_t *enclave, int task_index);
 
+/*
+ * Summary: whether to roll back to common invoking when asynchronous switchless invoking fails
+ * Parameters:
+ *      enclave: enclave
+ * Return:
+ *      true: yes
+ *      false: no
+ */
+bool uswitchless_need_rollback_to_common(cc_enclave_t *enclave);
+
 /*
  * Summary: fill a task
  * Parameters:
diff --git a/tools/codegener/Genheader.ml b/tools/codegener/Genheader.ml
index e00157d..0f244f3 100644
--- a/tools/codegener/Genheader.ml
+++ b/tools/codegener/Genheader.ml
@@ -31,6 +31,14 @@ let generate_args_include (ufs: untrusted_func list) =
     "#include \"enclave.h\"\n" ^
     error_include ^ "\n"
 
+let generate_function_id_ex (tf: trusted_func) =
+    let f = tf.tf_fdecl in
+    let f_name = f.fname in
+    if tf.tf_is_switchless then
+        "fid_sl_async_" ^ f_name
+    else
+        "fid_" ^ f_name
+
 let generate_function_id (f: func_decl) =
     let f_name = f.fname in
     "fid_" ^ f_name
@@ -77,7 +85,8 @@ let generate_rproxy_prototype_sl_async (tf: trusted_func) =
   else
     let fd = tf.tf_fdecl in
     let func_name = fd.fname ^ "_async" in
-    let enclave_decl = "(\n    cc_enclave_t *enclave,\n    int *task_id" in
+    let enclave_decl =
+        "(\n    " ^  (match fd.rtype with Void -> "cc_enclave_t *enclave,\n    int *task_id" | _ -> "cc_enclave_t *enclave,\n    int *task_id,\n    " ^ (get_tystr fd.rtype ^ " *retval")) in
     let func_args =
       let func_args_list =
           List.map (fun f -> gen_parm_str f) fd.plist
@@ -270,7 +279,7 @@ let generate_args_header (ec: enclave_content) =
     let trust_fid_body =
         let trust_fid_pre =
             List.mapi
-        (fun i f -> sprintf "    %s = %d," (generate_function_id f.tf_fdecl) (i + 2)) tfunc_decls
+        (fun i f -> sprintf "    %s = %d," (generate_function_id_ex f) (i + 2)) ec.tfunc_decls
     in
     String.concat "\n" trust_fid_pre
         in
diff --git a/tools/codegener/Gentrust.ml b/tools/codegener/Gentrust.ml
index 6b6fa00..d950899 100644
--- a/tools/codegener/Gentrust.ml
+++ b/tools/codegener/Gentrust.ml
@@ -146,6 +146,7 @@ let set_switchless_ecall_func (tf : trusted_func) =
         match tfd.rtype with
             | Void -> ""
             | _ -> "    (void)memcpy(retval, &ret, sizeof(ret));" in
+    if tf.tf_is_switchless then
     [
         sprintf "\nvoid sl_ecall_%s(void *task_buf)" tfd.fname;
         "{";
@@ -160,15 +161,15 @@ let set_switchless_ecall_func (tf : trusted_func) =
         write_back_retval;
         "}";
     ]
+    else ["";]
 
 let set_ecall_func (tf : trusted_func) =
-    if tf.tf_is_switchless then
-        set_switchless_ecall_func tf
-    else
+    let slfunc = String.concat " " (set_switchless_ecall_func tf) in
     let tfd = tf.tf_fdecl in 
     let params_point = set_parameters_point tfd in
     let out_params = set_out_params tfd in
     [
+        "" ^ slfunc;
         sprintf "cc_enclave_result_t ecall_%s (" tfd.fname;
         "    uint8_t* in_buf,";
         "    size_t in_buf_size,";
@@ -396,8 +397,7 @@ let gen_trusted(ec : enclave_content) =
             "    (cc_ecall_func_t) ecall_unregister_shared_memory,";
             "    " ^ concat ",\n    "
                 (List.map (fun (tf) ->
-                    sprintf "(cc_ecall_func_t) ecall_%s" tf.tf_fdecl.fname)
-                (List.filter (fun tf -> not tf.tf_is_switchless) trust_funcs));
+                    sprintf "(cc_ecall_func_t) ecall_%s" tf.tf_fdecl.fname) trust_funcs);
             "};";
             "";
             "size_t ecall_table_size = CC_ARRAY_LEN(cc_ecall_tables);\n";
diff --git a/tools/codegener/Genuntrust.ml b/tools/codegener/Genuntrust.ml
index 6fb4967..8bc8e03 100644
--- a/tools/codegener/Genuntrust.ml
+++ b/tools/codegener/Genuntrust.ml
@@ -80,6 +80,40 @@ let set_call_user_func (fd : func_decl) =
         "}";
     ]
 
+let sl_async_set_call_user_func (fd : func_decl) =
+    [
+        "/* Call the cc_enclave function */";
+        "if (!enclave) {";
+        "    ret = CC_ERROR_BAD_PARAMETERS;";
+        "    goto exit;";
+        "}";
+        "if (pthread_rwlock_rdlock(&enclave->rwlock)) {";
+        "    ret = CC_ERROR_BUSY;";
+        "    goto exit;";
+        "}";
+        "if (!enclave->list_ops_node || !enclave->list_ops_node->ops_desc ||";
+        "         !enclave->list_ops_node->ops_desc->ops ||";
+        "         !enclave->list_ops_node->ops_desc->ops->cc_ecall_enclave) {";
+        "    ret = CC_ERROR_BAD_PARAMETERS;";
+        "    goto exit;";
+        "}";
+        "if ((ret = enclave->list_ops_node->ops_desc->ops->cc_ecall_enclave(";
+        "         enclave,";
+        sprintf "         fid_sl_async_%s," fd.fname;
+        "         in_buf,";
+        "         in_buf_size,";
+        "         out_buf,";
+        "         out_buf_size,";
+        "         &ms,";
+        "         &ocall_table)) != CC_SUCCESS) {";
+        "    pthread_rwlock_unlock(&enclave->rwlock);";
+        "    goto exit; }";
+        "if (pthread_rwlock_unlock(&enclave->rwlock)) {";
+        "    ret = CC_ERROR_BUSY;";
+        "    goto exit;";
+        "}";
+    ]
+
 let set_ecall_func_arguments (fd : func_decl) =
     [
         sprintf "cc_enclave_result_t %s(\n    %s" fd.fname  (match fd.rtype with Void -> "cc_enclave_t *enclave" | _ -> "cc_enclave_t *enclave,\n    " ^ (get_tystr fd.rtype ^ "* retval"))
@@ -100,7 +134,7 @@ let set_ecall_func_arguments (fd : func_decl) =
 
 let set_sl_async_ecall_func_arguments (fd : func_decl) =
     [
-        sprintf "cc_enclave_result_t %s(\n    %s" (fd.fname ^ "_async")  "cc_enclave_t *enclave,\n    int *task_id"
+        sprintf "cc_enclave_result_t %s(\n    %s" (fd.fname ^ "_async") (match fd.rtype with Void -> "cc_enclave_t *enclave,\n    int *task_id" | _ -> "cc_enclave_t *enclave,\n    int *task_id,\n    " ^ (get_tystr fd.rtype ^ " *retval"))
         ^ (if fd.plist <> [] then
             ",\n    " ^
             concat ",\n    "
@@ -119,6 +153,7 @@ let set_sl_async_ecall_func_arguments (fd : func_decl) =
 let set_sl_ecall_func (tf : trusted_func) =
     let tfd = tf.tf_fdecl in
     let init_point = set_init_pointer tfd in
+    let arg_size = set_args_size tfd in
     let get_param_name (_, decl) = decl.identifier in
     (*let is_ptr_type (ptype) =
         match ptype with
@@ -224,6 +259,61 @@ let set_sl_ecall_func (tf : trusted_func) =
 
         "    pthread_rwlock_unlock(&enclave->rwlock);";
         (if tfd.plist <> [] then "    free(params_buf);" else "");
+        "    if (ret != CC_ERROR_SWITCHLESS_ROLLBACK2COMMON) {\n        return ret;\n    }";
+        "\n    /* rollback to common invoking when async invoking fails */";
+        "    ret = CC_FAIL;";
+        "    *task_id = -1;";
+        "";
+        "    /* Init buffer and size  */";
+        "    size_t in_buf_size = 0;";
+        "    size_t out_buf_size = 0;";
+        "    uint8_t* in_buf = NULL;";
+        "    uint8_t* out_buf = NULL;";
+        "    uint32_t ms = TEE_SECE_AGENT_ID;";
+        sprintf "    %s_size_t args_size;" tfd.fname;
+        "";
+        "    /* Init pointer */";
+        if init_point <> ["";"";""] then
+            concat "\n" init_point
+        else "    /* There is no pointer */";
+        "";
+        "    memset(&args_size, 0, sizeof(args_size));";
+        "    /* Fill argments size */";
+        if arg_size <> [""] then
+            "    " ^ concat "\n    " (set_args_size tfd)
+        else "/* There is no argments size */";
+        "";
+        sprintf "    in_buf_size += size_to_aligned_size(sizeof(%s_size_t));"
+          tfd.fname;
+
+        "    " ^ concat "\n    " (set_data_in tfd);
+        "";
+
+        "    " ^ concat "\n    " (set_data_out tfd);
+        "";
+        "    /* Allocate in_buf and out_buf */";
+        "    in_buf = (uint8_t*)malloc(in_buf_size);";
+        "    out_buf = (uint8_t*)malloc(out_buf_size);";
+        "    if (in_buf == NULL || out_buf == NULL) {";
+        "        ret = CC_ERROR_OUT_OF_MEMORY;";
+        "        goto exit;";
+        "    }";
+
+        "";
+        "    " ^ concat "\n    " (set_in_memcpy tfd);
+        "";
+        "    " ^ concat "\n    " (sl_async_set_call_user_func tfd);
+        "";
+        "    " ^ concat "\n    " (set_out_memcpy tfd);
+        "    ret = CC_SUCCESS;";
+        "";
+
+        "exit:";
+        "    if (in_buf)";
+        "        free(in_buf);";
+        "    if (out_buf)";
+        "        free(out_buf);";
+        "";
         "    return ret;";
         "}";
     ]
-- 
2.27.0

