simple dns request and parse response

dns.cpp

#include <WS2tcpip.h>
#include <WinSock2.h>

#include <algorithm>
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <iomanip>
#include <string_view>
#include <sstream>
#include <type_traits>
#include <vector>
#include <utility>

#pragma comment(lib, "ws2_32.lib")

constexpr char const* FQDN = "google.com";
constexpr char const* DNS_IP_ADDRESS = "8.8.8.8";

struct dns_header {
  std::uint16_t id;
  union {
    std::uint16_t flags;
    // :note: nameless struct is not standard
    struct {
      // Response Code
      // 0: No Error
      // 1: Format Error
      // 2: Server Failure
      // 3: Name Error
      // 4: Not Implemented
      // 5: Refused
      std::byte rcode : 4;
      // Checking Disabled
      std::byte cd : 1;
      // Authentic Data
      std::byte ad : 1;
      // Reversed
      // Always Set 0
      std::byte z : 1;
      // Recursion Available
      // 0: Recursion Unavailable
      // 1: Recursion Available
      std::byte ra : 1;
      // Recursion Desired
      // 0: Not Recursively
      // 1: Request Recursively
      std::byte rd : 1;
      // TrucCation
      // 0: Not Truncated
      // 1: Truncated
      std::byte tc : 1;
      // Authoritative Answer
      // 0: Not Authoritative Server
      // 1: Authoritative Server
      std::byte aa : 1;
      // OPeration Code
      // 0: a standard query
      // 1: an inverse query
      // 2: a server status request
      std::byte opcode : 4;
      // Query/Response
      // 0: Query
      // 1: Response
      std::byte qr : 1;
    };
  };
  std::uint16_t qdcount;
  std::uint16_t ancount;
  std::uint16_t nscount;
  std::uint16_t arcount;
};

enum class record_class : std::uint16_t {
  // Internet
  IN_ = 0x01,
  // CSNET
  CS = 0x02,
  // CHAOS
  CH = 0x03,
  // Hesiod
  HS = 0x04,
  // Any Class
  ANY = 0xff
};

enum class record_type : std::uint16_t {
  A = 0x01,
  NS = 0x02,
  MD = 0x03,
  MF = 0x04,
  CNAME = 0x05,
  SOA = 0x06,
  MB = 0x07,
  MG = 0x08,
  MR = 0x09,
  NULL_ = 0x0A,
  WKS = 0x0B,
  PTR = 0x0C,
  HINFO = 0x0D,
  MINFO = 0x0E,
  MX = 0x0F,
  TXT = 0x10,
  // RFC 3596
  AAAA = 0x1C,
};

struct question_section {
  std::string qname;
  record_type qtype{};
  record_class qclass{};
};

struct resource_record {
  std::string fqdn;
  record_type type{};
  record_class clazz{};
  // seconds
  std::uint32_t ttl{};
  // bytes
  std::uint16_t rdlength{};
  std::vector<std::byte> rdata;
};

template <typename Os>
Os& operator<<(Os& os, resource_record const& record) {
  os << "Name:\t\t" << record.fqdn << "\n";
  os << "Type:\t\t" << static_cast<std::underlying_type_t<record_type>>(record.type) << "\n";
  os << "Class:\t\t" << static_cast<std::underlying_type_t<record_class>>(record.clazz) << "\n";
  os << "TTL:\t\t" << record.ttl << "\n";
  os << "RDLength:\t" << record.rdlength << "\n";
  // assume IPv4
  os << "RData:\t\t";
  bool is_first = true;
  for (auto const& address : record.rdata) {
    if (!std::exchange(is_first, false)) {
      os << ".";
    }
    os << std::dec << std::to_integer<int>(address) << std::hex;
  }
  os << "\n";
  return os;
}

struct dns_response {
  dns_header header;
  std::vector<question_section> questions;
  std::vector<resource_record> answers;
  std::vector<resource_record> authorities;
  std::vector<resource_record> additionals;
};

template <typename Os>
Os& operator<<(Os& os, dns_response const& res) {
  // header
  os << "Header:\n";
  os << "ID:\t\t" << res.header.id << "\n";
  os << "QR:\t\t" << std::boolalpha << std::to_integer<int>(res.header.qr) << "\n";
  os << "OPCODE:\t\t" << std::to_integer<int>(res.header.opcode) << "\n";
  os << "AA:\t\t" << std::boolalpha << std::to_integer<int>(res.header.aa) << "\n";
  os << "TC:\t\t" << std::boolalpha << std::to_integer<int>(res.header.tc) << "\n";
  os << "RD:\t\t" << std::boolalpha << std::to_integer<int>(res.header.rd) << "\n";
  os << "RA:\t\t" << std::boolalpha << std::to_integer<int>(res.header.ra) << "\n";
  os << "Z:\t\t" << std::boolalpha << std::to_integer<int>(res.header.z) << "\n";
  os << "AD:\t\t" << std::boolalpha << std::to_integer<int>(res.header.ad) << "\n";
  os << "CD:\t\t" << std::boolalpha << std::to_integer<int>(res.header.cd) << "\n";
  os << "RCODE:\t\t" << std::to_integer<int>(res.header.rcode) << "\n";
  os << "QDCOUNT:\t" << res.header.qdcount << "\n";
  os << "ANCOUNT:\t" << res.header.ancount << "\n";
  os << "NSCOUNT:\t" << res.header.nscount << "\n";
  os << "ARCOUNT:\t" << res.header.arcount << "\n";
  os << "\n";

  // questions
  if (res.header.qdcount > 0) {
    os << "Questions:\n";
    for (auto const& q : res.questions) {
      os << "Name:\t\t" << q.qname << "\n";
      os << "Type:\t\t" << static_cast<std::underlying_type_t<record_type>>(q.qtype) << "\n";
      os << "Class:\t\t" << static_cast<std::underlying_type_t<record_class>>(q.qclass) << "\n";
      os << "\n";
    }
  }

  // answers
  if (res.header.ancount > 0) {
    os << "Answers:\n";
    for (auto const& a : res.answers) {
      os << a;
      os << "\n";
    }
    os << "\n";
  }
  // authorities
  if (res.header.nscount > 0) {
    os << "Authorities:\n";
    for (auto const& a : res.authorities) {
      os << a;
      os << "\n";
    }
    os << "\n";
  }
  // additionals
  if (res.header.arcount > 0) {
    os << "Additionals:\n";
    for (auto const& a : res.additionals) {
      os << a;
      os << "\n";
    }
    os << "\n";
  }
  return os;
}

std::vector<std::byte> build_request(std::string_view fqdn) {
  // Header
  dns_header header;
  std::memset(&header, 0, sizeof(header));
  header.id = htons(0x862a);
  header.rd = std::byte{1};
  header.ad = std::byte{1};
  header.flags = htons(header.flags);
  header.qdcount = htons(1);

  // Question
  // - QName
  std::istringstream iss(fqdn.data());
  std::ostringstream oss;
  std::string domain;
  while (std::getline(iss, domain, '.')) {
    const auto size = (std::uint8_t) domain.size();
    oss << size;
    oss << domain;
  }
  std::string qname = oss.str();
  // "\0"(1byte) + Type(2byte) + Class(2byte)
  std::vector<std::byte> question(qname.size() + 1 + 2 + 2, std::byte{'\0'});
  decltype(question)::pointer cursor = question.data();
  std::memcpy(cursor, qname.c_str(), qname.size());
  // offset 1byte(for '\0')
  cursor += qname.size() + 1;

  // - QType
  std::uint16_t qtype = htons(1);
  std::memcpy(cursor, &qtype, sizeof(qtype));
  cursor += sizeof(qtype);

  // - QClass
  std::uint16_t qclass = htons(1);
  std::memcpy(cursor, &qclass, sizeof(qclass));

  // build request
  std::vector<std::byte> request(sizeof(header) + question.size(), std::byte{});
  std::memcpy(request.data(), &header, sizeof(header));
  std::memcpy(request.data() + sizeof(header), question.data(), question.size());

  return request;
}

/**
 * @brief parse label
 * @param response - response from server
 * @param offset - position to start reading
 * @return read bytes and parsed label
 */
std::pair<std::ptrdiff_t, std::string> parse_label(std::vector<std::byte> const& response, std::size_t offset) {
  std::ostringstream label_ss;
  auto cursor = response.data() + offset;
  while (std::to_integer<std::uint8_t>(*cursor) != '\0') {
    auto const length = std::to_integer<std::uint8_t>(*cursor);
    ++cursor;
    std::vector<char> domain(length + 1, '\0');
    std::memcpy(domain.data(), cursor, length);
    cursor += length;
    label_ss << domain.data() << ".";
  }
  ++cursor;
  std::string label = label_ss.str();
  // trim trailing dot
  label = label.substr(0, label.size() - 1);
  return {cursor - (response.data() + offset), label};
}

/**
 * @brief parse Resource Record.
 * @param response - response from server
 * @param offset - position to start reading.
 * @return read bytes and parsed data
 */
std::pair<std::size_t, resource_record> parse_resource_record(std::vector<std::byte> const& response, std::size_t offset) {
  const auto flag_mask = 0b1100'0000'0000'0000;
  auto cursor = response.data() + offset;
  resource_record record;
  // Name
  // check message compression
  std::uint16_t raw_name;
  std::memcpy(&raw_name, cursor, sizeof(raw_name));
  raw_name = ntohs(raw_name);
  if ((raw_name & flag_mask) == flag_mask) {
    // compressed
    auto const pointer_offset = raw_name & ~flag_mask;
    auto const& [bytes, domain_name] = parse_label(response, pointer_offset);
    // flag(2bit) + offset(14bit)
    cursor += 2;
    record.fqdn = domain_name;
  } else {
    // not compressed. label is in place
    // TODO: labels ending with a pointer
    auto const& [bytes, domain_name] = parse_label(response, offset);
    cursor += bytes;
    record.fqdn = domain_name;
  }

  // Type
  std::memcpy(&record.type, cursor, sizeof(record.type));
  record.type = static_cast<record_type>(ntohs(static_cast<std::underlying_type_t<record_type>>(record.type)));
  cursor += sizeof(record.type);

  // Class
  std::memcpy(&record.clazz, cursor, sizeof(record.clazz));
  record.clazz = static_cast<record_class>(ntohs(static_cast<std::underlying_type_t<record_type>>(record.clazz)));
  cursor += sizeof(record.clazz);

  // TTL
  std::memcpy(&record.ttl, cursor, sizeof(record.ttl));
  record.ttl = ntohl(record.ttl);
  cursor += sizeof(record.ttl);

  // RDLength
  std::memcpy(&record.rdlength, cursor, sizeof(record.rdlength));
  record.rdlength = ntohs(record.rdlength);
  cursor += sizeof(record.rdlength);

  // RDATA
  record.rdata.resize(record.rdlength);
  std::memcpy(record.rdata.data(), cursor, record.rdlength);

  return {cursor - (response.data() + offset), record};
}

dns_response parse_response(std::vector<std::byte> const& response) {
  // Header
  dns_header header;
  auto cursor = response.data();
  std::memcpy(&header.id, cursor, sizeof(header.id));
  header.id = ntohs(header.id);
  cursor += sizeof(header.id);

  std::memcpy(&header.flags, cursor, sizeof(header.flags));
  header.flags = ntohs(header.flags);
  cursor += sizeof(header.flags);

  std::memcpy(&header.qdcount, cursor, sizeof(header.qdcount));
  header.qdcount = ntohs(header.qdcount);
  cursor += sizeof(header.qdcount);

  std::memcpy(&header.ancount, cursor, sizeof(header.ancount));
  header.ancount = ntohs(header.ancount);
  cursor += sizeof(header.ancount);

  std::memcpy(&header.nscount, cursor, sizeof(header.nscount));
  header.nscount = ntohs(header.nscount);
  cursor += sizeof(header.nscount);

  std::memcpy(&header.arcount, cursor, sizeof(header.arcount));
  header.arcount = ntohs(header.arcount);
  cursor += sizeof(header.arcount);

  // Question
  std::vector<question_section> questions;
  for (int i = 0; i < header.qdcount; ++i) {
    question_section question;
    // - QName
    auto const& [bytes, qname] = parse_label(response, cursor - response.data());
    question.qname = qname;
    cursor += bytes;

    // - Type
    std::memcpy(&question.qtype, cursor, sizeof(question.qtype));
    question.qtype = static_cast<record_type>(ntohs(static_cast<std::underlying_type_t<record_type>>(question.qtype)));
    cursor += sizeof(question.qtype);

    // - Class
    std::memcpy(&question.qclass, cursor, sizeof(question.qclass));
    question.qclass = static_cast<record_class>(ntohs(static_cast<std::underlying_type_t<record_type>>(question.qclass)));
    cursor += sizeof(question.qclass);

    questions.push_back(question);
  }

  // Answer
  std::vector<resource_record> answers;
  for (int i = 0; i < header.ancount; ++i) {
    auto const& [bytes, answer] = parse_resource_record(response, cursor - response.data());
    cursor += bytes;

    answers.push_back(answer);
  }

  // Authority
  std::vector<resource_record> authorities;
  for (int i = 0; i < header.nscount; ++i) {
    auto const& [bytes, authority] = parse_resource_record(response, cursor - response.data());
    cursor += bytes;

    authorities.push_back(authority);
  }

  // Additional
  std::vector<resource_record> additionals;
  for (int i = 0; i < header.arcount; ++i) {
    auto const& [bytes, additional] = parse_resource_record(response, cursor - response.data());
    cursor += bytes;

    additionals.push_back(additional);
  }

  return {
    header,
    questions,
    answers,
    authorities,
    additionals
  };
}

// for debug print
void dump(std::vector<std::byte> const& binary) {
  int counter = 0;
  for (auto const byte : binary) {
    ++counter;
    std::cout << std::hex << std::setfill('0') << std::setw(2) << std::to_integer<int>(byte) << " ";
    if (counter % 0x8 == 0) {
      std::cout << " ";
    }
    if (counter % 0x10 == 0) {
      std::cout << "\n";
    }
  }
  std::cout << "\n";
}

int fetch_dns() {
  WSAData wsaData;

  SOCKET sock;
  struct sockaddr_in addr;

  auto const result = WSAStartup(MAKEWORD(2, 0), &wsaData);
  if (result != 0) {
    std::cerr << "WSAStartup failed: " << result << std::endl;
    return 1;
  }

  sock = socket(AF_INET, SOCK_DGRAM, 0);

  inet_pton(AF_INET, DNS_IP_ADDRESS, &addr.sin_addr.S_un.S_addr);

  addr.sin_family = AF_INET;
  addr.sin_port = htons(53);

  auto request = build_request(FQDN);

  std::cout << "---raw request---\n";
  dump(request);

  sendto(sock, reinterpret_cast<char*>(request.data()), static_cast<int>(request.size()), 0, reinterpret_cast<struct sockaddr*>(&addr), sizeof(addr));

  std::vector<std::byte> buffer(512);

  memset(buffer.data(), 0, buffer.size());
  auto size = recv(sock, (char*) buffer.data(), (int) buffer.size(), 0);
  if (size != 0) {
    buffer.resize(size);
    std::cout << "---raw response---\n";
    dump(buffer);
    auto const& response = parse_response(buffer);
    std::cout << "---formatted response---\n";
    std::cout << response;
  }

  closesocket(sock);

  WSACleanup();
  return 0;
}

int main() {
  return fetch_dns();
}

/*
---raw request---
86 2a 01 20 00 01 00 00  00 00 00 00 06 67 6f 6f
67 6c 65 03 63 6f 6d 00  00 01 00 01
---raw response---
86 2a 81 80 00 01 00 01  00 00 00 00 06 67 6f 6f
67 6c 65 03 63 6f 6d 00  00 01 00 01 c0 0c 00 01
00 01 00 00 00 76 00 04  ac d9 1a ee
---formatted response---
Header:
ID:             862a
QR:             1
OPCODE:         0
AA:             0
TC:             0
RD:             1
RA:             1
Z:              0
AD:             0
CD:             0
RCODE:          0
QDCOUNT:        1
ANCOUNT:        1
NSCOUNT:        0
ARCOUNT:        0

Questions:
Name:           google.com
Type:           1
Class:          1

Answers:
Name:           google.com
Type:           1
Class:          1
TTL:            76
RDLength:       4
RData:          172.217.26.238
*/