From ce7585ac6cad8211a7e5af012da70ce221060aae Mon Sep 17 00:00:00 2001
From: Mylloon <kennel.anri@tutanota.com>
Date: Mon, 5 Jun 2023 21:44:44 +0200
Subject: [PATCH] Add query building

---
 lib/query.ml | 20 ++++++++++++++++++++
 test/test.ml |  7 ++++++-
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/lib/query.ml b/lib/query.ml
index 3af6ee1..d6218a3 100644
--- a/lib/query.ml
+++ b/lib/query.ml
@@ -32,3 +32,23 @@ let encode_dns_name domain_name =
   in
   Bytes.cat (Bytes.concat Bytes.empty encoded_parts) (Bytes.of_string "\x00")
 ;;
+
+let build ?(id = None) domain_name record_type =
+  let class_IN = 1 in
+  let header =
+    { id =
+        (match id with
+         | Some i -> i
+         | None -> Random.int 65535)
+    ; flags = 1 lsl 8
+    ; num_questions = 1
+    ; num_answers = 0
+    ; num_authorities = 0
+    ; num_additionals = 0
+    }
+  in
+  let question =
+    { name = encode_dns_name domain_name; type_ = record_type; class_ = class_IN }
+  in
+  Bytes.cat (header_to_bytes header) (question_to_bytes question)
+;;
diff --git a/test/test.ml b/test/test.ml
index 7326d50..b4d571e 100644
--- a/test/test.ml
+++ b/test/test.ml
@@ -17,5 +17,10 @@ let () =
   (* Dns name *)
   let dns_name = Query.encode_dns_name "google.com" in
   assert ("\\x06google\\x03com\\x00" = Utils.get_bytecode dns_name);
-  assert (6 = Char.code (Bytes.get dns_name 0))
+  assert (6 = Char.code (Bytes.get dns_name 0));
+  (* Query *)
+  let type_A = 1 in
+  assert (
+    "\\x15\\xB3\\x01\\x00\\x00\\x01\\x00\\x00\\x00\\x00\\x00\\x00\\x07example\\x03com\\x00\\x00\\x01\\x00\\x01"
+    = Utils.get_bytecode (Query.build ~id:(Some 5555) "example.com" type_A))
 ;;