diff --git a/enclaves/quartz/src/main.rs b/enclaves/quartz/src/main.rs index b8aa4c2..649fe71 100644 --- a/enclaves/quartz/src/main.rs +++ b/enclaves/quartz/src/main.rs @@ -18,7 +18,10 @@ mod cli; mod proto; mod server; -use std::time::Duration; +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; use clap::Parser; use quartz_cw::state::{Config, LightClientOpts}; @@ -56,8 +59,14 @@ async fn main() -> Result<(), Box> { light_client_opts, ); + let sk = Arc::new(Mutex::new(None)); + Server::builder() - .add_service(CoreServer::new(CoreService::new(config, EpidAttestor))) + .add_service(CoreServer::new(CoreService::new( + config, + sk.clone(), + EpidAttestor, + ))) .serve(args.rpc_addr) .await?; diff --git a/enclaves/quartz/src/server.rs b/enclaves/quartz/src/server.rs index 0dda79f..967393b 100644 --- a/enclaves/quartz/src/server.rs +++ b/enclaves/quartz/src/server.rs @@ -17,16 +17,15 @@ use quartz_proto::quartz::{ }; use quartz_relayer::types::{InstantiateResponse, SessionCreateResponse, SessionSetPubKeyResponse}; use rand::Rng; -use tonic::{Request, Response, Status}; +use tonic::{Request, Response, Result as TonicResult, Status}; use crate::attestor::Attestor; -type TonicResult = Result; - #[derive(Clone, Debug)] pub struct CoreService { config: Config, nonce: Arc>, + sk: Arc>>, attestor: A, } @@ -34,10 +33,11 @@ impl CoreService where A: Attestor, { - pub fn new(config: Config, attestor: A) -> Self { + pub fn new(config: Config, sk: Arc>>, attestor: A) -> Self { Self { config, nonce: Arc::new(Mutex::new([0u8; 32])), + sk, attestor, } } @@ -87,7 +87,9 @@ where ) -> TonicResult> { // FIXME(hu55a1n1) - disallow calling more than once let nonce = self.nonce.lock().unwrap(); + let sk = SigningKey::random(&mut rand::thread_rng()); + *self.sk.lock().unwrap() = Some(sk.clone()); let pk = sk.verifying_key(); let session_set_pub_key_msg = SessionSetPubKey::new(*nonce, *pk);