Track Stake inputs for post-stake invalidation

PIVX-Labs · Sep 4, 2023 · d8e61f3 · d8e61f3
1 parent 91af6e0
commit d8e61f3
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 12 deletions.
diff --git a/scripts/mempool.js b/scripts/mempool.js
@@ -4,6 +4,12 @@ import { sleep } from './misc.js';
 import { debug } from './settings.js';
 import { getEventEmitter } from './event_bus.js';
 
+/**
+ * @typedef {Object} cIn - An input of a Tx
+ * @property {string} txid - The transaction ID
+ * @property {number} vout - The output index
+ */
+
 /** An Unspent Transaction Output, used as Inputs of future transactions */
 export class UTXO {
     /**
@@ -12,16 +18,19 @@ export class UTXO {
      * @param {String} UTXO.path - If applicable, the HD Path of the owning address
      * @param {Number} UTXO.sats - Satoshi value in this UTXO
      * @param {String} UTXO.script - HEX encoded spending script
+     * @param {Array<cIn>?} UTXO.vin - The inputs of the transaction, if any
      * @param {Number} UTXO.vout - Output position of this transaction
      * @param {Number} UTXO.height - Block height of the UTXO
      * @param {Number} UTXO.status - UTXO status enum state
-     * @param {bool} UTXO.isDelegate - Whether the UTXO is a cold stake delegation
+     * @param {boolean} UTXO.isDelegate - Whether the UTXO is a cold stake delegation
+     * @param {boolean} UTXO.isReward - Whether the UTXO is a reward
      */
     constructor({
         id,
         path,
         sats,
         script,
+        vin = [],
         vout,
         height,
         status,
@@ -44,6 +53,10 @@ export class UTXO {
          *  @type {String} */
         this.script = script;
 
+        /** The inputs of the transaction, if any
+         *  @type {Array<cIn>} */
+        this.vin = vin;
+
         /** Output position of this transaction
          *  @type {Number} */
         this.vout = vout;
@@ -56,10 +69,12 @@ export class UTXO {
          *  @type {Number} */
         this.status = status;
 
-        /** If it's a delegation UTXO
-         * @type {bool} */
+        /** Whether it's a delegation UTXO
+         * @type {boolean} */
         this.isDelegate = isDelegate;
 
+        /** Whether it's a reward UTXO
+         * @type {boolean} */
         this.isReward = isReward;
     }
 
@@ -97,6 +112,17 @@ export class Mempool {
     /** The PENDING state (standard UTXO is in mempool, pending confirmation) */
     static PENDING = 2;
 
+    /**
+     * Fetch a UTXO by ID and Index
+     * @param {string} id - Transaction ID of the UTXO
+     * @param {number} out - Output position of the UTXO
+     */
+    getUTXO(id, out) {
+        return this.UTXOs.find(
+            (cUTXO) => cUTXO.id === id && cUTXO.vout === out
+        );
+    }
+
     /**
      * Remove a UTXO after a set amount of time
      * @param {Number} nBlocks - Estimated blocks to wait
@@ -153,21 +179,14 @@ export class Mempool {
 
     /**
      * Add a new UTXO to the wallet
-     * @param {Object} UTXO
-     * @param {String} UTXO.id - Transaction ID
-     * @param {String} UTXO.path - If applicable, the HD Path of the owning address
-     * @param {Number} UTXO.sats - Satoshi value in this UTXO
-     * @param {String} UTXO.script - HEX encoded spending script
-     * @param {Number} UTXO.vout - Output position of this transaction
-     * @param {Number} UTXO.height - Block height of the UTXO
-     * @param {Number} UTXO.status - UTXO status enum state
-     * @param {Boolean} UTXO.isDelegate - If this is a Cold Delegation
+     * @param {UTXO} UTXO
      */
     addUTXO({
         id,
         path,
         sats,
         script,
+        vin,
         vout,
         height,
         status,
@@ -179,6 +198,7 @@ export class Mempool {
             path,
             sats,
             script,
+            vin,
             vout,
             height,
             status,
@@ -189,8 +209,19 @@ export class Mempool {
         if (this.isAlreadyStored({ id, vout })) {
             this.updateUTXO({ id, vout });
         } else {
+            // If this new UTXO is a reward with one vin (i.e: a Stake), we'll backtrace the input
+            if (isReward && vin?.length === 1) {
+                const cStakeInput = this.getUTXO(vin[0].txid, vin[0].vout);
+
+                // And if the input is in our wallet, remove it
+                if (cStakeInput) {
+                    this.removeUTXO(cStakeInput);
+                }
+            }
             this.UTXOs.push(newUTXO);
         }
+
+        // Re-render the Balance UIs
         getBalance(true);
         getStakingBalance(true);
     }
@@ -337,10 +368,12 @@ export class Mempool {
     subscribeToNetwork() {
         getEventEmitter().on('utxo', async (utxos) => {
             for (const utxo of utxos) {
+                // If we have the UTXO, we update it's confirmation status
                 if (this.isAlreadyStored({ id: utxo.txid, vout: utxo.vout })) {
                     this.updateUTXO({ id: utxo.txid, vout: utxo.vout });
                     continue;
                 }
+                // If the UTXO is new, we'll process it and add it internally
                 this.addUTXO(await getNetwork().getUTXOFullInfo(utxo));
             }
         });

diff --git a/scripts/network.js b/scripts/network.js
@@ -322,6 +322,7 @@ export class ExplorerNetwork extends Network {
             path,
             sats: Math.round(cVout.value * COIN),
             script: cVout.scriptPubKey.hex,
+            vin: cTx?.vin || [],
             vout: cVout.n,
             height: this.cachedBlockCount - (cTx.confirmations - 1),
             status: cTx.confirmations < 1 ? Mempool.PENDING : Mempool.CONFIRMED,