darkirc: Persist RLN message counters

bin/darkirc/src/crypto/rln.rs

@@ -37,25 +37,16 @@ use tracing::info;
 pub const RLN_TRAPDOOR_DERIVATION_PATH: pallas::Base = pallas::Base::from_raw([4211, 0, 0, 0]);
 pub const RLN_NULLIFIER_DERIVATION_PATH: pallas::Base = pallas::Base::from_raw([4212, 0, 0, 0]);
 
-/// A user-side RLN identity: long-lived secrets plus an in-memory
-/// per-epoch send counter.
+/// A user-side RLN identity: long-lived secrets plus a per-epoch
+/// send counter.
 ///
-/// The struct is `Copy` so it can be cheaply duplicated when handed
-/// off to client tasks; the per-epoch counter (`message_id`,
-/// `last_epoch`) is therefore expected to be tracked by whichever
-/// task owns the canonical mutable copy. In the typical DarkIRC
-/// configuration that's the `IrcServer::rln_identity` field
-/// (`RwLock<Option<RlnIdentity>>`).
-///
-/// `message_id` and `last_epoch` are not persisted to disk on
-/// shutdown. An RLN epoch is `RLN_EPOCH_LEN` (10 minutes); the
-/// worst case of a node restart is that the counter resets to 0
-/// for the current epoch, which only risks a slash if the user
-/// actually sent distinct messages with the same message_id in the
-/// same epoch - which generally requires hot-restarting at sub-
-/// second cadence. If/when that becomes an operational concern,
-/// persist the counter through a `next_message_id_persisted`-style
-/// helper at the call site.
+/// The struct is `Copy` so it can be cheaply duplicated after a
+/// message slot has been reserved. The canonical mutable copy lives
+/// in `IrcServer::rln_identity` and outbound sends must reserve a
+/// slot through `IrcServer::reserve_rln_message_id()`, which persists
+/// `message_id` and `last_epoch` before proof creation. Persisting
+/// first means a crash can burn a slot, but cannot roll the counter
+/// back and self-slash the identity on restart.
 #[derive(Copy, Clone, SerialEncodable, SerialDecodable)]
 pub struct RlnIdentity {
     pub nullifier: pallas::Base,

bin/darkirc/src/irc/client.rs

@@ -41,13 +41,10 @@ use smol::{
 use tracing::{debug, error, warn};
 
 use super::{
-    server::{IrcServer, MAX_MSG_LEN},
+    server::{IrcServer, RlnMessageReservation, MAX_MSG_LEN},
     NickServ, SERVER_NAME,
 };
-use crate::{
-    crypto::rln::{RlnIdentity, RLN2_SIGNAL_ZKBIN},
-    Privmsg,
-};
+use crate::{crypto::rln::RLN2_SIGNAL_ZKBIN, Privmsg};
 
 const PENALTY_LIMIT: usize = 5;
 
@@ -213,20 +210,24 @@ impl Client {
                                 // honestly to syncing peers and that
                                 // peers would strict-reject anyway.
                                 let blob = {
-                                    let mut id_guard = self.server.rln_identity.write().await;
-                                    let Some(rln_identity) = id_guard.as_mut() else {
-                                        warn!(
-                                            "[IRC CLIENT] No RLN identity registered; \
-                                             refusing to send. Use \
-                                             `/msg NickServ REGISTER ...` to register."
-                                        );
-                                        continue
-                                    };
-                                    let mid = match rln_identity
-                                        .next_message_id(event.header.timestamp)
+                                    let (rln_identity, mid) = match self
+                                        .server
+                                        .reserve_rln_message_id(event.header.timestamp)
+                                        .await?
                                     {
-                                        Some(v) => v,
-                                        None => {
+                                        RlnMessageReservation::Reserved {
+                                            identity,
+                                            message_id,
+                                        } => (identity, message_id),
+                                        RlnMessageReservation::MissingIdentity => {
+                                            warn!(
+                                                "[IRC CLIENT] No RLN identity registered; \
+                                                 refusing to send. Use \
+                                                 `/msg NickServ REGISTER ...` to register."
+                                            );
+                                            continue
+                                        }
+                                        RlnMessageReservation::BudgetExhausted => {
                                             warn!(
                                                 "[IRC CLIENT] RLN message budget \
                                                  exhausted for this epoch; dropping \

bin/darkirc/src/irc/server.rs

@@ -24,7 +24,7 @@ use darkfi::{
     util::path::expand_path,
     Error, Result,
 };
-use darkfi_serial::deserialize_async;
+use darkfi_serial::{deserialize_async, serialize_async};
 use futures_rustls::{
     rustls::{self, pki_types::PrivateKeyDer},
     TlsAcceptor,
@@ -42,7 +42,7 @@ use url::Url;
 
 use super::{
     client::Client,
-    services::nickserv::{ACCOUNTS_DB_PREFIX, ACCOUNTS_KEY_RLN_IDENTITY},
+    services::nickserv::{ACCOUNTS_DB_PREFIX, ACCOUNTS_DEFAULT_TREE, ACCOUNTS_KEY_RLN_IDENTITY},
     IrcChannel, IrcContact,
 };
 use crate::{
@@ -58,6 +58,73 @@ pub const MAX_NICK_LEN: usize = 24;
 /// Max message length
 pub const MAX_MSG_LEN: usize = 512;
 
+/// Result of attempting to reserve the next RLN message slot.
+pub enum RlnMessageReservation {
+    /// No active RLN identity is configured.
+    MissingIdentity,
+    /// The active identity has already used its epoch budget.
+    BudgetExhausted,
+    /// A message slot was persisted and can be used to build a proof.
+    Reserved { identity: RlnIdentity, message_id: u64 },
+}
+
+/// Persist the active RLN counter to the default mirror and matching account tree.
+async fn persist_rln_identity_counter(sled_db: &sled::Db, identity: &RlnIdentity) -> Result<()> {
+    let encoded = serialize_async(identity).await;
+    let active_commitment = identity.commitment();
+    let mut updated_account = false;
+
+    for raw in sled_db.tree_names() {
+        let bytes: &[u8] = raw.as_ref();
+        let Ok(name) = std::str::from_utf8(bytes) else { continue };
+        let Some(account_name) = name.strip_prefix(ACCOUNTS_DB_PREFIX) else { continue };
+        if account_name == "default" || account_name.is_empty() {
+            continue
+        }
+
+        let tree = sled_db.open_tree(name)?;
+        let Some(blob) = tree.get(ACCOUNTS_KEY_RLN_IDENTITY)? else { continue };
+        let Ok(stored): std::result::Result<RlnIdentity, _> = deserialize_async(&blob).await else {
+            continue
+        };
+        if stored.commitment() == active_commitment {
+            tree.insert(ACCOUNTS_KEY_RLN_IDENTITY, encoded.clone())?;
+            updated_account = true;
+        }
+    }
+
+    if !updated_account {
+        warn!(
+            target: "darkirc::irc::server",
+            "active RLN identity has no matching account tree; persisting default mirror only",
+        );
+    }
+
+    let default_db = sled_db.open_tree(ACCOUNTS_DEFAULT_TREE)?;
+    default_db.insert(ACCOUNTS_KEY_RLN_IDENTITY, encoded)?;
+    sled_db.flush_async().await?;
+    Ok(())
+}
+
+/// Reserve the next RLN message ID and persist it before proof creation.
+pub(crate) async fn reserve_rln_message_id_in_store(
+    sled_db: &sled::Db,
+    active: &mut Option<RlnIdentity>,
+    now_millis: u64,
+) -> Result<RlnMessageReservation> {
+    let Some(current) = active else { return Ok(RlnMessageReservation::MissingIdentity) };
+
+    let mut updated = *current;
+    let Some(message_id) = updated.next_message_id(now_millis) else {
+        return Ok(RlnMessageReservation::BudgetExhausted)
+    };
+
+    persist_rln_identity_counter(sled_db, &updated).await?;
+    *current = updated;
+
+    Ok(RlnMessageReservation::Reserved { identity: updated, message_id })
+}
+
 /// IRC server instance
 pub struct IrcServer {
     /// DarkIrc instance
@@ -88,6 +155,12 @@ pub struct IrcServer {
 }
 
 impl IrcServer {
+    /// Reserve and persist the next RLN message slot before proof creation.
+    pub async fn reserve_rln_message_id(&self, now_millis: u64) -> Result<RlnMessageReservation> {
+        let mut active = self.rln_identity.write().await;
+        reserve_rln_message_id_in_store(&self.darkirc.sled, &mut active, now_millis).await
+    }
+
     /// Instantiate a new IRC server. This function will try to bind a TCP socket,
     /// and optionally load a TLS certificate and key. To start the listening loop,
     /// call `IrcServer::listen()`.
@@ -424,3 +497,71 @@ impl IrcServer {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use darkfi::event_graph::rln::epoch_of;
+    use darkfi_sdk::pasta::pallas;
+    use darkfi_serial::deserialize_async;
+
+    use super::*;
+
+    fn test_identity(limit: u64) -> RlnIdentity {
+        RlnIdentity {
+            nullifier: pallas::Base::from(0xabc_u64),
+            trapdoor: pallas::Base::from(0xdef_u64),
+            user_message_limit: limit,
+            message_id: 0,
+            last_epoch: 0,
+        }
+    }
+
+    #[test]
+    fn rln_message_reservation_persists_default_and_account_counters() {
+        smol::block_on(async {
+            let sled_db = sled::Config::new().temporary(true).open().unwrap();
+            let account = sled_db.open_tree(format!("{ACCOUNTS_DB_PREFIX}alice")).unwrap();
+            let default = sled_db.open_tree(ACCOUNTS_DEFAULT_TREE).unwrap();
+            let identity = test_identity(2);
+            let encoded = serialize_async(&identity).await;
+            account.insert(ACCOUNTS_KEY_RLN_IDENTITY, encoded.clone()).unwrap();
+            default.insert(ACCOUNTS_KEY_RLN_IDENTITY, encoded).unwrap();
+
+            let now = 1_704_067_800_000;
+            let mut active = Some(identity);
+            let reservation =
+                reserve_rln_message_id_in_store(&sled_db, &mut active, now).await.unwrap();
+            let RlnMessageReservation::Reserved { identity: reserved, message_id } = reservation
+            else {
+                panic!("expected reservation")
+            };
+            assert_eq!(message_id, 0);
+            assert_eq!(reserved.message_id, 1);
+            assert_eq!(reserved.last_epoch, epoch_of(now));
+
+            let stored_default: RlnIdentity =
+                deserialize_async(&default.get(ACCOUNTS_KEY_RLN_IDENTITY).unwrap().unwrap())
+                    .await
+                    .unwrap();
+            let stored_account: RlnIdentity =
+                deserialize_async(&account.get(ACCOUNTS_KEY_RLN_IDENTITY).unwrap().unwrap())
+                    .await
+                    .unwrap();
+            assert_eq!(stored_default.message_id, 1);
+            assert_eq!(stored_account.message_id, 1);
+            assert_eq!(stored_default.last_epoch, epoch_of(now));
+            assert_eq!(stored_account.last_epoch, epoch_of(now));
+
+            let reservation =
+                reserve_rln_message_id_in_store(&sled_db, &mut active, now).await.unwrap();
+            let RlnMessageReservation::Reserved { message_id, .. } = reservation else {
+                panic!("expected second reservation")
+            };
+            assert_eq!(message_id, 1);
+
+            let exhausted =
+                reserve_rln_message_id_in_store(&sled_db, &mut active, now).await.unwrap();
+            assert!(matches!(exhausted, RlnMessageReservation::BudgetExhausted));
+        })
+    }
+}

bin/darkirc/src/irc/services/nickserv.rs

@@ -473,7 +473,7 @@ impl NickServ {
         };
 
         // Create a new RLN identity. `last_epoch` is initialised to
-        // 0 deterministically - the first call to `next_message_id`
+        // 0 deterministically - the first persisted send reservation
         // will detect the rollover to the current wall-clock epoch.
         let new_rln_identity = RlnIdentity {
             nullifier: identity_nullifier,
@@ -670,12 +670,9 @@ impl NickServ {
         let db_default = self.server.darkirc.sled.open_tree(ACCOUNTS_DEFAULT_TREE)?;
         db_default.insert(ACCOUNTS_KEY_RLN_IDENTITY, blob.as_ref())?;
 
-        // Swap in-memory. The new identity comes in with
-        // message_id=0 / last_epoch=0; the first send reconciles
-        // last_epoch to the current wall-clock epoch and proceeds.
-        // The behaviour matches what happens at startup when the
-        // default tree is loaded by `IrcServer::new`, so this
-        // doesn't introduce a new failure mode.
+        // Swap in-memory. The loaded identity includes any persisted
+        // counter state from its account tree; future sends reserve and
+        // flush the next slot before proof creation.
         *self.server.rln_identity.write().await = Some(identity);
 
         Ok(notices(