Transfers app: query impl (#96)

Co-authored-by: Daniel Gushchyan <d.gushchyan@gmail.com>
Co-authored-by: Ethan Buchman <ethan@coinculture.info>
Co-authored-by: David Kajpust <kajpustd@gmail.com>
Co-authored-by: Daniel Gushchyan <39884512+dangush@users.noreply.github.com>
Co-authored-by: David Kajpust <davidkajpust@davids-mbp.home>
Co-authored-by: dave <davidkajpust@informal.systems>

apps/transfers/README.md

@@ -185,7 +185,9 @@ bash scripts/build.sh
 ### Configure and Run Gramine
 
 Setup and sign the Gramine config, and then start the gramine process, which will run the 
-grpc server that hosts the transfer application.
+grpc server that hosts the transfer application. The quartz port defaults to `11090`, but you can set it deliberately with 
+`export QUARTZ_PORT=XXXXX`. It is best to set it everytime, since if 2 people are sshing into the same machine to use
+the secure enclave, this could create undesired env conditions where your app is talking to the wrong enclave.
 
 ```
 bash scripts/start.sh

apps/transfers/contracts/Cargo.lock

@@ -1250,6 +1250,7 @@ dependencies = [
  "cw20-base",
  "getrandom",
  "quartz-cw",
+ "serde",
  "serde_json",
  "sha2 0.10.8",
  "thiserror",

apps/transfers/contracts/Cargo.toml

@@ -30,6 +30,7 @@ mock-sgx = ["quartz-cw/mock-sgx"]
 # external
 sha2 = "0.10.8"
 serde_json = "1.0.117"
+serde = { version = "1.0.137", default-features = false, features = ["derive"] }
 thiserror = { version = "1.0.49" }
 
 # cosmwasm

apps/transfers/contracts/src/contract.rs

@@ -1,13 +1,16 @@
-use cosmwasm_std::{entry_point, DepsMut, Env, HexBinary, MessageInfo, Response};
+use cosmwasm_std::{
+    entry_point, to_json_binary, Binary, Deps, DepsMut, Env, HexBinary, MessageInfo, Response,
+    StdResult,
+};
 use quartz_cw::handler::RawHandler;
 
 use crate::{
     error::ContractError,
     msg::{
-        execute::{Request, UpdateMsg},
-        ExecuteMsg, InstantiateMsg,
+        execute::{QueryResponseMsg, Request, UpdateMsg},
+        ExecuteMsg, InstantiateMsg, QueryMsg,
     },
-    state::{DENOM, REQUESTS, STATE},
+    state::{BALANCES, DENOM, REQUESTS, STATE},
 };
 
 #[cfg_attr(not(feature = "library"), entry_point)]
@@ -43,19 +46,58 @@ pub fn execute(
     use execute::*;
 
     match msg {
+        // Quartz msgs
         ExecuteMsg::Quartz(msg) => msg.handle_raw(deps, &env, &info).map_err(Into::into),
+
+        // Clear user msgs
+        ExecuteMsg::Deposit => deposit(deps, env, info),
+        ExecuteMsg::Withdraw => withdraw(deps, env, info),
+        ExecuteMsg::ClearTextTransferRequest(_) => unimplemented!(),
+        ExecuteMsg::QueryRequest(msg) => query_balance(deps, env, info, msg),
+
+        // Cipher user msgs
         ExecuteMsg::TransferRequest(msg) => transfer_request(deps, env, info, msg),
+
+        // Enclave msgs
         ExecuteMsg::Update(attested_msg) => {
             let _ = attested_msg
                 .clone()
                 .handle_raw(deps.branch(), &env, &info)?;
-
-            // Extract underlying UpdateMsg and pass to update()
-            update(deps, env, info, UpdateMsg(attested_msg.msg))
+            let UpdateMsg {
+                ciphertext,
+                quantity,
+                withdrawals,
+            } = attested_msg.msg.0;
+            update(
+                deps,
+                env,
+                info,
+                UpdateMsg {
+                    ciphertext,
+                    quantity,
+                    withdrawals,
+                },
+            )
+        }
+
+        ExecuteMsg::QueryResponse(attested_msg) => {
+            let _ = attested_msg
+                .clone()
+                .handle_raw(deps.branch(), &env, &info)?;
+            let QueryResponseMsg {
+                address,
+                encrypted_bal,
+            } = attested_msg.msg.0;
+            store_balance(
+                deps,
+                env,
+                info,
+                QueryResponseMsg {
+                    address,
+                    encrypted_bal,
+                },
+            )
         }
-        ExecuteMsg::Deposit => deposit(deps, env, info),
-        ExecuteMsg::Withdraw => withdraw(deps, env, info),
-        ExecuteMsg::ClearTextTransferRequest(_) => unimplemented!(),
     }
 }
 
@@ -65,10 +107,56 @@ pub mod execute {
 
     use crate::{
         error::ContractError,
-        msg::execute::{Request, TransferRequestMsg, UpdateMsg},
-        state::{DENOM, REQUESTS, STATE},
+        msg::execute::{QueryRequestMsg, QueryResponseMsg, Request, TransferRequestMsg, UpdateMsg},
+        state::{BALANCES, DENOM, REQUESTS, STATE},
     };
 
+    pub fn deposit(deps: DepsMut, _env: Env, info: MessageInfo) -> Result<Response, ContractError> {
+        let denom: String = DENOM.load(deps.storage)?;
+        let quantity = must_pay(&info, &denom)?;
+
+        let mut requests = REQUESTS.load(deps.storage)?;
+
+        requests.push(Request::Deposit(info.sender, quantity));
+
+        REQUESTS.save(deps.storage, &requests)?;
+
+        let event = Event::new("transfer").add_attribute("action", "user");
+        let resp = Response::new().add_event(event);
+
+        Ok(resp)
+    }
+
+    pub fn withdraw(
+        deps: DepsMut,
+        _env: Env,
+        info: MessageInfo,
+    ) -> Result<Response, ContractError> {
+        let mut requests = REQUESTS.load(deps.storage)?;
+
+        requests.push(Request::Withdraw(info.sender));
+
+        REQUESTS.save(deps.storage, &requests)?;
+
+        let event = Event::new("transfer").add_attribute("action", "user");
+        let resp = Response::new().add_event(event);
+
+        Ok(resp)
+    }
+
+    pub fn query_balance(
+        _deps: DepsMut,
+        _env: Env,
+        _info: MessageInfo,
+        msg: QueryRequestMsg,
+    ) -> Result<Response, ContractError> {
+        let event = Event::new("query_balance")
+            .add_attribute("query", "user")
+            .add_attribute("emphemeral_pubkey", msg.emphemeral_pubkey.to_string());
+        let resp = Response::new().add_event(event);
+        Ok(resp)
+    }
+
     pub fn transfer_request(
         deps: DepsMut,
         _env: Env,
@@ -94,12 +182,12 @@ pub mod execute {
         msg: UpdateMsg,
     ) -> Result<Response, ContractError> {
         // Store state
-        STATE.save(deps.storage, &msg.0.ciphertext)?;
+        STATE.save(deps.storage, &msg.ciphertext)?;
 
         // Clear queue
         let mut requests: Vec<Request> = REQUESTS.load(deps.storage)?;
 
-        requests.drain(0..msg.0.quantity as usize);
+        requests.drain(0..msg.quantity as usize);
 
         REQUESTS.save(deps.storage, &requests)?;
 
@@ -107,7 +195,6 @@ pub mod execute {
         let denom = DENOM.load(deps.storage)?;
 
         let messages = msg
-            .0
             .withdrawals
             .into_iter()
             .map(|(user, funds)| BankMsg::Send {
@@ -120,38 +207,36 @@ pub mod execute {
         Ok(resp)
     }
 
-    pub fn deposit(deps: DepsMut, _env: Env, info: MessageInfo) -> Result<Response, ContractError> {
-        let denom = DENOM.load(deps.storage)?;
-        let quantity = must_pay(&info, &denom)?;
-
-        let mut requests = REQUESTS.load(deps.storage)?;
-
-        requests.push(Request::Deposit(info.sender, quantity));
-
-        REQUESTS.save(deps.storage, &requests)?;
-
-        let event = Event::new("transfer").add_attribute("action", "user");
-        let resp = Response::new().add_event(event);
-
-        Ok(resp)
-    }
-
-    pub fn withdraw(
+    pub fn store_balance(
         deps: DepsMut,
         _env: Env,
-        info: MessageInfo,
+        _info: MessageInfo,
+        msg: QueryResponseMsg,
     ) -> Result<Response, ContractError> {
-        // TODO: verify denom
+        // Store state
+        BALANCES.save(deps.storage, &msg.address.to_string(), &msg.encrypted_bal)?;
 
-        let mut requests = REQUESTS.load(deps.storage)?;
-
-        requests.push(Request::Withdraw(info.sender));
-
-        REQUESTS.save(deps.storage, &requests)?;
-
-        let event = Event::new("transfer").add_attribute("action", "user");
+        // Emit event
+        let event = Event::new("store_balance")
+            .add_attribute("query", "enclave") // TODO Weird to name it enclave?
+            .add_attribute("address", msg.address.to_string())
+            .add_attribute("encrypted_balance", msg.encrypted_bal.to_string());
         let resp = Response::new().add_event(event);
-
         Ok(resp)
     }
 }
+
+#[cfg_attr(not(feature = "library"), entry_point)]
+pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> StdResult<Binary> {
+    match msg {
+        QueryMsg::GetBalance { address } => to_json_binary(&query::get_balance(deps, address)?),
+    }
+}
+mod query {
+    use super::*;
+
+    pub fn get_balance(deps: Deps, address: String) -> StdResult<HexBinary> {
+        let balance = BALANCES.may_load(deps.storage, &address)?;
+        Ok(balance.unwrap_or_default())
+    }
+}

apps/transfers/contracts/src/msg.rs

@@ -1,9 +1,11 @@
 use cosmwasm_schema::cw_serde;
-use cosmwasm_std::{Addr, HexBinary, Uint128};
 use quartz_cw::{
-    msg::execute::attested::{RawAttested, RawDefaultAttestation},
+    msg::execute::attested::{RawAttested, RawAttestedMsgSansHandler, RawDefaultAttestation},
     prelude::*,
 };
+use serde::{Deserialize, Serialize};
+
+type AttestedMsg<M, RA = RawDefaultAttestation> = RawAttested<RawAttestedMsgSansHandler<M>, RA>;
 
 #[cw_serde]
 pub struct InstantiateMsg<RA = RawDefaultAttestation> {
@@ -11,45 +13,37 @@ pub struct InstantiateMsg<RA = RawDefaultAttestation> {
     pub denom: String,
 }
 
+#[cw_serde]
+pub enum QueryMsg {
+    GetBalance { address: String },
+}
+
 #[cw_serde]
 #[allow(clippy::large_enum_variant)]
 pub enum ExecuteMsg<RA = RawDefaultAttestation> {
     // quartz initialization
     Quartz(QuartzExecuteMsg),
 
-    // ----- user txs
+    // User msgs
     // clear text
     Deposit,
     Withdraw,
-
+    ClearTextTransferRequest(execute::ClearTextTransferRequestMsg),
     // ciphertext
     TransferRequest(execute::TransferRequestMsg),
-    // ---- end user txs
-    ClearTextTransferRequest(execute::ClearTextTransferRequestMsg),
+    QueryRequest(execute::QueryRequestMsg),
 
-    // enclave msg
-    Update(RawAttested<execute::RawUpdateMsg, RA>),
+    // Enclave msgs
+    Update(AttestedMsg<execute::UpdateMsg, RA>),
+    QueryResponse(AttestedMsg<execute::QueryResponseMsg, RA>),
 }
 
 pub mod execute {
-    use cosmwasm_std::{DepsMut, Env, MessageInfo, Response, StdError};
-    use quartz_cw::{
-        error::Error,
-        handler::Handler,
-        msg::{execute::attested::HasUserData, HasDomainType},
-        state::UserData,
-    };
+    use cosmwasm_schema::cw_serde;
+    use cosmwasm_std::{Addr, HexBinary, Uint128};
+    use quartz_cw::{msg::execute::attested::HasUserData, state::UserData};
     use sha2::{Digest, Sha256};
 
-    use super::*;
-
-    #[cw_serde]
-    pub struct TransferRequestMsg {
-        pub ciphertext: HexBinary,
-        pub digest: HexBinary,
-        // pub proof: π
-    }
-
     #[cw_serde]
     pub struct ClearTextTransferRequestMsg {
         pub sender: Addr,
@@ -58,7 +52,18 @@ pub mod execute {
         // pub proof: π
     }
 
-    // Ciphertext of a transfer request
+    #[cw_serde]
+    pub struct QueryRequestMsg {
+        pub emphemeral_pubkey: HexBinary,
+    }
+
+    #[cw_serde]
+    pub struct TransferRequestMsg {
+        pub ciphertext: HexBinary,
+        pub digest: HexBinary,
+        // pub proof: π
+    }
+
     #[cw_serde]
     pub enum Request {
         Transfer(HexBinary),
@@ -67,20 +72,17 @@ pub mod execute {
     }
 
     #[cw_serde]
-    pub struct RawUpdateMsg {
+    pub struct UpdateMsg {
         pub ciphertext: HexBinary,
         pub quantity: u32,
         pub withdrawals: Vec<(Addr, Uint128)>,
         // pub proof: π
     }
 
-    #[derive(Clone, Debug, PartialEq)]
-    pub struct UpdateMsg(pub RawUpdateMsg);
-
     impl HasUserData for UpdateMsg {
         fn user_data(&self) -> UserData {
             let mut hasher = Sha256::new();
-            hasher.update(serde_json::to_string(&self.0).expect("infallible serializer"));
+            hasher.update(serde_json::to_string(&self).expect("infallible serializer"));
             let digest: [u8; 32] = hasher.finalize().into();
 
             let mut user_data = [0u8; 64];
@@ -89,33 +91,22 @@ pub mod execute {
         }
     }
 
-    impl HasDomainType for RawUpdateMsg {
-        type DomainType = UpdateMsg;
+    #[cw_serde]
+    pub struct QueryResponseMsg {
+        pub address: Addr,
+        pub encrypted_bal: HexBinary,
+        // pub proof: π
     }
 
-    impl TryFrom<RawUpdateMsg> for UpdateMsg {
-        type Error = StdError;
+    impl HasUserData for QueryResponseMsg {
+        fn user_data(&self) -> UserData {
+            let mut hasher = Sha256::new();
+            hasher.update(serde_json::to_string(&self).expect("infallible serializer"));
+            let digest: [u8; 32] = hasher.finalize().into();
 
-        fn try_from(value: RawUpdateMsg) -> Result<Self, Self::Error> {
-            Ok(Self(value))
-        }
-    }
-
-    impl From<UpdateMsg> for RawUpdateMsg {
-        fn from(value: UpdateMsg) -> Self {
-            value.0
-        }
-    }
-
-    impl Handler for UpdateMsg {
-        fn handle(
-            self,
-            _deps: DepsMut<'_>,
-            _env: &Env,
-            _info: &MessageInfo,
-        ) -> Result<Response, Error> {
-            // basically handle `transfer_request` here
-            Ok(Response::default())
+            let mut user_data = [0u8; 64];
+            user_data[0..32].copy_from_slice(&digest);
+            user_data
         }
     }
 }

apps/transfers/contracts/src/state.rs

@@ -1,9 +1,9 @@
 use cosmwasm_std::HexBinary;
-use cw_storage_plus::Item;
+use cw_storage_plus::{Item, Map};
 
 use crate::msg::execute::Request;
 
 pub const STATE: Item<HexBinary> = Item::new("state");
 pub const REQUESTS: Item<Vec<Request>> = Item::new("requests");
-
 pub const DENOM: Item<String> = Item::new("donation_denom");
+pub const BALANCES: Map<&str, HexBinary> = Map::new("balances");

apps/transfers/enclave/proto/transfers.proto

@@ -3,13 +3,22 @@ syntax = "proto3";
 package transfers;
 
 service Settlement {
-  rpc Run (RunTransfersRequest) returns (RunTransfersResponse) {}
+  rpc Run (UpdateRequest) returns (UpdateResponse) {}
+  rpc Query (QueryRequest) returns (QueryResponse) {}
 }
 
-message RunTransfersRequest {
+message UpdateRequest {
   string message = 1;
 }
 
-message RunTransfersResponse {
+message UpdateResponse {
+  string message = 1;
+}
+
+message QueryRequest {
+  string message = 1;
+}
+
+message QueryResponse {
   string message = 1;
 }

apps/transfers/enclave/src/prost/transfers.rs

@@ -1,13 +1,25 @@
 // This file is @generated by prost-build.
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]
-pub struct RunTransfersRequest {
+pub struct UpdateRequest {
     #[prost(string, tag = "1")]
     pub message: ::prost::alloc::string::String,
 }
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]
-pub struct RunTransfersResponse {
+pub struct UpdateResponse {
+    #[prost(string, tag = "1")]
+    pub message: ::prost::alloc::string::String,
+}
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct QueryRequest {
+    #[prost(string, tag = "1")]
+    pub message: ::prost::alloc::string::String,
+}
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct QueryResponse {
     #[prost(string, tag = "1")]
     pub message: ::prost::alloc::string::String,
 }
@@ -98,11 +110,8 @@ pub mod settlement_client {
         }
         pub async fn run(
             &mut self,
-            request: impl tonic::IntoRequest<super::RunTransfersRequest>,
-        ) -> std::result::Result<
-            tonic::Response<super::RunTransfersResponse>,
-            tonic::Status,
-        > {
+            request: impl tonic::IntoRequest<super::UpdateRequest>,
+        ) -> std::result::Result<tonic::Response<super::UpdateResponse>, tonic::Status> {
             self.inner
                 .ready()
                 .await
@@ -118,6 +127,28 @@ pub mod settlement_client {
             req.extensions_mut().insert(GrpcMethod::new("transfers.Settlement", "Run"));
             self.inner.unary(req, path, codec).await
         }
+        pub async fn query(
+            &mut self,
+            request: impl tonic::IntoRequest<super::QueryRequest>,
+        ) -> std::result::Result<tonic::Response<super::QueryResponse>, tonic::Status> {
+            self.inner
+                .ready()
+                .await
+                .map_err(|e| {
+                    tonic::Status::new(
+                        tonic::Code::Unknown,
+                        format!("Service was not ready: {}", e.into()),
+                    )
+                })?;
+            let codec = tonic::codec::ProstCodec::default();
+            let path = http::uri::PathAndQuery::from_static(
+                "/transfers.Settlement/Query",
+            );
+            let mut req = request.into_request();
+            req.extensions_mut()
+                .insert(GrpcMethod::new("transfers.Settlement", "Query"));
+            self.inner.unary(req, path, codec).await
+        }
     }
 }
 /// Generated server implementations.
@@ -129,11 +160,12 @@ pub mod settlement_server {
     pub trait Settlement: Send + Sync + 'static {
         async fn run(
             &self,
-            request: tonic::Request<super::RunTransfersRequest>,
-        ) -> std::result::Result<
-            tonic::Response<super::RunTransfersResponse>,
-            tonic::Status,
-        >;
+            request: tonic::Request<super::UpdateRequest>,
+        ) -> std::result::Result<tonic::Response<super::UpdateResponse>, tonic::Status>;
+        async fn query(
+            &self,
+            request: tonic::Request<super::QueryRequest>,
+        ) -> std::result::Result<tonic::Response<super::QueryResponse>, tonic::Status>;
     }
     #[derive(Debug)]
     pub struct SettlementServer<T: Settlement> {
@@ -217,18 +249,16 @@ pub mod settlement_server {
                 "/transfers.Settlement/Run" => {
                     #[allow(non_camel_case_types)]
                     struct RunSvc<T: Settlement>(pub Arc<T>);
-                    impl<
-                        T: Settlement,
-                    > tonic::server::UnaryService<super::RunTransfersRequest>
+                    impl<T: Settlement> tonic::server::UnaryService<super::UpdateRequest>
                     for RunSvc<T> {
-                        type Response = super::RunTransfersResponse;
+                        type Response = super::UpdateResponse;
                         type Future = BoxFuture<
                             tonic::Response<Self::Response>,
                             tonic::Status,
                         >;
                         fn call(
                             &mut self,
-                            request: tonic::Request<super::RunTransfersRequest>,
+                            request: tonic::Request<super::UpdateRequest>,
                         ) -> Self::Future {
                             let inner = Arc::clone(&self.0);
                             let fut = async move {
@@ -260,6 +290,50 @@ pub mod settlement_server {
                     };
                     Box::pin(fut)
                 }
+                "/transfers.Settlement/Query" => {
+                    #[allow(non_camel_case_types)]
+                    struct QuerySvc<T: Settlement>(pub Arc<T>);
+                    impl<T: Settlement> tonic::server::UnaryService<super::QueryRequest>
+                    for QuerySvc<T> {
+                        type Response = super::QueryResponse;
+                        type Future = BoxFuture<
+                            tonic::Response<Self::Response>,
+                            tonic::Status,
+                        >;
+                        fn call(
+                            &mut self,
+                            request: tonic::Request<super::QueryRequest>,
+                        ) -> Self::Future {
+                            let inner = Arc::clone(&self.0);
+                            let fut = async move {
+                                <T as Settlement>::query(&inner, request).await
+                            };
+                            Box::pin(fut)
+                        }
+                    }
+                    let accept_compression_encodings = self.accept_compression_encodings;
+                    let send_compression_encodings = self.send_compression_encodings;
+                    let max_decoding_message_size = self.max_decoding_message_size;
+                    let max_encoding_message_size = self.max_encoding_message_size;
+                    let inner = self.inner.clone();
+                    let fut = async move {
+                        let inner = inner.0;
+                        let method = QuerySvc(inner);
+                        let codec = tonic::codec::ProstCodec::default();
+                        let mut grpc = tonic::server::Grpc::new(codec)
+                            .apply_compression_config(
+                                accept_compression_encodings,
+                                send_compression_encodings,
+                            )
+                            .apply_max_message_size_config(
+                                max_decoding_message_size,
+                                max_encoding_message_size,
+                            );
+                        let res = grpc.unary(method, req).await;
+                        Ok(res)
+                    };
+                    Box::pin(fut)
+                }
                 _ => {
                     Box::pin(async move {
                         Ok(

apps/transfers/enclave/src/state.rs

@@ -32,3 +32,27 @@ impl TryFrom<RawState> for State {
         Ok(Self { state: o.state })
     }
 }
+
+#[derive(Clone, Debug)]
+pub struct Balance {
+    pub balance: Uint128,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)]
+pub struct RawBalance {
+    pub balance: Uint128,
+}
+
+impl From<Balance> for RawBalance {
+    fn from(o: Balance) -> Self {
+        Self { balance: o.balance }
+    }
+}
+
+impl TryFrom<RawBalance> for Balance {
+    type Error = anyhow::Error;
+
+    fn try_from(o: RawBalance) -> Result<Self, anyhow::Error> {
+        Ok(Self { balance: o.balance })
+    }
+}

apps/transfers/enclave/src/transfers_server.rs

@@ -17,8 +17,10 @@ use tonic::{Request, Response, Result as TonicResult, Status};
 use transfers_contract::msg::execute::{ClearTextTransferRequestMsg, Request as TransfersRequest};
 
 use crate::{
-    proto::{settlement_server::Settlement, RunTransfersRequest, RunTransfersResponse},
-    state::{RawState, State},
+    proto::{
+        settlement_server::Settlement, QueryRequest, QueryResponse, UpdateRequest, UpdateResponse,
+    },
+    state::{RawBalance, RawState, State},
 };
 
 pub type RawCipherText = HexBinary;
@@ -30,19 +32,32 @@ pub struct TransfersService<A> {
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
-pub struct RunTransfersRequestMessage {
+pub struct UpdateRequestMessage {
     state: HexBinary,
     requests: Vec<TransfersRequest>,
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
-pub struct RunTransfersResponseMessage {
+pub struct QueryRequestMessage {
+    state: HexBinary,
+    address: Addr,
+    ephemeral_pubkey: HexBinary,
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct QueryResponseMessage {
+    address: Addr,
+    encrypted_bal: HexBinary,
+}
+
+#[derive(PartialEq, Clone, Debug, Serialize, Deserialize)]
+pub struct UpdateMsg {
     ciphertext: HexBinary,
     quantity: u32,
     withdrawals: Vec<(Addr, Uint128)>,
 }
 
-impl HasUserData for RunTransfersResponseMessage {
+impl HasUserData for UpdateMsg {
     fn user_data(&self) -> UserData {
         let mut hasher = Sha256::new();
         hasher.update(serde_json::to_string(&self).expect("infallible serializer"));
@@ -54,6 +69,24 @@ impl HasUserData for RunTransfersResponseMessage {
     }
 }
 
+impl HasUserData for QueryResponseMessage {
+    fn user_data(&self) -> UserData {
+        let mut hasher = Sha256::new();
+        hasher.update(serde_json::to_string(&self).expect("infallible serializer"));
+        let digest: [u8; 32] = hasher.finalize().into();
+
+        let mut user_data = [0u8; 64];
+        user_data[0..32].copy_from_slice(&digest);
+        user_data
+    }
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+struct AttestedMsg<M> {
+    msg: M,
+    quote: Vec<u8>,
+}
+
 impl<A> TransfersService<A>
 where
     A: Attestor,
@@ -68,14 +101,11 @@ impl<A> Settlement for TransfersService<A>
 where
     A: Attestor + Send + Sync + 'static,
 {
-    async fn run(
-        &self,
-        request: Request<RunTransfersRequest>,
-    ) -> TonicResult<Response<RunTransfersResponse>> {
+    async fn run(&self, request: Request<UpdateRequest>) -> TonicResult<Response<UpdateResponse>> {
         // Request contains a serialized json string
 
         // Serialize request into struct containing State and the Requests vec
-        let message: RunTransfersRequestMessage = {
+        let message: UpdateRequestMessage = {
             let message = request.into_inner().message;
             serde_json::from_str(&message).map_err(|e| Status::invalid_argument(e.to_string()))?
         };
@@ -83,8 +113,6 @@ where
         // Decrypt and deserialize the state
         let mut state = {
             if message.state.len() == 1 && message.state[0] == 0 {
-                println!("{}", message.state);
-
                 State {
                     state: BTreeMap::<Addr, Uint128>::new(),
                 }
@@ -103,7 +131,7 @@ where
 
         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)>::new();
+        let mut withdrawals_response: Vec<(Addr, Uint128)> = Vec::<(Addr, Uint128)>::new();
 
         // Loop through requests, match on cases, and apply changes to state
         for req in message.requests {
@@ -170,7 +198,7 @@ where
         };
 
         // Prepare message to chain
-        let msg = RunTransfersResponseMessage {
+        let msg = UpdateMsg {
             ciphertext: state_enc,
             quantity: requests_len,
             withdrawals: withdrawals_response,
@@ -187,7 +215,69 @@ where
         let message =
             serde_json::to_string(&attested_msg).map_err(|e| Status::internal(e.to_string()))?;
 
-        Ok(Response::new(RunTransfersResponse { message }))
+        Ok(Response::new(UpdateResponse { message }))
+    }
+
+    async fn query(&self, request: Request<QueryRequest>) -> TonicResult<Response<QueryResponse>> {
+        // Serialize request into struct containing State and the Requests vec
+        let message: QueryRequestMessage = {
+            let message: String = request.into_inner().message;
+            serde_json::from_str(&message).map_err(|e| Status::invalid_argument(e.to_string()))?
+        };
+
+        // Decrypt and deserialize the state
+        let state = {
+            if message.state.len() == 1 && message.state[0] == 0 {
+                State {
+                    state: BTreeMap::<Addr, Uint128>::new(),
+                }
+            } else {
+                let sk_lock = self
+                    .sk
+                    .lock()
+                    .map_err(|e| Status::internal(e.to_string()))?;
+                let sk = sk_lock
+                    .as_ref()
+                    .ok_or(Status::internal("SigningKey unavailable"))?;
+                decrypt_state(sk, &message.state)?
+            }
+        };
+
+        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 = QueryResponseMessage {
+            address: message.address,
+            encrypted_bal: bal_enc,
+        };
+
+        // Attest to message
+        let attestation: HexBinary = self
+            .attestor
+            .quote(msg.clone())
+            .map_err(|e| Status::internal(e.to_string()))?
+            .into();
+
+        let attested_msg = RawAttested { msg, attestation };
+        let message =
+            serde_json::to_string(&attested_msg).map_err(|e| Status::internal(e.to_string()))?;
+
+        Ok(Response::new(QueryResponse { message }))
     }
 }
 
@@ -221,3 +311,12 @@ fn encrypt_state(state: RawState, enclave_pk: VerifyingKey) -> TonicResult<RawCi
         Err(e) => Err(Status::internal(format!("Encryption error: {}", e))),
     }
 }
+
+fn encrypt_balance(balance: RawBalance, ephemeral_pk: VerifyingKey) -> TonicResult<RawCipherText> {
+    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))),
+    }
+}

apps/transfers/scripts/decrypt_ephemeral.sh

@@ -0,0 +1,50 @@
+#!/bin/bash
+
+# Check if an encrypted balance is provided
+if [ "$#" -ne 1 ]; then
+    echo "Usage: $0 <encrypted_balance_hex>"
+    exit 1
+fi
+
+ENCRYPTED_BALANCE=$1
+
+# Check if eciespy is installed
+if ! pip list 2>/dev/null | grep -q eciespy; then
+    echo "eciespy is not installed. Installing now..."
+    pip install eciespy >/dev/null 2>&1
+fi
+
+# Extract the private key from wasmd
+EPHEMERAL_PRIVKEY=$( yes | wasmd keys export ephemeral_user --unsafe --unarmored-hex)
+
+if [ $? -ne 0 ]; then
+    echo "Failed to export private key. Make sure 'ephemeral_user' exists and you've entered the correct password."
+    exit 1
+fi
+
+# Create a temporary Python script for decryption
+TEMP_PYTHON_SCRIPT=$(mktemp)
+cat << EOF > "$TEMP_PYTHON_SCRIPT"
+from ecies import decrypt
+import binascii
+import sys
+
+private_key = bytes.fromhex(sys.argv[1])
+encrypted = binascii.unhexlify(sys.argv[2])
+
+try:
+    decrypted = decrypt(private_key, encrypted)
+    print(decrypted.decode())
+except Exception as e:
+    print(f"Decryption failed: {str(e)}")
+EOF
+
+# Run the Python script to decrypt
+DECRYPTED=$(python3 "$TEMP_PYTHON_SCRIPT" "$EPHEMERAL_PRIVKEY" "$ENCRYPTED_BALANCE")
+
+echo "---------------------------------------------------------"
+echo "Decrypted result:"
+echo "$DECRYPTED"
+
+# Clean up
+rm "$TEMP_PYTHON_SCRIPT"

apps/transfers/scripts/listen.sh

@@ -1,5 +1,3 @@
-
-
 ROOT=${ROOT:-$HOME}
 
 DEFAULT_NODE="127.0.0.1:26657"
@@ -19,58 +17,129 @@ CMD="wasmd --node http://$NODE_URL"
 
 WSURL="ws://$NODE_URL/websocket"
 
-SUBSCRIBE="{\"jsonrpc\":\"2.0\",\"method\":\"subscribe\",\"params\":[\"execute._contract_address = '$CONTRACT' AND wasm-transfer.action = 'user'\"],\"id\":1}"
+SUBSCRIBE_TRANSFER="{\"jsonrpc\":\"2.0\",\"method\":\"subscribe\",\"params\":[\"execute._contract_address = '$CONTRACT' AND wasm-transfer.action = 'user'\"],\"id\":1}"
+SUBSCRIBE_QUERY="{\"jsonrpc\":\"2.0\",\"method\":\"subscribe\",\"params\":[\"execute._contract_address = '$CONTRACT' AND wasm-query_balance.query = 'user'\"],\"id\":2}"
 
-echo $SUBSCRIBE
-
-echo "--------------------------------------------------------"
-echo "subscribe to events"
+# Attestation constants
+IAS_API_KEY="669244b3e6364b5888289a11d2a1726d"
+RA_CLIENT_SPID="51CAF5A48B450D624AEFE3286D314894"
+QUOTE_FILE="/tmp/${USER}_test.quote"
+REPORT_FILE="/tmp/${USER}_datareport"
+REPORT_SIG_FILE="/tmp/${USER}_datareportsig"
 
 # cat keeps the stdin open so websocat doesnt close
-(echo "$SUBSCRIBE"; cat) | websocat $WSURL | while read msg; do 
-    if [[ "$msg" == '{"jsonrpc":"2.0","id":1,"result":{}}' ]]; then
-        echo "... subscribed"
+(echo "$SUBSCRIBE_TRANSFER"; echo "$SUBSCRIBE_QUERY"; cat) | websocat $WSURL | while read msg; do 
+    if [[ "$msg" == '{"jsonrpc":"2.0","id":1,"result":{}}' ]] || [[ "$msg" == '{"jsonrpc":"2.0","id":2,"result":{}}' ]]; then
         echo "---------------------------------------------------------"
+        echo "... subscribed to $msg"
         echo "... waiting for event"
         continue
     fi 
 
-    if echo "$msg" | jq 'has("error")' > /dev/null; then
-        echo "... error msg $msg"
+    # TODO - Some reason this is saying ERROR when its fine, needs to be fixed or removed
+    #if echo "$msg" | sed 's/"log":"\[.*\]"/"log":"<invalid_json>"/' | jq 'has("error")' > /dev/null; then
+     #   echo "... error msg $msg"
+     #   echo "---------------------------------------------------------"
+     #   echo "... waiting for event"
+     #   continue
+    #fi 
+
+    CLEAN_MSG=$(echo "$msg" | sed 's/"log":"\[.*\]"/"log":"<invalid_json>"/' | jq '.result.events')
+
+    if echo "$CLEAN_MSG" | grep -q 'wasm-transfer'; then
         echo "---------------------------------------------------------"
-        echo "... waiting for event"
-        continue
-    fi 
-
-
-    echo "... received event! "
-    echo $msg 
+        echo "... received wasm-transfer event!"
 
         echo "... fetching requests"
-
         REQUESTS=$($CMD query wasm contract-state raw $CONTRACT $(printf '%s' "requests" | hexdump -ve '/1 "%02X"') -o json | jq -r .data | base64 -d)
         STATE=$($CMD query wasm contract-state raw $CONTRACT $(printf '%s' "state" | hexdump -ve '/1 "%02X"') -o json | jq -r .data | base64 -d)
-
         export ENCLAVE_REQUEST=$(jq -nc --argjson requests "$REQUESTS" --argjson state $STATE '$ARGS.named')
-    echo $ENCLAVE_REQUEST | jq .
-
         export REQUEST_MSG=$(jq -nc --arg message "$ENCLAVE_REQUEST" '$ARGS.named')
 
         cd $ROOT/cycles-quartz/apps/transfers/enclave
 
         echo "... executing transfer"
-    export UPDATE=$(grpcurl -plaintext -import-path ./proto/ -proto transfers.proto -d "$REQUEST_MSG" "127.0.0.1:$QUARTZ_PORT" transfers.Settlement/Run | jq .message | jq -R 'fromjson | fromjson' | jq -c )
-
+        export ATTESTED_MSG=$(grpcurl -plaintext -import-path ./proto/ -proto transfers.proto -d "$REQUEST_MSG" "127.0.0.1:$QUARTZ_PORT" transfers.Settlement/Run | jq .message | jq -R 'fromjson | fromjson' | jq -c )
+        QUOTE=$(echo "$ATTESTED_MSG" | jq -c '.attestation')
+        MSG=$(echo "$ATTESTED_MSG" | jq -c '.msg')
 
+        echo "... getting report"
+        echo -n "$QUOTE" | xxd -r -p - > "$QUOTE_FILE"
+        gramine-sgx-ias-request report -g "$RA_CLIENT_SPID" -k "$IAS_API_KEY" -q "$QUOTE_FILE" -r "$REPORT_FILE" -s "$REPORT_SIG_FILE" > /dev/null 2>&1
+        REPORT=$(cat "$REPORT_FILE")
+        REPORTSIG=$(cat "$REPORT_SIG_FILE" | tr -d '\r')
 
         echo "... submitting update"
-
-    $CMD tx wasm execute $CONTRACT "{\"update\": "$UPDATE" }" --chain-id testing --from admin --node http://$NODE_URL -y
-
+        export EXECUTE=$(jq -nc --argjson update "$(jq -nc --argjson msg "$MSG" --argjson attestation \
+            "$(jq -nc --argjson report "$(jq -nc --argjson report "$REPORT" --arg reportsig "$REPORTSIG" '$ARGS.named')" '$ARGS.named')" \
+            '$ARGS.named')" '$ARGS.named')
+        echo $EXECUTE | jq '.'
+        $CMD tx wasm execute "$CONTRACT" "$EXECUTE" --from admin --chain-id testing -y --gas 2000000
 
         echo " ... done"
         echo "---------------------------------------------------------"
         echo "... waiting for event"
+    elif echo "$CLEAN_MSG" | grep -q 'wasm-query_balance'; then
+        echo "... received wasm-query_balance event!"
+        echo "... fetching state"
+
+        STATE=$($CMD query wasm contract-state raw $CONTRACT $(printf '%s' "state" | hexdump -ve '/1 "%02X"') -o json | jq -r .data | base64 -d)
+
+        # Extract the address from the event
+        ADDRESS=$(echo "$msg" | sed 's/"log":"\[.*\]"/"log":"<invalid_json>"/' | jq -r '.result.events["message.sender"]'[0])
+
+        EPHEMERAL_PUBKEY=$(echo "$msg" | sed 's/"log":"\[.*\]"/"log":"<invalid_json>"/' | jq -r '.result.events["wasm-query_balance.emphemeral_pubkey"]'[0])
+
+        # Create the enclave request with state and address
+        export ENCLAVE_REQUEST=$(jq -nc --argjson state "$STATE" --arg address "$ADDRESS" --arg ephemeral_pubkey "$EPHEMERAL_PUBKEY" '$ARGS.named')
+        export REQUEST_MSG=$(jq -nc --arg message "$ENCLAVE_REQUEST" '$ARGS.named')
+
+        cd $ROOT/cycles-quartz/apps/transfers/enclave
+
+        echo "... executing query balance"
+        ATTESTED_MSG=$(grpcurl -plaintext -import-path ./proto/ -proto transfers.proto -d "$REQUEST_MSG" "127.0.0.1:$QUARTZ_PORT" transfers.Settlement/Query | jq -r '.message | fromjson')
+        echo "atts msg"
+        echo $ATTESTED_MSG
+        QUOTE=$(echo "$ATTESTED_MSG" | jq -c '.attestation')
+        MSG=$(echo "$ATTESTED_MSG" | jq -c '.msg')
+        echo "quote"
+        echo $QUOTE
+        echo "msg"
+        echo $MSG
+
+        echo -n "$QUOTE" | xxd -r -p - > "$QUOTE_FILE"
+        gramine-sgx-ias-request report -g "$RA_CLIENT_SPID" -k "$IAS_API_KEY" -q "$QUOTE_FILE" -r "$REPORT_FILE" -s "$REPORT_SIG_FILE" > /dev/null 2>&1
+        REPORT=$(cat "$REPORT_FILE")
+        REPORTSIG=$(cat "$REPORT_SIG_FILE" | tr -d '\r')
+
+        echo "... submitting update"
+
+        # Create the QueryResponseMsg structure with address inside the msg
+        export QUERY_RESPONSE_MSG=$(jq -n \
+            --arg address "$ADDRESS" \
+            --argjson msg "$MSG" \
+            '{address: $address, encrypted_bal: $msg.encrypted_bal}')
+
+
+        # Create the execute message for query_response
+        export EXECUTE=$(jq -nc \
+            --argjson query_response "$(jq -nc \
+                --argjson msg "$QUERY_RESPONSE_MSG" \
+                --argjson attestation "$(jq -nc \
+                    --argjson report "$(jq -nc \
+                        --argjson report "$REPORT" \
+                        --arg reportsig "$REPORTSIG" \
+                        '$ARGS.named')" \
+                    '$ARGS.named')" \
+                '$ARGS.named')" \
+            '{query_response: $query_response}')
+
+        echo $EXECUTE | jq '.'
+
+        $CMD tx wasm execute "$CONTRACT" "$EXECUTE" --from admin --chain-id testing -y --gas 2000000
+
+        echo " ... done"
+        echo "------------------------------------"
+        echo "... waiting for event"
+    fi
 done
-
-

apps/transfers/scripts/start.sh

@@ -11,8 +11,11 @@ DIR_QUARTZ_TM_PROVER="$DIR_QUARTZ/utils/tm-prover"
 NODE_URL=${NODE_URL:-127.0.0.1:26657}
 CMD="wasmd --node http://$NODE_URL"
 
+# Use the QUARTZ_PORT environment variable if set, otherwise default to 11090
+QUARTZ_PORT="${QUARTZ_PORT:-11090}"
 
 echo "--------------------------------------------------------"
+echo "QUARTZ_PORT is set to: $QUARTZ_PORT"
 echo "set trusted hash"
 
 cd "$DIR_QUARTZ_TM_PROVER"
@@ -51,6 +54,7 @@ gramine-manifest  \
 -Dquartz_dir="$(pwd)"  \
 -Dtrusted_height="$TRUSTED_HEIGHT"  \
 -Dtrusted_hash="$TRUSTED_HASH"  \
+-Dgramine_port="$QUARTZ_PORT" \
 quartz.manifest.template quartz.manifest
 
 echo "... sign manifest"