Vault: Fix issue with stale update vaults (#163)

coachchucksol · web-flow · commit 3fdcd8876ea1 · 2024-12-23T14:09:21.000-07:00
Forces the vault to update from the previous state, given missed update
epochs or partial updates that have not called
close_vault_update_state_tracker
diff --git a/clients/js/vault_client/accounts/vault.ts b/clients/js/vault_client/accounts/vault.ts
@@ -82,6 +82,7 @@ export type Vault = {
   programFeeBps: number;
   bump: number;
   isPaused: number;
+  lastStartStateUpdateSlot: bigint;
   reserved: Array<number>;
 };
 
@@ -122,6 +123,7 @@ export type VaultArgs = {
   programFeeBps: number;
   bump: number;
   isPaused: number;
+  lastStartStateUpdateSlot: number | bigint;
   reserved: Array<number>;
 };
 
@@ -163,7 +165,8 @@ export function getVaultEncoder(): Encoder<VaultArgs> {
     ['programFeeBps', getU16Encoder()],
     ['bump', getU8Encoder()],
     ['isPaused', getBoolEncoder()],
-    ['reserved', getArrayEncoder(getU8Encoder(), { size: 259 })],
+    ['lastStartStateUpdateSlot', getU64Encoder()],
+    ['reserved', getArrayEncoder(getU8Encoder(), { size: 251 })],
   ]);
 }
 
@@ -205,7 +208,8 @@ export function getVaultDecoder(): Decoder<Vault> {
     ['programFeeBps', getU16Decoder()],
     ['bump', getU8Decoder()],
     ['isPaused', getBoolDecoder()],
-    ['reserved', getArrayDecoder(getU8Decoder(), { size: 259 })],
+    ['lastStartStateUpdateSlot', getU64Decoder()],
+    ['reserved', getArrayDecoder(getU8Decoder(), { size: 251 })],
   ]);
 }
 
diff --git a/clients/rust/vault_client/src/generated/accounts/vault.rs b/clients/rust/vault_client/src/generated/accounts/vault.rs
@@ -104,8 +104,9 @@ pub struct Vault {
     pub program_fee_bps: u16,
     pub bump: u8,
     pub is_paused: bool,
+    pub last_start_state_update_slot: u64,
     #[cfg_attr(feature = "serde", serde(with = "serde_with::As::<serde_with::Bytes>"))]
-    pub reserved: [u8; 259],
+    pub reserved: [u8; 251],
 }
 
 impl Vault {
diff --git a/idl/jito_vault.json b/idl/jito_vault.json
@@ -1652,12 +1652,18 @@
               "defined": "PodBool"
             }
           },
+          {
+            "name": "lastStartStateUpdateSlot",
+            "type": {
+              "defined": "PodU64"
+            }
+          },
           {
             "name": "reserved",
             "type": {
               "array": [
                 "u8",
-                259
+                251
               ]
             }
           }
diff --git a/integration_tests/tests/vault/close_vault_update_state_tracker.rs b/integration_tests/tests/vault/close_vault_update_state_tracker.rs
@@ -0,0 +1,204 @@
+#[cfg(test)]
+mod tests {
+    use jito_vault_core::{
+        config::Config, vault_operator_delegation::VaultOperatorDelegation,
+        vault_update_state_tracker::VaultUpdateStateTracker,
+    };
+    use solana_sdk::signature::{Keypair, Signer};
+
+    use crate::fixtures::{
+        fixture::{ConfiguredVault, TestBuilder},
+        vault_client::VaultStakerWithdrawalTicketRoot,
+    };
+
+    #[tokio::test]
+    async fn test_dos_vault_if_state_tracker_not_closed_when_additional_assets_need_unstaking_is_zero(
+    ) {
+        let mut fixture = TestBuilder::new().await;
+
+        let deposit_fee_bps = 0;
+        let withdrawal_fee_bps = 0;
+        let reward_fee_bps = 0;
+        let num_operators = 1;
+        let slasher_amounts = vec![];
+
+        let ConfiguredVault {
+            mut vault_program_client,
+            vault_root,
+            operator_roots,
+            ..
+        } = fixture
+            .setup_vault_with_ncn_and_operators(
+                deposit_fee_bps,
+                withdrawal_fee_bps,
+                reward_fee_bps,
+                num_operators,
+                &slasher_amounts,
+            )
+            .await
+            .unwrap();
+
+        let depositor = Keypair::new();
+        vault_program_client
+            .configure_depositor(&vault_root, &depositor.pubkey(), 100_000)
+            .await
+            .unwrap();
+        vault_program_client
+            .do_mint_to(&vault_root, &depositor, 100_000, 100_000)
+            .await
+            .unwrap();
+
+        vault_program_client
+            .do_add_delegation(&vault_root, &operator_roots[0].operator_pubkey, 100_000)
+            .await
+            .unwrap();
+
+        let config = vault_program_client
+            .get_config(&Config::find_program_address(&jito_vault_program::id()).0)
+            .await
+            .unwrap();
+
+        let VaultStakerWithdrawalTicketRoot { base: _ } = vault_program_client
+            .do_enqueue_withdrawal(&vault_root, &depositor, 10_000)
+            .await
+            .unwrap();
+        let vault = vault_program_client
+            .get_vault(&vault_root.vault_pubkey)
+            .await
+            .unwrap();
+        assert_eq!(vault.vrt_enqueued_for_cooldown_amount(), 10_000);
+        assert_eq!(vault.vrt_cooling_down_amount(), 0);
+        assert_eq!(vault.vrt_ready_to_claim_amount(), 0);
+
+        fixture
+            .warp_slot_incremental(config.epoch_length())
+            .await
+            .unwrap();
+
+        let operator_pubkeys: Vec<_> = operator_roots
+            .iter()
+            .map(|root| root.operator_pubkey)
+            .collect();
+
+        //start of full vault update
+        let slot = fixture.get_current_slot().await.unwrap();
+
+        let ncn_epoch = slot / config.epoch_length();
+
+        let vault_update_state_tracker = VaultUpdateStateTracker::find_program_address(
+            &jito_vault_program::id(),
+            &vault_root.vault_pubkey,
+            ncn_epoch,
+        )
+        .0;
+        //initialize vault_update_state_tracker
+        vault_program_client
+            .initialize_vault_update_state_tracker(
+                &vault_root.vault_pubkey,
+                &vault_update_state_tracker,
+            )
+            .await
+            .unwrap();
+        //crank update all operators so addoitional_assets_need_unstaking=0
+        for i in 0..operator_pubkeys.len() {
+            let operator_index = (i + (ncn_epoch as usize)) % operator_pubkeys.len();
+            let operator = &operator_pubkeys[operator_index];
+            vault_program_client
+                .crank_vault_update_state_tracker(
+                    &vault_root.vault_pubkey,
+                    operator,
+                    &VaultOperatorDelegation::find_program_address(
+                        &jito_vault_program::id(),
+                        &vault_root.vault_pubkey,
+                        operator,
+                    )
+                    .0,
+                    &vault_update_state_tracker,
+                )
+                .await
+                .unwrap();
+        }
+        //fast forward to next epoch without closing update_state_tracker even though all operators have been updated
+        fixture
+            .warp_slot_incremental(config.epoch_length())
+            .await
+            .unwrap();
+        //close the update state tracker(not necessary though)
+        vault_program_client
+            .close_vault_update_state_tracker(
+                &vault_root.vault_pubkey,
+                &vault_update_state_tracker,
+                slot / config.epoch_length(),
+            )
+            .await
+            .unwrap();
+
+        //end of full vault update
+
+        //initialize another update state tracker
+        let slot = fixture.get_current_slot().await.unwrap();
+
+        let ncn_epoch = slot / config.epoch_length();
+
+        let vault_update_state_tracker = VaultUpdateStateTracker::find_program_address(
+            &jito_vault_program::id(),
+            &vault_root.vault_pubkey,
+            ncn_epoch,
+        )
+        .0;
+        //new epoch: initialize new vault update state tracker
+        vault_program_client
+            .initialize_vault_update_state_tracker(
+                &vault_root.vault_pubkey,
+                &vault_update_state_tracker,
+            )
+            .await
+            .unwrap();
+        //crank update all operators
+        for i in 0..operator_pubkeys.len() {
+            let operator_index = (i + (ncn_epoch as usize)) % operator_pubkeys.len();
+            let operator = &operator_pubkeys[operator_index];
+            vault_program_client
+                .crank_vault_update_state_tracker(
+                    &vault_root.vault_pubkey,
+                    operator,
+                    &VaultOperatorDelegation::find_program_address(
+                        &jito_vault_program::id(),
+                        &vault_root.vault_pubkey,
+                        operator,
+                    )
+                    .0,
+                    &vault_update_state_tracker,
+                )
+                .await
+                .unwrap();
+        }
+        //closing vault_update_state_tracker fails cos additional_assets_need_unstaking>0 as there were no more delegations to decrement from the operators
+        let result = vault_program_client
+            .close_vault_update_state_tracker(
+                &vault_root.vault_pubkey,
+                &vault_update_state_tracker,
+                slot / config.epoch_length(),
+            )
+            .await;
+
+        let vault = vault_program_client
+            .get_vault(&vault_root.vault_pubkey)
+            .await
+            .unwrap();
+
+        assert_eq!(vault.additional_assets_need_unstaking(), 0);
+
+        assert!(result.is_ok());
+
+        // Previous result
+        // //Error cos additional_assets_need_unstaking>0
+        // assert_vault_error(
+        //     result,
+        //     VaultError::NonZeroAdditionalAssetsNeededForWithdrawalAtEndOfUpdate,
+        // );
+        // //Now that vault_update_state_tracker can't be closed, all other operations are impossible:withdrawals, delegations, cooldowns...
+
+        // end of block
+    }
+}
diff --git a/integration_tests/tests/vault/initialize_vault_update_state_tracker.rs b/integration_tests/tests/vault/initialize_vault_update_state_tracker.rs
@@ -293,6 +293,12 @@ mod tests {
             .await
             .unwrap();
 
+        // Do full update
+        vault_program_client
+            .do_full_vault_update(&vault_root.vault_pubkey, &[operator_root.operator_pubkey])
+            .await
+            .unwrap();
+
         // Set fees
         let new_withdrawal_fee_bps = 10;
         let new_program_fee_bps = 11;
@@ -427,12 +433,11 @@ mod tests {
             .unwrap();
 
         fixture
-            .warp_slot_incremental(2 * config.epoch_length())
+            .warp_slot_incremental(1 * config.epoch_length())
             .await
             .unwrap();
 
         // enqueued for cool down assets are now cooling down
-
         let slot = fixture.get_current_slot().await.unwrap();
         let ncn_epoch = slot / config.epoch_length();
 
@@ -442,6 +447,7 @@ mod tests {
             ncn_epoch,
         )
         .0;
+
         vault_program_client
             .initialize_vault_update_state_tracker(
                 &vault_root.vault_pubkey,
@@ -460,6 +466,13 @@ mod tests {
             75_000 - Vault::DEFAULT_INITIALIZATION_TOKEN_AMOUNT
         );
 
+        // skip cranking operator 0, advance to next epoch
+
+        fixture
+            .warp_slot_incremental(config.epoch_length())
+            .await
+            .unwrap();
+
         // Update fees
         let new_withdrawal_fee_bps = 10;
         let new_program_fee_bps = 11;
@@ -479,13 +492,6 @@ mod tests {
             .await
             .unwrap();
 
-        // skip cranking operator 0, advance to next epoch
-
-        fixture
-            .warp_slot_incremental(config.epoch_length())
-            .await
-            .unwrap();
-
         let slot = fixture.get_current_slot().await.unwrap();
         let ncn_epoch = slot / config.epoch_length();
         // no assets cooled down, additional_assets_need_unstaking = 75_000
diff --git a/integration_tests/tests/vault/mod.rs b/integration_tests/tests/vault/mod.rs
@@ -1,6 +1,7 @@
 mod add_delegation;
 mod burn_withdrawal_ticket;
 mod close_update_state_tracker;
+mod close_vault_update_state_tracker;
 mod cooldown_delegation;
 mod cooldown_vault_ncn_ticket;
 mod crank_vault_update_state_tracker;
diff --git a/vault_core/src/vault.rs b/vault_core/src/vault.rs
diff --git a/vault_program/src/crank_vault_update_state_tracker.rs b/vault_program/src/crank_vault_update_state_tracker.rs
diff --git a/vault_program/src/initialize_vault_update_state_tracker.rs b/vault_program/src/initialize_vault_update_state_tracker.rs

Original file line number	Diff line number	Diff line change
`@@ -1652,12 +1652,18 @@`
`1652`	`1652`	`"defined": "PodBool"`
`1653`	`1653`	`}`
`1654`	`1654`	`},`
	`1655`	`+ {`
	`1656`	`+ "name": "lastStartStateUpdateSlot",`
	`1657`	`+ "type": {`
	`1658`	`+ "defined": "PodU64"`
	`1659`	`+ }`
	`1660`	`+ },`
`1655`	`1661`	`{`
`1656`	`1662`	`"name": "reserved",`
`1657`	`1663`	`"type": {`
`1658`	`1664`	`"array": [`
`1659`	`1665`	`"u8",`
`1660`		`- 259`
	`1666`	`+ 251`
`1661`	`1667`	`]`
`1662`	`1668`	`}`
`1663`	`1669`	`}`