LookupServer: Unify DNSRequest & DNSResponse

They're really the same thing: a DNS packet can contain both questions and
answers, and there's a single bit in the header that determines whether the
packet represents a query or a response. It'll be simpler for us to represent
both types of packets using the same class.

This class can be both serialized and deserialized to/from a raw DNS packet.

8 files changed, 145 insertions, 207 deletions

diff --git a/Userland/Services/LookupServer/CMakeLists.txt b/Userland/Services/LookupServer/CMakeLists.txt
index 0204aec286..f8d868cc27 100644
--- a/Userland/Services/LookupServer/CMakeLists.txt
+++ b/Userland/Services/LookupServer/CMakeLists.txt
@@ -3,8 +3,7 @@ compile_ipc(LookupClient.ipc LookupClientEndpoint.h)
 
 set(SOURCES
     DNSAnswer.cpp
-    DNSRequest.cpp
-    DNSResponse.cpp
+    DNSPacket.cpp
     LookupServer.cpp
     LookupServerEndpoint.h
     LookupClientEndpoint.h
diff --git a/Userland/Services/LookupServer/ClientConnection.cpp b/Userland/Services/LookupServer/ClientConnection.cpp
index a7d90b4680..4c16802dc5 100644
--- a/Userland/Services/LookupServer/ClientConnection.cpp
+++ b/Userland/Services/LookupServer/ClientConnection.cpp
@@ -25,7 +25,7 @@
  */
 
 #include "ClientConnection.h"
-#include "DNSRequest.h"
+#include "DNSPacket.h"
 #include "LookupServer.h"
 #include <AK/IPv4Address.h>
 
diff --git a/Userland/Services/LookupServer/DNSResponse.cpp b/Userland/Services/LookupServer/DNSPacket.cpp
index a7031906c3..7ed266a3ca 100644
--- a/Userland/Services/LookupServer/DNSResponse.cpp
+++ b/Userland/Services/LookupServer/DNSPacket.cpp
@@ -1,5 +1,6 @@
 /*
  * Copyright (c) 2018-2020, Andreas Kling <kling@serenityos.org>
+ * Copyright (c) 2021, Sergey Bugaev <bugaevc@serenityos.org>
  * All rights reserved.
  *
  * Redistribution and use in source and binary forms, with or without
@@ -24,15 +25,92 @@
  * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  */
 
-#include "DNSResponse.h"
+#include "DNSPacket.h"
 #include "DNSPacketHeader.h"
-#include "DNSRequest.h"
-#include <AK/Debug.h>
 #include <AK/IPv4Address.h>
+#include <AK/MemoryStream.h>
 #include <AK/StringBuilder.h>
+#include <arpa/inet.h>
+#include <ctype.h>
+#include <stdlib.h>
 
 namespace LookupServer {
 
+void DNSPacket::add_question(const String& name, u16 record_type, ShouldRandomizeCase should_randomize_case)
+{
+    ASSERT(m_questions.size() <= UINT16_MAX);
+
+    if (name.is_empty())
+        return;
+
+    StringBuilder builder;
+    for (size_t i = 0; i < name.length(); ++i) {
+        u8 ch = name[i];
+        if (should_randomize_case == ShouldRandomizeCase::Yes) {
+            // Randomize the 0x20 bit in every ASCII character.
+            if (isalpha(ch)) {
+                if (arc4random_uniform(2))
+                    ch |= 0x20;
+                else
+                    ch &= ~0x20;
+            }
+        }
+        builder.append(ch);
+    }
+
+    if (name[name.length() - 1] != '.')
+        builder.append('.');
+
+    m_questions.empend(builder.to_string(), record_type, (u16)C_IN);
+}
+
+ByteBuffer DNSPacket::to_byte_buffer() const
+{
+    DNSPacketHeader header;
+    header.set_id(m_id);
+    if (is_query())
+        header.set_is_query();
+    else
+        header.set_is_response();
+    // FIXME: What should this be?
+    header.set_opcode(0);
+    header.set_truncated(false); // hopefully...
+    header.set_recursion_desired(true);
+    // FIXME: what should the be for requests?
+    header.set_recursion_available(true);
+    header.set_question_count(m_questions.size());
+    header.set_answer_count(m_answers.size());
+
+    DuplexMemoryStream stream;
+
+    stream << ReadonlyBytes { &header, sizeof(header) };
+    for (auto& question : m_questions) {
+        auto parts = question.name().split('.');
+        for (auto& part : parts) {
+            stream << (u8)part.length();
+            stream << part.bytes();
+        }
+        stream << '\0';
+        stream << htons(question.record_type());
+        stream << htons(question.class_code());
+    }
+    for (auto& answer : m_answers) {
+        auto parts = answer.name().split('.');
+        for (auto& part : parts) {
+            stream << (u8)part.length();
+            stream << part.bytes();
+        }
+        stream << '\0';
+        stream << htons(answer.type());
+        stream << htons(answer.class_code());
+        stream << htonl(answer.ttl());
+        stream << htons(answer.record_data().length());
+        stream << answer.record_data().bytes();
+    }
+
+    return stream.copy_into_contiguous_buffer();
+}
+
 static String parse_dns_name(const u8* data, size_t& offset, size_t max_offset, size_t recursion_level = 0);
 
 class [[gnu::packed]] DNSRecordWithoutName {
@@ -56,47 +134,49 @@ private:
 
 static_assert(sizeof(DNSRecordWithoutName) == 10);
 
-Optional<DNSResponse> DNSResponse::from_raw_response(const u8* raw_data, size_t raw_size)
+Optional<DNSPacket> DNSPacket::from_raw_packet(const u8* raw_data, size_t raw_size)
 {
     if (raw_size < sizeof(DNSPacketHeader)) {
         dbgln("DNS response not large enough ({} out of {}) to be a DNS packet.", raw_size, sizeof(DNSPacketHeader));
         return {};
     }
 
-    auto& response_header = *(const DNSPacketHeader*)(raw_data);
-#if LOOKUPSERVER_DEBUG
-    dbgln("Got response (ID: {})", response_header.id());
-    dbgln("  Question count: {}", response_header.question_count());
-    dbgln("    Answer count: {}", response_header.answer_count());
-    dbgln(" Authority count: {}", response_header.authority_count());
-    dbgln("Additional count: {}", response_header.additional_count());
+    auto& header = *(const DNSPacketHeader*)(raw_data);
+#ifdef LOOKUPSERVER_DEBUG
+    dbgln("Got packet (ID: {})", header.id());
+    dbgln("  Question count: {}", header.question_count());
+    dbgln("    Answer count: {}", header.answer_count());
+    dbgln(" Authority count: {}", header.authority_count());
+    dbgln("Additional count: {}", header.additional_count());
 #endif
 
-    DNSResponse response;
-    response.m_id = response_header.id();
-    response.m_code = response_header.response_code();
+    DNSPacket packet;
+    packet.m_id = header.id();
+    packet.m_query_or_response = header.is_response();
+    packet.m_code = header.response_code();
 
-    if (response.code() != DNSResponse::Code::NOERROR)
-        return response;
+    // FIXME: Should we parse further in this case?
+    if (packet.code() != Code::NOERROR)
+        return packet;
 
     size_t offset = sizeof(DNSPacketHeader);
 
-    for (u16 i = 0; i < response_header.question_count(); ++i) {
+    for (u16 i = 0; i < header.question_count(); i++) {
         auto name = parse_dns_name(raw_data, offset, raw_size);
         struct RawDNSAnswerQuestion {
             NetworkOrdered<u16> record_type;
             NetworkOrdered<u16> class_code;
         };
         auto& record_and_class = *(const RawDNSAnswerQuestion*)&raw_data[offset];
-        response.m_questions.empend(name, record_and_class.record_type, record_and_class.class_code);
+        packet.m_questions.empend(name, record_and_class.record_type, record_and_class.class_code);
         offset += 4;
-#if LOOKUPSERVER_DEBUG
-        auto& question = response.m_questions.last();
+#ifdef LOOKUPSERVER_DEBUG
+        auto& question = packet.m_questions.last();
         dbgln("Question #{}: name=_{}_, type={}, class={}", i, question.name(), question.record_type(), question.class_code());
 #endif
     }
 
-    for (u16 i = 0; i < response_header.answer_count(); ++i) {
+    for (u16 i = 0; i < header.answer_count(); ++i) {
         auto name = parse_dns_name(raw_data, offset, raw_size);
 
         auto& record = *(const DNSRecordWithoutName*)(&raw_data[offset]);
@@ -104,6 +184,7 @@ Optional<DNSResponse> DNSResponse::from_raw_response(const u8* raw_data, size_t
         String data;
 
         offset += sizeof(DNSRecordWithoutName);
+
         if (record.type() == T_PTR) {
             size_t dummy_offset = offset;
             data = parse_dns_name(raw_data, dummy_offset, raw_size);
@@ -113,14 +194,14 @@ Optional<DNSResponse> DNSResponse::from_raw_response(const u8* raw_data, size_t
             // FIXME: Parse some other record types perhaps?
             dbgln("data=(unimplemented record type {})", record.type());
         }
-#if LOOKUPSERVER_DEBUG
+#ifdef LOOKUPSERVER_DEBUG
         dbgln("Answer   #{}: name=_{}_, type={}, ttl={}, length={}, data=_{}_", i, name, record.type(), record.ttl(), record.data_length(), data);
 #endif
-        response.m_answers.empend(name, record.type(), record.record_class(), record.ttl(), data);
+        packet.m_answers.empend(name, record.type(), record.record_class(), record.ttl(), data);
         offset += record.data_length();
     }
 
-    return response;
+    return packet;
 }
 
 String parse_dns_name(const u8* data, size_t& offset, size_t max_offset, size_t recursion_level)
diff --git a/Userland/Services/LookupServer/DNSResponse.h b/Userland/Services/LookupServer/DNSPacket.h
index bfffb23561..c0e5ccd8a3 100644
--- a/Userland/Services/LookupServer/DNSResponse.h
+++ b/Userland/Services/LookupServer/DNSPacket.h
@@ -1,5 +1,6 @@
 /*
  * Copyright (c) 2020, Andreas Kling <kling@serenityos.org>
+ * Copyright (c) 2021, Sergey Bugaev <bugaevc@serenityos.org>
  * All rights reserved.
  *
  * Redistribution and use in source and binary forms, with or without
@@ -32,13 +33,37 @@
 #include <AK/Types.h>
 #include <AK/Vector.h>
 
+#define T_A 1
+#define T_NS 2
+#define T_CNAME 5
+#define T_SOA 6
+#define T_PTR 12
+#define T_MX 15
+
+#define C_IN 1
+
 namespace LookupServer {
 
-class DNSResponse {
+enum class ShouldRandomizeCase {
+    No = 0,
+    Yes
+};
+
+class DNSPacket {
 public:
-    static Optional<DNSResponse> from_raw_response(const u8*, size_t);
+    DNSPacket() { }
+
+    static Optional<DNSPacket> from_raw_packet(const u8*, size_t);
+    ByteBuffer to_byte_buffer() const;
+
+    bool is_query() const { return !m_query_or_response; }
+    bool is_response() const { return m_query_or_response; }
+    void set_is_query() { m_query_or_response = false; }
+    void set_is_response() { m_query_or_response = true; }
 
     u16 id() const { return m_id; }
+    void set_id(u16 id) { m_id = id; }
+
     const Vector<DNSQuestion>& questions() const { return m_questions; }
     const Vector<DNSAnswer>& answers() const { return m_answers; }
 
@@ -54,6 +79,8 @@ public:
         return m_answers.size();
     }
 
+    void add_question(const String& name, u16 record_type, ShouldRandomizeCase);
+
     enum class Code : u8 {
         NOERROR = 0,
         FORMERR = 1,
@@ -68,12 +95,12 @@ public:
     };
 
     Code code() const { return (Code)m_code; }
+    void set_code(Code code) { m_code = (u8)code; }
 
 private:
-    DNSResponse() { }
-
     u16 m_id { 0 };
     u8 m_code { 0 };
+    bool m_query_or_response { false };
     Vector<DNSQuestion> m_questions;
     Vector<DNSAnswer> m_answers;
 };
diff --git a/Userland/Services/LookupServer/DNSRequest.cpp b/Userland/Services/LookupServer/DNSRequest.cpp
deleted file mode 100644
index 243c0f727c..0000000000
--- a/Userland/Services/LookupServer/DNSRequest.cpp
+++ /dev/null
@@ -1,100 +0,0 @@
-/*
- * Copyright (c) 2018-2020, Andreas Kling <kling@serenityos.org>
- * All rights reserved.
- *
- * Redistribution and use in source and binary forms, with or without
- * modification, are permitted provided that the following conditions are met:
- *
- * 1. Redistributions of source code must retain the above copyright notice, this
- *    list of conditions and the following disclaimer.
- *
- * 2. Redistributions in binary form must reproduce the above copyright notice,
- *    this list of conditions and the following disclaimer in the documentation
- *    and/or other materials provided with the distribution.
- *
- * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
- * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
- */
-
-#include "DNSRequest.h"
-#include "DNSPacketHeader.h"
-#include <AK/MemoryStream.h>
-#include <AK/StringBuilder.h>
-#include <arpa/inet.h>
-#include <ctype.h>
-#include <stdlib.h>
-
-namespace LookupServer {
-
-const u16 C_IN = 1;
-
-DNSRequest::DNSRequest()
-    : m_id(arc4random_uniform(UINT16_MAX))
-{
-}
-
-void DNSRequest::add_question(const String& name, u16 record_type, ShouldRandomizeCase should_randomize_case)
-{
-    ASSERT(m_questions.size() <= UINT16_MAX);
-
-    if (name.is_empty())
-        return;
-
-    StringBuilder builder;
-    for (size_t i = 0; i < name.length(); ++i) {
-        u8 ch = name[i];
-        if (should_randomize_case == ShouldRandomizeCase::Yes) {
-            // Randomize the 0x20 bit in every ASCII character.
-            if (isalpha(ch)) {
-                if (arc4random_uniform(2))
-                    ch |= 0x20;
-                else
-                    ch &= ~0x20;
-            }
-        }
-        builder.append(ch);
-    }
-
-    if (name[name.length() - 1] != '.')
-        builder.append('.');
-
-    m_questions.empend(builder.to_string(), record_type, C_IN);
-}
-
-ByteBuffer DNSRequest::to_byte_buffer() const
-{
-    DNSPacketHeader request_header;
-    request_header.set_id(m_id);
-    request_header.set_is_query();
-    request_header.set_opcode(0);
-    request_header.set_truncated(false);
-    request_header.set_recursion_desired(true);
-    request_header.set_question_count(m_questions.size());
-
-    DuplexMemoryStream stream;
-
-    stream << ReadonlyBytes { &request_header, sizeof(request_header) };
-
-    for (auto& question : m_questions) {
-        auto parts = question.name().split('.');
-        for (auto& part : parts) {
-            stream << (u8)part.length();
-            stream << part.bytes();
-        }
-        stream << '\0';
-        stream << htons(question.record_type());
-        stream << htons(question.class_code());
-    }
-
-    return stream.copy_into_contiguous_buffer();
-}
-
-}
diff --git a/Userland/Services/LookupServer/DNSRequest.h b/Userland/Services/LookupServer/DNSRequest.h
deleted file mode 100644
index 71da4acd14..0000000000
--- a/Userland/Services/LookupServer/DNSRequest.h
+++ /dev/null
@@ -1,69 +0,0 @@
-/*
- * Copyright (c) 2020, Andreas Kling <kling@serenityos.org>
- * All rights reserved.
- *
- * Redistribution and use in source and binary forms, with or without
- * modification, are permitted provided that the following conditions are met:
- *
- * 1. Redistributions of source code must retain the above copyright notice, this
- *    list of conditions and the following disclaimer.
- *
- * 2. Redistributions in binary form must reproduce the above copyright notice,
- *    this list of conditions and the following disclaimer in the documentation
- *    and/or other materials provided with the distribution.
- *
- * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
- * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
- * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
- * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
- * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
- * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
- * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
- * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
- * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
- * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
- */
-
-#pragma once
-
-#include "DNSQuestion.h"
-#include <AK/Types.h>
-#include <AK/Vector.h>
-
-#define T_A 1
-#define T_NS 2
-#define T_CNAME 5
-#define T_SOA 6
-#define T_PTR 12
-#define T_MX 15
-
-namespace LookupServer {
-
-enum class ShouldRandomizeCase {
-    No = 0,
-    Yes
-};
-
-class DNSRequest {
-public:
-    DNSRequest();
-
-    void add_question(const String& name, u16 record_type, ShouldRandomizeCase);
-
-    const Vector<DNSQuestion>& questions() const { return m_questions; }
-
-    u16 question_count() const
-    {
-        ASSERT(m_questions.size() < UINT16_MAX);
-        return m_questions.size();
-    }
-
-    u16 id() const { return m_id; }
-    ByteBuffer to_byte_buffer() const;
-
-private:
-    u16 m_id { 0 };
-    Vector<DNSQuestion> m_questions;
-};
-
-}
diff --git a/Userland/Services/LookupServer/LookupServer.cpp b/Userland/Services/LookupServer/LookupServer.cpp
index 7145f54d89..7831734471 100644
--- a/Userland/Services/LookupServer/LookupServer.cpp
+++ b/Userland/Services/LookupServer/LookupServer.cpp
@@ -26,8 +26,7 @@
 
 #include "LookupServer.h"
 #include "ClientConnection.h"
-#include "DNSRequest.h"
-#include "DNSResponse.h"
+#include "DNSPacket.h"
 #include <AK/ByteBuffer.h>
 #include <AK/Debug.h>
 #include <AK/HashMap.h>
@@ -173,7 +172,9 @@ Vector<String> LookupServer::lookup(const String& hostname, const String& namese
         m_lookup_cache.remove(it);
     }
 
-    DNSRequest request;
+    DNSPacket request;
+    request.set_is_query();
+    request.set_id(arc4random_uniform(UINT16_MAX));
     request.add_question(hostname, record_type, should_randomize_case);
 
     auto buffer = request.to_byte_buffer();
@@ -204,7 +205,7 @@ Vector<String> LookupServer::lookup(const String& hostname, const String& namese
 
     did_get_response = true;
 
-    auto o_response = DNSResponse::from_raw_response(response_buffer, nrecv);
+    auto o_response = DNSPacket::from_raw_packet(response_buffer, nrecv);
     if (!o_response.has_value())
         return {};
 
@@ -215,7 +216,7 @@ Vector<String> LookupServer::lookup(const String& hostname, const String& namese
         return {};
     }
 
-    if (response.code() == DNSResponse::Code::REFUSED) {
+    if (response.code() == DNSPacket::Code::REFUSED) {
         if (should_randomize_case == ShouldRandomizeCase::Yes) {
             // Retry with 0x20 case randomization turned off.
             return lookup(hostname, nameserver, did_get_response, record_type, ShouldRandomizeCase::No);
diff --git a/Userland/Services/LookupServer/LookupServer.h b/Userland/Services/LookupServer/LookupServer.h
index 3297c5b5c7..6a2b04e3e7 100644
--- a/Userland/Services/LookupServer/LookupServer.h
+++ b/Userland/Services/LookupServer/LookupServer.h
@@ -26,8 +26,7 @@
 
 #pragma once
 
-#include "DNSRequest.h"
-#include "DNSResponse.h"
+#include "DNSPacket.h"
 #include <LibCore/Object.h>
 
 namespace LookupServer {