cleanup

ethui · Feb 3, 2025 · 549ef3c · 549ef3c
1 parent 6cbf703
commit 549ef3c
Show file tree

Hide file tree

Showing 7 changed files with 78 additions and 43 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -10,9 +10,9 @@ alloy = { version = "0.11.0", features = [
   "rpc",
   "rpc-types",
 ] }
-rstest = "0.24.0"
 thiserror = "2.0.11"
 
 [dev-dependencies]
+rstest = "0.24.0"
 tokio = { version = "1.43.0", features = ["rt-multi-thread"] }
 lazy_static = "1.5.0"
diff --git a/src/comptroller.rs b/src/comptroller.rs
@@ -0,0 +1,31 @@
+use alloy::{
+    network::{Network, TransactionBuilder as _},
+    primitives::{b256, Address, B256},
+    providers::Provider,
+};
+
+use crate::{error::DetectProxyResult, utils::u256_to_address, ProxyType};
+
+const COMPTROLLER_INTERFACE: [B256; 1] = [
+    // bytes4(keccak256("comptrollerImplementation()")) padded to 32 bytes
+    b256!("0xbb82aa5e00000000000000000000000000000000000000000000000000000000"),
+];
+
+pub(crate) async fn detect_comptroller_proxy<N, P: Provider<N>>(
+    address: Address,
+    provider: P,
+) -> DetectProxyResult<Option<ProxyType>>
+where
+    N: Network,
+{
+    let call_0 = <N as Network>::TransactionRequest::default()
+        .with_to(address)
+        .with_input(COMPTROLLER_INTERFACE[0]);
+
+    if let Ok(value) = provider.call(&call_0).await {
+        let b256: B256 = B256::from_slice(&value);
+        return Ok(Some(ProxyType::Comptroller(u256_to_address(b256.into()))));
+    };
+
+    Ok(None)
+}
diff --git a/src/eip1167.rs b/src/eip1167.rs
@@ -1,46 +1,45 @@
-use alloy::primitives::{bytes, Address, Bytes};
+use alloy::{
+    network::Network,
+    primitives::{bytes, Address, Bytes},
+    providers::Provider,
+};
+
+use crate::{error::DetectProxyResult, ProxyType};
 
 const EIP1167_PREFIX: Bytes = bytes!("363d3d373d3d3d363d");
 const EIP1167_SUFFIX: Bytes = bytes!("57fd5bf3");
 const EIP1167_SUFFIX_OFFSET_FROM_ADDRESS_END: usize = 11;
 
-pub(crate) fn detect_eip1167_minimal_proxy(code: &Bytes) -> Option<Address> {
+pub(crate) async fn detect_eip1167_minimal_proxy<N, P: Provider<N>>(
+    address: Address,
+    provider: P,
+) -> DetectProxyResult<Option<ProxyType>>
+where
+    N: Network,
+{
+    let code = provider.get_code_at(address).await?;
+
     if !code.starts_with(&EIP1167_PREFIX) {
-        return None;
+        return Ok(None);
     }
 
     // detect length of address (20 bytes non-optimized, 0 < N < 20 bytes for vanity addresses)
     // push1 ... push20 use opcode 0x60 ... 0x73
     let address_len = code[EIP1167_PREFIX.len()] as usize - 0x5f;
 
     if !(1..=20).contains(&address_len) {
-        return None;
+        return Ok(None);
     }
 
     let address_pos = EIP1167_PREFIX.len() + 1;
     let suffix = &code[address_pos + address_len + EIP1167_SUFFIX_OFFSET_FROM_ADDRESS_END..];
 
     if !suffix.starts_with(&EIP1167_SUFFIX) {
-        return None;
+        return Ok(None);
     }
 
     let address_hex = &code[address_pos..address_pos + address_len];
     let address = Address::left_padding_from(address_hex);
 
-    Some(address)
-}
-
-#[cfg(test)]
-mod tests {
-    use alloy::primitives::address;
-
-    use super::*;
-
-    #[test]
-    fn parse_eip1167_code() {
-        let bytecode: Bytes = bytes!("363d3d373d3d3d363d73f62849f9a0b5bf2913b396098f7c7019b51a820a5af43d82803e903d91602b57fd5bf3000000000000000000000000000000000000000000000000000000000000007a6900000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000");
-        let address: Address = address!("f62849f9a0b5bf2913b396098f7c7019b51a820a");
-
-        assert_eq!(detect_eip1167_minimal_proxy(&bytecode), Some(address));
-    }
+    Ok(Some(ProxyType::Eip1167(address)))
 }
diff --git a/src/eip897.rs b/src/eip897.rs
@@ -4,7 +4,7 @@ use alloy::{
     providers::Provider,
 };
 
-use crate::{error::DetectProxyResult, utils::u256_to_address};
+use crate::{error::DetectProxyResult, utils::u256_to_address, ProxyType};
 
 const EIP_897_INTERFACE: [B256; 2] = [
     // bytes4(keccak256("implementation()")) padded to 32 bytes
@@ -16,7 +16,7 @@ const EIP_897_INTERFACE: [B256; 2] = [
 pub(crate) async fn detect_eip897_proxy<N, P: Provider<N>>(
     address: Address,
     provider: P,
-) -> DetectProxyResult<Option<Address>>
+) -> DetectProxyResult<Option<ProxyType>>
 where
     N: Network,
 {
@@ -26,7 +26,7 @@ where
 
     if let Ok(value) = provider.call(&call_0).await {
         let b256: B256 = B256::from_slice(&value);
-        return Ok(Some(u256_to_address(b256.into())));
+        return Ok(Some(ProxyType::Eip897(u256_to_address(b256.into()))));
     };
 
     Ok(None)

diff --git a/src/lib.rs b/src/lib.rs
@@ -1,3 +1,4 @@
+mod comptroller;
 mod eip1167;
 mod eip1967;
 mod eip897;
@@ -17,6 +18,7 @@ pub enum ProxyType {
     OpenZeppelin(Address),
     Eip897(Address),
     Safe(Address),
+    Comptroller(Address),
 }
 
 pub async fn detect_proxy<N, P: Provider<N>>(
@@ -26,10 +28,8 @@ pub async fn detect_proxy<N, P: Provider<N>>(
 where
     N: Network,
 {
-    let code = provider.get_code_at(address).await?;
-
-    if let Some(address) = eip1167::detect_eip1167_minimal_proxy(&code) {
-        return Ok(Some(ProxyType::Eip1167(address)));
+    if let Some(proxy_type) = eip1167::detect_eip1167_minimal_proxy(address, provider).await? {
+        return Ok(Some(proxy_type));
     }
 
     if let Some(proxy_type) = eip1967::detect_eip1967_direct_proxy(address, provider).await? {
@@ -40,16 +40,20 @@ where
         return Ok(Some(proxy_type));
     }
 
-    if let Some(address) = openzeppelin::detect_open_zeppelin_proxy(address, provider).await? {
-        return Ok(Some(ProxyType::OpenZeppelin(address)));
+    if let Some(proxy_type) = openzeppelin::detect_open_zeppelin_proxy(address, provider).await? {
+        return Ok(Some(proxy_type));
+    }
+
+    if let Some(proxy_type) = eip897::detect_eip897_proxy(address, provider).await? {
+        return Ok(Some(proxy_type));
     }
 
-    if let Some(address) = eip897::detect_eip897_proxy(address, provider).await? {
-        return Ok(Some(ProxyType::Eip897(address)));
+    if let Some(proxy_type) = safe::detect_safe_proxy(address, provider).await? {
+        return Ok(Some(proxy_type));
     }
 
-    if let Some(address) = safe::detect_safe_proxy(address, provider).await? {
-        return Ok(Some(ProxyType::Safe(address)));
+    if let Some(proxy_type) = comptroller::detect_comptroller_proxy(address, provider).await? {
+        return Ok(Some(proxy_type));
     }
 
     Ok(None)
@@ -79,6 +83,7 @@ mod tests {
     #[case::openzeppelin(address!("0xC986c2d326c84752aF4cC842E033B9ae5D54ebbB"), ProxyType::OpenZeppelin(address!("0x0656368c4934e56071056da375d4a691d22161f8")))]
     #[case::eip897(address!("0x8260b9eC6d472a34AD081297794d7Cc00181360a"), ProxyType::Eip1967Direct(address!("0xe4e4003afe3765aca8149a82fc064c0b125b9e5a")))]
     #[case::eip897(address!("0x0DA0C3e52C977Ed3cBc641fF02DD271c3ED55aFe"), ProxyType::Safe(address!("0xd9db270c1b5e3bd161e8c8503c55ceabee709552")))]
+    #[case::eip897(address!("0x3d9819210A31b4961b30EF54bE2aeD79B9c9Cd3B"), ProxyType::Comptroller(address!("0xbafe01ff935c7305907c33bf824352ee5979b526")))]
     #[tokio::test]
     async fn mainnet(#[case] proxy: Address, #[case] impl_: ProxyType) {
         let provider = ProviderBuilder::new().on_http(MAINNET_RPC.clone());

diff --git a/src/openzeppelin.rs b/src/openzeppelin.rs
@@ -4,21 +4,21 @@ use alloy::{
     providers::Provider,
 };
 
-use crate::{error::DetectProxyResult, utils::storage_slot_as_address};
+use crate::{error::DetectProxyResult, utils::storage_slot_as_address, ProxyType};
 
 const OPEN_ZEPPELIN_PREFIX: B256 =
     b256!("0x7050c9e0f4ca769c69bd3a8ef740bc37934f8e2c036e5a723fd8ee048ed3f8c3");
 
 pub(crate) async fn detect_open_zeppelin_proxy<N, P: Provider<N>>(
     address: Address,
     provider: P,
-) -> DetectProxyResult<Option<Address>>
+) -> DetectProxyResult<Option<ProxyType>>
 where
     N: Network,
 {
     if let Ok(Some(addr)) = storage_slot_as_address(&provider, address, OPEN_ZEPPELIN_PREFIX).await
     {
-        return Ok(Some(addr));
+        return Ok(Some(ProxyType::OpenZeppelin(addr)));
     }
 
     Ok(None)

diff --git a/src/safe.rs b/src/safe.rs
@@ -4,27 +4,27 @@ use alloy::{
     providers::Provider,
 };
 
-use crate::{error::DetectProxyResult, utils::u256_to_address};
+use crate::{error::DetectProxyResult, utils::u256_to_address, ProxyType};
 
-const EIP_SAFE_INTERFACE: [B256; 1] = [
+const SAFE_INTERFACE: [B256; 1] = [
     // bytes4(keccak256("masterCopy()")) padded to 32 bytes
     b256!("0xa619486e00000000000000000000000000000000000000000000000000000000"),
 ];
 
 pub(crate) async fn detect_safe_proxy<N, P: Provider<N>>(
     address: Address,
     provider: P,
-) -> DetectProxyResult<Option<Address>>
+) -> DetectProxyResult<Option<ProxyType>>
 where
     N: Network,
 {
     let call_0 = <N as Network>::TransactionRequest::default()
         .with_to(address)
-        .with_input(EIP_SAFE_INTERFACE[0]);
+        .with_input(SAFE_INTERFACE[0]);
 
     if let Ok(value) = provider.call(&call_0).await {
         let b256: B256 = B256::from_slice(&value);
-        return Ok(Some(u256_to_address(b256.into())));
+        return Ok(Some(ProxyType::Safe(u256_to_address(b256.into()))));
     };
 
     Ok(None)