unikernel/duniverse/ocaml-dns/mirage/stub/dns_stub_mirage.ml

@@ -0,0 +1,402 @@
+(* mirage stub resolver *)
+open Lwt.Infix
+
+open Dns
+
+let src = Logs.Src.create "dns_stub_mirage" ~doc:"effectful DNS stub layer"
+module Log = (val Logs.src_log src : Logs.LOG)
+
+module Make (S : Tcpip.Stack.V4V6) = struct
+
+  (* data in the wild:
+     - a request comes in hdr, q
+       - q to be found in cache
+       - q not found in cache (to be forwarded to the recursive resolver)
+         - unless q in transit (this to-be-done if it is worth it (is it?))
+         - a fresh hdr, q is generated and sent to the recursive resolver
+         - now hdr, q is registered to be awaited for
+         -- we can either signal the request task once we found something,
+            or preserve the original hdr, q together with ip and port
+     - a reply goes out hdr, q, answer
+
+     the "Client" is only concerned about the connection to the resolver, with
+     multiplexing.
+
+     the current API is:
+      dns_client calls connect .. -> flow
+                       send flow data
+                       recv flow (* potentially multiple times *)
+
+     i.e. our flow being (int * _):
+       connect <nothing>
+       send (id, _) data <- id <- data[2..3]
+       recv (id, _) <- registers condition in N[id] ; waits ; removes condition
+
+     or phrased differently:
+       a recv_loop reads continously, whenever a full packet is received,
+        N[id] is woken up with the packet
+  *)
+
+  let metrics =
+    let f = function
+      | `Udp_queries -> "udp-queries"
+      | `Tcp_queries -> "tcp-queries"
+      | `Ocaml_queries -> "ocaml-queries"
+      | `Tcp_connections -> "tcp-connections"
+      | `Authoritative_answers -> "authoritative-answers"
+      | `Authoritative_errors -> "authoritative-errors"
+      | `Reserved_answers -> "reserved-answers"
+      | `On_update -> "on-update"
+      | `Resolver_queries -> "resolver-queries"
+      | `Resolver_answers -> "resolver-answers"
+      | `Resolver_nodata -> "resolver-nodata"
+      | `Resolver_nodomain -> "resolver-nodomain"
+      | `Resolver_servfail -> "resolver-servfail"
+      | `Resolver_notimp -> "resolver-notimplemented"
+    in
+    let metrics = Dns.counter_metrics ~f "stub-resolver" in
+    (fun x -> Metrics.add metrics (fun x -> x) (fun d -> d x))
+
+  module H = Happy_eyeballs_mirage.Make(S)
+  module Client = Dns_client_mirage.Make(S)(H)
+  module TLS = Tls_mirage.Make(S.TCP)
+
+  (* likely this should contain:
+     - a primary server (handling updates)
+     - a client on steroids: multiplexing on connections
+     - listening for DNS requests from clients:
+        first find them in primary server
+        if not authoritative, use the client
+  *)
+
+  (* task management
+     - multiple requests for the same name, type can be done at the same "time"
+     -> need to remember outstanding requests and signal to clients
+  *)
+
+  (* take multiple resolver IPs and round-robin / ask both (take first answer,
+     ignoring ServFail etc.) *)
+
+  (* timeout of resolver, retransmission (to another resolver / another flow) *)
+
+  module Dns_flow = Dns_mirage.Make(S)
+
+  type t = {
+    client : Client.t ;
+    reserved : Dns_server.t ;
+    mutable server : Dns_server.t ;
+    on_update : old:Dns_trie.t -> ?authenticated_key:[`raw] Domain_name.t -> update_source:Ipaddr.t -> Dns_trie.t -> unit Lwt.t ;
+    push : (Ipaddr.t * int * string * (int32 * string) Lwt.u) option -> unit ;
+    mutable update_tls : Tls.Config.server -> unit ;
+  }
+
+  let primary_data { server ; _ } =
+    server.Dns_server.data
+
+  let update_primary_data t trie =
+    let server = Dns_server.with_data t.server trie in
+    t.server <- server
+
+  let resolve_external { push ; _ } (ip, port) data =
+    let th, wk = Lwt.wait () in
+    push (Some (ip, port, data, wk));
+    th
+
+  let update_tls { update_tls ; _ } tls = update_tls tls
+
+  let build_reply header question proto ?additional data =
+    let ttl = Packet.minimum_ttl data in
+    let packet = Packet.create ?additional header question data in
+    ttl, fst (Packet.encode proto packet)
+
+  let query_server trie question data header proto =
+    match Dns_server.handle_question trie question with
+    | Ok (_flags, answer, additional) ->
+      (* TODO do sth with flags *)
+      metrics `Authoritative_answers;
+      let data = `Answer answer in
+      let ttl = Packet.minimum_ttl data in
+      let packet = Packet.create ?additional header question data in
+      let packet =
+        match Dns_block.edns packet with
+        | None -> packet
+        | Some edns ->
+          Dns_resolver_metrics.resolver_stats `Blocked;
+          Dns.Packet.with_edns packet (Some edns)
+      in
+      let reply = ttl, fst (Packet.encode proto packet) in
+      Some reply
+    | Error (Rcode.NotAuth, _) -> None
+    | Error (rcode, answer) ->
+      metrics `Authoritative_errors;
+      let data = `Rcode_error (rcode, Packet.opcode_data data, answer) in
+      let reply = build_reply header question proto data in
+      Some reply
+
+  let tsig_decode_sign server proto packet buf header question =
+    let now = Mirage_ptime.now () in
+    match Dns_server.handle_tsig server now packet buf with
+    | Error _ ->
+      let data =
+        `Rcode_error (Rcode.Refused, Packet.opcode_data packet.Packet.data, None)
+      in
+      let reply = build_reply header question proto data in
+      Error reply
+    | Ok k ->
+      let key =
+        match k with None -> None | Some (keyname, _, _, _) -> Some keyname
+      in
+      let sign data =
+        let ttl = Packet.minimum_ttl data in
+        let packet = Packet.create header question data in
+        match k with
+        | None -> Some (ttl, fst (Packet.encode proto packet))
+        | Some (keyname, _tsig, mac, dnskey) ->
+          match Dns_tsig.encode_and_sign ~proto ~mac packet now dnskey keyname with
+          | Error s -> Log.err (fun m -> m "error %a while signing answer" Dns_tsig.pp_s s); None
+          | Ok (cs, _) -> Some (ttl, cs)
+      in
+      Ok (key, sign)
+
+  let axfr_server server proto packet question buf header =
+    match tsig_decode_sign server proto packet buf header question with
+    | Error e -> Some e
+    | Ok (key, sign) ->
+      match Dns_server.handle_axfr_request server proto key question with
+      | Error rcode ->
+        let err = `Rcode_error (rcode, Packet.opcode_data packet.Packet.data, None) in
+        let reply = build_reply header question proto err in
+        Some reply
+      | Ok axfr ->
+        sign (`Axfr_reply axfr)
+
+  let update_server t proto ip packet question u buf header =
+    let server = t.server in
+    match tsig_decode_sign server proto packet buf header question with
+    | Error e -> Lwt.return (Some e)
+    | Ok (key, sign) ->
+      match Dns_server.handle_update server proto key question u with
+      | Ok (trie, _) ->
+        let old = server.data in
+        let server' = Dns_server.with_data server trie in
+        t.server <- server';
+        metrics `On_update;
+        t.on_update ~old ?authenticated_key:key ~update_source:ip trie >|= fun () ->
+        sign `Update_ack
+      | Error rcode ->
+        Lwt.return (sign (`Rcode_error (rcode, Opcode.Update, None)))
+
+  let server t proto ip packet header question data buf =
+    match data with
+    | `Query -> Lwt.return (query_server t.server question data header proto)
+    | `Axfr_request ->
+      Lwt.return (axfr_server t.server proto packet question buf header)
+    | `Update u ->
+      update_server t proto ip packet question u buf header
+    | _ ->
+      let data =
+        `Rcode_error (Rcode.NotImp, Packet.opcode_data packet.Packet.data, None)
+      in
+      let pkt = build_reply header question proto data in
+      Lwt.return (Some pkt)
+
+  let resolve t question data header proto =
+    metrics `Resolver_queries;
+    let name = fst question in
+    match data, snd question with
+    | `Query, `K Rr_map.K key ->
+      begin Client.get_resource_record t.client key name >|= function
+        | Error `Msg msg ->
+          Log.err (fun m -> m "couldn't resolve %s" msg);
+          let data = `Rcode_error (Rcode.ServFail, Opcode.Query, None) in
+          metrics `Resolver_servfail;
+          let reply = build_reply header question proto data in
+          Some reply
+        | Error `No_data (domain, soa) ->
+          let answer = (Name_rr_map.empty, Name_rr_map.singleton domain Soa soa) in
+          let data = `Answer answer in
+          metrics `Resolver_nodata;
+          let reply = build_reply header question proto data in
+          Some reply
+        | Error `No_domain (domain, soa) ->
+          let answer = (Name_rr_map.empty, Name_rr_map.singleton domain Soa soa) in
+          let data = `Rcode_error (Rcode.NXDomain, Opcode.Query, Some answer) in
+          metrics `Resolver_nodomain;
+          let reply = build_reply header question proto data in
+          Some reply
+        | Ok reply ->
+          let answer = (Name_rr_map.singleton name key reply, Name_rr_map.empty) in
+          let data = `Answer answer in
+          metrics `Resolver_answers;
+          let reply = build_reply header question proto data in
+          Some reply
+      end
+    | _ ->
+      Log.err (fun m -> m "not implemented %a, data %a"
+                   Dns.Packet.Question.pp question
+                   Dns.Packet.pp_data data);
+      let data = `Rcode_error (Rcode.NotImp, Packet.opcode_data data, None) in
+      metrics `Resolver_notimp;
+      let reply = build_reply header question proto data in
+      Lwt.return (Some reply)
+
+  (* we're now doing up to three lookups for each request:
+    - in authoritative server (Dns_trie)
+    - in reserved trie (Dns_trie)
+    - in resolver cache (Dns_cache)
+    - asking a remote resolver
+
+     instead, on startup authoritative (from external) could be merged with
+     reserved (but that makes data store very big and not easy to understand
+     (lots of files for the reserved zones)) *)
+  let handle t proto ip buf =
+    match Packet.decode buf with
+    | Error err ->
+      Log.err (fun m -> m "couldn't decode %a" Packet.pp_err err);
+      Dns_resolver_metrics.response_metric 0L;
+      Dns_resolver_metrics.resolver_stats `Error;
+      let answer = Packet.raw_error buf Rcode.FormErr in
+      Lwt.return (Option.map (fun r -> 0l, r) answer)
+    | Ok packet ->
+      Dns_resolver_metrics.resolver_stats `Queries;
+      let start = Mirage_mtime.elapsed_ns () in
+      let header, question, data = packet.Packet.header, packet.question, packet.data in
+      (* check header flags: recursion desired (and send recursion available) *)
+      (server t proto ip packet header question data buf >>= function
+        | Some data -> Lwt.return (Some data)
+        | None ->
+          (* next look in reserved trie! *)
+          match query_server t.reserved question data header proto with
+          | Some data -> metrics `Reserved_answers ; Lwt.return (Some data)
+          | None -> resolve t question data header proto) >|= fun reply ->
+      let stop = Mirage_mtime.elapsed_ns () in
+      Dns_resolver_metrics.response_metric (Int64.sub stop start);
+      reply
+
+  let send_tls flow data =
+    let len = Cstruct.create 2 in
+    Cstruct.BE.set_uint16 len 0 (Cstruct.length data);
+    TLS.writev flow [len; data] >>= function
+    | Ok () -> Lwt.return (Ok ())
+    | Error e ->
+      Log.err (fun m -> m "tls error %a while writing" TLS.pp_write_error e);
+      TLS.close flow >|= fun () ->
+      Error ()
+
+  type tls_flow = { tls_flow : TLS.flow ; mutable linger : Cstruct.t }
+
+  let rec read_tls ({ tls_flow ; linger } as f) length =
+    if Cstruct.length linger >= length then
+      let a, b = Cstruct.split linger length in
+      f.linger <- b;
+      Lwt.return (Ok a)
+    else
+      TLS.read tls_flow >>= function
+      | Ok `Eof -> Log.debug (fun m -> m "end of file while reading"); TLS.close tls_flow >|= fun () -> Error ()
+      | Error e -> Log.warn (fun m -> m "error reading TLS: %a" TLS.pp_error e); TLS.close tls_flow >|= fun () -> Error ()
+      | Ok (`Data d) ->
+        f.linger <- Cstruct.append linger d;
+        read_tls f length
+
+  let read_tls_packet f =
+    read_tls f 2 >>= function
+    | Error () -> Lwt.return (Error ())
+    | Ok k ->
+      let len = Cstruct.BE.get_uint16 k 0 in
+      read_tls f len
+
+  let create ?(cache_size = 10000) ?(udp = true) ?(tcp = true) ?(port = 53) ?tls ?(tls_port = 853) ?edns ?nameservers ?timeout ?(on_update = fun ~old:_ ?authenticated_key:_ ~update_source:_ _trie -> Lwt.return_unit) primary ~happy_eyeballs stack : t Lwt.t =
+    Client.connect ~cache_size ?edns ?nameservers ?timeout (stack, happy_eyeballs) >|= fun client ->
+    let server = Dns_server.Primary.server primary in
+    let stream, push = Lwt_stream.create () in
+    let reserved = Dns_server.create Dns_resolver_root.reserved Mirage_crypto_rng.generate in
+    let update_tls _ = () in
+    let t = { client ; reserved ; server ; on_update ; push ; update_tls } in
+    let udp_cb ~src ~dst:_ ~src_port buf =
+      let buf = Cstruct.to_string buf in
+      metrics `Udp_queries;
+      handle t `Udp src buf >>= function
+      | None -> Lwt.return_unit
+      | Some (_ttl, data) ->
+        let data = Cstruct.of_string data in
+        S.UDP.write ~src_port:port ~dst:src ~dst_port:src_port (S.udp stack) data >|= function
+        | Error e -> Log.warn (fun m -> m "udp: failure %a while sending to %a:%d"
+                                  S.UDP.pp_error e Ipaddr.pp src src_port)
+        | Ok () -> ()
+    in
+    if udp then
+      S.UDP.listen (S.udp stack) ~port udp_cb ;
+    let tcp_cb flow =
+      metrics `Tcp_connections;
+      let dst_ip, dst_port = S.TCP.dst flow in
+      Log.debug (fun m -> m "tcp connection from %a:%d" Ipaddr.pp dst_ip dst_port) ;
+      let f = Dns_flow.of_flow flow in
+      let rec loop () =
+        Dns_flow.read_tcp f >>= function
+        | Error () -> Lwt.return_unit
+        | Ok data ->
+          metrics `Tcp_queries;
+          let data = Cstruct.to_string data in
+          handle t `Tcp dst_ip data >>= function
+          | None ->
+            Log.warn (fun m -> m "no TCP output") ;
+            loop ()
+          | Some (_ttl, data) ->
+            let data = Cstruct.of_string data in
+            Dns_flow.send_tcp flow data >>= function
+            | Ok () -> loop ()
+            | Error () -> Lwt.return_unit
+      in
+      loop ()
+    in
+    if tcp then
+      S.TCP.listen (S.tcp stack) ~port tcp_cb;
+    let rec ocaml_cb () =
+      Lwt_stream.get stream >>= function
+      | Some (dst_ip, _dst_port, data, wk) ->
+        metrics `Ocaml_queries;
+        begin
+          handle t `Tcp dst_ip data >|= function
+          | None ->
+            Log.warn (fun m -> m "no TCP output")
+          | Some (ttl, data) ->
+            Lwt.wakeup wk (ttl, data);
+        end >>= fun () ->
+        ocaml_cb ()
+      | None -> Lwt.return_unit in
+    Lwt.async ocaml_cb;
+    let tls_cb cfg flow =
+      let dst_ip, dst_port = S.TCP.dst flow in
+      TLS.server_of_flow cfg flow >>= function
+      | Error e ->
+        Log.warn (fun m -> m "TLS error (from %a:%d): %a" Ipaddr.pp dst_ip dst_port
+          TLS.pp_write_error e);
+        Lwt.return_unit
+      | Ok tls ->
+        Log.debug (fun m -> m "tls connection from %a:%d" Ipaddr.pp dst_ip dst_port);
+        let tls_and_linger = { tls_flow = tls ; linger = Cstruct.empty } in
+        let rec loop () =
+          read_tls_packet tls_and_linger >>= function
+          | Error () ->
+            Lwt.return_unit
+          | Ok data ->
+            let data = Cstruct.to_string data in
+            handle t `Tcp dst_ip data >>= function
+            | None ->
+              Log.warn (fun m -> m "no TLS output") ;
+              loop ()
+            | Some (_ttl, data) ->
+              let data = Cstruct.of_string data in
+              send_tls tls data >>= function
+              | Ok () -> loop ()
+              | Error () -> Lwt.return_unit
+        in
+        loop ()
+    in
+    let update_tls tls_cfg =
+      S.TCP.listen (S.tcp stack) ~port:tls_port (tls_cb tls_cfg);
+    in
+    t.update_tls <- update_tls;
+    (match tls with None -> () | Some cfg -> update_tls cfg);
+    t
+end