argus_rs/
lib.rs

1//! Argus is a minimal, blazing fast contract storage introspection tool written in rust.
2//!
3//! It is designed to extract storage mapping slots from a contract, without needing to know the
4//! contract's source code.
5
6use alloy::{
7    network::TransactionBuilder,
8    primitives::{address, bytes, Address, Bytes, U256},
9    providers::{ext::TraceApi, ProviderBuilder},
10    rpc::types::{trace::parity::TraceType, TransactionRequest},
11};
12use eyre::{bail, OptionExt, Result};
13use std::{collections::VecDeque, fmt::Debug};
14
15/// The `Introspector` struct is used to introspect a contract's storage and determine which
16/// standards it adheres to.
17#[derive(Debug, Clone)]
18pub struct Introspector {
19    /// The contract address to introspect.
20    pub contract_address: Address,
21    /// The provider to use for querying the contract. Supports WebSocket, IPC, and HTTP
22    /// transports.
23    rpc_url: String,
24}
25
26/// The `IntrospectResult` struct is used to store the results of an introspection.
27#[derive(Debug, Clone, Default, Eq, PartialEq)]
28pub struct IntrospectResult {
29    /// The balance slot of the contract.
30    pub balance_slot: Option<U256>,
31    /// The allowance slot of the contract.
32    pub allowance_slot: Option<U256>,
33    /// The token approvals slot of the contract.
34    pub token_approvals_slot: Option<U256>,
35    /// The operator approvals slot of the contract.
36    pub operator_approvals_slot: Option<U256>,
37    /// The erc1155 balance slot of the contract.
38    pub erc_1155_balance_slot: Option<U256>,
39}
40
41const INTROSPECT_ADDRESS: Address = address!("000000000000000000696c6f76656f7474657273");
42
43impl Introspector {
44    /// Creates a new `Introspector` instance with the given `contract_address` and `provider`.
45    pub fn try_new(contract_address: Address, rpc_url: impl Into<String>) -> Self {
46        Self { contract_address, rpc_url: rpc_url.into() }
47    }
48
49    /// Run an introspection on the contract and return the results.
50    ///
51    /// Performs the following calls, returning the base slot for each:
52    /// - balanceOf(address)
53    /// - allowance(address, address)
54    /// - getApproved(address)
55    /// - isApprovedForAll(address, address)
56    /// - balanceOf(address, uint256)
57    pub async fn run(&self) -> Result<IntrospectResult> {
58        let balance_slot_future = self.get_balance_slot();
59        let allowance_slot_future = self.get_allowance_slot();
60        let token_approvals_slot_future = self.get_token_approvals_slot();
61        let operator_approvals_slot_future = self.get_operator_approvals_slot();
62        let erc_1155_balance_slot_future = self.get_erc_1155_balance_slot();
63
64        // Run all calls concurrently.
65        let (
66            balance_slot_result,
67            allowance_slot_result,
68            token_approvals_slot_result,
69            operator_approvals_slot_result,
70            erc_1155_balance_slot_result,
71        ) = futures::future::join5(
72            balance_slot_future,
73            allowance_slot_future,
74            token_approvals_slot_future,
75            operator_approvals_slot_future,
76            erc_1155_balance_slot_future,
77        )
78        .await;
79
80        let maybe_erc_20 = balance_slot_result.is_ok() && allowance_slot_result.is_ok();
81        let maybe_erc_721 = balance_slot_result.is_ok() &&
82            token_approvals_slot_result.is_ok() &&
83            operator_approvals_slot_result.is_ok();
84        let maybe_erc_1155 =
85            operator_approvals_slot_result.is_ok() && erc_1155_balance_slot_result.is_ok();
86
87        if maybe_erc_20 {
88            Ok(IntrospectResult {
89                balance_slot: balance_slot_result.ok(),
90                allowance_slot: allowance_slot_result.ok(),
91                ..Default::default()
92            })
93        } else if maybe_erc_721 {
94            Ok(IntrospectResult {
95                balance_slot: balance_slot_result.ok(),
96                token_approvals_slot: token_approvals_slot_result.ok(),
97                operator_approvals_slot: operator_approvals_slot_result.ok(),
98                ..Default::default()
99            })
100        } else if maybe_erc_1155 {
101            Ok(IntrospectResult {
102                operator_approvals_slot: operator_approvals_slot_result.ok(),
103                erc_1155_balance_slot: erc_1155_balance_slot_result.ok(),
104                ..Default::default()
105            })
106        } else {
107            Ok(IntrospectResult {
108                balance_slot: balance_slot_result.ok(),
109                allowance_slot: allowance_slot_result.ok(),
110                token_approvals_slot: token_approvals_slot_result.ok(),
111                operator_approvals_slot: operator_approvals_slot_result.ok(),
112                erc_1155_balance_slot: erc_1155_balance_slot_result.ok(),
113            })
114        }
115    }
116
117    /// Get the balance slot of the contract by tracing a call to `balanceOf(address)`.
118    pub async fn get_balance_slot(&self) -> Result<U256> {
119        let calldata =
120            bytes!("70a08231000000000000000000000000000000000000000000696c6f76656f7474657273"); // balanceOf(INTROSPECT_ADDRESS)
121        return self.extract_slot(calldata).await;
122    }
123
124    /// Get the allowance slot of the contract by tracing a call to `allowance(address, address)`.
125    pub async fn get_allowance_slot(&self) -> Result<U256> {
126        let calldata =
127            bytes!("dd62ed3e000000000000000000000000000000000000000000696c6f76656f7474657273000000000000000000000000000000000000000000696c6f76656f7474657273"); // allowance(INTROSPECT_ADDRESS, INTROSPECT_ADDRESS)
128        return self.extract_slot(calldata).await;
129    }
130
131    /// Get the token approvals slot of the contract by tracing a call to `getApproved(address)`.
132    pub async fn get_token_approvals_slot(&self) -> Result<U256> {
133        let calldata =
134            bytes!("081812fc000000000000000000000000000000000000000000696c6f76656f7474657273"); // getApproved(INTROSPECT_ADDRESS)
135        return self.extract_slot(calldata).await;
136    }
137
138    /// Get the operator approvals slot of the contract by tracing a call to
139    /// `isApprovedForAll(address, address)`.
140    pub async fn get_operator_approvals_slot(&self) -> Result<U256> {
141        let calldata =
142            bytes!("e985e9c5000000000000000000000000000000000000000000696c6f76656f7474657273000000000000000000000000000000000000000000696c6f76656f7474657273"); // isApprovedForAll(INTROSPECT_ADDRESS, INTROSPECT_ADDRESS)
143        return self.extract_slot(calldata).await;
144    }
145
146    /// Get the erc1155 balance slot of the contract by tracing a call to `balanceOf(address,
147    /// uint256)`.
148    pub async fn get_erc_1155_balance_slot(&self) -> Result<U256> {
149        let calldata =
150            bytes!("00fdd58e000000000000000000000000000000000000000000696c6f76656f7474657273000000000000000000000000000000000000000000696c6f76656f7474657273"); // balanceOf(INTROSPECT_ADDRESS, INTROSPECT_ADDRESS)
151        return self.extract_slot(calldata).await;
152    }
153
154    /// Extract the slot from the given calldata.
155    ///
156    /// This function traces the call to the contract and extracts the slot
157    /// from the trace by monitoring the stack and memory, and looking for sha3 operations.
158    ///
159    /// Mapping slots are stored in the format:
160    /// `keccak256(bytes32(mapping_key) + bytes32(slot))`
161    ///
162    /// For example, the mapping slot for `balanceOf(0x000...123)` is stored as:
163    /// `keccak256(0x00000000000000000000000000000000000000000000000000000000000001230000000000000000000000000000000000000000000000000000000000000001)`
164    ///
165    /// In this case, the mapping key (address) is `0x000...123` and the slot is `1`, so we return
166    /// `1`.
167    async fn extract_slot(&self, calldata: Bytes) -> Result<U256> {
168        let provider = ProviderBuilder::new().on_builtin(&self.rpc_url).await?;
169
170        let mut tx = TransactionRequest::default()
171            .with_from(INTROSPECT_ADDRESS)
172            .with_to(self.contract_address)
173            .with_input(calldata.clone());
174        tx.input.data = Some(calldata);
175
176        let result = provider
177            .trace_call(&tx, &[TraceType::VmTrace, TraceType::Trace])
178            .await?
179            .vm_trace
180            .ok_or_eyre("vm trace not found")?
181            .ops;
182
183        // A naive stack impl. We don't track this accurately, but it's good enough for our
184        // purposes.
185        let mut stack: VecDeque<U256> = VecDeque::new();
186        let mut memory: Vec<u8> = Vec::new();
187
188        for instruction in result {
189            if let Some(executed) = &instruction.ex {
190                if let Some(memory_delta) = &executed.mem {
191                    let offset = memory_delta.off;
192                    let data = memory_delta.data.to_vec();
193
194                    // expand memory if necessary
195                    if memory.len() < offset + data.len() {
196                        memory.resize(offset + data.len(), 0);
197                    }
198
199                    // copy data into memory at the given offset
200                    memory[offset..offset + data.len()].copy_from_slice(&data);
201                }
202
203                stack.extend(executed.push.iter());
204            }
205
206            if let Some(opcode) = &instruction.op {
207                // If we see a SHA3, we can extract the contents that are being hashed.
208                // We store these contents in a
209                if ["KECCAK256", "SHA3"].contains(&opcode.as_str()) {
210                    let _ = stack.pop_back().ok_or_eyre("stack underflow")?; // pop off the hashed value
211                    let offset: usize =
212                        stack.pop_back().ok_or_eyre("stack underflow")?.try_into()?;
213                    let size: usize = stack.pop_back().ok_or_eyre("stack underflow")?.try_into()?;
214
215                    let data = memory[offset..offset + size].to_vec();
216
217                    // is data of the format bytes32(INTROSPECT_ADDRESS).concat(bytes32(N))?
218                    if data.len() == 64 && &data[12..32] == INTROSPECT_ADDRESS.as_slice() {
219                        // get the second 32 bytes, converted to U256. this is the mapping slot.
220                        let balance_slot: U256 = U256::from_be_slice(&data[32..64]);
221                        return Ok(balance_slot);
222                    }
223                }
224            }
225        }
226
227        bail!("failed to get balance slot")
228    }
229}
230
231#[cfg(test)]
232mod tests {
233    use super::*;
234    use alloy::primitives::address;
235
236    #[tokio::test]
237    async fn test_introspect_erc20() {
238        let rpc_url = std::env::var("RPC_URL").unwrap_or_else(|_| {
239            println!("RPC_URL not set, skipping test");
240            std::process::exit(0);
241        });
242
243        let contract_address = address!("c02aaa39b223fe8d0a0e5c4f27ead9083c756cc2");
244        let introspector = Introspector::try_new(contract_address, rpc_url);
245
246        let res = introspector.run().await.expect("failed to run introspection");
247        assert_eq!(
248            res,
249            IntrospectResult {
250                balance_slot: Some(U256::from(3)),
251                allowance_slot: Some(U256::from(4)),
252                ..Default::default()
253            }
254        );
255    }
256
257    #[tokio::test]
258    async fn test_introspect_erc721() {
259        let rpc_url = std::env::var("RPC_URL").unwrap_or_else(|_| {
260            println!("RPC_URL not set, skipping test");
261            std::process::exit(0);
262        });
263
264        let contract_address = address!("bc4ca0eda7647a8ab7c2061c2e118a18a936f13d");
265        let introspector = Introspector::try_new(contract_address, rpc_url);
266
267        let res = introspector.run().await.expect("failed to run introspection");
268        assert_eq!(
269            res,
270            IntrospectResult {
271                balance_slot: Some(U256::from(1)),
272                token_approvals_slot: Some(U256::from(3)),
273                operator_approvals_slot: Some(U256::from(5)),
274                ..Default::default()
275            }
276        );
277    }
278
279    #[tokio::test]
280    async fn test_introspect_erc1155() {
281        let rpc_url = std::env::var("RPC_URL").unwrap_or_else(|_| {
282            println!("RPC_URL not set, skipping test");
283            std::process::exit(0);
284        });
285
286        let contract_address = address!("c552292732f7a9a4a494da557b47bc01e01722df");
287        let introspector = Introspector::try_new(contract_address, rpc_url);
288
289        let res = introspector.run().await.expect("failed to run introspection");
290        assert_eq!(
291            res,
292            IntrospectResult {
293                operator_approvals_slot: Some(U256::from(1)),
294                erc_1155_balance_slot: Some(U256::from(0)),
295                ..Default::default()
296            }
297        );
298    }
299}
argus_rs/lib.rs

argus_rs/
lib.rs
