fix(wallet-storage): restore Platform and Core from one wallet snapshot

packages/rs-platform-wallet-storage/src/sqlite/persister.rs

@@ -9,8 +9,8 @@ use rusqlite::{Connection, OptionalExtension, Transaction};
 use dpp::prelude::Identifier;
 use platform_wallet::changeset::changeset::SweepBatch;
 use platform_wallet::changeset::{
-    ClientStartState, IdentityChangeSet, PersistenceCapabilities, PersistenceError,
-    PlatformWalletChangeSet, PlatformWalletPersistence,
+    ClientStartState, ClientWalletStartState, IdentityChangeSet, PersistenceCapabilities,
+    PersistenceError, PlatformWalletChangeSet, PlatformWalletPersistence,
 };
 use platform_wallet::wallet::identity::ManagedIdentity;
 use platform_wallet::wallet::platform_wallet::WalletId;
@@ -1533,46 +1533,26 @@ impl PlatformWalletPersistence for SqlitePersister {
                 .map_err(PersistenceError::from)?;
         }
 
-        let addrs_all = schema::platform_addrs::load_all(&conn).map_err(PersistenceError::from)?;
+        schema::platform_addrs::validate_wallet_ids(&conn).map_err(PersistenceError::from)?;
         let mut addresses_loaded: usize = 0;
-        // Wallets whose platform-address rows could not be read. They are
-        // already counted here, so the loop below skips them rather than
-        // rebuilding a wallet whose platform balance would be understated.
-        let mut unreadable: std::collections::BTreeSet<WalletId> =
-            std::collections::BTreeSet::new();
-        for (wallet_id, entry) in addrs_all {
-            let (addrs, count) = match entry {
-                Ok(entry) => entry,
-                Err(original) => {
-                    degrade_whole_wallet(&ctx, wallet_id, &original)
-                        .map_err(|_| PersistenceError::from(original))?;
-                    unreadable.insert(wallet_id);
-                    continue;
-                }
-            };
-            // Skip a wallet with no platform state at all (no addresses,
-            // no registrations, all sync watermarks zero).
-            if count > 0
-                || !addrs.per_account.is_empty()
-                || addrs.sync_height > 0
-                || addrs.sync_timestamp > 0
-                || addrs.last_known_recent_block > 0
-            {
-                addresses_loaded += count;
-                state.platform_addresses.insert(wallet_id, addrs);
-            }
-        }
 
         // Per-wallet keyless rehydration payload; the manager rebuilds each
         // wallet watch-only and derives signing keys later on demand.
         let wallet_ids = schema::wallets::list_ids(&conn).map_err(PersistenceError::from)?;
         let wallets_seen = wallet_ids.len();
         for wallet_id in wallet_ids {
-            if unreadable.contains(&wallet_id) {
-                continue;
-            }
             match load_one_wallet(&conn, wallet_id, &ctx, self.config.load_policy) {
-                Ok(wallet_state) => {
+                Ok((wallet_state, (addrs, count))) => {
+                    // Publish both halves only after the wallet's transaction succeeds.
+                    if count > 0
+                        || !addrs.per_account.is_empty()
+                        || addrs.sync_height > 0
+                        || addrs.sync_timestamp > 0
+                        || addrs.last_known_recent_block > 0
+                    {
+                        addresses_loaded += count;
+                        state.platform_addresses.insert(wallet_id, addrs);
+                    }
                     state.wallets.insert(wallet_id, wallet_state);
                 }
                 // The isolation boundary. One wallet's failure is recorded
@@ -1689,11 +1669,13 @@ fn load_one_wallet(
     wallet_id: WalletId,
     ctx: &LoadCtx,
     policy: LoadPolicy,
-) -> Result<platform_wallet::changeset::ClientWalletStartState, PersistenceError> {
+) -> Result<(ClientWalletStartState, schema::platform_addrs::LoadAllEntry), PersistenceError> {
     let tx = rusqlite::Transaction::new_unchecked(conn, rusqlite::TransactionBehavior::Immediate)
         .map_err(WalletStorageError::from)
         .map_err(PersistenceError::from)?;
     let counted_before = ctx.counts_checkpoint();
+    let platform_state =
+        schema::platform_addrs::load_entry(&tx, &wallet_id).map_err(PersistenceError::from)?;
     let (mut state, sweeps) = load_wallet_snapshot(&tx, wallet_id, ctx)?;
     if !sweeps.is_empty() {
         schema::core_state::apply_replay_sweeps(&tx, &wallet_id, sweeps)
@@ -1722,7 +1704,7 @@ fn load_one_wallet(
     }
     .map_err(WalletStorageError::from)
     .map_err(PersistenceError::from)?;
-    Ok(state)
+    Ok((state, platform_state))
 }
 
 fn load_wallet_snapshot(
@@ -1974,16 +1956,6 @@ fn record_wallet_degradation(
     Ok(())
 }
 
-/// [`record_wallet_degradation`] for a failure that is already a typed
-/// storage error, so its own kind tag is the cause.
-fn degrade_whole_wallet(
-    ctx: &LoadCtx,
-    wallet_id: WalletId,
-    original: &WalletStorageError,
-) -> Result<(), WalletStorageError> {
-    record_wallet_degradation(ctx, wallet_id, original.error_kind_str(), original)
-}
-
 /// The kind tag of the typed storage error inside a persistence error.
 ///
 /// The boundary attributes a dropped wallet to its cause, and the cause is

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

@@ -84,7 +84,6 @@ fn decode_platform_payment_row(
 
 /// Every `platform_payment` registration for one wallet, decoded into
 /// `(account_index, xpub)`.
-#[cfg(any(test, feature = "__test-helpers"))]
 pub(crate) fn list_platform_payment_registrations(
     conn: &Connection,
     wallet_id: &WalletId,
@@ -109,8 +108,7 @@ pub(crate) fn list_platform_payment_registrations(
 
 /// Bulk variant of [`list_platform_payment_registrations`]: every
 /// wallet's `platform_payment` registrations in one scan, grouped by
-/// `wallet_id`. Used by `load()` to avoid a per-wallet registrations
-/// query.
+/// `wallet_id`.
 pub(crate) fn all_platform_payment_registrations(
     conn: &Connection,
 ) -> Result<

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

@@ -112,36 +112,28 @@ pub struct PlatformAddressRow {
     pub funds: AddressFunds,
 }
 
-#[cfg(any(test, feature = "__test-helpers"))]
+/// Decode every address row for one wallet, including unregistered accounts.
 pub fn list_per_wallet(
     conn: &Connection,
     wallet_id: &WalletId,
 ) -> Result<Vec<PlatformAddressRow>, WalletStorageError> {
     let mut stmt = conn.prepare(
-        "SELECT account_index, address_index, address, balance, nonce, as_of_height \
+        "SELECT account_index, address_index, length(address), address, balance, nonce, as_of_height \
          FROM platform_addresses WHERE wallet_id = ?1 \
          ORDER BY account_index, address_index, address",
     )?;
-    let rows = stmt.query_map(params![wallet_id.as_slice()], |row| {
-        Ok((
-            row.get::<_, i64>(0)?,
-            row.get::<_, i64>(1)?,
-            row.get::<_, Vec<u8>>(2)?,
-            row.get::<_, i64>(3)?,
-            row.get::<_, i64>(4)?,
-            row.get::<_, i64>(5)?,
-        ))
-    })?;
+    let mut rows = stmt.query(params![wallet_id.as_slice()])?;
     let mut out = Vec::new();
-    for r in rows {
-        let (account_index, address_index, address_bytes, balance, nonce, as_of_height) = r?;
+    while let Some(row) = rows.next()? {
+        blob::check_fixed_width(row.get::<_, i64>(2)?, 20, "platform_addresses.address")?;
+        let address_bytes: Vec<u8> = row.get(3)?;
         out.push(decode_address_row(
-            account_index,
-            address_index,
+            row.get(0)?,
+            row.get(1)?,
             &address_bytes,
-            balance,
-            nonce,
-            as_of_height,
+            row.get(4)?,
+            row.get(5)?,
+            row.get(6)?,
         )?);
     }
     Ok(out)
@@ -196,14 +188,19 @@ fn account_state_from_rows(
 /// Build `PlatformAddressSyncStartState` for one wallet: the
 /// network-scoped sync watermark plus the per-account committed state
 /// reconstructed from registrations + address rows.
-///
-/// `load()` uses the grouped [`load_all`] path; this per-wallet form is
-/// retained for this crate's integration tests.
 #[cfg(any(test, feature = "__test-helpers"))]
 pub fn load_state(
     conn: &Connection,
     wallet_id: &WalletId,
 ) -> Result<PlatformAddressSyncStartState, WalletStorageError> {
+    load_entry(conn, wallet_id).map(|(state, _)| state)
+}
+
+/// Read one wallet inside the caller's snapshot, retaining the raw address count.
+pub(crate) fn load_entry(
+    conn: &Connection,
+    wallet_id: &WalletId,
+) -> Result<LoadAllEntry, WalletStorageError> {
     let row: Option<(i64, i64, i64)> = conn
         .query_row(
             "SELECT sync_height, sync_timestamp, last_known_recent_block \
@@ -222,12 +219,54 @@ pub fn load_state(
     };
     let registrations = accounts::list_platform_payment_registrations(conn, wallet_id)?;
     let address_rows = list_per_wallet(conn, wallet_id)?;
-    Ok(PlatformAddressSyncStartState {
-        per_account: build_per_account(&registrations, &address_rows),
-        sync_height: h,
-        sync_timestamp: t,
-        last_known_recent_block: r,
-    })
+    Ok((
+        PlatformAddressSyncStartState {
+            per_account: build_per_account(&registrations, &address_rows),
+            sync_height: h,
+            sync_timestamp: t,
+            last_known_recent_block: r,
+        },
+        address_rows.len(),
+    ))
+}
+
+/// Reject unattributable corrupt IDs before per-wallet isolation skips them.
+pub(crate) fn validate_wallet_ids(conn: &Connection) -> Result<(), WalletStorageError> {
+    for (sql, column) in [
+        (
+            "SELECT length(wallet_id), wallet_id FROM platform_address_sync \
+             WHERE length(wallet_id) != 32 OR typeof(wallet_id) != 'blob' LIMIT 1",
+            "platform_addresses.wallet_id",
+        ),
+        (
+            "SELECT length(wallet_id), wallet_id FROM platform_addresses \
+             WHERE length(wallet_id) != 32 OR typeof(wallet_id) != 'blob' LIMIT 1",
+            "platform_addresses.wallet_id",
+        ),
+        (
+            "SELECT length(wallet_id), wallet_id FROM account_registrations \
+             WHERE account_type = 'platform_payment' \
+             AND (length(wallet_id) != 32 OR typeof(wallet_id) != 'blob') LIMIT 1",
+            "account_registrations.wallet_id",
+        ),
+    ] {
+        let mut stmt = conn.prepare(sql)?;
+        let mut rows = stmt.query([])?;
+        if let Some(row) = rows.next()? {
+            let length: i64 = row.get(0)?;
+            if column == "account_registrations.wallet_id" {
+                blob::check_fixed_width(length, 32, column)?;
+            } else if length != 32 {
+                return Err(WalletStorageError::InvalidWalletIdLength {
+                    column,
+                    actual: usize::try_from(length).unwrap_or(usize::MAX),
+                });
+            }
+            let bytes: Vec<u8> = row.get(1)?;
+            super::id32(column, &bytes)?;
+        }
+    }
+    Ok(())
 }
 
 /// Total `platform_addresses` row count per wallet — used by tests
@@ -279,7 +318,7 @@ fn fail_wallet<T>(out: &mut PerWallet<T>, wallet_id: WalletId, err: WalletStorag
     }
 }
 
-/// Bulk reader for `load()`. Cost is a fixed number of grouped scans —
+/// Bulk reader. Cost is a fixed number of grouped scans —
 /// one over `platform_address_sync`, one over `platform_addresses`, and
 /// one over the `platform_payment` `account_registrations` — regardless
 /// of wallet count, rather than a per-wallet fan-out.