1 files changed, 19 insertions, 10 deletions

diff --git a/src/server/resolver.c b/src/server/resolver.c
index e05f365..a1fa82a 100644
--- a/src/server/resolver.c
+++ b/src/server/resolver.c
@@ -43,9 +43,9 @@ static bool lookup(
     return true;
 }
 
-static bool search(Question* question, Packet* result, BindingType type) {
+static bool search(Question* question, Packet* result, BindingType type, const RecordMap* map) {
     IpAddr addr;
-    char ip[4] = {1, 1, 1, 1};
+    uint8_t ip[4] = {1, 1, 1, 1};
     create_ip_addr(ip, &addr);
 
     uint16_t port = 53;
@@ -53,6 +53,10 @@ static bool search(Question* question, Packet* result, BindingType type) {
     create_socket_addr(port, addr, &saddr);
 
     while(1) {
+        if (record_map_get(map, question, result)) {
+            return true;
+        } 
+
         if (!lookup(question, result, type, saddr)) {
             return false;
         }
@@ -75,7 +79,7 @@ static bool search(Question* question, Packet* result, BindingType type) {
         }
 
         Packet recurse;
-        if (!search(&new_question, &recurse, type)) {
+        if (!search(&new_question, &recurse, type, map)) {
             return false;
         }
 
@@ -104,7 +108,7 @@ static void push_questions(Question* from, uint8_t from_len, Question** to, uint
     memcpy(*to + to_len, from, from_len * sizeof(Question));
 }
 
-void handle_query(Packet* request, Packet* response, BindingType type) {
+void handle_query(const Packet* request, Packet* response, BindingType type, const RecordMap* map) {
     memset(response, 0, sizeof(Packet));
     response->header.id = request->header.id;
     response->header.opcode = request->header.opcode;
@@ -121,7 +125,7 @@ void handle_query(Packet* request, Packet* response, BindingType type) {
         Packet result;
         memset(&result, 0, sizeof(Packet));
         result.header.id = response->header.id;
-        if (!search(&request->questions[i], &result, type)) {
+        if (!search(&request->questions[i], &result, type, map)) {
             response->header.response = SERVFAIL;
             break;
         }
@@ -158,9 +162,14 @@ void handle_query(Packet* request, Packet* response, BindingType type) {
         );        
         response->header.resource_entries += result.header.resource_entries;
 
-        free(result.questions);
-        free(result.answers);
-        free(result.authorities);
-        free(result.resources);
+        if (result.header.z == false) {
+            // Not from cache
+            free(result.questions);
+            free(result.answers);
+            free(result.authorities);
+            free(result.resources);
+        } else {
+            response->header.z = true;
+        }
     }
-}
\ No newline at end of file
+}