fix: Multiple potential issues with the tcp DNS client (#32636)

johanandren · leviramsey · web-flow · commit e17d1cd09a68 · 2025-02-03T17:10:42.000+01:00
* fix: improve TCP DNS client interactions with TCP actor
* ensure that TCP connection is closed when TcpDnsClient fails
* Fix for potential re-order of response bytes during high load
* Deathwatch the connection to make sure we are never stuck in a broken state
---------

Co-authored-by: Levi Ramsey &lt;levi.ramsey@alum.cs.umass.edu&gt;
diff --git a/akka-actor-tests/src/test/scala/akka/io/dns/internal/DnsClientSpec.scala b/akka-actor-tests/src/test/scala/akka/io/dns/internal/DnsClientSpec.scala
@@ -6,17 +6,17 @@ package akka.io.dns.internal
 
 import java.net.InetSocketAddress
 import java.util.concurrent.atomic.AtomicBoolean
-
 import scala.collection.immutable.Seq
 import scala.concurrent.duration._
-
 import akka.actor.Props
 import akka.io.Udp
 import akka.io.dns.{ RecordClass, RecordType }
 import akka.io.dns.internal.DnsClient.{ Answer, Question4 }
+import akka.testkit.WithLogCapturing
 import akka.testkit.{ AkkaSpec, ImplicitSender, TestProbe }
 
-class DnsClientSpec extends AkkaSpec with ImplicitSender {
+class DnsClientSpec extends AkkaSpec("""akka.loglevel = DEBUG
+      akka.loggers = ["akka.testkit.SilenceAllTestEventListener"]""") with ImplicitSender with WithLogCapturing {
   "The async DNS client" should {
     val exampleRequest = Question4(42, "akka.io")
     val exampleRequestMessage =
diff --git a/akka-actor-tests/src/test/scala/akka/io/dns/internal/TcpDnsClientSpec.scala b/akka-actor-tests/src/test/scala/akka/io/dns/internal/TcpDnsClientSpec.scala
@@ -5,17 +5,19 @@
 package akka.io.dns.internal
 
 import java.net.InetSocketAddress
-
-import scala.collection.immutable.Seq
-
+import akka.actor.ActorRef
+import akka.actor.PoisonPill
 import akka.actor.Props
 import akka.io.Tcp
 import akka.io.Tcp.{ Connected, PeerClosed, Register }
 import akka.io.dns.{ RecordClass, RecordType }
 import akka.io.dns.internal.DnsClient.Answer
+import akka.testkit.EventFilter
+import akka.testkit.WithLogCapturing
 import akka.testkit.{ AkkaSpec, ImplicitSender, TestProbe }
 
-class TcpDnsClientSpec extends AkkaSpec with ImplicitSender {
+class TcpDnsClientSpec extends AkkaSpec("""akka.loglevel = DEBUG
+      akka.loggers = ["akka.testkit.SilenceAllTestEventListener"]""") with ImplicitSender with WithLogCapturing {
   import TcpDnsClient._
 
   "The async TCP DNS client" should {
@@ -51,6 +53,27 @@ class TcpDnsClientSpec extends AkkaSpec with ImplicitSender {
       tcpExtensionProbe.expectMsg(Tcp.Connect(dnsServerAddress))
     }
 
+    "terminated if the connection terminates unexpectedly" in {
+      val tcpExtensionProbe = TestProbe()
+      val answerProbe = TestProbe()
+
+      val client = system.actorOf(Props(new TcpDnsClient(tcpExtensionProbe.ref, dnsServerAddress, answerProbe.ref)))
+
+      client ! exampleRequestMessage
+
+      tcpExtensionProbe.expectMsg(Tcp.Connect(dnsServerAddress))
+      tcpExtensionProbe.lastSender ! Connected(dnsServerAddress, localAddress)
+      expectMsgType[Register]
+      val registered = tcpExtensionProbe.lastSender
+
+      expectMsgType[Tcp.Write]
+
+      answerProbe.watch(client)
+
+      registered ! PoisonPill
+      answerProbe.expectTerminated(client)
+    }
+
     "accept a fragmented TCP response" in {
       val tcpExtensionProbe = TestProbe()
       val answerProbe = TestProbe()
@@ -72,22 +95,80 @@ class TcpDnsClientSpec extends AkkaSpec with ImplicitSender {
       answerProbe.expectMsg(Answer(42, Nil))
     }
 
-    "accept merged TCP responses" in {
+    "accept multiple fragmented TCP responses" in {
       val tcpExtensionProbe = TestProbe()
       val answerProbe = TestProbe()
 
       val client = system.actorOf(Props(new TcpDnsClient(tcpExtensionProbe.ref, dnsServerAddress, answerProbe.ref)))
 
       client ! exampleRequestMessage
-      client ! exampleRequestMessage.copy(id = 43)
 
       tcpExtensionProbe.expectMsg(Tcp.Connect(dnsServerAddress))
       tcpExtensionProbe.lastSender ! Connected(dnsServerAddress, localAddress)
       expectMsgType[Register]
       val registered = tcpExtensionProbe.lastSender
 
+      // pretend write+ack+write happened, so three requests written, now both coming back.
+      // (we need to make sure buffer is not reordered by sandwitched responses)
       expectMsgType[Tcp.Write]
+      val fullResponse1 = encodeLength(exampleResponseMessage.write().length) ++ exampleResponseMessage.write()
+      val exampleResponseMessage2 = exampleResponseMessage.copy(id = 43)
+      val fullResponse2 = encodeLength(exampleResponseMessage2.write().length) ++ exampleResponseMessage2.write()
+      val exampleResponseMessage3 = exampleResponseMessage.copy(id = 44)
+      val fullResponse3 = encodeLength(exampleResponseMessage3.write().length) ++ exampleResponseMessage3.write()
+      registered ! Tcp.Received(fullResponse1.take(8))
+      Thread.sleep(30) // give things some time to go wrong
+      registered ! Tcp.Received(fullResponse1.drop(8) ++ fullResponse2.take(8))
+      Thread.sleep(30)
+      registered ! Tcp.Received(fullResponse2.drop(8) ++ fullResponse3.take(8))
+      Thread.sleep(30)
+      registered ! Tcp.Received(fullResponse3.drop(8))
+
+      answerProbe.expectMsg(Answer(42, Nil))
+      answerProbe.expectMsg(Answer(43, Nil))
+      answerProbe.expectMsg(Answer(44, Nil))
+    }
+
+    "respect backpressure from the TCP actor" in {
+      val tcpExtensionProbe = TestProbe()
+
+      val client = system.actorOf(Props(new TcpDnsClient(tcpExtensionProbe.ref, dnsServerAddress, ActorRef.noSender)))
+
+      client ! exampleRequestMessage
+      client ! exampleRequestMessage.copy(id = 43)
+
+      tcpExtensionProbe.expectMsg(Tcp.Connect(dnsServerAddress))
+      tcpExtensionProbe.lastSender ! Connected(dnsServerAddress, localAddress)
+      expectMsgType[Register]
+      val registered = tcpExtensionProbe.lastSender
+
+      val ack = expectMsgType[Tcp.Write].ack
+      expectNoMessage()
+      registered ! ack
+
       expectMsgType[Tcp.Write]
+    }
+
+    "accept merged TCP responses" in {
+      val tcpExtensionProbe = TestProbe()
+      val answerProbe = TestProbe()
+
+      val client = system.actorOf(Props(new TcpDnsClient(tcpExtensionProbe.ref, dnsServerAddress, answerProbe.ref)))
+
+      client ! exampleRequestMessage
+      client ! exampleRequestMessage.copy(id = 43)
+
+      tcpExtensionProbe.expectMsg(Tcp.Connect(dnsServerAddress))
+      tcpExtensionProbe.lastSender ! Connected(dnsServerAddress, localAddress)
+      expectMsgType[Register]
+      val registered = tcpExtensionProbe.lastSender
+
+      var ack = expectMsgType[Tcp.Write].ack
+
+      registered ! ack
+
+      ack = expectMsgType[Tcp.Write].ack
+
       val fullResponse =
         encodeLength(exampleResponseMessage.write().length) ++ exampleResponseMessage.write() ++
         encodeLength(exampleResponseMessage.write().length) ++ exampleResponseMessage.copy(id = 43).write()
@@ -97,5 +178,71 @@ class TcpDnsClientSpec extends AkkaSpec with ImplicitSender {
       answerProbe.expectMsg(Answer(42, Nil))
       answerProbe.expectMsg(Answer(43, Nil))
     }
+
+    "report its failure to the outer client" in {
+      val tcpExtensionProbe = TestProbe()
+      val answerProbe = TestProbe()
+
+      val failToConnectClient =
+        system.actorOf(Props(new TcpDnsClient(tcpExtensionProbe.ref, dnsServerAddress, answerProbe.ref)))
+
+      failToConnectClient ! exampleRequestMessage
+
+      val connect = tcpExtensionProbe.expectMsg(Tcp.Connect(dnsServerAddress))
+
+      failToConnectClient ! Tcp.CommandFailed(connect)
+
+      answerProbe.expectMsg(DnsClient.TcpDropped)
+
+      val closesWithErrorClient =
+        system.actorOf(Props(new TcpDnsClient(tcpExtensionProbe.ref, dnsServerAddress, answerProbe.ref)))
+
+      closesWithErrorClient ! exampleRequestMessage
+
+      tcpExtensionProbe.expectMsg(connect)
+      tcpExtensionProbe.lastSender ! Connected(dnsServerAddress, localAddress)
+      expectMsgType[Register]
+      val registered = tcpExtensionProbe.lastSender
+
+      registered ! Tcp.ErrorClosed("BOOM!")
+
+      answerProbe.expectMsg(DnsClient.TcpDropped)
+    }
+
+    "should drop older requests when TCP connection backpressures" in {
+      val tcpExtensionProbe = TestProbe()
+      val connectionProbe = TestProbe()
+
+      val client = system.actorOf(Props(new TcpDnsClient(tcpExtensionProbe.ref, dnsServerAddress, ActorRef.noSender)))
+
+      // initial request
+      val initialRequest = exampleRequestMessage.copy(id = 1)
+      client ! initialRequest
+
+      tcpExtensionProbe.expectMsg(Tcp.Connect(dnsServerAddress))
+      tcpExtensionProbe.lastSender.tell(Connected(dnsServerAddress, localAddress), connectionProbe.ref)
+      connectionProbe.expectMsgType[Register]
+      val registered = connectionProbe.lastSender
+
+      var write = connectionProbe.expectMsgType[Tcp.Write]
+      write.data.drop(2) shouldBe initialRequest.write()
+
+      // 1 in flight, 14 more, buffer fits 10, should drop 5 oldest (id 2 - 6)
+      EventFilter.warning(occurrences = 5, pattern = "Dropping oldest buffered DNS request").intercept {
+        (2 to 16).foreach { i =>
+          client ! exampleRequestMessage.copy(id = i.toShort)
+        }
+      }
+
+      // initial write is acked
+      registered ! write.ack
+
+      // the rest of the buffered should be handled
+      (7 to 16).foreach { i =>
+        write = connectionProbe.expectMsgType[Tcp.Write]
+        write.data.drop(2) shouldBe exampleRequestMessage.copy(id = i.toShort).write()
+        registered ! write.ack // next ack
+      }
+    }
   }
 }
diff --git a/akka-actor/src/main/scala/akka/io/dns/internal/DnsClient.scala b/akka-actor/src/main/scala/akka/io/dns/internal/DnsClient.scala
@@ -6,7 +6,6 @@ package akka.io.dns.internal
 
 import java.net.{ InetAddress, InetSocketAddress }
 
-import scala.annotation.nowarn
 import scala.collection.{ immutable => im }
 import scala.concurrent.duration._
 import scala.util.Try
@@ -35,6 +34,8 @@ import akka.pattern.{ BackoffOpts, BackoffSupervisor }
 
   final case class DropRequest(question: DnsQuestion)
 
+  case object TcpDropped
+
   // sent as an indication that as of the time of sending, `id` is not being used
   //  by an active question.  Useful for a questioner which tracks which ids it can use
   final case class Dropped(id: Short) extends NoSerializationVerificationNeeded
@@ -47,6 +48,8 @@ import akka.pattern.{ BackoffOpts, BackoffSupervisor }
     if (question.name.last == '.') Seq(unchangedPair, question.name.dropRight(1) -> question.qType)
     else Seq(unchangedPair, (question.name + '.') -> question.qType)
   }
+
+  private final case class InFlight(replyTo: ActorRef, message: Message, tcpRequest: Boolean = false)
 }
 
 /**
@@ -62,7 +65,7 @@ import akka.pattern.{ BackoffOpts, BackoffSupervisor }
   val udp = IO(Udp)
   val tcp = IO(Tcp)
 
-  private[internal] var inflightRequests: Map[Short, (ActorRef, Message)] = Map.empty
+  private var inflightRequests: Map[Short, InFlight] = Map.empty
 
   lazy val tcpDnsClient: ActorRef = createTcpClient()
 
@@ -90,12 +93,11 @@ import akka.pattern.{ BackoffOpts, BackoffSupervisor }
   /**
    * Silent to allow map update syntax
    */
-  @nowarn()
   def ready(socket: ActorRef): Receive = {
     case DropRequest(question) =>
       val id = question.id
       inflightRequests.get(id) match {
-        case Some((_, sentMsg)) =>
+        case Some(InFlight(_, sentMsg, _)) =>
           val sentQs = sentMsg.questions.map { question =>
             question.name -> question.qType
           }
@@ -129,7 +131,7 @@ import akka.pattern.{ BackoffOpts, BackoffSupervisor }
         log.debug("Resolving [{}] (A)", name)
 
         val msg = message(name, id, RecordType.A)
-        inflightRequests += (id -> (sender() -> msg))
+        inflightRequests += (id -> InFlight(sender(), msg))
         log.debug("Message [{}] to [{}]: [{}]", id, ns, msg)
         socket ! Udp.Send(msg.write(), ns)
       }
@@ -143,7 +145,7 @@ import akka.pattern.{ BackoffOpts, BackoffSupervisor }
         log.debug("Resolving [{}] (AAAA)", name)
 
         val msg = message(name, id, RecordType.AAAA)
-        inflightRequests += (id -> (sender() -> msg))
+        inflightRequests += (id -> InFlight(sender(), msg))
         log.debug("Message to [{}]: [{}]", ns, msg)
         socket ! Udp.Send(msg.write(), ns)
       }
@@ -156,7 +158,7 @@ import akka.pattern.{ BackoffOpts, BackoffSupervisor }
       } else {
         log.debug("Resolving [{}] (SRV)", name)
         val msg = message(name, id, RecordType.SRV)
-        inflightRequests += (id -> (sender() -> msg))
+        inflightRequests += (id -> InFlight(sender(), msg))
         log.debug("Message to [{}]: [{}]", ns, msg)
         socket ! Udp.Send(msg.write(), ns)
       }
@@ -169,7 +171,7 @@ import akka.pattern.{ BackoffOpts, BackoffSupervisor }
           Try {
             val msg = Message.parse(send.payload)
             inflightRequests.get(msg.id).foreach {
-              case (s, _) =>
+              case InFlight(s, _, _) =>
                 s ! Failure(new RuntimeException("Send failed to nameserver"))
                 inflightRequests -= msg.id
             }
@@ -185,8 +187,10 @@ import akka.pattern.{ BackoffOpts, BackoffSupervisor }
       if (msg.flags.isTruncated) {
         log.debug("DNS response truncated, falling back to TCP")
         inflightRequests.get(msg.id) match {
-          case Some((_, msg)) =>
-            tcpDnsClient ! msg
+          case Some(inFlight) =>
+            inflightRequests = inflightRequests.updated(msg.id, inFlight.copy(tcpRequest = true))
+            tcpDnsClient ! inFlight.message
+
           case _ =>
             log.debug("Client for id {} not found. Discarding unsuccessful response.", msg.id)
         }
@@ -197,7 +201,7 @@ import akka.pattern.{ BackoffOpts, BackoffSupervisor }
       }
     case UdpAnswer(questions, response) =>
       inflightRequests.get(response.id) match {
-        case Some((reply, sentMsg)) =>
+        case Some(InFlight(reply, sentMsg, _)) =>
           val sentQs = sentMsg.questions.flatMap(withAndWithoutTrailingDots).toSet
           val answeredQs = questions.flatMap(withAndWithoutTrailingDots).toSet
 
@@ -220,19 +224,26 @@ import akka.pattern.{ BackoffOpts, BackoffSupervisor }
     // for TCP, we don't have to use the question for correlation
     case response: Answer =>
       inflightRequests.get(response.id) match {
-        case Some((reply, sentMsg)) =>
-          reply ! response
+        case Some(InFlight(replyTo, _, _)) =>
+          replyTo ! response
           inflightRequests -= response.id
 
         case None =>
           log.debug("Client for id [{}] not found. Discarding response.", response.id)
       }
 
+    case TcpDropped =>
+      log.warning("TCP client failed, clearing inflight resolves which were being resolved by TCP")
+
+      inflightRequests = inflightRequests.filterNot {
+        case (_, inFlight) => inFlight.tcpRequest
+      }
+
     case Udp.Unbind  => socket ! Udp.Unbind
     case Udp.Unbound => context.stop(self)
   }
 
-  def createTcpClient() = {
+  def createTcpClient(): ActorRef = {
     context.actorOf(
       BackoffSupervisor.props(
         BackoffOpts.onFailure(
diff --git a/akka-actor/src/main/scala/akka/io/dns/internal/TcpDnsClient.scala b/akka-actor/src/main/scala/akka/io/dns/internal/TcpDnsClient.scala
