Merge pull request #352 from dinosaure/miou

Provide a Miou implementation of dns-client

dns-client-miou-unix.opam

@@ -0,0 +1,29 @@
+opam-version: "2.0"
+maintainer: "team AT robur dot coop"
+authors: ["Robur <[email protected]>"]
+homepage: "https://github.com/mirage/ocaml-dns"
+bug-reports: "https://github.com/mirage/ocaml-dns/issues"
+dev-repo: "git+https://github.com/mirage/ocaml-dns.git"
+license: "BSD-2-Clause"
+
+build: [
+  [ "dune" "subst"] {dev}
+  [ "dune" "build" "-p" name "-j" jobs ]
+  [ "dune" "runtest" "-p" name "-j" jobs ] {with-test}
+]
+
+depends: [
+  "dune" {>="2.0.0"}
+  "ocaml" {>= "5.0.0"}
+  "dns-client" {= version}
+  "domain-name" {>= "0.4.0"}
+  "ipaddr" {>= "5.3.0"}
+  "miou" {>= "0.1.0"}
+  "tls-miou-unix"
+  "happy-eyeballs" {>= "0.6.0"}
+  "happy-eyeballs-miou-unix"
+]
+synopsis: "DNS client API for Miou"
+description: """
+A client implementation using uDNS using Miou.
+"""

miou/client/dns_client_miou_unix.ml

@@ -0,0 +1,189 @@
+let error_msgf fmt = Fmt.kstr (fun msg -> Error (`Msg msg)) fmt
+
+let src = Logs.Src.create "dns-client-miou-unix"
+
+module Log = (val Logs.src_log src : Logs.LOG)
+
+module Transport = struct
+  open Happy_eyeballs_miou_unix
+
+  type +'a io = 'a
+
+  type io_addr =
+    [ `Plaintext of Ipaddr.t * int | `Tls of Tls.Config.client * Ipaddr.t * int ]
+
+  type t = {
+      nameservers: io_addr list
+    ; proto: Dns.proto
+    ; timeout: float
+    ; happy: stack
+  }
+  and stack = Happy_eyeballs_miou_unix.t
+
+  type context =
+    { fd : [ `Udp of Miou_unix.file_descr
+           | `Tcp of Miou_unix.file_descr
+           | `Tls of Tls_miou_unix.t ]
+    ; timeout : float }
+
+  let clock = Mtime_clock.elapsed_ns
+
+  let same_address ipaddr' port' = function
+    | `Plaintext (ipaddr, port) -> Ipaddr.compare ipaddr ipaddr' = 0 && port = port'
+    | `Tls (_, ipaddr, port) -> Ipaddr.compare ipaddr ipaddr' = 0 && port = port'
+
+  exception Timeout
+
+  let with_timeout ~timeout:ts fn =
+    let timeout () = Miou_unix.sleep ts; raise Timeout in
+    let prm1 = Miou.async timeout in
+    let prm0 = Miou.async fn in
+    Miou.await_first [ prm0; prm1 ]
+
+  let connect_to_nameservers t =
+    let ( let* ) = Result.bind in
+    match t.proto with
+    | `Tcp ->
+      let ip_of_nameserver = function
+        | `Plaintext (ipaddr, port) -> (ipaddr, port)
+        | `Tls (_, ipaddr, port) -> (ipaddr, port) in
+      let ips = List.map ip_of_nameserver t.nameservers in
+      let* ((ipaddr, port) as addr), fd = connect_ip t.happy ips in
+      begin match List.find (same_address ipaddr port) t.nameservers with
+      | `Plaintext _ -> Ok (addr, `Tcp fd)
+      | `Tls (config, _, _) ->
+        try let fd = Tls_miou_unix.client_of_fd config fd in
+            Ok (addr, `Tls fd)
+        with End_of_file ->
+            Miou_unix.close fd;
+            error_msgf "Connection to nameservers (via TLS) impossible" end
+    | `Udp ->
+      let is_plaintext = function `Plaintext v -> Either.Left v | _ -> Either.Right () in
+      let[@warning "-8"] (ipaddr, port) :: _, _ = List.partition_map is_plaintext t.nameservers in
+      let proto_number, socket_type = Unix.((getprotobyname "udp").p_proto, SOCK_DGRAM) in
+      let domain = match ipaddr with
+        | Ipaddr.V4 _ -> Unix.PF_INET
+        | Ipaddr.V6 _ -> Unix.PF_INET6 in
+      let fd = Unix.socket domain socket_type proto_number in
+      let addr = Unix.ADDR_INET (Ipaddr_unix.to_inet_addr ipaddr, port) in
+      let connect () =
+        Unix.connect fd addr;
+        ((ipaddr, port), `Udp (Miou_unix.of_file_descr fd)) in
+      match with_timeout ~timeout:t.timeout connect with
+      | Ok value -> Ok value
+      | Error Timeout ->
+          Unix.close fd;
+          error_msgf "Connection to nameservers (via UDP) timeout"
+      | Error exn ->
+          Unix.close fd;
+          error_msgf "Unexpected error: %S" (Printexc.to_string exn)
+
+  let nameservers { nameservers; proto; _ } = (proto, nameservers)
+  let bind x f = f x
+  let lift = Fun.id
+  let rng = Mirage_crypto_rng.generate ?g:None
+
+  let connect t =
+    let ( let* ) = Result.bind in
+    let* ((addr, port), fd) = connect_to_nameservers t in
+    Log.debug (fun m -> m "Connected to a nameserver %a:%d" Ipaddr.pp addr port);
+    match fd with
+    | `Tcp _ | `Tls _ -> Ok (`Tcp, { fd; timeout= t.timeout })
+    | `Udp _ -> Ok (`Udp, { fd; timeout= t.timeout })
+
+  let send_recv_tls ~timeout ~id fd str =
+    let send () = Tls_miou_unix.write fd str in
+    let recv () =
+      let rec go buf rx_len =
+        let expected_len =
+          if rx_len >= 2 then Some (Bytes.get_uint16_be buf 0) else None in
+        match expected_len with
+        | None ->
+          let len = Tls_miou_unix.read fd buf ~off:rx_len in
+          if rx_len + len >= 2 && len > 0 then go buf (rx_len + len)
+          else failwith "TLS connection closed by nameserver"
+        | Some expected_len when rx_len >= expected_len + 2 ->
+          let id' = Bytes.get_uint16_be buf 2 in
+          if id = id'
+          then Bytes.sub_string buf 0 (expected_len + 2)
+          else
+            let buf' = Bytes.make 2048 '\000' in
+            let rx_len' = rx_len - (expected_len + 2) in
+            Bytes.blit buf (expected_len + 2) buf' 0 rx_len';
+            go buf' rx_len'
+        | Some expected_len when Bytes.length buf >= expected_len + 2 ->
+          let len = (expected_len + 2) - rx_len in
+          Tls_miou_unix.really_read fd buf ~off:rx_len ~len;
+          go buf (rx_len + len)
+        | Some expected_len ->
+          (* NOTE(dinosaure): in this branch, [buf] is not large enough to store
+             the DNS packet. We allocate a new buffer which can store the actual
+             DNS packet and use it for the next [go] iteration. *)
+          let buf' = Bytes.make (expected_len + 2) '\000' in
+          Bytes.blit buf 0 buf' 0 rx_len;
+          go buf' rx_len in
+      go (Bytes.make 2048 '\000') 0 in
+    let ( >>= ) = Result.bind in
+    match with_timeout ~timeout send >>= fun () ->
+          with_timeout ~timeout recv with
+    | Ok _ as rx -> rx
+    | Error Timeout -> error_msgf "DNS request timeout"
+    | Error (Failure msg) -> Error (`Msg msg)
+    | Error (End_of_file | Tls_miou_unix.Closed_by_peer) ->
+      error_msgf "End of file reading from nameserver"
+    | Error exn ->
+      error_msgf "Got an unexpected exception: %s"
+        (Printexc.to_string exn)
+
+  let send_recv { fd; timeout } str =
+    if String.length str > 4 then begin
+      match fd with
+      | `Tls fd ->
+        let id = String.get_int16_be str 2 in
+        send_recv_tls ~timeout ~id fd str 
+      | `Udp fd | `Tcp fd ->
+          let fd = Miou_unix.to_file_descr fd in
+          Unix.clear_nonblock fd;
+          let send () =
+            Log.debug (fun m -> m "sending a dns packet to resolver");
+            Unix.setsockopt_float fd Unix.SO_SNDTIMEO timeout;
+            let len = Unix.send_substring fd str 0 (String.length str) [] in
+            if len <> String.length str
+            then failwith "Broken write to upstream nameserver" in
+          let recv () =
+            let buffer = Bytes.make 2048 '\000' in
+            Unix.setsockopt_float fd Unix.SO_RCVTIMEO timeout;
+            let len = Unix.recv fd buffer 0 (Bytes.length buffer) [] in
+            (* TODO(dinosaure): should we check rx_len and continue until we got
+               the full packet (only for tcp/ip)? *)
+            if len > 0 && len <= Bytes.length buffer
+            then Bytes.sub_string buffer 0 len
+            else failwith "Reading from nameserver socket failed" in
+          let ( >>= ) = Result.bind in
+          match with_timeout ~timeout send >>= fun () ->
+                with_timeout ~timeout recv with
+          | Ok _ as rx -> rx
+          | Error Timeout -> error_msgf "DNS request timeout"
+          | Error (Failure msg) -> Error (`Msg msg)
+          | Error exn ->
+              error_msgf "Got an unexpected exception: %s"
+                (Printexc.to_string exn)
+    end
+    else error_msgf "Invalid context (data length <= 4)"
+
+  let close { fd; _ } = match fd with
+    | `Tcp fd | `Udp fd -> Miou_unix.close fd
+    | `Tls fd -> Tls_miou_unix.close fd
+
+  let of_ns ns = Int64.to_float ns /. 1_000_000_000.
+
+  let create ?nameservers ~timeout happy =
+    let proto, nameservers =
+      match nameservers with
+      | None -> (`Udp, [ `Plaintext (Ipaddr.of_string_exn "8.8.8.8", 53) ])
+      | Some (a, nss) -> (a, nss)
+    in
+    { nameservers; proto; timeout= of_ns timeout; happy }
+end
+
+include Dns_client.Make (Transport)

miou/client/dns_client_miou_unix.mli

@@ -0,0 +1,6 @@
+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
+   and type stack = Happy_eyeballs_miou_unix.t
+
+include module type of Dns_client.Make (Transport)

miou/client/dune

@@ -0,0 +1,5 @@
+(library
+ (name dns_client_miou_unix)
+ (modules dns_client_miou_unix)
+ (public_name dns-client-miou-unix)
+ (libraries dns-client tls-miou-unix happy-eyeballs-miou-unix))