diff --git a/src/iocore/dns/DNS.cc b/src/iocore/dns/DNS.cc index 4c6cecdcca3..77871ae69ec 100644 --- a/src/iocore/dns/DNS.cc +++ b/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; } diff --git a/src/iocore/dns/P_DNSConnection.h b/src/iocore/dns/P_DNSConnection.h index 14473797221..9fd607ddfbf 100644 --- a/src/iocore/dns/P_DNSConnection.h +++ b/src/iocore/dns/P_DNSConnection.h @@ -30,6 +30,8 @@ #pragma once +#include +#include #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(length_buf[0]) << 8) | length_buf[1]; + } + return consumed; + } + void reset() { diff --git a/src/iocore/dns/unit_tests/CMakeLists.txt b/src/iocore/dns/unit_tests/CMakeLists.txt index 602eb638dd0..78bc5fcc1d1 100644 --- a/src/iocore/dns/unit_tests/CMakeLists.txt +++ b/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 $) + +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 $) diff --git a/src/iocore/dns/unit_tests/test_DNSTCPData.cc b/src/iocore/dns/unit_tests/test_DNSTCPData.cc new file mode 100644 index 00000000000..5455796e1f0 --- /dev/null +++ b/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 + +#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); +}