Add chain state indexing and RPC endpoint

Cargo.toml

@@ -36,6 +36,7 @@ blake2 = "0.10"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = { version = "1.0", features = ["arbitrary_precision"] }
 hex = "0.4"
+bincode = "1.3"
 
 # BigInt
 num-bigint = { version = "0.4", features = ["serde"] }

crates/raito-bridge-node/Cargo.toml

@@ -22,6 +22,7 @@ tower-http = { version = "0.5", features = ["trace", "cors", "compression-gzip"]
 bitcoin.workspace = true
 # Storage
 sqlx.workspace = true
+bincode.workspace = true
 # CLI
 clap.workspace = true
 dotenv.workspace = true

crates/raito-bridge-node/src/chain_state.rs

@@ -0,0 +1,111 @@
+use std::str::FromStr;
+use std::sync::Arc;
+
+use accumulators::store::StoreError;
+use async_trait::async_trait;
+use bitcoin::block::BlockHash;
+use bitcoin::block::Header as BlockHeader;
+use bitcoin::Target;
+use bitcoin::Work;
+use raito_spv_verify::ChainState;
+
+const BLOCKS_PER_EPOCH: u32 = 2016;
+
+#[async_trait]
+pub trait ChainStateStore: Send + Sync {
+    async fn add_block_header(
+        &self,
+        height: u32,
+        block_header: &BlockHeader,
+    ) -> Result<(), StoreError>;
+    async fn get_block_headers(
+        &self,
+        start_height: u32,
+        num_blocks: u32,
+    ) -> Result<Vec<BlockHeader>, StoreError>;
+    async fn get_block_height(&self, block_hash: &BlockHash) -> Result<u32, StoreError>;
+    async fn add_chain_state(
+        &self,
+        height: u32,
+        chain_state: &ChainState,
+    ) -> Result<(), StoreError>;
+    async fn get_chain_state(&self, height: u32) -> Result<ChainState, StoreError>;
+}
+
+pub struct ChainStateManager {
+    current_state: ChainState,
+    store: Arc<dyn ChainStateStore>,
+}
+
+impl ChainStateManager {
+    pub async fn restore(
+        store: Arc<dyn ChainStateStore>,
+        height: u32,
+    ) -> Result<Self, anyhow::Error> {
+        let current_state = if height == 0 {
+            Self::genesis_state()
+        } else {
+            store.get_chain_state(height - 1).await?
+        };
+        Ok(Self {
+            current_state,
+            store,
+        })
+    }
+
+    pub async fn update(
+        &mut self,
+        block_height: u32,
+        block_header: &BlockHeader,
+    ) -> Result<(), anyhow::Error> {
+        let new_state = if block_height == 0 {
+            self.current_state.clone()
+        } else {
+            let mut prev_timestamps = self.current_state.prev_timestamps.clone();
+            prev_timestamps.push(block_header.time);
+            if prev_timestamps.len() > 11 {
+                prev_timestamps.remove(0);
+            }
+
+            let epoch_start_time = if block_height % BLOCKS_PER_EPOCH == 0 {
+                block_header.time
+            } else {
+                self.current_state.epoch_start_time
+            };
+
+            ChainState {
+                block_height,
+                total_work: self.current_state.total_work + block_header.work(),
+                best_block_hash: block_header.block_hash(),
+                current_target: block_header.target(),
+                epoch_start_time,
+                prev_timestamps,
+            }
+        };
+
+        self.store.add_chain_state(block_height, &new_state).await?;
+        self.store
+            .add_block_header(block_height, block_header)
+            .await?;
+        self.current_state = new_state;
+
+        Ok(())
+    }
+
+    pub fn genesis_state() -> ChainState {
+        ChainState {
+            block_height: 0,
+            total_work: Work::from_hex("0x100010001").unwrap(),
+            best_block_hash: BlockHash::from_str(
+                "000000000019d6689c085ae165831e934ff763ae46a2a6c172b3f1b60a8ce26f",
+            )
+            .unwrap(),
+            current_target: Target::from_hex(
+                "0xffff0000000000000000000000000000000000000000000000000000",
+            )
+            .unwrap(),
+            epoch_start_time: 1231006505,
+            prev_timestamps: vec![1231006505],
+        }
+    }
+}

crates/raito-bridge-node/src/indexer.rs

@@ -9,7 +9,7 @@ use tracing::{error, info};
 
 use raito_bitcoin_client::BitcoinClient;
 
-use crate::store::AppStore;
+use crate::{chain_state::ChainStateManager, store::AppStore};
 
 /// Bitcoin block indexer that builds MMR accumulator and generates sparse roots
 pub struct Indexer {
@@ -60,13 +60,18 @@ impl Indexer {
         let mut next_block_height = mmr.get_block_count().await?;
         info!("Current MMR blocks count: {}", next_block_height);
 
+        let mut chain_state_mgr =
+            ChainStateManager::restore(store.clone(), next_block_height).await?;
+        info!("Chain state manager initialized");
+
         loop {
             tokio::select! {
                 res = bitcoin_client.wait_block_header(next_block_height, self.config.indexing_lag) => {
                     match res {
                         Ok((block_header, block_hash)) => {
                             store.begin().await?;
-                            mmr.add_block_header(next_block_height, &block_header).await?;
+                            mmr.add_block_header(&block_header).await.map_err(|e| anyhow::anyhow!("Failed to add block header to MMR: {}", e))?;
+                            chain_state_mgr.update(next_block_height, &block_header).await.map_err(|e| anyhow::anyhow!("Failed to update chain state: {}", e))?;
                             store.commit().await?;
                             info!("Block #{} {} processed", next_block_height, block_hash);
                             next_block_height += 1;

crates/raito-bridge-node/src/main.rs

@@ -13,6 +13,7 @@ use crate::{
     shutdown::Shutdown,
 };
 
+mod chain_state;
 mod indexer;
 mod rpc;
 mod shutdown;

crates/raito-bridge-node/src/rpc.rs

@@ -21,9 +21,9 @@ use raito_spv_mmr::{
     block_mmr::{BlockInclusionProof, BlockMMR},
     sparse_roots::SparseRoots,
 };
-use raito_spv_verify::TransactionInclusionProof;
+use raito_spv_verify::{ChainState, TransactionInclusionProof};
 
-use crate::store::AppStore;
+use crate::{chain_state::ChainStateStore, store::AppStore};
 
 /// Query parameters for block inclusion proof generation and roots retrieval
 #[derive(Debug, Deserialize)]
@@ -62,21 +62,25 @@ pub struct RpcServer {
 #[derive(Debug, Clone)]
 pub struct AppState {
     mmr: Arc<BlockMMR>,
+    store: Arc<AppStore>,
     bitcoin_client: Arc<BitcoinClient>,
 }
 
 impl AppState {
-    pub async fn new(config: RpcConfig) -> Result<Self, anyhow::Error> {
+    pub fn new(config: RpcConfig) -> Result<Self, anyhow::Error> {
         let mmr_id = Some(config.mmr_id.clone());
-        let store =
-            AppStore::multiple_concurrent_readers(&config.mmr_db_path, mmr_id.clone()).await?;
+        let store = Arc::new(AppStore::multiple_concurrent_readers(
+            &config.mmr_db_path,
+            mmr_id.clone(),
+        ));
         let hasher = StarkBlakeHasher::default();
-        let mmr = BlockMMR::new(Arc::new(store), Arc::new(hasher), mmr_id);
+        let mmr = BlockMMR::new(store.clone(), Arc::new(hasher), mmr_id);
         let bitcoin_client =
             BitcoinClient::new(config.rpc_url.clone(), config.rpc_userpwd.clone())?;
         Ok(Self {
             mmr: Arc::new(mmr),
             bitcoin_client: Arc::new(bitcoin_client),
+            store: store.clone(),
         })
     }
 }
@@ -93,7 +97,6 @@ impl RpcServer {
         info!("Starting RPC server on {}", self.config.rpc_host);
 
         let app_state = AppState::new(self.config.clone())
-            .await
             .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
 
         let app = Router::new()
@@ -103,6 +106,7 @@ impl RpcServer {
             .route("/headers", get(get_block_headers))
             .route("/transaction-proof/:tx_id", get(get_transaction_proof))
             .route("/block-header/:block_height", get(get_block_header))
+            .route("/chain-state/:block_height", get(get_chain_state))
             .with_state(app_state)
             .layer(CompressionLayer::new())
             .layer(CorsLayer::permissive())
@@ -210,7 +214,7 @@ pub async fn get_block_header(
     Path(block_height): Path<u32>,
 ) -> Result<Json<BlockHeader>, StatusCode> {
     let block_header = state
-        .mmr
+        .store
         .get_block_headers(block_height, 1)
         .await
         .map_err(|e| {
@@ -220,9 +224,8 @@ pub async fn get_block_header(
             );
             StatusCode::INTERNAL_SERVER_ERROR
         })?
-        .get(0)
-        .ok_or(StatusCode::INTERNAL_SERVER_ERROR)?
-        .clone();
+        .pop()
+        .ok_or(StatusCode::INTERNAL_SERVER_ERROR)?;
 
     Ok(Json(block_header))
 }
@@ -242,7 +245,7 @@ pub async fn get_block_headers(
     let offset = query.offset.unwrap_or(0);
     let size = query.size.unwrap_or(10);
     let block_headers = state
-        .mmr
+        .store
         .get_block_headers(offset, size)
         .await
         .map_err(|e| {
@@ -282,13 +285,17 @@ pub async fn get_transaction_proof(
         })?;
 
     let block_hash = block_header.block_hash();
-    let block_height = state.mmr.get_block_height(&block_hash).await.map_err(|e| {
-        error!(
-            "Failed to get block height for block hash {}: {}",
-            block_hash, e
-        );
-        StatusCode::INTERNAL_SERVER_ERROR
-    })?;
+    let block_height = state
+        .store
+        .get_block_height(&block_hash)
+        .await
+        .map_err(|e| {
+            error!(
+                "Failed to get block height for block hash {}: {}",
+                block_hash, e
+            );
+            StatusCode::INTERNAL_SERVER_ERROR
+        })?;
 
     let transaction = state
         .bitcoin_client
@@ -307,3 +314,26 @@ pub async fn get_transaction_proof(
     };
     Ok(Json(transaction_proof.into()))
 }
+
+/// Get the chain state for a specific block height
+///
+/// # Returns
+/// * `Json<ChainState>` - The chain state in JSON format
+/// * `StatusCode::INTERNAL_SERVER_ERROR` - If fetching the chain state fails
+pub async fn get_chain_state(
+    State(state): State<AppState>,
+    Path(block_height): Path<u32>,
+) -> Result<Json<ChainState>, StatusCode> {
+    let chain_state = state
+        .store
+        .get_chain_state(block_height)
+        .await
+        .map_err(|e| {
+            error!(
+                "Failed to get chain state for height {}: {}",
+                block_height, e
+            );
+            StatusCode::INTERNAL_SERVER_ERROR
+        })?;
+    Ok(Json(chain_state))
+}