Przeglądaj źródła

event_graph: Add sled storage backend for the SMT

x 1 miesiąc temu
rodzic
commit
2695b2315d
3 zmienionych plików z 223 dodań i 35 usunięć
  1. 39 7
      src/event_graph/mod.rs
  2. 143 28
      src/event_graph/rln.rs
  3. 41 0
      src/event_graph/tests_rln.rs

+ 39 - 7
src/event_graph/mod.rs

@@ -800,6 +800,7 @@ impl EventGraph {
     ) -> Result<EventGraphPtr> {
     ) -> Result<EventGraphPtr> {
         config.validate()?;
         config.validate()?;
         let identity_state = IdentityState::new(&sled_db)?;
         let identity_state = IdentityState::new(&sled_db)?;
+        log_memory("after RLN identity state initialization");
         let rln_app_id = rln::RlnAppId::from_genesis(&config.genesis_contents);
         let rln_app_id = rln::RlnAppId::from_genesis(&config.genesis_contents);
         let current_genesis = generate_genesis(&config)?;
         let current_genesis = generate_genesis(&config)?;
         let (pregenerated_identity_commitments, pregenerated_identity_commitment_reprs) =
         let (pregenerated_identity_commitments, pregenerated_identity_commitment_reprs) =
@@ -955,9 +956,9 @@ impl EventGraph {
 
 
         let expected_leaves = expected_commitments.len();
         let expected_leaves = expected_commitments.len();
         let expected_slashed_count = expected_slashed.len();
         let expected_slashed_count = expected_slashed.len();
-        let (actual_commitments, actual_slashed) = {
+        let (actual_commitments, actual_slashed, actual_root) = {
             let state = self.identity_state.read().await;
             let state = self.identity_state.read().await;
-            (state.commitment_reprs(), state.slashed_commitment_reprs())
+            (state.commitment_reprs(), state.slashed_commitment_reprs(), state.root())
         };
         };
         let (actual_leaves, leaves_consistent) = match actual_commitments {
         let (actual_leaves, leaves_consistent) = match actual_commitments {
             Ok(commitments) => {
             Ok(commitments) => {
@@ -989,19 +990,31 @@ impl EventGraph {
 
 
         let static_count = events.len();
         let static_count = events.len();
         let historical_roots_consistent = self.historical_roots_index_consistent(static_count)?;
         let historical_roots_consistent = self.historical_roots_index_consistent(static_count)?;
+        let expected_current_root = if historical_roots_consistent {
+            Some(self.historical_roots_current_root(static_count)?)
+        } else {
+            None
+        };
+        let smt_nodes_consistent = match expected_current_root {
+            Some(root) => actual_root == root,
+            None => false,
+        };
         let recorded_count = self.rln_historical_roots_ordered.len();
         let recorded_count = self.rln_historical_roots_ordered.len();
         let by_value_count = self.rln_historical_roots_by_value.len();
         let by_value_count = self.rln_historical_roots_by_value.len();
-        let consistent = historical_roots_consistent && leaves_consistent && slashed_consistent;
+        let consistent = historical_roots_consistent &&
+            leaves_consistent &&
+            slashed_consistent &&
+            smt_nodes_consistent;
 
 
         info!(
         info!(
             target: "event_graph::new",
             target: "event_graph::new",
             concat!(
             concat!(
                 "[EVENTGRAPH] RLN state audit: static_count={} recorded_count={} ",
                 "[EVENTGRAPH] RLN state audit: static_count={} recorded_count={} ",
                 "by_value_count={} actual_leaves={} expected_leaves={} actual_slashed={} ",
                 "by_value_count={} actual_leaves={} expected_leaves={} actual_slashed={} ",
-                "expected_slashed={} consistent={}",
+                "expected_slashed={} smt_nodes_consistent={} consistent={}",
             ),
             ),
             static_count, recorded_count, by_value_count, actual_leaves, expected_leaves,
             static_count, recorded_count, by_value_count, actual_leaves, expected_leaves,
-            actual_slashed_count, expected_slashed_count, consistent,
+            actual_slashed_count, expected_slashed_count, smt_nodes_consistent, consistent,
         );
         );
 
 
         if consistent {
         if consistent {
@@ -1012,11 +1025,13 @@ impl EventGraph {
             target: "event_graph::new",
             target: "event_graph::new",
             concat!(
             concat!(
                 "[EVENTGRAPH] Rebuilding RLN state: {} static events, {} recorded roots, ",
                 "[EVENTGRAPH] Rebuilding RLN state: {} static events, {} recorded roots, ",
-                "{} by-value roots, {} leaves (expected {}), {} slashed (expected {})",
+                "{} by-value roots, {} leaves (expected {}), {} slashed (expected {}), ",
+                "smt_nodes_consistent={}",
             ),
             ),
             static_count, recorded_count, by_value_count, actual_leaves, expected_leaves,
             static_count, recorded_count, by_value_count, actual_leaves, expected_leaves,
-            actual_slashed_count, expected_slashed_count,
+            actual_slashed_count, expected_slashed_count, smt_nodes_consistent,
         );
         );
+        log_memory("before RLN state rebuild");
 
 
         self.rln_historical_roots_ordered.clear()?;
         self.rln_historical_roots_ordered.clear()?;
         self.rln_historical_roots_by_value.clear()?;
         self.rln_historical_roots_by_value.clear()?;
@@ -1025,11 +1040,13 @@ impl EventGraph {
             let mut state = self.identity_state.write().await;
             let mut state = self.identity_state.write().await;
             state.clear_for_rebuild()?;
             state.clear_for_rebuild()?;
         }
         }
+        log_memory("after RLN state clear for rebuild");
 
 
         for (ev, rln_node) in events {
         for (ev, rln_node) in events {
             let _ = self.apply_rln_static_event(&ev, &rln_node).await?;
             let _ = self.apply_rln_static_event(&ev, &rln_node).await?;
         }
         }
 
 
+        log_memory("after RLN state rebuild");
         info!(
         info!(
             target: "event_graph::new",
             target: "event_graph::new",
             "[EVENTGRAPH] RLN state rebuild complete",
             "[EVENTGRAPH] RLN state rebuild complete",
@@ -1066,6 +1083,21 @@ impl EventGraph {
         Ok(true)
         Ok(true)
     }
     }
 
 
+    fn historical_roots_current_root(&self, expected_count: usize) -> Result<pallas::Base> {
+        if expected_count == 0 {
+            return Ok(IdentityState::empty_root())
+        }
+
+        let mut current_root = None;
+        for item in self.rln_historical_roots_ordered.iter() {
+            let (_, value_bytes) = item?;
+            let (root, _) = decode_historical_root_value(&value_bytes)?;
+            current_root = Some(root);
+        }
+
+        current_root.ok_or_else(|| Error::Custom("missing RLN historical current root".into()))
+    }
+
     /// After header sync, event content can be fetched lazily via local
     /// After header sync, event content can be fetched lazily via local
     /// [`fetch_page`] or peer [`RangeReq`] responses with aligned blobs - the
     /// [`fetch_page`] or peer [`RangeReq`] responses with aligned blobs - the
     /// application pulls the events it actually wants to display or process,
     /// application pulls the events it actually wants to display or process,

+ 143 - 28
src/event_graph/rln.rs

@@ -34,15 +34,17 @@ use darkfi_sdk::{
     crypto::{
     crypto::{
         pasta_prelude::{FromUniformBytes, PrimeField},
         pasta_prelude::{FromUniformBytes, PrimeField},
         poseidon_hash,
         poseidon_hash,
-        smt::{MemoryStorageFp, PoseidonFp, SmtMemoryFp, EMPTY_NODES_FP, SMT_FP_DEPTH},
+        smt::{PoseidonFp, SparseMerkleTree, StorageAdapter, EMPTY_NODES_FP, SMT_FP_DEPTH},
     },
     },
+    error::{ContractError, ContractResult},
     pasta::pallas,
     pasta::pallas,
 };
 };
 use darkfi_serial::{async_trait, FutAsyncWriteExt, SerialDecodable, SerialEncodable};
 use darkfi_serial::{async_trait, FutAsyncWriteExt, SerialDecodable, SerialEncodable};
 use halo2_proofs::{arithmetic::Field, circuit::Value};
 use halo2_proofs::{arithmetic::Field, circuit::Value};
+use num_bigint::BigUint;
 use rand::rngs::OsRng;
 use rand::rngs::OsRng;
 use sled_overlay::sled;
 use sled_overlay::sled;
-use tracing::info;
+use tracing::{error, info};
 
 
 use super::Event;
 use super::Event;
 use crate::{
 use crate::{
@@ -264,6 +266,88 @@ const SIGNAL_VK_KEY: &str = "rlnv2-diff-signal-vk";
 const SIGNAL_PK_KEY: &str = "rlnv2-diff-signal-pk";
 const SIGNAL_PK_KEY: &str = "rlnv2-diff-signal-pk";
 const SLASH_VK_KEY: &str = "rlnv2-diff-slash-vk";
 const SLASH_VK_KEY: &str = "rlnv2-diff-slash-vk";
 const SLASH_PK_KEY: &str = "rlnv2-diff-slash-pk";
 const SLASH_PK_KEY: &str = "rlnv2-diff-slash-pk";
+const IDENTITY_LEAVES_TREE: &str = "rln-identity-leaves";
+const SLASHED_IDENTITY_LEAVES_TREE: &str = "rln-slashed-identity-leaves";
+const IDENTITY_SMT_NODES_TREE: &str = "rln-identity-smt-nodes";
+
+type SmtSledFp = SparseMerkleTree<
+    'static,
+    SMT_FP_DEPTH,
+    { SMT_FP_DEPTH + 1 },
+    pallas::Base,
+    PoseidonFp,
+    SledStorageFp,
+>;
+
+/// Sled-backed storage for RLN identity SMT nodes.
+///
+/// The leaf/tombstone trees remain the durable source of truth. This
+/// tree stores the SMT's internal node cache so membership proofs and
+/// root lookups do not require keeping the full sparse tree in heap.
+#[derive(Clone)]
+struct SledStorageFp {
+    tree: sled::Tree,
+}
+
+impl SledStorageFp {
+    fn new(tree: sled::Tree) -> Self {
+        Self { tree }
+    }
+}
+
+impl StorageAdapter for SledStorageFp {
+    type Value = pallas::Base;
+
+    fn put(&mut self, key: BigUint, value: pallas::Base) -> ContractResult {
+        if let Err(e) = self.tree.insert(key.to_bytes_le(), value.to_repr().as_ref()) {
+            error!(
+                target: "event_graph::rln",
+                "[RLN] SMT sled put failed: {e}",
+            );
+            return Err(ContractError::SmtPutFailed)
+        }
+
+        Ok(())
+    }
+
+    fn get(&self, key: &BigUint) -> Option<pallas::Base> {
+        let value = match self.tree.get(key.to_bytes_le()) {
+            Ok(value) => value?,
+            Err(e) => {
+                error!(
+                    target: "event_graph::rln",
+                    "[RLN] SMT sled get failed: {e}",
+                );
+                return None
+            }
+        };
+
+        if value.len() != 32 {
+            error!(
+                target: "event_graph::rln",
+                "[RLN] SMT sled node must be 32 bytes, got {}",
+                value.len(),
+            );
+            return None
+        }
+
+        let mut repr = [0u8; 32];
+        repr.copy_from_slice(&value);
+        pallas::Base::from_repr(repr).into()
+    }
+
+    fn del(&mut self, key: &BigUint) -> ContractResult {
+        if let Err(e) = self.tree.remove(key.to_bytes_le()) {
+            error!(
+                target: "event_graph::rln",
+                "[RLN] SMT sled del failed: {e}",
+            );
+            return Err(ContractError::SmtDelFailed)
+        }
+
+        Ok(())
+    }
+}
 
 
 /// ZK key cache.
 /// ZK key cache.
 pub struct ZkKeys {
 pub struct ZkKeys {
@@ -477,53 +561,87 @@ pub enum StaticEventCheck {
 /// The set of currently registered RLN identities, stored as a Sparse
 /// The set of currently registered RLN identities, stored as a Sparse
 /// Merkle Tree (SMT).
 /// Merkle Tree (SMT).
 ///
 ///
-/// Persistence model: leaf commitments are stored in a dedicated sled
-/// tree (`rln-identity-leaves`). The in-memory SMT is rebuilt from
-/// these leaves on startup.
+/// Persistence model: leaf commitments and slash tombstones are stored in
+/// dedicated sled trees. SMT internal nodes are also sled-backed, but are
+/// treated as rebuildable derived state from the canonical static DAG and the
+/// leaf/tombstone side tables.
 pub struct IdentityState {
 pub struct IdentityState {
-    smt: SmtMemoryFp,
+    smt: SmtSledFp,
     leaves: sled::Tree,
     leaves: sled::Tree,
     slashed: sled::Tree,
     slashed: sled::Tree,
+    smt_nodes: sled::Tree,
     recent_roots: VecDeque<pallas::Base>,
     recent_roots: VecDeque<pallas::Base>,
 }
 }
 
 
 impl IdentityState {
 impl IdentityState {
     pub fn new(sled_db: &sled::Db) -> Result<Self> {
     pub fn new(sled_db: &sled::Db) -> Result<Self> {
+        let leaves = sled_db.open_tree(IDENTITY_LEAVES_TREE)?;
+        let slashed = sled_db.open_tree(SLASHED_IDENTITY_LEAVES_TREE)?;
+        let smt_nodes = sled_db.open_tree(IDENTITY_SMT_NODES_TREE)?;
+        let smt = Self::new_smt(smt_nodes.clone());
+
+        let mut state = Self {
+            smt,
+            leaves,
+            slashed,
+            smt_nodes,
+            recent_roots: VecDeque::with_capacity(ROOT_HISTORY_SIZE),
+        };
+
+        if state.smt_nodes.is_empty() && !state.leaves.is_empty() {
+            state.restore_smt_nodes_from_leaves()?;
+        }
+
+        state.recent_roots.push_back(state.smt.root());
+        Ok(state)
+    }
+
+    fn new_smt(smt_nodes: sled::Tree) -> SmtSledFp {
         let hasher = PoseidonFp::new();
         let hasher = PoseidonFp::new();
-        let store = MemoryStorageFp::new();
-        let mut smt = SmtMemoryFp::new(store, hasher, &EMPTY_NODES_FP);
+        let store = SledStorageFp::new(smt_nodes);
+        SmtSledFp::new(store, hasher, &EMPTY_NODES_FP)
+    }
+
+    /// Return the root of an empty RLN identity SMT.
+    pub fn empty_root() -> pallas::Base {
+        EMPTY_NODES_FP[0]
+    }
 
 
-        let leaves = sled_db.open_tree("rln-identity-leaves")?;
-        let slashed = sled_db.open_tree("rln-slashed-identity-leaves")?;
+    /// Recreate the sled-backed SMT node tree from persisted leaves.
+    ///
+    /// This is a migration and repair path for databases that have the old
+    /// leaf side table but no `rln-identity-smt-nodes` tree yet. It restores
+    /// one leaf at a time to avoid allocating a full leaf batch in memory.
+    fn restore_smt_nodes_from_leaves(&mut self) -> Result<()> {
+        let mut restored = 0usize;
 
 
-        let mut batch = vec![];
-        for item in leaves.iter() {
+        for item in self.leaves.iter() {
             let (key, val) = item?;
             let (key, val) = item?;
             if key.len() != 32 || val.len() != 32 || key.as_ref() != val.as_ref() {
             if key.len() != 32 || val.len() != 32 || key.as_ref() != val.as_ref() {
                 continue
                 continue
             }
             }
+
             let mut repr = [0u8; 32];
             let mut repr = [0u8; 32];
             repr.copy_from_slice(&val);
             repr.copy_from_slice(&val);
-            if slashed.contains_key(repr)? {
+            if self.slashed.contains_key(repr)? {
                 continue
                 continue
             }
             }
-            if let Some(c) = pallas::Base::from_repr(repr).into() {
-                batch.push((c, c));
+
+            if let Some(commitment) = pallas::Base::from_repr(repr).into() {
+                self.smt.insert_batch(vec![(commitment, commitment)])?;
+                restored += 1;
             }
             }
         }
         }
 
 
-        if !batch.is_empty() {
+        if restored > 0 {
             info!(
             info!(
                 target: "event_graph::rln",
                 target: "event_graph::rln",
-                "[RLN] Restoring {} identities from sled", batch.len(),
+                "[RLN] Restored {} sled-backed identity SMT nodes from leaves",
+                restored,
             );
             );
-            smt.insert_batch(batch)?;
         }
         }
 
 
-        let mut recent_roots = VecDeque::with_capacity(ROOT_HISTORY_SIZE);
-        recent_roots.push_back(smt.root());
-
-        Ok(Self { smt, leaves, slashed, recent_roots })
+        Ok(())
     }
     }
 
 
     /// Returns true if the commitment is already a leaf in the tree.
     /// Returns true if the commitment is already a leaf in the tree.
@@ -650,18 +768,15 @@ impl IdentityState {
         self.recent_roots.push_back(root);
         self.recent_roots.push_back(root);
     }
     }
 
 
-    /// Reset the in-memory SMT and the persistent leaves tree to an
+    /// Reset the sled-backed SMT and the persistent leaves tree to an
     /// empty state, in preparation for replaying the canonical
     /// empty state, in preparation for replaying the canonical
     /// static-DAG history.
     /// static-DAG history.
     pub fn clear_for_rebuild(&mut self) -> Result<()> {
     pub fn clear_for_rebuild(&mut self) -> Result<()> {
         // Drop every derived identity state from sled.
         // Drop every derived identity state from sled.
         self.leaves.clear()?;
         self.leaves.clear()?;
         self.slashed.clear()?;
         self.slashed.clear()?;
-
-        // Replace the in-memory SMT with a fresh empty one.
-        let hasher = PoseidonFp::new();
-        let store = MemoryStorageFp::new();
-        self.smt = SmtMemoryFp::new(store, hasher, &EMPTY_NODES_FP);
+        self.smt_nodes.clear()?;
+        self.smt = Self::new_smt(self.smt_nodes.clone());
 
 
         // Reset the recent-roots cache. The empty SMT root is the
         // Reset the recent-roots cache. The empty SMT root is the
         // current state.
         // current state.

+ 41 - 0
src/event_graph/tests_rln.rs

@@ -217,6 +217,47 @@ fn rln_identity_state_persists_across_reopen() {
     assert!(s2.register(slashed).is_err());
     assert!(s2.register(slashed).is_err());
 }
 }
 
 
+#[test]
+fn rln_identity_state_restores_sled_smt_nodes_from_leaves() {
+    let db = sled::Config::new().temporary(true).open().unwrap();
+    let c = pallas::Base::from(0xbeefu64);
+
+    let original_root = {
+        let mut s = IdentityState::new(&db).unwrap();
+        s.register(c).unwrap();
+        s.root()
+    };
+
+    db.open_tree("rln-identity-smt-nodes").unwrap().clear().unwrap();
+
+    let s2 = IdentityState::new(&db).unwrap();
+    assert!(s2.contains(&c));
+    assert_eq!(s2.root(), original_root);
+    assert!(s2.prove_membership(&c).verify(&original_root, &c, &c));
+}
+
+#[test]
+fn rln_rebuild_detects_stale_sled_smt_nodes() {
+    smol::block_on(async {
+        let config = EventGraphConfig { hours_rotation: 1, ..test_config() };
+        let eg = make_eg_with_config(config).await;
+        let commitment = genesis_commitment_at(&eg, 0);
+        let original_root = eg.identity_state.read().await.root();
+
+        eg.sled_db.open_tree("rln-identity-smt-nodes").unwrap().clear().unwrap();
+
+        assert!(eg.rln_contains(&commitment).await);
+        assert_ne!(eg.identity_state.read().await.root(), original_root);
+        assert_eq!(eg.rln_historical_roots_ordered.len(), 1);
+
+        eg.rebuild_historical_roots_if_needed().await.unwrap();
+
+        assert!(eg.rln_contains(&commitment).await);
+        assert_eq!(eg.identity_state.read().await.root(), original_root);
+        assert_eq!(eg.rln_historical_roots_ordered.len(), 1);
+    })
+}
+
 #[test]
 #[test]
 fn rln_cross_app_isolation_on_internal_nullifier() {
 fn rln_cross_app_isolation_on_internal_nullifier() {
     // Two apps with different RlnAppId, same identity_secret_hash,
     // Two apps with different RlnAppId, same identity_secret_hash,