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