fix: prevent removal of delegated/undelegating accounts on eviction

.gitignore

@@ -19,6 +19,7 @@ target/
 # Ledger
 test-ledger/
 test-ledger-magicblock/
+magicblock-test-storage/
 
 # Mac
 **/.DS_Store

magicblock-account-cloner/src/bpf_loader_v1.rs

@@ -51,6 +51,7 @@ impl BpfUpgradableProgramModifications {
                 rent_epoch: Some(u64::MAX),
                 delegated: Some(false),
                 confined: Some(false),
+                remote_slot: Some(loaded_program.remote_slot),
             }
         };
 
@@ -71,6 +72,7 @@ impl BpfUpgradableProgramModifications {
                 rent_epoch: Some(u64::MAX),
                 delegated: Some(false),
                 confined: Some(false),
+                remote_slot: Some(loaded_program.remote_slot),
             }
         };
 

magicblock-account-cloner/src/lib.rs

@@ -100,6 +100,7 @@ impl ChainlinkCloner {
             executable: Some(request.account.executable()),
             delegated: Some(request.account.delegated()),
             confined: Some(request.account.confined()),
+            remote_slot: Some(request.account.remote_slot()),
         };
 
         let modify_ix = InstructionUtils::modify_accounts_instruction(vec![
@@ -222,6 +223,7 @@ impl ChainlinkCloner {
                 // and then deploy it and finally set the authority to match the
                 // one on chain
                 let slot = self.accounts_db.slot();
+                let program_remote_slot = program.remote_slot;
                 let DeployableV4Program {
                     pre_deploy_loader_state,
                     deploy_instruction,
@@ -248,6 +250,7 @@ impl ChainlinkCloner {
                         executable: Some(true),
                         data: Some(pre_deploy_loader_state),
                         confined: Some(false),
+                        remote_slot: Some(program_remote_slot),
                         ..Default::default()
                     }];
                     InstructionUtils::modify_accounts_instruction(
@@ -260,6 +263,7 @@ impl ChainlinkCloner {
                         pubkey: program_id,
                         data: Some(post_deploy_loader_state),
                         confined: Some(false),
+                        remote_slot: Some(program_remote_slot),
                         ..Default::default()
                     }];
                     InstructionUtils::modify_accounts_instruction(

magicblock-chainlink/src/chainlink/fetch_cloner.rs

@@ -217,7 +217,20 @@ where
                         if in_bank.undelegating() {
                             // We expect the account to still be delegated, but with the delegation
                             // program owner
-                            debug!("Received update for undelegating account {pubkey} delegated in bank={} delegated on chain={}", in_bank.delegated(), account.delegated());
+                            debug!("Received update for undelegating account {pubkey} \
+                                in_bank.delegated={}, \
+                                in_bank.owner={}, \
+                                in_bank.remote_slot={}, \
+                                chain.delegated={}, \
+                                chain.owner={}, \
+                                chain.remote_slot={}",
+                                in_bank.delegated(),
+                                in_bank.owner(),
+                                in_bank.remote_slot(),
+                                account.delegated(),
+                                account.owner(),
+                                account.remote_slot()
+                            );
 
                             // This will only be true in the following case:
                             // 1. a commit was triggered for the account
@@ -414,7 +427,7 @@ where
                                 if log::log_enabled!(log::Level::Trace) {
                                     trace!("Delegation record found for {pubkey}: {delegation_record:?}");
                                     trace!(
-                                        "Cloning delegated account: {pubkey} (remote slot {}, owner: {})",
+                                        "Resolving delegated account: {pubkey} (remote slot {}, owner: {})",
                                         account.remote_slot(),
                                         delegation_record.owner
                                     );
@@ -1262,7 +1275,16 @@ where
                     }
                     RefreshDecision::No => {
                         // Account is in bank and subscribed correctly - no fetch needed
-                        trace!("Account {pubkey} found in bank in valid state, no fetch needed");
+                        if log::log_enabled!(log::Level::Trace) {
+                            let undelegating = account_in_bank.undelegating();
+                            let delegated = account_in_bank.delegated();
+                            let owner = account_in_bank.owner().to_string();
+                            trace!("Account {pubkey} found in bank in valid state, no fetch needed \
+                                    undelegating = {undelegating}, \
+                                    delegated = {delegated}, \
+                                    owner={owner}"
+                            );
+                        }
                         in_bank.push(*pubkey);
                     }
                 }

magicblock-chainlink/src/chainlink/mod.rs

@@ -219,7 +219,37 @@ Kept: {} delegated, {} blacklisted",
 
         task::spawn(async move {
             while let Some(pubkey) = removed_accounts_rx.recv().await {
-                accounts_bank.remove_account(&pubkey);
+                accounts_bank.remove_account_conditionally(
+                    &pubkey,
+                    |account| {
+                        // Accounts that are still undelegating need to be kept in the bank
+                        // until the undelegation completes on chain.
+                        // Otherwise we might loose data in case the undelegation fails to
+                        // complete.
+                        // Another issue we avoid this way is that an account update received
+                        // before the account completes undelegation would overwrite the in-bank
+                        // account and thus also set unedelegating to false.
+                        let undelegating = account.undelegating();
+                        let delegated = account.delegated();
+                        let remove = !undelegating && !delegated;
+                        if log::log_enabled!(log::Level::Trace) {
+                            if remove {
+                                trace!(
+                                    "Removing unsubscribed account '{pubkey}' from bank"
+                                );
+                            } else {
+                                let owner = account.owner().to_string();
+                                trace!(
+                                    "Keeping unsubscribed account {pubkey} in bank \
+                                    undelegating = {undelegating}, \
+                                    delegated = {delegated}, \
+                                    owner={owner}"
+                                );
+                            }
+                        }
+                        remove
+                    },
+                );
             }
             warn!("Removed accounts channel closed, stopping subscription");
         })

magicblock-chainlink/tests/07_redeleg_us_same_slot.rs

@@ -23,7 +23,13 @@ async fn setup(slot: Slot) -> TestContext {
     TestContext::init(slot).await
 }
 
+// NOTE: disabled for now since we detect redelegation as follows:
+// - if `account.remote_slot >= undelegation_slot` assume it is the first delegation, i.e. it was
+//   not re-delegated
+// - in order to support same slot re-delegation we need an delegation index or similar to
+//   distinguish the two delegations
 #[tokio::test]
+#[ignore = "Same slot redelegation is currently not possible and currently it is detected as a single delegation operation."]
 async fn test_undelegate_redelegate_to_us_in_same_slot() {
     let mut slot: u64 = 11;
 

magicblock-core/src/traits.rs

@@ -15,4 +15,16 @@ pub trait AccountsBank: Send + Sync + 'static {
         &self,
         predicate: impl Fn(&Pubkey, &AccountSharedData) -> bool,
     ) -> usize;
+
+    fn remove_account_conditionally(
+        &self,
+        pubkey: &Pubkey,
+        predicate: impl Fn(&AccountSharedData) -> bool,
+    ) {
+        if let Some(acc) = self.get_account(pubkey) {
+            if predicate(&acc) {
+                self.remove_account(pubkey);
+            }
+        }
+    }
 }

magicblock-magic-program-api/src/instruction.rs

@@ -120,6 +120,7 @@ pub struct AccountModification {
     pub rent_epoch: Option<u64>,
     pub delegated: Option<bool>,
     pub confined: Option<bool>,
+    pub remote_slot: Option<u64>,
 }
 
 #[derive(Default, Clone, Serialize, Deserialize, Debug, PartialEq, Eq)]
@@ -133,4 +134,5 @@ pub struct AccountModificationForInstruction {
     pub rent_epoch: Option<u64>,
     pub delegated: Option<bool>,
     pub confined: Option<bool>,
+    pub remote_slot: Option<u64>,
 }

programs/magicblock/src/mutate_accounts/process_mutate_accounts.rs

@@ -223,6 +223,14 @@ pub(crate) fn process_mutate_accounts(
             );
             account.borrow_mut().set_confined(confined);
         }
+        if let Some(remote_slot) = modification.remote_slot {
+            ic_msg!(
+                invoke_context,
+                "MutateAccounts: setting remote_slot to {}",
+                remote_slot
+            );
+            account.borrow_mut().set_remote_slot(remote_slot);
+        }
     }
 
     if lamports_to_debit != 0 {
@@ -329,6 +337,7 @@ mod tests {
             rent_epoch: None,
             delegated: Some(true),
             confined: Some(true),
+            remote_slot: None,
         };
         let ix = InstructionUtils::modify_accounts_instruction(vec![
             modification.clone(),
@@ -623,4 +632,46 @@ mod tests {
             }
         );
     }
+
+    #[test]
+    fn test_mod_remote_slot() {
+        init_logger!();
+
+        let mod_key = Pubkey::new_unique();
+        let remote_slot = 12345u64;
+        let mut account_data = {
+            let mut map = HashMap::new();
+            map.insert(mod_key, AccountSharedData::new(100, 0, &mod_key));
+            map
+        };
+        ensure_started_validator(&mut account_data);
+
+        let ix = InstructionUtils::modify_accounts_instruction(vec![
+            AccountModification {
+                pubkey: mod_key,
+                remote_slot: Some(remote_slot),
+                ..Default::default()
+            },
+        ]);
+        let transaction_accounts = ix
+            .accounts
+            .iter()
+            .flat_map(|acc| {
+                account_data
+                    .remove(&acc.pubkey)
+                    .map(|shared_data| (acc.pubkey, shared_data))
+            })
+            .collect();
+
+        let mut accounts = process_instruction(
+            ix.data.as_slice(),
+            transaction_accounts,
+            ix.accounts,
+            Ok(()),
+        );
+
+        let _account_authority = accounts.drain(0..1).next().unwrap();
+        let modified_account = accounts.drain(0..1).next().unwrap();
+        assert_eq!(modified_account.remote_slot(), remote_slot);
+    }
 }

programs/magicblock/src/utils/instruction_utils.rs

@@ -182,6 +182,7 @@ impl InstructionUtils {
                     rent_epoch: account_modification.rent_epoch,
                     delegated: account_modification.delegated,
                     confined: account_modification.confined,
+                    remote_slot: account_modification.remote_slot,
                 };
             account_mods.insert(
                 account_modification.pubkey,

test-integration/configs/schedulecommit-conf-fees.ephem.toml

@@ -5,21 +5,23 @@ listen = "0.0.0.0:8899"
 
 [accountsdb]
 # size of the main storage, we have to preallocate in advance
-# it's advised to set this value based on formula 1KB * N * 3, 
-# where N is the number of accounts expected to be stored in 
+# it's advised to set this value based on formula 1KB * N * 3,
+# where N is the number of accounts expected to be stored in
 # database, e.g. for million accounts this would be 3GB
 database-size = 1048576000 # 1GB
 # minimal indivisible unit of addressing in main storage
 # offsets are calculated in terms of blocks
 block-size = "block256" # possible values block128 | block256 | block512
-# size of index file, we have to preallocate, 
+# size of index file, we have to preallocate,
 # can be as low as 1% of main storage size, but setting it to higher values won't hurt
 index-size = 2048576
 # max number of snapshots to keep around
 max-snapshots = 7
 # how frequently (slot-wise) we should take snapshots
 snapshot-frequency = 1024
 
+reset = true
+
 [ledger]
 reset = true
 

test-integration/programs/schedulecommit/src/api.rs

@@ -256,6 +256,16 @@ pub fn increase_count_instruction(committee: Pubkey) -> Instruction {
     )
 }
 
+pub fn set_count_instruction(committee: Pubkey, count: u64) -> Instruction {
+    let program_id = crate::id();
+    let account_metas = vec![AccountMeta::new(committee, false)];
+    Instruction::new_with_borsh(
+        program_id,
+        &ScheduleCommitInstruction::SetCount(count),
+        account_metas,
+    )
+}
+
 // -----------------
 // PDA
 // -----------------

test-integration/programs/schedulecommit/src/lib.rs

@@ -32,6 +32,8 @@ pub mod api;
 pub mod magicblock_program;
 mod utils;
 
+pub const FAIL_UNDELEGATION_COUNT: u64 = u64::MAX - 1;
+
 declare_id!("9hgprgZiRWmy8KkfvUuaVkDGrqo9GzeXMohwq6BazgUY");
 
 #[cfg(not(feature = "no-entrypoint"))]
@@ -113,6 +115,13 @@ pub enum ScheduleCommitInstruction {
     /// # Account references:
     /// - **0.** `[WRITE]` PDA Account to increase count of
     IncreaseCount,
+
+    /// Sets the count of a PDA of this program to the provided arbitrary u64 value.
+    /// This instruction can only run on the ephemeral after the account was
+    /// delegated or on chain while it is undelegated.
+    /// # Account references:
+    /// - **0.** `[WRITE]` PDA Account to write the counter to set to the supplied value
+    SetCount(u64),
     // This is invoked by the delegation program when we request to undelegate
     // accounts.
     // # Account references:
@@ -170,6 +179,7 @@ pub fn process_instruction<'a>(
             )
         }
         IncreaseCount => process_increase_count(accounts),
+        SetCount(value) => process_set_count(accounts, value),
     }
 }
 
@@ -398,6 +408,27 @@ fn process_increase_count(accounts: &[AccountInfo]) -> ProgramResult {
     Ok(())
 }
 
+fn process_set_count(accounts: &[AccountInfo], value: u64) -> ProgramResult {
+    msg!("Processing set_count instruction");
+    // NOTE: we don't check if the player owning the PDA is signer here for simplicity
+    let accounts_iter = &mut accounts.iter();
+    let account = next_account_info(accounts_iter)?;
+    msg!("Counter account key {}", account.key);
+    let mut main_account = {
+        let main_account_data = account.try_borrow_data()?;
+        MainAccount::try_from_slice(&main_account_data)?
+    };
+    msg!("Owner: {}", account.owner);
+    msg!("Counter account {:#?}", main_account);
+    main_account.count = value;
+    msg!("Set count to {:#?}", main_account);
+    let mut mut_data = account.try_borrow_mut_data()?;
+    let mut as_mut: &mut [u8] = mut_data.as_mut();
+    msg!("Mutating buffer of len: {}", as_mut.len());
+    main_account.serialize(&mut as_mut)?;
+    msg!("Serialized counter");
+    Ok(())
+}
 // -----------------
 // process_schedulecommit_and_undelegation_cpi_with_mod_after
 // -----------------
@@ -586,5 +617,18 @@ fn process_undelegate_request(
         system_program,
         account_seeds,
     )?;
+
+    {
+        let data = delegated_account.try_borrow_data()?;
+        match MainAccount::try_from_slice(&data) {
+            Ok(counter) => {
+                msg!("counter: {:?}", counter);
+                if counter.count == FAIL_UNDELEGATION_COUNT {
+                    return Err(ProgramError::Custom(111));
+                }
+            }
+            Err(err) => msg!("Failed to deserialize: {:?}", err),
+        }
+    };
     Ok(())
 }