Fix for Wallet confirmed states (#1782)

* Fix for Wallet confirmed states

* Only process if we do not have the header

app/server/src/main/scala/org/bitcoins/server/Main.scala

@@ -182,7 +182,7 @@ object Main extends App with BitcoinSLogger {
       if (headers.isEmpty) {
         FutureUtil.unit
       } else {
-        wallet.updateUtxoPendingStates(headers.last).map(_ => ())
+        wallet.updateUtxoPendingStates().map(_ => ())
       }
     }
     if (nodeConf.isSPVEnabled) {

node/src/main/scala/org/bitcoins/node/networking/peer/DataMessageHandler.scala

@@ -230,11 +230,33 @@ case class DataMessageHandler(
           this.copy(chainApi = newApi, syncing = newSyncing)
         }
       case msg: BlockMessage =>
+        val block = msg.block
         logger.info(
-          s"Received block message with hash ${msg.block.blockHeader.hash.flip}")
-        callbacks
-          .executeOnBlockReceivedCallbacks(logger, msg.block)
-          .map(_ => this)
+          s"Received block message with hash ${block.blockHeader.hash.flip}")
+
+        val newApiF = {
+          chainApi
+            .getHeader(block.blockHeader.hashBE)
+            .flatMap { headerOpt =>
+              if (headerOpt.isEmpty) {
+                for {
+                  processedApi <- chainApi.processHeader(block.blockHeader)
+                  _ <- callbacks.executeOnBlockHeadersReceivedCallbacks(
+                    logger,
+                    Vector(block.blockHeader))
+                } yield processedApi
+              } else Future.successful(chainApi)
+            }
+        }
+
+        for {
+          newApi <- newApiF
+          _ <-
+            callbacks
+              .executeOnBlockReceivedCallbacks(logger, block)
+        } yield {
+          this.copy(chainApi = newApi)
+        }
       case TransactionMessage(tx) =>
         MerkleBuffers.putTx(tx, callbacks).flatMap { belongsToMerkle =>
           if (belongsToMerkle) {

wallet-test/src/test/scala/org/bitcoins/wallet/UTXOLifeCycleTest.scala

@@ -6,6 +6,7 @@ import org.bitcoins.core.protocol.script.{EmptyScriptPubKey, P2PKHScriptPubKey}
 import org.bitcoins.core.protocol.transaction.TransactionOutput
 import org.bitcoins.core.wallet.fee.{SatoshisPerByte, SatoshisPerVirtualByte}
 import org.bitcoins.core.wallet.utxo.TxoState
+import org.bitcoins.core.wallet.utxo.TxoState._
 import org.bitcoins.crypto.ECPublicKey
 import org.bitcoins.testkit.wallet.{
   BitcoinSWalletTest,
@@ -46,6 +47,36 @@ class UTXOLifeCycleTest extends BitcoinSWalletTest {
     }
   }
 
+  it should "track a utxo state change to confirmed spent" in { param =>
+    val WalletWithBitcoindRpc(wallet, bitcoind) = param
+
+    for {
+      oldTransactions <- wallet.listTransactions()
+      tx <- wallet.sendToAddress(testAddr,
+                                 Satoshis(3000),
+                                 Some(SatoshisPerByte(Satoshis(3))))
+
+      updatedCoins <- wallet.spendingInfoDAO.findOutputsBeingSpent(tx)
+      newTransactions <- wallet.listTransactions()
+      _ = assert(updatedCoins.forall(_.state == PendingConfirmationsSpent))
+      _ = assert(!oldTransactions.map(_.transaction).contains(tx))
+      _ = assert(newTransactions.map(_.transaction).contains(tx))
+
+      // Give tx a fake hash so it can appear as it's in a block
+      hash <- bitcoind.getBestBlockHash
+      _ <- wallet.spendingInfoDAO.upsertAll(
+        updatedCoins.map(_.copyWithBlockHash(hash)).toVector)
+
+      // Put confirmations on top of the tx's block
+      _ <- bitcoind.getNewAddress.flatMap(
+        bitcoind.generateToAddress(wallet.walletConfig.requiredConfirmations,
+                                   _))
+      // Need to call this to actually update the state, normally a node callback would do this
+      _ <- wallet.updateUtxoPendingStates()
+      confirmedCoins <- wallet.spendingInfoDAO.findOutputsBeingSpent(tx)
+    } yield assert(confirmedCoins.forall(_.state == ConfirmedSpent))
+  }
+
   it should "track a utxo state change to pending recieved" in { param =>
     val WalletWithBitcoindRpc(wallet, bitcoind) = param
 
@@ -73,6 +104,44 @@ class UTXOLifeCycleTest extends BitcoinSWalletTest {
     }
   }
 
+  it should "track a utxo state change to confirmed received" in { param =>
+    val WalletWithBitcoindRpc(wallet, bitcoind) = param
+
+    for {
+      oldTransactions <- wallet.listTransactions()
+      addr <- wallet.getNewAddress()
+
+      blockHash <- bitcoind.getBestBlockHash
+
+      txId <- bitcoind.sendToAddress(addr, Satoshis(3000))
+      tx <- bitcoind.getRawTransactionRaw(txId)
+      _ <- wallet.processOurTransaction(
+        transaction = tx,
+        feeRate = SatoshisPerByte(Satoshis(3)),
+        inputAmount = Satoshis(4000),
+        sentAmount = Satoshis(3000),
+        blockHashOpt = Some(blockHash), // give fake hash
+        newTags = Vector.empty
+      )
+
+      updatedCoin <-
+        wallet.spendingInfoDAO.findByScriptPubKey(addr.scriptPubKey)
+      newTransactions <- wallet.listTransactions()
+      _ = assert(updatedCoin.forall(_.state == PendingConfirmationsReceived))
+      _ = assert(!oldTransactions.map(_.transaction).contains(tx))
+      _ = assert(newTransactions.map(_.transaction).contains(tx))
+
+      // Put confirmations on top of the tx's block
+      _ <- bitcoind.getNewAddress.flatMap(
+        bitcoind.generateToAddress(wallet.walletConfig.requiredConfirmations,
+                                   _))
+      // Need to call this to actually update the state, normally a node callback would do this
+      _ <- wallet.updateUtxoPendingStates()
+      confirmedCoins <-
+        wallet.spendingInfoDAO.findByScriptPubKey(addr.scriptPubKey)
+    } yield assert(confirmedCoins.forall(_.state == ConfirmedReceived))
+  }
+
   it should "track a utxo state change to reserved" in { param =>
     val WalletWithBitcoindRpc(wallet, _) = param
 

wallet/src/main/scala/org/bitcoins/wallet/api/WalletApi.scala

@@ -8,7 +8,6 @@ import org.bitcoins.core.config.NetworkParameters
 import org.bitcoins.core.currency.CurrencyUnit
 import org.bitcoins.core.hd.AddressType
 import org.bitcoins.core.protocol.BitcoinAddress
-import org.bitcoins.core.protocol.blockchain.BlockHeader
 import org.bitcoins.core.protocol.transaction.{
   Transaction,
   TransactionOutPoint,
@@ -71,8 +70,7 @@ trait WalletApi extends WalletLogger {
     * Takes in a block header and updates our TxoStates to the new chain tip
     * @param blockHeader Block header we are processing
     */
-  def updateUtxoPendingStates(
-      blockHeader: BlockHeader): Future[Vector[SpendingInfoDb]]
+  def updateUtxoPendingStates(): Future[Vector[SpendingInfoDb]]
 
   /** Gets the sum of all UTXOs in this wallet */
   def getBalance()(implicit ec: ExecutionContext): Future[CurrencyUnit] = {

wallet/src/main/scala/org/bitcoins/wallet/internal/TransactionProcessing.scala

@@ -313,7 +313,7 @@ private[wallet] trait TransactionProcessing extends WalletLogger {
           }
 
           // Update Txo State
-          val updateF = updateUtxoConfirmedState(unreservedTxo, blockHash)
+          val updateF = updateUtxoConfirmedState(unreservedTxo)
 
           updateF.foreach(tx =>
             logger.debug(

wallet/src/main/scala/org/bitcoins/wallet/internal/UtxoHandling.scala

@@ -4,7 +4,6 @@ import org.bitcoins.core.compat._
 import org.bitcoins.core.hd.HDAccount
 import org.bitcoins.core.number.UInt32
 import org.bitcoins.core.protocol.BitcoinAddress
-import org.bitcoins.core.protocol.blockchain.BlockHeader
 import org.bitcoins.core.protocol.script.{
   P2WPKHWitnessSPKV0,
   P2WPKHWitnessV0,
@@ -15,7 +14,7 @@ import org.bitcoins.core.protocol.transaction.{
   TransactionOutPoint,
   TransactionOutput
 }
-import org.bitcoins.core.util.{EitherUtil, FutureUtil}
+import org.bitcoins.core.util.EitherUtil
 import org.bitcoins.core.wallet.utxo.{AddressTag, TxoState}
 import org.bitcoins.crypto.DoubleSha256DigestBE
 import org.bitcoins.wallet.api.{AddUtxoError, AddUtxoResult, AddUtxoSuccess}
@@ -27,7 +26,7 @@ import scala.util.{Failure, Success, Try}
 
 /**
   * Provides functionality related to handling UTXOs in our wallet.
-  * The most notable examples of functioanlity here are enumerating
+  * The most notable examples of functionality here are enumerating
   * UTXOs in the wallet and importing a UTXO into the wallet for later
   * spending.
   */
@@ -85,25 +84,25 @@ private[wallet] trait UtxoHandling extends WalletLogger {
   }
 
   protected def updateUtxoConfirmedState(
-      txo: SpendingInfoDb,
-      blockHash: DoubleSha256DigestBE): Future[SpendingInfoDb] = {
-    updateUtxoConfirmedStates(Vector(txo), blockHash).map(_.head)
+      txo: SpendingInfoDb): Future[SpendingInfoDb] = {
+    updateUtxoConfirmedStates(Vector(txo)).map(_.head)
   }
 
   protected def updateUtxoConfirmedStates(
-      txos: Vector[SpendingInfoDb],
-      blockHash: DoubleSha256DigestBE): Future[Vector[SpendingInfoDb]] = {
-    for {
-      confsOpt <- chainQueryApi.getNumberOfConfirmations(blockHash)
-      stateChanges <- {
-        confsOpt match {
+      spendingInfoDbs: Vector[SpendingInfoDb]): Future[
+    Vector[SpendingInfoDb]] = {
+
+    val byBlock = spendingInfoDbs.groupBy(_.blockHash)
+
+    val toUpdateFs = byBlock.map {
+      case (Some(blockHash), txos) =>
+        chainQueryApi.getNumberOfConfirmations(blockHash).map {
           case None =>
-            Future.successful(txos)
+            Vector.empty
           case Some(confs) =>
-            val updatedTxos = txos.map { txo =>
+            txos.map { txo =>
               txo.state match {
-                case TxoState.PendingConfirmationsReceived |
-                    TxoState.DoesNotExist =>
+                case TxoState.PendingConfirmationsReceived =>
                   if (confs >= walletConfig.requiredConfirmations) {
                     txo.copyWithState(TxoState.ConfirmedReceived)
                   } else {
@@ -119,14 +118,20 @@ private[wallet] trait UtxoHandling extends WalletLogger {
                   // We should keep the utxo as reserved so it is not used in
                   // a future transaction that it should not be in
                   txo
-                case TxoState.ConfirmedReceived | TxoState.ConfirmedSpent =>
+                case TxoState.DoesNotExist | TxoState.ConfirmedReceived |
+                    TxoState.ConfirmedSpent =>
                   txo
               }
             }
-            spendingInfoDAO.upsertAll(updatedTxos)
         }
-      }
-    } yield stateChanges
+      case (None, _) =>
+        Future.successful(Vector.empty)
+    }
+
+    for {
+      toUpdate <- Future.sequence(toUpdateFs)
+      updated <- spendingInfoDAO.upsertAll(toUpdate.flatten.toVector)
+    } yield updated
   }
 
   /**
@@ -281,23 +286,18 @@ private[wallet] trait UtxoHandling extends WalletLogger {
     val mempoolUtxos = Try(groupedUtxos(None)).getOrElse(Vector.empty)
 
     // get the ones in blocks
-    val utxosInBlocks = groupedUtxos.map {
-      case (Some(hash), utxos) =>
-        Some(hash, utxos)
+    val utxosInBlocks = groupedUtxos.flatMap {
+      case (Some(_), utxos) =>
+        utxos
       case (None, _) =>
         None
-    }.flatten
+    }.toVector
 
     for {
       updatedMempoolUtxos <- spendingInfoDAO.updateAll(mempoolUtxos)
       // update the confirmed ones
-      updatedBlockUtxos <-
-        FutureUtil
-          .sequentially(utxosInBlocks.toVector) {
-            case (hash, utxos) =>
-              updateUtxoConfirmedStates(utxos, hash)
-          }
-      updated = updatedMempoolUtxos ++ updatedBlockUtxos.flatten
+      updatedBlockUtxos <- updateUtxoConfirmedStates(utxosInBlocks)
+      updated = updatedMempoolUtxos ++ updatedBlockUtxos
       _ <- walletCallbacks.executeOnReservedUtxos(logger, updated)
     } yield updated
   }
@@ -316,11 +316,10 @@ private[wallet] trait UtxoHandling extends WalletLogger {
   }
 
   /** @inheritdoc */
-  override def updateUtxoPendingStates(
-      blockHeader: BlockHeader): Future[Vector[SpendingInfoDb]] = {
+  override def updateUtxoPendingStates(): Future[Vector[SpendingInfoDb]] = {
     for {
       infos <- spendingInfoDAO.findAllPendingConfirmation
-      updatedInfos <- updateUtxoConfirmedStates(infos, blockHeader.hashBE)
+      updatedInfos <- updateUtxoConfirmedStates(infos)
     } yield updatedInfos
   }
 }