Source file dns_client_unix.ml

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
(* {!Transport} provides the implementation of the underlying flow
   that is in turn used by {!Dns_client.Make} to provide the
   blocking Unix convenience module:
*)

module Transport : Dns_client.S
  with type io_addr = Ipaddr.t * int
   and type stack = unit
   and type +'a io = 'a
= struct
  type io_addr = Ipaddr.t * int
  type ns_addr = [`TCP | `UDP] * io_addr
  type stack = unit
  type t = {
    nameserver : ns_addr ;
    timeout_ns : int64 ;
  }
  type context = { t : t ; fd : Unix.file_descr ; timeout_ns : int64 ref }
  type +'a io = 'a

  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 create ?nameserver ~timeout () =
    let nameserver =
      Rresult.R.(get_ok (of_option ~none:(fun () ->
          let ip =
            match
              read_file "/etc/resolv.conf" >>= fun data ->
              Dns_resolvconf.parse data >>= fun nameservers ->
              List.fold_left (fun acc ns ->
                  match acc, ns with
                  | Ok ip, _ -> Ok ip
                  | _, `Nameserver ip -> Ok ip)
                (Error (`Msg "no nameserver")) nameservers
            with
            | Error _ -> Ipaddr.(V4 (V4.of_string_exn (fst Dns_client.default_resolver)))
            | Ok ip -> ip
          in
          Ok (`TCP, (ip, 53)))
          nameserver))
    in
    { nameserver ; timeout_ns = timeout }

  let nameserver { nameserver ; _ } = nameserver
  let clock = Mtime_clock.elapsed_ns
  let rng = Mirage_crypto_rng.generate ?g:None

  open Rresult

  let bind a b = b a
  let lift v = v

  let close { fd ; _ } = try Unix.close fd with _ -> ()

  let with_timeout ctx f =
    let start = clock () in
    (* TODO cancel execution of f when time_left is 0 *)
    let r = f ctx.fd in
    let stop = clock () in
    ctx.timeout_ns := Int64.sub !(ctx.timeout_ns) (Int64.sub stop start);
    if !(ctx.timeout_ns) <= 0L then
      Error (`Msg "DNS resolution timed out.")
    else
      r

  (* there is no connect timeouts, just a request timeout (unix: receive timeout) *)
  let connect ?nameserver:ns t =
    let proto, (server, port) =
      match ns with None -> nameserver t | Some x -> x
    in
    try
      begin match proto with
        | `UDP -> Ok Unix.((getprotobyname "udp").p_proto)
        | `TCP -> Ok Unix.((getprotobyname "tcp").p_proto)
      end >>= fun proto_number ->
      let fam = match server with Ipaddr.V4 _ -> Unix.PF_INET | Ipaddr.V6 _ -> Unix.PF_INET6 in
      let socket = Unix.socket fam Unix.SOCK_STREAM proto_number in
      let time_left = ref t.timeout_ns in
      let addr = Unix.ADDR_INET (Ipaddr_unix.to_inet_addr server, port) in
      let ctx = { t ; fd = socket ; timeout_ns = time_left } in
      try
        with_timeout ctx (fun fd ->
          Unix.connect fd addr;
          Ok ctx)
      with e ->
        close ctx;
        Error (`Msg (Printexc.to_string e))
    with e ->
      Error (`Msg (Printexc.to_string e))

  let send ctx (tx : Cstruct.t) =
    let str = Cstruct.to_string tx in
    try
      with_timeout ctx (fun fd ->
        Unix.setsockopt_float fd Unix.SO_SNDTIMEO (Duration.to_f !(ctx.timeout_ns));
        let res = Unix.send_substring fd str 0 (String.length str) [] in
        if res <> String.length str
        then
          Error (`Msg ("Broken write to upstream NS" ^ (string_of_int res)))
        else Ok ())
   with e ->
     Error (`Msg (Printexc.to_string e))

  let recv ctx =
    let buffer = Bytes.make 2048 '\000' in
    try
      with_timeout ctx (fun fd ->
        Unix.setsockopt_float fd Unix.SO_RCVTIMEO (Duration.to_f !(ctx.timeout_ns));
        let x = Unix.recv fd buffer 0 (Bytes.length buffer) [] in
        if x > 0 && x <= Bytes.length buffer then
          Ok (Cstruct.of_bytes buffer ~len:x)
        else
          Error (`Msg "Reading from NS socket failed"))
    with e ->
      Error (`Msg (Printexc.to_string e))
end

(* Now that we have our {!Transport} implementation we can include the logic
   that goes on top of it: *)
include Dns_client.Make(Transport)

(* initialize the RNG *)
let () = Mirage_crypto_rng_unix.initialize ()