From d20833783fb7fa3da3166c9a62f7b7a5f49fac46 Mon Sep 17 00:00:00 2001
From: Qiwei Yang <yangqiwei97@gmail.com>
Date: Wed, 23 Oct 2024 16:38:57 +0800
Subject: [PATCH] add rpcs

---
 Networking/Sources/Networking/Peer.swift       |  5 +++++
 .../Node/NetworkingProtocol/Network.swift      |  4 ++++
 .../NetworkingProtocol/NetworkManager.swift    |  4 ++++
 Node/Sources/Node/NodeDataSource.swift         |  3 +--
 RPC/Sources/RPC/Handlers/ChainHandler.swift    | 18 ++++++++++++++++++
 5 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/Networking/Sources/Networking/Peer.swift b/Networking/Sources/Networking/Peer.swift
index 61513ad1..15a50c3b 100644
--- a/Networking/Sources/Networking/Peer.swift
+++ b/Networking/Sources/Networking/Peer.swift
@@ -157,6 +157,11 @@ public final class Peer<Handler: StreamHandler>: Sendable {
             }
         }
     }
+
+    // there should be only one connection per peer
+    public func getPeersCount() -> Int {
+        impl.connections.value.byId.values.count
+    }
 }
 
 final class PeerImpl<Handler: StreamHandler>: Sendable {
diff --git a/Node/Sources/Node/NetworkingProtocol/Network.swift b/Node/Sources/Node/NetworkingProtocol/Network.swift
index 6a94abb2..5a0ee275 100644
--- a/Node/Sources/Node/NetworkingProtocol/Network.swift
+++ b/Node/Sources/Node/NetworkingProtocol/Network.swift
@@ -77,6 +77,10 @@ public final class Network: Sendable {
     public func listenAddress() throws -> NetAddr {
         try peer.listenAddress()
     }
+
+    public func getPeersCount() -> Int {
+        peer.getPeersCount()
+    }
 }
 
 struct HandlerDef: StreamHandler {
diff --git a/Node/Sources/Node/NetworkingProtocol/NetworkManager.swift b/Node/Sources/Node/NetworkingProtocol/NetworkManager.swift
index dc8cc320..216e4eec 100644
--- a/Node/Sources/Node/NetworkingProtocol/NetworkManager.swift
+++ b/Node/Sources/Node/NetworkingProtocol/NetworkManager.swift
@@ -104,6 +104,10 @@ public final class NetworkManager: Sendable {
             )
         }
     }
+
+    public func getPeersCount() -> Int {
+        network.getPeersCount()
+    }
 }
 
 struct HandlerImpl: NetworkProtocolHandler {
diff --git a/Node/Sources/Node/NodeDataSource.swift b/Node/Sources/Node/NodeDataSource.swift
index b94b279e..996f3e8c 100644
--- a/Node/Sources/Node/NodeDataSource.swift
+++ b/Node/Sources/Node/NodeDataSource.swift
@@ -30,7 +30,6 @@ public final class NodeDataSource: DataSource {
     }
 
     public func getPeersCount() async throws -> Int {
-        // TODO: implememnt
-        0
+        networkManager.getPeersCount()
     }
 }
diff --git a/RPC/Sources/RPC/Handlers/ChainHandler.swift b/RPC/Sources/RPC/Handlers/ChainHandler.swift
index 2ec2f895..776e4794 100644
--- a/RPC/Sources/RPC/Handlers/ChainHandler.swift
+++ b/RPC/Sources/RPC/Handlers/ChainHandler.swift
@@ -10,6 +10,7 @@ struct ChainHandler {
 
         return [
             "chain_getBlock": handler.getBlock,
+            "chain_getState": handler.getState,
         ]
     }
 
@@ -32,4 +33,21 @@ struct ChainHandler {
             ]
         }
     }
+
+    func getState(request: JSONRequest) async throws -> any Encodable {
+        let hash = request.params?["hash"] as? String
+        if let hash {
+            guard let data = Data(fromHexString: hash), let data32 = Data32(data) else {
+                throw JSONError(code: -32602, message: "Invalid block hash")
+            }
+            let state = try await source.getState(hash: data32)
+            // return state root for now
+            return state?.stateRoot.description
+        } else {
+            // return best block state by default
+            let block = try await source.getBestBlock()
+            let state = try await source.getState(hash: block.hash)
+            return state?.stateRoot.description
+        }
+    }
 }