1 files changed, 21 insertions, 6 deletions

diff --git a/src/packet/packet.c b/src/packet/packet.c
index 1d96e38..07bfccc 100644
--- a/src/packet/packet.c
+++ b/src/packet/packet.c
@@ -1,34 +1,49 @@
+#include <stdbool.h>
+#include <stdio.h>
+#include <string.h>
+#include <stdlib.h>
+
 #include "packet.h"
 #include "buffer.h"
 #include "header.h"
 #include "question.h"
 #include "record.h"
 
-#include <stdio.h>
-#include <string.h>
-#include <stdlib.h>
-
 void read_packet(PacketBuffer* buffer, Packet* packet) {
     read_header(buffer, &packet->header);
     
     packet->questions = malloc(sizeof(Question) * packet->header.questions);
     for(uint16_t i = 0; i < packet->header.questions; i++) {
-        read_question(buffer, &packet->questions[i]);
+        if (!read_question(buffer, &packet->questions[i])) {
+            i--;
+            packet->header.questions--;
+        }
     }
 
     packet->answers = malloc(sizeof(Record) * packet->header.answers);
     for(uint16_t i = 0; i < packet->header.answers; i++) {
-        read_record(buffer, &packet->answers[i]);
+        if (!read_record(buffer, &packet->answers[i])) {
+            i--;
+            packet->header.answers--;
+        }
     }
 
     packet->authorities = malloc(sizeof(Record) * packet->header.authoritative_entries);
     for(uint16_t i = 0; i < packet->header.authoritative_entries; i++) {
         read_record(buffer, &packet->authorities[i]);
+        if (!read_record(buffer, &packet->authorities[i])) {
+            i--;
+            packet->header.authoritative_entries--;
+        }
     }
 
     packet->resources = malloc(sizeof(Record) * packet->header.resource_entries);
     for(uint16_t i = 0; i < packet->header.resource_entries; i++) {
         read_record(buffer, &packet->resources[i]);
+        if (!read_record(buffer, &packet->resources[i])) {
+            i--;
+            packet->header.resource_entries--;
+        }
     }
 }