Skip to content

Commit bc24db6

Browse files
hu55a1n1hu55a1n1
authored
fix: add check for matching proof key (#251)
1 parent d37355e commit bc24db6

File tree

19 files changed

+146
-67
lines changed

19 files changed

+146
-67
lines changed

‎crates/cli/src/handler/handshake.rs

+5-3
Original file line numberDiff line numberDiff line change
@@ -42,9 +42,11 @@ async fn handshake(args: HandshakeRequest, config: Config) -> Result<String> {
4242

4343
info!("Running SessionCreate");
4444

45-
let res: serde_json::Value = RelayMessage::SessionCreate
46-
.run_relay(config.enclave_rpc())
47-
.await?;
45+
let res: serde_json::Value = RelayMessage::SessionCreate {
46+
contract: args.contract.clone(),
47+
}
48+
.run_relay(config.enclave_rpc())
49+
.await?;
4850

4951
let output: WasmdTxResponse = serde_json::from_str(
5052
cw_client

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

+6-3
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
use color_eyre::{eyre::eyre, Result};
2+
use cosmrs::AccountId;
23
use quartz_common::proto::{
34
core_client::CoreClient, InstantiateRequest, SessionCreateRequest, SessionSetPubKeyRequest,
45
};
@@ -8,7 +9,7 @@ use serde_json::{json, Value as JsonValue};
89
#[derive(Debug)]
910
pub enum RelayMessage {
1011
Instantiate { init_msg: JsonValue },
11-
SessionCreate,
12+
SessionCreate { contract: AccountId },
1213
SessionSetPubKey { proof: ProofOutput },
1314
}
1415

@@ -37,8 +38,10 @@ impl RelayMessage {
3738
init_msg["quartz"] = msg;
3839
init_msg.to_string()
3940
})?,
40-
RelayMessage::SessionCreate => qc_client
41-
.session_create(tonic::Request::new(SessionCreateRequest {}))
41+
RelayMessage::SessionCreate { contract } => qc_client
42+
.session_create(tonic::Request::new(SessionCreateRequest {
43+
message: serde_json::to_string(&contract)?,
44+
}))
4245
.await
4346
.map_err(|e| {
4447
eyre!(

‎crates/contracts/core/src/error.rs

+2
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@ pub enum Error {
1919
TcbInfoQueryError(String),
2020
#[error("DCAP verification query error: {0}")]
2121
DcapVerificationQueryError(String),
22+
#[error("contract address mismatch")]
23+
ContractAddrMismatch,
2224
}
2325

2426
impl From<K256Error> for Error {

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

+8-2
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,16 @@ use crate::{
88
};
99

1010
impl Handler for SessionCreate {
11-
fn handle(self, deps: DepsMut<'_>, _env: &Env, _info: &MessageInfo) -> Result<Response, Error> {
11+
fn handle(self, deps: DepsMut<'_>, env: &Env, _info: &MessageInfo) -> Result<Response, Error> {
1212
// TODO(hu55a1n1): overwrite previous session?
13+
14+
let addr = deps.api.addr_validate(self.contract())?;
15+
if addr != env.contract.address {
16+
return Err(Error::ContractAddrMismatch);
17+
}
18+
1319
SESSION
14-
.save(deps.storage, &Session::create(self.into_nonce()))
20+
.save(deps.storage, &Session::create(self.nonce()))
1521
.map_err(Error::Std)?;
1622

1723
Ok(Response::new().add_attribute("action", "session_create"))

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

+12-4
Original file line numberDiff line numberDiff line change
@@ -10,36 +10,44 @@ use crate::{
1010
#[derive(Clone, Debug, PartialEq)]
1111
pub struct SessionCreate {
1212
nonce: Nonce,
13+
contract: String,
1314
}
1415

1516
impl SessionCreate {
16-
pub fn new(nonce: Nonce) -> Self {
17-
Self { nonce }
17+
pub fn new(nonce: Nonce, contract: String) -> Self {
18+
Self { nonce, contract }
1819
}
1920

20-
pub fn into_nonce(self) -> Nonce {
21+
pub fn nonce(&self) -> Nonce {
2122
self.nonce
2223
}
24+
25+
pub fn contract(&self) -> &str {
26+
self.contract.as_str()
27+
}
2328
}
2429

2530
#[cw_serde]
2631
pub struct RawSessionCreate {
2732
nonce: HexBinary,
33+
contract: String,
2834
}
2935

3036
impl TryFrom<RawSessionCreate> for SessionCreate {
3137
type Error = StdError;
3238

3339
fn try_from(value: RawSessionCreate) -> Result<Self, Self::Error> {
3440
let nonce = value.nonce.to_array()?;
35-
Ok(Self { nonce })
41+
let contract = value.contract;
42+
Ok(Self { nonce, contract })
3643
}
3744
}
3845

3946
impl From<SessionCreate> for RawSessionCreate {
4047
fn from(value: SessionCreate) -> Self {
4148
Self {
4249
nonce: value.nonce.into(),
50+
contract: value.contract,
4351
}
4452
}
4553
}

‎crates/contracts/core/src/state.rs

+7-4
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,13 @@ pub type Hash = [u8; 32];
1212
pub type Height = u64;
1313
pub type TrustThreshold = (u64, u64);
1414

15+
pub const CONFIG_KEY: &str = "quartz_config";
16+
pub const SESSION_KEY: &str = "quartz_session";
17+
pub const EPOCH_COUNTER_KEY: &str = "epoch_counter";
18+
pub const CONFIG: Item<RawConfig> = Item::new(CONFIG_KEY);
19+
pub const SESSION: Item<Session> = Item::new(SESSION_KEY);
20+
pub const EPOCH_COUNTER: Item<Uint64> = Item::new(EPOCH_COUNTER_KEY);
21+
1522
#[derive(Clone, Debug, PartialEq)]
1623
pub struct Config {
1724
mr_enclave: MrEnclave,
@@ -247,7 +254,3 @@ impl Session {
247254
self.nonce.to_array().expect("correct by construction")
248255
}
249256
}
250-
251-
pub const CONFIG: Item<RawConfig> = Item::new("quartz_config");
252-
pub const SESSION: Item<Session> = Item::new("quartz_session");
253-
pub const EPOCH_COUNTER: Item<Uint64> = Item::new("epoch_counter");

‎crates/enclave/core/src/server.rs

+47-7
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ use std::{
55
time::Duration,
66
};
77

8+
use cosmrs::AccountId;
89
use futures_util::StreamExt;
910
use k256::ecdsa::SigningKey;
1011
use quartz_contract_core::{
@@ -15,10 +16,11 @@ use quartz_contract_core::{
1516
},
1617
instantiate::CoreInstantiate,
1718
},
18-
state::{Config, LightClientOpts, Nonce, Session},
19+
state::{Config, LightClientOpts, Nonce, Session, SESSION_KEY},
1920
};
2021
use quartz_cw_proof::proof::{
2122
cw::{CwProof, RawCwProof},
23+
key::CwAbciKey,
2224
Proof,
2325
};
2426
use quartz_proto::quartz::{
@@ -106,13 +108,19 @@ impl QuartzServer {
106108
pub fn new<A>(
107109
config: Config,
108110
sk: Arc<Mutex<Option<SigningKey>>>,
111+
contract: Arc<Mutex<Option<AccountId>>>,
109112
attestor: A,
110113
ws_config: WsListenerConfig,
111114
) -> Self
112115
where
113116
A: Attestor + Clone,
114117
{
115-
let core_service = CoreServer::new(CoreService::new(config, sk.clone(), attestor.clone()));
118+
let core_service = CoreServer::new(CoreService::new(
119+
config,
120+
contract.clone(),
121+
sk.clone(),
122+
attestor.clone(),
123+
));
116124

117125
Self {
118126
router: Server::builder().add_service(core_service),
@@ -184,6 +192,7 @@ impl QuartzServer {
184192
pub struct CoreService<A> {
185193
config: Config,
186194
nonce: Arc<Mutex<Nonce>>,
195+
contract: Arc<Mutex<Option<AccountId>>>,
187196
sk: Arc<Mutex<Option<SigningKey>>>,
188197
attestor: A,
189198
}
@@ -192,10 +201,16 @@ impl<A> CoreService<A>
192201
where
193202
A: Attestor,
194203
{
195-
pub fn new(config: Config, sk: Arc<Mutex<Option<SigningKey>>>, attestor: A) -> Self {
204+
pub fn new(
205+
config: Config,
206+
contract: Arc<Mutex<Option<AccountId>>>,
207+
sk: Arc<Mutex<Option<SigningKey>>>,
208+
attestor: A,
209+
) -> Self {
196210
Self {
197211
config,
198212
nonce: Arc::new(Mutex::new([0u8; 32])),
213+
contract,
199214
sk,
200215
attestor,
201216
}
@@ -226,12 +241,19 @@ where
226241

227242
async fn session_create(
228243
&self,
229-
_request: Request<RawSessionCreateRequest>,
244+
request: Request<RawSessionCreateRequest>,
230245
) -> TonicResult<Response<RawSessionCreateResponse>> {
231246
// FIXME(hu55a1n1) - disallow calling more than once
247+
let deployed_contract: AccountId = serde_json::from_str(&request.into_inner().message)
248+
.map_err(|e| Status::invalid_argument(e.to_string()))?;
249+
250+
let mut contract = self.contract.lock().unwrap();
251+
*contract = Some(deployed_contract.clone());
252+
232253
let mut nonce = self.nonce.lock().unwrap();
233254
*nonce = rand::thread_rng().gen::<Nonce>();
234-
let msg = SessionCreate::new(*nonce);
255+
256+
let msg = SessionCreate::new(*nonce, deployed_contract.to_string());
235257

236258
let attestation = self
237259
.attestor
@@ -253,8 +275,15 @@ where
253275
serde_json::from_str(&request.into_inner().message)
254276
.map_err(|e| Status::invalid_argument(e.to_string()))?;
255277

278+
let contract = self.contract.lock().unwrap().clone();
279+
256280
let (value, _msg) = proof
257-
.verify(self.config.light_client_opts())
281+
.verify(
282+
self.config.light_client_opts(),
283+
contract.expect("contract not set"),
284+
SESSION_KEY.to_string(),
285+
None,
286+
)
258287
.map_err(Status::failed_precondition)?;
259288

260289
let session: Session = serde_json::from_slice(&value).unwrap();
@@ -290,7 +319,13 @@ pub struct ProofOfPublication<M> {
290319
}
291320

292321
impl<M> ProofOfPublication<M> {
293-
pub fn verify(self, light_client_opts: &LightClientOpts) -> Result<(Vec<u8>, M), String> {
322+
pub fn verify(
323+
self,
324+
light_client_opts: &LightClientOpts,
325+
contract_address: AccountId,
326+
storage_key: String,
327+
storage_namespace: Option<String>,
328+
) -> Result<(Vec<u8>, M), String> {
294329
let config_trust_threshold = light_client_opts.trust_threshold();
295330
let trust_threshold =
296331
TrustThreshold::new(config_trust_threshold.0, config_trust_threshold.1).unwrap();
@@ -322,6 +357,11 @@ impl<M> ProofOfPublication<M> {
322357
.and_then(|mut primary| primary.verify_to_height(target_height))
323358
.map_err(|e| e.to_string())?;
324359

360+
let key = CwAbciKey::new(contract_address, storage_key, storage_namespace);
361+
if key.into_vec() != self.merkle_proof.key() {
362+
return Err("Merkle proof key mismatch".to_string());
363+
}
364+
325365
let proof = CwProof::from(self.merkle_proof);
326366
proof
327367
.verify(

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

+6
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,12 @@ pub struct RawCwProof {
7979
proof: ProofOps,
8080
}
8181

82+
impl RawCwProof {
83+
pub fn key(&self) -> &[u8] {
84+
self.key.as_ref()
85+
}
86+
}
87+
8288
impl From<RawCwProof> for CwProof {
8389
fn from(RawCwProof { key, value, proof }: RawCwProof) -> Self {
8490
Self {

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

+4
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,10 @@ impl CwAbciKey {
8282
}
8383
}
8484

85+
pub fn into_vec(self) -> Vec<u8> {
86+
self.into()
87+
}
88+
8589
fn into_tuple(self) -> (AccountId, String, Option<String>) {
8690
match self {
8791
CwAbciKey::Item {

‎crates/enclave/proto/proto/quartz.proto

+3-1
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,9 @@ message InstantiateResponse {
1414
string message = 1;
1515
}
1616

17-
message SessionCreateRequest {}
17+
message SessionCreateRequest {
18+
string message = 1;
19+
}
1820

1921
message SessionCreateResponse {
2022
string message = 1;

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

+5-2
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,11 @@ pub struct InstantiateResponse {
66
#[prost(string, tag = "1")]
77
pub message: ::prost::alloc::string::String,
88
}
9-
#[derive(Clone, Copy, PartialEq, ::prost::Message)]
10-
pub struct SessionCreateRequest {}
9+
#[derive(Clone, PartialEq, ::prost::Message)]
10+
pub struct SessionCreateRequest {
11+
#[prost(string, tag = "1")]
12+
pub message: ::prost::alloc::string::String,
13+
}
1114
#[derive(Clone, PartialEq, ::prost::Message)]
1215
pub struct SessionCreateResponse {
1316
#[prost(string, tag = "1")]

‎examples/transfers/contracts/Cargo.lock

-10
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎examples/transfers/contracts/Cargo.toml

+3
Original file line numberDiff line numberDiff line change
@@ -57,3 +57,6 @@ getrandom = { version = "0.2.15", features = ["js"] }
5757
[dev-dependencies]
5858
cw-multi-test = { version = "2.1.0", default-features = false }
5959
serde_json = "1.0.122"
60+
61+
[patch.crates-io]
62+
quartz-common = { path = "../../../crates/common" }

‎examples/transfers/contracts/src/state.rs

+2-1
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,8 @@ use cw_storage_plus::{Item, Map};
33

44
use crate::msg::execute::Request;
55

6+
pub const REQUESTS_KEY: &str = "requests";
67
pub const STATE: Item<HexBinary> = Item::new("state");
7-
pub const REQUESTS: Item<Vec<Request>> = Item::new("requests");
8+
pub const REQUESTS: Item<Vec<Request>> = Item::new(REQUESTS_KEY);
89
pub const DENOM: Item<String> = Item::new("donation_denom");
910
pub const BALANCES: Map<&str, HexBinary> = Map::new("balances");

0 commit comments

Comments
 (0)