Merge pull request #94 from vacp2p/schouhy/add-get-account-balance-rpc-method

Add get account balance method to sequencer RPC server

Cargo.lock

@@ -196,11 +196,13 @@ checksum = "a27e8fe9ba4ae613c21f677c2cfaf0696c3744030c6f485b34634e502d6bb379"
 dependencies = [
  "actix-codec",
  "actix-http",
+ "actix-macros",
  "actix-router",
  "actix-rt",
  "actix-server",
  "actix-service",
  "actix-utils",
+ "actix-web-codegen",
  "ahash 0.7.8",
  "bytes",
  "bytestring",
@@ -225,6 +227,18 @@ dependencies = [
  "url",
 ]
 
+[[package]]
+name = "actix-web-codegen"
+version = "4.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f591380e2e68490b5dfaf1dd1aa0ebe78d84ba7067078512b4ea6e4492d622b8"
+dependencies = [
+ "actix-router",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.101",
+]
+
 [[package]]
 name = "actix_derive"
 version = "0.6.2"
@@ -4379,6 +4393,7 @@ dependencies = [
  "serde",
  "serde_json",
  "storage",
+ "tempfile",
  "tokio",
 ]
 

Cargo.toml

@@ -68,6 +68,7 @@ version = "1.0.60"
 
 [workspace.dependencies.actix-web]
 default-features = false
+features = ["macros"]
 version = "=4.1.0"
 
 [workspace.dependencies.clap]

common/src/rpc_primitives/requests.rs

@@ -34,12 +34,18 @@ pub struct GetGenesisIdRequest {}
 #[derive(Serialize, Deserialize, Debug)]
 pub struct GetLastBlockRequest {}
 
+#[derive(Serialize, Deserialize, Debug)]
+pub struct GetAccountBalanceRequest {
+    pub address: String,
+}
+
 parse_request!(HelloRequest);
 parse_request!(RegisterAccountRequest);
 parse_request!(SendTxRequest);
 parse_request!(GetBlockDataRequest);
 parse_request!(GetGenesisIdRequest);
 parse_request!(GetLastBlockRequest);
+parse_request!(GetAccountBalanceRequest);
 
 #[derive(Serialize, Deserialize, Debug)]
 pub struct HelloResponse {
@@ -70,3 +76,8 @@ pub struct GetGenesisIdResponse {
 pub struct GetLastBlockResponse {
     pub last_block: u64,
 }
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct GetAccountBalanceResponse {
+    pub balance: u64,
+}

sequencer_rpc/Cargo.toml

@@ -13,6 +13,7 @@ actix.workspace = true
 actix-cors.workspace = true
 futures.workspace = true
 hex.workspace = true
+tempfile.workspace = true
 
 actix-web.workspace = true
 tokio.workspace = true

sequencer_rpc/src/net_utils.rs

@@ -18,7 +18,7 @@ pub const SHUTDOWN_TIMEOUT_SECS: u64 = 10;
 
 pub const NETWORK: &str = "network";
 
-fn rpc_handler(
+pub(crate) fn rpc_handler(
     message: web::Json<Message>,
     handler: web::Data<JsonHandler>,
 ) -> impl Future<Output = Result<HttpResponse, HttpError>> {

sequencer_rpc/src/process.rs

@@ -5,6 +5,7 @@ use common::rpc_primitives::{
     errors::RpcError,
     message::{Message, Request},
     parser::RpcRequest,
+    requests::{GetAccountBalanceRequest, GetAccountBalanceResponse},
 };
 
 use common::rpc_primitives::requests::{
@@ -21,6 +22,7 @@ pub const SEND_TX: &str = "send_tx";
 pub const GET_BLOCK: &str = "get_block";
 pub const GET_GENESIS: &str = "get_genesis";
 pub const GET_LAST_BLOCK: &str = "get_last_block";
+pub const GET_ACCOUNT_BALANCE: &str = "get_account_balance";
 
 pub const HELLO_FROM_SEQUENCER: &str = "HELLO_FROM_SEQUENCER";
 
@@ -131,6 +133,26 @@ impl JsonHandler {
         respond(helperstruct)
     }
 
+    /// Returns the balance of the account at the given address.
+    /// The address must be a valid hex string of the correct length.
+    async fn process_get_account_balance(&self, request: Request) -> Result<Value, RpcErr> {
+        let get_account_req = GetAccountBalanceRequest::parse(Some(request.params))?;
+        let address_bytes = hex::decode(get_account_req.address)
+            .map_err(|_| RpcError::invalid_params("invalid hex".to_string()))?;
+        let address = address_bytes
+            .try_into()
+            .map_err(|_| RpcError::invalid_params("invalid length".to_string()))?;
+        let balance = {
+            let state = self.sequencer_state.lock().await;
+            state.store.acc_store.get_account_balance(&address)
+        }
+        .unwrap_or(0);
+
+        let helperstruct = GetAccountBalanceResponse { balance };
+
+        respond(helperstruct)
+    }
+
     pub async fn process_request_internal(&self, request: Request) -> Result<Value, RpcErr> {
         match request.method.as_ref() {
             HELLO => self.process_temp_hello(request).await,
@@ -139,7 +161,170 @@ impl JsonHandler {
             GET_BLOCK => self.process_get_block_data(request).await,
             GET_GENESIS => self.process_get_genesis(request).await,
             GET_LAST_BLOCK => self.process_get_last_block(request).await,
+            GET_ACCOUNT_BALANCE => self.process_get_account_balance(request).await,
             _ => Err(RpcErr(RpcError::method_not_found(request.method))),
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use crate::{rpc_handler, JsonHandler};
+    use common::rpc_primitives::RpcPollingConfig;
+    use sequencer_core::{
+        config::{AccountInitialData, SequencerConfig},
+        SequencerCore,
+    };
+    use serde_json::Value;
+    use tempfile::tempdir;
+    use tokio::sync::Mutex;
+
+    fn sequencer_config_for_tests() -> SequencerConfig {
+        let tempdir = tempdir().unwrap();
+        let home = tempdir.path().to_path_buf();
+        let initial_accounts = vec![
+            AccountInitialData {
+                addr: "cafe".repeat(16).to_string(),
+                balance: 100,
+            },
+            AccountInitialData {
+                addr: "feca".repeat(16).to_string(),
+                balance: 200,
+            },
+        ];
+
+        SequencerConfig {
+            home,
+            override_rust_log: Some("info".to_string()),
+            genesis_id: 1,
+            is_genesis_random: false,
+            max_num_tx_in_block: 10,
+            block_create_timeout_millis: 1000,
+            port: 8080,
+            initial_accounts,
+        }
+    }
+
+    fn json_handler_for_tests() -> JsonHandler {
+        let config = sequencer_config_for_tests();
+        let sequencer_core = Arc::new(Mutex::new(SequencerCore::start_from_config(config)));
+
+        JsonHandler {
+            polling_config: RpcPollingConfig::default(),
+            sequencer_state: sequencer_core,
+        }
+    }
+
+    async fn call_rpc_handler_with_json(handler: JsonHandler, request_json: Value) -> Value {
+        use actix_web::{test, web, App};
+
+        let app = test::init_service(
+            App::new()
+                .app_data(web::Data::new(handler))
+                .route("/", web::post().to(rpc_handler)),
+        )
+        .await;
+
+        let req = test::TestRequest::post()
+            .uri("/")
+            .set_json(request_json)
+            .to_request();
+
+        let resp = test::call_service(&app, req).await;
+        let body = test::read_body(resp).await;
+
+        serde_json::from_slice(&body).unwrap()
+    }
+
+    #[actix_web::test]
+    async fn test_get_account_balance_for_non_existent_account() {
+        let json_handler = json_handler_for_tests();
+        let request = serde_json::json!({
+            "jsonrpc": "2.0",
+            "method": "get_account_balance",
+            "params": { "address": "efac".repeat(16) },
+            "id": 1
+        });
+        let expected_response = serde_json::json!({
+            "id": 1,
+            "jsonrpc": "2.0",
+            "result": {
+                "balance": 0
+            }
+        });
+
+        let response = call_rpc_handler_with_json(json_handler, request).await;
+
+        assert_eq!(response, expected_response);
+    }
+
+    #[actix_web::test]
+    async fn test_get_account_balance_for_invalid_hex() {
+        let json_handler = json_handler_for_tests();
+        let request = serde_json::json!({
+            "jsonrpc": "2.0",
+            "method": "get_account_balance",
+            "params": { "address": "not_a_valid_hex" },
+            "id": 1
+        });
+        let expected_response = serde_json::json!({
+            "jsonrpc": "2.0",
+            "id": 1,
+            "error": {
+                "code": -32602,
+                "message": "Invalid params",
+                "data": "invalid hex"
+            }
+        });
+        let response = call_rpc_handler_with_json(json_handler, request).await;
+
+        assert_eq!(response, expected_response);
+    }
+
+    #[actix_web::test]
+    async fn test_get_account_balance_for_invalid_length() {
+        let json_handler = json_handler_for_tests();
+        let request = serde_json::json!({
+            "jsonrpc": "2.0",
+            "method": "get_account_balance",
+            "params": { "address": "cafecafe" },
+            "id": 1
+        });
+        let expected_response = serde_json::json!({
+            "jsonrpc": "2.0",
+            "id": 1,
+            "error": {
+                "code": -32602,
+                "message": "Invalid params",
+                "data": "invalid length"
+            }
+        });
+        let response = call_rpc_handler_with_json(json_handler, request).await;
+
+        assert_eq!(response, expected_response);
+    }
+
+    #[actix_web::test]
+    async fn test_get_account_balance_for_existing_account() {
+        let json_handler = json_handler_for_tests();
+        let request = serde_json::json!({
+            "jsonrpc": "2.0",
+            "method": "get_account_balance",
+            "params": { "address": "cafe".repeat(16) },
+            "id": 1
+        });
+        let expected_response = serde_json::json!({
+            "id": 1,
+            "jsonrpc": "2.0",
+            "result": {
+                "balance": 100
+            }
+        });
+
+        let response = call_rpc_handler_with_json(json_handler, request).await;
+
+        assert_eq!(response, expected_response);
+    }
+}