Add epoch query API to authorityAPI (MystenLabs#3725)

* Add epoch query API to authorityAPI

* Address feedback

crates/sui-core/src/authority.rs

@@ -999,6 +999,14 @@ impl AuthorityState {
         }
     }
 
+    pub fn handle_epoch_request(&self, request: &EpochRequest) -> SuiResult<EpochResponse> {
+        let epoch_info = match &request.epoch_id {
+            Some(id) => self.database.get_authenticated_epoch(id)?,
+            None => Some(self.database.get_latest_authenticated_epoch()),
+        };
+        Ok(EpochResponse { epoch_info })
+    }
+
     // TODO: This function takes both committee and genesis as parameter.
     // Technically genesis already contains committee information. Could consider merging them.
     pub async fn new(

crates/sui-core/src/authority/authority_store.rs

@@ -1354,6 +1354,13 @@ impl<S: Eq + Debug + Serialize + for<'de> Deserialize<'de>> SuiDataStore<S> {
         Ok(())
     }
 
+    pub fn get_authenticated_epoch(
+        &self,
+        epoch_id: &EpochId,
+    ) -> SuiResult<Option<AuthenticatedEpoch>> {
+        Ok(self.tables.epochs.get(epoch_id)?)
+    }
+
     pub fn get_latest_authenticated_epoch(&self) -> AuthenticatedEpoch {
         self.tables
             .epochs

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

@@ -22,8 +22,8 @@ use sui_types::crypto::{get_key_pair, AuthorityKeyPair, AuthorityPublicKeyBytes}
 use sui_types::error::SuiError;
 use sui_types::messages::{
     AccountInfoRequest, AccountInfoResponse, BatchInfoRequest, BatchInfoResponseItem,
-    CertifiedTransaction, ObjectInfoRequest, ObjectInfoResponse, Transaction,
-    TransactionInfoRequest, TransactionInfoResponse,
+    CertifiedTransaction, EpochRequest, EpochResponse, ObjectInfoRequest, ObjectInfoResponse,
+    Transaction, TransactionInfoRequest, TransactionInfoResponse,
 };
 use sui_types::messages_checkpoint::{CheckpointRequest, CheckpointResponse};
 use sui_types::object::Object;
@@ -201,9 +201,15 @@ impl AuthorityAPI for ConfigurableBatchActionClient {
 
     async fn handle_checkpoint(
         &self,
-        _request: CheckpointRequest,
+        request: CheckpointRequest,
     ) -> Result<CheckpointResponse, SuiError> {
-        todo!();
+        let state = self.state.clone();
+        state.handle_checkpoint_request(&request)
+    }
+
+    async fn handle_epoch(&self, request: EpochRequest) -> Result<EpochResponse, SuiError> {
+        let state = self.state.clone();
+        state.handle_epoch_request(&request)
     }
 }
 

crates/sui-core/src/authority_client.rs

@@ -67,6 +67,8 @@ pub trait AuthorityAPI {
         &self,
         request: CheckpointRequest,
     ) -> Result<CheckpointResponse, SuiError>;
+
+    async fn handle_epoch(&self, request: EpochRequest) -> Result<EpochResponse, SuiError>;
 }
 
 pub type BatchInfoResponseItemStream = BoxStream<'static, Result<BatchInfoResponseItem, SuiError>>;
@@ -195,6 +197,14 @@ impl AuthorityAPI for NetworkAuthorityClient {
             .map(tonic::Response::into_inner)
             .map_err(Into::into)
     }
+
+    async fn handle_epoch(&self, request: EpochRequest) -> Result<EpochResponse, SuiError> {
+        self.client()
+            .epoch_info(request)
+            .await
+            .map(tonic::Response::into_inner)
+            .map_err(Into::into)
+    }
 }
 
 #[derive(Clone, Copy, Default)]
@@ -309,6 +319,12 @@ impl AuthorityAPI for LocalAuthorityClient {
 
         state.handle_checkpoint_request(&request)
     }
+
+    async fn handle_epoch(&self, request: EpochRequest) -> Result<EpochResponse, SuiError> {
+        let state = self.state.clone();
+
+        state.handle_epoch_request(&request)
+    }
 }
 
 impl LocalAuthorityClient {

crates/sui-core/src/authority_server.rs

@@ -438,4 +438,18 @@ impl Validator for ValidatorService {
 
         return Ok(tonic::Response::new(response));
     }
+
+    async fn epoch_info(
+        &self,
+        request: tonic::Request<EpochRequest>,
+    ) -> Result<tonic::Response<EpochResponse>, tonic::Status> {
+        let request = request.into_inner();
+
+        let response = self
+            .state
+            .handle_epoch_request(&request)
+            .map_err(|e| tonic::Status::internal(e.to_string()))?;
+
+        return Ok(tonic::Response::new(response));
+    }
 }

crates/sui-core/src/safe_client.rs

@@ -15,6 +15,7 @@ use sui_types::{
     error::{SuiError, SuiResult},
     messages::*,
 };
+use tap::TapFallible;
 use tracing::info;
 
 #[derive(Clone)]
@@ -264,6 +265,7 @@ impl<C> SafeClient<C> {
 
     /// This function is used by the higher level authority logic to report an
     /// error that could be due to this authority.
+    /// TODO: Get rid of this. https://github.com/MystenLabs/sui/issues/3740
     pub fn report_client_error(&self, error: &SuiError) {
         info!(?error, authority =? self.address, "Client error");
     }
@@ -581,4 +583,56 @@ where
         ));
         Ok(Box::pin(stream))
     }
+
+    fn verify_epoch(
+        &self,
+        requested_epoch_id: Option<EpochId>,
+        response: &EpochResponse,
+    ) -> SuiResult {
+        if let Some(epoch) = &response.epoch_info {
+            fp_ensure!(
+                requested_epoch_id.is_none() || requested_epoch_id == Some(epoch.epoch()),
+                SuiError::InvalidEpochResponse("Responded epoch number mismatch".to_string())
+            );
+        }
+        match (requested_epoch_id, &response.epoch_info) {
+            (None, None) => Err(SuiError::InvalidEpochResponse(
+                "Latest epoch must not be None".to_string(),
+            )),
+            (Some(epoch_id), None) => {
+                fp_ensure!(
+                    epoch_id != 0,
+                    SuiError::InvalidEpochResponse("Genesis epoch must be available".to_string())
+                );
+                Ok(())
+            }
+            (_, Some(AuthenticatedEpoch::Genesis(_))) => {
+                // TODO: Verify the epoch data using genesis committee
+                Ok(())
+            }
+            (_, Some(AuthenticatedEpoch::Signed(_))) => {
+                // TODO: Verify the epoch data using previous committee
+                Ok(())
+            }
+            (_, Some(AuthenticatedEpoch::Certified(_))) => {
+                // TODO: Verify the epoch data using previous committee
+                Ok(())
+            }
+        }
+    }
+
+    pub async fn handle_epoch(&self, request: EpochRequest) -> Result<EpochResponse, SuiError> {
+        let epoch_id = request.epoch_id;
+        let authority = self.address;
+        let response = self.authority_client.handle_epoch(request).await?;
+        self.verify_epoch(epoch_id, &response)
+            .map_err(|err| SuiError::ByzantineAuthoritySuspicion {
+                authority,
+                reason: err.to_string(),
+            })
+            .tap_err(|err| {
+                self.report_client_error(err);
+            })?;
+        Ok(response)
+    }
 }

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

@@ -952,6 +952,10 @@ async fn test_quorum_once_with_timeout() {
         ) -> Result<CheckpointResponse, SuiError> {
             unreachable!();
         }
+
+        async fn handle_epoch(&self, _request: EpochRequest) -> Result<EpochResponse, SuiError> {
+            unreachable!()
+        }
     }
 
     let count = Arc::new(Mutex::new(0));

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

@@ -27,8 +27,8 @@ use std::fs;
 use std::sync::Arc;
 use sui_types::messages::{
     AccountInfoRequest, AccountInfoResponse, BatchInfoRequest, BatchInfoResponseItem,
-    CertifiedTransaction, ObjectInfoRequest, ObjectInfoResponse, Transaction,
-    TransactionInfoRequest, TransactionInfoResponse,
+    CertifiedTransaction, EpochRequest, EpochResponse, ObjectInfoRequest, ObjectInfoResponse,
+    Transaction, TransactionInfoRequest, TransactionInfoResponse,
 };
 use sui_types::object::Object;
 
@@ -565,6 +565,10 @@ impl AuthorityAPI for TrustworthyAuthorityClient {
         unimplemented!();
     }
 
+    async fn handle_epoch(&self, _request: EpochRequest) -> Result<EpochResponse, SuiError> {
+        unimplemented!()
+    }
+
     /// Handle Batch information requests for this authority.
     async fn handle_batch_stream(
         &self,
@@ -677,7 +681,11 @@ impl AuthorityAPI for ByzantineAuthorityClient {
         &self,
         _request: CheckpointRequest,
     ) -> Result<CheckpointResponse, SuiError> {
-        unimplemented!();
+        unimplemented!()
+    }
+
+    async fn handle_epoch(&self, _request: EpochRequest) -> Result<EpochResponse, SuiError> {
+        unimplemented!()
     }
 
     /// Handle Batch information requests for this authority.

crates/sui-network/build.rs

@@ -83,6 +83,15 @@ fn main() -> Result<()> {
                 .codec_path(codec_path)
                 .build(),
         )
+        .method(
+            Method::builder()
+                .name("epoch_info")
+                .route_name("Epoch")
+                .input_type("sui_types::messages::EpochRequest")
+                .output_type("sui_types::messages::EpochResponse")
+                .codec_path(codec_path)
+                .build(),
+        )
         .build();
 
     Builder::new()

crates/sui-types/src/error.rs

@@ -397,6 +397,12 @@ pub enum SuiError {
 
     #[error("Invalid committee composition")]
     InvalidCommittee(String),
+
+    #[error("Invalid authenticated epoch: {0}")]
+    InvalidAuthenticatedEpoch(String),
+
+    #[error("Invalid epoch request response: {0}")]
+    InvalidEpochResponse(String),
 }
 
 pub type SuiResult<T = ()> = Result<T, SuiError>;

crates/sui-types/src/messages.rs

@@ -1842,6 +1842,24 @@ impl GenesisEpoch {
             auth_sign_info: EmptySignInfo {},
         }
     }
+
+    pub fn verify(&self, genesis_committee: Committee) -> SuiResult {
+        fp_ensure!(
+            self.epoch_info.first_checkpoint == 0,
+            SuiError::InvalidAuthenticatedEpoch(
+                "Genesis epoch's first checkpoint must be 0".to_string()
+            )
+        );
+        fp_ensure!(
+            self.epoch_info.epoch() == 0,
+            SuiError::InvalidAuthenticatedEpoch("Genesis epoch must be epoch 0".to_string())
+        );
+        fp_ensure!(
+            self.epoch_info.committee == genesis_committee,
+            SuiError::InvalidAuthenticatedEpoch("Genesis epoch committee mismatch".to_string())
+        );
+        Ok(())
+    }
 }
 
 impl SignedEpoch {
@@ -1874,7 +1892,9 @@ impl SignedEpoch {
         let epoch = self.epoch_info.epoch();
         fp_ensure!(
             epoch != 0 && epoch - 1 == self.auth_sign_info.epoch,
-            SuiError::from("Epoch number in the committee inconsistent with signature")
+            SuiError::InvalidAuthenticatedEpoch(
+                "Epoch number in the committee inconsistent with signature".to_string()
+            )
         );
         self.auth_sign_info
             .verify(&self.epoch_info, prev_epoch_committee)?;
@@ -1889,7 +1909,9 @@ impl CertifiedEpoch {
         let epoch = self.epoch_info.epoch();
         fp_ensure!(
             epoch != 0 && epoch - 1 == self.auth_sign_info.epoch,
-            SuiError::from("Epoch number in the committee inconsistent with signature")
+            SuiError::InvalidAuthenticatedEpoch(
+                "Epoch number in the committee inconsistent with signature".to_string()
+            )
         );
         self.auth_sign_info.verify(&self.epoch_info, committee)?;
         Ok(())
@@ -1928,3 +1950,13 @@ impl AuthenticatedEpoch {
         }
     }
 }
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct EpochRequest {
+    pub epoch_id: Option<EpochId>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct EpochResponse {
+    pub epoch_info: Option<AuthenticatedEpoch>,
+}