chore: small refactor

alchemyplatform · Nov 5, 2024 · 8356075 · 8356075
1 parent d4c23aa
commit 8356075
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 16 deletions.
diff --git a/crates/builder/src/bundle_proposer.rs b/crates/builder/src/bundle_proposer.rs
@@ -1380,13 +1380,7 @@ impl<UO: UserOperation> ProposalContext<UO> {
         let is_sender_staked = entity_infos.sender.is_staked;
         if is_factory_staked || is_sender_staked {
             if let Some(paymaster) = entity_infos.paymaster {
-                self.entity_updates.insert(
-                    paymaster.address(),
-                    EntityUpdate {
-                        entity: paymaster.entity,
-                        update_type: EntityUpdateType::PaymasterAmendment,
-                    },
-                );
+                self.add_erep_015_paymaster_amendment(paymaster.entity)
             }
         }
 
@@ -1489,13 +1483,7 @@ impl<UO: UserOperation> ProposalContext<UO> {
 
         if paymaster_ammendment_required {
             if let Some(paymaster) = entity_infos.paymaster {
-                self.entity_updates.insert(
-                    paymaster.address(),
-                    EntityUpdate {
-                        entity: paymaster.entity,
-                        update_type: EntityUpdateType::PaymasterAmendment,
-                    },
-                );
+                self.add_erep_015_paymaster_amendment(paymaster.entity)
             };
         }
     }
@@ -1559,6 +1547,23 @@ impl<UO: UserOperation> ProposalContext<UO> {
             }
         }
     }
+
+    fn add_erep_015_paymaster_amendment(&mut self, entity: Entity) {
+        if !entity.is_paymaster() {
+            warn!(
+                "Attempted to add EREP-015 paymaster amendment for non-paymaster entity: {:?}",
+                entity
+            );
+            return;
+        }
+        self.entity_updates.insert(
+            entity.address,
+            EntityUpdate {
+                entity,
+                update_type: EntityUpdateType::PaymasterAmendment,
+            },
+        );
+    }
 }
 
 #[cfg(test)]

diff --git a/crates/pool/src/mempool/uo_pool.rs b/crates/pool/src/mempool/uo_pool.rs
@@ -34,7 +34,7 @@ use rundler_types::{
 use rundler_utils::emit::WithEntryPoint;
 use tokio::sync::broadcast;
 use tonic::async_trait;
-use tracing::info;
+use tracing::{info, warn};
 
 use super::{
     paymaster::PaymasterTracker, pool::PoolInner, reputation::AddressReputation, Mempool,
@@ -733,7 +733,14 @@ where
                 self.reputation.handle_srep_050_penalty(entity.address);
             }
             EntityUpdateType::PaymasterAmendment => {
-                self.reputation.remove_seen(entity.address);
+                if entity.is_paymaster() {
+                    self.reputation.remove_seen(entity.address);
+                } else {
+                    warn!(
+                        "Received paymaster amendment for non-paymaster entity: {:?}",
+                        entity
+                    );
+                }
             }
         }
 

diff --git a/crates/types/src/entity.rs b/crates/types/src/entity.rs
@@ -100,6 +100,11 @@ impl Entity {
         Self::new(EntityType::Paymaster, address)
     }
 
+    /// Check if the entity is a paymaster
+    pub fn is_paymaster(&self) -> bool {
+        self.kind == EntityType::Paymaster
+    }
+
     /// Create a new aggregator entity at address
     pub fn aggregator(address: Address) -> Self {
         Self::new(EntityType::Aggregator, address)