refactor: forester: (#1271)

1. disable timeout checks
2. add derivation pubkey

forester-utils/src/forester_epoch.rs

@@ -42,8 +42,13 @@ impl Forester {
         &mut self,
         rpc: &mut impl RpcConnection,
         forester_keypair: &Keypair,
+        derivation: &Pubkey,
     ) -> Result<Signature, RpcError> {
-        let ix = create_report_work_instruction(&forester_keypair.pubkey(), self.report_work.epoch);
+        let ix = create_report_work_instruction(
+            &forester_keypair.pubkey(),
+            derivation,
+            self.report_work.epoch,
+        );
         rpc.create_and_send_transaction(&[ix], &forester_keypair.pubkey(), &[forester_keypair])
             .await
     }
@@ -322,6 +327,7 @@ impl Epoch {
         rpc: &mut R,
         protocol_config: &ProtocolConfig,
         authority: &Keypair,
+        derivation: &Pubkey,
     ) -> Result<Option<Epoch>, RpcError> {
         let epoch_registration =
             Self::slots_until_next_epoch_registration(rpc, protocol_config).await?;
@@ -333,6 +339,7 @@ impl Epoch {
 
         let instruction = create_register_forester_epoch_pda_instruction(
             &authority.pubkey(),
+            derivation,
             epoch_registration.epoch,
         );
         let signature = rpc
@@ -345,7 +352,7 @@ impl Epoch {
             .await?
             .unwrap();
         let forester_epoch_pda_pubkey =
-            get_forester_epoch_pda_from_authority(&authority.pubkey(), epoch_registration.epoch).0;
+            get_forester_epoch_pda_from_authority(derivation, epoch_registration.epoch).0;
 
         let phases = get_epoch_phases(protocol_config, epoch_pda.epoch);
         Ok(Some(Self {

forester-utils/src/registry.rs

@@ -181,6 +181,7 @@ pub async fn get_rent_exemption_for_state_merkle_tree_and_queue<R: RpcConnection
 pub async fn create_rollover_address_merkle_tree_instructions<R: RpcConnection>(
     rpc: &mut R,
     authority: &Pubkey,
+    derivation: &Pubkey,
     new_nullifier_queue_keypair: &Keypair,
     new_address_merkle_tree_keypair: &Keypair,
     merkle_tree_pubkey: &Pubkey,
@@ -220,6 +221,7 @@ pub async fn create_rollover_address_merkle_tree_instructions<R: RpcConnection>(
     let instruction = light_registry::account_compression_cpi::sdk::create_rollover_address_merkle_tree_instruction(
         CreateRolloverMerkleTreeInstructionInputs {
             authority: *authority,
+            derivation: *derivation,
             new_queue: new_nullifier_queue_keypair.pubkey(),
             new_merkle_tree: new_address_merkle_tree_keypair.pubkey(),
             old_queue: *nullifier_queue_pubkey,
@@ -239,6 +241,7 @@ pub async fn create_rollover_address_merkle_tree_instructions<R: RpcConnection>(
 pub async fn perform_state_merkle_tree_roll_over<R: RpcConnection>(
     rpc: &mut R,
     authority: &Keypair,
+    derivation: &Pubkey,
     new_nullifier_queue_keypair: &Keypair,
     new_state_merkle_tree_keypair: &Keypair,
     merkle_tree_pubkey: &Pubkey,
@@ -249,6 +252,7 @@ pub async fn perform_state_merkle_tree_roll_over<R: RpcConnection>(
     let instructions = create_rollover_address_merkle_tree_instructions(
         rpc,
         &authority.pubkey(),
+        derivation,
         new_nullifier_queue_keypair,
         new_state_merkle_tree_keypair,
         merkle_tree_pubkey,
@@ -273,6 +277,7 @@ pub async fn perform_state_merkle_tree_roll_over<R: RpcConnection>(
 pub async fn create_rollover_state_merkle_tree_instructions<R: RpcConnection>(
     rpc: &mut R,
     authority: &Pubkey,
+    derivation: &Pubkey,
     new_nullifier_queue_keypair: &Keypair,
     new_state_merkle_tree_keypair: &Keypair,
     new_cpi_context_keypair: &Keypair,
@@ -320,6 +325,7 @@ pub async fn create_rollover_state_merkle_tree_instructions<R: RpcConnection>(
     let instruction = create_rollover_state_merkle_tree_instruction(
         CreateRolloverMerkleTreeInstructionInputs {
             authority: *authority,
+            derivation: *derivation,
             new_queue: new_nullifier_queue_keypair.pubkey(),
             new_merkle_tree: new_state_merkle_tree_keypair.pubkey(),
             old_queue: *nullifier_queue_pubkey,

forester/src/cli.rs

@@ -34,6 +34,9 @@ pub struct StartArgs {
     #[arg(long, env = "FORESTER_PAYER")]
     pub payer: Option<String>,
 
+    #[arg(long, env = "FORESTER_DERIVATION_PUBKEY")]
+    pub derivation: Option<String>,
+
     #[arg(long, env = "FORESTER_PHOTON_API_KEY")]
     pub photon_api_key: Option<String>,
 

forester/src/config.rs

@@ -20,6 +20,7 @@ pub struct ForesterConfig {
     pub general_config: GeneralConfig,
     pub registry_pubkey: Pubkey,
     pub payer_keypair: Keypair,
+    pub derivation_pubkey: Pubkey,
     pub address_tree_data: Vec<TreeAccounts>,
     pub state_tree_data: Vec<TreeAccounts>,
 }
@@ -104,6 +105,20 @@ impl ForesterConfig {
         let payer =
             Keypair::from_bytes(&payer).map_err(|e| ForesterError::ConfigError(e.to_string()))?;
 
+        let derivation: Vec<u8> = match &args.derivation {
+            Some(derivation_str) => serde_json::from_str(derivation_str)
+                .map_err(|e| ForesterError::ConfigError(e.to_string()))?,
+            None => {
+                return Err(ForesterError::ConfigError(
+                    "Derivation is required".to_string(),
+                ))
+            }
+        };
+        let derivation_array: [u8; 32] = derivation
+            .try_into()
+            .map_err(|_| ForesterError::ConfigError("Derivation must be 32 bytes".to_string()))?;
+        let derivation = Pubkey::from(derivation_array);
+
         let rpc_url = args
             .rpc_url
             .clone()
@@ -147,6 +162,7 @@ impl ForesterConfig {
             registry_pubkey: Pubkey::from_str(&registry_pubkey)
                 .map_err(|e| ForesterError::ConfigError(e.to_string()))?,
             payer_keypair: payer,
+            derivation_pubkey: derivation,
             address_tree_data: vec![],
             state_tree_data: vec![],
         })
@@ -179,6 +195,7 @@ impl ForesterConfig {
             },
             registry_pubkey: Pubkey::default(),
             payer_keypair: Keypair::new(),
+            derivation_pubkey: Pubkey::default(),
             address_tree_data: vec![],
             state_tree_data: vec![],
         })
@@ -195,6 +212,7 @@ impl Clone for ForesterConfig {
             general_config: self.general_config.clone(),
             registry_pubkey: self.registry_pubkey,
             payer_keypair: self.payer_keypair.insecure_clone(),
+            derivation_pubkey: self.derivation_pubkey,
             address_tree_data: self.address_tree_data.clone(),
             state_tree_data: self.state_tree_data.clone(),
         }

forester/src/epoch_manager.rs

@@ -231,7 +231,7 @@ impl<R: RpcConnection, I: Indexer<R>> EpochManager<R, I> {
         // Check if we're currently in the active phase
         if current_slot >= phases.active.start && current_slot < phases.active.end {
             info!("Currently in active phase. Attempting to process the new tree immediately.");
-            info!("Recovering regitration info...");
+            info!("Recovering registration info...");
             if let Ok(mut epoch_info) = self.recover_registration_info(current_epoch).await {
                 info!("Recovered registration info for current epoch");
                 let tree_schedule = TreeForesterSchedule::new_with_schedule(
@@ -335,8 +335,9 @@ impl<R: RpcConnection, I: Indexer<R>> EpochManager<R, I> {
 
     async fn recover_registration_info(&self, epoch: u64) -> Result<ForesterEpochInfo> {
         debug!("Recovering registration info for epoch {}", epoch);
+
         let forester_epoch_pda_pubkey =
-            get_forester_epoch_pda_from_authority(&self.config.payer_keypair.pubkey(), epoch).0;
+            get_forester_epoch_pda_from_authority(&self.config.derivation_pubkey, epoch).0;
         let mut rpc = self.rpc_pool.get_connection().await?;
         let existing_pda = rpc
             .get_anchor_account::<ForesterEpochPda>(&forester_epoch_pda_pubkey)
@@ -497,7 +498,7 @@ impl<R: RpcConnection, I: Indexer<R>> EpochManager<R, I> {
 
         if slot < phases.registration.end {
             let forester_epoch_pda_pubkey =
-                get_forester_epoch_pda_from_authority(&self.config.payer_keypair.pubkey(), epoch).0;
+                get_forester_epoch_pda_from_authority(&self.config.derivation_pubkey, epoch).0;
             let existing_registration = rpc
                 .get_anchor_account::<ForesterEpochPda>(&forester_epoch_pda_pubkey)
                 .await?;
@@ -523,6 +524,7 @@ impl<R: RpcConnection, I: Indexer<R>> EpochManager<R, I> {
                     &mut rpc,
                     &self.protocol_config,
                     &self.config.payer_keypair,
+                    &self.config.derivation_pubkey,
                 )
                 .await
                 {
@@ -649,14 +651,13 @@ impl<R: RpcConnection, I: Indexer<R>> EpochManager<R, I> {
         epoch_info: &ForesterEpochInfo,
     ) -> Result<ForesterEpochInfo> {
         info!("Waiting for active phase");
-
         let mut rpc = self.rpc_pool.get_connection().await?;
 
         let active_phase_start_slot = epoch_info.epoch.phases.active.start;
         wait_until_slot_reached(&mut *rpc, &self.slot_tracker, active_phase_start_slot).await?;
 
         let forester_epoch_pda_pubkey = get_forester_epoch_pda_from_authority(
-            &self.config.payer_keypair.pubkey(),
+            &self.config.derivation_pubkey,
             epoch_info.epoch.epoch,
         )
         .0;
@@ -669,6 +670,7 @@ impl<R: RpcConnection, I: Indexer<R>> EpochManager<R, I> {
                 // TODO: we can put this ix into every tx of the first batch of the current active phase
                 let ix = create_finalize_registration_instruction(
                     &self.config.payer_keypair.pubkey(),
+                    &self.config.derivation_pubkey,
                     epoch_info.epoch.epoch,
                 );
                 rpc.create_and_send_transaction(
@@ -853,6 +855,7 @@ impl<R: RpcConnection, I: Indexer<R>> EpochManager<R, I> {
                 let start_time = Instant::now();
                 let batch_tx_future = send_batched_transactions(
                     &self.config.payer_keypair,
+                    &self.config.derivation_pubkey,
                     self.rpc_pool.clone(),
                     &batched_tx_config, // TODO: define config in epoch manager
                     tree.tree_accounts,
@@ -942,7 +945,7 @@ impl<R: RpcConnection, I: Indexer<R>> EpochManager<R, I> {
             SolanaRpcConnection::new(self.config.external_services.rpc_url.as_str(), None);
 
         let forester_epoch_pda_pubkey = get_forester_epoch_pda_from_authority(
-            &self.config.payer_keypair.pubkey(),
+            &self.config.derivation_pubkey,
             epoch_info.epoch.epoch,
         )
         .0;
@@ -962,6 +965,7 @@ impl<R: RpcConnection, I: Indexer<R>> EpochManager<R, I> {
 
         let ix = create_report_work_instruction(
             &self.config.payer_keypair.pubkey(),
+            &self.config.derivation_pubkey,
             epoch_info.epoch.epoch,
         );
         match rpc
@@ -1080,7 +1084,7 @@ pub async fn run_service<R: RpcConnection, I: Indexer<R>>(
 
             let trees = {
                 let rpc = rpc_pool.get_connection().await?;
-                fetch_trees(&*rpc).await
+                fetch_trees(&*rpc).await?
             };
             info!("Fetched initial trees: {:?}", trees);
 

forester/src/main.rs

@@ -63,7 +63,7 @@ async fn main() -> Result<(), ForesterError> {
             debug!("Fetching trees...");
             debug!("RPC URL: {}", config.external_services.rpc_url);
             let rpc = SolanaRpcConnection::new(config.external_services.rpc_url.clone(), None);
-            let trees = fetch_trees(&rpc).await;
+            let trees = fetch_trees(&rpc).await?;
             if trees.is_empty() {
                 warn!("No trees found. Exiting.");
             }

forester/src/rollover/operations.rs

@@ -130,6 +130,7 @@ pub async fn rollover_state_merkle_tree<R: RpcConnection, I: Indexer<R>>(
 
     let rollover_signature = perform_state_merkle_tree_rollover_forester(
         &config.payer_keypair,
+        &config.derivation_pubkey,
         rpc,
         &new_nullifier_queue_keypair,
         &new_merkle_tree_keypair,
@@ -161,6 +162,7 @@ pub async fn rollover_state_merkle_tree<R: RpcConnection, I: Indexer<R>>(
 #[allow(clippy::too_many_arguments)]
 pub async fn perform_state_merkle_tree_rollover_forester<R: RpcConnection>(
     payer: &Keypair,
+    derivation: &Pubkey,
     context: &mut R,
     new_queue_keypair: &Keypair,
     new_address_merkle_tree_keypair: &Keypair,
@@ -172,6 +174,7 @@ pub async fn perform_state_merkle_tree_rollover_forester<R: RpcConnection>(
     let instructions = create_rollover_state_merkle_tree_instructions(
         context,
         &payer.pubkey(),
+        derivation,
         new_queue_keypair,
         new_address_merkle_tree_keypair,
         new_cpi_context_keypair,
@@ -205,6 +208,7 @@ pub async fn rollover_address_merkle_tree<R: RpcConnection, I: Indexer<R>>(
     let new_merkle_tree_keypair = Keypair::new();
     let rollover_signature = perform_address_merkle_tree_rollover(
         &config.payer_keypair,
+        &config.derivation_pubkey,
         rpc,
         &new_nullifier_queue_keypair,
         &new_merkle_tree_keypair,
@@ -224,6 +228,7 @@ pub async fn rollover_address_merkle_tree<R: RpcConnection, I: Indexer<R>>(
 
 pub async fn perform_address_merkle_tree_rollover<R: RpcConnection>(
     payer: &Keypair,
+    derivation: &Pubkey,
     context: &mut R,
     new_queue_keypair: &Keypair,
     new_address_merkle_tree_keypair: &Keypair,
@@ -233,6 +238,7 @@ pub async fn perform_address_merkle_tree_rollover<R: RpcConnection>(
     let instructions = create_rollover_address_merkle_tree_instructions(
         context,
         &payer.pubkey(),
+        derivation,
         new_queue_keypair,
         new_address_merkle_tree_keypair,
         old_merkle_tree_pubkey,
@@ -252,6 +258,7 @@ pub async fn perform_address_merkle_tree_rollover<R: RpcConnection>(
 pub async fn create_rollover_address_merkle_tree_instructions<R: RpcConnection>(
     rpc: &mut R,
     authority: &Pubkey,
+    derivation: &Pubkey,
     new_nullifier_queue_keypair: &Keypair,
     new_address_merkle_tree_keypair: &Keypair,
     merkle_tree_pubkey: &Pubkey,
@@ -290,6 +297,7 @@ pub async fn create_rollover_address_merkle_tree_instructions<R: RpcConnection>(
     let instruction = create_rollover_address_merkle_tree_instruction(
         CreateRolloverMerkleTreeInstructionInputs {
             authority: *authority,
+            derivation: *derivation,
             new_queue: new_nullifier_queue_keypair.pubkey(),
             new_merkle_tree: new_address_merkle_tree_keypair.pubkey(),
             old_queue: *nullifier_queue_pubkey,
@@ -310,6 +318,7 @@ pub async fn create_rollover_address_merkle_tree_instructions<R: RpcConnection>(
 pub async fn create_rollover_state_merkle_tree_instructions<R: RpcConnection>(
     rpc: &mut R,
     authority: &Pubkey,
+    derivation: &Pubkey,
     new_nullifier_queue_keypair: &Keypair,
     new_state_merkle_tree_keypair: &Keypair,
     new_cpi_context_keypair: &Keypair,
@@ -359,6 +368,7 @@ pub async fn create_rollover_state_merkle_tree_instructions<R: RpcConnection>(
     let instruction = create_rollover_state_merkle_tree_instruction(
         CreateRolloverMerkleTreeInstructionInputs {
             authority: *authority,
+            derivation: *derivation,
             new_queue: new_nullifier_queue_keypair.pubkey(),
             new_merkle_tree: new_state_merkle_tree_keypair.pubkey(),
             old_queue: *nullifier_queue_pubkey,