From 706c9fe961eac5cf43f273daeade3da24d4fb803 Mon Sep 17 00:00:00 2001
From: Chris Stewart <stewart.chris1234@gmail.com>
Date: Tue, 17 Aug 2021 07:51:12 -0500
Subject: [PATCH] Add check for valid broadcast states (#3560)

---
 .../org/bitcoins/dlc/wallet/DLCWallet.scala   | 29 ++++++++++++++++++-
 1 file changed, 28 insertions(+), 1 deletion(-)

diff --git a/dlc-wallet/src/main/scala/org/bitcoins/dlc/wallet/DLCWallet.scala b/dlc-wallet/src/main/scala/org/bitcoins/dlc/wallet/DLCWallet.scala
index 31478bc720..b565e00b84 100644
--- a/dlc-wallet/src/main/scala/org/bitcoins/dlc/wallet/DLCWallet.scala
+++ b/dlc-wallet/src/main/scala/org/bitcoins/dlc/wallet/DLCWallet.scala
@@ -1155,8 +1155,18 @@ abstract class DLCWallet
 
   override def broadcastDLCFundingTx(
       contractId: ByteVector): Future[Transaction] = {
+    val dlcDbOptF = dlcDAO.findByContractId(contractId)
+    val fundingTxF = getDLCFundingTx(contractId)
     for {
-      tx <- getDLCFundingTx(contractId)
+      dlcDbOpt <- dlcDbOptF
+      _ = dlcDbOpt match {
+        case None =>
+          sys.error(
+            s"Cannot broadcast DLC when we don't know the contract, given contractId=${contractId}")
+        case Some(dlcDb) =>
+          isValidBroadcastState(dlcDb)
+      }
+      tx <- fundingTxF
       _ <- updateDLCState(contractId, DLCState.Broadcasted)
       _ = logger.info(
         s"Broadcasting funding transaction ${tx.txIdBE.hex} for contract ${contractId.toHex}")
@@ -1164,6 +1174,23 @@ abstract class DLCWallet
     } yield tx
   }
 
+  /** Checks if the DLC is in a valid state to broadcast the funding tx.
+    * This is particurarily useful for situations when users want to
+    * re-broadcast funding txs. You should only be able to re-broadcast
+    * a funding tx in two states, [[DLCState.Signed]] or [[DLCState.Broadcasted]]
+    * The reason accepted is needed is that is the state the DLC is in
+    * when a user gives us their sign message
+    */
+  private def isValidBroadcastState(dlcDb: DLCDb): DLCDb = {
+    dlcDb.state match {
+      case DLCState.Broadcasted | DLCState.Signed => dlcDb
+      case state @ (DLCState.Offered | DLCState.Confirmed | DLCState.Accepted |
+          DLCState.Claimed | DLCState.RemoteClaimed | DLCState.Refunded) =>
+        sys.error(
+          s"Cannot broadcast the dlc when it is in the state=${state} contractId=${dlcDb.contractIdOpt}")
+    }
+  }
+
   override def executeDLC(
       contractId: ByteVector,
       sigs: Seq[OracleAttestmentTLV]): Future[Transaction] = {