feat: distribute public keys

integration-tests/tests/docker/mod.rs

@@ -8,7 +8,6 @@ use bollard::{
     service::{HostConfig, Ipam, PortBinding},
     Docker,
 };
-use curv::elliptic::curves::{Ed25519, Point};
 use futures::{lock::Mutex, StreamExt};
 use hyper::{Body, Client, Method, Request, StatusCode, Uri};
 use mpc_recovery::msg::{AddKeyRequest, AddKeyResponse, NewAccountRequest, NewAccountResponse};
@@ -285,7 +284,6 @@ impl SignNode {
         docker: &Docker,
         network: &str,
         node_id: u64,
-        pk_set: &Vec<Point<Ed25519>>,
         sk_share: &ExpandedKeyPair,
         datastore_url: &str,
         gcp_project_id: &str,
@@ -297,8 +295,6 @@ impl SignNode {
             "start-sign".to_string(),
             "--node-id".to_string(),
             node_id.to_string(),
-            "--pk-set".to_string(),
-            serde_json::to_string(&pk_set)?,
             "--sk-share".to_string(),
             serde_json::to_string(&sk_share)?,
             "--web-port".to_string(),

integration-tests/tests/lib.rs

@@ -80,7 +80,6 @@ where
             &docker,
             NETWORK,
             i as u64,
-            &pk_set,
             share,
             &datastore.address,
             GCP_PROJECT_ID,

mpc-recovery/src/leader_node/mod.rs

@@ -1,5 +1,8 @@
 use crate::key_recovery::get_user_recovery_pk;
-use crate::msg::{AddKeyRequest, AddKeyResponse, NewAccountRequest, NewAccountResponse};
+use crate::msg::{
+    AcceptNodePublicKeysRequest, AddKeyRequest, AddKeyResponse, NewAccountRequest,
+    NewAccountResponse,
+};
 use crate::oauth::{OAuthTokenVerifier, UniversalTokenVerifier};
 use crate::relayer::error::RelayerError;
 use crate::relayer::msg::RegisterAccountRequest;
@@ -10,6 +13,7 @@ use crate::transaction::{
 };
 use crate::{nar, NodeId};
 use axum::{http::StatusCode, routing::post, Extension, Json, Router};
+use curv::elliptic::curves::{Ed25519, Point};
 use near_crypto::{ParseKeyError, PublicKey, SecretKey};
 use near_primitives::account::id::ParseAccountError;
 use near_primitives::types::AccountId;
@@ -76,6 +80,24 @@ pub async fn run(config: Config) {
         pagoda_firebase_audience_id,
     };
 
+    // Get keys from all sign nodes, and broadcast them out as a set.
+    let pk_set = match gather_sign_node_pks(&state).await {
+        Ok(pk_set) => pk_set,
+        Err(err) => {
+            tracing::error!("Unable to gather public keys: {err}");
+            return;
+        }
+    };
+    tracing::debug!(?pk_set, "Gathered public keys");
+    let messages = match broadcast_pk_set(&state, pk_set).await {
+        Ok(messages) => messages,
+        Err(err) => {
+            tracing::error!("Unable to broadcast public keys: {err}");
+            return;
+        }
+    };
+    tracing::debug!(?messages, "broadcasted public key statuses");
+
     //TODO: not secure, allow only for testnet, whitelist endpoint etc. for mainnet
     let cors_layer = tower_http::cors::CorsLayer::permissive();
 
@@ -439,6 +461,56 @@ async fn add_key<T: OAuthTokenVerifier>(
         }
     }
 }
+
+async fn gather_sign_node_pks(state: &LeaderState) -> anyhow::Result<Vec<Point<Ed25519>>> {
+    let fut = nar::retry_every(std::time::Duration::from_millis(250), || async {
+        let results: anyhow::Result<Vec<(usize, Point<Ed25519>)>> = crate::transaction::call(
+            &state.reqwest_client,
+            &state.sign_nodes,
+            "public_key_node",
+            (),
+        )
+        .await;
+        let mut results = match results {
+            Ok(results) => results,
+            Err(err) => {
+                tracing::debug!("failed to gather pk: {err}");
+                return Err(err);
+            }
+        };
+
+        results.sort_by_key(|(index, _)| *index);
+        let results: Vec<Point<Ed25519>> =
+            results.into_iter().map(|(_index, point)| point).collect();
+
+        anyhow::Result::Ok(results)
+    });
+
+    let results = tokio::time::timeout(std::time::Duration::from_secs(10), fut)
+        .await
+        .map_err(|_| anyhow::anyhow!("timeout gathering sign node pks"))??;
+    Ok(results)
+}
+
+async fn broadcast_pk_set(
+    state: &LeaderState,
+    pk_set: Vec<Point<Ed25519>>,
+) -> anyhow::Result<Vec<String>> {
+    let request = AcceptNodePublicKeysRequest {
+        public_keys: pk_set,
+    };
+
+    let messages: Vec<String> = crate::transaction::call(
+        &state.reqwest_client,
+        &state.sign_nodes,
+        "accept_pk_set",
+        request,
+    )
+    .await?;
+
+    Ok(messages)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

mpc-recovery/src/main.rs

@@ -1,5 +1,4 @@
 use clap::Parser;
-use curv::elliptic::curves::{Ed25519, Point};
 use mpc_recovery::{gcp::GcpService, LeaderConfig};
 use multi_party_eddsa::protocols::ExpandedKeyPair;
 use near_primitives::types::AccountId;
@@ -61,9 +60,6 @@ enum Cli {
         /// Node ID
         #[arg(long, env("MPC_RECOVERY_NODE_ID"))]
         node_id: u64,
-        /// Root public key
-        #[arg(long, env("MPC_RECOVERY_PK_SET"))]
-        pk_set: String,
         /// Secret key share, will be pulled from GCP Secret Manager if omitted
         #[arg(long, env("MPC_RECOVERY_SK_SHARE"))]
         sk_share: Option<String>,
@@ -166,7 +162,6 @@ async fn main() -> anyhow::Result<()> {
         }
         Cli::StartSign {
             node_id,
-            pk_set,
             sk_share,
             web_port,
             gcp_project_id,
@@ -175,12 +170,10 @@ async fn main() -> anyhow::Result<()> {
             let gcp_service = GcpService::new(gcp_project_id, gcp_datastore_url).await?;
             let sk_share = load_sh_skare(&gcp_service, node_id, sk_share).await?;
 
-            // TODO put these in a better defined format
-            let pk_set: Vec<Point<Ed25519>> = serde_json::from_str(&pk_set).unwrap();
             // TODO Import just the private key and derive the rest
             let sk_share: ExpandedKeyPair = serde_json::from_str(&sk_share).unwrap();
 
-            mpc_recovery::run_sign_node(gcp_service, node_id, pk_set, sk_share, web_port).await;
+            mpc_recovery::run_sign_node(gcp_service, node_id, sk_share, web_port).await;
         }
     }
 

mpc-recovery/src/msg.rs

@@ -1,3 +1,4 @@
+use curv::elliptic::curves::{Ed25519, Point};
 use ed25519_dalek::Signature;
 use serde::{Deserialize, Serialize};
 
@@ -77,6 +78,11 @@ pub struct SigShareRequest {
     pub payload: Vec<u8>,
 }
 
+#[derive(Serialize, Deserialize, Debug)]
+pub struct AcceptNodePublicKeysRequest {
+    pub public_keys: Vec<Point<Ed25519>>,
+}
+
 mod hex_sig_share {
     use ed25519_dalek::Signature;
     use serde::{Deserialize, Deserializer, Serializer};

mpc-recovery/src/nar.rs

@@ -6,6 +6,7 @@
 use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 use std::sync::atomic::{AtomicU64, Ordering};
+use std::time::Duration;
 
 use near_crypto::PublicKey;
 use near_jsonrpc_client::errors::{JsonRpcError, JsonRpcServerError};
@@ -23,6 +24,16 @@ use crate::relayer::error::RelayerError;
 
 pub(crate) type CachedAccessKeyNonces = RwLock<HashMap<(AccountId, PublicKey), AtomicU64>>;
 
+pub(crate) async fn retry_every<R, E, T, F>(interval: Duration, task: F) -> T::Output
+where
+    F: FnMut() -> T,
+    T: core::future::Future<Output = core::result::Result<R, E>>,
+{
+    let retry_strategy = std::iter::repeat_with(|| interval);
+    let task = Retry::spawn(retry_strategy, task);
+    task.await
+}
+
 pub(crate) async fn retry<R, E, T, F>(task: F) -> T::Output
 where
     F: FnMut() -> T,

mpc-recovery/src/sign_node/aggregate_signer.rs

@@ -1,5 +1,6 @@
 use std::collections::HashMap;
 use std::hash::{Hash, Hasher};
+use std::sync::Arc;
 
 use curv::arithmetic::Converter;
 use curv::cryptographic_primitives::commitments::{
@@ -13,6 +14,7 @@ use multi_party_eddsa::protocols::aggsig::{self, KeyAgg, SignSecondMsg};
 use rand8::rngs::OsRng;
 use rand8::Rng;
 use serde::{Deserialize, Serialize};
+use tokio::sync::RwLock;
 
 use crate::transaction::{to_dalek_public_key, to_dalek_signature};
 
@@ -59,7 +61,7 @@ impl SigningState {
         Ok(commitment)
     }
 
-    pub fn get_reveal(
+    pub async fn get_reveal(
         &mut self,
         node_info: NodeInfo,
         recieved_commitments: Vec<SignedCommitment>,
@@ -79,7 +81,7 @@ impl SigningState {
             .remove(&our_c.commitment)
             .ok_or(format!("Committment {:?} not found", &our_c.commitment))?;
 
-        let (reveal, state) = state.reveal(&node_info, recieved_commitments)?;
+        let (reveal, state) = state.reveal(&node_info, recieved_commitments).await?;
         let reveal = Reveal(reveal);
         self.revealed.insert(reveal.clone(), state);
         Ok(reveal)
@@ -166,13 +168,14 @@ impl Committed {
         Ok((sc, s))
     }
 
-    pub fn reveal(
+    pub async fn reveal(
         self,
         node_info: &NodeInfo,
         commitments: Vec<SignedCommitment>,
     ) -> Result<(SignSecondMsg, Revealed), String> {
         let (commitments, signing_public_keys) = node_info
-            .signed_by_every_node(commitments)?
+            .signed_by_every_node(commitments)
+            .await?
             .into_iter()
             .unzip();
         Ok((
@@ -227,16 +230,27 @@ impl Revealed {
 // Stores info about the other nodes we're interacting with
 #[derive(Clone)]
 pub struct NodeInfo {
-    pub nodes_public_keys: Vec<Point<Ed25519>>,
+    pub nodes_public_keys: Arc<RwLock<Option<Vec<Point<Ed25519>>>>>,
     pub our_index: usize,
 }
 
 impl NodeInfo {
-    fn signed_by_every_node(
+    pub fn new(our_index: usize, nodes_public_keys: Option<Vec<Point<Ed25519>>>) -> Self {
+        Self {
+            our_index,
+            nodes_public_keys: Arc::new(RwLock::new(nodes_public_keys)),
+        }
+    }
+
+    async fn signed_by_every_node(
         &self,
         signed: Vec<SignedCommitment>,
     ) -> Result<Vec<(AggrCommitment, Point<Ed25519>)>, String> {
         self.nodes_public_keys
+            .read()
+            .await
+            .as_ref()
+            .ok_or_else(|| "No nodes public keys available to sign".to_string())?
             .iter()
             .zip(signed.iter())
             .map(|(public_key, signed)| signed.verify(public_key))
@@ -320,8 +334,8 @@ mod tests {
     use ed25519_dalek::{SignatureError, Verifier};
     use multi_party_eddsa::protocols::ExpandedKeyPair;
 
-    #[test]
-    fn aggregate_signatures() {
+    #[tokio::test]
+    async fn aggregate_signatures() {
         pub fn verify_dalek(
             pk: &Point<Ed25519>,
             sig: &protocols::Signature,
@@ -349,10 +363,7 @@ mod tests {
             n3.public_key.clone(),
         ];
 
-        let ni = |n| NodeInfo {
-            nodes_public_keys: nodes_public_keys.clone(),
-            our_index: n,
-        };
+        let ni = |n| NodeInfo::new(n, Some(nodes_public_keys.clone()));
 
         // Set up nodes with that config
         let mut s1 = SigningState::new();
@@ -368,9 +379,9 @@ mod tests {
         ];
 
         let reveals = vec![
-            s1.get_reveal(ni(0), commitments.clone()).unwrap(),
-            s2.get_reveal(ni(1), commitments.clone()).unwrap(),
-            s3.get_reveal(ni(2), commitments.clone()).unwrap(),
+            s1.get_reveal(ni(0), commitments.clone()).await.unwrap(),
+            s2.get_reveal(ni(1), commitments.clone()).await.unwrap(),
+            s3.get_reveal(ni(2), commitments.clone()).await.unwrap(),
         ];
 
         let sig_shares = vec![