Make DLCDataHandler more type safe (#4123)

dlc-node-test/src/test/scala/org/bitcoins/dlc/node/DLCNegotiationTest.scala

@@ -63,7 +63,7 @@ class DLCNegotiationTest extends BitcoinSDualWalletTest {
         accept <- walletA.acceptDLCOffer(offer, None, None)
 
         // Send accept message to begin p2p
-        _ = handler ! accept.toMessage
+        _ = handler ! DLCDataHandler.Received(accept.toMessage)
 
         _ <- TestAsyncUtil.awaitConditionF(
           () =>

dlc-node-test/src/test/scala/org/bitcoins/dlc/node/DLCServerTest.scala

@@ -71,18 +71,18 @@ class DLCServerTest extends BitcoinSActorFixtureWithDLCWallet {
                   ByteVector.fromValidHex("00112233445566778899aabbccddeeff"))
         clientConnectionHandler = clientConnectionHandlerOpt.get
         _ = clientProbe.send(clientConnectionHandler, pingTLV)
-        _ = serverProbe.expectMsg(LnMessage(pingTLV))
+        _ = serverProbe.expectMsg(DLCDataHandler.Received(LnMessage(pingTLV)))
         pongTLV = PongTLV.forIgnored(
           ByteVector.fromValidHex("00112233445566778899aabbccddeeff"))
         serverConnectionHandler = serverConnectionHandlerOpt.get
         _ = serverProbe.send(serverConnectionHandler, pongTLV)
-        _ = clientProbe.expectMsg(LnMessage(pongTLV))
+        _ = clientProbe.expectMsg(DLCDataHandler.Received(LnMessage(pongTLV)))
         // 131063 - is a magic size for OS X when this test case starts failing (131073 overall TLV size)
         ignored = ByteVector.fill(65000)(0x55)
         bigTLV =
           PongTLV.forIgnored(ignored)
         _ = clientProbe.send(clientConnectionHandler, bigTLV)
-        _ = serverProbe.expectMsg(LnMessage(bigTLV))
+        _ = serverProbe.expectMsg(DLCDataHandler.Received(LnMessage(bigTLV)))
         _ = clientProbe.send(clientConnectionHandler,
                              DLCConnectionHandler.CloseConnection)
         _ = clientProbe.send(clientConnectionHandler, pingTLV)

dlc-node/src/main/scala/org/bitcoins/dlc/node/DLCConnectionHandler.scala

@@ -82,7 +82,7 @@ class DLCConnectionHandler(
         log.debug(s"Unaligned bytes: ${newUnalignedBytes.toHex}")
       }
 
-      messages.foreach(m => handler ! m)
+      messages.foreach(m => handler ! DLCDataHandler.Received(m))
 
       connection ! Tcp.ResumeReading
       context.become(connected(newUnalignedBytes))

dlc-node/src/main/scala/org/bitcoins/dlc/node/DLCDataHandler.scala

@@ -11,18 +11,15 @@ class DLCDataHandler(dlcWalletApi: DLCWalletApi, connectionHandler: ActorRef)
     extends Actor
     with ActorLogging {
   implicit val ec: ExecutionContextExecutor = context.system.dispatcher
-  import DLCDataHandler.{Send, SendLnMessage}
 
   override def preStart(): Unit = {
     val _ = context.watch(connectionHandler)
   }
 
   override def receive: Receive = LoggingReceive {
-    case Send(tlv) =>
-      connectionHandler ! tlv
-    case SendLnMessage(lnMessage) =>
-      connectionHandler ! lnMessage
-    case lnMessage: LnMessage[TLV] =>
+    case DLCDataHandler.Send(lnMessage) =>
+      send(lnMessage)
+    case DLCDataHandler.Received(lnMessage) =>
       log.info(s"Received LnMessage ${lnMessage.typeName}")
       val f: Future[Unit] = handleTLVMessage(lnMessage)
       f.failed.foreach(err =>
@@ -45,7 +42,7 @@ class DLCDataHandler(dlcWalletApi: DLCWalletApi, connectionHandler: ActorRef)
         Future.unit //is this right?
       case ping: PingTLV =>
         val pong = PongTLV.forIgnored(ping.ignored)
-        connectionHandler ! LnMessage(pong)
+        send(LnMessage(pong))
         Future.unit
       case pong: PongTLV =>
         log.debug(s"Received pong message $pong")
@@ -53,14 +50,14 @@ class DLCDataHandler(dlcWalletApi: DLCWalletApi, connectionHandler: ActorRef)
       case dlcOffer: DLCOfferTLV =>
         val f = for {
           accept <- dlcWalletApi.acceptDLCOffer(dlcOffer, None, None)
-          _ = connectionHandler ! accept.toMessage
-        } yield ()
+        } yield send(accept.toMessage)
         f
       case dlcAccept: DLCAcceptTLV =>
         val f = for {
           sign <- dlcWalletApi.signDLC(dlcAccept)
-          _ = connectionHandler ! sign.toMessage
-        } yield ()
+        } yield {
+          send(sign.toMessage)
+        }
         f
       case dlcSign: DLCSignTLV =>
         val f = for {
@@ -70,6 +67,10 @@ class DLCDataHandler(dlcWalletApi: DLCWalletApi, connectionHandler: ActorRef)
         f
     }
   }
+
+  private def send(lnMessage: LnMessage[TLV]): Unit = {
+    connectionHandler ! lnMessage
+  }
 }
 
 object DLCDataHandler {
@@ -77,9 +78,8 @@ object DLCDataHandler {
   type Factory = (DLCWalletApi, ActorContext, ActorRef) => ActorRef
 
   sealed trait Command
-  case class Received(tlv: TLV) extends Command
-  case class Send(tlv: TLV) extends Command
-  case class SendLnMessage(message: LnMessage[TLV]) extends Command
+  case class Received(message: LnMessage[TLV]) extends Command
+  case class Send(message: LnMessage[TLV]) extends Command
 
   def defaultFactory(
       dlcWalletApi: DLCWalletApi,

dlc-node/src/main/scala/org/bitcoins/dlc/node/DLCNode.scala

@@ -76,7 +76,7 @@ case class DLCNode(wallet: DLCWalletApi)(implicit
                                       externalPayoutAddress,
                                       externalChangeAddress)
     } yield {
-      handler ! DLCDataHandler.SendLnMessage(accept.toMessage)
+      handler ! DLCDataHandler.Send(accept.toMessage)
       accept
     }
   }