From 0e0917f4fef33f35ec3152825cff29541b367161 Mon Sep 17 00:00:00 2001
From: Hannes Mehnert <hannes@mehnert.org>
Date: Fri, 11 Nov 2022 12:07:06 +0100
Subject: [PATCH] DNS: start task reading Lwt_mvar and distributing DNS replies
 to clients

Before, a DNS request was sent and the first thing appearing in the Lwt_mvar
was taken as reply. The issue with this was two-fold:
- it could be a reply for a different request
- there could be DNS replies being sent to the uplink stack leading to
  Lwt_mvar.put being called, which blocks if there is already a value in the
  mvar.

No, the separate task is a loop reading the mvar, using a Lwt_condition to
signal the receive of that ID (potentially discarding if there's no client
waiting). The DNS query registers itself (using the ID) in the map with a
Lwt_condition, and waits to be notified (or a timeout occurs).
---
 my_dns.ml | 31 ++++++++++++++++++++++++++-----
 1 file changed, 26 insertions(+), 5 deletions(-)

diff --git a/my_dns.ml b/my_dns.ml
index 9f3c877..372c29a 100644
--- a/my_dns.ml
+++ b/my_dns.ml
@@ -5,11 +5,14 @@ module Transport (R : Mirage_random.S) (C : Mirage_clock.MCLOCK) (Time : Mirage_
   type io_addr = Ipaddr.V4.t * int
   type stack = Router.t * (src_port:int -> dst:Ipaddr.V4.t -> dst_port:int -> Cstruct.t -> (unit, [ `Msg of string ]) result Lwt.t) * (Udp_packet.t * Cstruct.t) Lwt_mvar.t
 
+  module IM = Map.Make(Int)
+
   type t = {
     protocol : Dns.proto ;
     nameserver : io_addr ;
     stack : stack ;
     timeout_ns : int64 ;
+    mutable requests : Cstruct.t Lwt_condition.t IM.t ;
   }
   type context = t
 
@@ -17,12 +20,26 @@ module Transport (R : Mirage_random.S) (C : Mirage_clock.MCLOCK) (Time : Mirage_
   let rng = R.generate ?g:None
   let clock = C.elapsed_ns
 
+  let rec read t =
+    let _, _, answer = t.stack in
+    Lwt_mvar.take answer >>= fun (_, data) ->
+    if Cstruct.length data > 2 then begin
+      match IM.find_opt (Cstruct.BE.get_uint16 data 0) t.requests with
+      | Some cond -> Lwt_condition.broadcast cond data
+      | None -> ()
+    end;
+    read t
+
   let create ?nameservers ~timeout stack =
     let protocol, nameserver = match nameservers with
       | None | Some (_, []) -> invalid_arg "no nameserver found"
       | Some (proto, ns :: _) -> proto, ns
     in
-    { protocol ; nameserver ; stack ; timeout_ns = timeout }
+    let t =
+      { protocol ; nameserver ; stack ; timeout_ns = timeout ; requests = IM.empty }
+    in
+    Lwt.async (fun () -> read t);
+    t
 
   let with_timeout timeout_ns f =
     let timeout = Time.sleep_ns timeout_ns >|= fun () -> Error (`Msg "DNS request timeout") in
@@ -33,14 +50,18 @@ module Transport (R : Mirage_random.S) (C : Mirage_clock.MCLOCK) (Time : Mirage_
   let send_recv (ctx : context) buf : (Cstruct.t, [> `Msg of string ]) result Lwt.t =
     let open Router in
     let dst, dst_port = ctx.nameserver in
-    let router, send_udp, answer = ctx.stack in
+    let router, send_udp, _ = ctx.stack in
     let src_port, evict =
       My_nat.free_udp_port router.nat ~src:router.uplink#my_ip ~dst ~dst_port:53
     in
+    let id = Cstruct.BE.get_uint16 buf 0 in
     with_timeout ctx.timeout_ns
-      ((send_udp ~src_port ~dst ~dst_port buf >|= Rresult.R.open_error_msg) >>= function
-        | Ok () -> (Lwt_mvar.take answer >|= fun (_, dns_response) -> Ok dns_response)
-        | Error _ as e -> Lwt.return e) >|= fun result ->
+      (let cond = Lwt_condition.create () in
+       ctx.requests <- IM.add id cond ctx.requests;
+       (send_udp ~src_port ~dst ~dst_port buf >|= Rresult.R.open_error_msg) >>= function
+       | Ok () -> Lwt_condition.wait cond >|= fun dns_response -> Ok dns_response
+       | Error _ as e -> Lwt.return e) >|= fun result ->
+    ctx.requests <- IM.remove id ctx.requests;
     evict ();
     result