WebDNS/util/src/builder.cpp

122 lines
2.7 KiB
C++

#include <cstring>
#include "teadns/builder.h"
using namespace ts::dns;
using namespace ts::dns::builder;
inline bool write_dn(char *&buffer, size_t& max_size, std::string dn) {
size_t index{0};
size_t length;
do {
auto next = dn.find('.', index);
length = next == -1 ? dn.length() - index : next - index;
if(max_size + 1 < length)
return false;
*buffer = length;
memcpy(buffer + 1, dn.data() + index, length);
buffer += 1 + length;
max_size -= 1 + length;
index = next;
} while(++index);
if(max_size < 1)
return false;
*buffer = 0;
buffer++;
max_size--;
return true;
}
size_t DNSBuilder::build(char *buffer, size_t max_size, std::string &error) {
size_t begin = max_size;
if(max_size < this->_header.buffer_size()) {
error = "buffer too short for header";
return 0;
}
this->_header.set_query_count(this->_queries.size());
this->_header.set_answer_count(this->_answers.size());
this->_header.set_authority_count(0);
this->_header.set_additional_count(0);
memcpy(buffer, this->_header.buffer(), this->_header.buffer_size());
max_size -= this->_header.buffer_size();
buffer += this->_header.buffer_size();
for(auto& query : this->_queries) {
if(!write_dn(buffer, max_size, query.name)) {
error = "buffer too short for query dn";
return 0;
}
if(max_size < 4) {
error = "buffer too short for query data";
return 0;
}
*(uint16_t*) buffer = htons(query.type);
buffer += 2;
*(uint16_t*) buffer = htons(query.klass);
buffer += 2;
max_size -= 4;
}
for(auto& answer : this->_answers) {
if(!answer.build(buffer, max_size, error)) {
error = "failed to build answer: " + error;
return 0;
}
}
return begin - max_size;
}
bool DNSResourceRecords::build(char *&buffer, size_t &max_size, std::string &error) {
if(!write_dn(buffer, max_size, this->name)) {
error = "failed to write dn";
return false;
}
if(max_size < 10) {
error = "buffer too small";
return false;
}
*(uint16_t*) buffer = htons(this->type);
max_size -= 2; buffer += 2;
*(uint16_t*) buffer = htons(this->klass);
max_size -= 2; buffer += 2;
*(uint32_t*) buffer = htonl(this->ttl);
max_size -= 4; buffer += 4;
auto& length = *(uint16_t*) buffer;
max_size -= 2; buffer += 2;
if(this->payload_builder) {
auto start = max_size;
if(!this->payload_builder->build(buffer, max_size, error)) {
error = "failed to write payload: " + error;
return false;
}
length = htons(start - max_size);
} else {
length = 0;
}
return true;
}
bool rrbuilder::A::build(char *&buffer, size_t &max_size, std::string &error) {
if(max_size < 4) {
error = "buffer too small";
return false;
}
memcpy(buffer, &this->address, 4);
buffer += 4;
max_size -= 4;
return true;
}