From fdec0f459911613909fee64c650bdc0ec51d3905 Mon Sep 17 00:00:00 2001
From: Tony Ke <tonyke@fuzz.land>
Date: Tue, 23 Jan 2024 12:29:30 +0800
Subject: [PATCH] add chain id to the cache key

---
 src/main.rs  | 36 ++++++++++++++++++++++--------------
 src/utils.rs | 19 +++++++++++++++++++
 2 files changed, 41 insertions(+), 14 deletions(-)
 create mode 100644 src/utils.rs

diff --git a/src/main.rs b/src/main.rs
index cd1e24e..13bb003 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -14,16 +14,19 @@ use crate::rpc_cache_handler::RpcCacheHandler;
 
 mod cli;
 mod rpc_cache_handler;
+mod utils;
 
 struct ChainState {
     rpc_url: Url,
+    id: u64,
     cache_entries: HashMap<String, CacheEntry>,
 }
 
 impl ChainState {
-    fn new(rpc_url: Url) -> Self {
+    fn new(rpc_url: Url, chain_id: u64) -> Self {
         Self {
             rpc_url,
+            id: chain_id,
             cache_entries: Default::default(),
         }
     }
@@ -76,6 +79,7 @@ async fn request_rpc(
 
 fn read_cache(
     redis_con: &mut r2d2::PooledConnection<redis::Client>,
+    chain_id: u64,
     handler: &dyn RpcCacheHandler,
     method: &str,
     params: &Value,
@@ -85,7 +89,7 @@ fn read_cache(
         .context("fail to extract cache key")?;
 
     let cache_key = match cache_key {
-        Some(cache_key) => format!("{}:{}", method, cache_key),
+        Some(cache_key) => format!("{chain_id}:{method}:{cache_key}"),
         None => return Ok(CacheStatus::NotAvailable),
     };
 
@@ -102,11 +106,11 @@ fn read_cache(
 
 #[actix_web::post("/{chain}")]
 async fn rpc_call(
-    path: web::Path<(String, )>,
+    path: web::Path<(String,)>,
     data: web::Data<AppState>,
     body: web::Json<Value>,
 ) -> Result<HttpResponse, Error> {
-    let (chain, ) = path.into_inner();
+    let (chain,) = path.into_inner();
     let chain_state = data
         .chains
         .get(&chain.to_uppercase())
@@ -122,8 +126,6 @@ async fn rpc_call(
     let mut uncached_requests = HashMap::new();
     let mut ids_in_original_order = vec![];
 
-
-
     for mut request in requests {
         let id = match request["id"].take() {
             Value::Number(n) if n.as_u64().is_some() => n.as_u64().unwrap(),
@@ -154,6 +156,7 @@ async fn rpc_call(
 
         let result = read_cache(
             &mut redis_con,
+            chain_state.id,
             cache_entry.handler.as_ref(),
             &method,
             &params,
@@ -199,11 +202,11 @@ async fn rpc_call(
             chain_state.rpc_url.clone(),
             &request_body,
         )
-            .await
-            .map_err(|err| {
-                tracing::error!("fail to make rpc request because: {}", err);
-                error::ErrorInternalServerError(format!("fail to make rpc request because: {}", err))
-            })?;
+        .await
+        .map_err(|err| {
+            tracing::error!("fail to make rpc request because: {}", err);
+            error::ErrorInternalServerError(format!("fail to make rpc request because: {}", err))
+        })?;
 
         let result_values = match rpc_result {
             Value::Array(v) => v,
@@ -289,8 +292,9 @@ async fn main() -> std::io::Result<()> {
 
     let redis_client = redis::Client::open(arg.redis_url).expect("Failed to create Redis client");
     let redis_con_pool = r2d2::Pool::builder()
-            .max_size(300)
-            .build(redis_client).expect("Failed to create Redis connection pool");
+        .max_size(300)
+        .build(redis_client)
+        .expect("Failed to create Redis connection pool");
 
     let mut app_state = AppState {
         chains: Default::default(),
@@ -305,7 +309,11 @@ async fn main() -> std::io::Result<()> {
     for (name, rpc_url) in arg.endpoints.iter() {
         tracing::info!("Adding endpoint {} linked to {}", name, rpc_url);
 
-        let mut chain_state = ChainState::new(rpc_url.clone());
+        let chain_id = utils::get_chain_id(rpc_url.as_str())
+            .await
+            .expect("fail to get chain id");
+
+        let mut chain_state = ChainState::new(rpc_url.clone(), chain_id);
 
         for factory in &handler_factories {
             let handler = factory();
diff --git a/src/utils.rs b/src/utils.rs
new file mode 100644
index 0000000..427fd3b
--- /dev/null
+++ b/src/utils.rs
@@ -0,0 +1,19 @@
+use serde_json::{json, Value};
+
+pub async fn get_chain_id(rpc_url: &str) -> anyhow::Result<u64> {
+    let request_payload = json!({
+        "jsonrpc": "2.0",
+        "method": "eth_chainId",
+        "params": [],
+        "id": 1
+    });
+
+    let client = reqwest::Client::new();
+    let response = client.post(rpc_url).json(&request_payload).send().await?;
+
+    let json: Value = response.json().await?;
+    match json["result"].as_str() {
+        Some(chain_id) => Ok(u64::from_str_radix(&chain_id[2..], 16)?),
+        None => Err(anyhow::anyhow!("fail to get chain id: {json}")),
+    }
+}