1 files changed, 20 insertions, 25 deletions

diff --git a/src/libknot/packet/rrset-wire.c b/src/libknot/packet/rrset-wire.c
index a822cfe..1d4f78e 100644
--- a/src/libknot/packet/rrset-wire.c
+++ b/src/libknot/packet/rrset-wire.c
@@ -73,6 +73,7 @@ static bool dname_equal_wire(const knot_dname_t *d1, const knot_dname_t *d2,
 {
 	assert(d1);
 	assert(d2);
+	assert(wire);
 
 	d2 = knot_wire_seek_label(d2, wire);
 
@@ -80,7 +81,7 @@ static bool dname_equal_wire(const knot_dname_t *d1, const knot_dname_t *d2,
 		if (!label_is_equal(d1, d2)) {
 			return false;
 		}
-		d1 = knot_wire_next_label(d1, NULL);
+		d1 = knot_dname_next_label(d1);
 		d2 = knot_wire_next_label(d2, wire);
 	}
 
@@ -170,7 +171,7 @@ static int write_rdata_naptr_header(const uint8_t **src, size_t *src_avail,
 		written += (len); \
 	}
 
-#define CHECK_NEXT_LABEL(res) \
+#define CHECK_WIRE_NEXT_LABEL(res) \
 	if (res == NULL) { return KNOT_EINVAL; }
 
 /*!
@@ -201,7 +202,7 @@ static int compr_put_dname(const knot_dname_t *dname, uint8_t *dst, uint16_t max
 	int suffix_labels = compr->suffix.labels;
 	while (suffix_labels > name_labels) {
 		suffix = knot_wire_next_label(suffix, compr->wire);
-		CHECK_NEXT_LABEL(suffix);
+		CHECK_WIRE_NEXT_LABEL(suffix);
 		--suffix_labels;
 	}
 
@@ -210,8 +211,7 @@ static int compr_put_dname(const knot_dname_t *dname, uint8_t *dst, uint16_t max
 	uint16_t written = 0;
 	while (name_labels > suffix_labels) {
 		WRITE_LABEL(dst, written, dname, max, (*dname + 1));
-		dname = knot_wire_next_label(dname, NULL);
-		CHECK_NEXT_LABEL(dname);
+		dname = knot_dname_next_label(dname);
 		--name_labels;
 	}
 
@@ -221,10 +221,9 @@ static int compr_put_dname(const knot_dname_t *dname, uint8_t *dst, uint16_t max
 	const knot_dname_t *compr_ptr = suffix;
 	while (dname[0] != '\0') {
 		// Next labels.
-		const knot_dname_t *next_dname = knot_wire_next_label(dname, NULL);
-		CHECK_NEXT_LABEL(next_dname);
+		const knot_dname_t *next_dname = knot_dname_next_label(dname);
 		const knot_dname_t *next_suffix = knot_wire_next_label(suffix, compr->wire);
-		CHECK_NEXT_LABEL(next_suffix);
+		CHECK_WIRE_NEXT_LABEL(next_suffix);
 
 		// Two labels match, extend suffix length.
 		if (!label_is_equal(dname, suffix)) {
@@ -324,11 +323,11 @@ static int write_owner(const knot_rrset_t *rrset, uint8_t **dst, size_t *dst_ava
 	return KNOT_EOK;
 }
 
-static int write_fixed_header(const knot_rrset_t *rrset, uint16_t rrset_index,
+static int write_fixed_header(const knot_rrset_t *rrset, const knot_rdata_t *rdata,
                               uint8_t **dst, size_t *dst_avail, uint16_t flags)
 {
 	assert(rrset);
-	assert(rrset_index < rrset->rrs.count);
+	assert(rdata);
 	assert(dst && *dst);
 	assert(dst_avail);
 
@@ -339,10 +338,8 @@ static int write_fixed_header(const knot_rrset_t *rrset, uint16_t rrset_index,
 	wire_ctx_write_u16(&write, rrset->rclass);
 
 	if ((flags & KNOT_PF_ORIGTTL) && rrset->type == KNOT_RRTYPE_RRSIG) {
-		const knot_rdata_t *rdata = knot_rdataset_at(&rrset->rrs, rrset_index);
 		wire_ctx_write_u32(&write, knot_rrsig_original_ttl(rdata));
 	} else if ((flags & KNOT_PF_SOAMINTTL) && rrset->type == KNOT_RRTYPE_SOA) {
-		const knot_rdata_t *rdata = knot_rdataset_at(&rrset->rrs, rrset_index);
 		wire_ctx_write_u32(&write, MIN(knot_soa_minimum(rdata), rrset->ttl));
 	} else {
 		wire_ctx_write_u32(&write, rrset->ttl);
@@ -433,15 +430,15 @@ static int rdata_traverse_write(const uint8_t **src, size_t *src_avail,
 }
 
 static int write_rdata(const knot_rrset_t *rrset, uint16_t rrset_index,
-                       uint8_t **dst, size_t *dst_avail, knot_compr_t *compr)
+                       const knot_rdata_t *rdata, uint8_t **dst, size_t *dst_avail,
+                       knot_compr_t *compr)
 {
 	assert(rrset);
 	assert(rrset_index < rrset->rrs.count);
+	assert(rdata);
 	assert(dst && *dst);
 	assert(dst_avail);
 
-	const knot_rdata_t *rdata = knot_rdataset_at(&rrset->rrs, rrset_index);
-
 	// Reserve space for RDLENGTH.
 	if (sizeof(uint16_t) > *dst_avail) {
 		return KNOT_ESPACE;
@@ -478,7 +475,8 @@ static int write_rdata(const knot_rrset_t *rrset, uint16_t rrset_index,
 	return KNOT_EOK;
 }
 
-static int write_rr(const knot_rrset_t *rrset, uint16_t rrset_index, uint8_t **dst,
+static int write_rr(const knot_rrset_t *rrset, uint16_t rrset_index,
+                    const knot_rdata_t *rdata, uint8_t **dst,
                     size_t *dst_avail, knot_compr_t *compr, uint16_t flags)
 {
 	int ret = write_owner(rrset, dst, dst_avail, compr);
@@ -486,17 +484,17 @@ static int write_rr(const knot_rrset_t *rrset, uint16_t rrset_index, uint8_t **d
 		return ret;
 	}
 
-	ret = write_fixed_header(rrset, rrset_index, dst, dst_avail, flags);
+	ret = write_fixed_header(rrset, rdata, dst, dst_avail, flags);
 	if (ret != KNOT_EOK) {
 		return ret;
 	}
 
-	return write_rdata(rrset, rrset_index, dst, dst_avail, compr);
+	return write_rdata(rrset, rrset_index, rdata, dst, dst_avail, compr);
 }
 
 _public_
 int knot_rrset_to_wire_extra(const knot_rrset_t *rrset, uint8_t *wire,
-                             uint16_t max_size, uint16_t rotate,
+                             uint32_t max_size, uint16_t rotate,
                              knot_compr_t *compr, uint16_t flags)
 {
 	if (rrset == NULL || wire == NULL) {
@@ -512,15 +510,12 @@ int knot_rrset_to_wire_extra(const knot_rrset_t *rrset, uint8_t *wire,
 	uint8_t *write = wire;
 	size_t capacity = max_size;
 
-	// FIXME remove this and make the max_size parameter uint32_t in next major libknot release!
-	if ((flags & KNOT_PF_BUFENOUGH)) {
-		capacity = SIZE_MAX;
-	}
-
 	uint16_t count = rrset->rrs.count;
+	knot_rdata_t *rdata = rotate > 1 ? knot_rdataset_at(&rrset->rrs, rotate - 1) : rrset->rrs.rdata;
 	for (int i = rotate; i < count + rotate; i++) {
 		uint16_t pos = (i < count) ? i : (i - count);
-		int ret = write_rr(rrset, pos, &write, &capacity, compr, flags);
+		rdata = pos ? knot_rdataset_next(rdata) : rrset->rrs.rdata;
+		int ret = write_rr(rrset, pos, rdata, &write, &capacity, compr, flags);
 		if (ret != KNOT_EOK) {
 			return ret;
 		}