Impl query and update request handlers

examples/transfers/enclave/src/request.rs

@@ -1,6 +1,8 @@
+use cosmrs::AccountId;
+use k256::ecdsa::{SigningKey, VerifyingKey};
 use quartz_common::{
     contract::msg::execute::attested::{HasUserData, RawMsgSansHandler},
-    enclave::{attestor::Attestor, handler::Handler, Enclave},
+    enclave::{attestor::Attestor, handler::Handler, key_manager::KeyManager, Enclave},
 };
 use tonic::Status;
 use transfers_contract::msg::{AttestedMsg, ExecuteMsg};
@@ -31,7 +33,11 @@ fn attested_msg<T: HasUserData + Clone, A: Attestor>(
 }
 
 #[async_trait::async_trait]
-impl<E: Enclave> Handler<E> for EnclaveRequest {
+impl<E: Enclave> Handler<E> for EnclaveRequest
+where
+    E: Enclave<Contract = AccountId>,
+    E::KeyManager: KeyManager<PubKey = VerifyingKey, PrivKey = SigningKey>,
+{
     type Error = Status;
     type Response = ExecuteMsg<<E::Attestor as Attestor>::RawAttestation>;
 

examples/transfers/enclave/src/request/query.rs

@@ -1,15 +1,87 @@
-use quartz_common::enclave::{handler::Handler, Enclave};
+use cosmrs::AccountId;
+use cosmwasm_std::Uint128;
+use ecies::{decrypt, encrypt};
+use k256::ecdsa::{SigningKey, VerifyingKey};
+use quartz_common::enclave::{handler::Handler, key_manager::KeyManager, Enclave};
 use tonic::Status;
 use transfers_contract::msg::execute;
 
-use crate::proto::QueryRequest;
+use crate::{
+    proto::QueryRequest,
+    state::{RawBalance, State},
+    transfers_server::{QueryRequestMessage, RawCipherText},
+};
 
 #[async_trait::async_trait]
-impl<E: Enclave> Handler<E> for QueryRequest {
+impl<E: Enclave> Handler<E> for QueryRequest
+where
+    E: Enclave<Contract = AccountId>,
+    E::KeyManager: KeyManager<PubKey = VerifyingKey, PrivKey = SigningKey>,
+{
     type Error = Status;
     type Response = execute::QueryResponseMsg;
 
-    async fn handle(self, _ctx: &E) -> Result<Self::Response, Self::Error> {
-        todo!()
+    async fn handle(self, ctx: &E) -> Result<Self::Response, Self::Error> {
+        let message: QueryRequestMessage = {
+            let message: String = self.message;
+            serde_json::from_str(&message).map_err(|e| Status::invalid_argument(e.to_string()))?
+        };
+
+        // Decrypt and deserialize the state
+        let state = match &message.state.to_vec()[..] {
+            &[0] => State::default(),
+            state_bytes => {
+                let sk = ctx
+                    .key_manager()
+                    .await
+                    .priv_key()
+                    .await
+                    .ok_or_else(|| Status::internal("failed to get private key"))?;
+                decrypt_state(&sk, state_bytes)?
+            }
+        };
+
+        let bal = match state.state.get(&message.address) {
+            Some(balance) => RawBalance { balance: *balance },
+            None => RawBalance {
+                balance: Uint128::new(0),
+            },
+        };
+
+        // Parse the ephemeral public key
+        let ephemeral_pubkey =
+            VerifyingKey::from_sec1_bytes(&message.ephemeral_pubkey).map_err(|e| {
+                Status::invalid_argument(format!("Invalid ephemeral public key: {}", e))
+            })?;
+
+        // Encrypt the balance using the ephemeral public key
+        let bal_enc = encrypt_balance(bal, ephemeral_pubkey)
+            .map_err(|e| Status::internal(format!("Encryption error: {}", e)))?;
+
+        // Prepare message to chain
+        let msg = execute::QueryResponseMsg {
+            address: message.address,
+            encrypted_bal: bal_enc,
+        };
+
+        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: RawBalance,
+    ephemeral_pk: VerifyingKey,
+) -> Result<RawCipherText, 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))),
     }
 }

examples/transfers/enclave/src/request/update.rs

@@ -1,15 +1,216 @@
-use quartz_common::enclave::{handler::Handler, Enclave};
+use std::collections::btree_map::Entry;
+
+use cosmrs::AccountId;
+use cosmwasm_std::{Addr, HexBinary, Uint128};
+use ecies::{decrypt, encrypt};
+use k256::ecdsa::{SigningKey, VerifyingKey};
+use quartz_common::enclave::{
+    handler::Handler,
+    key_manager::KeyManager,
+    kv_store::{ConfigKey, ConfigKeyName, ContractKey, ContractKeyName, KvStore},
+    server::ProofOfPublication,
+    Enclave,
+};
 use tonic::Status;
-use transfers_contract::msg::execute;
+use transfers_contract::{
+    msg::{
+        execute,
+        execute::{ClearTextTransferRequestMsg, Request as TransfersRequest},
+    },
+    state::REQUESTS_KEY,
+};
 
-use crate::proto::UpdateRequest;
+use crate::{
+    proto::UpdateRequest,
+    state::State,
+    transfers_server::{RawCipherText, UpdateRequestMessage},
+};
 
 #[async_trait::async_trait]
-impl<E: Enclave> Handler<E> for UpdateRequest {
+impl<E: Enclave> Handler<E> for UpdateRequest
+where
+    E: Enclave<Contract = AccountId>,
+    E::KeyManager: KeyManager<PubKey = VerifyingKey, PrivKey = SigningKey>,
+{
     type Error = Status;
     type Response = execute::UpdateMsg;
 
-    async fn handle(self, _ctx: &E) -> Result<Self::Response, Self::Error> {
-        todo!()
+    async fn handle(self, ctx: &E) -> Result<Self::Response, Self::Error> {
+        // verify proof
+        let proof: ProofOfPublication<UpdateRequestMessage> = {
+            let message = self.message;
+            serde_json::from_str(&message).map_err(|e| Status::invalid_argument(e.to_string()))?
+        };
+        let contract = ctx
+            .store()
+            .await
+            .get(ContractKey::new(ContractKeyName))
+            .await
+            .map_err(|e| Status::internal(e.to_string()))?
+            .ok_or_else(|| Status::not_found("contract not found"))?;
+        let config = ctx
+            .store()
+            .await
+            .get(ConfigKey::new(ConfigKeyName))
+            .await
+            .map_err(|e| Status::internal(e.to_string()))?
+            .ok_or_else(|| Status::not_found("config not found"))?;
+        let (proof_value, message) = proof
+            .verify(
+                config.light_client_opts(),
+                contract,
+                REQUESTS_KEY.to_string(),
+                None,
+            )
+            .map_err(Status::failed_precondition)?;
+
+        let proof_value_matches_msg =
+            serde_json::to_string(&message.requests).is_ok_and(|s| s.as_bytes() == proof_value);
+        if !proof_value_matches_msg {
+            return Err(Status::failed_precondition("proof verification"));
+        }
+
+        // Decrypt and deserialize the state
+        let mut state = match &message.state.to_vec()[..] {
+            &[0] => State::default(),
+            state_bytes => {
+                let sk = ctx
+                    .key_manager()
+                    .await
+                    .priv_key()
+                    .await
+                    .ok_or_else(|| Status::internal("failed to get private key"))?;
+                decrypt_state(&sk, state_bytes)?
+            }
+        };
+
+        let requests_len = message.requests.len() as u32;
+
+        // 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 {
+                TransfersRequest::Transfer(ciphertext) => {
+                    // TODO: ensure_seq_num_consistency(message.seq_num, pending_sequenced_requests)?;
+
+                    // Decrypt transfer ciphertext into cleartext struct (acquires lock on enclave sk to do so)
+                    let transfer: ClearTextTransferRequestMsg = {
+                        let sk = ctx
+                            .key_manager()
+                            .await
+                            .priv_key()
+                            .await
+                            .ok_or_else(|| Status::internal("failed to get private key"))?;
+
+                        decrypt_transfer(&sk, &ciphertext)?
+                    };
+                    if let Entry::Occupied(mut entry) = state.state.entry(transfer.sender) {
+                        let balance = entry.get();
+                        if balance >= &transfer.amount {
+                            entry.insert(balance - transfer.amount);
+
+                            state
+                                .state
+                                .entry(transfer.receiver)
+                                .and_modify(|bal| *bal += transfer.amount)
+                                .or_insert(transfer.amount);
+                        }
+                        // TODO: handle errors
+                    }
+                }
+                TransfersRequest::Withdraw(receiver) => {
+                    // If a user with no balance requests withdraw, withdraw request for 0 coins gets processed
+                    // TODO: A no-op seems like a bad design choice in a privacy system
+                    if let Some(withdraw_bal) = state.state.remove(&receiver) {
+                        withdrawals_response.push((receiver, withdraw_bal));
+                    }
+                }
+                TransfersRequest::Deposit(sender, amount) => {
+                    state
+                        .state
+                        .entry(sender)
+                        .and_modify(|bal| *bal += amount)
+                        .or_insert(amount);
+                }
+            }
+        }
+
+        // Encrypt state
+        // Gets lock on PrivKey, generates PubKey to encrypt with
+        let state_enc = {
+            let pk = ctx
+                .key_manager()
+                .await
+                .pub_key()
+                .await
+                .ok_or_else(|| Status::internal("failed to get public key"))?;
+
+            encrypt_state(state, pk).map_err(|e| Status::invalid_argument(e.to_string()))?
+        };
+
+        // Prepare message to chain
+        let msg = execute::UpdateMsg {
+            ciphertext: state_enc,
+            quantity: requests_len,
+            withdrawals: withdrawals_response,
+        };
+
+        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<RawCipherText, 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(())
+// }

examples/transfers/enclave/src/state.rs

@@ -4,7 +4,7 @@ use anyhow;
 use cosmwasm_std::{Addr, HexBinary, Uint128};
 use serde::{Deserialize, Serialize};
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, PartialEq, Eq, Default, Deserialize, Serialize)]
 pub struct State {
     pub state: BTreeMap<Addr, Uint128>,
 }