Move DNSClient to separate file

vlasovsoft1979 · Mar 31, 2024 · 113f0a9 · 113f0a9
1 parent f9424a9
commit 113f0a9
Show file tree

Hide file tree

Showing 6 changed files with 166 additions and 139 deletions.
diff --git a/libdns/CMakeLists.txt b/libdns/CMakeLists.txt
@@ -11,10 +11,10 @@ add_library(dns STATIC
     dns_package.cpp dns_package.h
     dns_selector.cpp dns_selector.h
     dns_socket.cpp dns_socket.h
+    dns_client.cpp dns_client.h
     dns.cpp dns.h
 )
 
-
 target_include_directories(dns PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})
 target_link_libraries(dns PRIVATE JsonCpp::JsonCpp)
 

diff --git a/libdns/dns.cpp b/libdns/dns.cpp
@@ -462,128 +462,3 @@ void DNSServer::join()
 {
     impl->join();
 }
-
-DNSClient::DNSClient(const std::string& host, int port)
-    : host(host)
-    , port(port)
-{}
-
-DNSPackage DNSClient::requestUdp(uint16_t id, DNSRecordType type, const std::string& host)
-{
-    DNSPackage package;
-    package.header.ID = id;
-    package.header.flags.RD = 1;
-    package.header.QDCOUNT = 1;
-    package.requests.emplace_back(DNSRequest{ type, host });
-    DNSBuffer buf;
-    package.append(buf);
-    if (buf.result.size() > UDP_SIZE)
-    {
-        throw std::runtime_error("UDP request too big");
-    }
-
-    SOCKET s = socket(AF_INET, SOCK_DGRAM, 0);
-    if (s == INVALID_SOCKET)
-    {
-        throw std::runtime_error("Can't create UDP socket");
-    }
-    sockaddr_in server = { 0 };
-    server.sin_family = AF_INET;
-    server.sin_port = htons(port);
-    inet_pton(AF_INET, this->host.c_str(), &server.sin_addr);
-
-    int bytes_sent = sendto(s, reinterpret_cast<const char*>(&buf.result[0]), static_cast<int>(buf.result.size()), 0, reinterpret_cast<sockaddr*>(&server), static_cast<int>(sizeof(server)));
-    if (bytes_sent < buf.result.size())
-    {
-        throw std::runtime_error("Error sending UDP data");
-    }
-
-    std::vector<uint8_t> in_buf(UDP_SIZE, 0);
-    int bytes_received = recvfrom(s, reinterpret_cast<char*>(&in_buf[0]), static_cast<int>(in_buf.size()), 0, nullptr, nullptr);
-    if (bytes_received < 0)
-    {
-        throw std::runtime_error("Error receiving UDP data");
-    }
-
-    closesocket(s);
-
-    DNSPackage response(&in_buf[0]);
-
-    return response;
-}
-
-DNSPackage DNSClient::requestTcp(uint16_t id, DNSRecordType type, const std::string& host)
-{
-    DNSPackage package;
-    package.header.ID = id;
-    package.header.flags.RD = 1;
-    package.header.QDCOUNT = 1;
-    package.requests.emplace_back(DNSRequest{ type, host });
-    DNSBuffer buf;
-    buf.append(static_cast<uint16_t>(0u));  // SIZE (will be calculated later)
-    buf.data_start = buf.result.size();
-    package.append(buf);
-    buf.overwrite_uint16(0, static_cast<uint16_t>(buf.result.size() - sizeof(uint16_t)));
-
-    SOCKET s = socket(AF_INET, SOCK_STREAM, 0);
-    if (s == INVALID_SOCKET)
-    {
-        throw std::runtime_error("Can't create TCP socket");
-    }
-    sockaddr_in server = { 0 };
-    server.sin_family = AF_INET;
-    server.sin_port = htons(port);
-    inet_pton(AF_INET, this->host.c_str(), &server.sin_addr);
-
-    int result = connect(s, reinterpret_cast<sockaddr*>(&server), sizeof(server));
-    if (result == SOCKET_ERROR) {
-        throw std::runtime_error("Can't connect to server");
-    }
-
-    int bytes_sent = send(s, reinterpret_cast<const char*>(&buf.result[0]), static_cast<int>(buf.result.size()), 0);
-    if (bytes_sent < buf.result.size())
-    {
-        throw std::runtime_error("Error sending TCP data");
-    }
-
-    std::vector<uint8_t> in_buf(sizeof(uint16_t), 0);
-    int bytes_received = recv(s, reinterpret_cast<char*>(&in_buf[0]), static_cast<int>(in_buf.size()), 0);
-    if (bytes_received < 0)
-    {
-        throw std::runtime_error("Error receiving TCP data");
-    }
-
-    const uint8_t* ptr = &in_buf[0];
-    uint16_t size = get_uint16(ptr);
-    in_buf.resize(in_buf.size() + size, 0);
-    bytes_received = recv(s, reinterpret_cast<char*>(&in_buf[sizeof(uint16_t)]), static_cast<int>(size), 0);
-    if (bytes_received < 0)
-    {
-        throw std::runtime_error("Error receiving TCP data");
-    }
-
-    closesocket(s);
-
-    DNSPackage response(&in_buf[sizeof(uint16_t)]);
-
-    return response;
-}
-
-bool DNSClient::command(const std::string& cmd)
-{
-    SOCKET s = socket(AF_INET, SOCK_DGRAM, 0);
-    if (s == INVALID_SOCKET)
-    {
-        throw std::runtime_error("Can't create UDP socket");
-    }
-    sockaddr_in server = {0};
-    server.sin_family = AF_INET;
-    server.sin_port = htons(port);
-    inet_pton(AF_INET, host.c_str(), &server.sin_addr);
-
-    int length = sendto(s, cmd.c_str(), static_cast<int>(cmd.size()), 0, reinterpret_cast<sockaddr*>(&server), static_cast<int>(sizeof(server)));
-
-    closesocket(s);
-
-    return length == cmd.size();
-}
diff --git a/libdns/dns.h b/libdns/dns.h
@@ -25,16 +25,3 @@ class DNSServer
     std::unique_ptr<DNSServerImpl> impl;
 };
 
-class DNSClient
-{
-public:
-    DNSClient(const std::string& host, int port);
-
-    bool command(const std::string& cmd);
-    DNSPackage requestUdp(uint16_t id, DNSRecordType type, const std::string& host);
-    DNSPackage requestTcp(uint16_t id, DNSRecordType type, const std::string& host);
-
-private:
-    std::string host;
-    int port;
-};
diff --git a/libdns/dns_client.cpp b/libdns/dns_client.cpp
@@ -0,0 +1,143 @@
+#include "dns_client.h"
+
+#if defined(_WIN32)
+#include <WinSock2.h>
+#include <ws2tcpip.h>
+#else
+#include <sys/socket.h>
+#include <arpa/inet.h>
+#include <unistd.h>
+#include <fcntl.h>
+#include <string.h>
+#endif
+
+#include "dns_socket.h"
+#include "dns_request.h"
+#include "dns_buffer.h"
+#include "dns_utils.h"
+
+DNSClient::DNSClient(const std::string& host, int port)
+    : host(host)
+    , port(port)
+{}
+
+DNSPackage DNSClient::requestUdp(uint16_t id, DNSRecordType type, const std::string& host)
+{
+    DNSPackage package;
+    package.header.ID = id;
+    package.header.flags.RD = 1;
+    package.header.QDCOUNT = 1;
+    package.requests.emplace_back(DNSRequest{ type, host });
+    DNSBuffer buf;
+    package.append(buf);
+    if (buf.result.size() > UDP_SIZE)
+    {
+        throw std::runtime_error("UDP request too big");
+    }
+
+    SOCKET s = socket(AF_INET, SOCK_DGRAM, 0);
+    if (s == INVALID_SOCKET)
+    {
+        throw std::runtime_error("Can't create UDP socket");
+    }
+    sockaddr_in server = { 0 };
+    server.sin_family = AF_INET;
+    server.sin_port = htons(port);
+    inet_pton(AF_INET, this->host.c_str(), &server.sin_addr);
+
+    int bytes_sent = sendto(s, reinterpret_cast<const char*>(&buf.result[0]), static_cast<int>(buf.result.size()), 0, reinterpret_cast<sockaddr*>(&server), static_cast<int>(sizeof(server)));
+    if (bytes_sent < buf.result.size())
+    {
+        throw std::runtime_error("Error sending UDP data");
+    }
+
+    std::vector<uint8_t> in_buf(UDP_SIZE, 0);
+    int bytes_received = recvfrom(s, reinterpret_cast<char*>(&in_buf[0]), static_cast<int>(in_buf.size()), 0, nullptr, nullptr);
+    if (bytes_received < 0)
+    {
+        throw std::runtime_error("Error receiving UDP data");
+    }
+
+    closesocket(s);
+
+    DNSPackage response(&in_buf[0]);
+
+    return response;
+}
+
+DNSPackage DNSClient::requestTcp(uint16_t id, DNSRecordType type, const std::string& host)
+{
+    DNSPackage package;
+    package.header.ID = id;
+    package.header.flags.RD = 1;
+    package.header.QDCOUNT = 1;
+    package.requests.emplace_back(DNSRequest{ type, host });
+    DNSBuffer buf;
+    buf.append(static_cast<uint16_t>(0u));  // SIZE (will be calculated later)
+    buf.data_start = buf.result.size();
+    package.append(buf);
+    buf.overwrite_uint16(0, static_cast<uint16_t>(buf.result.size() - sizeof(uint16_t)));
+
+    SOCKET s = socket(AF_INET, SOCK_STREAM, 0);
+    if (s == INVALID_SOCKET)
+    {
+        throw std::runtime_error("Can't create TCP socket");
+    }
+    sockaddr_in server = { 0 };
+    server.sin_family = AF_INET;
+    server.sin_port = htons(port);
+    inet_pton(AF_INET, this->host.c_str(), &server.sin_addr);
+
+    int result = connect(s, reinterpret_cast<sockaddr*>(&server), sizeof(server));
+    if (result == SOCKET_ERROR) {
+        throw std::runtime_error("Can't connect to server");
+    }
+
+    int bytes_sent = send(s, reinterpret_cast<const char*>(&buf.result[0]), static_cast<int>(buf.result.size()), 0);
+    if (bytes_sent < buf.result.size())
+    {
+        throw std::runtime_error("Error sending TCP data");
+    }
+
+    std::vector<uint8_t> in_buf(sizeof(uint16_t), 0);
+    int bytes_received = recv(s, reinterpret_cast<char*>(&in_buf[0]), static_cast<int>(in_buf.size()), 0);
+    if (bytes_received < 0)
+    {
+        throw std::runtime_error("Error receiving TCP data");
+    }
+
+    const uint8_t* ptr = &in_buf[0];
+    uint16_t size = get_uint16(ptr);
+    in_buf.resize(in_buf.size() + size, 0);
+    bytes_received = recv(s, reinterpret_cast<char*>(&in_buf[sizeof(uint16_t)]), static_cast<int>(size), 0);
+    if (bytes_received < 0)
+    {
+        throw std::runtime_error("Error receiving TCP data");
+    }
+
+    closesocket(s);
+
+    DNSPackage response(&in_buf[sizeof(uint16_t)]);
+
+    return response;
+}
+
+bool DNSClient::command(const std::string& cmd)
+{
+    SOCKET s = socket(AF_INET, SOCK_DGRAM, 0);
+    if (s == INVALID_SOCKET)
+    {
+        throw std::runtime_error("Can't create UDP socket");
+    }
+    sockaddr_in server = {0};
+    server.sin_family = AF_INET;
+    server.sin_port = htons(port);
+    inet_pton(AF_INET, host.c_str(), &server.sin_addr);
+
+    int length = sendto(s, cmd.c_str(), static_cast<int>(cmd.size()), 0, reinterpret_cast<sockaddr*>(&server), static_cast<int>(sizeof(server)));
+
+    closesocket(s);
+
+    return length == cmd.size();
+}
+
diff --git a/libdns/dns_client.h b/libdns/dns_client.h
@@ -0,0 +1,21 @@
+#pragma once
+
+#include <string>
+
+#include "dns_consts.h"
+#include "dns_package.h"
+
+class DNSClient
+{
+public:
+    DNSClient(const std::string& host, int port);
+
+    bool command(const std::string& cmd);
+    DNSPackage requestUdp(uint16_t id, DNSRecordType type, const std::string& host);
+    DNSPackage requestTcp(uint16_t id, DNSRecordType type, const std::string& host);
+
+private:
+    std::string host;
+    int port;
+};
+
diff --git a/tests/tst_dns.cpp b/tests/tst_dns.cpp
@@ -6,6 +6,7 @@
 #include "dns_buffer.h"
 #include "dns_header.h"
 #include "dns_package.h"
+#include "dns_client.h"
 
 static const std::string HOST = "127.0.0.1";
 static const int PORT = 10000;