fix: add check for matching proof key

crates/cli/src/handler/handshake.rs

@@ -42,9 +42,11 @@ async fn handshake(args: HandshakeRequest, config: Config) -> Result<String> {
 
     info!("Running SessionCreate");
 
-    let res: serde_json::Value = RelayMessage::SessionCreate
-        .run_relay(config.enclave_rpc())
-        .await?;
+    let res: serde_json::Value = RelayMessage::SessionCreate {
+        contract: args.contract.clone(),
+    }
+    .run_relay(config.enclave_rpc())
+    .await?;
 
     let output: WasmdTxResponse = serde_json::from_str(
         cw_client

crates/cli/src/handler/utils/relay.rs

@@ -1,4 +1,5 @@
 use color_eyre::{eyre::eyre, Result};
+use cosmrs::AccountId;
 use quartz_common::proto::{
     core_client::CoreClient, InstantiateRequest, SessionCreateRequest, SessionSetPubKeyRequest,
 };
@@ -8,7 +9,7 @@ use serde_json::{json, Value as JsonValue};
 #[derive(Debug)]
 pub enum RelayMessage {
     Instantiate { init_msg: JsonValue },
-    SessionCreate,
+    SessionCreate { contract: AccountId },
     SessionSetPubKey { proof: ProofOutput },
 }
 
@@ -37,8 +38,10 @@ impl RelayMessage {
                     init_msg["quartz"] = msg;
                     init_msg.to_string()
                 })?,
-            RelayMessage::SessionCreate => qc_client
-                .session_create(tonic::Request::new(SessionCreateRequest {}))
+            RelayMessage::SessionCreate { contract } => qc_client
+                .session_create(tonic::Request::new(SessionCreateRequest {
+                    message: serde_json::to_string(&contract)?,
+                }))
                 .await
                 .map_err(|e| {
                     eyre!(

crates/contracts/core/src/error.rs

@@ -19,6 +19,8 @@ pub enum Error {
     TcbInfoQueryError(String),
     #[error("DCAP verification query error: {0}")]
     DcapVerificationQueryError(String),
+    #[error("contract address mismatch")]
+    ContractAddrMismatch,
 }
 
 impl From<K256Error> for Error {

crates/contracts/core/src/handler/execute/session_create.rs

@@ -8,10 +8,16 @@ use crate::{
 };
 
 impl Handler for SessionCreate {
-    fn handle(self, deps: DepsMut<'_>, _env: &Env, _info: &MessageInfo) -> Result<Response, Error> {
+    fn handle(self, deps: DepsMut<'_>, env: &Env, _info: &MessageInfo) -> Result<Response, Error> {
         // TODO(hu55a1n1): overwrite previous session?
+
+        let addr = deps.api.addr_validate(self.contract())?;
+        if addr != env.contract.address {
+            return Err(Error::ContractAddrMismatch);
+        }
+
         SESSION
-            .save(deps.storage, &Session::create(self.into_nonce()))
+            .save(deps.storage, &Session::create(self.nonce()))
             .map_err(Error::Std)?;
 
         Ok(Response::new().add_attribute("action", "session_create"))

crates/contracts/core/src/msg/execute/session_create.rs

@@ -10,36 +10,44 @@ use crate::{
 #[derive(Clone, Debug, PartialEq)]
 pub struct SessionCreate {
     nonce: Nonce,
+    contract: String,
 }
 
 impl SessionCreate {
-    pub fn new(nonce: Nonce) -> Self {
-        Self { nonce }
+    pub fn new(nonce: Nonce, contract: String) -> Self {
+        Self { nonce, contract }
     }
 
-    pub fn into_nonce(self) -> Nonce {
+    pub fn nonce(&self) -> Nonce {
         self.nonce
     }
+
+    pub fn contract(&self) -> &str {
+        self.contract.as_str()
+    }
 }
 
 #[cw_serde]
 pub struct RawSessionCreate {
     nonce: HexBinary,
+    contract: String,
 }
 
 impl TryFrom<RawSessionCreate> for SessionCreate {
     type Error = StdError;
 
     fn try_from(value: RawSessionCreate) -> Result<Self, Self::Error> {
         let nonce = value.nonce.to_array()?;
-        Ok(Self { nonce })
+        let contract = value.contract;
+        Ok(Self { nonce, contract })
     }
 }
 
 impl From<SessionCreate> for RawSessionCreate {
     fn from(value: SessionCreate) -> Self {
         Self {
             nonce: value.nonce.into(),
+            contract: value.contract,
         }
     }
 }

crates/contracts/core/src/state.rs

@@ -12,6 +12,13 @@ pub type Hash = [u8; 32];
 pub type Height = u64;
 pub type TrustThreshold = (u64, u64);
 
+pub const CONFIG_KEY: &str = "quartz_config";
+pub const SESSION_KEY: &str = "quartz_session";
+pub const EPOCH_COUNTER_KEY: &str = "epoch_counter";
+pub const CONFIG: Item<RawConfig> = Item::new(CONFIG_KEY);
+pub const SESSION: Item<Session> = Item::new(SESSION_KEY);
+pub const EPOCH_COUNTER: Item<Uint64> = Item::new(EPOCH_COUNTER_KEY);
+
 #[derive(Clone, Debug, PartialEq)]
 pub struct Config {
     mr_enclave: MrEnclave,
@@ -247,7 +254,3 @@ impl Session {
         self.nonce.to_array().expect("correct by construction")
     }
 }
-
-pub const CONFIG: Item<RawConfig> = Item::new("quartz_config");
-pub const SESSION: Item<Session> = Item::new("quartz_session");
-pub const EPOCH_COUNTER: Item<Uint64> = Item::new("epoch_counter");

crates/enclave/core/src/server.rs

@@ -5,6 +5,7 @@ use std::{
     time::Duration,
 };
 
+use cosmrs::AccountId;
 use futures_util::StreamExt;
 use k256::ecdsa::SigningKey;
 use quartz_contract_core::{
@@ -15,10 +16,11 @@ use quartz_contract_core::{
         },
         instantiate::CoreInstantiate,
     },
-    state::{Config, LightClientOpts, Nonce, Session},
+    state::{Config, LightClientOpts, Nonce, Session, SESSION_KEY},
 };
 use quartz_cw_proof::proof::{
     cw::{CwProof, RawCwProof},
+    key::CwAbciKey,
     Proof,
 };
 use quartz_proto::quartz::{
@@ -106,13 +108,19 @@ impl QuartzServer {
     pub fn new<A>(
         config: Config,
         sk: Arc<Mutex<Option<SigningKey>>>,
+        contract: Arc<Mutex<Option<AccountId>>>,
         attestor: A,
         ws_config: WsListenerConfig,
     ) -> Self
     where
         A: Attestor + Clone,
     {
-        let core_service = CoreServer::new(CoreService::new(config, sk.clone(), attestor.clone()));
+        let core_service = CoreServer::new(CoreService::new(
+            config,
+            contract.clone(),
+            sk.clone(),
+            attestor.clone(),
+        ));
 
         Self {
             router: Server::builder().add_service(core_service),
@@ -184,6 +192,7 @@ impl QuartzServer {
 pub struct CoreService<A> {
     config: Config,
     nonce: Arc<Mutex<Nonce>>,
+    contract: Arc<Mutex<Option<AccountId>>>,
     sk: Arc<Mutex<Option<SigningKey>>>,
     attestor: A,
 }
@@ -192,10 +201,16 @@ impl<A> CoreService<A>
 where
     A: Attestor,
 {
-    pub fn new(config: Config, sk: Arc<Mutex<Option<SigningKey>>>, attestor: A) -> Self {
+    pub fn new(
+        config: Config,
+        contract: Arc<Mutex<Option<AccountId>>>,
+        sk: Arc<Mutex<Option<SigningKey>>>,
+        attestor: A,
+    ) -> Self {
         Self {
             config,
             nonce: Arc::new(Mutex::new([0u8; 32])),
+            contract,
             sk,
             attestor,
         }
@@ -226,12 +241,19 @@ where
 
     async fn session_create(
         &self,
-        _request: Request<RawSessionCreateRequest>,
+        request: Request<RawSessionCreateRequest>,
     ) -> TonicResult<Response<RawSessionCreateResponse>> {
         // FIXME(hu55a1n1) - disallow calling more than once
+        let deployed_contract: AccountId = serde_json::from_str(&request.into_inner().message)
+            .map_err(|e| Status::invalid_argument(e.to_string()))?;
+
+        let mut contract = self.contract.lock().unwrap();
+        *contract = Some(deployed_contract.clone());
+
         let mut nonce = self.nonce.lock().unwrap();
         *nonce = rand::thread_rng().gen::<Nonce>();
-        let msg = SessionCreate::new(*nonce);
+
+        let msg = SessionCreate::new(*nonce, deployed_contract.to_string());
 
         let attestation = self
             .attestor
@@ -253,8 +275,15 @@ where
             serde_json::from_str(&request.into_inner().message)
                 .map_err(|e| Status::invalid_argument(e.to_string()))?;
 
+        let contract = self.contract.lock().unwrap().clone();
+
         let (value, _msg) = proof
-            .verify(self.config.light_client_opts())
+            .verify(
+                self.config.light_client_opts(),
+                contract.expect("contract not set"),
+                SESSION_KEY.to_string(),
+                None,
+            )
             .map_err(Status::failed_precondition)?;
 
         let session: Session = serde_json::from_slice(&value).unwrap();
@@ -290,7 +319,13 @@ pub struct ProofOfPublication<M> {
 }
 
 impl<M> ProofOfPublication<M> {
-    pub fn verify(self, light_client_opts: &LightClientOpts) -> Result<(Vec<u8>, M), String> {
+    pub fn verify(
+        self,
+        light_client_opts: &LightClientOpts,
+        contract_address: AccountId,
+        storage_key: String,
+        storage_namespace: Option<String>,
+    ) -> Result<(Vec<u8>, M), String> {
         let config_trust_threshold = light_client_opts.trust_threshold();
         let trust_threshold =
             TrustThreshold::new(config_trust_threshold.0, config_trust_threshold.1).unwrap();
@@ -322,6 +357,11 @@ impl<M> ProofOfPublication<M> {
         .and_then(|mut primary| primary.verify_to_height(target_height))
         .map_err(|e| e.to_string())?;
 
+        let key = CwAbciKey::new(contract_address, storage_key, storage_namespace);
+        if key.into_vec() != self.merkle_proof.key() {
+            return Err("Merkle proof key mismatch".to_string());
+        }
+
         let proof = CwProof::from(self.merkle_proof);
         proof
             .verify(

crates/enclave/cw-proof/src/proof/cw.rs

@@ -79,6 +79,12 @@ pub struct RawCwProof {
     proof: ProofOps,
 }
 
+impl RawCwProof {
+    pub fn key(&self) -> &[u8] {
+        self.key.as_ref()
+    }
+}
+
 impl From<RawCwProof> for CwProof {
     fn from(RawCwProof { key, value, proof }: RawCwProof) -> Self {
         Self {

crates/enclave/cw-proof/src/proof/key.rs

@@ -82,6 +82,10 @@ impl CwAbciKey {
         }
     }
 
+    pub fn into_vec(self) -> Vec<u8> {
+        self.into()
+    }
+
     fn into_tuple(self) -> (AccountId, String, Option<String>) {
         match self {
             CwAbciKey::Item {

crates/enclave/proto/proto/quartz.proto

@@ -14,7 +14,9 @@ message InstantiateResponse {
   string message = 1;
 }
 
-message SessionCreateRequest {}
+message SessionCreateRequest {
+  string message = 1;
+}
 
 message SessionCreateResponse {
   string message = 1;

crates/enclave/proto/src/prost/quartz.rs

@@ -6,8 +6,11 @@ pub struct InstantiateResponse {
     #[prost(string, tag = "1")]
     pub message: ::prost::alloc::string::String,
 }
-#[derive(Clone, Copy, PartialEq, ::prost::Message)]
-pub struct SessionCreateRequest {}
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct SessionCreateRequest {
+    #[prost(string, tag = "1")]
+    pub message: ::prost::alloc::string::String,
+}
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct SessionCreateResponse {
     #[prost(string, tag = "1")]

examples/transfers/contracts/src/state.rs

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

examples/transfers/enclave/src/main.rs

@@ -86,8 +86,6 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         admin_sk,
     };
 
-    let sk = Arc::new(Mutex::new(None));
-
     // Event queue
     let (tx, mut rx) = mpsc::channel::<TransfersOp<DefaultAttestor>>(1);
     // Consumer task: dequeue and process events
@@ -99,10 +97,19 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         }
     });
 
-    QuartzServer::new(config.clone(), sk.clone(), attestor.clone(), ws_config)
-        .add_service(TransfersService::new(config, sk, attestor, tx))
-        .serve(args.rpc_addr)
-        .await?;
+    let contract = Arc::new(Mutex::new(None));
+    let sk = Arc::new(Mutex::new(None));
+
+    QuartzServer::new(
+        config.clone(),
+        contract.clone(),
+        sk.clone(),
+        attestor.clone(),
+        ws_config,
+    )
+    .add_service(TransfersService::new(config, sk, contract, attestor, tx))
+    .serve(args.rpc_addr)
+    .await?;
 
     Ok(())
 }