Refactor signed batch (MystenLabs#3299)

crates/sui-core/src/authority.rs

@@ -80,6 +80,7 @@ mod authority_store;
 pub use authority_store::{
     AuthorityStore, GatewayStore, ResolverWrapper, SuiDataStore, UpdateType,
 };
+use sui_types::committee::EpochId;
 use sui_types::messages_checkpoint::{
     CheckpointRequest, CheckpointRequestType, CheckpointResponse,
 };
@@ -312,6 +313,10 @@ impl AuthorityState {
         self.batch_channels.subscribe()
     }
 
+    pub fn epoch(&self) -> EpochId {
+        self.committee.load().epoch
+    }
+
     async fn handle_transaction_impl(
         &self,
         transaction: Transaction,
@@ -344,12 +349,8 @@ impl AuthorityState {
 
         let owned_objects = input_objects.filter_owned_objects();
 
-        let signed_transaction = SignedTransaction::new(
-            self.committee.load().epoch,
-            transaction,
-            self.name,
-            &*self.secret,
-        );
+        let signed_transaction =
+            SignedTransaction::new(self.epoch(), transaction, self.name, &*self.secret);
 
         // Check and write locks, to signed transaction, into the database
         // The call to self.set_transaction_lock checks the lock is not conflicting,
@@ -642,7 +643,7 @@ impl AuthorityState {
             &self.move_vm,
             &self._native_functions,
             gas_status,
-            self.committee.load().epoch,
+            self.epoch(),
         );
 
         self.metrics.total_effects.inc();
@@ -651,8 +652,7 @@ impl AuthorityState {
             .inc_by(effects.events.len() as u64);
 
         // TODO: Distribute gas charge and rebate, which can be retrieved from effects.
-        let signed_effects =
-            effects.to_sign_effects(self.committee.load().epoch, &self.name, &*self.secret);
+        let signed_effects = effects.to_sign_effects(self.epoch(), &self.name, &*self.secret);
 
         Ok((temporary_store, signed_effects))
     }

crates/sui-core/src/authority_active/gossip/configurable_batch_action_client.rs

@@ -149,7 +149,12 @@ impl AuthorityAPI for ConfigurableBatchActionClient {
         let name = self.state.name;
         let mut items: Vec<Result<BatchInfoResponseItem, SuiError>> = Vec::new();
         let mut seq = 0;
-        let zero_batch = SignedBatch::new(AuthorityBatch::initial(), &*secret, name);
+        let zero_batch = SignedBatch::new(
+            self.state.epoch(),
+            AuthorityBatch::initial(),
+            &*secret,
+            name,
+        );
         items.push(Ok(BatchInfoResponseItem(UpdateItem::Batch(zero_batch))));
         let _ = actions.iter().for_each(|action| {
             match action {
@@ -167,7 +172,12 @@ impl AuthorityAPI for ConfigurableBatchActionClient {
                     let new_batch = AuthorityBatch::make_next(&last_batch, &transactions).unwrap();
                     last_batch = new_batch;
                     items.push({
-                        let item = SignedBatch::new(last_batch.clone(), &*secret, name);
+                        let item = SignedBatch::new(
+                            self.state.epoch(),
+                            last_batch.clone(),
+                            &*secret,
+                            name,
+                        );
                         Ok(BatchInfoResponseItem(UpdateItem::Batch(item)))
                     });
                 }

crates/sui-core/src/authority_batch.rs

@@ -86,8 +86,12 @@ impl crate::authority::AuthorityState {
             Some((_, last_batch)) => last_batch.batch,
             None => {
                 // Make a batch at zero
-                let zero_batch =
-                    SignedBatch::new(AuthorityBatch::initial(), &*self.secret, self.name);
+                let zero_batch = SignedBatch::new(
+                    self.epoch(),
+                    AuthorityBatch::initial(),
+                    &*self.secret,
+                    self.name,
+                );
                 self.db().batches.insert(&0, &zero_batch)?;
                 zero_batch.batch
             }
@@ -104,6 +108,7 @@ impl crate::authority::AuthorityState {
         if !transactions.is_empty() {
             // Make a new batch, to put the old transactions not in a batch in.
             let last_signed_batch = SignedBatch::new(
+                self.epoch(),
                 // Unwrap safe due to check not empty
                 AuthorityBatch::make_next(&last_batch, &transactions)?,
                 &*self.secret,
@@ -190,6 +195,7 @@ impl crate::authority::AuthorityState {
 
                 // Make and store a new batch.
                 let new_batch = SignedBatch::new(
+                    self.epoch(),
                     // Unwrap safe since we tested above it is not empty
                     AuthorityBatch::make_next(&prev_batch, &current_batch).unwrap(),
                     &*self.secret,

crates/sui-core/src/epoch/tests/reconfiguration_tests.rs

@@ -139,7 +139,7 @@ async fn test_start_epoch_change() {
         &state.move_vm,
         &state._native_functions,
         SuiGasStatus::new_with_budget(1000, 1, 1),
-        state.committee.load().epoch,
+        state.epoch(),
     );
     let signed_effects = effects.to_sign_effects(0, &state.name, &*state.secret);
     assert_eq!(
@@ -202,7 +202,7 @@ async fn test_finish_epoch_change() {
 
     // Verify that epoch changed in every authority state.
     for active in actives {
-        assert_eq!(active.state.committee.load().epoch, 1);
+        assert_eq!(active.state.epoch(), 1);
         assert_eq!(active.net.load().committee.epoch, 1);
         assert_eq!(
             active

crates/sui-core/src/safe_client.rs

@@ -198,9 +198,7 @@ impl<C> SafeClient<C> {
         )>,
     ) -> SuiResult {
         // check the signature of the batch
-        signed_batch
-            .signature
-            .verify(&signed_batch.batch, signed_batch.authority)?;
+        signed_batch.verify(&self.committee)?;
 
         // ensure transactions enclosed match requested range
 

crates/sui-core/src/unit_tests/batch_tests.rs

@@ -568,7 +568,7 @@ impl AuthorityAPI for TrustworthyAuthorityClient {
         let mut items = Vec::new();
         let mut last_batch = AuthorityBatch::initial();
         items.push({
-            let item = SignedBatch::new(last_batch.clone(), &*secret, name);
+            let item = SignedBatch::new_with_zero_epoch(last_batch.clone(), &*secret, name);
             BatchInfoResponseItem(UpdateItem::Batch(item))
         });
         let mut seq = 0;
@@ -584,7 +584,7 @@ impl AuthorityAPI for TrustworthyAuthorityClient {
             let new_batch = AuthorityBatch::make_next(&last_batch, &transactions).unwrap();
             last_batch = new_batch;
             items.push({
-                let item = SignedBatch::new(last_batch.clone(), &*secret, name);
+                let item = SignedBatch::new_with_zero_epoch(last_batch.clone(), &*secret, name);
                 BatchInfoResponseItem(UpdateItem::Batch(item))
             });
         }
@@ -684,7 +684,7 @@ impl AuthorityAPI for ByzantineAuthorityClient {
         let mut items = Vec::new();
         let mut last_batch = AuthorityBatch::initial();
         items.push({
-            let item = SignedBatch::new(last_batch.clone(), &*secret, name);
+            let item = SignedBatch::new_with_zero_epoch(last_batch.clone(), &*secret, name);
             BatchInfoResponseItem(UpdateItem::Batch(item))
         });
         let mut seq = 0;
@@ -706,7 +706,7 @@ impl AuthorityAPI for ByzantineAuthorityClient {
             let new_batch = AuthorityBatch::make_next(&last_batch, &transactions).unwrap();
             last_batch = new_batch;
             items.push({
-                let item = SignedBatch::new(last_batch.clone(), &*secret, name);
+                let item = SignedBatch::new_with_zero_epoch(last_batch.clone(), &*secret, name);
                 BatchInfoResponseItem(UpdateItem::Batch(item))
             });
         }

crates/sui-core/tests/staged/sui.yaml

@@ -26,6 +26,13 @@ AuthorityBatch:
         TUPLEARRAY:
           CONTENT: U8
           SIZE: 32
+AuthoritySignInfo:
+  STRUCT:
+    - epoch: U64
+    - authority:
+        TYPENAME: PublicKeyBytes
+    - signature:
+        TYPENAME: AuthoritySignature
 AuthoritySignature:
   NEWTYPESTRUCT:
     TUPLEARRAY:
@@ -320,10 +327,8 @@ SignedBatch:
   STRUCT:
     - batch:
         TYPENAME: AuthorityBatch
-    - authority:
-        TYPENAME: PublicKeyBytes
-    - signature:
-        TYPENAME: AuthoritySignature
+    - auth_signature:
+        TYPENAME: AuthoritySignInfo
 SingleTransactionKind:
   ENUM:
     0:

crates/sui-types/src/batch.rs

@@ -2,8 +2,9 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::base_types::{AuthorityName, ExecutionDigests};
-use crate::crypto::{sha3_hash, AuthoritySignature, BcsSignable};
-use crate::error::SuiError;
+use crate::committee::{Committee, EpochId};
+use crate::crypto::{sha3_hash, AuthoritySignInfo, AuthoritySignature, BcsSignable};
+use crate::error::{SuiError, SuiResult};
 use serde::{Deserialize, Serialize};
 
 pub type TxSequenceNumber = u64;
@@ -98,29 +99,39 @@ impl AuthorityBatch {
 }
 
 /// An transaction signed by a single authority
-#[derive(Eq, Clone, Debug, Serialize, Deserialize)]
+#[derive(Eq, PartialEq, Clone, Debug, Serialize, Deserialize)]
 pub struct SignedBatch {
     pub batch: AuthorityBatch,
-    pub authority: AuthorityName,
-    pub signature: AuthoritySignature,
+    pub auth_signature: AuthoritySignInfo,
 }
 
 impl SignedBatch {
     pub fn new(
+        epoch: EpochId,
         batch: AuthorityBatch,
         secret: &dyn signature::Signer<AuthoritySignature>,
         authority: AuthorityName,
     ) -> SignedBatch {
+        let signature = AuthoritySignature::new(&batch, secret);
         SignedBatch {
-            signature: AuthoritySignature::new(&batch, secret),
             batch,
-            authority,
+            auth_signature: AuthoritySignInfo {
+                epoch,
+                authority,
+                signature,
+            },
         }
     }
-}
 
-impl PartialEq for SignedBatch {
-    fn eq(&self, other: &Self) -> bool {
-        self.batch == other.batch && self.authority == other.authority
+    pub fn new_with_zero_epoch(
+        batch: AuthorityBatch,
+        secret: &dyn signature::Signer<AuthoritySignature>,
+        authority: AuthorityName,
+    ) -> SignedBatch {
+        Self::new(0, batch, secret, authority)
+    }
+
+    pub fn verify(&self, committee: &Committee) -> SuiResult {
+        self.auth_signature.verify(&self.batch, committee)
     }
 }