allow multiple chain

sentioxyz · Aug 13, 2024 · f72d956 · f72d956
1 parent db362d1
commit f72d956
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 7 deletions.
diff --git a/aptos-move/aptos-tracer/src/config.rs b/aptos-move/aptos-tracer/src/config.rs
@@ -1,7 +1,7 @@
 // Copyright © Aptos Foundation
 // SPDX-License-Identifier: Apache-2.0
 
-use std::path::PathBuf;
+use std::{collections::HashMap, path::PathBuf};
 use serde::{Deserialize, Serialize};
 
 #[derive(Clone, Debug, Deserialize, Serialize)]
@@ -26,6 +26,9 @@ pub struct DebuggerServerConfig {
     #[serde(default = "DebuggerServerConfig::default_rest_endpoint")]
     pub rest_endpoint: String,
 
+    #[serde(default = "DebuggerServerConfig::default_rest_endpoint_map")]
+    pub rest_endpoint_map: HashMap<u16, String>,
+
     /// use db or not
     #[serde(default = "DebuggerServerConfig::default_use_db")]
     pub use_db: bool,
@@ -43,6 +46,7 @@ impl DebuggerServerConfig {
             listen_port: DebuggerServerConfig::default_listen_port(),
             db_path: DebuggerServerConfig::default_db_path(),
             rest_endpoint: DebuggerServerConfig::default_rest_endpoint(),
+            rest_endpoint_map: DebuggerServerConfig::default_rest_endpoint_map(),
             use_db: DebuggerServerConfig::default_use_db(),
             sentio_endpoint: DebuggerServerConfig::default_sentio_endpoint(),
         }
@@ -56,6 +60,10 @@ impl DebuggerServerConfig {
         self.rest_endpoint = rest_endpoint
     }
 
+    pub fn set_rest_endpoints(&mut self, rest_endpoint_map: HashMap<u16, String>) {
+        self.rest_endpoint_map = rest_endpoint_map
+    }
+
     pub fn set_use_db(&mut self, use_db: bool) {
         self.use_db = use_db
     }
@@ -82,6 +90,13 @@ impl DebuggerServerConfig {
         "https://fullnode.mainnet.aptoslabs.com/v1".to_string()
     }
 
+    fn default_rest_endpoint_map() -> HashMap<u16, String> {
+        HashMap::from([
+            (1, "https://fullnode.mainnet.aptoslabs.com/v1".to_string()),
+            (2001, "https://aptos.testnet.suzuka.movementlabs.xyz/v1".to_string())
+        ])
+    }
+
     fn default_use_db() -> bool {
         false
     }

diff --git a/aptos-move/aptos-tracer/src/lib.rs b/aptos-move/aptos-tracer/src/lib.rs
@@ -212,7 +212,7 @@ impl SyncAptosTracer {
     pub fn trace_transaction(
         &self,
         txn_hash: String,
-        chain_id: u8,
+        chain_id: u16,
     ) -> Result<CallTraceWithSource> {
         let txn_data = self.debugger.get_transaction_by_hash(txn_hash)?;
         let txn = txn_data.transaction;

diff --git a/aptos-move/aptos-tracer/src/main.rs b/aptos-move/aptos-tracer/src/main.rs
@@ -7,6 +7,7 @@ use aptos_rest_client::Client;
 use clap::{Parser, Subcommand};
 use std::path::PathBuf;
 use url::Url;
+use std::collections::HashMap;
 
 #[derive(Subcommand)]
 pub enum Target {
@@ -15,7 +16,7 @@ pub enum Target {
     /// Use a local db instance to serve as query endpoint.
     DB { path: PathBuf, listen_address: Option<String>, listen_port: Option<u16> },
     /// Use a full node's rest api as query endpoint and run as a server.
-    ServerBasedOnRest { endpoint: String, listen_address: Option<String>, listen_port: Option<u16> },
+    ServerBasedOnRest { endpoints: String, listen_address: Option<String>, listen_port: Option<u16> },
 }
 #[derive(Parser)]
 pub struct Argument {
@@ -56,12 +57,19 @@ async fn main() -> Result<()> {
             }
             run_debugger_server(config).await
         },
-        Target::ServerBasedOnRest { endpoint, listen_address, listen_port } => {
+        Target::ServerBasedOnRest { endpoints, listen_address, listen_port } => {
             // run as a server if the target is DB
             let mut config = DebuggerServerConfig::default();
-            config.set_rest_endpoint(endpoint);
             config.set_use_db(false);
             config.set_sentio_endpoint(args.sentio_endpoint);
+
+            let mut endpoint_map: HashMap<u16, String> = HashMap::new(); 
+            let chain_to_endpoint: Vec<&str> = endpoints.split(',').collect();
+            for i in chain_to_endpoint.iter() {
+                let t: Vec<&str> = i.split('=').collect();
+                endpoint_map.insert(t[0].parse::<u16>().unwrap(), t[1].to_string());
+            }
+            config.set_rest_endpoints(endpoint_map);
             if let Some(address) = listen_address {
                 config.listen_address = address;
             }

diff --git a/aptos-move/aptos-tracer/src/server.rs b/aptos-move/aptos-tracer/src/server.rs
@@ -14,13 +14,17 @@ use aptos_logger::{error, info};
 use aptos_rest_client::Client;
 
 #[handler]
-async fn call_trace(Path((chain_id, hash)): Path<(u8, String)>, config: Data<&DebuggerServerConfig>) -> Result<Json<CallTraceWithSource>> {
+async fn call_trace(Path((chain_id, hash)): Path<(u16, String)>, config: Data<&DebuggerServerConfig>) -> Result<Json<CallTraceWithSource>> {
     let mut tracer;
     if config.use_db {
         tracer = SyncAptosTracer::db(config.db_path.to_str().unwrap(), config.clone().sentio_endpoint);
     } else {
+        let endpoint = match config.rest_endpoint_map.get(&chain_id) {
+            Some(endpoint) => endpoint,
+            None => &config.rest_endpoint
+        };
         tracer = SyncAptosTracer::rest_client(
-            Client::new(Url::parse(config.rest_endpoint.as_str()).unwrap()),
+            Client::new(Url::parse(endpoint.as_str()).unwrap()),
             config.clone().sentio_endpoint);
     }