AB upgrade app part (crustio#240)

* AB upgrade app part

* Delete useless declare

Co-authored-by: LowEntropyBody <jszyyx@163.com>

src/app/App.cpp

@@ -1,7 +1,7 @@
 #include "App.h"
 
 bool offline_chain_mode = false;
-bool g_init_upgrade = false;
+bool g_upgrade_flag = false;
 extern std::string config_file_path;
 crust::Log *p_log = crust::Log::get_instance();
 
@@ -34,7 +34,7 @@ int SGX_CDECL main(int argc, char *argv[])
         }
         else if (strcmp(argv[i], "--upgrade") == 0)
         {
-            g_init_upgrade = true;
+            g_upgrade_flag = true;
             i++;
         }
         else if (strcmp(argv[i], "-v") == 0 || strcmp(argv[i], "--version") == 0)

src/app/ecalls/ECalls.cpp

@@ -25,6 +25,8 @@ std::unordered_map<std::string, int> g_task_priority_um = {
     {"Ecall_handle_report_result", 0},
     {"Ecall_gen_upgrade_data", 0},
     {"Ecall_restore_from_upgrade", 0},
+    {"Ecall_enable_upgrade", 0},
+    {"Ecall_disable_upgrade", 0},
     {"Ecall_seal_file", 1},
     {"Ecall_unseal_file", 1},
     {"Ecall_srd_decrease", 1},
@@ -59,6 +61,15 @@ std::unordered_map<std::string, std::unordered_set<std::string>> g_block_tasks_u
         }
     },
 };
+// Upgrade blocks task set1
+std::unordered_set<std::string> g_upgrade_block_task = {
+    "Ecall_seal_file",
+    "Ecall_unseal_file",
+    "Ecall_srd_decrease",
+    "Ecall_srd_increase",
+    "Ecall_confirm_file",
+    "Ecall_delete_file",
+};
 // Task info, invoked enclave function to enclave task_info mapping
 tbb::concurrent_unordered_map<std::string, enclave_task_t> g_running_task_um;
 // Record running task number
@@ -107,6 +118,14 @@ void task_sleep(int priority)
 sgx_status_t try_get_enclave(const char *name)
 {
     std::string tname(name);
+    if (UPGRADE_STATUS_COMPLETE == get_g_upgrade_status())
+    {
+        if (tname.compare("Ecall_disable_upgrade") != 0)
+        {
+            return SGX_ERROR_SERVICE_UNAVAILABLE;
+        }
+    }
+
     std::thread::id tid = std::this_thread::get_id();
     std::stringstream ss;
     ss << tid;
@@ -223,6 +242,21 @@ void free_enclave(const char *name)
     g_invoked_ecalls_um[tname]--;
 }
 
+/**
+ * @description: Get blocking upgrade ecalls' number
+ * @return: Blocking ecalls' number
+ */
+int get_upgrade_ecalls_num()
+{
+    int block_task_num = 0;
+    for (auto task : g_upgrade_block_task)
+    {
+        block_task_num += g_invoked_ecalls_um[task];
+    }
+
+    return block_task_num;
+}
+
 /**
  * @description: Show enclave thread info
  * @return: Task information
@@ -608,15 +642,52 @@ sgx_status_t Ecall_gen_upgrade_data(sgx_enclave_id_t eid, crust_status_t *status
  * @description: Generate upgrade metadata
  * @param status -> Pointer to metadata
  */
-sgx_status_t Ecall_restore_from_upgrade(sgx_enclave_id_t eid, crust_status_t *status, const char *meta, size_t meta_len)
+sgx_status_t Ecall_restore_from_upgrade(sgx_enclave_id_t eid, crust_status_t *status, const char *meta, size_t meta_len, size_t total_size, bool transfer_end)
+{
+    sgx_status_t ret = SGX_SUCCESS;
+    if (SGX_SUCCESS != (ret = try_get_enclave(__FUNCTION__)))
+    {
+        return ret;
+    }
+
+    ret = ecall_restore_from_upgrade(eid, status, meta, meta_len, total_size, transfer_end);
+
+    free_enclave(__FUNCTION__);
+
+    return ret;
+}
+
+/**
+ * @description: Enable upgrade
+ * @param block_height -> Current block height
+ */
+sgx_status_t Ecall_enable_upgrade(sgx_enclave_id_t eid, crust_status_t *status, size_t block_height)
+{
+    sgx_status_t ret = SGX_SUCCESS;
+    if (SGX_SUCCESS != (ret = try_get_enclave(__FUNCTION__)))
+    {
+        return ret;
+    }
+
+    ret = ecall_enable_upgrade(eid, status, block_height);
+
+    free_enclave(__FUNCTION__);
+
+    return ret;
+}
+
+/**
+ * @description: Disable upgrade
+ */
+sgx_status_t Ecall_disable_upgrade(sgx_enclave_id_t eid)
 {
     sgx_status_t ret = SGX_SUCCESS;
     if (SGX_SUCCESS != (ret = try_get_enclave(__FUNCTION__)))
     {
         return ret;
     }
 
-    ret = ecall_restore_from_upgrade(eid, status, meta, meta_len);
+    ret = ecall_disable_upgrade(eid);
 
     free_enclave(__FUNCTION__);
 

src/app/ecalls/ECalls.h

@@ -13,6 +13,7 @@
 #include <sgx_eid.h>
 #include "Enclave_u.h"
 #include "Log.h"
+#include "Data.h"
 #include "CrustStatus.h"
 #include "Json.hpp"
 
@@ -70,8 +71,12 @@ sgx_status_t Ecall_id_get_info(sgx_enclave_id_t eid);
 
 sgx_status_t Ecall_get_workload(sgx_enclave_id_t eid);
 
+sgx_status_t Ecall_enable_upgrade(sgx_enclave_id_t eid, crust_status_t *status, size_t block_height);
+sgx_status_t Ecall_disable_upgrade(sgx_enclave_id_t eid);
 sgx_status_t Ecall_gen_upgrade_data(sgx_enclave_id_t eid, crust_status_t *status, size_t block_height);
-sgx_status_t Ecall_restore_from_upgrade(sgx_enclave_id_t eid, crust_status_t *status, const char *meta, size_t meta_len);
+sgx_status_t Ecall_restore_from_upgrade(sgx_enclave_id_t eid, crust_status_t *status, const char *meta, size_t meta_len, size_t total_size, bool transfer_end);
+
+int get_upgrade_ecalls_num();
 
 std::string show_enclave_thread_info();