Skip to content

Commit

Permalink
Enfornce rate limits on decrease_take
Browse files Browse the repository at this point in the history
  • Loading branch information
gztensor committed Jul 11, 2024
1 parent 852fa4c commit 095b618
Show file tree
Hide file tree
Showing 2 changed files with 80 additions and 2 deletions.
11 changes: 10 additions & 1 deletion pallets/subtensor/src/staking.rs
Original file line number Diff line number Diff line change
Expand Up @@ -155,8 +155,17 @@ impl<T: Config> Pallet<T> {
let min_take = MinTake::<T>::get();
ensure!(take >= min_take, Error::<T>::DelegateTakeTooLow);

// Set last block for rate limiting
// --- 5. Enforce the rate limit (independently on do_add_stake rate limits)
let block: u64 = Self::get_current_block_as_u64();
ensure!(
!Self::exceeds_tx_delegate_take_rate_limit(
Self::get_last_tx_block_delegate_take(&coldkey),
block
),
Error::<T>::DelegateTxRateLimitExceeded
);

// Set last block for rate limiting
Self::set_last_tx_block_delegate_take(&coldkey, block);

// --- 4. Set the new take value.
Expand Down
71 changes: 70 additions & 1 deletion pallets/subtensor/tests/staking.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2824,6 +2824,9 @@ fn test_can_set_min_take_ok() {
u16::MAX / 10
));

// Wait so that we can decrease take
step_block(1 + InitialTxDelegateTakeRateLimit::get() as u16);

// Coldkey / hotkey 0 decreases take to min
assert_ok!(SubtensorModule::do_decrease_take(
<<Test as Config>::RuntimeOrigin>::signed(coldkey0),
Expand Down Expand Up @@ -3136,7 +3139,7 @@ fn test_rate_limits_enforced_on_increase_take() {
});
}

// Test rate-limiting on increase_take
// Test rate-limiting on decrease_take + increase_take
#[test]
fn test_rate_limits_enforced_on_increase_take_after_decrease() {
new_test_ext(1).execute_with(|| {
Expand All @@ -3159,6 +3162,9 @@ fn test_rate_limits_enforced_on_increase_take_after_decrease() {
pallet_subtensor::MaxTake::<Test>::get()
));

// Wait so that we can decrease take
step_block(1 + InitialTxDelegateTakeRateLimit::get() as u16);

// Coldkey / hotkey 0 decreases take
assert_ok!(SubtensorModule::do_decrease_take(
<<Test as Config>::RuntimeOrigin>::signed(coldkey0),
Expand Down Expand Up @@ -3196,6 +3202,69 @@ fn test_rate_limits_enforced_on_increase_take_after_decrease() {
});
}

// Test rate-limiting on decrease_take
#[test]
fn test_rate_limits_enforced_on_decrease_take() {
new_test_ext(1).execute_with(|| {
// Make account
let hotkey0 = U256::from(1);
let coldkey0 = U256::from(3);

// Add balance
SubtensorModule::add_balance_to_coldkey_account(&coldkey0, 100000);

// Register the neuron to a new network
let netuid = 1;
add_network(netuid, 0, 0);
register_ok_neuron(netuid, hotkey0, coldkey0, 124124);

// Coldkey / hotkey 0 become delegates with max take
assert_ok!(SubtensorModule::do_become_delegate(
<<Test as Config>::RuntimeOrigin>::signed(coldkey0),
hotkey0,
pallet_subtensor::MaxTake::<Test>::get()
));

// Wait so that we can decrease take first time
step_block(1 + InitialTxDelegateTakeRateLimit::get() as u16);

// Coldkey / hotkey 0 decreases take
assert_ok!(SubtensorModule::do_decrease_take(
<<Test as Config>::RuntimeOrigin>::signed(coldkey0),
hotkey0,
SubtensorModule::get_min_take() + 1
));

// Attempt to immediately decrease again fails due to rate limit
assert_eq!(
SubtensorModule::do_decrease_take(
<<Test as Config>::RuntimeOrigin>::signed(coldkey0),
hotkey0,
SubtensorModule::get_min_take()
),
Err(Error::<Test>::DelegateTxRateLimitExceeded.into())
);
assert_eq!(
SubtensorModule::get_hotkey_take(&hotkey0),
SubtensorModule::get_min_take() + 1
);

// After number of blocks, take can be decreased
step_block(1 + InitialTxDelegateTakeRateLimit::get() as u16);

// Can increase after waiting
assert_ok!(SubtensorModule::do_decrease_take(
<<Test as Config>::RuntimeOrigin>::signed(coldkey0),
hotkey0,
SubtensorModule::get_min_take()
));
assert_eq!(
SubtensorModule::get_hotkey_take(&hotkey0),
SubtensorModule::get_min_take()
);
});
}

// Helper function to set up a test environment
fn setup_test_environment() -> (AccountId, AccountId, AccountId) {
let current_coldkey = U256::from(1);
Expand Down

0 comments on commit 095b618

Please sign in to comment.