From c95f48d424940fdd83076be34aa5c09fceb54d58 Mon Sep 17 00:00:00 2001
From: Louise Poole <louise@datarevenue.com>
Date: Tue, 15 Oct 2024 13:59:54 +0200
Subject: [PATCH] feat: implement get_code_for_address util function

---
 src/protocol/mod.rs      |   1 +
 src/protocol/vm/mod.rs   |   1 +
 src/protocol/vm/utils.rs | 128 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 130 insertions(+)
 create mode 100644 src/protocol/vm/mod.rs
 create mode 100644 src/protocol/vm/utils.rs

diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs
index 55f161fd..9c373cc9 100644
--- a/src/protocol/mod.rs
+++ b/src/protocol/mod.rs
@@ -10,6 +10,7 @@ pub mod state;
 pub mod tycho;
 pub mod uniswap_v2;
 pub mod uniswap_v3;
+mod vm;
 
 /// A trait for converting types to and from `Bytes`.
 ///
diff --git a/src/protocol/vm/mod.rs b/src/protocol/vm/mod.rs
new file mode 100644
index 00000000..95adf51c
--- /dev/null
+++ b/src/protocol/vm/mod.rs
@@ -0,0 +1 @@
+mod utils;
diff --git a/src/protocol/vm/utils.rs b/src/protocol/vm/utils.rs
new file mode 100644
index 00000000..5912ba4f
--- /dev/null
+++ b/src/protocol/vm/utils.rs
@@ -0,0 +1,128 @@
+// TODO: remove skip for clippy dead_code check
+#![allow(dead_code)]
+use reqwest::{blocking::Client, StatusCode};
+use serde_json::json;
+use std::{env, time::Duration};
+
+#[derive(Debug)]
+pub enum RpcError {
+    Http(reqwest::Error),
+    Rpc(String, StatusCode),
+    InvalidResponse(String),
+}
+
+fn exec_rpc_method(
+    url: &str,
+    method: &str,
+    params: Vec<serde_json::Value>,
+    timeout: u64,
+) -> Result<serde_json::Value, RpcError> {
+    let client = Client::new();
+    let payload = json!({
+        "jsonrpc": "2.0",
+        "method": method,
+        "params": params,
+        "id": 1,
+    });
+
+    let response = client
+        .post(url)
+        .json(&payload)
+        .timeout(Duration::from_secs(timeout))
+        .send()
+        .map_err(RpcError::Http)?;
+
+    if response.status().is_client_error() || response.status().is_server_error() {
+        let status = response.status();
+        let body = response
+            .text()
+            .unwrap_or_else(|_| "Unknown error".to_string());
+        return Err(RpcError::Rpc(body, status));
+    }
+
+    let data: serde_json::Value = response
+        .json()
+        .map_err(RpcError::Http)?;
+
+    if let Some(result) = data.get("result") {
+        Ok(result.clone())
+    } else if let Some(error) = data.get("error") {
+        Err(RpcError::InvalidResponse(format!(
+            "RPC failed to call {} with Error: {}",
+            method, error
+        )))
+    } else {
+        Ok(serde_json::Value::Null)
+    }
+}
+
+pub fn get_code_for_address(
+    address: &str,
+    connection_string: Option<String>,
+) -> Result<Option<Vec<u8>>, RpcError> {
+    // Get the connection string, defaulting to the RPC_URL environment variable
+    let connection_string = connection_string.or_else(|| env::var("RPC_URL").ok());
+
+    let connection_string = match connection_string {
+        Some(url) => url,
+        None => {
+            return Err(RpcError::InvalidResponse(
+                "RPC_URL environment variable is not set".to_string(),
+            ))
+        }
+    };
+
+    let method = "eth_getCode";
+    let params = vec![json!(address), json!("latest")];
+
+    match exec_rpc_method(&connection_string, method, params, 240) {
+        Ok(code) => {
+            if let Some(code_str) = code.as_str() {
+                let code_bytes = hex::decode(&code_str[2..]).map_err(|_| {
+                    RpcError::InvalidResponse(format!(
+                        "Failed to decode hex string for address {}",
+                        address
+                    ))
+                })?;
+                Ok(Some(code_bytes))
+            } else {
+                Ok(None)
+            }
+        }
+        Err(e) => {
+            println!("Error fetching code for address {}: {:?}", address, e);
+            Err(e)
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use dotenv::dotenv;
+
+    use super::*;
+
+    #[test]
+    #[cfg_attr(not(feature = "network_tests"), ignore)]
+    fn test_get_code_for_address() {
+        let rpc_url = env::var("ETH_RPC_URL").unwrap_or_else(|_| {
+            dotenv().expect("Missing .env file");
+            env::var("ETH_RPC_URL").expect("Missing ETH_RPC_URL in .env file")
+        });
+
+        let address = "0x88e6A0c2dDD26FEEb64F039a2c41296FcB3f5640";
+        let result = get_code_for_address(address, Some(rpc_url));
+
+        assert!(result.is_ok(), "Network call should not fail");
+
+        let code_bytes = result.unwrap();
+        match code_bytes {
+            Some(bytes) => {
+                assert!(!bytes.is_empty(), "Code should not be empty");
+            }
+            None => {
+                panic!("There should be some code for the address");
+            }
+        }
+    }
+}