From 976c1cacf05108b85006937d055babdd5b0cf507 Mon Sep 17 00:00:00 2001
From: hu55a1n1 <sufialhussaini@gmail.com>
Date: Wed, 10 Jan 2024 07:10:24 -0800
Subject: [PATCH] Impl RegisterEpochKeyMsg and handler

---
 .../contracts/cw-tee-mtcs/src/contract.rs     | 25 ++++++++++++++++---
 .../contracts/cw-tee-mtcs/src/msg.rs          |  6 +++++
 .../contracts/cw-tee-mtcs/src/state.rs        |  8 +++++-
 3 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/bisenzone-cw-mvp/contracts/cw-tee-mtcs/src/contract.rs b/bisenzone-cw-mvp/contracts/cw-tee-mtcs/src/contract.rs
index 5063f10..55f68d4 100644
--- a/bisenzone-cw-mvp/contracts/cw-tee-mtcs/src/contract.rs
+++ b/bisenzone-cw-mvp/contracts/cw-tee-mtcs/src/contract.rs
@@ -4,7 +4,7 @@ use cosmwasm_std::{
 use cw2::set_contract_version;
 
 use crate::error::ContractError;
-use crate::msg::execute::{BootstrapKeyManagerMsg, JoinComputeNodeMsg};
+use crate::msg::execute::{BootstrapKeyManagerMsg, JoinComputeNodeMsg, RegisterEpochKeyMsg};
 use crate::msg::{ExecuteMsg, InstantiateMsg, QueryMsg};
 use crate::state::{State, STATE};
 
@@ -45,6 +45,9 @@ pub fn execute(
         }) => {
             execute::bootstrap_key_manger(deps, compute_mrenclave, key_manager_mrenclave, tcb_info)
         }
+        ExecuteMsg::RegisterEpochKey(RegisterEpochKeyMsg { epoch_key }) => {
+            execute::register_epoch_key(deps, epoch_key)
+        }
         ExecuteMsg::JoinComputeNode(JoinComputeNodeMsg {
             io_exchange_key,
             address,
@@ -58,8 +61,8 @@ pub mod execute {
     use k256::ecdsa::VerifyingKey;
 
     use crate::state::{
-        Mrenclave, RawAddress, RawMrenclave, RawNonce, RawPublicKey, RawTcbInfo, SgxState,
-        SGX_STATE,
+        EpochState, Mrenclave, RawAddress, RawMrenclave, RawNonce, RawPublicKey, RawTcbInfo,
+        SgxState, EPOCH_STATE, SGX_STATE,
     };
     use crate::state::{Request, REQUESTS};
     use crate::ContractError;
@@ -98,6 +101,22 @@ pub mod execute {
             .add_attribute("tcb_info", tcb_info))
     }
 
+    pub fn register_epoch_key(
+        deps: DepsMut,
+        epoch_key: RawPublicKey,
+    ) -> Result<Response, ContractError> {
+        let _ = VerifyingKey::from_sec1_bytes(&hex::decode(&epoch_key)?)?;
+
+        let epoch_state = EpochState {
+            epoch_key: epoch_key.clone(),
+        };
+        EPOCH_STATE.save(deps.storage, &epoch_state)?;
+
+        Ok(Response::new()
+            .add_attribute("action", "register_epoch_key")
+            .add_attribute("epoch_key", epoch_key))
+    }
+
     pub fn enqueue_join_request(
         deps: DepsMut,
         io_exchange_key: RawPublicKey,
diff --git a/bisenzone-cw-mvp/contracts/cw-tee-mtcs/src/msg.rs b/bisenzone-cw-mvp/contracts/cw-tee-mtcs/src/msg.rs
index b8b69c4..f442b26 100644
--- a/bisenzone-cw-mvp/contracts/cw-tee-mtcs/src/msg.rs
+++ b/bisenzone-cw-mvp/contracts/cw-tee-mtcs/src/msg.rs
@@ -6,6 +6,7 @@ pub struct InstantiateMsg;
 #[cw_serde]
 pub enum ExecuteMsg {
     BootstrapKeyManager(execute::BootstrapKeyManagerMsg),
+    RegisterEpochKey(execute::RegisterEpochKeyMsg),
     JoinComputeNode(execute::JoinComputeNodeMsg),
 }
 
@@ -19,6 +20,11 @@ pub mod execute {
         pub tcb_info: String,
     }
 
+    #[cw_serde]
+    pub struct RegisterEpochKeyMsg {
+        pub epoch_key: String,
+    }
+
     #[cw_serde]
     pub struct JoinComputeNodeMsg {
         pub io_exchange_key: String,
diff --git a/bisenzone-cw-mvp/contracts/cw-tee-mtcs/src/state.rs b/bisenzone-cw-mvp/contracts/cw-tee-mtcs/src/state.rs
index 98acc2b..c618384 100644
--- a/bisenzone-cw-mvp/contracts/cw-tee-mtcs/src/state.rs
+++ b/bisenzone-cw-mvp/contracts/cw-tee-mtcs/src/state.rs
@@ -26,6 +26,12 @@ pub struct SgxState {
     pub tcb_info: RawTcbInfo,
 }
 
+#[cw_serde]
+pub struct EpochState {
+    pub epoch_key: RawPublicKey,
+}
+
 pub const STATE: Item<State> = Item::new("state");
 pub const REQUESTS: Item<Vec<(RawNonce, Request)>> = Item::new("requests");
-pub const SGX_STATE: Item<SgxState> = Item::new("sgxstate");
+pub const SGX_STATE: Item<SgxState> = Item::new("sgx_state");
+pub const EPOCH_STATE: Item<EpochState> = Item::new("epoch_state");