Add randomness-related admin handlers for disaster recovery use (#18500)

crates/sui-network/src/randomness/mod.rs

@@ -6,7 +6,10 @@ use anemo::PeerId;
 use anyhow::Result;
 use fastcrypto::groups::bls12381;
 use fastcrypto_tbls::{
-    dkg, nodes::PartyId, tbls::ThresholdBls, types::ShareIndex, types::ThresholdBls12381MinSig,
+    dkg,
+    nodes::PartyId,
+    tbls::ThresholdBls,
+    types::{ShareIndex, ThresholdBls12381MinSig},
 };
 use mysten_metrics::spawn_monitored_task;
 use mysten_network::anemo_ext::NetworkExt;
@@ -24,7 +27,9 @@ use sui_types::{
     committee::EpochId,
     crypto::{RandomnessPartialSignature, RandomnessRound, RandomnessSignature},
 };
-use tokio::sync::{mpsc, OnceCell};
+use tokio::sync::{
+    OnceCell, {mpsc, oneshot},
+};
 use tracing::{debug, error, info, instrument, warn};
 
 mod auth;
@@ -101,6 +106,54 @@ impl Handle {
             .expect("RandomnessEventLoop mailbox should not overflow or be closed")
     }
 
+    /// Admin interface handler: generates partial signatures for the given round at the
+    /// current epoch.
+    pub fn admin_get_partial_signatures(
+        &self,
+        round: RandomnessRound,
+        tx: oneshot::Sender<Vec<u8>>,
+    ) {
+        self.sender
+            .try_send(RandomnessMessage::AdminGetPartialSignatures(round, tx))
+            .expect("RandomnessEventLoop mailbox should not overflow or be closed")
+    }
+
+    /// Admin interface handler: injects partial signatures for the given round at the
+    /// current epoch, skipping validity checks.
+    pub fn admin_inject_partial_signatures(
+        &self,
+        authority_name: AuthorityName,
+        round: RandomnessRound,
+        sigs: Vec<RandomnessPartialSignature>,
+        result_channel: oneshot::Sender<Result<()>>,
+    ) {
+        self.sender
+            .try_send(RandomnessMessage::AdminInjectPartialSignatures(
+                authority_name,
+                round,
+                sigs,
+                result_channel,
+            ))
+            .expect("RandomnessEventLoop mailbox should not overflow or be closed")
+    }
+
+    /// Admin interface handler: injects full signature for the given round at the
+    /// current epoch, skipping validity checks.
+    pub fn admin_inject_full_signature(
+        &self,
+        round: RandomnessRound,
+        sig: RandomnessSignature,
+        result_channel: oneshot::Sender<Result<()>>,
+    ) {
+        self.sender
+            .try_send(RandomnessMessage::AdminInjectFullSignature(
+                round,
+                sig,
+                result_channel,
+            ))
+            .expect("RandomnessEventLoop mailbox should not overflow or be closed")
+    }
+
     // For testing.
     pub fn new_stub() -> Self {
         let (sender, mut receiver) = mpsc::channel(1);
@@ -120,7 +173,7 @@ impl Handle {
     }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Debug)]
 enum RandomnessMessage {
     UpdateEpoch(
         EpochId,
@@ -138,6 +191,18 @@ enum RandomnessMessage {
         Vec<Vec<u8>>,
         Option<RandomnessSignature>,
     ),
+    AdminGetPartialSignatures(RandomnessRound, oneshot::Sender<Vec<u8>>),
+    AdminInjectPartialSignatures(
+        AuthorityName,
+        RandomnessRound,
+        Vec<RandomnessPartialSignature>,
+        oneshot::Sender<Result<()>>,
+    ),
+    AdminInjectFullSignature(
+        RandomnessRound,
+        RandomnessSignature,
+        oneshot::Sender<Result<()>>,
+    ),
 }
 
 struct RandomnessEventLoop {
@@ -220,6 +285,24 @@ impl RandomnessEventLoop {
                     self.receive_partial_signatures(peer_id, epoch, round, partial_sigs)
                 }
             }
+            RandomnessMessage::AdminGetPartialSignatures(round, tx) => {
+                self.admin_get_partial_signatures(round, tx)
+            }
+            RandomnessMessage::AdminInjectPartialSignatures(
+                authority_name,
+                round,
+                sigs,
+                result_channel,
+            ) => {
+                let _ = result_channel.send(self.admin_inject_partial_signatures(
+                    authority_name,
+                    round,
+                    sigs,
+                ));
+            }
+            RandomnessMessage::AdminInjectFullSignature(round, sig, result_channel) => {
+                let _ = result_channel.send(self.admin_inject_full_signature(round, sig));
+            }
         }
     }
 
@@ -875,4 +958,56 @@ impl RandomnessEventLoop {
         }
         self.metrics.set_num_rounds_pending(num_rounds_pending);
     }
+
+    fn admin_get_partial_signatures(&self, round: RandomnessRound, tx: oneshot::Sender<Vec<u8>>) {
+        let shares = if let Some(shares) = self.dkg_output.as_ref().and_then(|d| d.shares.as_ref())
+        {
+            shares
+        } else {
+            let _ = tx.send(Vec::new()); // no error handling needed if receiver is already dropped
+            return;
+        };
+
+        let partial_sigs =
+            ThresholdBls12381MinSig::partial_sign_batch(shares.iter(), &round.signature_message());
+        // no error handling needed if receiver is already dropped
+        let _ = tx.send(bcs::to_bytes(&partial_sigs).expect("serialization should not fail"));
+    }
+
+    fn admin_inject_partial_signatures(
+        &mut self,
+        authority_name: AuthorityName,
+        round: RandomnessRound,
+        sigs: Vec<RandomnessPartialSignature>,
+    ) -> Result<()> {
+        let peer_id = self
+            .authority_info
+            .get(&authority_name)
+            .map(|(peer_id, _)| *peer_id)
+            .ok_or(anyhow::anyhow!("unknown AuthorityName {authority_name:?}"))?;
+        self.received_partial_sigs.insert((round, peer_id), sigs);
+        self.maybe_aggregate_partial_signatures(self.epoch, round);
+        Ok(())
+    }
+
+    fn admin_inject_full_signature(
+        &mut self,
+        round: RandomnessRound,
+        sig: RandomnessSignature,
+    ) -> Result<()> {
+        let vss_pk = {
+            let Some(dkg_output) = &self.dkg_output else {
+                return Err(anyhow::anyhow!(
+                    "called admin_inject_full_signature before DKG completed"
+                ));
+            };
+            &dkg_output.vss_pk
+        };
+
+        ThresholdBls12381MinSig::verify(vss_pk.c0(), &round.signature_message(), &sig)
+            .map_err(|e| anyhow::anyhow!("invalid full signature: {e:?}"))?;
+
+        self.process_valid_full_signature(self.epoch, round, sig);
+        Ok(())
+    }
 }

crates/sui-node/Cargo.toml

@@ -12,6 +12,8 @@ anemo-tower.workspace = true
 arc-swap.workspace = true
 axum.workspace = true
 anyhow.workspace = true
+base64.workspace = true
+bcs.workspace = true
 clap.workspace = true
 prometheus.workspace = true
 tokio = { workspace = true, features = ["full"] }

crates/sui-node/src/admin.rs

@@ -8,12 +8,21 @@ use axum::{
     routing::{get, post},
     Router,
 };
+use base64::Engine;
 use humantime::parse_duration;
 use serde::Deserialize;
-use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 use std::sync::Arc;
-use sui_types::error::SuiError;
+use std::{
+    net::{IpAddr, Ipv4Addr, SocketAddr},
+    str::FromStr,
+};
+use sui_types::{
+    base_types::AuthorityName,
+    crypto::{RandomnessPartialSignature, RandomnessRound, RandomnessSignature},
+    error::SuiError,
+};
 use telemetry_subscribers::TracingHandle;
+use tokio::sync::oneshot;
 use tracing::info;
 
 // Example commands:
@@ -47,6 +56,18 @@ use tracing::info;
 // Reset tracing to the TRACE_FILTER env var.
 //
 //   $ curl -X POST 'http://127.0.0.1:1337/reset-tracing'
+//
+// Get the node's randomness partial signatures for round 123.
+//
+//  $ curl 'http://127.0.0.1:1337/randomness-partial-sigs?round=123'
+//
+// Inject a randomness partial signature from another node, bypassing validity checks.
+//
+//  $ curl 'http://127.0.0.1:1337/randomness-inject-partial-sigs?authority_name=hexencodedname&round=123&sigs=base64encodedsigs'
+//
+// Inject a full signature from another node, bypassing validity checks.
+//
+//  $ curl 'http://127.0.0.1:1337/randomness-inject-full-sig?round=123&sigs=base64encodedsig'
 
 const LOGGING_ROUTE: &str = "/logging";
 const TRACING_ROUTE: &str = "/enable-tracing";
@@ -56,6 +77,9 @@ const CLEAR_BUFFER_STAKE_ROUTE: &str = "/clear-override-buffer-stake";
 const FORCE_CLOSE_EPOCH: &str = "/force-close-epoch";
 const CAPABILITIES: &str = "/capabilities";
 const NODE_CONFIG: &str = "/node-config";
+const RANDOMNESS_PARTIAL_SIGS_ROUTE: &str = "/randomness-partial-sigs";
+const RANDOMNESS_INJECT_PARTIAL_SIGS_ROUTE: &str = "/randomness-inject-partial-sigs";
+const RANDOMNESS_INJECT_FULL_SIG_ROUTE: &str = "/randomness-inject-full-sig";
 
 struct AppState {
     node: Arc<SuiNode>,
@@ -86,6 +110,15 @@ pub async fn run_admin_server(node: Arc<SuiNode>, port: u16, tracing_handle: Tra
         .route(FORCE_CLOSE_EPOCH, post(force_close_epoch))
         .route(TRACING_ROUTE, post(enable_tracing))
         .route(TRACING_RESET_ROUTE, post(reset_tracing))
+        .route(RANDOMNESS_PARTIAL_SIGS_ROUTE, get(randomness_partial_sigs))
+        .route(
+            RANDOMNESS_INJECT_PARTIAL_SIGS_ROUTE,
+            post(randomness_inject_partial_sigs),
+        )
+        .route(
+            RANDOMNESS_INJECT_FULL_SIG_ROUTE,
+            post(randomness_inject_full_sig),
+        )
         .with_state(Arc::new(app_state));
 
     let socket_address = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port);
@@ -288,3 +321,114 @@ async fn force_close_epoch(
         Err(err) => (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()),
     }
 }
+
+#[derive(Deserialize)]
+struct Round {
+    round: u64,
+}
+
+async fn randomness_partial_sigs(
+    State(state): State<Arc<AppState>>,
+    round: Query<Round>,
+) -> (StatusCode, String) {
+    let Query(Round { round }) = round;
+
+    let (tx, rx) = oneshot::channel();
+    state
+        .node
+        .randomness_handle()
+        .admin_get_partial_signatures(RandomnessRound(round), tx);
+
+    let sigs = match rx.await {
+        Ok(sigs) => sigs,
+        Err(err) => return (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()),
+    };
+
+    let output = format!(
+        "{}\n",
+        base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(sigs)
+    );
+
+    (StatusCode::OK, output)
+}
+
+#[derive(Deserialize)]
+struct PartialSigsToInject {
+    hex_authority_name: String,
+    round: u64,
+    base64_sigs: String,
+}
+
+async fn randomness_inject_partial_sigs(
+    State(state): State<Arc<AppState>>,
+    args: Query<PartialSigsToInject>,
+) -> (StatusCode, String) {
+    let Query(PartialSigsToInject {
+        hex_authority_name,
+        round,
+        base64_sigs,
+    }) = args;
+
+    let authority_name = match AuthorityName::from_str(hex_authority_name.as_str()) {
+        Ok(authority_name) => authority_name,
+        Err(err) => return (StatusCode::BAD_REQUEST, err.to_string()),
+    };
+
+    let sigs: Vec<u8> = match base64::engine::general_purpose::URL_SAFE_NO_PAD.decode(base64_sigs) {
+        Ok(sigs) => sigs,
+        Err(err) => return (StatusCode::BAD_REQUEST, err.to_string()),
+    };
+
+    let sigs: Vec<RandomnessPartialSignature> = match bcs::from_bytes(&sigs) {
+        Ok(sigs) => sigs,
+        Err(err) => return (StatusCode::BAD_REQUEST, err.to_string()),
+    };
+
+    let (tx_result, rx_result) = oneshot::channel();
+    state
+        .node
+        .randomness_handle()
+        .admin_inject_partial_signatures(authority_name, RandomnessRound(round), sigs, tx_result);
+
+    match rx_result.await {
+        Ok(Ok(())) => (StatusCode::OK, "partial signatures injected\n".to_string()),
+        Ok(Err(e)) => (StatusCode::BAD_REQUEST, e.to_string()),
+        Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()),
+    }
+}
+
+#[derive(Deserialize)]
+struct FullSigToInject {
+    round: u64,
+    base64_sig: String,
+}
+
+async fn randomness_inject_full_sig(
+    State(state): State<Arc<AppState>>,
+    args: Query<FullSigToInject>,
+) -> (StatusCode, String) {
+    let Query(FullSigToInject { round, base64_sig }) = args;
+
+    let sig: Vec<u8> = match base64::engine::general_purpose::URL_SAFE_NO_PAD.decode(base64_sig) {
+        Ok(sig) => sig,
+        Err(err) => return (StatusCode::BAD_REQUEST, err.to_string()),
+    };
+
+    let sig: RandomnessSignature = match bcs::from_bytes(&sig) {
+        Ok(sig) => sig,
+        Err(err) => return (StatusCode::BAD_REQUEST, err.to_string()),
+    };
+
+    let (tx_result, rx_result) = oneshot::channel();
+    state.node.randomness_handle().admin_inject_full_signature(
+        RandomnessRound(round),
+        sig,
+        tx_result,
+    );
+
+    match rx_result.await {
+        Ok(Ok(())) => (StatusCode::OK, "full signature injected\n".to_string()),
+        Ok(Err(e)) => (StatusCode::BAD_REQUEST, e.to_string()),
+        Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()),
+    }
+}