Merge branch 'main' into hu55a1n1/v0_2-improvements

Author: hu55a1n1

Cargo.lock

@@ -13,7 +13,7 @@ impl Handler for SessionSetPubKey {
         let (nonce, pub_key) = self.into_tuple();
 
         let session = session
-            .with_pub_key(nonce, pub_key)
+            .with_pub_key(nonce, pub_key.clone())
             .ok_or(Error::BadSessionTransition)?;
         SESSION.save(deps.storage, &session).map_err(Error::Std)?;
 
@@ -24,9 +24,6 @@ impl Handler for SessionSetPubKey {
 
         Ok(Response::new()
             .add_attribute("action", "session_set_pub_key")
-            .add_attribute(
-                "pub_key",
-                HexBinary::from(pub_key.to_sec1_bytes().into_vec()).to_hex(),
-            ))
+            .add_attribute("pub_key", HexBinary::from(pub_key).to_hex()))
     }
 }

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

@@ -1,26 +1,24 @@
 use cosmwasm_schema::cw_serde;
 use cosmwasm_std::{HexBinary, StdError};
-use k256::ecdsa::VerifyingKey;
 use sha2::{Digest, Sha256};
 
 use crate::{
-    error::Error,
     msg::{execute::attested::HasUserData, HasDomainType},
     state::{Nonce, UserData},
 };
 
 #[derive(Clone, Debug, PartialEq)]
 pub struct SessionSetPubKey {
     nonce: Nonce,
-    pub_key: VerifyingKey,
+    pub_key: Vec<u8>,
 }
 
 impl SessionSetPubKey {
-    pub fn new(nonce: Nonce, pub_key: VerifyingKey) -> Self {
+    pub fn new(nonce: Nonce, pub_key: Vec<u8>) -> Self {
         Self { nonce, pub_key }
     }
 
-    pub fn into_tuple(self) -> (Nonce, VerifyingKey) {
+    pub fn into_tuple(self) -> (Nonce, Vec<u8>) {
         (self.nonce, self.pub_key)
     }
 }
@@ -42,18 +40,18 @@ impl TryFrom<RawSessionSetPubKey> for SessionSetPubKey {
 
     fn try_from(value: RawSessionSetPubKey) -> Result<Self, Self::Error> {
         let nonce = value.nonce.to_array()?;
-        let pub_key = VerifyingKey::from_sec1_bytes(&value.pub_key)
-            .map_err(Error::from)
-            .map_err(|e| StdError::generic_err(e.to_string()))?;
-        Ok(Self { nonce, pub_key })
+        Ok(Self {
+            nonce,
+            pub_key: value.pub_key.into(),
+        })
     }
 }
 
 impl From<SessionSetPubKey> for RawSessionSetPubKey {
     fn from(value: SessionSetPubKey) -> Self {
         Self {
             nonce: value.nonce.into(),
-            pub_key: value.pub_key.to_sec1_bytes().into_vec().into(),
+            pub_key: value.pub_key.into(),
         }
     }
 }

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

@@ -1,7 +1,6 @@
 use cosmwasm_schema::cw_serde;
 use cosmwasm_std::{HexBinary, StdError, Uint64};
 use cw_storage_plus::Item;
-use k256::ecdsa::VerifyingKey;
 
 pub type MrEnclave = [u8; 32];
 pub type Nonce = [u8; 32];
@@ -233,9 +232,9 @@ impl Session {
         }
     }
 
-    pub fn with_pub_key(mut self, nonce: Nonce, pub_key: VerifyingKey) -> Option<Self> {
+    pub fn with_pub_key(mut self, nonce: Nonce, pub_key: Vec<u8>) -> Option<Self> {
         if self.nonce == nonce && self.pub_key.is_none() {
-            self.pub_key = Some(pub_key.to_sec1_bytes().into_vec().into());
+            self.pub_key = Some(pub_key.into());
             Some(self)
         } else {
             None

crates/contracts/core/src/state.rs

@@ -24,7 +24,7 @@ async-trait.workspace = true
 displaydoc.workspace = true
 futures-util.workspace = true
 hex.workspace = true
-k256.workspace = true
+k256 = { workspace = true, features = ["pem", "serde"] }
 log.workspace = true
 rand.workspace = true
 reqwest = { workspace = true, features = ["blocking"] }

crates/enclave/core/Cargo.toml

@@ -1,5 +1,4 @@
 use cosmrs::AccountId;
-use k256::ecdsa::VerifyingKey;
 use quartz_contract_core::{
     msg::execute::{attested::Attested, session_set_pub_key::SessionSetPubKey},
     state::{Session, SESSION_KEY},
@@ -24,7 +23,7 @@ use crate::{
 impl<E> Handler<E> for RawSessionSetPubKeyRequest
 where
     E: Enclave,
-    E::KeyManager: KeyManager<PubKey = VerifyingKey>,
+    E::KeyManager: KeyManager,
     E::Store: Store<Contract = AccountId>,
 {
     type Error = Status;
@@ -71,13 +70,10 @@ where
         }
 
         // generate enclave key
-        ctx.key_manager().await.keygen().await;
-        let pk = ctx
-            .key_manager()
-            .await
-            .pub_key()
-            .await
-            .ok_or_else(|| Status::internal("failed to get public key"))?;
+        let pk = {
+            let pk = ctx.key_manager().await.pub_key().await;
+            serde_json::to_vec(&pk).expect("pubkey serialization failure")
+        };
 
         // create `SessionSetPubKey` msg and attest to it
         let msg = SessionSetPubKey::new(nonce, pk);

crates/enclave/core/src/handler/session_set_pubkey.rs

@@ -1,12 +1,11 @@
+use serde::Serialize;
+
 pub mod default;
 pub mod shared;
 
 #[async_trait::async_trait]
 pub trait KeyManager: Send + Sync + 'static {
-    type PubKey;
-    type PrivKey;
+    type PubKey: Serialize;
 
-    async fn keygen(&mut self);
-    async fn pub_key(&self) -> Option<Self::PubKey>;
-    async fn priv_key(&self) -> Option<Self::PrivKey>;
+    async fn pub_key(&self) -> Self::PubKey;
 }

crates/enclave/core/src/key_manager.rs

@@ -2,25 +2,24 @@ use k256::ecdsa::{SigningKey, VerifyingKey};
 
 use crate::key_manager::KeyManager;
 
-#[derive(Clone, Default)]
+#[derive(Clone)]
 pub struct DefaultKeyManager {
-    sk: Option<SigningKey>,
+    pub sk: SigningKey,
+}
+
+impl Default for DefaultKeyManager {
+    fn default() -> Self {
+        Self {
+            sk: SigningKey::random(&mut rand::thread_rng()),
+        }
+    }
 }
 
 #[async_trait::async_trait]
 impl KeyManager for DefaultKeyManager {
     type PubKey = VerifyingKey;
-    type PrivKey = SigningKey;
-
-    async fn keygen(&mut self) {
-        self.sk = Some(SigningKey::random(&mut rand::thread_rng()));
-    }
-
-    async fn pub_key(&self) -> Option<Self::PubKey> {
-        self.sk.clone().map(|sk| VerifyingKey::from(&sk))
-    }
 
-    async fn priv_key(&self) -> Option<Self::PrivKey> {
-        self.sk.clone()
+    async fn pub_key(&self) -> Self::PubKey {
+        self.sk.clone().into()
     }
 }

crates/enclave/core/src/key_manager/default.rs

@@ -1,12 +1,12 @@
 use std::sync::Arc;
 
-use tokio::sync::RwLock;
+use tokio::sync::{RwLock, RwLockReadGuard};
 
 use crate::key_manager::KeyManager;
 
 #[derive(Clone, Debug)]
 pub struct SharedKeyManager<K> {
-    inner: Arc<RwLock<K>>,
+    pub inner: Arc<RwLock<K>>,
 }
 
 impl<K> SharedKeyManager<K> {
@@ -15,22 +15,17 @@ impl<K> SharedKeyManager<K> {
             inner: Arc::new(RwLock::new(key_manager)),
         }
     }
+
+    pub async fn read_lock(&self) -> RwLockReadGuard<'_, K> {
+        self.inner.read().await
+    }
 }
 
 #[async_trait::async_trait]
 impl<K: KeyManager> KeyManager for SharedKeyManager<K> {
     type PubKey = K::PubKey;
-    type PrivKey = K::PrivKey;
-
-    async fn keygen(&mut self) {
-        self.inner.write().await.keygen().await
-    }
 
-    async fn pub_key(&self) -> Option<Self::PubKey> {
+    async fn pub_key(&self) -> Self::PubKey {
         self.inner.read().await.pub_key().await
     }
-
-    async fn priv_key(&self) -> Option<Self::PrivKey> {
-        self.inner.read().await.priv_key().await
-    }
 }

crates/enclave/core/src/key_manager/shared.rs

@@ -54,7 +54,7 @@ pub struct DefaultEnclave<C, A = DefaultAttestor, K = DefaultKeyManager, S = Def
     pub ctx: C,
 }
 
-impl<C> DefaultSharedEnclave<C> {
+impl<C: Send + Sync + 'static> DefaultSharedEnclave<C> {
     pub fn shared(attestor: DefaultAttestor, config: Config, ctx: C) -> DefaultSharedEnclave<C> {
         DefaultSharedEnclave {
             attestor,
@@ -63,6 +63,18 @@ impl<C> DefaultSharedEnclave<C> {
             ctx,
         }
     }
+
+    pub fn with_key_manager<K: KeyManager>(
+        self,
+        key_manager: K,
+    ) -> DefaultEnclave<C, <Self as Enclave>::Attestor, K, <Self as Enclave>::Store> {
+        DefaultEnclave {
+            attestor: self.attestor,
+            key_manager,
+            store: self.store,
+            ctx: self.ctx,
+        }
+    }
 }
 
 #[async_trait::async_trait]

crates/enclave/core/src/lib.rs

@@ -1,8 +1,6 @@
 use cosmwasm_std::{Addr, HexBinary, Uint128};
 use k256::ecdsa::VerifyingKey;
-use quartz_common::enclave::{
-    handler::Handler, key_manager::KeyManager, DefaultSharedEnclave, Enclave,
-};
+use quartz_common::enclave::{handler::Handler, DefaultSharedEnclave};
 use serde::{Deserialize, Serialize};
 use tonic::Status;
 use transfers_contract::msg::execute;
@@ -35,12 +33,7 @@ impl Handler<DefaultSharedEnclave<()>> for QueryRequest {
         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"))?;
+                let sk = ctx.key_manager.read_lock().await.sk.clone();
                 decrypt_state(&sk, state_bytes)?
             }
         };