More cleanup

informalsystems · Jan 17, 2025 · d1cf669 · d1cf669
1 parent fe4e95f
commit d1cf669
Show file tree

Hide file tree

Showing 6 changed files with 75 additions and 101 deletions.
diff --git a/examples/transfers/enclave/src/event.rs b/examples/transfers/enclave/src/event.rs
@@ -1,4 +1,6 @@
-use anyhow::Error as AnyhowError;
+use std::collections::BTreeMap;
+
+use anyhow::{anyhow, Error as AnyhowError};
 use cosmrs::AccountId;
 use quartz_common::enclave::{chain_client::ChainClient, handler::Handler};
 use tendermint_rpc::event::Event as TmEvent;
@@ -46,3 +48,14 @@ where
         }
     }
 }
+
+fn first_event_with_key<'a>(
+    events: &'a BTreeMap<String, Vec<String>>,
+    key: &str,
+) -> Result<&'a String, AnyhowError> {
+    events
+        .get(key)
+        .ok_or_else(|| anyhow!("missing execute._contract_address in events"))?
+        .first()
+        .ok_or_else(|| anyhow!("execute._contract_address is empty"))
+}
diff --git a/examples/transfers/enclave/src/event/query.rs b/examples/transfers/enclave/src/event/query.rs
@@ -6,7 +6,9 @@ use serde_json::json;
 use tendermint_rpc::event::Event as TmEvent;
 use transfers_contract::msg::QueryMsg::GetState;
 
-use crate::{proto::QueryRequest, request::query::QueryRequestMessage};
+use crate::{
+    event::first_event_with_key, proto::QueryRequest, request::query::QueryRequestMessage,
+};
 
 #[derive(Clone, Debug)]
 pub struct QueryEvent {
@@ -27,27 +29,14 @@ impl TryFrom<TmEvent> for QueryEvent {
             return Err(anyhow!("irrelevant event"));
         };
 
-        let contract = events
-            .get("execute._contract_address")
-            .ok_or_else(|| anyhow!("missing execute._contract_address in events"))?
-            .first()
-            .ok_or_else(|| anyhow!("execute._contract_address is empty"))?
+        let contract = first_event_with_key(events, "execute._contract_address")?
             .parse::<AccountId>()
             .map_err(|e| anyhow!("failed to parse contract address: {}", e))?;
 
-        let sender = events
-            .get("message.sender")
-            .ok_or_else(|| anyhow!("Missing message.sender in events"))?
-            .first()
-            .ok_or_else(|| anyhow!("execute.sender is empty"))?
-            .to_owned();
+        let sender = first_event_with_key(events, "message.sender")?.to_owned();
 
-        let ephemeral_pubkey = events
-            .get("wasm-query_balance.emphemeral_pubkey")
-            .ok_or_else(|| anyhow!("Missing wasm-query_balance.emphemeral_pubkey in events"))?
-            .first()
-            .ok_or_else(|| anyhow!("execute.query_balance.emphemeral_pubkey is empty"))?
-            .to_owned();
+        let ephemeral_pubkey =
+            first_event_with_key(events, "wasm-query_balance.emphemeral_pubkey")?.to_owned();
 
         Ok(QueryEvent {
             contract,

diff --git a/examples/transfers/enclave/src/event/transfer.rs b/examples/transfers/enclave/src/event/transfer.rs
@@ -13,7 +13,9 @@ use transfers_contract::msg::{
     QueryMsg::{GetRequests, GetState},
 };
 
-use crate::{proto::UpdateRequest, request::update::UpdateRequestMessage};
+use crate::{
+    event::first_event_with_key, proto::UpdateRequest, request::update::UpdateRequestMessage,
+};
 
 #[derive(Clone, Debug)]
 pub struct TransferEvent {
@@ -32,11 +34,7 @@ impl TryFrom<TmEvent> for TransferEvent {
             return Err(anyhow!("irrelevant event"));
         };
 
-        let contract = events
-            .get("execute._contract_address")
-            .ok_or_else(|| anyhow!("missing execute._contract_address in events"))?
-            .first()
-            .ok_or_else(|| anyhow!("execute._contract_address is empty"))?
+        let contract = first_event_with_key(events, "execute._contract_address")?
             .parse::<AccountId>()
             .map_err(|e| anyhow!("failed to parse contract address: {}", e))?;
 

diff --git a/examples/transfers/enclave/src/request.rs b/examples/transfers/enclave/src/request.rs
@@ -1,11 +1,17 @@
+use cosmwasm_std::HexBinary;
+use ecies::{decrypt, encrypt};
+use k256::ecdsa::{SigningKey, VerifyingKey};
 use quartz_common::{
     contract::msg::execute::attested::{HasUserData, RawMsgSansHandler},
     enclave::{attestor::Attestor, handler::Handler, DefaultSharedEnclave, Enclave},
 };
 use tonic::Status;
-use transfers_contract::msg::{AttestedMsg, ExecuteMsg};
+use transfers_contract::msg::{execute::ClearTextTransferRequestMsg, AttestedMsg, ExecuteMsg};
 
-use crate::proto::{QueryRequest, UpdateRequest};
+use crate::{
+    proto::{QueryRequest, UpdateRequest},
+    state::{Balance, State},
+};
 
 pub mod query;
 pub mod update;
@@ -52,3 +58,38 @@ impl Handler<DefaultSharedEnclave<()>> for EnclaveRequest {
         }
     }
 }
+
+fn decrypt_transfer(
+    sk: &SigningKey,
+    ciphertext: &HexBinary,
+) -> Result<ClearTextTransferRequestMsg, Status> {
+    let o =
+        decrypt(&sk.to_bytes(), ciphertext).map_err(|e| Status::invalid_argument(e.to_string()))?;
+
+    serde_json::from_slice(&o)
+        .map_err(|e| Status::internal(format!("Could not deserialize transfer {}", e)))
+}
+
+fn decrypt_state(sk: &SigningKey, ciphertext: &[u8]) -> Result<State, Status> {
+    let o =
+        decrypt(&sk.to_bytes(), ciphertext).map_err(|e| Status::invalid_argument(e.to_string()))?;
+    serde_json::from_slice(&o).map_err(|e| Status::invalid_argument(e.to_string()))
+}
+
+fn encrypt_state(state: State, enclave_pk: VerifyingKey) -> Result<HexBinary, Status> {
+    let serialized_state = serde_json::to_string(&state).expect("infallible serializer");
+
+    match encrypt(&enclave_pk.to_sec1_bytes(), serialized_state.as_bytes()) {
+        Ok(encrypted_state) => Ok(encrypted_state.into()),
+        Err(e) => Err(Status::internal(format!("Encryption error: {}", e))),
+    }
+}
+
+fn encrypt_balance(balance: Balance, ephemeral_pk: VerifyingKey) -> Result<HexBinary, Status> {
+    let serialized_balance = serde_json::to_string(&balance).expect("infallible serializer");
+
+    match encrypt(&ephemeral_pk.to_sec1_bytes(), serialized_balance.as_bytes()) {
+        Ok(encrypted_balance) => Ok(encrypted_balance.into()),
+        Err(e) => Err(Status::internal(format!("Encryption error: {}", e))),
+    }
+}
diff --git a/examples/transfers/enclave/src/request/query.rs b/examples/transfers/enclave/src/request/query.rs
@@ -1,6 +1,5 @@
 use cosmwasm_std::{Addr, HexBinary, Uint128};
-use ecies::{decrypt, encrypt};
-use k256::ecdsa::{SigningKey, VerifyingKey};
+use k256::ecdsa::VerifyingKey;
 use quartz_common::enclave::{
     handler::Handler, key_manager::KeyManager, DefaultSharedEnclave, Enclave,
 };
@@ -10,6 +9,7 @@ use transfers_contract::msg::execute;
 
 use crate::{
     proto::QueryRequest,
+    request::{decrypt_state, encrypt_balance},
     state::{Balance, State},
 };
 
@@ -71,18 +71,3 @@ impl Handler<DefaultSharedEnclave<()>> for QueryRequest {
         Ok(msg)
     }
 }
-
-fn decrypt_state(sk: &SigningKey, ciphertext: &[u8]) -> Result<State, Status> {
-    let o =
-        decrypt(&sk.to_bytes(), ciphertext).map_err(|e| Status::invalid_argument(e.to_string()))?;
-    serde_json::from_slice(&o).map_err(|e| Status::invalid_argument(e.to_string()))
-}
-
-fn encrypt_balance(balance: Balance, ephemeral_pk: VerifyingKey) -> Result<HexBinary, Status> {
-    let serialized_balance = serde_json::to_string(&balance).expect("infallible serializer");
-
-    match encrypt(&ephemeral_pk.to_sec1_bytes(), serialized_balance.as_bytes()) {
-        Ok(encrypted_balance) => Ok(encrypted_balance.into()),
-        Err(e) => Err(Status::internal(format!("Encryption error: {}", e))),
-    }
-}
diff --git a/examples/transfers/enclave/src/request/update.rs b/examples/transfers/enclave/src/request/update.rs
@@ -1,8 +1,6 @@
 use std::collections::btree_map::Entry;
 
 use cosmwasm_std::{Addr, HexBinary, Uint128};
-use ecies::{decrypt, encrypt};
-use k256::ecdsa::{SigningKey, VerifyingKey};
 use quartz_common::enclave::{
     handler::Handler,
     key_manager::KeyManager,
@@ -17,7 +15,11 @@ use transfers_contract::{
     state::REQUESTS_KEY,
 };
 
-use crate::{proto::UpdateRequest, state::State};
+use crate::{
+    proto::UpdateRequest,
+    request::{decrypt_state, decrypt_transfer, encrypt_state},
+    state::State,
+};
 
 #[derive(Clone, Debug, Serialize, Deserialize, Default)]
 pub struct UpdateRequestMessage {
@@ -85,12 +87,6 @@ impl Handler<DefaultSharedEnclave<()>> for UpdateRequest {
         // Instantiate empty withdrawals map to include in response (Update message to smart contract)
         let mut withdrawals_response: Vec<(Addr, Uint128)> = Vec::<(Addr, Uint128)>::new();
 
-        // let pending_sequenced_requests = message
-        //     .requests
-        //     .iter()
-        //     .filter(|req| matches!(req, TransfersRequest::Transfer(_)))
-        //     .count();
-
         // Loop through requests, match on cases, and apply changes to state
         for req in message.requests {
             match req {
@@ -162,51 +158,3 @@ impl Handler<DefaultSharedEnclave<()>> for UpdateRequest {
         Ok(msg)
     }
 }
-
-fn decrypt_transfer(
-    sk: &SigningKey,
-    ciphertext: &HexBinary,
-) -> Result<ClearTextTransferRequestMsg, Status> {
-    let o =
-        decrypt(&sk.to_bytes(), ciphertext).map_err(|e| Status::invalid_argument(e.to_string()))?;
-
-    serde_json::from_slice(&o)
-        .map_err(|e| Status::internal(format!("Could not deserialize transfer {}", e)))
-}
-
-fn decrypt_state(sk: &SigningKey, ciphertext: &[u8]) -> Result<State, Status> {
-    let o =
-        decrypt(&sk.to_bytes(), ciphertext).map_err(|e| Status::invalid_argument(e.to_string()))?;
-    serde_json::from_slice(&o).map_err(|e| Status::invalid_argument(e.to_string()))
-}
-
-fn encrypt_state(state: State, enclave_pk: VerifyingKey) -> Result<HexBinary, Status> {
-    let serialized_state = serde_json::to_string(&state).expect("infallible serializer");
-
-    match encrypt(&enclave_pk.to_sec1_bytes(), serialized_state.as_bytes()) {
-        Ok(encrypted_state) => Ok(encrypted_state.into()),
-        Err(e) => Err(Status::internal(format!("Encryption error: {}", e))),
-    }
-}
-
-// fn ensure_seq_num_consistency(
-//     seq_num_in_store: &mut u64,
-//     seq_num_on_chain: u64,
-//     pending_sequenced_requests: usize,
-// ) -> Result<(), Status> {
-//     if seq_num_on_chain < *seq_num_in_store {
-//         return Err(Status::failed_precondition("replay attempted"));
-//     }
-//
-//     // make sure number of pending requests are equal to the diff b/w on-chain v/s in-mem seq num
-//     let seq_num_diff = seq_num_on_chain - *seq_num_in_store;
-//     if seq_num_diff != pending_sequenced_requests as u64 {
-//         return Err(Status::failed_precondition(format!(
-//             "seq_num_diff mismatch: num({seq_num_diff}) v/s diff({pending_sequenced_requests})"
-//         )));
-//     }
-//
-//     *seq_num_in_store = seq_num_on_chain;
-//
-//     Ok(())
-// }