fix(platform-wallet-storage): persist sentinel-scoped token balances of wallet-owned identities

packages/rs-platform-wallet-storage/src/sqlite/schema/token_balances.rs

@@ -2,11 +2,15 @@
 //!
 //! # Precondition
 //!
-//! Every `identity_id` MUST already exist in `identities` and belong to the
-//! flush's `wallet_id` (or have NULL `wallet_id` for the all-zero sentinel
-//! scope). The FK enforces existence; the wallet match is checked here and
-//! propagates [`WalletStorageError::WalletIdMismatch`] on a mis-attributed
-//! caller.
+//! Every `identity_id` MUST already exist in `identities`. Under a real
+//! wallet scope it must also belong to the flush's `wallet_id`. The FK
+//! enforces existence; the wallet match is checked here and propagates
+//! [`WalletStorageError::WalletIdMismatch`] on a mis-attributed caller.
+//!
+//! The all-zero sentinel scope skips the wallet match: `IdentitySyncManager`
+//! has no wallet linkage and writes every balance under it, including those
+//! of wallet-owned identities. Rows carry no wallet column, so there is
+//! nothing to attribute.
 
 use rusqlite::{params, Transaction};
 
@@ -16,27 +20,32 @@ use platform_wallet::wallet::platform_wallet::WalletId;
 use crate::sqlite::error::WalletStorageError;
 use crate::sqlite::util::safe_cast;
 
-/// Keyed by `(identity_id, token_id)`. `wallet_id` feeds the precondition
-/// check only — no column — since cascade flows
-/// `wallets → identities → token_balances` via the nullable
-/// `identities.wallet_id` FK. No auto-prune: hosts deleting identities
-/// out-of-band must prune `token_balances` themselves.
+/// Keyed by `(identity_id, token_id)`. `wallet_id` feeds the ownership
+/// check only (skipped for the all-zero sentinel scope) — no column —
+/// since cascade flows `wallets → identities → token_balances` via the
+/// nullable `identities.wallet_id` FK. No auto-prune: hosts deleting
+/// identities out-of-band must prune `token_balances` themselves.
 pub fn apply(
     tx: &Transaction<'_>,
     wallet_id: &WalletId,
     cs: &TokenBalanceChangeSet,
 ) -> Result<(), WalletStorageError> {
-    let touched: std::collections::BTreeSet<dpp::prelude::Identifier> = cs
-        .balances
-        .keys()
-        .map(|(identity_id, _)| *identity_id)
-        .chain(
-            cs.removed_balances
-                .iter()
-                .map(|(identity_id, _)| *identity_id),
-        )
-        .collect();
-    super::assert_identities_belong_to_wallet(tx, wallet_id, &touched)?;
+    // The sentinel scope is identity-keyed, not wallet-attributed (see the
+    // module docs), so only a real wallet scope is checked for ownership.
+    let scope_is_sentinel = wallet_id.iter().all(|b| *b == 0);
+    if !scope_is_sentinel {
+        let touched: std::collections::BTreeSet<dpp::prelude::Identifier> = cs
+            .balances
+            .keys()
+            .map(|(identity_id, _)| *identity_id)
+            .chain(
+                cs.removed_balances
+                    .iter()
+                    .map(|(identity_id, _)| *identity_id),
+            )
+            .collect();
+        super::assert_identities_belong_to_wallet(tx, wallet_id, &touched)?;
+    }
     if !cs.balances.is_empty() {
         let now = chrono::Utc::now().timestamp();
         let mut stmt = tx.prepare_cached(

packages/rs-platform-wallet-storage/tests/integration/sqlite_structural_hardening.rs

@@ -8,13 +8,14 @@
 
 use crate::common;
 
-use common::{ensure_wallet_meta, fresh_persister, wid};
+use common::{ensure_identity, ensure_wallet_meta, fresh_persister, wid};
 
 use dashcore::{Address, Network, OutPoint, TxOut, Txid};
+use dpp::prelude::Identifier;
 use key_wallet::Utxo;
 use platform_wallet::changeset::{
     CoreChangeSet, PlatformAddressBalanceEntry, PlatformAddressChangeSet, PlatformWalletChangeSet,
-    PlatformWalletPersistence,
+    PlatformWalletPersistence, TokenBalanceChangeSet,
 };
 use platform_wallet::wallet::platform_wallet::WalletId;
 use platform_wallet_storage::WalletStorageError;
@@ -97,6 +98,112 @@ fn platform_addr_mixed_wallet_rejected() {
     );
 }
 
+/// `IdentitySyncManager` has no wallet linkage, so it stores token
+/// balances under the all-zero sentinel scope (`token_balances` rows are
+/// keyed by `(identity_id, token_id)` alone). A wallet-owned identity's
+/// balance must persist under that scope rather than fail the flush with
+/// `WalletIdMismatch` on every sync pass.
+#[test]
+fn sentinel_scoped_token_balance_of_wallet_owned_identity_persists() {
+    let (persister, _tmp, _path) = fresh_persister();
+    let owner_wallet = wid(0xB3);
+    ensure_wallet_meta(&persister, &owner_wallet);
+    let identity = [0x21u8; 32];
+    ensure_identity(&persister, &identity, Some(&owner_wallet));
+    let token = [0x22u8; 32];
+    let pair = (Identifier::from(identity), Identifier::from(token));
+
+    let mut balances = TokenBalanceChangeSet::default();
+    balances.balances.insert(pair, 7);
+    persister
+        .store(WalletId::default(), balances.into())
+        .expect("sentinel-scoped token balance of a wallet-owned identity must persist");
+
+    let stored: i64 = persister
+        .lock_conn_for_test()
+        .query_row(
+            "SELECT balance FROM token_balances WHERE identity_id = ?1 AND token_id = ?2",
+            params![identity.as_slice(), token.as_slice()],
+            |row| row.get(0),
+        )
+        .expect("token balance row written");
+    assert_eq!(stored, 7);
+
+    let mut removal = TokenBalanceChangeSet::default();
+    removal.removed_balances.insert(pair);
+    persister
+        .store(WalletId::default(), removal.into())
+        .expect("sentinel-scoped token balance removal must persist");
+    let remaining: i64 = persister
+        .lock_conn_for_test()
+        .query_row(
+            "SELECT COUNT(*) FROM token_balances WHERE identity_id = ?1",
+            params![identity.as_slice()],
+            |row| row.get(0),
+        )
+        .expect("count token balance rows");
+    assert_eq!(remaining, 0);
+}
+
+/// The sentinel scope skips only the ownership check: a balance for an
+/// identity with no `identities` row is still rejected by the FK.
+#[test]
+fn sentinel_scoped_token_balance_of_missing_identity_rejected() {
+    let (persister, _tmp, _path) = fresh_persister();
+    let missing_identity = [0x25u8; 32];
+
+    let mut balances = TokenBalanceChangeSet::default();
+    balances.balances.insert(
+        (
+            Identifier::from(missing_identity),
+            Identifier::from([0x26u8; 32]),
+        ),
+        7,
+    );
+    let err = persister
+        .store(WalletId::default(), balances.into())
+        .expect_err("a balance for a missing identity must be rejected");
+    assert!(
+        format!("{err:?}").contains("FOREIGN KEY"),
+        "expected FOREIGN KEY failure, got `{err:?}`"
+    );
+    let rows: i64 = persister
+        .lock_conn_for_test()
+        .query_row(
+            "SELECT COUNT(*) FROM token_balances WHERE identity_id = ?1",
+            params![missing_identity.as_slice()],
+            |row| row.get(0),
+        )
+        .expect("count token balance rows");
+    assert_eq!(rows, 0);
+}
+
+/// A real wallet scope still may not write token balances for another
+/// wallet's identity.
+#[test]
+fn token_balance_of_another_wallets_identity_rejected() {
+    let (persister, _tmp, _path) = fresh_persister();
+    let owner_wallet = wid(0xB4);
+    let other_wallet = wid(0xB5);
+    ensure_wallet_meta(&persister, &owner_wallet);
+    ensure_wallet_meta(&persister, &other_wallet);
+    let identity = [0x23u8; 32];
+    ensure_identity(&persister, &identity, Some(&owner_wallet));
+
+    let mut balances = TokenBalanceChangeSet::default();
+    balances.balances.insert(
+        (Identifier::from(identity), Identifier::from([0x24u8; 32])),
+        7,
+    );
+    let err = persister
+        .store(other_wallet, balances.into())
+        .expect_err("another wallet's scope must be rejected");
+    assert!(
+        format!("{err}").contains("wallet id mismatch"),
+        "expected wallet id mismatch, got `{err}`"
+    );
+}
+
 fn p2pkh(byte: u8) -> Address {
     use dashcore::address::Payload;
     use dashcore::hashes::Hash;