Store sk in CoreService and instantiate it in main.rs

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<dyn std::error::Error>> {
         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?;
 

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<T> = Result<T, Status>;
-
 #[derive(Clone, Debug)]
 pub struct CoreService<A> {
     config: Config,
     nonce: Arc<Mutex<Nonce>>,
+    sk: Arc<Mutex<Option<SigningKey>>>,
     attestor: A,
 }
 
@@ -34,10 +33,11 @@ impl<A> CoreService<A>
 where
     A: Attestor,
 {
-    pub fn new(config: Config, attestor: A) -> Self {
+    pub fn new(config: Config, sk: Arc<Mutex<Option<SigningKey>>>, attestor: A) -> Self {
         Self {
             config,
             nonce: Arc::new(Mutex::new([0u8; 32])),
+            sk,
             attestor,
         }
     }
@@ -87,7 +87,9 @@ where
     ) -> TonicResult<Response<RawSessionSetPubKeyResponse>> {
         // 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);