Skip to content

Commit

Permalink
feat: add sequence number for replay protection (#252)
Browse files Browse the repository at this point in the history
  • Loading branch information
hu55a1n1 authored Nov 7, 2024
1 parent ea08f47 commit bbc8b8a
Show file tree
Hide file tree
Showing 17 changed files with 171 additions and 198 deletions.
112 changes: 12 additions & 100 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions crates/contracts/core/src/handler/execute.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
pub mod attested;
pub mod sequenced;
pub mod session_create;
pub mod session_set_pub_key;

Expand Down
5 changes: 2 additions & 3 deletions crates/contracts/core/src/handler/execute/attested.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,7 @@ use crate::{
error::Error,
handler::Handler,
msg::execute::attested::{
Attestation, Attested, AttestedMsgSansHandler, DcapAttestation, HasUserData,
MockAttestation, Quote,
Attestation, Attested, DcapAttestation, HasUserData, MockAttestation, MsgSansHandler, Quote,
},
state::CONFIG,
};
Expand Down Expand Up @@ -181,7 +180,7 @@ where
}
}

impl<T> Handler for AttestedMsgSansHandler<T> {
impl<T> Handler for MsgSansHandler<T> {
fn handle(
self,
_deps: DepsMut<'_>,
Expand Down
16 changes: 16 additions & 0 deletions crates/contracts/core/src/handler/execute/sequenced.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
use cosmwasm_std::{DepsMut, Env, MessageInfo, Response, StdResult, Uint64};

use crate::{
error::Error, handler::Handler, msg::execute::sequenced::SequencedMsg, state::SEQUENCE_NUM,
};

impl<T: Handler> Handler for SequencedMsg<T> {
fn handle(self, deps: DepsMut<'_>, env: &Env, info: &MessageInfo) -> Result<Response, Error> {
SEQUENCE_NUM.update(deps.storage, |mut counter| -> StdResult<_> {
counter += Uint64::one();
Ok(counter)
})?;

self.0.handle(deps, env, info)
}
}
15 changes: 11 additions & 4 deletions crates/contracts/core/src/handler/execute/session_set_pub_key.rs
Original file line number Diff line number Diff line change
@@ -1,20 +1,27 @@
use cosmwasm_std::{DepsMut, Env, HexBinary, MessageInfo, Response};
use cosmwasm_std::{DepsMut, Env, HexBinary, MessageInfo, Response, Uint64};

use crate::{
error::Error, handler::Handler, msg::execute::session_set_pub_key::SessionSetPubKey,
state::SESSION,
error::Error,
handler::Handler,
msg::execute::session_set_pub_key::SessionSetPubKey,
state::{SEQUENCE_NUM, SESSION},
};

impl Handler for SessionSetPubKey {
fn handle(self, deps: DepsMut<'_>, _env: &Env, _info: &MessageInfo) -> Result<Response, Error> {
let session = SESSION.load(deps.storage).map_err(Error::Std)?;
let (nonce, pub_key) = self.into_tuple();

let session = session
.with_pub_key(nonce, pub_key)
.ok_or(Error::BadSessionTransition)?;

SESSION.save(deps.storage, &session).map_err(Error::Std)?;

let sequence_num = Uint64::new(0);
SEQUENCE_NUM
.save(deps.storage, &sequence_num)
.map_err(Error::Std)?;

Ok(Response::new()
.add_attribute("action", "session_set_pub_key")
.add_attribute(
Expand Down
1 change: 1 addition & 0 deletions crates/contracts/core/src/msg/execute.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
pub mod attested;
pub mod sequenced;
pub mod session_create;
pub mod session_set_pub_key;

Expand Down
18 changes: 9 additions & 9 deletions crates/contracts/core/src/msg/execute/attested.rs
Original file line number Diff line number Diff line change
Expand Up @@ -222,16 +222,16 @@ impl Attestation for MockAttestation {
}

#[derive(Clone, Debug, PartialEq)]
pub struct AttestedMsgSansHandler<T>(pub T);
pub struct MsgSansHandler<T>(pub T);

#[cw_serde]
pub struct RawAttestedMsgSansHandler<T>(pub T);
pub struct RawMsgSansHandler<T>(pub T);

impl<T: Serialize> HasDomainType for RawAttestedMsgSansHandler<T> {
type DomainType = AttestedMsgSansHandler<T>;
impl<T: Serialize> HasDomainType for RawMsgSansHandler<T> {
type DomainType = MsgSansHandler<T>;
}

impl<T> HasUserData for AttestedMsgSansHandler<T>
impl<T> HasUserData for MsgSansHandler<T>
where
T: HasUserData,
{
Expand All @@ -240,16 +240,16 @@ where
}
}

impl<T> TryFrom<RawAttestedMsgSansHandler<T>> for AttestedMsgSansHandler<T> {
impl<T> TryFrom<RawMsgSansHandler<T>> for MsgSansHandler<T> {
type Error = StdError;

fn try_from(value: RawAttestedMsgSansHandler<T>) -> Result<Self, Self::Error> {
fn try_from(value: RawMsgSansHandler<T>) -> Result<Self, Self::Error> {
Ok(Self(value.0))
}
}

impl<T> From<AttestedMsgSansHandler<T>> for RawAttestedMsgSansHandler<T> {
fn from(value: AttestedMsgSansHandler<T>) -> Self {
impl<T> From<MsgSansHandler<T>> for RawMsgSansHandler<T> {
fn from(value: MsgSansHandler<T>) -> Self {
Self(value.0)
}
}
29 changes: 29 additions & 0 deletions crates/contracts/core/src/msg/execute/sequenced.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
use cosmwasm_schema::cw_serde;
use cosmwasm_std::StdError;
use serde::Serialize;

use crate::msg::HasDomainType;

#[derive(Clone, Debug, PartialEq)]
pub struct SequencedMsg<D>(pub D);

#[cw_serde]
pub struct RawSequencedMsg<R>(pub R);

impl<R: Serialize + HasDomainType> HasDomainType for RawSequencedMsg<R> {
type DomainType = SequencedMsg<R::DomainType>;
}

impl<R: HasDomainType> TryFrom<RawSequencedMsg<R>> for SequencedMsg<R::DomainType> {
type Error = StdError;

fn try_from(value: RawSequencedMsg<R>) -> Result<Self, Self::Error> {
Ok(Self(value.0.try_into()?))
}
}

impl<R: HasDomainType> From<SequencedMsg<R::DomainType>> for RawSequencedMsg<R> {
fn from(value: SequencedMsg<R::DomainType>) -> Self {
Self(value.0.into())
}
}
2 changes: 2 additions & 0 deletions crates/contracts/core/src/state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,9 +15,11 @@ 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 SEQUENCE_NUM_KEY: &str = "quartz_seq_num";
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);
pub const SEQUENCE_NUM: Item<Uint64> = Item::new(SEQUENCE_NUM_KEY);

#[derive(Clone, Debug, PartialEq)]
pub struct Config {
Expand Down
4 changes: 0 additions & 4 deletions crates/utils/cw-client/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,3 @@ tonic.workspace = true
cosmrs = { workspace = true, default-features = false, features = ["cosmwasm"] }
cosmos-sdk-proto = { workspace = true, default-features = false, features = ["grpc", "grpc-transport"] }
tendermint = { workspace = true, default-features = false }

[dev-dependencies]
tokio.workspace = true
transfers-contract = { path = "../../../examples/transfers/contracts" }
2 changes: 1 addition & 1 deletion crates/utils/cw-client/src/cli.rs
Original file line number Diff line number Diff line change
Expand Up @@ -106,7 +106,7 @@ impl CwClient for CliClient {
Ok(query_result)
}

fn query_raw<R: DeserializeOwned + Default>(
async fn query_raw<R: DeserializeOwned + Default>(
&self,
contract: &Self::Address,
query: Self::RawQuery,
Expand Down
Loading

0 comments on commit bbc8b8a

Please sign in to comment.