Refactor commitment

src/commit.rs

@@ -1,203 +1,206 @@
+use std::str::FromStr;
+
 use anyhow::{Error, Result};
 
 use hex::FromHex;
 use sdk_common::pcd::Pcd;
 use sdk_common::silentpayments::create_transaction;
 use sdk_common::sp_client::spclient::Recipient;
-use sdk_common::{error::AnkError, network::CommitMessage};
+use sdk_common::network::CommitMessage;
 use sdk_common::sp_client::bitcoin::consensus::deserialize;
-use sdk_common::sp_client::bitcoin::{Amount, Transaction, Txid, OutPoint};
+use sdk_common::sp_client::bitcoin::{Amount, Transaction, OutPoint};
 use sdk_common::process::{Process, ProcessState, CACHEDPROCESSES};
-use serde_json::{json, Map, Value};
+use serde_json::json;
 
 use crate::{lock_freezed_utxos, MutexExt, DAEMON, WALLET};
 
 pub(crate) fn handle_commit_request(commit_msg: CommitMessage) -> Result<OutPoint> {
-    // Attempt to deserialize `init_tx` as a `Transaction`
+    // Attempt to process the initial transaction or reference
-    if let Ok(tx) = deserialize::<Transaction>(&Vec::from_hex(&commit_msg.init_tx)?) {
+    match parse_initial_tx(&commit_msg.init_tx)? {
-        // This is the first transaction of a chain of commitments
+        InitialTx::Transaction(tx) => handle_initial_transaction(tx, &commit_msg),
-        // Create the root commitment outpoint
+        InitialTx::OutPoint(outpoint) => handle_existing_commitment(outpoint, commit_msg),
-        let root_commitment = OutPoint::new(tx.txid(), 0);
-
-        // TODO: Check that the output pays us
-
-        // Validation tokens must be empty for the initial transaction
-        if !commit_msg.validation_tokens.is_empty() {
-            return Err(AnkError::GenericError(
-                "Validation tokens must be empty".to_string(),
-            ))?;
-        }
-
-        // Obtain the daemon instance and broadcast the transaction
-        let daemon = DAEMON.get().unwrap().lock_anyhow()?;
-        daemon.broadcast(&tx)?;
-
-        let roles = Value::Object(commit_msg.roles.iter().map(|(name, def)| (name.to_owned(), Value::String(serde_json::to_string(def).unwrap()))).collect());
-
-        // put roles in a map
-        let roles_only_map = json!({
-            "roles": roles
-        });
-
-        // We check that for testing but it's useless
-        assert!(commit_msg.roles == roles.extract_roles()?);
-
-        let roles_commitment = roles_only_map.hash_fields(root_commitment)?;
-
-        assert!(roles_commitment.get("roles") == commit_msg.pcd_commitment.get("roles"));
-
-        // We always keep an empty state as the last state
-        let empty_state = ProcessState {
-            commited_in: root_commitment,
-            ..Default::default()
-        };
-
-        // Initialize the process state
-        let mut init_state = empty_state.clone();
-        init_state.encrypted_pcd = roles_only_map;
-        init_state.pcd_commitment = commit_msg.pcd_commitment;
-
-        // Access the cached processes and insert the new commitment
-        let mut commitments = CACHEDPROCESSES
-            .get()
-            .ok_or(Error::msg("CACHEDPROCESSES not initialized"))?
-            .lock_anyhow()?;
-        // We are confident that `root_commitment` doesn't exist in the map
-        commitments.insert(
-            root_commitment,
-            Process::new(vec![init_state, empty_state], vec![]),
-        );
-
-        // Add the outpoint to the list of frozen UTXOs
-        lock_freezed_utxos()?.insert(root_commitment);
-
-        // Wait for validation tokens to spend the new output and commit the hash
-        Ok(root_commitment)
     }
-    // Attempt to deserialize `init_tx` as an `OutPoint`
+}
-    else if let Ok(outpoint) = deserialize::<OutPoint>(&Vec::from_hex(&commit_msg.init_tx)?) {
-        // Reference the first transaction of a chain
-        let mut commitments = CACHEDPROCESSES
-            .get()
-            .ok_or(Error::msg("CACHEDPROCESSES not initialized"))?
-            .lock_anyhow()?;
-        let commitment = commitments
-            .get_mut(&outpoint)
-            .ok_or(Error::msg("Commitment not found"))?;
 
-        if commit_msg.validation_tokens.is_empty() {
+// Enum to differentiate between parsed transaction and outpoint
-            // Register a new state if validation tokens are empty
+enum InitialTx {
+    Transaction(Transaction),
+    OutPoint(OutPoint),
+}
 
-            // Get all the latest concurrent states
+// Parse the initial transaction as either a `Transaction` or `OutPoint`
-            let concurrent_states = commitment.get_latest_concurrent_states()?;
+fn parse_initial_tx(init_tx: &str) -> Result<InitialTx> {
-
+    if let Ok(tx_hex) = Vec::from_hex(init_tx) {
-            let (empty_state, actual_states) = concurrent_states.split_last().unwrap(); // We necessary have 1 state
+        let tx = deserialize::<Transaction>(&tx_hex)?;
-            let current_outpoint = empty_state.commited_in;
+        Ok(InitialTx::Transaction(tx))
-
+    } else if let Ok(outpoint) = OutPoint::from_str(init_tx) {
-            // Check for existing states with the same PCD hash
+        Ok(InitialTx::OutPoint(outpoint))
-            if actual_states
-                .into_iter()
-                .any(|state| state.pcd_commitment == commit_msg.pcd_commitment)
-            {
-                return Err(anyhow::Error::msg("Proposed state already exists"));
-            }
-
-            let roles = Value::Object(commit_msg.roles.iter().map(|(name, def)| (name.to_owned(), Value::String(serde_json::to_string(def).unwrap()))).collect());
-
-            let roles_only_map = json!({
-                "roles": roles
-            });
-
-            let new_state = ProcessState {
-                commited_in: current_outpoint,
-                pcd_commitment: commit_msg.pcd_commitment,
-                encrypted_pcd: roles_only_map,
-                ..Default::default()
-            };
-
-            // Insert the new process state
-            commitment.insert_concurrent_state(new_state)?;
-
-            Ok(current_outpoint)
-        } else {
-            // Validation tokens are provided; process the pending state
-
-            // Clone the state, we'll need it for validation purpose
-            let mut state_to_validate = commitment.get_latest_concurrent_states()?
-                .into_iter()
-                .find(|state| {
-                    state.pcd_commitment == commit_msg.pcd_commitment
-                })
-                .ok_or(anyhow::Error::msg("Unknown state"))?
-                .clone();
-
-            // We update the validation tokens for our clone
-            state_to_validate.validation_tokens = commit_msg.validation_tokens;
-            // We test the validity of the state with the provided proofs
-            state_to_validate.is_valid(commitment.get_latest_commited_state())?;
-
-            // If the new state is valid we commit it in a new transaction that spends the last commited_in
-            // By spending it we also know the next outpoint to monitor for the next state
-            // We add a placeholder state with that information at the tip of the chain
-            // We also remove all concurrent states that didn't get validated
-            let mut freezed_utxos = lock_freezed_utxos()?;
-            let mandatory_input = if freezed_utxos.remove(&state_to_validate.commited_in) {
-                state_to_validate.commited_in
-            } else {
-                // This shoudln't happen, except if we send the commitment message to another relay
-                // We need to think about the case a relay is down
-                // Probably relays should have backups of their keys so that it can be bootstrapped again and resume commiting
-                return Err(Error::msg("Commitment utxo doesn't exist"))
-            };
-
-            // TODO we should make this sequence more atomic by handling errors in a way that we get back to the current state if any step fails
-
-            let sp_wallet = WALLET.get().ok_or(Error::msg("Wallet not initialized"))?.get_wallet()?;
-
-            let recipient = Recipient {
-                address: sp_wallet.get_client().get_receiving_address(),
-                amount: Amount::from_sat(1000),
-                nb_outputs: 1
-            };
-
-            let daemon = DAEMON.get().unwrap().lock_anyhow()?;
-            let fee_rate = daemon.estimate_fee(6)?;
-
-            let psbt = create_transaction(
-                vec![mandatory_input], 
-                &freezed_utxos, 
-                &sp_wallet, 
-                vec![recipient], 
-                None, 
-                fee_rate, 
-                None
-            )?;
-
-            let new_tx = psbt.extract_tx()?;
-
-            daemon.test_mempool_accept(&new_tx)?;
-
-            // We're ready to commit, we first update our process states 
-            // We remove all concurrent states
-            let _ = commitment.remove_all_concurrent_states()?;
-            // debug!("removed states: {:?}", rm_states);
-            // We push the validated state back
-            commitment.insert_concurrent_state(state_to_validate.clone())?;
-
-            // We broadcast transaction
-            let txid = daemon.broadcast(&new_tx)?;
-
-            // We push a new, empty state commited in the newly created output
-            let commited_in = OutPoint::new(txid, 0);
-
-            // Add the newly created outpoint to our list of freezed utxos
-            freezed_utxos.insert(commited_in);
-            
-            commitment.update_states_tip(commited_in)?;
-
-            Ok(commited_in)
-        }
     } else {
         Err(Error::msg("init_tx must be a valid transaction or txid"))
     }
 }
+
+// Handle the case where `init_tx` is a new transaction
+fn handle_initial_transaction(tx: Transaction, commit_msg: &CommitMessage) -> Result<OutPoint> {
+    let root_commitment = OutPoint::new(tx.txid(), 0);
+
+    // Validation tokens must be empty for the initial transaction
+    if !commit_msg.validation_tokens.is_empty() {
+        return Err(anyhow::Error::msg(
+            "Validation tokens must be empty".to_string(),
+        ));
+    }
+
+    // Broadcast the transaction
+    let daemon = DAEMON.get().unwrap().lock_anyhow()?;
+    daemon.broadcast(&tx)?;
+
+    // Process roles and commitments
+    let roles_only_map = json!({ "roles": serde_json::to_value(&commit_msg.roles)? });
+    let roles_commitment = roles_only_map.hash_fields(root_commitment)?;
+
+    if roles_commitment.get("roles") != commit_msg.pcd_commitment.get("roles") {
+        return Err(Error::msg("Role commitment mismatch"));
+    }
+
+    // Create the initial process state
+    let empty_state = ProcessState {
+        commited_in: root_commitment,
+        ..Default::default()
+    };
+    let mut init_state = empty_state.clone();
+    init_state.encrypted_pcd = roles_only_map;
+    init_state.pcd_commitment = commit_msg.pcd_commitment.clone();
+
+    // Cache the process
+    let mut commitments = CACHEDPROCESSES
+        .get()
+        .ok_or(Error::msg("CACHEDPROCESSES not initialized"))?
+        .lock_anyhow()?;
+    commitments.insert(
+        root_commitment,
+        Process::new(vec![init_state, empty_state], vec![]),
+    );
+
+    // Add to frozen UTXOs
+    lock_freezed_utxos()?.insert(root_commitment);
+
+    Ok(root_commitment)
+}
+
+// Handle the case where `init_tx` is a reference to an existing outpoint
+fn handle_existing_commitment(outpoint: OutPoint, commit_msg: CommitMessage) -> Result<OutPoint> {
+    let mut commitments = CACHEDPROCESSES
+        .get()
+        .ok_or(Error::msg("CACHEDPROCESSES not initialized"))?
+        .lock_anyhow()?;
+    let commitment = commitments
+        .get_mut(&outpoint)
+        .ok_or(Error::msg("Commitment not found"))?;
+
+    if commit_msg.validation_tokens.is_empty() {
+        register_new_state(commitment, commit_msg)
+    } else {
+        process_validation(commitment, commit_msg)
+    }
+}
+
+// Register a new state when validation tokens are empty
+fn register_new_state(commitment: &mut Process, commit_msg: CommitMessage) -> Result<OutPoint> {
+    let concurrent_states = commitment.get_latest_concurrent_states()?;
+    let (empty_state, actual_states) = concurrent_states.split_last().unwrap();
+    let current_outpoint = empty_state.commited_in;
+
+    // Ensure no duplicate states
+    if actual_states
+        .iter()
+        .any(|state| state.pcd_commitment == commit_msg.pcd_commitment)
+    {
+        return Err(Error::msg("Proposed state already exists"));
+    }
+
+    // Add the new state
+    let roles_only_map = json!({ "roles": serde_json::to_value(&commit_msg.roles)? });
+    let new_state = ProcessState {
+        commited_in: current_outpoint,
+        pcd_commitment: commit_msg.pcd_commitment,
+        encrypted_pcd: roles_only_map,
+        ..Default::default()
+    };
+    commitment.insert_concurrent_state(new_state)?;
+
+    Ok(current_outpoint)
+}
+
+// Process validation for a state with validation tokens
+fn process_validation(commitment: &mut Process, commit_msg: CommitMessage) -> Result<OutPoint> {
+    let mut state_to_validate = commitment
+        .get_latest_concurrent_states()?
+        .into_iter()
+        .find(|state| state.pcd_commitment == commit_msg.pcd_commitment)
+        .ok_or(Error::msg("Unknown state"))?
+        .clone();
+
+    state_to_validate.validation_tokens = commit_msg.validation_tokens;
+    state_to_validate.is_valid(commitment.get_latest_commited_state())?;
+
+    let mandatory_input = prepare_next_transaction(&state_to_validate)?;
+    let commited_in = commit_new_transaction(commitment, mandatory_input)?;
+
+    Ok(commited_in)
+}
+
+// Prepare the next transaction based on the validated state
+fn prepare_next_transaction(state_to_validate: &ProcessState) -> Result<OutPoint> {
+    let mut freezed_utxos = lock_freezed_utxos()?;
+    if freezed_utxos.remove(&state_to_validate.commited_in) {
+        Ok(state_to_validate.commited_in)
+    } else {
+        Err(Error::msg("Commitment UTXO doesn't exist"))
+    }
+}
+
+// Commit the new transaction and update the process state
+fn commit_new_transaction(
+    commitment: &mut Process,
+    mandatory_input: OutPoint,
+) -> Result<OutPoint> {
+    let sp_wallet = WALLET
+        .get()
+        .ok_or(Error::msg("Wallet not initialized"))?
+        .get_wallet()?;
+
+    let recipient = Recipient {
+        address: sp_wallet.get_client().get_receiving_address(),
+        amount: Amount::from_sat(1000),
+        nb_outputs: 1,
+    };
+
+    let daemon = DAEMON.get().unwrap().lock_anyhow()?;
+    let fee_rate = daemon.estimate_fee(6)?;
+
+    let freezed_utxos = lock_freezed_utxos()?;
+
+    let psbt = create_transaction(
+        vec![mandatory_input],
+        &freezed_utxos,
+        &sp_wallet,
+        vec![recipient],
+        None,
+        fee_rate,
+        None,
+    )?;
+
+    let new_tx = psbt.extract_tx()?;
+    daemon.test_mempool_accept(&new_tx)?;
+
+    let txid = daemon.broadcast(&new_tx)?;
+    let commited_in = OutPoint::new(txid, 0);
+
+    lock_freezed_utxos()?.insert(commited_in);
+    commitment.update_states_tip(commited_in)?;
+
+    Ok(commited_in)
+}
+
+}