From 4951920f673607950b0bf77c9268b09af50307b8 Mon Sep 17 00:00:00 2001
From: Mylloon <kennel.anri@tutanota.com>
Date: Wed, 14 Jun 2023 18:41:43 +0200
Subject: [PATCH] use of a reader so we always store the full data content when
 parsing response, using a pointer logic

---
 bin/main.ml     |  6 +++---
 lib/response.ml | 23 +++++++++++++----------
 lib/types.ml    |  5 +++++
 3 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/bin/main.ml b/bin/main.ml
index c7ba09e..aa85030 100644
--- a/bin/main.ml
+++ b/bin/main.ml
@@ -3,9 +3,9 @@ open Dnstoy
 let () =
   let response = Network.send_request "8.8.8.8" "www.example.com" in
   print_endline (Utils.get_bytecode response);
-  let response', dns_header = Response.parse_header response in
-  let response'', dns_question = Response.parse_question response' in
-  let dns_record = Response.parse_record response'' in
+  let reader, dns_header = Response.parse_header { data = response; pointer = 0 } in
+  let reader', dns_question = Response.parse_question reader in
+  let dns_record = Response.parse_record reader' in
   Debug.print_dns_header dns_header;
   Debug.print_dns_question dns_question;
   Debug.print_dns_record dns_record
diff --git a/lib/response.ml b/lib/response.ml
index 04b7890..9fa0dbf 100644
--- a/lib/response.ml
+++ b/lib/response.ml
@@ -3,36 +3,39 @@ open Utils
 
 let parse_header reader =
   let max_size = 6 in
-  match List.init max_size (fun offset -> unpack_short_be reader (offset * 2)) with
+  match List.init max_size (fun offset -> unpack_short_be reader.data (offset * 2)) with
   | [ id; flags; num_questions; num_answers; num_authorities; num_additionals ] ->
-    ( bytes_forward reader (max_size * 2)
+    ( { reader with pointer = reader.pointer + (max_size * 2) }
     , { id; flags; num_questions; num_answers; num_authorities; num_additionals } )
   | _ -> failwith "Invalid number of fields"
 ;;
 
 let rec parse_question reader =
-  let reader', name_b = decode_name_simple reader in
-  let name = String.to_bytes name_b in
+  let name, offset_name = decode_name_simple reader in
+  let data = bytes_forward reader.data (reader.pointer + offset_name) in
   let max_size = 2 in
-  match List.init max_size (fun offset -> unpack_short_be reader' (offset * 2)) with
-  | [ type_; class_ ] -> bytes_forward reader' (max_size * 2), { name; type_; class_ }
+  match List.init max_size (fun offset -> unpack_short_be data (offset * 2)) with
+  | [ type_; class_ ] ->
+    ( { reader with pointer = reader.pointer + offset_name + (max_size * 2) }
+    , { name; type_; class_ } )
   | _ -> failwith "Invalid number of fields"
 
 and decode_name_simple reader =
   let rec read_parts parts pos =
-    let length = int_of_char (Bytes.get reader pos) in
+    let length = int_of_char (Bytes.get reader.data (reader.pointer + pos)) in
     if length = 0
     then 1, parts
     else (
-      let part = Bytes.sub_string reader (pos + 1) length in
+      let part = Bytes.sub_string reader.data (reader.pointer + pos + 1) length in
       let last_length, parts' = read_parts (parts @ [ part ]) (pos + 1 + length) in
       last_length + length + 1, parts')
   in
   let offset, parts = read_parts [] 0 in
-  bytes_forward reader offset, String.concat "." parts
+  String.to_bytes (String.concat "." parts), offset
 ;;
 
 let parse_record _reader =
-  let name, type_, class_, ttl, data = Bytes.empty, 0, 0, 0, Bytes.empty in
+  let name, _offset_name = Bytes.empty, 0 (* decode_name_simple reader *) in
+  let type_, class_, ttl, data = 0, 0, 0, Bytes.empty in
   { name; type_; class_; ttl; data }
 ;;
diff --git a/lib/types.ml b/lib/types.ml
index de0d1d0..521b6ff 100644
--- a/lib/types.ml
+++ b/lib/types.ml
@@ -20,3 +20,8 @@ type dns_record =
   ; ttl : int
   ; data : bytes
   }
+
+type reader =
+  { data : bytes
+  ; pointer : int
+  }