feat(coprocessor): gw-listener, catchup for KMSGeneration contract

add --catchup-kms-generation-from-block

Author: rudy-6-4

coprocessor/README.md

@@ -160,6 +160,8 @@ Options:
           [default: 1]
       --error-sleep-max-secs <ERROR_SLEEP_MAX_SECS>
           [default: 10]
+      --catchup-kms-generation-from-block <BLOCK_NUMBER OR BLOCKS IN PAST>
+          [default: None]
   -h, --help
           Print help
   -V, --version

coprocessor/fhevm-engine/gw-listener/src/bin/gw_listener.rs

@@ -72,6 +72,9 @@ struct Conf {
     /// gw-listener service name in OTLP traces
     #[arg(long, default_value = "gw-listener")]
     pub service_name: String,
+
+    #[arg(long, default_value = None, help = "Can be negative from last processed block", allow_hyphen_values = true)]
+    pub catchup_kms_generation_from_block: Option<i64>,
 }
 
 fn install_signal_handlers(cancel_token: CancellationToken) -> anyhow::Result<()> {
@@ -142,7 +145,7 @@ async fn main() -> anyhow::Result<()> {
     let cancel_token = CancellationToken::new();
 
     let Some(host_chain_id) = conf.host_chain_id.or_else(chain_id_from_env) else {
-        anyhow::bail!("--chain-id or CHAIN_ID env var is missing.")
+        anyhow::bail!("--host-chain-id or CHAIN_ID env var is missing.")
     };
     let config = ConfigSettings {
         host_chain_id,
@@ -156,6 +159,7 @@ async fn main() -> anyhow::Result<()> {
         health_check_timeout: conf.health_check_timeout,
         get_logs_poll_interval: conf.get_logs_poll_interval,
         get_logs_block_batch_size: conf.get_logs_block_batch_size,
+        catchup_kms_generation_from_block: conf.catchup_kms_generation_from_block,
     };
 
     let gw_listener = GatewayListener::new(

coprocessor/fhevm-engine/gw-listener/src/gw_listener.rs

@@ -111,9 +111,13 @@ impl<P: Provider<Ethereum> + Clone + 'static, A: AwsS3Interface + Clone + 'stati
             let s = self.clone();
             let d = db_pool.clone();
             tokio::spawn(async move {
+                let mut catchup_kms_generation_from = s.conf.catchup_kms_generation_from_block;
                 let mut sleep_duration = s.conf.error_sleep_initial_secs as u64;
                 loop {
-                    match s.run_get_logs(&d, &mut sleep_duration).await {
+                    match s
+                        .run_get_logs(&d, &mut sleep_duration, &mut catchup_kms_generation_from)
+                        .await
+                    {
                         Ok(_) => {
                             info!("run_get_logs() stopped");
                             break;
@@ -182,9 +186,26 @@ impl<P: Provider<Ethereum> + Clone + 'static, A: AwsS3Interface + Clone + 'stati
         &self,
         db_pool: &Pool<Postgres>,
         sleep_duration: &mut u64,
+        catchup_kms_generation_from: &mut Option<i64>,
     ) -> anyhow::Result<()> {
         let mut ticker = tokio::time::interval(self.conf.get_logs_poll_interval);
         let mut last_processed_block_num = self.get_last_processed_block_num(db_pool).await?;
+        if let Some(from_block) = *catchup_kms_generation_from {
+            info!(from_block, "Catchup on KMSGeneration, start");
+            let from_block = if from_block >= 0 {
+                // start from specified block
+                from_block
+            } else if let Some(last) = last_processed_block_num {
+                // go N block in past from processed
+                last as i64 + from_block
+            } else {
+                // go N block in past
+                let current_block = self.provider.get_block_number().await?;
+                current_block as i64 + from_block
+            };
+            // note, we cannot catchup block 0
+            last_processed_block_num = Some(from_block.try_into().unwrap_or(0));
+        }
 
         loop {
             tokio::select! {
@@ -217,6 +238,9 @@ impl<P: Provider<Ethereum> + Clone + 'static, A: AwsS3Interface + Clone + 'stati
                         .to_block(to_block);
 
                     let logs = self.provider.get_logs(&filter).await?;
+                    if catchup_kms_generation_from.is_some() && from_block < current_block {
+                        info!(from_block, to_block, nb_events=logs.len(), "Catchup on KMSGeneration, get_logs");
+                    }
                     for log in logs {
                         if let Ok(event) = KMSGeneration::KMSGenerationEvents::decode_log(&log.inner) {
                             match event.data {
@@ -246,6 +270,16 @@ impl<P: Provider<Ethereum> + Clone + 'static, A: AwsS3Interface + Clone + 'stati
                         }
                     }
                     last_processed_block_num = Some(to_block);
+                    if catchup_kms_generation_from.is_some() {
+                        if to_block == current_block {
+                            info!("Catchup on KMSGeneration, finished");
+                            *catchup_kms_generation_from = None;
+                        } else {
+                            // if an error happens catchup will restart here
+                            *catchup_kms_generation_from = Some(to_block as i64 + 1);
+                            info!(catchup_kms_generation_from, "Catchup on KMSGeneration continue");
+                        }
+                    }
                     self.update_last_block_num(db_pool, last_processed_block_num).await?;
                     if to_block < current_block {
                         debug!(to_block = to_block,

coprocessor/fhevm-engine/gw-listener/src/lib.rs

@@ -48,6 +48,7 @@ pub struct ConfigSettings {
 
     pub get_logs_poll_interval: Duration,
     pub get_logs_block_batch_size: u64,
+    pub catchup_kms_generation_from_block: Option<i64>,
 }
 
 pub fn chain_id_from_env() -> Option<ChainId> {
@@ -74,6 +75,7 @@ impl Default for ConfigSettings {
             health_check_timeout: Duration::from_secs(4),
             get_logs_poll_interval: Duration::from_secs(1),
             get_logs_block_batch_size: 100,
+            catchup_kms_generation_from_block: None,
         }
     }
 }

coprocessor/fhevm-engine/gw-listener/tests/gw_listener_tests.rs

@@ -459,6 +459,135 @@ async fn keygen_ok() -> anyhow::Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+#[serial(db)]
+async fn keygen_ok_catchup_positive() -> anyhow::Result<()> {
+    keygen_ok_catchup_gen(true).await
+}
+
+#[tokio::test]
+#[serial(db)]
+async fn keygen_ok_catchup_negative() -> anyhow::Result<()> {
+    keygen_ok_catchup_gen(false).await
+}
+
+async fn keygen_ok_catchup_gen(positive: bool) -> anyhow::Result<()> {
+    use aws_sdk_s3::operation::get_object::GetObjectOutput;
+    use aws_sdk_s3::primitives::ByteStream;
+    use aws_sdk_s3::Client;
+    use aws_smithy_mocks::RuleMode;
+    use aws_smithy_mocks::{mock, mock_client};
+    use gw_listener::KeyType;
+
+    // see ../contracts/KMSGeneration.sol
+    let buckets = [
+        "test-bucket1/PUB-P1",
+        "test-bucket2/PUB-P2",
+        "test-bucket3/PUB-P3",
+        "test-bucket4/PUB-P4",
+    ];
+
+    let keys_digests = [KeyType::PublicKey, KeyType::ServerKey];
+
+    let key_id = U256::from(16);
+
+    let mut rules = vec![];
+    for &bucket in &buckets {
+        for key_type in &keys_digests {
+            let key_type_str: &str = to_bucket_key_prefix(*key_type);
+            let key_id_no_0x = key_id_to_key_bucket(key_id);
+            let key = format!("{}/{}", key_type_str, key_id_no_0x);
+            eprintln!("Adding {}/{}", bucket, key);
+            let get_object_rule = mock!(Client::get_object)
+                .match_requests(move |req| req.bucket() == Some(bucket) && req.key() == Some(&key))
+                .then_output(|| {
+                    GetObjectOutput::builder()
+                        .body(ByteStream::from_static(b"key_bytes"))
+                        .build()
+                });
+            rules.push(get_object_rule);
+        }
+    }
+    for &bucket in &buckets {
+        let key_id_no_0x = &format!("{key_id:064X}");
+        let key = format!("PUB/CRS/{key_id_no_0x}");
+        eprintln!("Adding {}/{}", bucket, key);
+        let get_object_rule = mock!(Client::get_object)
+            .match_requests(move |req| req.bucket() == Some(bucket) && req.key() == Some(&key))
+            .then_output(|| {
+                GetObjectOutput::builder()
+                    .body(ByteStream::from_static(b"key_bytes"))
+                    .build()
+            });
+        rules.push(get_object_rule);
+    }
+    let rules_ref: Vec<_> = rules.iter().collect();
+
+    // Create a mocked client with the rule
+    let s3 = mock_client!(aws_sdk_s3, RuleMode::MatchAny, &rules_ref);
+
+    let env = TestEnvironment::new().await?;
+    let provider = ProviderBuilder::new()
+        .wallet(env.wallet)
+        .connect_ws(WsConnect::new(env.anvil.ws_endpoint_url()))
+        .await?;
+    let aws_s3_client = AwsS3ClientMocked(s3);
+    let input_verification = InputVerification::deploy(&provider).await?;
+    let kms_generation = KMSGeneration::deploy(&provider).await?;
+
+    assert!(provider.get_block_number().await? > 0);
+
+    let txn_req = kms_generation
+        .keygen_public_key()
+        .into_transaction_request();
+    let pending_txn = provider.send_transaction(txn_req).await?;
+    let receipt = pending_txn.get_receipt().await?;
+    assert!(receipt.status());
+
+    let txn_req = kms_generation
+        .keygen_server_key()
+        .into_transaction_request();
+    let pending_txn = provider.send_transaction(txn_req).await?;
+    let receipt = pending_txn.get_receipt().await?;
+    assert!(receipt.status());
+
+    let txn_req = kms_generation.crsgen().into_transaction_request();
+    let pending_txn = provider.send_transaction(txn_req).await?;
+    let receipt = pending_txn.get_receipt().await?;
+    assert!(receipt.status());
+
+    assert!(has_not_public_key(&env.db_pool.clone()).await?);
+    assert!(has_not_server_key(&env.db_pool.clone()).await?);
+    assert!(has_not_crs(&env.db_pool.clone()).await?);
+
+    let catchup_kms_generation_from_block = if positive {
+        Some(0)
+    } else {
+        Some(-(provider.get_block_number().await? as i64))
+    };
+    let conf = ConfigSettings {
+        catchup_kms_generation_from_block,
+        ..env.conf.clone()
+    };
+    let gw_listener = GatewayListener::new(
+        *input_verification.address(),
+        *kms_generation.address(),
+        conf,
+        env.cancel_token.clone(),
+        provider.clone(),
+        aws_s3_client.clone(),
+    );
+    let listener = tokio::spawn(async move { gw_listener.run().await });
+
+    assert!(has_public_key(&env.db_pool.clone()).await?);
+    assert!(has_server_key(&env.db_pool.clone()).await?);
+    assert!(has_crs(&env.db_pool.clone()).await?);
+
+    env.cancel_token.cancel();
+    listener.abort();
+    Ok(())
+}
+
 #[tokio::test]
 #[serial(db)]
 async fn keygen_compromised_key() -> anyhow::Result<()> {