220 lines
7.5 KiB
C++
220 lines
7.5 KiB
C++
#pragma once
|
|
|
|
#include <cstdint>
|
|
#include <netinet/in.h>
|
|
#include <vector>
|
|
#include <memory>
|
|
#include "types.h"
|
|
|
|
namespace ts::dns {
|
|
class DNSBuilder;
|
|
namespace builder {
|
|
namespace rrbuilder {
|
|
class base;
|
|
}
|
|
|
|
class DNSHeader {
|
|
friend class ts::dns::DNSBuilder;
|
|
|
|
template <size_t index>
|
|
constexpr static size_t be2leIndex() {
|
|
int result = index;
|
|
if(result >= 8)
|
|
result -= 8;
|
|
else
|
|
result += 8;
|
|
return result;
|
|
}
|
|
public:
|
|
DNSHeader() {
|
|
this->_buffer = (uint16_t*) &this->internal_buffer;
|
|
}
|
|
|
|
explicit DNSHeader(uint16_t* buffer) : _buffer{buffer} { }
|
|
|
|
inline void id(uint16_t id) { this->_buffer[0] = id; }
|
|
|
|
inline void set_answer(bool flag) { this->set_option_bits<15>(flag, 1); }
|
|
|
|
inline void set_authoritative_answer(bool flag) { this->set_option_bits<10>(flag, 1); }
|
|
inline void set_truncation(bool flag) { this->set_option_bits<9>(flag, 1); }
|
|
inline void set_recursion_desired(bool flag) { this->set_option_bits<8>(flag, 1); }
|
|
inline void set_recursion_available(bool flag) { this->set_option_bits<7>(flag, 1); }
|
|
|
|
inline void set_query_type(uint8_t type) { this->set_option_bits<11>(type, 4); }
|
|
inline void set_response_code(uint8_t code) { this->set_option_bits<0>(code, 4); }
|
|
|
|
inline void set_query_count(uint16_t count) { this->_buffer[2] = htons(count); }
|
|
inline void set_answer_count(uint16_t count) { this->_buffer[3] = htons(count); }
|
|
inline void set_authority_count(uint16_t count) { this->_buffer[4] = htons(count); }
|
|
inline void set_additional_count(uint16_t count) { this->_buffer[5] = htons(count); }
|
|
|
|
[[nodiscard]] inline const void* buffer() const { return this->_buffer; }
|
|
[[nodiscard]] constexpr size_t buffer_size() const { return 12; }
|
|
private:
|
|
template <size_t index, size_t le_index = be2leIndex<index>()>
|
|
inline void set_option_bits(uint8_t value, uint8_t bit_count) {
|
|
auto clear = (1UL << bit_count) - 1;
|
|
this->_buffer[1] &= ~(clear << le_index);
|
|
this->_buffer[1] |= ((value & clear) << le_index);
|
|
}
|
|
|
|
uint16_t* _buffer{};
|
|
uint16_t internal_buffer[6]{0,0,0,0,0,0};
|
|
};
|
|
|
|
class DNSQuery {
|
|
friend class ts::dns::DNSBuilder;
|
|
public:
|
|
inline void set_qname(const std::string& value) { this->name = value; }
|
|
inline void set_qtype(const rrtype::value& value) { this->type = value; }
|
|
inline void set_qclass(const rrclass::value& value) { this->klass = value; }
|
|
|
|
DNSQuery(std::string name, rrtype::value type, rrclass::value klass) : name{std::move(name)}, type{type}, klass{klass} {}
|
|
DNSQuery() = default;
|
|
private:
|
|
|
|
std::string name;
|
|
rrtype::value type{rrtype::A};
|
|
rrclass::value klass{rrclass::IN};
|
|
};
|
|
|
|
class DNSResourceRecords {
|
|
friend class ts::dns::DNSBuilder;
|
|
public:
|
|
[[nodiscard]] inline std::string get_name() { return this->name; }
|
|
|
|
inline void set_ttl(uint32_t value) { this->ttl = value; }
|
|
inline void set_type(rrtype::value value) { this->type = value; }
|
|
inline void set_class(rrclass::value value) { this->klass = value; }
|
|
|
|
bool build(char*& /* buffer */, size_t& /* max length */, std::string& /* error */);
|
|
|
|
template <typename T>
|
|
[[nodiscard]] inline T& builder() {
|
|
if(T::type != this->type)
|
|
throw std::logic_error{"type mismatch"};
|
|
if(!this->payload_builder)
|
|
this->payload_builder.reset(new T{this});
|
|
auto entry = dynamic_cast<T*>(&*this->payload_builder);
|
|
if(!entry)
|
|
throw std::logic_error{"type mismatch to previous initialized"};
|
|
return *entry;
|
|
}
|
|
private:
|
|
uint32_t ttl{0};
|
|
|
|
std::string name;
|
|
rrtype::value type{rrtype::A};
|
|
rrclass::value klass{rrclass::IN};
|
|
|
|
std::unique_ptr<rrbuilder::base> payload_builder;
|
|
};
|
|
|
|
namespace rrbuilder {
|
|
class base {
|
|
public:
|
|
virtual bool build(char*& /* buffer */, size_t& /* max length */, std::string& /* error */) = 0;
|
|
|
|
protected:
|
|
explicit base(DNSResourceRecords* handle) : handle{handle} {}
|
|
private:
|
|
DNSResourceRecords* handle;
|
|
};
|
|
|
|
#define define_builder(name, base, ...) \
|
|
struct name : public base { \
|
|
friend class builder::DNSResourceRecords; \
|
|
public: \
|
|
static constexpr auto type = rrtype::name; \
|
|
bool build(char*&, size_t&, std::string&) override; \
|
|
__VA_ARGS__ \
|
|
private: \
|
|
explicit name(DNSResourceRecords* handle) : base{handle} {} \
|
|
}
|
|
|
|
define_builder(A, base,
|
|
void set_address(uint32_t address) { this->address = address; }
|
|
inline void set_address(const uint8_t (&address)[4]) {
|
|
this->set_address((address[3] << 24UL) | (address[2] << 16UL) | (address[1] << 8UL) | address[0]);
|
|
}
|
|
|
|
private:
|
|
uint32_t address{0};
|
|
);
|
|
|
|
define_builder(AAAA, base,
|
|
inline void set_address(const uint32_t (&address)[4]) {
|
|
this->set_address(address[0], address[1], address[2], address[3]);
|
|
}
|
|
|
|
inline void set_address(uint32_t a, uint32_t b, uint32_t c, uint32_t d) {
|
|
this->address[0] = a;
|
|
this->address[1] = b;
|
|
this->address[2] = c;
|
|
this->address[3] = d;
|
|
}
|
|
|
|
inline void set_address(const uint8_t (&address)[16]) {
|
|
this->set_address(
|
|
(uint32_t) htonl((address[ 0] << 24UL) | (address[ 1] << 16UL) | (address[ 2] << 8UL) | address[ 3]),
|
|
(uint32_t) htonl((address[ 4] << 24UL) | (address[ 5] << 16UL) | (address[ 6] << 8UL) | address[ 7]),
|
|
(uint32_t) htonl((address[ 8] << 24UL) | (address[ 8] << 16UL) | (address[10] << 8UL) | address[11]),
|
|
(uint32_t) htonl((address[12] << 24UL) | (address[13] << 16UL) | (address[14] << 8UL) | address[15])
|
|
);
|
|
}
|
|
|
|
private:
|
|
uint32_t address[4]{0, 0, 0, 0};
|
|
);
|
|
|
|
define_builder(TXT, base,
|
|
inline void set_text(const std::string& text) { this->_text = text; }
|
|
private:
|
|
std::string _text;
|
|
);
|
|
|
|
#undef define_builder
|
|
}
|
|
}
|
|
|
|
class DNSBuilder {
|
|
public:
|
|
builder::DNSHeader& header() { return this->_header; }
|
|
|
|
builder::DNSQuery& push_query() {
|
|
this->_queries.emplace_back();
|
|
return this->_queries.back();
|
|
}
|
|
|
|
builder::DNSQuery& query(size_t index) {
|
|
if(this->_queries.size() <= index)
|
|
std::__throw_out_of_range("index out of range");
|
|
|
|
return this->_queries[index];
|
|
}
|
|
|
|
builder::DNSResourceRecords& push_answer(const std::string& domain) {
|
|
auto& entry = this->_answers.emplace_back();
|
|
entry.name = domain;
|
|
return this->_answers.back();
|
|
}
|
|
|
|
builder::DNSResourceRecords& answer(size_t index) {
|
|
if(this->_answers.size() <= index)
|
|
std::__throw_out_of_range("index out of range");
|
|
|
|
return this->_answers[index];
|
|
}
|
|
|
|
inline size_t answer_count() const { return this->_answers.size(); }
|
|
inline size_t query_count() const { return this->_queries.size(); }
|
|
|
|
size_t build(char* /* buffer */, size_t /* max length */, std::string& /* error */);
|
|
private:
|
|
builder::DNSHeader _header{};
|
|
|
|
std::vector<builder::DNSQuery> _queries{};
|
|
std::vector<builder::DNSResourceRecords> _answers{};
|
|
};
|
|
} |