[1396] Add TreenodeRRset::getLength() implementation

src/lib/datasrc/memory/treenode_rrset.cc

@@ -100,6 +100,19 @@ TreeNodeRRset::toText() const {
 }
 
 namespace {
+void
+sizeupName(const LabelSequence& name_labels, RdataNameAttributes,
+           size_t* length)
+{
+    *length += name_labels.getDataLength();
+}
+
+void
+sizeupData(const void*, size_t data_len, size_t* length)
+{
+    *length += data_len;
+}
+
 void
 renderName(const LabelSequence& name_labels, RdataNameAttributes attr,
            AbstractMessageRenderer* renderer)
@@ -114,6 +127,35 @@ renderData(const void* data, size_t data_len,
     renderer->writeData(data, data_len);
 }
 
+// Helper for calculating wire data length of a single (etiher main or
+// RRSIG) RRset.
+uint16_t
+getLengthHelper(size_t* rlength, size_t rr_count, uint16_t name_labels_size,
+                RdataReader& reader, bool (RdataReader::* rdata_iterate_fn)())
+{
+    uint16_t length = 0;
+
+    for (size_t i = 0; i < rr_count; ++i) {
+        size_t rrlen = 0;
+
+        rrlen += name_labels_size;
+        rrlen += 2; // TYPE field
+        rrlen += 2; // CLASS field
+        rrlen += 2; // TTL field
+        rrlen += 2; // RDLENGTH field
+
+        *rlength = 0;
+        const bool rendered = (reader.*rdata_iterate_fn)();
+        assert(rendered == true);
+
+        rrlen += *rlength;
+        assert(length + rrlen < 65536);
+        length += rrlen;
+    }
+
+    return (length);
+}
+
 // Common code logic for rendering a single (either main or RRSIG) RRset.
 size_t
 writeRRs(AbstractMessageRenderer& renderer, size_t rr_count,
@@ -149,6 +191,39 @@ writeRRs(AbstractMessageRenderer& renderer, size_t rr_count,
 }
 }
 
+uint16_t
+TreeNodeRRset::getLength() const {
+    size_t rlength = 0;
+    RdataReader reader(rrclass_, rdataset_->type, rdataset_->getDataBuf(),
+                       rdataset_->getRdataCount(), rrsig_count_,
+                       boost::bind(sizeupName, _1, _2, &rlength),
+                       boost::bind(sizeupData, _1, _2, &rlength));
+
+    // Get the owner name of the RRset in the form of LabelSequence.
+    uint8_t labels_buf[LabelSequence::MAX_SERIALIZED_LENGTH];
+    const LabelSequence name_labels = getOwnerLabels(labels_buf);
+    const uint16_t name_labels_size = name_labels.getDataLength();
+
+    // Find the length of the main (non RRSIG) RRs
+    const uint16_t rrset_length =
+        getLengthHelper(&rlength, rdataset_->getRdataCount(), name_labels_size,
+                        reader, &RdataReader::iterateRdata);
+
+    rlength = 0;
+    const bool rendered = reader.iterateRdata();
+    assert(rendered == false); // we should've reached the end
+
+    // Find the length of any RRSIGs, if we supposed to do so
+    const uint16_t rrsig_length = dnssec_ok_ ?
+        getLengthHelper(&rlength, rrsig_count_, name_labels_size,
+                        reader, &RdataReader::iterateSingleSig) : 0;
+
+    // the uint16_ts are promoted to ints during addition below, so it
+    // won't overflow a 16-bit register.
+    assert(rrset_length + rrsig_length < 65536);
+    return (rrset_length + rrsig_length);
+}
+
 unsigned int
 TreeNodeRRset::toWire(AbstractMessageRenderer& renderer) const {
     RdataReader reader(rrclass_, rdataset_->type, rdataset_->getDataBuf(),

src/lib/datasrc/memory/treenode_rrset.h

@@ -155,7 +155,7 @@ public:
         node_(node), rdataset_(rdataset),
         rrsig_count_(rdataset_->getSigRdataCount()), rrclass_(rrclass),
         dnssec_ok_(dnssec_ok), name_(NULL), realname_(new dns::Name(realname)),
-	ttl_data_(rdataset->getTTLData()), ttl_(NULL)
+        ttl_data_(rdataset->getTTLData()), ttl_(NULL)
     {}
 
     virtual ~TreeNodeRRset() {
@@ -168,6 +168,8 @@ public:
         return (rdataset_->getRdataCount());
     }
 
+    virtual uint16_t getLength() const;
+
     virtual const dns::Name& getName() const;
     virtual const dns::RRClass& getClass() const {
         return (rrclass_);

src/lib/datasrc/tests/memory/treenode_rrset_unittest.cc

@@ -333,6 +333,37 @@ checkToWireResult(OutputType& expected_output, OutputType& actual_output,
                   actual_output.getData(), actual_output.getLength());
 }
 
+TEST_F(TreeNodeRRsetTest, getLength) {
+    // A RR
+    // www.example.com = 1 + 3 + 1 + 7 + 1 + 3 + 1 = 17 octets
+    // TYPE field = 2 octets
+    // CLASS field = 2 octets
+    // TTL field = 2 octets
+    // RDLENGTH field = 2 octets
+    // A RDATA = 4 octets
+    // Total = 17 + 2 + 2 + 2 + 2 + 4 = 29 octets
+
+    // RRSIG RR
+    // www.example.com = 1 + 4 + 1 + 7 + 1 + 3 + 1 = 17 octets
+    // TYPE field = 2 octets
+    // CLASS field = 2 octets
+    // TTL field = 2 octets
+    // RDLENGTH field = 2 octets
+    // RRSIG RDATA = 18 + [1 + 7 + 1 + 3 + 1 (example.com)] + 3 (base64
+    //               decode of FAKE) octets
+    // Total = 17 + 2 + 2 + 2 + 2 + 34 = 59 octets
+
+    // 1. with RRSIG, DNSSEC not OK
+    // ` 2 A RRs + 0 RRSIG RRs
+    const TreeNodeRRset rrset1(rrclass_, www_node_, a_rdataset_, false);
+    EXPECT_EQ(29 + 29, rrset1.getLength());
+
+    // 2. with RRSIG, DNSSEC OK
+    // ` 2 A RRs + 1 RRSIG RR
+    const TreeNodeRRset rrset2(rrclass_, www_node_, a_rdataset_, true);
+    EXPECT_EQ(29 + 29 + 59, rrset2.getLength());
+}
+
 TEST_F(TreeNodeRRsetTest, toWire) {
     MessageRenderer expected_renderer, actual_renderer;
     OutputBuffer expected_buffer(0), actual_buffer(0);