fix(platform-wallet): keep the cached identity in step with the chain around updates

packages/rs-platform-wallet-ffi/src/identity_update.rs

@@ -311,6 +311,58 @@ pub unsafe extern "C" fn platform_wallet_update_identity_with_signer(
     disable_public_key_ids: *const u32,
     disable_public_key_ids_count: usize,
     signer_handle: *mut SignerHandle,
+) -> PlatformWalletFFIResult {
+    update_identity_with_signer(
+        wallet_handle,
+        identity_id,
+        add_public_keys,
+        add_public_keys_count,
+        disable_public_key_ids,
+        disable_public_key_ids_count,
+        None,
+        signer_handle,
+    )
+}
+
+/// [`platform_wallet_update_identity_with_signer`], with the revision the
+/// update must carry. When the cached identity would give the transition
+/// another revision (it moved since the caller read it), the call fails
+/// before anything is signed or broadcast, so a caller that recorded
+/// `revision` knows the only revision this update can land at.
+#[no_mangle]
+#[allow(clippy::too_many_arguments)]
+pub unsafe extern "C" fn platform_wallet_update_identity_with_signer_at_revision(
+    wallet_handle: Handle,
+    identity_id: *const u8,
+    add_public_keys: *const IdentityPubkeyFFI,
+    add_public_keys_count: usize,
+    disable_public_key_ids: *const u32,
+    disable_public_key_ids_count: usize,
+    revision: u64,
+    signer_handle: *mut SignerHandle,
+) -> PlatformWalletFFIResult {
+    update_identity_with_signer(
+        wallet_handle,
+        identity_id,
+        add_public_keys,
+        add_public_keys_count,
+        disable_public_key_ids,
+        disable_public_key_ids_count,
+        Some(revision),
+        signer_handle,
+    )
+}
+
+#[allow(clippy::too_many_arguments)]
+unsafe fn update_identity_with_signer(
+    wallet_handle: Handle,
+    identity_id: *const u8,
+    add_public_keys: *const IdentityPubkeyFFI,
+    add_public_keys_count: usize,
+    disable_public_key_ids: *const u32,
+    disable_public_key_ids_count: usize,
+    revision: Option<u64>,
+    signer_handle: *mut SignerHandle,
 ) -> PlatformWalletFFIResult {
     check_ptr!(signer_handle);
 
@@ -375,7 +427,14 @@ pub unsafe extern "C" fn platform_wallet_update_identity_with_signer(
         block_on_worker(async move {
             let signer: &VTableSigner = &*(signer_addr as *const VTableSigner);
             identity_wallet
-                .update_identity_with_external_signer(&id, add_keys, disable_ids, signer, None)
+                .update_identity_with_external_signer_at_revision(
+                    &id,
+                    add_keys,
+                    disable_ids,
+                    revision,
+                    signer,
+                    None,
+                )
                 .await
         })
     });

packages/rs-platform-wallet/src/error.rs

@@ -77,6 +77,20 @@ pub enum PlatformWalletError {
     #[error("Invalid identity data: {0}")]
     InvalidIdentityData(String),
 
+    /// An identity update was asked to carry a revision the cached identity
+    /// would not give it: the cache moved since the caller read the
+    /// revision, or the caller read it elsewhere. Raised before the
+    /// transition is signed, so nothing was broadcast.
+    #[error(
+        "identity {identity_id}: the update was asked to carry revision {expected}, \
+         but the cached identity would give it {actual}; nothing was broadcast"
+    )]
+    IdentityRevisionMismatch {
+        identity_id: Identifier,
+        expected: u64,
+        actual: u64,
+    },
+
     /// A DashPay invitation link's voucher key is for the other network (a
     /// testnet link opened in a mainnet wallet, or the reverse). Definitive:
     /// the claim refuses the link before any network work.

packages/rs-platform-wallet/src/wallet/identity/network/update.rs

@@ -95,6 +95,42 @@ impl IdentityWallet {
         signer: &S,
         settings: Option<PutSettings>,
     ) -> Result<(), PlatformWalletError>
+    where
+        S: Signer<IdentityPublicKey> + Send + Sync,
+    {
+        self.update_identity_with_external_signer_at_revision(
+            identity_id,
+            add_public_keys,
+            disable_public_keys,
+            None,
+            signer,
+            settings,
+        )
+        .await
+    }
+
+    /// [`Self::update_identity_with_external_signer`], with the revision
+    /// the transition must carry.
+    ///
+    /// The transition takes the cached identity's revision plus one, read
+    /// when the update starts. A caller that records the revision before
+    /// submitting (to tell later whether the update can still land: Drive
+    /// takes an identity update only at one above the identity's revision)
+    /// cannot see that read, and a refresh in between moves it. With
+    /// `revision` set, an update that would carry another one fails with
+    /// [`PlatformWalletError::IdentityRevisionMismatch`] before anything is
+    /// signed or broadcast, so the revision the caller recorded is the only
+    /// one this update can land at. `None` takes whatever the cache gives,
+    /// as [`Self::update_identity_with_external_signer`] does.
+    pub async fn update_identity_with_external_signer_at_revision<S>(
+        &self,
+        identity_id: &Identifier,
+        add_public_keys: Vec<IdentityPublicKey>,
+        disable_public_keys: Vec<u32>,
+        revision: Option<u64>,
+        signer: &S,
+        settings: Option<PutSettings>,
+    ) -> Result<(), PlatformWalletError>
     where
         S: Signer<IdentityPublicKey> + Send + Sync,
     {
@@ -131,8 +167,8 @@ impl IdentityWallet {
         };
 
         // Increment revision for the update transition.
-        let original_revision = identity.revision();
-        identity.set_revision(original_revision + 1);
+        let transition_revision = update_revision(identity_id, identity.revision(), revision)?;
+        identity.set_revision(transition_revision);
 
         // Pick the MASTER signing key — DPP requires identity update
         // transitions to be authorized by MASTER specifically.
@@ -223,10 +259,15 @@ impl IdentityWallet {
                 )
             })?;
             if let Some(managed) = info.identity_manager.managed_identity_mut(identity_id) {
-                // Bump the cached revision so a subsequent update
-                // doesn't reuse the pre-broadcast value.
+                // The identity is at the transition's revision now, so a
+                // subsequent update must not reuse the pre-broadcast value.
+                // The cache may have been refreshed since the snapshot (to
+                // this revision or past it), so it takes the higher of the
+                // two instead of adding one to whatever it holds.
                 let cached_revision = managed.identity.revision();
-                managed.identity.set_revision(cached_revision + 1);
+                managed
+                    .identity
+                    .set_revision(cached_revision.max(transition_revision));
 
                 for key in added_keys_for_local_apply {
                     let breadcrumb = identity_index.map(|idx| (self.wallet_id, idx, key.id()));
@@ -332,3 +373,62 @@ impl IdentityWallet {
         Ok(result)
     }
 }
+
+/// The revision an update of an identity cached at `cached` carries: one
+/// above it. With `expected` set, any other value is refused.
+fn update_revision(
+    identity_id: &Identifier,
+    cached: u64,
+    expected: Option<u64>,
+) -> Result<u64, PlatformWalletError> {
+    let next = cached.checked_add(1).ok_or_else(|| {
+        PlatformWalletError::InvalidIdentityData(format!(
+            "identity {identity_id} is at the last revision"
+        ))
+    })?;
+    match expected {
+        Some(expected) if expected != next => Err(PlatformWalletError::IdentityRevisionMismatch {
+            identity_id: *identity_id,
+            expected,
+            actual: next,
+        }),
+        _ => Ok(next),
+    }
+}
+
+#[cfg(test)]
+mod update_revision_tests {
+    use super::*;
+
+    #[test]
+    fn should_carry_one_above_the_cached_revision() {
+        let id = Identifier::from([7u8; 32]);
+        assert_eq!(update_revision(&id, 4, None).unwrap(), 5);
+        assert_eq!(update_revision(&id, 4, Some(5)).unwrap(), 5);
+    }
+
+    #[test]
+    fn should_refuse_an_update_asked_for_another_revision() {
+        let id = Identifier::from([7u8; 32]);
+        // The cache moved past the revision the caller read.
+        match update_revision(&id, 5, Some(5)) {
+            Err(PlatformWalletError::IdentityRevisionMismatch {
+                expected, actual, ..
+            }) => assert_eq!((expected, actual), (5, 6)),
+            other => panic!("unexpected: {other:?}"),
+        }
+        assert!(matches!(
+            update_revision(&id, 3, Some(5)),
+            Err(PlatformWalletError::IdentityRevisionMismatch { .. })
+        ));
+    }
+
+    #[test]
+    fn should_refuse_an_identity_at_the_last_revision() {
+        let id = Identifier::from([7u8; 32]);
+        assert!(matches!(
+            update_revision(&id, u64::MAX, None),
+            Err(PlatformWalletError::InvalidIdentityData(_))
+        ));
+    }
+}

packages/swift-sdk/Sources/SwiftDashSDK/PlatformWallet/ManagedPlatformWallet.swift

@@ -3729,10 +3729,21 @@ extension ManagedPlatformWallet {
     /// otherwise subsequent operations that try to sign with the
     /// newly-added keys will fail. The signer here only signs the
     /// update transition itself with an existing MASTER key.
+    ///
+    /// The transition carries the cached identity's revision plus one,
+    /// read when the Rust side starts the update. With `revision` set, an
+    /// update that would carry another one throws before anything is
+    /// signed or broadcast (as `.unknown`, with a message that says the
+    /// update was asked for another revision and nothing was broadcast),
+    /// so a caller that recorded `revision` beforehand knows the only
+    /// revision this update can land at: Drive takes an identity update
+    /// only at one above the identity's revision. `nil` takes whatever the
+    /// cache gives.
     public func updateIdentity(
         identityId: Identifier,
         addPublicKeys: [ManagedPlatformWallet.IdentityPubkey] = [],
         disablePublicKeyIds: [UInt32] = [],
+        revision: UInt64? = nil,
         signer: KeychainSigner
     ) async throws {
         let handle = self.handle
@@ -3751,30 +3762,25 @@ extension ManagedPlatformWallet {
             let result = idBytes.withUnsafeBufferPointer {
                 idBp -> PlatformWalletFFIResult in
                 disableIds.withUnsafeBufferPointer { disableBp in
+                    let update = { (rows: UnsafePointer<IdentityPubkeyFFI>?, count: Int) -> PlatformWalletFFIResult in
+                        let disable = disableIds.isEmpty ? nil : disableBp.baseAddress
+                        guard let revision else {
+                            return platform_wallet_update_identity_with_signer(
+                                handle, idBp.baseAddress!, rows, UInt(count),
+                                disable, UInt(disableIds.count), signerHandle)
+                        }
+                        return platform_wallet_update_identity_with_signer_at_revision(
+                            handle, idBp.baseAddress!, rows, UInt(count),
+                            disable, UInt(disableIds.count), revision, signerHandle)
+                    }
                     if addPubkeys.isEmpty {
-                        return platform_wallet_update_identity_with_signer(
-                            handle,
-                            idBp.baseAddress!,
-                            nil,
-                            0,
-                            disableIds.isEmpty ? nil : disableBp.baseAddress,
-                            UInt(disableIds.count),
-                            signerHandle
-                        )
+                        return update(nil, 0)
                     }
                     return ManagedPlatformWallet.withPubkeyFFIArray(
                         addPubkeys,
                         buffers: pubkeyBuffers
                     ) { ffiRowsPtr, ffiRowsCount in
-                        platform_wallet_update_identity_with_signer(
-                            handle,
-                            idBp.baseAddress!,
-                            ffiRowsPtr,
-                            UInt(ffiRowsCount),
-                            disableIds.isEmpty ? nil : disableBp.baseAddress,
-                            UInt(disableIds.count),
-                            signerHandle
-                        )
+                        update(ffiRowsPtr, ffiRowsCount)
                     }
                 }
             }