unikernel/duniverse/ocaml-dns/lwt/client/dns_client_lwt.ml

@@ -0,0 +1,368 @@
+open Lwt.Infix
+
+module IM = Map.Make(Int)
+
+let src = Logs.Src.create "dns_client_lwt" ~doc:"effectful DNS lwt layer"
+module Log = (val Logs.src_log src : Logs.LOG)
+
+module Transport : Dns_client.S
+ with type io_addr = [ `Plaintext of Ipaddr.t * int | `Tls of Tls.Config.client * Ipaddr.t * int ]
+ and type +'a io = 'a Lwt.t
+ and type stack = Happy_eyeballs_lwt.t
+= struct
+  type io_addr = [ `Plaintext of Ipaddr.t * int | `Tls of Tls.Config.client * Ipaddr.t * int ]
+  type +'a io = 'a Lwt.t
+  type stack = Happy_eyeballs_lwt.t
+  type nameservers =
+    | Static of io_addr list
+    | Resolv_conf of {
+        mutable nameservers : io_addr list;
+        mutable digest : Digest.t option
+      }
+  type t = {
+    nameservers : nameservers;
+    timeout_ns : int64 ;
+    (* TODO: avoid race, use a mvar instead of condition *)
+    mutable fd : [ `Plain of Lwt_unix.file_descr | `Tls of Tls_lwt.Unix.t ] option ;
+    mutable connected_condition : (unit, [ `Msg of string ]) result Lwt_condition.t option ;
+    mutable requests : (string * (string, [ `Msg of string ]) result Lwt_condition.t) IM.t ;
+    he : Happy_eyeballs_lwt.t ;
+  }
+  type context = t
+
+  let nameserver_ips = function
+    | Static nameservers -> nameservers
+    | Resolv_conf { nameservers; _ } -> nameservers
+
+  let read_file file =
+    try
+      let fh = open_in file in
+      try
+        let content = really_input_string fh (in_channel_length fh) in
+        close_in_noerr fh ;
+        Ok content
+      with _ ->
+        close_in_noerr fh;
+        Error (`Msg ("Error reading file: " ^ file))
+    with _ -> Error (`Msg ("Error opening file " ^ file))
+
+  let clock = Mtime_clock.elapsed_ns
+
+  let close_socket fd =
+    Lwt.catch (fun () -> Lwt_unix.close fd) (fun _ -> Lwt.return_unit)
+
+  let authenticator =
+    let authenticator_ref = ref None in
+    fun () ->
+      match !authenticator_ref with
+      | Some x -> x
+      | None -> match Ca_certs.authenticator () with
+        | Ok a -> authenticator_ref := Some a ; a
+        | Error `Msg m -> invalid_arg ("failed to load trust anchors: " ^ m)
+
+  let decode_resolv_conf data =
+    let ( let* ) = Result.bind in
+    let authenticator = authenticator () in
+    let* ns = Dns_resolvconf.parse data in
+    match
+      List.flatten
+        (List.map
+           (fun (`Nameserver ip) ->
+              match Tls.Config.client ~authenticator ~ip () with
+              | Ok tls -> [ `Tls (tls, ip, 853) ; `Plaintext (ip, 53) ]
+              | Error `Msg msg ->
+                Log.err (fun m -> m "creating TLS configuratio for %a: %s"
+                            Ipaddr.pp ip msg);
+                [ `Plaintext (ip, 53) ])
+           ns)
+    with
+    | [] -> Error (`Msg "no nameservers in resolv.conf")
+    | ns -> Ok ns
+
+  let resolv_conf () =
+    let ( let* ) = Result.bind in
+    let* data = read_file "/etc/resolv.conf" in
+    let* ns =
+      Result.map_error
+        (function `Msg msg ->
+           Log.warn (fun m -> m "error %s decoding resolv.conf %S" msg data);
+           `Msg msg)
+        (decode_resolv_conf data)
+    in
+    Ok (ns, Digest.string data)
+
+  let default_resolver () =
+    let authenticator = authenticator () in
+    let peer_name = Dns_client.default_resolver_hostname in
+    let tls_config =
+      match Tls.Config.client ~authenticator ~peer_name () with
+      | Ok cfg -> cfg
+      | Error `Msg msg -> invalid_arg msg
+    in
+    List.map (fun ip -> `Tls (tls_config, ip, 853)) Dns_client.default_resolvers
+
+  let maybe_resolv_conf t =
+    match t.nameservers with
+    | Static _ -> ()
+    | Resolv_conf resolv_conf ->
+      let needs_update =
+        match read_file "/etc/resolv.conf", resolv_conf.digest with
+        | Ok data, Some dgst ->
+          let dgst' = Digest.string data in
+          if Digest.equal dgst' dgst then
+            `No
+          else
+            `Data (data, dgst')
+        | Ok data, None ->
+          let digest = Digest.string data in
+          `Data (data, digest)
+        | Error _, None ->
+          `No
+        | Error `Msg msg, Some _ ->
+          Log.warn (fun m -> m "error reading /etc/resolv.conf: %s" msg);
+          `Default
+      in
+      match needs_update with
+      | `No -> ()
+      | `Default ->
+        resolv_conf.digest <- None;
+        resolv_conf.nameservers <- default_resolver ()
+      | `Data (data, dgst) ->
+        match decode_resolv_conf data with
+        | Ok ns ->
+          resolv_conf.digest <- Some dgst;
+          resolv_conf.nameservers <- ns
+        | Error `Msg msg ->
+          Log.warn (fun m -> m "error %s decoding resolv.conf: %S" msg data);
+          resolv_conf.digest <- None;
+          resolv_conf.nameservers <- default_resolver ()
+
+  let create ?nameservers ~timeout happy_eyeballs =
+    let nameservers =
+      match nameservers with
+      | Some (`Udp, _) -> invalid_arg "UDP is not supported"
+      | Some (`Tcp, ns) -> Static ns
+      | None ->
+        match resolv_conf () with
+        | Error _ -> Resolv_conf { nameservers = default_resolver (); digest = None }
+        | Ok (ips, digest) -> Resolv_conf { nameservers = ips; digest = Some digest }
+    in
+    {
+      nameservers ;
+      timeout_ns = timeout ;
+      fd = None ;
+      connected_condition = None ;
+      requests = IM.empty ;
+      he = happy_eyeballs ;
+    }
+
+  let nameservers { nameservers; _ } = `Tcp, nameserver_ips nameservers
+
+  let rng = Mirage_crypto_rng.generate ?g:None
+
+  let with_timeout timeout f =
+    let timeout =
+      Lwt_unix.sleep (Duration.to_f timeout) >|= fun () ->
+      Error (`Msg "DNS request timeout")
+    in
+    Lwt.pick [ f ; timeout ]
+
+  let close _ = Lwt.return_unit
+
+  let send_query fd tx =
+    Lwt.catch (fun () ->
+      match fd with
+      | `Plain fd ->
+        Lwt_unix.send fd (Bytes.unsafe_of_string tx) 0
+          (String.length tx) [] >>= fun res ->
+        if res <> String.length tx then
+          Lwt_result.fail (`Msg ("oops" ^ (string_of_int res)))
+        else
+          Lwt_result.return ()
+      | `Tls fd ->
+        Lwt_result.ok (Tls_lwt.Unix.write fd tx))
+      (fun e -> Lwt.return (Error (`Msg (Printexc.to_string e))))
+
+  let send_recv (t : context) tx =
+    if String.length tx > 4 then
+      match t.fd with
+      | None -> Lwt.return (Error (`Msg "no connection to the nameserver established"))
+      | Some fd ->
+        let id = String.get_uint16_be tx 2 in
+        with_timeout t.timeout_ns
+          (let open Lwt_result.Infix in
+           send_query fd tx >>= fun () ->
+           let cond = Lwt_condition.create () in
+           t.requests <- IM.add id (tx, cond) t.requests;
+           let open Lwt.Infix in
+           Lwt_condition.wait cond >|= fun data ->
+           match data with Ok _ | Error `Msg _ as r -> r) >|= fun r ->
+        t.requests <- IM.remove id t.requests;
+        r
+    else
+      Lwt.return (Error (`Msg "invalid DNS packet (data length <= 4)"))
+
+  let bind = Lwt.bind
+  let lift = Lwt.return
+
+  let rec read_loop ?(linger = "") (t : t) fd =
+    Lwt.catch (fun () ->
+      match fd with
+      | `Plain fd ->
+        let recv_buffer = Bytes.create 2048 in
+        Lwt_unix.recv fd recv_buffer 0 (Bytes.length recv_buffer) [] >|= fun r ->
+        (r, recv_buffer)
+      | `Tls fd ->
+        let recv_buffer = Bytes.create 2048 in
+        Tls_lwt.Unix.read fd recv_buffer >|= fun r ->
+        (r, recv_buffer))
+     (fun e ->
+      Log.err (fun m -> m "error %s reading from resolver" (Printexc.to_string e));
+      Lwt.return (0, Bytes.empty)) >>= function
+     | (0, _) ->
+       (match fd with
+       | `Plain fd -> close_socket fd
+       | `Tls fd -> Tls_lwt.Unix.close fd) >|= fun () ->
+       t.fd <- None;
+       if not (IM.is_empty t.requests) then
+         Log.info (fun m -> m "end of file reading from resolver")
+     | (read_len, cs) ->
+       let rec handle_data data =
+         let cs_len = String.length data in
+         if cs_len > 2 then
+           let len = String.get_uint16_be data 0 in
+           if cs_len - 2 >= len then
+             let packet, rest =
+               if cs_len - 2 = len
+               then data, ""
+               else String.sub data 0 (len + 2), String.sub data (len + 2) (String.length data - len - 2)
+             in
+             let id = String.get_uint16_be packet 2 in
+             (match IM.find_opt id t.requests with
+              | None -> Log.warn (fun m -> m "received unsolicited data, ignoring")
+              | Some (_, cond) ->
+                Lwt_condition.broadcast cond (Ok packet));
+             handle_data rest
+           else
+             read_loop ~linger:data t fd
+         else
+           read_loop ~linger:data t fd
+       in
+       let cs = String.sub (Bytes.unsafe_to_string cs) 0 read_len in
+       handle_data (if String.length linger = 0 then cs else linger ^ cs)
+
+  let req_all fd t =
+    IM.fold (fun _id (data, _) r ->
+        r >>= function
+        | Error _ as e -> Lwt.return e
+        | Ok () -> send_query fd data)
+      t.requests (Lwt.return (Ok ()))
+
+  let to_pairs =
+    List.map (function `Plaintext (ip, port) | `Tls (_, ip, port) -> ip, port)
+
+  let find_ns ns (addr, port) =
+    List.find (function `Plaintext (ip, p) | `Tls (_, ip, p) ->
+        Ipaddr.compare ip addr = 0 && p = port)
+      ns
+
+  let rec connect_to_ns_list (t : t) connected_condition nameservers =
+    let ns = to_pairs nameservers in
+    (* The connect_timeout given here is a bit too much, since it should
+       be (a) connect to the remote NS (b) send query, receive answer.
+
+       At the moment, how this is done, is that we use the connect_timeout
+       for (a) and another separate one for (b). Since we do connection
+       pooling, it is slightly tricky to use only a single connect_timeout. *)
+    Happy_eyeballs_lwt.connect_ip ~connect_timeout:t.timeout_ns t.he ns >>= function
+    | Error `Msg msg ->
+      let err =
+        Error (`Msg (Fmt.str "error %s connecting to resolver %a"
+                       msg
+                       Fmt.(list ~sep:(any ", ") (pair ~sep:(any ":") Ipaddr.pp int))
+                       (to_pairs (nameserver_ips t.nameservers))))
+      in
+      Lwt_condition.broadcast connected_condition err;
+      t.connected_condition <- None;
+      Lwt.return err
+    | Ok (addr, socket) ->
+      let continue socket =
+        t.fd <- Some socket;
+        Lwt.async (fun () ->
+            read_loop t socket >>= fun () ->
+            if IM.is_empty t.requests then
+              Lwt.return_unit
+            else
+              connect_via_tcp_to_ns t >|= function
+              | Error (`Msg msg) ->
+                Log.err (fun m -> m "error while connecting to resolver: %s"  msg)
+              | Ok () -> ());
+        Lwt_condition.broadcast connected_condition (Ok ());
+        t.connected_condition <- None;
+        req_all socket t
+      in
+      let config = find_ns (nameserver_ips t.nameservers) addr in
+      match config with
+      | `Plaintext _ -> continue (`Plain socket)
+      | `Tls (tls_cfg, _, _) ->
+        Lwt.catch (fun () ->
+            Tls_lwt.Unix.client_of_fd tls_cfg socket >>= fun f ->
+            continue (`Tls f))
+          (fun e ->
+             Log.warn (fun m -> m "TLS handshake with %a:%d failed: %s"
+                          Ipaddr.pp (fst addr) (snd addr) (Printexc.to_string e));
+             let ns' =
+               List.filter
+                 (function
+                   | `Tls (_, ip, port) ->
+                     not (Ipaddr.compare ip (fst addr) = 0 && port = snd addr)
+                   | _ -> true)
+                 nameservers
+             in
+             if ns' = [] then begin
+               let err = Error (`Msg "no further nameservers configured") in
+               Lwt_condition.broadcast connected_condition err;
+               t.connected_condition <- None;
+               Lwt.return err
+             end else
+               connect_to_ns_list t connected_condition ns')
+
+  and connect_via_tcp_to_ns (t : t) =
+    match t.fd, t.connected_condition with
+    | Some _, _ -> Lwt.return (Ok ())
+    | None, Some w -> Lwt_condition.wait w
+    | None, None ->
+      let connected_condition = Lwt_condition.create () in
+      t.connected_condition <- Some connected_condition ;
+      maybe_resolv_conf t;
+      connect_to_ns_list t connected_condition (nameserver_ips t.nameservers)
+
+  let connect t =
+    connect_via_tcp_to_ns t >|= function
+    | Ok () -> Ok (`Tcp, t)
+    | Error `Msg msg -> Error (`Msg msg)
+end
+
+(* Now that we have our {!Transport} implementation we can include the logic
+   that goes on top of it: *)
+include Dns_client.Make(Transport)
+
+let create ?cache_size ?edns ?nameservers ?timeout happy_eyeballs =
+  let dns = create ?cache_size ?edns ?nameservers ?timeout happy_eyeballs in
+  let getaddrinfo record domain_name =
+    let open Lwt_result.Infix in
+    match record with
+    | `A ->
+      getaddrinfo dns Dns.Rr_map.A domain_name >|= fun (_ttl, set) ->
+      Ipaddr.V4.Set.fold (fun ipv4 -> Ipaddr.Set.add (Ipaddr.V4 ipv4))
+        set Ipaddr.Set.empty
+    | `AAAA ->
+      getaddrinfo dns Dns.Rr_map.Aaaa domain_name >|= fun (_ttl, set) ->
+      Ipaddr.V6.Set.fold (fun ipv6 -> Ipaddr.Set.add (Ipaddr.V6 ipv6))
+        set Ipaddr.Set.empty
+  in
+  Happy_eyeballs_lwt.inject happy_eyeballs getaddrinfo;
+  dns
+
+(* initialize the RNG *)
+let () = Mirage_crypto_rng_unix.use_default ()