1 files changed, 16 insertions, 13 deletions

diff --git a/ffd/tlv.c b/ffd/tlv.c
index 684a0cc..0a1950d 100644
--- a/ffd/tlv.c
+++ b/ffd/tlv.c
@@ -32,14 +32,14 @@
 #include <arpa/inet.h>
 
 
-bool ffd_tlv_parse(const ffd_packet_t *packet, ffd_tlv_cb cb) {
+bool ffd_tlv_parse(const ffd_packet_t *packet, ffd_tlv_cb cb, void *arg) {
 	if (packet->version_magic != htons(FFD_VERSION_MAGIC))
 		return false;
 
-	const uint8_t *data = packet->tlv + 2, *end = ((uint8_t*)packet->tlv) + packet->len;
+	const uint8_t *data = packet->tlv + 2, *end = packet->tlv + packet->len;
 
 	while (data <= end) {
-		if (!data[-2]) {
+		if (data[-2] == TLV_PAD1) {
 			/* Pad1 */
 			data++;
 			continue;
@@ -50,24 +50,27 @@ bool ffd_tlv_parse(const ffd_packet_t *packet, ffd_tlv_cb cb) {
 			break;
 		}
 
-		cb(data[-2], data, data[-1]);
+		if (data[-2] != TLV_PADN)
+			cb(data[-2], data, data[-1], arg);
+
 		data += data[-1] + 2;
 	}
 
 	return true;
 }
 
-bool ffd_tlv_add(ffd_packet_t *packet, size_t max_len, ffd_tlv_type_t type, void *data, size_t len) {
-	if (ntohs(packet->len)+len+2 > max_len)
-		return false;
+void* ffd_tlv_add(ffd_packet_t *packet, size_t max_len, ffd_tlv_type_t type, size_t len) {
+	size_t pktlen = ntohs(packet->len);
 
-	packet->len = htons(ntohs(packet->len)+len+2);
+	if (pktlen+len+2 > max_len)
+		return NULL;
 
-	uint8_t *pkgdata = packet->tlv;
+	uint8_t *data = packet->tlv+pktlen+2;
 
-	pkgdata[0] = type;
-	pkgdata[1] = len;
-	memcpy(pkgdata+2, data, len);
+	data[-2] = type;
+	data[-1] = len;
 
-	return true;
+	packet->len = htons(pktlen+len+2);
+
+	return data;
 }