feat(rpc): get state proof

Author: snormore

Cargo.lock

@@ -7548,6 +7548,7 @@ name = "lightning-types"
 version = "0.1.0"
 dependencies = [
  "anyhow",
+ "atomo",
  "bincode",
  "cid 0.10.1",
  "compile-time-run",

core/application/src/state/query.rs

@@ -29,6 +29,8 @@ use lightning_interfaces::types::{
     Service,
     ServiceId,
     ServiceRevenue,
+    StateProofKey,
+    StateProofValue,
     TotalServed,
     TransactionRequest,
     TransactionResponse,
@@ -268,4 +270,24 @@ impl SyncQueryRunnerInterface for QueryRunner {
     fn get_state_root(&self) -> Result<StateRootHash> {
         self.run(|ctx| ApplicationStateTree::get_state_root(ctx))
     }
+
+    /// Returns the state proof for a given key from the application state using the state tree.
+    fn get_state_proof(
+        &self,
+        key: StateProofKey,
+    ) -> Result<(
+        Option<StateProofValue>,
+        <ApplicationStateTree as StateTree>::Proof,
+    )> {
+        type Serde = <ApplicationStateTree as StateTree>::Serde;
+
+        self.run(|ctx| {
+            let (table, serialized_key) = key.raw::<Serde>();
+            let proof = ApplicationStateTree::get_state_proof(ctx, &table, serialized_key.clone())?;
+            let value = self
+                .run(|ctx| ctx.get_raw_value(table, &serialized_key))
+                .map(|value| key.value::<Serde>(value));
+            Ok((value, proof))
+        })
+    }
 }

core/interfaces/src/application.rs

@@ -13,10 +13,13 @@ use lightning_types::{
     ChainId,
     Committee,
     NodeIndex,
+    StateProofKey,
+    StateProofValue,
     TransactionRequest,
     TxHash,
     Value,
 };
+use merklize::trees::mpt::MptStateProof;
 use merklize::StateRootHash;
 use serde::{Deserialize, Serialize};
 
@@ -191,6 +194,12 @@ pub trait SyncQueryRunnerInterface: Clone + Send + Sync + 'static {
 
     /// Returns the state root hash from the application state.
     fn get_state_root(&self) -> Result<StateRootHash>;
+
+    /// Returns the state proof for a given key from the application state using the state tree.
+    fn get_state_proof(
+        &self,
+        key: StateProofKey,
+    ) -> Result<(Option<StateProofValue>, MptStateProof)>;
 }
 
 #[derive(Clone, Debug)]

core/rpc/Cargo.toml

@@ -32,6 +32,7 @@ reqwest.workspace = true
 once_cell = "1.19"
 clap = { version = "4.4.10", features = ["derive"] }
 
+lightning-application = { path = "../application" }
 lightning-firewall = { path = "../firewall" }
 lightning-types = { path = "../types" }
 lightning-interfaces = { path = "../interfaces" }
@@ -52,7 +53,6 @@ hex = "0.4.3"
 [dev-dependencies]
 reqwest.workspace = true
 lightning-test-utils = { path = "../test-utils" }
-lightning-application = { path = "../application" }
 lightning-fetcher = { path = "../fetcher" }
 lightning-blockstore = { path = "../blockstore" }
 lightning-blockstore-server = { path = "../blockstore-server" }

core/rpc/src/api/flk.rs

@@ -4,6 +4,7 @@ use fleek_crypto::{EthAddress, NodePublicKey};
 use hp_fixed::unsigned::HpUfixed;
 use jsonrpsee::core::{RpcResult, SubscriptionResult};
 use jsonrpsee::proc_macros::rpc;
+use lightning_application::env::ApplicationStateTree;
 use lightning_interfaces::types::{
     AccountInfo,
     Blake3Hash,
@@ -23,7 +24,8 @@ use lightning_interfaces::types::{
 };
 use lightning_interfaces::PagingParams;
 use lightning_openrpc_macros::open_rpc;
-use merklize::StateRootHash;
+use lightning_types::{StateProofKey, StateProofValue};
+use merklize::{StateRootHash, StateTree};
 
 #[open_rpc(namespace = "flk", tag = "1.0.0")]
 #[rpc(client, server, namespace = "flk")]
@@ -189,6 +191,16 @@ pub trait FleekApi {
     #[method(name = "get_state_root")]
     async fn get_state_root(&self, epoch: Option<u64>) -> RpcResult<StateRootHash>;
 
+    #[method(name = "get_state_proof")]
+    async fn get_state_proof(
+        &self,
+        key: StateProofKey,
+        epoch: Option<u64>,
+    ) -> RpcResult<(
+        Option<StateProofValue>,
+        <ApplicationStateTree as StateTree>::Proof,
+    )>;
+
     #[method(name = "send_txn")]
     async fn send_txn(&self, tx: TransactionRequest) -> RpcResult<()>;
 

core/rpc/src/logic/flk_impl.rs

@@ -5,6 +5,7 @@ use fleek_crypto::{EthAddress, NodePublicKey};
 use hp_fixed::unsigned::HpUfixed;
 use jsonrpsee::core::{RpcResult, SubscriptionResult};
 use jsonrpsee::{PendingSubscriptionSink, SubscriptionMessage};
+use lightning_application::env::ApplicationStateTree;
 use lightning_interfaces::prelude::*;
 use lightning_interfaces::types::{
     AccountInfo,
@@ -29,8 +30,9 @@ use lightning_interfaces::types::{
     Value,
 };
 use lightning_interfaces::PagingParams;
+use lightning_types::{StateProofKey, StateProofValue};
 use lightning_utils::application::QueryRunnerExt;
-use merklize::StateRootHash;
+use merklize::{StateRootHash, StateTree};
 
 use crate::api::FleekApiServer;
 use crate::error::RPCError;
@@ -388,6 +390,7 @@ impl<C: Collection> FleekApiServer for FleekApi<C> {
         Ok((sub_dag_index, self.data.query_runner.get_epoch_info().epoch))
     }
 
+    /// Returns the state root for the given epoch.
     async fn get_state_root(&self, epoch: Option<u64>) -> RpcResult<StateRootHash> {
         Ok(self
             .data
@@ -397,6 +400,24 @@ impl<C: Collection> FleekApiServer for FleekApi<C> {
             .map_err(|e| RPCError::custom(e.to_string()))?)
     }
 
+    /// Returns the state proof for a given key and epoch.
+    async fn get_state_proof(
+        &self,
+        key: StateProofKey,
+        epoch: Option<u64>,
+    ) -> RpcResult<(
+        Option<StateProofValue>,
+        <ApplicationStateTree as StateTree>::Proof,
+    )> {
+        let (value, proof) = self
+            .data
+            .query_runner(epoch)
+            .await?
+            .get_state_proof(key)
+            .map_err(|e| RPCError::custom(e.to_string()))?;
+        Ok((value, proof))
+    }
+
     async fn send_txn(&self, tx: TransactionRequest) -> RpcResult<()> {
         Ok(self
             .data

core/rpc/src/tests.rs

@@ -13,6 +13,7 @@ use hp_fixed::unsigned::HpUfixed;
 use jsonrpsee::http_client::{HttpClient, HttpClientBuilder};
 use lightning_application::app::Application;
 use lightning_application::config::Config as AppConfig;
+use lightning_application::env::ApplicationStateTree;
 use lightning_application::genesis::{Genesis, GenesisAccount, GenesisNode, GenesisNodeServed};
 use lightning_application::state::QueryRunner;
 use lightning_blockstore::blockstore::Blockstore;
@@ -38,8 +39,9 @@ use lightning_rep_collector::ReputationAggregator;
 use lightning_signer::Signer;
 use lightning_test_utils::json_config::JsonConfigProvider;
 use lightning_test_utils::keys::EphemeralKeystore;
-use lightning_types::FirewallConfig;
+use lightning_types::{AccountInfo, FirewallConfig, StateProofKey, StateProofValue};
 use lightning_utils::application::QueryRunnerExt;
+use merklize::StateProof;
 use reqwest::Client;
 use resolved_pathbuf::ResolvedPathBuf;
 use serde::{Deserialize, Serialize};
@@ -1181,3 +1183,68 @@ async fn test_rpc_get_state_root() -> Result<()> {
     node.shutdown().await;
     Ok(())
 }
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_rpc_get_state_proof() -> Result<()> {
+    let temp_dir = tempdir()?;
+
+    // Create keys
+    let owner_secret_key = AccountOwnerSecretKey::generate();
+    let owner_public_key = owner_secret_key.to_pk();
+    let owner_eth_address: EthAddress = owner_public_key.into();
+
+    // Init application service
+    let mut genesis = Genesis::default();
+    genesis.account.push(GenesisAccount {
+        public_key: owner_public_key.into(),
+        flk_balance: 1000u64.into(),
+        stables_balance: 0,
+        bandwidth_balance: 0,
+    });
+
+    let genesis_path = genesis
+        .write_to_dir(temp_dir.path().to_path_buf().try_into().unwrap())
+        .unwrap();
+
+    let port = 30025;
+    let node = init_rpc(&temp_dir, genesis_path, port).await;
+
+    wait_for_server_start(port).await?;
+
+    let client = RpcClient::new_no_auth(&format!("http://127.0.0.1:{port}/rpc/v0"))?;
+    let state_key = StateProofKey::Accounts(owner_eth_address);
+    let (value, proof) =
+        FleekApiClient::get_state_proof(&client, StateProofKey::Accounts(owner_eth_address), None)
+            .await?;
+
+    assert!(value.is_some());
+    let value = value.unwrap();
+    assert_eq!(
+        value.clone(),
+        StateProofValue::Accounts(AccountInfo {
+            flk_balance: 1000u64.into(),
+            stables_balance: HpUfixed::zero(),
+            bandwidth_balance: 0,
+            nonce: 0,
+        })
+    );
+
+    // Verify proof.
+    let root_hash = FleekApiClient::get_state_root(&client, None).await?;
+    proof
+        .verify_membership::<_, _, ApplicationStateTree>(
+            state_key.table(),
+            owner_eth_address,
+            AccountInfo {
+                flk_balance: 1000u64.into(),
+                stables_balance: HpUfixed::zero(),
+                bandwidth_balance: 0,
+                nonce: 0,
+            },
+            root_hash,
+        )
+        .unwrap();
+
+    node.shutdown().await;
+    Ok(())
+}

core/types/Cargo.toml

@@ -7,6 +7,7 @@ edition = "2021"
 
 [dependencies]
 anyhow.workspace = true
+atomo.workspace = true
 cid.workspace = true
 serde.workspace = true
 ink-quill.workspace = true

core/types/src/lib.rs

@@ -18,6 +18,7 @@ mod reputation;
 mod response;
 mod rpc;
 mod state;
+mod state_proof;
 mod transaction;
 
 pub use application::*;
@@ -38,6 +39,7 @@ pub use reputation::*;
 pub use response::*;
 pub use rpc::*;
 pub use state::*;
+pub use state_proof::*;
 pub use transaction::*;
 
 /// The physical address of a node where it can be reached, the port numbers are

core/types/src/state.rs

@@ -37,13 +37,13 @@ impl Tokens {
     }
 }
 
-#[derive(Debug, PartialEq, Serialize, Deserialize, Clone, Default, schemars::JsonSchema)]
+#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, Clone, Default, schemars::JsonSchema)]
 pub struct NodeServed {
     pub served: CommodityServed,
     pub stables_revenue: HpUfixed<6>,
 }
 
-#[derive(Debug, PartialEq, Serialize, Deserialize, Clone, Default, schemars::JsonSchema)]
+#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, Clone, Default, schemars::JsonSchema)]
 pub struct TotalServed {
     pub served: CommodityServed,
     pub reward_pool: HpUfixed<6>,
@@ -75,14 +75,14 @@ pub enum CommodityTypes {
     Gpu = 2,
 }
 
-#[derive(Clone, Debug, Hash, Serialize, Deserialize, schemars::JsonSchema)]
+#[derive(Clone, Debug, Hash, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
 pub struct ReportedReputationMeasurements {
     pub reporting_node: NodeIndex,
     pub measurements: ReputationMeasurements,
 }
 
 /// Metadata, state stored in the blockchain that applies to the current block
-#[derive(Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, schemars::JsonSchema)]
 pub enum Metadata {
     ChainId,
     Epoch,
@@ -100,7 +100,7 @@ pub enum Metadata {
 }
 
 /// The Value enum is a data type used to represent values in a key-value pair for a metadata table
-#[derive(Serialize, Deserialize, Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
 pub enum Value {
     ChainId(u32),
     Epoch(u64),
@@ -378,7 +378,19 @@ pub struct Service {
     pub slashing: (),
 }
 
-#[derive(Debug, Hash, PartialEq, PartialOrd, Ord, Eq, Serialize, Deserialize, Clone, Default)]
+#[derive(
+    Debug,
+    Hash,
+    PartialEq,
+    PartialOrd,
+    Ord,
+    Eq,
+    Serialize,
+    Deserialize,
+    Clone,
+    Default,
+    schemars::JsonSchema,
+)]
 pub struct Committee {
     pub members: Vec<NodeIndex>,
     pub ready_to_change: Vec<NodeIndex>,