Test DNS TCP split length reads

src/iocore/dns/DNS.cc

@@ -876,21 +876,19 @@ DNSHandler::recv_dns(int /* event ATS_UNUSED */, Event * /* e ATS_UNUSED */)
         }
         if (dnsc->tcp_data.total_length == 0) {
           // Read the 2-byte length prefix incrementally
-          res = dnsc->sock.recv(dnsc->tcp_data.length_buf + dnsc->tcp_data.length_read,
-                                sizeof(dnsc->tcp_data.length_buf) - dnsc->tcp_data.length_read, 0);
+          unsigned char length_buf[sizeof(dnsc->tcp_data.length_buf)];
+
+          res = dnsc->sock.recv(length_buf, dnsc->tcp_data.length_prefix_bytes_remaining(), 0);
           if (res == -EAGAIN) {
             break;
           }
           if (res <= 0) {
             goto Lerror;
           }
-          dnsc->tcp_data.length_read += res;
-          if (dnsc->tcp_data.length_read < sizeof(dnsc->tcp_data.length_buf)) {
+          dnsc->tcp_data.append_length_prefix_bytes(length_buf, res);
+          if (!dnsc->tcp_data.length_prefix_is_complete()) {
             continue;
           }
-          uint16_t net_length;
-          memcpy(&net_length, dnsc->tcp_data.length_buf, sizeof(net_length));
-          dnsc->tcp_data.total_length = ntohs(net_length);
           if (dnsc->tcp_data.total_length == 0) {
             goto Lerror;
           }

src/iocore/dns/P_DNSConnection.h

@@ -30,6 +30,8 @@
 
 #pragma once
 
+#include <cstddef>
+#include <cstring>
 #include "iocore/dns/DNSEventIO.h"
 #include "iocore/dns/DNSProcessor.h"
 
@@ -98,6 +100,31 @@ struct DNSConnection {
     unsigned short length_read   = 0;
     unsigned short total_length  = 0;
     unsigned short done_reading  = 0;
+    size_t
+    length_prefix_bytes_remaining() const
+    {
+      return sizeof(length_buf) - length_read;
+    }
+
+    bool
+    length_prefix_is_complete() const
+    {
+      return length_read == sizeof(length_buf);
+    }
+
+    size_t
+    append_length_prefix_bytes(void const *bytes, size_t nbytes)
+    {
+      size_t const remaining = length_prefix_bytes_remaining();
+      size_t const consumed  = nbytes < remaining ? nbytes : remaining;
+      memcpy(length_buf + length_read, bytes, consumed);
+      length_read += consumed;
+      if (length_prefix_is_complete()) {
+        total_length = (static_cast<unsigned short>(length_buf[0]) << 8) | length_buf[1];
+      }
+      return consumed;
+    }
+
     void
     reset()
     {

src/iocore/dns/unit_tests/CMakeLists.txt

@@ -21,3 +21,10 @@ add_executable(test_HostEnt test_HostEnt.cc ../HostEnt.cc)
 target_include_directories(test_HostEnt PRIVATE "${PROJECT_SOURCE_DIR}/include")
 target_link_libraries(test_HostEnt PRIVATE Catch2::Catch2WithMain ts::tscore ts::tsutil ts::inkevent)
 add_catch2_test(NAME test_dns_HostEnt COMMAND $<TARGET_FILE:test_HostEnt>)
+
+add_executable(test_DNSTCPData test_DNSTCPData.cc)
+target_include_directories(
+  test_DNSTCPData PRIVATE "${PROJECT_SOURCE_DIR}/include" "${PROJECT_SOURCE_DIR}/src/iocore/dns"
+)
+target_link_libraries(test_DNSTCPData PRIVATE Catch2::Catch2WithMain ts::tscore ts::tsutil ts::inkevent)
+add_catch2_test(NAME test_dns_DNSTCPData COMMAND $<TARGET_FILE:test_DNSTCPData>)

src/iocore/dns/unit_tests/test_DNSTCPData.cc

@@ -0,0 +1,66 @@
+/** @file
+
+  Unit tests for DNS-over-TCP read state.
+
+  @section license License
+
+    Licensed to the Apache Software Foundation (ASF) under one
+  or more contributor license agreements.  See the NOTICE file
+  distributed with this work for additional information
+  regarding copyright ownership.  The ASF licenses this file
+  to you under the Apache License, Version 2.0 (the
+  "License"); you may not use this file except in compliance
+  with the License.  You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+ */
+
+#include <catch2/catch_test_macros.hpp>
+
+#include "P_DNSConnection.h"
+
+TEST_CASE("DNS TCP length prefix can be read one byte at a time", "[dns][tcp]")
+{
+  DNSConnection::TCPData tcp_data;
+  unsigned char const    prefix[] = {0x12, 0x34};
+
+  REQUIRE(tcp_data.length_prefix_bytes_remaining() == 2);
+  REQUIRE_FALSE(tcp_data.length_prefix_is_complete());
+
+  REQUIRE(tcp_data.append_length_prefix_bytes(prefix, 1) == 1);
+  CHECK(tcp_data.length_read == 1);
+  CHECK(tcp_data.length_prefix_bytes_remaining() == 1);
+  CHECK_FALSE(tcp_data.length_prefix_is_complete());
+  CHECK(tcp_data.total_length == 0);
+
+  REQUIRE(tcp_data.append_length_prefix_bytes(prefix + 1, 1) == 1);
+  CHECK(tcp_data.length_read == 2);
+  CHECK(tcp_data.length_prefix_bytes_remaining() == 0);
+  CHECK(tcp_data.length_prefix_is_complete());
+  CHECK(tcp_data.total_length == 0x1234);
+}
+
+TEST_CASE("DNS TCP length prefix state clamps and resets", "[dns][tcp]")
+{
+  DNSConnection::TCPData tcp_data;
+  unsigned char const    prefix[] = {0x00, 0x08, 0xff};
+
+  REQUIRE(tcp_data.append_length_prefix_bytes(prefix, sizeof(prefix)) == 2);
+  CHECK(tcp_data.length_prefix_bytes_remaining() == 0);
+  CHECK(tcp_data.length_prefix_is_complete());
+  CHECK(tcp_data.total_length == 8);
+
+  tcp_data.done_reading = 4;
+  tcp_data.reset();
+  CHECK(tcp_data.length_read == 0);
+  CHECK(tcp_data.length_prefix_bytes_remaining() == 2);
+  CHECK_FALSE(tcp_data.length_prefix_is_complete());
+  CHECK(tcp_data.total_length == 0);
+  CHECK(tcp_data.done_reading == 0);
+}