Move DKG message size limit check into SuiTxValidator (#17735)

MystenLabs · May 28, 2024 · 60aed80 · 60aed80
1 parent 7b19a35
commit 60aed80
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 32 deletions.
diff --git a/crates/sui-core/src/authority/authority_per_epoch_store.rs b/crates/sui-core/src/authority/authority_per_epoch_store.rs
@@ -3283,22 +3283,13 @@ impl AuthorityPerEpochStore {
                             "Received RandomnessDkgMessage from {:?}",
                             authority.concise()
                         );
-                        if bytes.len() > dkg::DKG_MESSAGES_MAX_SIZE {
-                            warn!(
-                                "Ignoring RandomnessDkgMessage from {:?} because it is too large",
-                                authority.concise()
-                            );
-                        } else {
-                            match bcs::from_bytes(bytes) {
-                                Ok(message) => {
-                                    randomness_manager.add_message(authority, message)?
-                                }
-                                Err(e) => {
-                                    warn!(
+                        match bcs::from_bytes(bytes) {
+                            Ok(message) => randomness_manager.add_message(authority, message)?,
+                            Err(e) => {
+                                warn!(
                                     "Failed to deserialize RandomnessDkgMessage from {:?}: {e:?}",
                                     authority.concise(),
                                 );
-                                }
                             }
                         }
                     } else {
@@ -3325,24 +3316,17 @@ impl AuthorityPerEpochStore {
                             "Received RandomnessDkgConfirmation from {:?}",
                             authority.concise()
                         );
-                        if bytes.len() > dkg::DKG_MESSAGES_MAX_SIZE {
-                            warn!(
-                                "Ignoring RandomnessDkgConfirmation from {:?} because it is too large",
-                                authority.concise()
-                            );
-                        } else {
-                            match bcs::from_bytes(bytes) {
-                                Ok(confirmation) => randomness_manager.add_confirmation(
-                                    batch,
-                                    authority,
-                                    confirmation,
-                                )?,
-                                Err(e) => {
-                                    warn!(
+                        match bcs::from_bytes(bytes) {
+                            Ok(confirmation) => randomness_manager.add_confirmation(
+                                batch,
+                                authority,
+                                confirmation,
+                            )?,
+                            Err(e) => {
+                                warn!(
                                     "Failed to deserialize RandomnessDkgMessage from {:?}: {e:?}",
                                     authority.concise(),
                                 );
-                                }
                             }
                         }
                     } else {

diff --git a/crates/sui-core/src/consensus_validator.rs b/crates/sui-core/src/consensus_validator.rs
@@ -5,12 +5,16 @@ use std::sync::Arc;
 
 use consensus_core::{TransactionVerifier, ValidationError};
 use eyre::WrapErr;
+use fastcrypto_tbls::dkg;
 use mysten_metrics::monitored_scope;
 use narwhal_types::{validate_batch_version, BatchAPI};
 use narwhal_worker::TransactionValidator;
 use prometheus::{register_int_counter_with_registry, IntCounter, Registry};
 use sui_protocol_config::ProtocolConfig;
-use sui_types::messages_consensus::{ConsensusTransaction, ConsensusTransactionKind};
+use sui_types::{
+    error::SuiError,
+    messages_consensus::{ConsensusTransaction, ConsensusTransactionKind},
+};
 use tap::TapFallible;
 use tracing::{info, warn};
 
@@ -69,12 +73,22 @@ impl SuiTxValidator {
                     ckpt_messages.push(signature.clone());
                     ckpt_batch.push(signature.summary);
                 }
+                ConsensusTransactionKind::RandomnessDkgMessage(_, bytes) => {
+                    if bytes.len() > dkg::DKG_MESSAGES_MAX_SIZE {
+                        warn!("batch verification error: DKG Message too large");
+                        return Err(SuiError::InvalidDkgMessageSize.into());
+                    }
+                }
+                ConsensusTransactionKind::RandomnessDkgConfirmation(_, bytes) => {
+                    if bytes.len() > dkg::DKG_MESSAGES_MAX_SIZE {
+                        warn!("batch verification error: DKG Confirmation too large");
+                        return Err(SuiError::InvalidDkgMessageSize.into());
+                    }
+                }
                 ConsensusTransactionKind::EndOfPublish(_)
                 | ConsensusTransactionKind::CapabilityNotification(_)
                 | ConsensusTransactionKind::NewJWKFetched(_, _, _)
-                | ConsensusTransactionKind::RandomnessStateUpdate(_, _)
-                | ConsensusTransactionKind::RandomnessDkgMessage(_, _)
-                | ConsensusTransactionKind::RandomnessDkgConfirmation(_, _) => {}
+                | ConsensusTransactionKind::RandomnessStateUpdate(_, _) => {}
             }
         }
 

diff --git a/crates/sui-types/src/error.rs b/crates/sui-types/src/error.rs
@@ -407,6 +407,8 @@ pub enum SuiError {
 
     #[error("Invalid digest length. Expected {expected}, got {actual}")]
     InvalidDigestLength { expected: usize, actual: usize },
+    #[error("Invalid DKG message size")]
+    InvalidDkgMessageSize,
 
     #[error("Unexpected message.")]
     UnexpectedMessage,