Browse Source

Introduce ProofsWithState for proof state management (#1488)

* Introduce ProofsWithState for atomic proof state management

Replace individual proof state operations with a unified ProofsWithState type
that enforces the invariant that all proofs in a set share the same state. This
shifts responsibility for state consistency to the database layer and
simplifies state transition logic in the saga implementations.

Changes to ProofsTransaction trait:
- add_proofs now returns Acquired<ProofsWithState>
- update_proofs_states(ys, state) -> update_proofs(&mut
  Acquired<ProofsWithState>)
- get_proofs_states(ys) -> get_proofs(ys) returning Acquired<ProofsWithState>

Benefits:
- State transitions validated in memory before persisting
- Eliminates scattered state checking in saga code
- Database layer guarantees proof state consistency
- Cleaner API with state encapsulated in the type

* Add test for get_proofs when some proofs not found

Add get_proofs_fails_when_some_not_found test that verifies the database
returns an error when not all requested proofs exist, rather than silently
returning a partial result.

Fix get_proofs implementation to check that the number of rows returned matches
the number of requested Y values, returning ProofNotFound if any proofs are
missing.
C 2 tuần trước cách đây
mục cha
commit
7d135270b3

+ 18 - 9
crates/cdk-common/src/database/mint/mod.rs

@@ -8,7 +8,9 @@ use cashu::Amount;
 
 use super::{DbTransactionFinalizer, Error};
 use crate::database::Acquired;
-use crate::mint::{self, MeltQuote, MintKeySetInfo, MintQuote as MintMintQuote, Operation};
+use crate::mint::{
+    self, MeltQuote, MintKeySetInfo, MintQuote as MintMintQuote, Operation, ProofsWithState,
+};
 use crate::nuts::{
     BlindSignature, BlindedMessage, CurrencyUnit, Id, MeltQuoteState, Proof, Proofs, PublicKey,
     State,
@@ -297,19 +299,23 @@ pub trait ProofsTransaction {
         proof: Proofs,
         quote_id: Option<QuoteId>,
         operation: &Operation,
-    ) -> Result<(), Self::Err>;
-    /// Updates the proofs to a given states and return the previous states
-    async fn update_proofs_states(
+    ) -> Result<Acquired<ProofsWithState>, Self::Err>;
+
+    /// Updates the proofs to the given state in the database.
+    ///
+    /// Also updates the `state` field on the [`ProofsWithState`] wrapper to reflect
+    /// the new state after the database update succeeds.
+    async fn update_proofs_state(
         &mut self,
-        ys: &[PublicKey],
-        proofs_state: State,
-    ) -> Result<Vec<Option<State>>, Self::Err>;
+        proofs: &mut Acquired<ProofsWithState>,
+        new_state: State,
+    ) -> Result<(), Self::Err>;
 
     /// get proofs states
-    async fn get_proofs_states(
+    async fn get_proofs(
         &mut self,
         ys: &[PublicKey],
-    ) -> Result<Vec<Option<State>>, Self::Err>;
+    ) -> Result<Acquired<ProofsWithState>, Self::Err>;
 
     /// Remove [`Proofs`]
     async fn remove_proofs(
@@ -521,3 +527,6 @@ pub trait Database<Error>:
 
 /// Type alias for Mint Database
 pub type DynMintDatabase = std::sync::Arc<dyn Database<Error> + Send + Sync>;
+
+/// Type alias for Mint Transaction
+pub type DynMintTransaction = Box<dyn Transaction<Error> + Send + Sync>;

+ 3 - 0
crates/cdk-common/src/database/mint/test/mod.rs

@@ -255,6 +255,9 @@ macro_rules! mint_db_test {
             get_blind_signatures_in_transaction,
             reject_duplicate_payment_ids,
             remove_spent_proofs_should_fail,
+            get_proofs_with_inconsistent_states_fails,
+            get_proofs_fails_when_some_not_found,
+            update_proofs_state_updates_proofs_with_state,
         );
     };
     ($make_db_fn:ident, $($name:ident),+ $(,)?) => {

+ 263 - 14
crates/cdk-common/src/database/mint/test/proofs.rs

@@ -8,6 +8,7 @@ use cashu::{Amount, Id, SecretKey};
 use crate::database::mint::test::setup_keyset;
 use crate::database::mint::{Database, Error, KeysDatabase, Proof, QuoteId};
 use crate::mint::Operation;
+use crate::state::check_state_transition;
 
 /// Test get proofs by keyset id
 pub async fn get_proofs_by_keyset_id<DB>(db: DB)
@@ -217,7 +218,11 @@ where
 
     // Update to pending
     let mut tx = Database::begin_transaction(&db).await.unwrap();
-    let _old_states = tx.update_proofs_states(&ys, State::Pending).await.unwrap();
+    let mut proofs = tx.get_proofs(&ys).await.unwrap();
+    check_state_transition(proofs.state, State::Pending).unwrap();
+    tx.update_proofs_state(&mut proofs, State::Pending)
+        .await
+        .unwrap();
     tx.commit().await.unwrap();
 
     // Verify new state
@@ -227,8 +232,11 @@ where
 
     // Update to spent
     let mut tx = Database::begin_transaction(&db).await.unwrap();
-    let old_states = tx.update_proofs_states(&ys, State::Spent).await.unwrap();
-    assert_eq!(old_states, vec![Some(State::Pending), Some(State::Pending)]);
+    let mut proofs = tx.get_proofs(&ys).await.unwrap();
+    check_state_transition(proofs.state, State::Spent).unwrap();
+    tx.update_proofs_state(&mut proofs, State::Spent)
+        .await
+        .unwrap();
     tx.commit().await.unwrap();
 
     // Verify final state
@@ -237,6 +245,66 @@ where
     assert_eq!(states[1], Some(State::Spent));
 }
 
+/// Test that update_proofs_state updates the ProofsWithState.state field
+pub async fn update_proofs_state_updates_proofs_with_state<DB>(db: DB)
+where
+    DB: Database<Error> + KeysDatabase<Err = Error>,
+{
+    use cashu::State;
+
+    let keyset_id = setup_keyset(&db).await;
+    let quote_id = QuoteId::new_uuid();
+
+    let proofs = vec![Proof {
+        amount: Amount::from(100),
+        keyset_id,
+        secret: Secret::generate(),
+        c: SecretKey::generate().public_key(),
+        witness: None,
+        dleq: None,
+    }];
+
+    let ys: Vec<_> = proofs.iter().map(|p| p.y().unwrap()).collect();
+
+    // Add proofs and verify initial state
+    let mut tx = Database::begin_transaction(&db).await.unwrap();
+    let mut proofs = tx
+        .add_proofs(
+            proofs.clone(),
+            Some(quote_id),
+            &Operation::new_swap(Amount::ZERO, Amount::ZERO, Amount::ZERO),
+        )
+        .await
+        .unwrap();
+    assert_eq!(proofs.state, State::Unspent);
+
+    // Update to Pending and verify ProofsWithState.state is updated
+    tx.update_proofs_state(&mut proofs, State::Pending)
+        .await
+        .unwrap();
+    assert_eq!(
+        proofs.state,
+        State::Pending,
+        "ProofsWithState.state should be updated to Pending after update_proofs_state"
+    );
+    tx.commit().await.unwrap();
+
+    // Get proofs again and update to Spent
+    let mut tx = Database::begin_transaction(&db).await.unwrap();
+    let mut proofs = tx.get_proofs(&ys).await.unwrap();
+    assert_eq!(proofs.state, State::Pending);
+
+    tx.update_proofs_state(&mut proofs, State::Spent)
+        .await
+        .unwrap();
+    assert_eq!(
+        proofs.state,
+        State::Spent,
+        "ProofsWithState.state should be updated to Spent after update_proofs_state"
+    );
+    tx.commit().await.unwrap();
+}
+
 /// Test removing proofs
 pub async fn remove_proofs<DB>(db: DB)
 where
@@ -346,14 +414,18 @@ where
 
     // First update to Pending (valid state transition)
     let mut tx = Database::begin_transaction(&db).await.unwrap();
-    tx.update_proofs_states(&[ys[0], ys[1]], State::Pending)
+    let mut proofs = tx.get_proofs(&[ys[0], ys[1]]).await.unwrap();
+    check_state_transition(proofs.state, State::Pending).unwrap();
+    tx.update_proofs_state(&mut proofs, State::Pending)
         .await
         .unwrap();
     tx.commit().await.unwrap();
 
     // Then mark some as spent
     let mut tx = Database::begin_transaction(&db).await.unwrap();
-    tx.update_proofs_states(&[ys[0], ys[1]], State::Spent)
+    let mut proofs = tx.get_proofs(&[ys[0], ys[1]]).await.unwrap();
+    check_state_transition(proofs.state, State::Spent).unwrap();
+    tx.update_proofs_state(&mut proofs, State::Spent)
         .await
         .unwrap();
     tx.commit().await.unwrap();
@@ -705,26 +777,30 @@ where
 
     // Transition proofs to Pending state
     let mut tx = Database::begin_transaction(&db).await.unwrap();
-    let _records = tx
-        .get_proof_ys_by_quote_id(&quote_id)
+    let mut records = tx.get_proofs(&ys).await.expect("valid records");
+    check_state_transition(records.state, State::Pending).unwrap();
+    tx.update_proofs_state(&mut records, State::Pending)
         .await
-        .expect("valid records");
-    tx.update_proofs_states(&ys, State::Pending).await.unwrap();
+        .unwrap();
     tx.commit().await.unwrap();
 
     // Removing Pending proofs should also succeed
     let mut tx = Database::begin_transaction(&db).await.unwrap();
     let result = tx.remove_proofs(&[ys[0]], Some(quote_id.clone())).await;
-    assert!(result.is_ok(), "Removing Pending proof should succeed");
+    assert!(
+        result.is_ok(),
+        "Removing Pending proof should succeed: {:?}",
+        result,
+    );
     tx.rollback().await.unwrap(); // Rollback to keep proofs for next test
 
     // Now transition proofs to Spent state
     let mut tx = Database::begin_transaction(&db).await.unwrap();
-    let _records = tx
-        .get_proof_ys_by_quote_id(&quote_id)
+    let mut records = tx.get_proofs(&ys).await.expect("valid records");
+    check_state_transition(records.state, State::Spent).unwrap();
+    tx.update_proofs_state(&mut records, State::Spent)
         .await
-        .expect("valid records");
-    tx.update_proofs_states(&ys, State::Spent).await.unwrap();
+        .unwrap();
     tx.commit().await.unwrap();
 
     // Verify proofs are now in Spent state
@@ -762,3 +838,176 @@ where
         "Second proof should still exist"
     );
 }
+
+/// Test that get_proofs fails when proofs have inconsistent states
+///
+/// This validates the database layer's responsibility to ensure all proofs
+/// returned by get_proofs share the same state. The mint never needs proofs
+/// with different states, so this is an invariant the database must enforce.
+pub async fn get_proofs_with_inconsistent_states_fails<DB>(db: DB)
+where
+    DB: Database<Error> + KeysDatabase<Err = Error>,
+{
+    use cashu::State;
+
+    let keyset_id = setup_keyset(&db).await;
+    let quote_id = QuoteId::new_uuid();
+
+    // Create three proofs
+    let proofs = vec![
+        Proof {
+            amount: Amount::from(100),
+            keyset_id,
+            secret: Secret::generate(),
+            c: SecretKey::generate().public_key(),
+            witness: None,
+            dleq: None,
+        },
+        Proof {
+            amount: Amount::from(200),
+            keyset_id,
+            secret: Secret::generate(),
+            c: SecretKey::generate().public_key(),
+            witness: None,
+            dleq: None,
+        },
+        Proof {
+            amount: Amount::from(300),
+            keyset_id,
+            secret: Secret::generate(),
+            c: SecretKey::generate().public_key(),
+            witness: None,
+            dleq: None,
+        },
+    ];
+
+    let ys: Vec<_> = proofs.iter().map(|p| p.y().unwrap()).collect();
+
+    // Add all proofs (initial state is Unspent)
+    let mut tx = Database::begin_transaction(&db).await.unwrap();
+    tx.add_proofs(
+        proofs,
+        Some(quote_id),
+        &Operation::new_swap(Amount::ZERO, Amount::ZERO, Amount::ZERO),
+    )
+    .await
+    .unwrap();
+    tx.commit().await.unwrap();
+
+    // Transition only the first two proofs to Pending state
+    let mut tx = Database::begin_transaction(&db).await.unwrap();
+    let mut first_two_proofs = tx.get_proofs(&ys[0..2]).await.unwrap();
+    check_state_transition(first_two_proofs.state, State::Pending).unwrap();
+    tx.update_proofs_state(&mut first_two_proofs, State::Pending)
+        .await
+        .unwrap();
+    tx.commit().await.unwrap();
+
+    // Verify the states are now inconsistent via get_proofs_states
+    let states = db.get_proofs_states(&ys).await.unwrap();
+    assert_eq!(
+        states[0],
+        Some(State::Pending),
+        "First proof should be Pending"
+    );
+    assert_eq!(
+        states[1],
+        Some(State::Pending),
+        "Second proof should be Pending"
+    );
+    assert_eq!(
+        states[2],
+        Some(State::Unspent),
+        "Third proof should be Unspent"
+    );
+
+    // Now try to get all three proofs via get_proofs - this should fail
+    // because the proofs have inconsistent states
+    let mut tx = Database::begin_transaction(&db).await.unwrap();
+    let result = tx.get_proofs(&ys).await;
+
+    assert!(
+        result.is_err(),
+        "get_proofs should fail when proofs have inconsistent states"
+    );
+
+    tx.rollback().await.unwrap();
+}
+
+/// Test that get_proofs fails when some requested proofs don't exist
+///
+/// This validates that the database returns an error when not all requested
+/// proofs are found, rather than silently returning a partial result.
+pub async fn get_proofs_fails_when_some_not_found<DB>(db: DB)
+where
+    DB: Database<Error> + KeysDatabase<Err = Error>,
+{
+    let keyset_id = setup_keyset(&db).await;
+    let quote_id = QuoteId::new_uuid();
+
+    // Create two proofs that will be stored
+    let stored_proofs = vec![
+        Proof {
+            amount: Amount::from(100),
+            keyset_id,
+            secret: Secret::generate(),
+            c: SecretKey::generate().public_key(),
+            witness: None,
+            dleq: None,
+        },
+        Proof {
+            amount: Amount::from(200),
+            keyset_id,
+            secret: Secret::generate(),
+            c: SecretKey::generate().public_key(),
+            witness: None,
+            dleq: None,
+        },
+    ];
+
+    // Create a third proof that will NOT be stored
+    let non_existent_proof = Proof {
+        amount: Amount::from(300),
+        keyset_id,
+        secret: Secret::generate(),
+        c: SecretKey::generate().public_key(),
+        witness: None,
+        dleq: None,
+    };
+
+    let stored_ys: Vec<_> = stored_proofs.iter().map(|p| p.y().unwrap()).collect();
+    let non_existent_y = non_existent_proof.y().unwrap();
+
+    // Add only the first two proofs
+    let mut tx = Database::begin_transaction(&db).await.unwrap();
+    tx.add_proofs(
+        stored_proofs,
+        Some(quote_id),
+        &Operation::new_swap(Amount::ZERO, Amount::ZERO, Amount::ZERO),
+    )
+    .await
+    .unwrap();
+    tx.commit().await.unwrap();
+
+    // Verify the stored proofs exist
+    let states = db.get_proofs_states(&stored_ys).await.unwrap();
+    assert_eq!(states.len(), 2);
+    assert!(states[0].is_some(), "First proof should exist");
+    assert!(states[1].is_some(), "Second proof should exist");
+
+    // Verify the non-existent proof doesn't exist
+    let states = db.get_proofs_states(&[non_existent_y]).await.unwrap();
+    assert_eq!(states[0], None, "Third proof should not exist");
+
+    // Now try to get all three proofs (2 exist, 1 doesn't) - this should fail
+    let all_ys = vec![stored_ys[0], stored_ys[1], non_existent_y];
+    let mut tx = Database::begin_transaction(&db).await.unwrap();
+    let result = tx.get_proofs(&all_ys).await;
+
+    assert!(
+        result.is_err(),
+        "get_proofs should fail when some proofs don't exist (got 2 of 3)"
+    );
+
+    tx.rollback().await.unwrap();
+}

+ 5 - 4
crates/cdk-common/src/database/mod.rs

@@ -20,10 +20,11 @@ pub type DynKVStore = std::sync::Arc<dyn KVStore<Err = Error> + Send + Sync>;
 
 #[cfg(feature = "mint")]
 pub use mint::{
-    Database as MintDatabase, DynMintDatabase, KeysDatabase as MintKeysDatabase,
-    KeysDatabaseTransaction as MintKeyDatabaseTransaction, ProofsDatabase as MintProofsDatabase,
-    ProofsTransaction as MintProofsTransaction, QuotesDatabase as MintQuotesDatabase,
-    QuotesTransaction as MintQuotesTransaction, SignaturesDatabase as MintSignaturesDatabase,
+    Database as MintDatabase, DynMintDatabase, DynMintTransaction,
+    KeysDatabase as MintKeysDatabase, KeysDatabaseTransaction as MintKeyDatabaseTransaction,
+    ProofsDatabase as MintProofsDatabase, ProofsTransaction as MintProofsTransaction,
+    QuotesDatabase as MintQuotesDatabase, QuotesTransaction as MintQuotesTransaction,
+    SignaturesDatabase as MintSignaturesDatabase,
     SignaturesTransaction as MintSignatureTransaction, Transaction as MintTransaction,
 };
 #[cfg(all(feature = "mint", feature = "auth"))]

+ 63 - 1
crates/cdk-common/src/mint.rs

@@ -1,6 +1,7 @@
 //! Mint types
 
 use std::fmt;
+use std::ops::Deref;
 use std::str::FromStr;
 
 use bitcoin::bip32::DerivationPath;
@@ -8,7 +9,7 @@ use cashu::quote_id::QuoteId;
 use cashu::util::unix_time;
 use cashu::{
     Bolt11Invoice, MeltOptions, MeltQuoteBolt11Response, MintQuoteBolt11Response,
-    MintQuoteBolt12Response, PaymentMethod,
+    MintQuoteBolt12Response, PaymentMethod, Proofs, State,
 };
 use lightning::offers::offer::Offer;
 use serde::{Deserialize, Serialize};
@@ -31,6 +32,67 @@ pub enum OperationKind {
     Melt,
 }
 
+/// A collection of proofs that share a common state.
+///
+/// This type enforces the invariant that all proofs in the collection have the same state.
+/// The mint never needs to operate on a set of proofs with different states - proofs are
+/// always processed together as a unit (e.g., during swap, melt, or mint operations).
+///
+/// # Database Layer Responsibility
+///
+/// This design shifts the responsibility of ensuring state consistency to the database layer.
+/// When the database retrieves proofs via [`get_proofs`](crate::database::mint::ProofsTransaction::get_proofs),
+/// it must verify that all requested proofs share the same state and return an error if they don't.
+/// This prevents invalid proof sets from propagating through the system.
+///
+/// # State Transitions
+///
+/// State transitions are validated using [`check_state_transition`](crate::state::check_state_transition)
+/// before updating. The database layer then persists the new state for all proofs in a single transaction
+/// via [`update_proofs_state`](crate::database::mint::ProofsTransaction::update_proofs_state).
+///
+/// # Example
+///
+/// ```ignore
+/// // Database layer ensures all proofs have the same state
+/// let mut proofs = tx.get_proofs(&ys).await?;
+///
+/// // Validate the state transition
+/// check_state_transition(proofs.state, State::Spent)?;
+///
+/// // Persist the state change
+/// tx.update_proofs_state(&mut proofs, State::Spent).await?;
+/// ```
+#[derive(Debug)]
+pub struct ProofsWithState {
+    proofs: Proofs,
+    /// The current state of the proofs
+    pub state: State,
+}
+
+impl Deref for ProofsWithState {
+    type Target = Proofs;
+
+    fn deref(&self) -> &Self::Target {
+        &self.proofs
+    }
+}
+
+impl ProofsWithState {
+    /// Creates a new `ProofsWithState` with the given proofs and their shared state.
+    ///
+    /// # Note
+    ///
+    /// This constructor assumes all proofs share the given state. It is typically
+    /// called by the database layer after verifying state consistency.
+    pub fn new(proofs: Proofs, current_state: State) -> Self {
+        Self {
+            proofs,
+            state: current_state,
+        }
+    }
+}
+
 impl fmt::Display for OperationKind {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {

+ 111 - 44
crates/cdk-sql-common/src/mint/proofs.rs

@@ -4,11 +4,12 @@ use std::collections::HashMap;
 use std::str::FromStr;
 
 use async_trait::async_trait;
-use cdk_common::database::{self, Error, MintProofsDatabase};
-use cdk_common::mint::Operation;
+use cdk_common::database::{self, Acquired, Error, MintProofsDatabase};
+use cdk_common::mint::{Operation, ProofsWithState};
 use cdk_common::nut00::ProofsMethods;
 use cdk_common::quote_id::QuoteId;
 use cdk_common::secret::Secret;
+use cdk_common::util::unix_time;
 use cdk_common::{Amount, Id, Proof, Proofs, PublicKey, State};
 
 use super::{SQLMintDatabase, SQLTransaction};
@@ -69,9 +70,7 @@ pub(super) fn sql_row_to_proof(row: Vec<Column>) -> Result<Proof, Error> {
     })
 }
 
-pub(super) fn sql_row_to_proof_with_state(
-    row: Vec<Column>,
-) -> Result<(Proof, Option<State>), Error> {
+pub(super) fn sql_row_to_proof_with_state(row: Vec<Column>) -> Result<(Proof, State), Error> {
     unpack_into!(
         let (
             keyset_id, amount, secret, c, witness, state
@@ -79,7 +78,9 @@ pub(super) fn sql_row_to_proof_with_state(
     );
 
     let amount: u64 = column_as_number!(amount);
-    let state = column_as_nullable_string!(state).and_then(|s| State::from_str(&s).ok());
+    let state = column_as_nullable_string!(state)
+        .and_then(|s| State::from_str(&s).ok())
+        .unwrap_or(State::Pending);
 
     Ok((
         Proof {
@@ -116,13 +117,22 @@ where
 {
     type Err = Error;
 
+    /// Adds proofs to the database with initial state `Unspent`.
+    ///
+    /// This method first checks if any of the proofs already exist in the database.
+    /// If a proof exists and is spent, returns [`Error::AttemptUpdateSpentProof`].
+    /// If a proof exists in any other state, returns [`Error::Duplicate`].
+    ///
+    /// On success, returns the proofs wrapped in [`Acquired<ProofsWithState>`] with
+    /// state set to `Unspent`, indicating the rows are locked for the duration of
+    /// the transaction.
     async fn add_proofs(
         &mut self,
         proofs: Proofs,
         quote_id: Option<QuoteId>,
         operation: &Operation,
-    ) -> Result<(), Self::Err> {
-        let current_time = cdk_common::util::unix_time();
+    ) -> Result<Acquired<ProofsWithState>, Self::Err> {
+        let current_time = unix_time();
 
         // Check any previous proof, this query should return None in order to proceed storing
         // Any result here would error
@@ -144,7 +154,7 @@ where
             None => Ok(()), // no previous record
         }?;
 
-        for proof in proofs {
+        for proof in &proofs {
             let y = proof.y()?;
 
             query(
@@ -162,7 +172,7 @@ where
             .bind("c", proof.c.to_bytes().to_vec())
             .bind(
                 "witness",
-                proof.witness.and_then(|w| serde_json::to_string(&w).inspect_err(|e| tracing::error!("Failed to serialize witness: {:?}", e)).ok()),
+                proof.witness.clone().and_then(|w| serde_json::to_string(&w).inspect_err(|e| tracing::error!("Failed to serialize witness: {:?}", e)).ok()),
             )
             .bind("state", "UNSPENT".to_string())
             .bind("quote_id", quote_id.clone().map(|q| q.to_string()))
@@ -173,24 +183,27 @@ where
             .await?;
         }
 
-        Ok(())
+        Ok(ProofsWithState::new(proofs, State::Unspent).into())
     }
 
-    async fn update_proofs_states(
+    /// Updates all proofs to the given state in the database.
+    ///
+    /// Also updates the `state` field on the [`ProofsWithState`] wrapper to reflect
+    /// the new state after the database update succeeds.
+    ///
+    /// When the new state is `Spent`, this method also updates the `keyset_amounts`
+    /// table to track the total redeemed amount per keyset for analytics purposes.
+    ///
+    /// # Prerequisites
+    ///
+    /// The proofs must have been previously acquired via `add_proofs`
+    /// or `get_proofs` to ensure they are locked within the current transaction.
+    async fn update_proofs_state(
         &mut self,
-        ys: &[PublicKey],
+        proofs: &mut Acquired<ProofsWithState>,
         new_state: State,
-    ) -> Result<Vec<Option<State>>, Self::Err> {
-        let mut current_states = get_current_states(&self.inner, ys, true).await?;
-
-        if current_states.len() != ys.len() {
-            tracing::warn!(
-                "Attempted to update state of non-existent proof {} {}",
-                current_states.len(),
-                ys.len()
-            );
-            return Err(database::Error::ProofNotFound);
-        }
+    ) -> Result<(), Self::Err> {
+        let ys = proofs.ys()?;
 
         query(r#"UPDATE proof SET state = :new_state WHERE y IN (:ys)"#)?
             .bind("new_state", new_state.to_string())
@@ -200,22 +213,24 @@ where
 
         if new_state == State::Spent {
             query(
-                r#"
-                INSERT INTO keyset_amounts (keyset_id, total_issued, total_redeemed)
-                SELECT keyset_id, 0, COALESCE(SUM(amount), 0)
-                FROM proof
-                WHERE y IN (:ys)
-                GROUP BY keyset_id
-                ON CONFLICT (keyset_id)
-                DO UPDATE SET total_redeemed = keyset_amounts.total_redeemed + EXCLUDED.total_redeemed
-                "#,
-            )?
-            .bind_vec("ys", ys.iter().map(|y| y.to_bytes().to_vec()).collect())
-            .execute(&self.inner)
-            .await?;
+                    r#"
+                    INSERT INTO keyset_amounts (keyset_id, total_issued, total_redeemed)
+                    SELECT keyset_id, 0, COALESCE(SUM(amount), 0)
+                    FROM proof
+                    WHERE y IN (:ys)
+                    GROUP BY keyset_id
+                    ON CONFLICT (keyset_id)
+                    DO UPDATE SET total_redeemed = keyset_amounts.total_redeemed + EXCLUDED.total_redeemed
+                    "#,
+                )?
+                .bind_vec("ys", ys.iter().map(|y| y.to_bytes().to_vec()).collect())
+                .execute(&self.inner)
+                .await?;
         }
 
-        Ok(ys.iter().map(|y| current_states.remove(y)).collect())
+        proofs.state = new_state;
+
+        Ok(())
     }
 
     async fn remove_proofs(
@@ -332,13 +347,62 @@ where
         .collect::<Result<Vec<_>, _>>()?)
     }
 
-    async fn get_proofs_states(
+    async fn get_proofs(
         &mut self,
         ys: &[PublicKey],
-    ) -> Result<Vec<Option<State>>, Self::Err> {
-        let mut current_states = get_current_states(&self.inner, ys, true).await?;
+    ) -> Result<Acquired<ProofsWithState>, Self::Err> {
+        if ys.is_empty() {
+            return Err(database::Error::ProofNotFound);
+        }
 
-        Ok(ys.iter().map(|y| current_states.remove(y)).collect())
+        let rows = query(
+            r#"
+             SELECT
+                 keyset_id,
+                 amount,
+                 secret,
+                 c,
+                 witness,
+                 state
+             FROM
+                 proof
+             WHERE
+                 y IN (:ys)
+             FOR UPDATE
+             "#,
+        )?
+        .bind_vec("ys", ys.iter().map(|y| y.to_bytes().to_vec()).collect())
+        .fetch_all(&self.inner)
+        .await?;
+
+        if rows.is_empty() || rows.len() != ys.len() {
+            return Err(database::Error::ProofNotFound);
+        }
+
+        let results: Vec<(Proof, State)> = rows
+            .into_iter()
+            .map(sql_row_to_proof_with_state)
+            .collect::<Result<Vec<_>, _>>()?;
+
+        let mut proofs = Vec::with_capacity(results.len());
+        let mut first_state: Option<State> = None;
+
+        for (proof, state) in results {
+            if let Some(first) = first_state {
+                if first != state {
+                    return Err(database::Error::Internal(
+                        "Proofs have inconsistent states".to_string(),
+                    ));
+                }
+            } else {
+                first_state = Some(state);
+            }
+
+            proofs.push(proof);
+        }
+
+        let state = first_state.unwrap_or(State::Unspent);
+        Ok(ProofsWithState::new(proofs, state).into())
     }
 }
 
@@ -425,7 +489,8 @@ where
         keyset_id: &Id,
     ) -> Result<(Proofs, Vec<Option<State>>), Self::Err> {
         let conn = self.pool.get().map_err(|e| Error::Database(Box::new(e)))?;
-        Ok(query(
+
+        let (proofs, states): (Vec<Proof>, Vec<State>) = query(
             r#"
             SELECT
                keyset_id,
@@ -447,7 +512,9 @@ where
         .map(sql_row_to_proof_with_state)
         .collect::<Result<Vec<_>, _>>()?
         .into_iter()
-        .unzip())
+        .unzip();
+
+        Ok((proofs, states.into_iter().map(Some).collect()))
     }
 
     /// Get total proofs redeemed by keyset id

+ 4 - 41
crates/cdk/src/mint/melt/melt_saga/mod.rs

@@ -7,7 +7,6 @@ use cdk_common::database::DynMintDatabase;
 use cdk_common::mint::{MeltSagaState, Operation, Saga, SagaStateEnum};
 use cdk_common::nut00::KnownMethod;
 use cdk_common::nuts::MeltQuoteState;
-use cdk_common::state::check_state_transition;
 use cdk_common::{Amount, Error, ProofsMethods, PublicKey, QuoteId, State};
 #[cfg(feature = "prometheus")]
 use cdk_prometheus::METRICS;
@@ -21,6 +20,7 @@ use crate::mint::melt::shared;
 use crate::mint::subscription::PubSubManager;
 use crate::mint::verification::Verification;
 use crate::mint::{MeltQuoteBolt11Response, MeltRequest};
+use crate::Mint;
 
 mod compensation;
 mod state;
@@ -246,48 +246,11 @@ impl MeltSaga<Initial> {
 
         let input_ys = melt_request.inputs().ys()?;
 
-        for current_state in tx
-            .get_proofs_states(&input_ys)
-            .await?
-            .into_iter()
-            .collect::<Option<Vec<_>>>()
-            .ok_or(Error::UnexpectedProofState)?
-        {
-            check_state_transition(current_state, State::Pending)
-                .map_err(|_| Error::UnexpectedProofState)?;
-        }
-
-        // Update proof states to Pending
-        let original_states = match tx.update_proofs_states(&input_ys, State::Pending).await {
-            Ok(states) => states,
-            Err(cdk_common::database::Error::AttemptUpdateSpentProof)
-            | Err(cdk_common::database::Error::AttemptRemoveSpentProof) => {
-                tx.rollback().await?;
-                return Err(Error::TokenAlreadySpent);
-            }
-            Err(err) => {
-                tx.rollback().await?;
-                return Err(err.into());
-            }
-        };
-
-        // Check for forbidden states (Pending or Spent)
-        let has_forbidden_state = original_states
-            .iter()
-            .any(|state| matches!(state, Some(State::Pending) | Some(State::Spent)));
+        let mut proofs = tx.get_proofs(&input_ys).await?;
 
-        if has_forbidden_state {
+        if let Err(err) = Mint::update_proofs_state(&mut tx, &mut proofs, State::Pending).await {
             tx.rollback().await?;
-            return Err(
-                if original_states
-                    .iter()
-                    .any(|s| matches!(s, Some(State::Pending)))
-                {
-                    Error::TokenPending
-                } else {
-                    Error::TokenAlreadySpent
-                },
-            );
+            return Err(err);
         }
 
         let previous_state = quote.state;

+ 3 - 22
crates/cdk/src/mint/melt/shared.rs

@@ -8,12 +8,12 @@
 
 use cdk_common::database::{self, Acquired, DynMintDatabase};
 use cdk_common::nuts::{BlindSignature, BlindedMessage, MeltQuoteState, State};
-use cdk_common::state::check_state_transition;
 use cdk_common::{Amount, Error, PublicKey, QuoteId};
 use cdk_signatory::signatory::SignatoryKeySet;
 
 use crate::mint::subscription::PubSubManager;
 use crate::mint::MeltQuote;
+use crate::Mint;
 
 /// Retrieves fee and amount configuration for the keyset matching the change outputs.
 ///
@@ -393,28 +393,9 @@ pub async fn finalize_melt_core(
             .await?;
     }
 
-    for current_state in tx
-        .get_proofs_states(input_ys)
-        .await?
-        .into_iter()
-        .collect::<Option<Vec<_>>>()
-        .ok_or(Error::UnexpectedProofState)?
-    {
-        check_state_transition(current_state, State::Spent)
-            .map_err(|_| Error::UnexpectedProofState)?;
-    }
+    let mut proofs = tx.get_proofs(input_ys).await?;
 
-    // Mark input proofs as spent
-    match tx.update_proofs_states(input_ys, State::Spent).await {
-        Ok(_) => {}
-        Err(database::Error::AttemptUpdateSpentProof) => {
-            tracing::info!("Proofs for quote {} already marked as spent", quote.id);
-            return Ok(());
-        }
-        Err(err) => {
-            return Err(err.into());
-        }
-    }
+    Mint::update_proofs_state(tx, &mut proofs, State::Spent).await?;
 
     // Publish proof state changes
     for pk in input_ys.iter() {

+ 1 - 0
crates/cdk/src/mint/mod.rs

@@ -39,6 +39,7 @@ mod issue;
 mod keysets;
 mod ln;
 mod melt;
+mod proofs;
 mod start_up_check;
 mod subscription;
 mod swap;

+ 267 - 0
crates/cdk/src/mint/proofs.rs

@@ -0,0 +1,267 @@
+use cdk_common::database::{Acquired, DynMintTransaction};
+use cdk_common::mint::ProofsWithState;
+use cdk_common::state::{self, check_state_transition};
+use cdk_common::{Error, State};
+
+use crate::Mint;
+
+impl Mint {
+    /// Updates the state of proofs with validation and error handling.
+    ///
+    /// This method:
+    /// 1. Validates the state transition is allowed via `check_state_transition`
+    /// 2. Persists the new state to the database
+    /// 3. Updates the `ProofsWithState.state` field on success
+    ///
+    /// # Errors
+    ///
+    /// - [`Error::UnexpectedProofState`] if the state transition is invalid
+    /// - [`Error::TokenAlreadySpent`] if the database rejects the update (proofs already spent)
+    pub async fn update_proofs_state(
+        tx: &mut DynMintTransaction,
+        proofs: &mut Acquired<ProofsWithState>,
+        new_state: State,
+    ) -> Result<(), Error> {
+        check_state_transition(proofs.state, new_state).map_err(|err| match err {
+            state::Error::AlreadySpent => Error::TokenAlreadySpent,
+            state::Error::Pending => Error::TokenPending,
+            _ => Error::UnexpectedProofState,
+        })?;
+
+        tx.update_proofs_state(proofs, new_state)
+            .await
+            .map_err(|err| match err {
+                cdk_common::database::Error::AttemptUpdateSpentProof
+                | cdk_common::database::Error::AttemptRemoveSpentProof => Error::TokenAlreadySpent,
+                err => err.into(),
+            })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use cdk_common::mint::Operation;
+    use cdk_common::nuts::ProofsMethods;
+    use cdk_common::{Amount, Error, State};
+
+    use crate::test_helpers::mint::{create_test_mint, mint_test_proofs};
+    use crate::Mint;
+
+    /// Test successful transition from Unspent to Pending
+    #[tokio::test]
+    async fn test_update_proofs_state_unspent_to_pending() {
+        let mint = create_test_mint().await.unwrap();
+        let proofs = mint_test_proofs(&mint, Amount::from(100)).await.unwrap();
+        let ys = proofs.ys().unwrap();
+
+        let db = mint.localstore();
+
+        // Add proofs to the database first
+        {
+            let mut tx = db.begin_transaction().await.unwrap();
+            tx.add_proofs(
+                proofs.clone(),
+                None,
+                &Operation::new_swap(Amount::ZERO, Amount::ZERO, Amount::ZERO),
+            )
+            .await
+            .unwrap();
+            tx.commit().await.unwrap();
+        }
+
+        let mut tx = db.begin_transaction().await.unwrap();
+        let mut acquired = tx.get_proofs(&ys).await.unwrap();
+
+        assert_eq!(acquired.state, State::Unspent);
+
+        Mint::update_proofs_state(&mut tx, &mut acquired, State::Pending)
+            .await
+            .unwrap();
+
+        assert_eq!(acquired.state, State::Pending);
+        tx.commit().await.unwrap();
+
+        // Verify state persisted to database
+        let states = db.get_proofs_states(&ys).await.unwrap();
+        assert!(states.iter().all(|s| *s == Some(State::Pending)));
+    }
+
+    /// Test successful transition from Pending to Spent
+    #[tokio::test]
+    async fn test_update_proofs_state_pending_to_spent() {
+        let mint = create_test_mint().await.unwrap();
+        let proofs = mint_test_proofs(&mint, Amount::from(100)).await.unwrap();
+        let ys = proofs.ys().unwrap();
+
+        let db = mint.localstore();
+
+        // Add proofs and transition to Pending
+        {
+            let mut tx = db.begin_transaction().await.unwrap();
+            tx.add_proofs(
+                proofs.clone(),
+                None,
+                &Operation::new_swap(Amount::ZERO, Amount::ZERO, Amount::ZERO),
+            )
+            .await
+            .unwrap();
+            tx.commit().await.unwrap();
+        }
+
+        {
+            let mut tx = db.begin_transaction().await.unwrap();
+            let mut acquired = tx.get_proofs(&ys).await.unwrap();
+            Mint::update_proofs_state(&mut tx, &mut acquired, State::Pending)
+                .await
+                .unwrap();
+            tx.commit().await.unwrap();
+        }
+
+        // Now test Pending -> Spent transition
+        let mut tx = db.begin_transaction().await.unwrap();
+        let mut acquired = tx.get_proofs(&ys).await.unwrap();
+
+        assert_eq!(acquired.state, State::Pending);
+
+        Mint::update_proofs_state(&mut tx, &mut acquired, State::Spent)
+            .await
+            .unwrap();
+
+        assert_eq!(acquired.state, State::Spent);
+        tx.commit().await.unwrap();
+
+        // Verify state persisted to database
+        let states = db.get_proofs_states(&ys).await.unwrap();
+        assert!(states.iter().all(|s| *s == Some(State::Spent)));
+    }
+
+    /// Test that update_proofs_state rejects same-state transition (Pending -> Pending)
+    #[tokio::test]
+    async fn test_update_proofs_state_rejects_same_state_transition() {
+        let mint = create_test_mint().await.unwrap();
+        let proofs = mint_test_proofs(&mint, Amount::from(100)).await.unwrap();
+        let ys = proofs.ys().unwrap();
+
+        let db = mint.localstore();
+
+        // Add proofs and transition to Pending
+        {
+            let mut tx = db.begin_transaction().await.unwrap();
+            tx.add_proofs(
+                proofs.clone(),
+                None,
+                &Operation::new_swap(Amount::ZERO, Amount::ZERO, Amount::ZERO),
+            )
+            .await
+            .unwrap();
+            tx.commit().await.unwrap();
+        }
+
+        {
+            let mut tx = db.begin_transaction().await.unwrap();
+            let mut acquired = tx.get_proofs(&ys).await.unwrap();
+            Mint::update_proofs_state(&mut tx, &mut acquired, State::Pending)
+                .await
+                .unwrap();
+            tx.commit().await.unwrap();
+        }
+
+        // Try invalid transition: Pending -> Pending (same state)
+        let mut tx = db.begin_transaction().await.unwrap();
+        let mut acquired = tx.get_proofs(&ys).await.unwrap();
+
+        assert_eq!(acquired.state, State::Pending);
+
+        let result = Mint::update_proofs_state(&mut tx, &mut acquired, State::Pending).await;
+
+        assert!(matches!(result, Err(Error::TokenPending)));
+    }
+
+    /// Test that update_proofs_state rejects invalid transition from Spent
+    #[tokio::test]
+    async fn test_update_proofs_state_invalid_transition_from_spent() {
+        let mint = create_test_mint().await.unwrap();
+        let proofs = mint_test_proofs(&mint, Amount::from(100)).await.unwrap();
+        let ys = proofs.ys().unwrap();
+
+        let db = mint.localstore();
+
+        // Add proofs and transition to Spent
+        {
+            let mut tx = db.begin_transaction().await.unwrap();
+            tx.add_proofs(
+                proofs.clone(),
+                None,
+                &Operation::new_swap(Amount::ZERO, Amount::ZERO, Amount::ZERO),
+            )
+            .await
+            .unwrap();
+            tx.commit().await.unwrap();
+        }
+
+        {
+            let mut tx = db.begin_transaction().await.unwrap();
+            let mut acquired = tx.get_proofs(&ys).await.unwrap();
+            Mint::update_proofs_state(&mut tx, &mut acquired, State::Pending)
+                .await
+                .unwrap();
+            Mint::update_proofs_state(&mut tx, &mut acquired, State::Spent)
+                .await
+                .unwrap();
+            tx.commit().await.unwrap();
+        }
+
+        // Try invalid transition: Spent -> Pending (not allowed)
+        let mut tx = db.begin_transaction().await.unwrap();
+        let mut acquired = tx.get_proofs(&ys).await.unwrap();
+
+        assert_eq!(acquired.state, State::Spent);
+
+        let result = Mint::update_proofs_state(&mut tx, &mut acquired, State::Pending).await;
+
+        assert!(matches!(result, Err(Error::TokenAlreadySpent)));
+    }
+
+    /// Test that ProofsWithState.state is updated after successful update
+    #[tokio::test]
+    async fn test_update_proofs_state_updates_wrapper_state() {
+        let mint = create_test_mint().await.unwrap();
+        let proofs = mint_test_proofs(&mint, Amount::from(100)).await.unwrap();
+        let ys = proofs.ys().unwrap();
+
+        let db = mint.localstore();
+
+        // Add proofs to the database first
+        {
+            let mut tx = db.begin_transaction().await.unwrap();
+            tx.add_proofs(
+                proofs.clone(),
+                None,
+                &Operation::new_swap(Amount::ZERO, Amount::ZERO, Amount::ZERO),
+            )
+            .await
+            .unwrap();
+            tx.commit().await.unwrap();
+        }
+
+        let mut tx = db.begin_transaction().await.unwrap();
+        let mut acquired = tx.get_proofs(&ys).await.unwrap();
+
+        // Before update
+        assert_eq!(acquired.state, State::Unspent);
+
+        // After update
+        Mint::update_proofs_state(&mut tx, &mut acquired, State::Pending)
+            .await
+            .unwrap();
+
+        // The wrapper's state field should be updated
+        assert_eq!(
+            acquired.state,
+            State::Pending,
+            "ProofsWithState.state should be updated after successful update_proofs_state"
+        );
+
+        tx.commit().await.unwrap();
+    }
+}

+ 34 - 66
crates/cdk/src/mint/swap/swap_saga/mod.rs

@@ -4,7 +4,6 @@ use std::sync::Arc;
 use cdk_common::database::DynMintDatabase;
 use cdk_common::mint::{Operation, Saga, SwapSagaState};
 use cdk_common::nuts::BlindedMessage;
-use cdk_common::state::check_state_transition;
 use cdk_common::{database, Amount, Error, Proofs, ProofsMethods, PublicKey, QuoteId, State};
 use tokio::sync::Mutex;
 use tracing::instrument;
@@ -12,6 +11,7 @@ use tracing::instrument;
 use self::compensation::{CompensatingAction, RemoveSwapSetup};
 use self::state::{Initial, SetupComplete, Signed};
 use crate::mint::subscription::PubSubManager;
+use crate::Mint;
 
 pub mod compensation;
 mod state;
@@ -177,54 +177,30 @@ impl<'a> SwapSaga<'a, Initial> {
         );
 
         // Add input proofs to DB
-        if let Err(err) = tx
+        let mut new_proofs = match tx
             .add_proofs(input_proofs.clone(), quote_id.clone(), &operation)
             .await
         {
-            tx.rollback().await?;
-            return Err(match err {
-                database::Error::Duplicate => Error::TokenPending,
-                database::Error::AttemptUpdateSpentProof => Error::TokenAlreadySpent,
-                _ => Error::Database(err),
-            });
-        }
+            Ok(proofs) => proofs,
+            Err(err) => {
+                tx.rollback().await?;
+                return Err(match err {
+                    database::Error::Duplicate => Error::TokenPending,
+                    database::Error::AttemptUpdateSpentProof => Error::TokenAlreadySpent,
+                    _ => Error::Database(err),
+                });
+            }
+        };
 
         let ys = match input_proofs.ys() {
             Ok(ys) => ys,
             Err(err) => return Err(Error::NUT00(err)),
         };
 
-        // Update input proof states to Pending
-        let original_proof_states = match tx.update_proofs_states(&ys, State::Pending).await {
-            Ok(states) => states,
-            Err(database::Error::AttemptUpdateSpentProof)
-            | Err(database::Error::AttemptRemoveSpentProof) => {
-                tx.rollback().await?;
-                return Err(Error::TokenAlreadySpent);
-            }
-            Err(err) => {
-                tx.rollback().await?;
-                return Err(err.into());
-            }
-        };
-
-        // Verify proofs weren't already pending or spent
-        if ys.len() != original_proof_states.len() {
-            tracing::error!("Mismatched proof states");
+        if let Err(err) = Mint::update_proofs_state(&mut tx, &mut new_proofs, State::Pending).await
+        {
             tx.rollback().await?;
-            return Err(Error::Internal);
-        }
-
-        let forbidden_states = [State::Pending, State::Spent];
-        for original_state in original_proof_states.iter().flatten() {
-            if forbidden_states.contains(original_state) {
-                tx.rollback().await?;
-                return Err(if *original_state == State::Pending {
-                    Error::TokenPending
-                } else {
-                    Error::TokenAlreadySpent
-                });
-            }
+            return Err(err);
         }
 
         // Add output blinded messages
@@ -423,33 +399,19 @@ impl SwapSaga<'_, Signed> {
             }
         }
 
-        for current_state in tx
-            .get_proofs_states(&self.state_data.ys)
-            .await?
-            .into_iter()
-            .collect::<Option<Vec<_>>>()
-            .ok_or(Error::UnexpectedProofState)?
-        {
-            check_state_transition(current_state, State::Spent)
-                .map_err(|_| Error::UnexpectedProofState)?;
-        }
-
-        match tx
-            .update_proofs_states(&self.state_data.ys, State::Spent)
-            .await
-        {
-            Ok(_) => {}
-            Err(database::Error::AttemptUpdateSpentProof)
-            | Err(database::Error::AttemptRemoveSpentProof) => {
-                tx.rollback().await?;
-                self.compensate_all().await?;
-                return Err(Error::TokenAlreadySpent);
-            }
+        let mut proofs = match tx.get_proofs(&self.state_data.ys).await {
+            Ok(proofs) => proofs,
             Err(err) => {
                 tx.rollback().await?;
                 self.compensate_all().await?;
                 return Err(err.into());
             }
+        };
+
+        if let Err(err) = Mint::update_proofs_state(&mut tx, &mut proofs, State::Spent).await {
+            tx.rollback().await?;
+            self.compensate_all().await?;
+            return Err(err);
         }
 
         // Publish proof state changes
@@ -457,11 +419,17 @@ impl SwapSaga<'_, Signed> {
             self.pubsub.proof_state((*pk, State::Spent));
         }
 
-        tx.add_completed_operation(
-            &self.state_data.operation,
-            &self.state_data.fee_breakdown.per_keyset,
-        )
-        .await?;
+        if let Err(err) = tx
+            .add_completed_operation(
+                &self.state_data.operation,
+                &self.state_data.fee_breakdown.per_keyset,
+            )
+            .await
+        {
+            tx.rollback().await?;
+            self.compensate_all().await?;
+            return Err(err.into());
+        }
 
         // Delete saga - swap completed successfully (best-effort, atomic with TX2)
         // Don't fail the swap if saga deletion fails - orphaned saga will be