WebDNS/util/include/teadns/builder.h

#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};
			);

			#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];
			}

			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{};
	};
}