mte/unikernel/duniverse/ocaml-dns/mirage/stub/dns_stub_mirage.ml
2025-11-11 02:07:51 +01:00

402 lines
16 KiB
OCaml

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