Skip to content

Commit

Permalink
fix(pool): fix race condition in paymaster tracking
Browse files Browse the repository at this point in the history
  • Loading branch information
@dancoombs
dancoombs committed Jun 14, 2024
1 parent 98f8e2c commit ef374b8
Show file tree
Hide file tree
Showing 3 changed files with 42 additions and 100 deletions.
104 changes: 16 additions & 88 deletions crates/pool/src/mempool/paymaster.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ use rundler_types::{
use rundler_utils::cache::LruMap;

use super::MempoolResult;
use crate::chain::{BalanceUpdate, MinedOp};
use crate::chain::MinedOp;

/// Keeps track of current and pending paymaster balances
#[derive(Debug)]
Expand Down Expand Up @@ -94,28 +94,6 @@ where
Ok(stake_status)
}

pub(crate) fn update_paymaster_balances_after_update<'a>(
&self,
entity_balance_updates: impl Iterator<Item = &'a BalanceUpdate>,
unmined_entity_balance_updates: impl Iterator<Item = &'a BalanceUpdate>,
) {
for balance_update in entity_balance_updates {
self.state.write().update_paymaster_balance_from_event(
balance_update.address,
balance_update.amount,
balance_update.is_addition,
)
}

for unmined_balance_update in unmined_entity_balance_updates {
self.state.write().update_paymaster_balance_from_event(
unmined_balance_update.address,
unmined_balance_update.amount,
!unmined_balance_update.is_addition,
)
}
}

pub(crate) async fn paymaster_balance(
&self,
paymaster: Address,
Expand Down Expand Up @@ -174,6 +152,19 @@ where
self.state.write().set_tracking(tracking_enabled);
}

pub(crate) async fn reset_confirmed_balances_for(
&self,
addresses: Vec<Address>,
) -> MempoolResult<()> {
let balances = self.entry_point.get_balances(addresses.clone()).await?;

self.state
.write()
.set_confimed_balances(&addresses, &balances);

Ok(())
}

pub(crate) async fn reset_confimed_balances(&self) -> MempoolResult<()> {
let paymaster_addresses = self.paymaster_addresses();

Expand All @@ -194,6 +185,7 @@ where
.write()
.update_paymaster_balance_from_mined_op(mined_op);
}

pub(crate) fn remove_operation(&self, id: &UserOperationId) {
self.state.write().remove_operation(id);
}
Expand Down Expand Up @@ -324,21 +316,6 @@ impl PaymasterTrackerInner {
keys
}

fn update_paymaster_balance_from_event(
&mut self,
paymaster: Address,
amount: U256,
should_add: bool,
) {
if let Some(paymaster_balance) = self.paymaster_balances.get(&paymaster) {
if should_add {
paymaster_balance.confirmed = paymaster_balance.confirmed.saturating_add(amount);
} else {
paymaster_balance.confirmed = paymaster_balance.confirmed.saturating_sub(amount);
}
}
}

fn paymaster_metadata(&self, paymaster: Address) -> Option<PaymasterMetadata> {
if let Some(paymaster_balance) = self.paymaster_balances.peek(&paymaster) {
return Some(PaymasterMetadata {
Expand Down Expand Up @@ -530,7 +507,7 @@ mod tests {
};

use super::*;
use crate::{chain::BalanceUpdate, mempool::paymaster::PaymasterTracker};
use crate::mempool::paymaster::PaymasterTracker;

fn demo_pool_op(uo: UserOperation) -> PoolOperation {
PoolOperation {
Expand Down Expand Up @@ -609,55 +586,6 @@ mod tests {
assert!(res.is_err());
}

#[tokio::test]
async fn test_update_balance() {
let paymaster_tracker = new_paymaster_tracker();

let paymaster = Address::random();
let pending_op_cost = U256::from(100);
let confirmed_balance = U256::from(1000);

paymaster_tracker.add_new_paymaster(paymaster, confirmed_balance, pending_op_cost);

let deposit = BalanceUpdate {
address: paymaster,
entrypoint: Address::random(),
amount: 100.into(),
is_addition: true,
};

// deposit
paymaster_tracker
.update_paymaster_balances_after_update(vec![&deposit].into_iter(), vec![].into_iter());

let balance = paymaster_tracker
.paymaster_balance(paymaster)
.await
.unwrap();

assert_eq!(balance.confirmed_balance, 1100.into());

let withdrawal = BalanceUpdate {
address: paymaster,
entrypoint: Address::random(),
amount: 50.into(),
is_addition: false,
};

// withdrawal
paymaster_tracker.update_paymaster_balances_after_update(
vec![&withdrawal].into_iter(),
vec![].into_iter(),
);

let balance = paymaster_tracker
.paymaster_balance(paymaster)
.await
.unwrap();

assert_eq!(balance.confirmed_balance, 1050.into());
}

#[tokio::test]
async fn new_uo_not_enough_balance_tracking_disabled() {
let paymaster_tracker = new_paymaster_tracker();
Expand Down
35 changes: 24 additions & 11 deletions crates/pool/src/mempool/uo_pool.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -151,15 +151,24 @@ where
.iter()
.filter(|op| op.entry_point == self.config.entry_point);

let entity_balance_updates = update
.entity_balance_updates
.iter()
.filter(|u| u.entrypoint == self.config.entry_point);
let entity_balance_updates = update.entity_balance_updates.iter().filter_map(|u| {
if u.entrypoint == self.config.entry_point {
Some(u.address)
} else {
None
}
});

let unmined_entity_balance_updates = update
.unmined_entity_balance_updates
.iter()
.filter(|u| u.entrypoint == self.config.entry_point);
.filter_map(|u| {
if u.entrypoint == self.config.entry_point {
Some(u.address)
} else {
None
}
});

let unmined_ops = deduped_ops
.unmined_ops
Expand All @@ -169,14 +178,18 @@ where
let mut unmined_op_count = 0;

if update.reorg_larger_than_history {
let _ = self.reset_confirmed_paymaster_balances().await;
if let Err(e) = self.reset_confirmed_paymaster_balances().await {
tracing::error!("Failed to reset confirmed paymaster balances: {:?}", e);
}
} else {
let addresses = entity_balance_updates
.chain(unmined_entity_balance_updates)
.collect::<Vec<_>>();
if let Err(e) = self.paymaster.reset_confirmed_balances_for(addresses).await {
tracing::error!("Failed to reset confirmed paymaster balances: {:?}", e);
}
}

self.paymaster.update_paymaster_balances_after_update(
entity_balance_updates,
unmined_entity_balance_updates,
);

for op in mined_ops {
if op.entry_point != self.config.entry_point {
continue;
Expand Down
3 changes: 2 additions & 1 deletion crates/rpc/src/eth/error.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ const THROTTLED_OR_BANNED_CODE: i32 = -32504;
const STAKE_TOO_LOW_CODE: i32 = -32505;
const UNSUPORTED_AGGREGATOR_CODE: i32 = -32506;
const SIGNATURE_CHECK_FAILED_CODE: i32 = -32507;
const PAYMASTER_DEPOSIT_TOO_LOW: i32 = -32508;
const EXECUTION_REVERTED: i32 = -32521;

pub(crate) type EthResult<T> = Result<T, EthRpcError>;
Expand Down Expand Up @@ -383,14 +384,14 @@ impl From<EthRpcError> for ErrorObjectOwned {
EthRpcError::PaymasterValidationRejected(data) => {
rpc_err_with_data(PAYMASTER_VALIDATION_REJECTED_CODE, msg, data)
}
EthRpcError::PaymasterBalanceTooLow(_, _) => rpc_err(PAYMASTER_DEPOSIT_TOO_LOW, msg),
EthRpcError::OpcodeViolation(_, _)
| EthRpcError::OpcodeViolationMap(_)
| EthRpcError::OutOfGas(_)
| EthRpcError::UnstakedAggregator
| EthRpcError::MultipleRolesViolation(_)
| EthRpcError::UnstakedPaymasterContext
| EthRpcError::SenderAddressUsedAsAlternateEntity(_)
| EthRpcError::PaymasterBalanceTooLow(_, _)
| EthRpcError::AssociatedStorageIsAlternateSender
| EthRpcError::AssociatedStorageDuringDeploy(_, _, _)
| EthRpcError::InvalidStorageAccess(_, _, _) => rpc_err(OPCODE_VIOLATION_CODE, msg),
Expand Down

0 comments on commit ef374b8

Please sign in to comment.