Fixed some bugs in DNS. Added some helper functions.

include/dns.h

@@ -144,6 +144,13 @@ namespace Tins {
          */
         DNS(const uint8_t *buffer, uint32_t total_sz);
         
+        /**
+         * \brief Copy constructor.
+         * 
+         * \param other The DNS object to be copied.
+         */
+        DNS(const DNS &other);
+        
         /**
          * \brief Destructor.
          */
@@ -368,6 +375,13 @@ namespace Tins {
          */
         void add_query(const std::string &name, QueryType type, QueryClass qclass);
         
+        /**
+         * \brief Add a query to perform.
+         * 
+         * \param query The query to be added.
+         */
+        void add_query(const Query &query);
+        
         /**
          * \brief Add a query response.
          * 
@@ -379,6 +393,31 @@ namespace Tins {
          */
         void add_answer(const std::string &name, QueryType type, QueryClass qclass,
                         uint32_t ttl, uint32_t ip);
+                        
+        /**
+         * \brief Add a query response.
+         * 
+         * \param name The resolved name.
+         * \param type The type of this answer.
+         * \param qclass The class of this answer.
+         * \param ttl The time-to-live of this answer.
+         * \param dname The domain of the resolved name.
+         */
+        void add_answer(const std::string &name, QueryType type, QueryClass qclass,
+                        uint32_t ttl, const std::string &dname);
+                        
+        /**
+         * \brief Add a query response.
+         * 
+         * \param name The resolved name.
+         * \param type The type of this answer.
+         * \param qclass The class of this answer.
+         * \param ttl The time-to-live of this answer.
+         * \param data The data of this option.
+         * \param sz The size of the data.
+         */
+        void add_answer(const std::string &name, QueryType type, QueryClass qclass,
+                        uint32_t ttl, const uint8_t *data, uint32_t sz);
         
         /**
          * \brief Add an authority record.
@@ -387,10 +426,11 @@ namespace Tins {
          * \param type The type of this record.
          * \param qclass The class of this record.
          * \param ttl The time-to-live of this record.
-         * \param ip The ip address of the resolved name.
+         * \param data The data of this option.
+         * \param sz The size of the data.
          */
         void add_authority(const std::string &name, QueryType type, QueryClass qclass,
-                        uint32_t ttl, uint32_t ip);
+                        uint32_t ttl, const uint8_t *data, uint32_t sz);
         
         /**
          * \brief Add an additional record.
@@ -418,6 +458,13 @@ namespace Tins {
          * record.
          */
         std::list<Resource> dns_answers();
+        
+        /**
+         * \brief Clones this PDU.
+         *
+         * \sa PDU::clone_pdu
+         */
+        PDU *clone_pdu() const;
     private:
         struct dnshdr {
             uint16_t id;
@@ -445,22 +492,17 @@ namespace Tins {
             uint8_t *data;
             uint16_t data_sz;
             
-            ResourceRecord(uint8_t *d = 0, uint16_t len = 0) : data_sz(len) {
+            ResourceRecord(const uint8_t *d = 0, uint16_t len = 0) : data_sz(len) {
-                if(d)
+                if(d && len) {
+                    data = new uint8_t[len];
                     std::memcpy(data, d, data_sz);
+                }
             }
             
             virtual ~ResourceRecord() {}
-            uint32_t write(uint8_t *buffer) const {
+            virtual ResourceRecord *clone() const = 0;
-                uint32_t sz(do_write(buffer));
+            void copy_fields(ResourceRecord *other) const;
-                buffer += sz;
+            uint32_t write(uint8_t *buffer) const;
-                std::memcpy(buffer, &info, sizeof(info));
-                buffer += sizeof(info);
-                *((uint16_t*)buffer) = Utils::net_to_host_s(data_sz);
-                buffer += sizeof(uint16_t);
-                std::memcpy(buffer, data, data_sz);
-                return sz + sizeof(info) + sizeof(uint16_t) + data_sz;
-            }
             virtual uint32_t do_write(uint8_t *buffer) const = 0;
             virtual uint32_t size() const = 0;
             virtual bool matches(const std::string &dname) { return false; }
@@ -476,19 +518,20 @@ namespace Tins {
         struct OffsetedResourceRecord : public ResourceRecord {
             uint16_t offset;
             
-            OffsetedResourceRecord(uint16_t off, uint8_t *data = 0, uint16_t len = 0) : ResourceRecord(data,len), offset(off | 0xc0) {}
+            OffsetedResourceRecord(uint16_t off, const uint8_t *data = 0, uint16_t len = 0) : ResourceRecord(data,len), offset(off | 0xc0) {}
             
             uint32_t do_write(uint8_t *buffer) const {
                 std::memcpy(buffer, &offset, sizeof(offset));
                 return sizeof(offset);
             }
             uint32_t size() const { return sizeof(ResourceRecord::info) + sizeof(offset) + data_sz + sizeof(uint16_t); }
+            ResourceRecord *clone() const;
         };
         
         struct NamedResourceRecord : public ResourceRecord {
             std::string name;
             
-            NamedResourceRecord(const std::string &nm,  uint8_t *data = 0, uint16_t len = 0) : ResourceRecord(data,len), name(nm) {}
+            NamedResourceRecord(const std::string &nm,  const uint8_t *data = 0, uint16_t len = 0) : ResourceRecord(data,len), name(nm) {}
             
             uint32_t do_write(uint8_t *buffer) const {
                 std::memcpy(buffer, name.c_str(), name.size() + 1);
@@ -504,9 +547,11 @@ namespace Tins {
             const std::string *dname_pointer() const {
                 return &name;
             }
+            ResourceRecord *clone() const;
         };
         
         typedef std::map<uint16_t, std::string> SuffixMap;
+        typedef std::map<uint16_t, uint16_t> SuffixIndices;
         
         const uint8_t *build_resource_list(std::list<ResourceRecord*> &lst, const uint8_t *ptr, uint32_t &sz, uint16_t nrecs);
         uint32_t find_domain_name(const std::string &dname);
@@ -519,17 +564,24 @@ namespace Tins {
         void compose_name(const uint8_t *ptr, uint32_t sz, std::string &out);
         void convert_resources(const std::list<ResourceRecord*> &lst, std::list<Resource> &res);
         ResourceRecord *make_record(const std::string &name, QueryType type, QueryClass qclass, uint32_t ttl, uint32_t ip);
+        ResourceRecord *make_record(const std::string &name, QueryType type, QueryClass qclass, uint32_t ttl, const std::string &dname);
+        ResourceRecord *make_record(const std::string &name, QueryType type, QueryClass qclass, uint32_t ttl, const uint8_t *ptr, uint32_t len);
         void add_suffix(uint32_t index, const uint8_t *data, uint32_t sz);
         uint32_t build_suffix_map(uint32_t index, const std::list<ResourceRecord*> &lst);
         uint32_t build_suffix_map(uint32_t index, const std::list<Query> &lst);
         void build_suffix_map();
+        void copy_list(const std::list<ResourceRecord*> &from, std::list<ResourceRecord*> &to) const;
+        
+        void copy_fields(const DNS *other);
         
         dnshdr dns;
         uint32_t extra_size;
         std::list<Query> queries;
         std::list<ResourceRecord*> ans, arity, addit;
         SuffixMap suffixes;
+        SuffixIndices suffix_indices;
     };
 };
 
 #endif // __DNS_H
+

src/dns.cpp

@@ -62,6 +62,11 @@ Tins::DNS::DNS(const uint8_t *buffer, uint32_t total_sz) : PDU(255), extra_size(
     build_resource_list(addit, buffer, total_sz, additional());
 }
 
+Tins::DNS::DNS(const DNS &other) : PDU(255) {
+    copy_fields(&other);
+    copy_inner_pdu(other);
+}
+
 const uint8_t *Tins::DNS::build_resource_list(list<ResourceRecord*> &lst, const uint8_t *ptr, uint32_t &sz, uint16_t nrecs) {
     const uint8_t *ptr_end(ptr + sz);
     const uint8_t *parse_start(ptr);
@@ -188,14 +193,39 @@ void Tins::DNS::add_query(const string &name, QueryType type, QueryClass qclass)
     dns.questions = Utils::net_to_host_s(queries.size());
 }
 
+void Tins::DNS::add_query(const Query &query) {
+    add_query(
+        query.name, 
+        static_cast<QueryType>(query.type), 
+        static_cast<QueryClass>(query.qclass)
+    );
+}
+
 void Tins::DNS::add_answer(const string &name, QueryType type, QueryClass qclass, uint32_t ttl, uint32_t ip) {
     ResourceRecord *res = make_record(name, type, qclass, ttl, ip);
     ans.push_back(res);
     dns.answers = Utils::net_to_host_s(ans.size());
 }
 
-void Tins::DNS::add_authority(const string &name, QueryType type, QueryClass qclass, uint32_t ttl, uint32_t ip) {
+void Tins::DNS::add_answer(const std::string &name, QueryType type, QueryClass qclass,
-    ResourceRecord *res = make_record(name, type, qclass, ttl, ip);
+                        uint32_t ttl, const std::string &dname) {
+    string new_str;
+    parse_domain_name(dname, new_str);
+    ResourceRecord *res = make_record(name, type, qclass, ttl, new_str);
+    ans.push_back(res);
+    dns.answers = Utils::net_to_host_s(ans.size());
+}
+
+void Tins::DNS::add_answer(const std::string &name, QueryType type, QueryClass qclass,
+                        uint32_t ttl, const uint8_t *data, uint32_t sz) {
+    ResourceRecord *res = make_record(name, type, qclass, ttl, data, sz);
+    ans.push_back(res);
+    dns.answers = Utils::net_to_host_s(ans.size());
+}
+
+void Tins::DNS::add_authority(const string &name, QueryType type, 
+  QueryClass qclass, uint32_t ttl, const uint8_t *data, uint32_t sz) {
+    ResourceRecord *res = make_record(name, type, qclass, ttl, data, sz);
     arity.push_back(res);
     dns.authority = Utils::net_to_host_s(arity.size());
 }
@@ -207,15 +237,23 @@ void Tins::DNS::add_additional(const string &name, QueryType type, QueryClass qc
 }
 
 Tins::DNS::ResourceRecord *Tins::DNS::make_record(const std::string &name, QueryType type, QueryClass qclass, uint32_t ttl, uint32_t ip) {
+    ip = Utils::net_to_host_l(ip);
+    return make_record(name, type, qclass, ttl, reinterpret_cast<uint8_t*>(&ip), sizeof(ip));
+}
+
+Tins::DNS::ResourceRecord *Tins::DNS::make_record(const std::string &name, QueryType type, QueryClass qclass, uint32_t ttl, const std::string &dname) {
+    return make_record(name, type, qclass, ttl, reinterpret_cast<const uint8_t*>(dname.c_str()), dname.size() + 1);
+}
+
+Tins::DNS::ResourceRecord *Tins::DNS::make_record(const std::string &name, QueryType type, QueryClass qclass, uint32_t ttl, const uint8_t *ptr, uint32_t len) {
     string nm;
     parse_domain_name(name, nm);
     uint16_t index = find_domain_name(nm);
     ResourceRecord *res;
-    ip = Utils::net_to_host_l(ip);
     if(index)
-        res = new OffsetedResourceRecord(Utils::net_to_host_s(index), (uint8_t*)&ip, sizeof(uint32_t));
+        res = new OffsetedResourceRecord(Utils::net_to_host_s(index), ptr, len);
     else
-        res = new NamedResourceRecord(nm, (uint8_t*)&ip, sizeof(uint32_t));
+        res = new NamedResourceRecord(nm, ptr, len);
     res->info.type = Utils::net_to_host_s(type);
     res->info.qclass = Utils::net_to_host_s(qclass);
     res->info.ttl = Utils::net_to_host_l(ttl);
@@ -303,10 +341,18 @@ uint8_t *Tins::DNS::serialize_list(const std::list<ResourceRecord*> &lst, uint8_
 
 void Tins::DNS::add_suffix(uint32_t index, const uint8_t *data, uint32_t sz) {
     uint32_t i(0), suff_sz(data[0]);
-    while(i + suff_sz + 1 <= sz && suff_sz) {
+    SuffixMap::iterator it;
-        i++;
+    while((i + suff_sz + 1 <= sz || (suff_sz == 0xc0 && i + 1 < sz)) && suff_sz) {
-        suffixes.insert(std::make_pair(index + i - 1, string(data + i, data + i + suff_sz)));
+        if((suff_sz & 0xc0)) {
-        i += suff_sz;
+            if((it = suffixes.find(data[i+1])) != suffixes.end())
+                suffix_indices[index + i] = data[i+1];
+            i += sizeof(uint16_t);
+        }
+        else {
+            ++i;
+            suffixes.insert(std::make_pair(index + i - 1, string(data + i, data + i + suff_sz)));
+            i += suff_sz;
+        }        
         if(i < sz)
             suff_sz = data[i];
     }
@@ -331,7 +377,6 @@ uint32_t Tins::DNS::build_suffix_map(uint32_t index, const list<ResourceRecord*>
     return index;
 }
 
-
 uint32_t Tins::DNS::build_suffix_map(uint32_t index, const list<Query> &lst) {
     for(list<Query>::const_iterator it(lst.begin()); it != lst.end(); ++it) {
         add_suffix(index, (uint8_t*)it->name.c_str(), it->name.size());
@@ -357,16 +402,24 @@ void Tins::DNS::compose_name(const uint8_t *ptr, uint32_t sz, std::string &out)
             uint16_t index = Utils::net_to_host_s(*((uint16_t*)(ptr + i)));
             index &= 0x3fff;
             SuffixMap::iterator it(suffixes.find(index));
-            assert(it != suffixes.end());
+            SuffixIndices::iterator suff_it(suffix_indices.find(index));
+            assert(it != suffixes.end() && suff_it == suffix_indices.end());
             bool first(true);
             do {
-                if(!first)
+                if(it != suffixes.end()) {
-                    out.push_back('.');
+                    if(!first)
-                first = false;
+                        out.push_back('.');
-                out += it->second;
+                    first = false;
-                index += it->second.size() + 1;
+                    out += it->second;
+                    index += it->second.size() + 1;
+                }
+                else
+                    index = suff_it->second;
                 it = suffixes.find(index);
-            } while(it != suffixes.end());
+                if(it == suffixes.end())
+                    suff_it = suffix_indices.find(index);
+                
+            } while(it != suffixes.end() || suff_it != suffix_indices.end());
             break;
         }
         else {
@@ -422,4 +475,56 @@ list<Tins::DNS::Resource> Tins::DNS::dns_answers() {
     return res;
 }
 
+Tins::PDU *Tins::DNS::clone_pdu() const {
+    DNS *new_pdu = new DNS();
+    new_pdu->copy_fields(this);
+    new_pdu->copy_inner_pdu(*this);
+    return new_pdu;
+}
 
+void Tins::DNS::copy_fields(const DNS *other) {
+    std::memcpy(&dns, &other->dns, sizeof(dns));
+    extra_size = other->extra_size;
+    queries = other->queries;
+    copy_list(other->ans, ans);
+    copy_list(other->arity, arity);
+    copy_list(other->addit, addit);
+}
+
+void Tins::DNS::copy_list(const list<ResourceRecord*> &from, list<ResourceRecord*> &to) const {
+    for(list<ResourceRecord*>::const_iterator it(from.begin()); it != from.end(); ++it) {
+        to.push_back((*it)->clone());
+    }
+}
+
+// ResourceRecord
+
+void Tins::DNS::ResourceRecord::copy_fields(ResourceRecord *other) const {
+    std::memcpy(&other->info, &info, sizeof(info));
+    other->data_sz = data_sz;
+    other->data = new uint8_t[data_sz];
+    std::memcpy(other->data, data, data_sz);
+}
+
+uint32_t Tins::DNS::ResourceRecord::write(uint8_t *buffer) const {
+    uint32_t sz(do_write(buffer));
+    buffer += sz;
+    std::memcpy(buffer, &info, sizeof(info));
+    buffer += sizeof(info);
+    *((uint16_t*)buffer) = Utils::net_to_host_s(data_sz);
+    buffer += sizeof(uint16_t);
+    std::memcpy(buffer, data, data_sz);
+    return sz + sizeof(info) + sizeof(uint16_t) + data_sz;
+}
+
+Tins::DNS::ResourceRecord *Tins::DNS::OffsetedResourceRecord::clone() const {
+    ResourceRecord *r = new OffsetedResourceRecord(offset);
+    copy_fields(r);
+    return r;
+}
+
+Tins::DNS::ResourceRecord *Tins::DNS::NamedResourceRecord::clone() const {
+    ResourceRecord *r = new NamedResourceRecord(name);
+    copy_fields(r);
+    return r;
+}