add register threshold-enc-key functionality

Author: alxiong

Cargo.lock

@@ -215,7 +215,6 @@ test-dyn-comm: build_release_until build-test-utils
                  --stamp-dir /tmp \
                  --output test-configs/c1" \
     --run   "7:target/release/register \
-                 -i 0 \
                  -u http://localhost:8545 \
                  -k 0x2bbf15bc655c4cc157b769cfcb1ea9924b9e1a35 \
                  -c test-configs/c1/committee.toml" \

justfile

@@ -20,8 +20,10 @@ bincode = { workspace = true }
 clap = { workspace = true }
 futures = { workspace = true }
 rand = { workspace = true }
+reqwest = { workspace = true }
 serde = { workspace = true }
 timeboost-config = { path = "../timeboost-config" }
+timeboost-crypto = { path = "../timeboost-crypto" }
 timeboost-types = { path = "../timeboost-types" }
 timeboost-utils = { path = "../timeboost-utils" }
 tokio = { workspace = true }
@@ -30,4 +32,4 @@ tracing = { workspace = true }
 url = { workspace = true }
 
 [build-dependencies]
-alloy = { workspace = true }
+alloy = { workspace = true }

timeboost-contract/Cargo.toml

@@ -4,13 +4,17 @@ use alloy::{
     primitives::Address,
     providers::{Provider, WalletProvider},
 };
-use anyhow::{Context, Result, bail};
-use clap::Parser;
+use anyhow::{Context, Result, anyhow, bail};
+use clap::{Parser, ValueEnum};
+use reqwest::Client;
 use std::path::PathBuf;
+use std::time::Duration;
 use timeboost_config::CommitteeConfig;
 use timeboost_contract::{CommitteeMemberSol, KeyManager, provider::build_provider};
+use timeboost_crypto::prelude::ThresholdEncKey;
+use timeboost_utils::enc_key::ThresholdEncKeyCellAccumulator;
 use timeboost_utils::types::logging;
-use tracing::info;
+use tracing::{info, warn};
 use url::Url;
 
 #[derive(Clone, Debug, Parser)]
@@ -22,7 +26,7 @@ struct Args {
     )]
     mnemonic: String,
 
-    #[clap(short, long)]
+    #[clap(short, long, default_value_t = 0)]
     index: u32,
 
     #[clap(short, long)]
@@ -35,6 +39,20 @@ struct Args {
     /// Path to the committee.toml config for the next committee
     #[clap(short, long)]
     config: PathBuf,
+
+    /// What to register (new committee or threshold enc key?)
+    #[clap(long, short, default_value = "new-committee")]
+    action: Action,
+}
+
+/// Specific register action
+#[derive(Clone, Copy, Debug, Default, ValueEnum)]
+enum Action {
+    /// register the next committee
+    #[default]
+    NewCommittee,
+    /// register the threshold encryption key (when ready)
+    ThresholdEncKey,
 }
 
 #[tokio::main]
@@ -46,8 +64,6 @@ async fn main() -> Result<()> {
         .await
         .context(format!("Failed to read config file: {:?}", &args.config))?;
 
-    info!("Start committee registration");
-
     let provider = build_provider(args.mnemonic.clone(), args.index, args.url.clone())?;
     let addr = args.key_manager_addr;
     if provider
@@ -63,40 +79,79 @@ async fn main() -> Result<()> {
 
     let contract = KeyManager::new(addr, provider);
 
-    // prepare input argument from config file
-    let members = config
-        .members
-        .iter()
-        .map(|m| {
-            Ok::<_, anyhow::Error>(CommitteeMemberSol {
-                sigKey: m.signing_key.to_bytes().into(),
-                dhKey: m.dh_key.as_bytes().into(),
-                dkgKey: m.dkg_enc_key.to_bytes()?.into(),
-                networkAddress: m.public_address.to_string(),
-            })
-        })
-        .collect::<Result<Vec<_>, _>>()?;
-
-    let timestamp: u64 = config
-        .effective_timestamp
-        .as_second()
-        .try_into()
-        .with_context(|| {
-            format!(
-                "failed to convert timestamp {} to u64",
-                config.effective_timestamp
-            )
-        })?;
-
-    // send tx and invoke the contract
-    let _tx_receipt = contract
-        .setNextCommittee(timestamp, members)
-        .send()
-        .await?
-        .get_receipt()
-        .await?;
-
-    let registered_cid = contract.nextCommitteeId().call().await? - 1;
-    info!("Registered new committee with id: {registered_cid}");
+    match args.action {
+        Action::NewCommittee => {
+            info!("Start committee registration");
+            // prepare input argument from config file
+            let members = config
+                .members
+                .iter()
+                .map(|m| {
+                    Ok::<_, anyhow::Error>(CommitteeMemberSol {
+                        sigKey: m.signing_key.to_bytes().into(),
+                        dhKey: m.dh_key.as_bytes().into(),
+                        dkgKey: m.dkg_enc_key.to_bytes()?.into(),
+                        networkAddress: m.public_address.to_string(),
+                    })
+                })
+                .collect::<Result<Vec<_>, _>>()?;
+
+            let timestamp: u64 = config
+                .effective_timestamp
+                .as_second()
+                .try_into()
+                .with_context(|| {
+                    format!(
+                        "failed to convert timestamp {} to u64",
+                        config.effective_timestamp
+                    )
+                })?;
+
+            // send tx and invoke the contract
+            let _tx_receipt = contract
+                .setNextCommittee(timestamp, members)
+                .send()
+                .await?
+                .get_receipt()
+                .await?;
+
+            let registered_cid = contract.nextCommitteeId().call().await? - 1;
+            info!("Registered new committee with id: {registered_cid}");
+        }
+        Action::ThresholdEncKey => {
+            info!("Start threshold encryption key registration");
+            let client = Client::builder().timeout(Duration::from_secs(1)).build()?;
+            let urls = config
+                .members
+                .iter()
+                .map(|m| {
+                    let addr = m.http_api.clone();
+                    Url::parse(&format!("http://{addr}/v1/encryption-key"))
+                        .with_context(|| format!("parsing {addr} into a url"))
+                })
+                .collect::<Result<Vec<_>, _>>()?;
+
+            let mut acc = ThresholdEncKeyCellAccumulator::new(client, urls.into_iter());
+            let Some(key) = acc.enc_key().await else {
+                warn!("encryption key not available yet");
+                return Err(anyhow!(
+                    "threshold enc key not available on enough nodes, try later"
+                ));
+            };
+
+            let _tx_receipt = contract
+                .setThresholdEncryptionKey(key.to_owned().to_bytes()?.into())
+                .send()
+                .await?
+                .get_receipt()
+                .await?;
+            assert_eq!(
+                &ThresholdEncKey::from_bytes(&contract.thresholdEncryptionKey().call().await?.0)?,
+                key
+            );
+
+            info!("Registered threshold encryption key");
+        }
+    }
     Ok(())
 }

timeboost-contract/src/binaries/register.rs

@@ -5,6 +5,7 @@ use anyhow::anyhow;
 use ark_ec::{AffineRepr, CurveGroup, hashing::HashToCurve};
 use ark_ff::{PrimeField, UniformRand};
 use ark_poly::{DenseUVPolynomial, Polynomial, polynomial::univariate::DensePolynomial};
+use ark_serialize::SerializationError;
 use ark_std::rand::Rng;
 use ark_std::rand::rngs::OsRng;
 use derive_more::From;
@@ -342,6 +343,19 @@ pub struct PublicKey<C: CurveGroup> {
     key: C,
 }
 
+impl<C: CurveGroup> PublicKey<C> {
+    pub fn to_bytes(&self) -> Result<Vec<u8>, SerializationError> {
+        let mut v = Vec::new();
+        self.key.serialize_compressed(&mut v)?;
+        Ok(v)
+    }
+
+    pub fn from_bytes(value: &[u8]) -> Result<Self, SerializationError> {
+        let key = C::deserialize_compressed(value)?;
+        Ok(Self { key })
+    }
+}
+
 #[serde_as]
 #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Zeroize, ZeroizeOnDrop, From)]
 pub struct KeyShare<C: CurveGroup> {

timeboost-crypto/src/sg_encryption.rs

@@ -1,6 +1,6 @@
 use reqwest::{Client, Url};
 use std::collections::HashMap;
-use timeboost::crypto::prelude::ThresholdEncKey;
+use timeboost_crypto::prelude::ThresholdEncKey;
 
 use tracing::warn;
 
@@ -28,7 +28,7 @@ async fn fetch_encryption_key(client: &Client, enckey_url: &Url) -> Option<Thres
 }
 
 /// helper struct to keep track of sufficient quorum of DKG keys
-pub(crate) struct ThresholdEncKeyCellAccumulator {
+pub struct ThresholdEncKeyCellAccumulator {
     client: Client,
     // DKG results on individual node
     results: HashMap<Url, Option<ThresholdEncKey>>,
@@ -40,7 +40,7 @@ pub(crate) struct ThresholdEncKeyCellAccumulator {
 
 impl ThresholdEncKeyCellAccumulator {
     /// give a list of TimeboostApi's endpoint to query `/enckey` status
-    pub(crate) fn new(client: Client, urls: impl Iterator<Item = Url>) -> Self {
+    pub fn new(client: Client, urls: impl Iterator<Item = Url>) -> Self {
         let results: HashMap<Url, Option<ThresholdEncKey>> = urls.map(|url| (url, None)).collect();
         let threshold = results.len().div_ceil(3);
         Self {
@@ -53,7 +53,7 @@ impl ThresholdEncKeyCellAccumulator {
 
     /// try to get the threshold encryption key, only available after a threshold of nodes
     /// finish their DKG processes.
-    pub(crate) async fn enc_key(&mut self) -> Option<&ThresholdEncKey> {
+    pub async fn enc_key(&mut self) -> Option<&ThresholdEncKey> {
         // if result is already available, directly return
         if self.output.is_some() {
             self.output.as_ref()

yapper/src/enc_key.rs renamed to timeboost-utils/src/enc_key.rs

@@ -1,3 +1,4 @@
+pub mod enc_key;
 pub mod load_generation;
 pub mod types;
 pub mod until;

timeboost-utils/src/lib.rs

@@ -23,7 +23,6 @@ use crate::config::YapperConfig;
 use crate::yapper::Yapper;
 
 mod config;
-mod enc_key;
 mod yapper;
 
 #[derive(Parser, Debug)]

yapper/src/main.rs

@@ -11,11 +11,12 @@ use anyhow::{Context, Result};
 use futures::future::join_all;
 use reqwest::{Client, Url};
 use timeboost::types::BundleVariant;
+use timeboost_utils::enc_key::ThresholdEncKeyCellAccumulator;
 use timeboost_utils::load_generation::{TxInfo, make_bundle, make_dev_acct_bundle, tps_to_millis};
 use tokio::time::interval;
 use tracing::warn;
 
-use crate::{config::YapperConfig, enc_key::ThresholdEncKeyCellAccumulator};
+use crate::config::YapperConfig;
 
 /// This is the address of the prefunded dev account for nitro chain
 /// https://docs.arbitrum.io/run-arbitrum-node/run-local-full-chain-simulation#default-endpoints-and-addresses