mte/unikernel/duniverse/ocaml-dns/lwt/client/dns_client_lwt.ml
2025-11-11 02:07:51 +01:00

368 lines
13 KiB
OCaml

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 ()