Impl SessionSetPubKey handling

Cargo.lock

@@ -2358,6 +2358,7 @@ dependencies = [
  "clap",
  "color-eyre",
  "cosmwasm-std",
+ "k256 0.13.3",
  "prost 0.12.3",
  "quartz-cw",
  "quartz-proto",
@@ -2391,6 +2392,7 @@ dependencies = [
  "cosmwasm-std",
  "displaydoc",
  "ecies",
+ "k256 0.13.3",
  "quartz-cw",
  "quartz-proto",
  "quartz-tee-ra",

enclaves/quartz/Cargo.toml

@@ -7,6 +7,7 @@ edition = "2021"
 clap = { version = "4.1.8", features = ["derive"] }
 color-eyre = "0.6.2"
 cosmwasm-std = "1.4.0"
+k256 = { version = "0.13.2", default-features = false, features = ["ecdsa", "alloc"] }
 prost = "0.12"
 rand = "0.8.5"
 serde = { version = "1.0.189", features = ["derive"] }

enclaves/quartz/src/main.rs

@@ -4,7 +4,6 @@
     clippy::checked_conversions,
     clippy::panic,
     clippy::panic_in_result_fn,
-    clippy::unwrap_used,
     missing_docs,
     trivial_casts,
     trivial_numeric_casts,

enclaves/quartz/src/server.rs

@@ -1,13 +1,21 @@
+use std::sync::{Arc, Mutex};
+
+use k256::ecdsa::SigningKey;
 use quartz_cw::{
-    msg::{execute::session_create::SessionCreate, instantiate::CoreInstantiate},
+    msg::{
+        execute::{session_create::SessionCreate, session_set_pub_key::SessionSetPubKey},
+        instantiate::CoreInstantiate,
+    },
     state::{Config, Nonce},
 };
 use quartz_proto::quartz::{
     core_server::Core, InstantiateRequest as RawInstantiateRequest,
     InstantiateResponse as RawInstantiateResponse, SessionCreateRequest as RawSessionCreateRequest,
     SessionCreateResponse as RawSessionCreateResponse,
+    SessionSetPubKeyRequest as RawSessionSetPubKeyRequest,
+    SessionSetPubKeyResponse as RawSessionSetPubKeyResponse,
 };
-use quartz_relayer::types::{InstantiateResponse, SessionCreateResponse};
+use quartz_relayer::types::{InstantiateResponse, SessionCreateResponse, SessionSetPubKeyResponse};
 use rand::Rng;
 use tonic::{Request, Response, Status};
 
@@ -15,9 +23,10 @@ use crate::attestor::Attestor;
 
 type TonicResult<T> = Result<T, Status>;
 
-#[derive(Clone, PartialEq, Debug)]
+#[derive(Clone, Debug)]
 pub struct CoreService<A> {
     config: Config,
+    nonce: Arc<Mutex<Nonce>>,
     attestor: A,
 }
 
@@ -26,7 +35,11 @@ where
     A: Attestor,
 {
     pub fn new(config: Config, attestor: A) -> Self {
-        Self { config, attestor }
+        Self {
+            config,
+            nonce: Arc::new(Mutex::new([0u8; 32])),
+            attestor,
+        }
     }
 }
 
@@ -53,15 +66,36 @@ where
         &self,
         _request: Request<RawSessionCreateRequest>,
     ) -> TonicResult<Response<RawSessionCreateResponse>> {
-        let nonce = rand::thread_rng().gen::<Nonce>();
+        let mut nonce = self.nonce.lock().unwrap();
-        let session_create_msg = SessionCreate::new(nonce);
+        *nonce = rand::thread_rng().gen::<Nonce>();
+
+        let session_create_msg = SessionCreate::new(*nonce);
 
         let quote = self
             .attestor
             .quote(session_create_msg)
             .map_err(|e| Status::internal(e.to_string()))?;
 
-        let response = SessionCreateResponse::new(nonce, quote);
+        let response = SessionCreateResponse::new(*nonce, quote);
+        Ok(Response::new(response.into()))
+    }
+
+    async fn session_set_pub_key(
+        &self,
+        _request: Request<RawSessionSetPubKeyRequest>,
+    ) -> TonicResult<Response<RawSessionSetPubKeyResponse>> {
+        let nonce = self.nonce.lock().unwrap();
+        let sk = SigningKey::random(&mut rand::thread_rng());
+        let pk = sk.verifying_key();
+
+        let session_set_pub_key_msg = SessionSetPubKey::new(*nonce, *pk);
+
+        let quote = self
+            .attestor
+            .quote(session_set_pub_key_msg)
+            .map_err(|e| Status::internal(e.to_string()))?;
+
+        let response = SessionSetPubKeyResponse::new(*nonce, *pk, quote);
         Ok(Response::new(response.into()))
     }
 }

utils/quartz-proto/proto/quartz.proto

@@ -5,6 +5,7 @@ package quartz;
 service Core {
   rpc Instantiate (InstantiateRequest) returns (InstantiateResponse) {}
   rpc SessionCreate (SessionCreateRequest) returns (SessionCreateResponse) {}
+  rpc SessionSetPubKey (SessionSetPubKeyRequest) returns (SessionSetPubKeyResponse) {}
 }
 
 message InstantiateRequest {}
@@ -18,3 +19,9 @@ message SessionCreateRequest {}
 message SessionCreateResponse {
   string message = 1;
 }
+
+message SessionSetPubKeyRequest {}
+
+message SessionSetPubKeyResponse {
+  string message = 1;
+}

utils/quartz-proto/src/prost/quartz.rs

@@ -16,6 +16,15 @@ pub struct SessionCreateResponse {
     #[prost(string, tag = "1")]
     pub message: ::prost::alloc::string::String,
 }
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct SessionSetPubKeyRequest {}
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct SessionSetPubKeyResponse {
+    #[prost(string, tag = "1")]
+    pub message: ::prost::alloc::string::String,
+}
 /// Generated client implementations.
 pub mod core_client {
     #![allow(unused_variables, dead_code, missing_docs, clippy::let_unit_value)]
@@ -147,6 +156,31 @@ pub mod core_client {
             req.extensions_mut().insert(GrpcMethod::new("quartz.Core", "SessionCreate"));
             self.inner.unary(req, path, codec).await
         }
+        pub async fn session_set_pub_key(
+            &mut self,
+            request: impl tonic::IntoRequest<super::SessionSetPubKeyRequest>,
+        ) -> std::result::Result<
+            tonic::Response<super::SessionSetPubKeyResponse>,
+            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(
+                "/quartz.Core/SessionSetPubKey",
+            );
+            let mut req = request.into_request();
+            req.extensions_mut()
+                .insert(GrpcMethod::new("quartz.Core", "SessionSetPubKey"));
+            self.inner.unary(req, path, codec).await
+        }
     }
 }
 /// Generated server implementations.
@@ -170,6 +204,13 @@ pub mod core_server {
             tonic::Response<super::SessionCreateResponse>,
             tonic::Status,
         >;
+        async fn session_set_pub_key(
+            &self,
+            request: tonic::Request<super::SessionSetPubKeyRequest>,
+        ) -> std::result::Result<
+            tonic::Response<super::SessionSetPubKeyResponse>,
+            tonic::Status,
+        >;
     }
     #[derive(Debug)]
     pub struct CoreServer<T: Core> {
@@ -340,6 +381,52 @@ pub mod core_server {
                     };
                     Box::pin(fut)
                 }
+                "/quartz.Core/SessionSetPubKey" => {
+                    #[allow(non_camel_case_types)]
+                    struct SessionSetPubKeySvc<T: Core>(pub Arc<T>);
+                    impl<
+                        T: Core,
+                    > tonic::server::UnaryService<super::SessionSetPubKeyRequest>
+                    for SessionSetPubKeySvc<T> {
+                        type Response = super::SessionSetPubKeyResponse;
+                        type Future = BoxFuture<
+                            tonic::Response<Self::Response>,
+                            tonic::Status,
+                        >;
+                        fn call(
+                            &mut self,
+                            request: tonic::Request<super::SessionSetPubKeyRequest>,
+                        ) -> Self::Future {
+                            let inner = Arc::clone(&self.0);
+                            let fut = async move {
+                                <T as Core>::session_set_pub_key(&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 = SessionSetPubKeySvc(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(

utils/quartz-relayer/Cargo.toml

@@ -10,6 +10,7 @@ cosmrs = { version = "=0.11.0", features = ["cosmwasm"] }
 cosmwasm-std = "1.4.0"
 displaydoc = { version = "0.2.3", default-features = false }
 ecies = "0.2.6"
+k256 = { version = "0.13.2", default-features = false, features = ["ecdsa", "alloc"] }
 serde = { version = "1.0.189", features = ["derive"] }
 serde_json = "1.0.94"
 subtle-encoding = { version = "0.5.1", features = ["bech32-preview"] }

utils/quartz-relayer/src/types.rs

@@ -1,8 +1,13 @@
 use cosmwasm_std::{HexBinary, StdError};
-use quartz_cw::state::{Config, Nonce, RawConfig};
+use k256::ecdsa::VerifyingKey;
+use quartz_cw::{
+    error::Error as QuartzCwError,
+    state::{Config, Nonce, RawConfig},
+};
 use quartz_proto::quartz::{
     InstantiateResponse as RawInstantiateResponse,
     SessionCreateResponse as RawSessionCreateResponse,
+    SessionSetPubKeyResponse as RawSessionSetPubKeyResponse,
 };
 use serde::{Deserialize, Serialize};
 
@@ -165,3 +170,94 @@ impl From<SessionCreateResponseMsg> for RawSessionCreateResponseMsg {
         }
     }
 }
+
+#[derive(Clone, Debug, PartialEq)]
+pub struct SessionSetPubKeyResponse {
+    message: SessionSetPubKeyResponseMsg,
+}
+
+impl SessionSetPubKeyResponse {
+    pub fn new(nonce: Nonce, pub_key: VerifyingKey, quote: Vec<u8>) -> Self {
+        Self {
+            message: SessionSetPubKeyResponseMsg {
+                nonce,
+                pub_key,
+                quote,
+            },
+        }
+    }
+
+    pub fn quote(&self) -> &[u8] {
+        &self.message.quote
+    }
+
+    pub fn into_message(self) -> SessionSetPubKeyResponseMsg {
+        self.message
+    }
+}
+
+impl TryFrom<RawSessionSetPubKeyResponse> for SessionSetPubKeyResponse {
+    type Error = StdError;
+
+    fn try_from(value: RawSessionSetPubKeyResponse) -> Result<Self, Self::Error> {
+        let raw_message: RawSessionSetPubKeyResponseMsg = serde_json::from_str(&value.message)
+            .map_err(|e| StdError::parse_err("RawSessionSetPubKeyResponseMsg", e))?;
+        Ok(Self {
+            message: raw_message.try_into()?,
+        })
+    }
+}
+
+impl From<SessionSetPubKeyResponse> for RawSessionSetPubKeyResponse {
+    fn from(value: SessionSetPubKeyResponse) -> Self {
+        let raw_message: RawSessionSetPubKeyResponseMsg = value.message.into();
+        Self {
+            message: serde_json::to_string(&raw_message).expect("infallible serializer"),
+        }
+    }
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub struct SessionSetPubKeyResponseMsg {
+    nonce: Nonce,
+    pub_key: VerifyingKey,
+    quote: Vec<u8>,
+}
+
+impl SessionSetPubKeyResponseMsg {
+    pub fn into_tuple(self) -> (VerifyingKey, Vec<u8>) {
+        (self.pub_key, self.quote)
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub struct RawSessionSetPubKeyResponseMsg {
+    nonce: HexBinary,
+    pub_key: HexBinary,
+    quote: HexBinary,
+}
+
+impl TryFrom<RawSessionSetPubKeyResponseMsg> for SessionSetPubKeyResponseMsg {
+    type Error = StdError;
+
+    fn try_from(value: RawSessionSetPubKeyResponseMsg) -> Result<Self, Self::Error> {
+        let pub_key = VerifyingKey::from_sec1_bytes(&value.pub_key)
+            .map_err(QuartzCwError::from)
+            .map_err(|e| StdError::generic_err(e.to_string()))?;
+        Ok(Self {
+            nonce: value.nonce.to_array()?,
+            pub_key,
+            quote: value.quote.into(),
+        })
+    }
+}
+
+impl From<SessionSetPubKeyResponseMsg> for RawSessionSetPubKeyResponseMsg {
+    fn from(value: SessionSetPubKeyResponseMsg) -> Self {
+        Self {
+            nonce: value.nonce.into(),
+            pub_key: value.pub_key.to_sec1_bytes().into_vec().into(),
+            quote: value.quote.into(),
+        }
+    }
+}