mirror of
https://github.com/mirage/qubes-mirage-firewall.git
synced 2024-10-01 01:05:39 -04:00
DNS: start task reading Lwt_mvar and distributing DNS replies to clients
Before, a DNS request was sent and the first thing appearing in the Lwt_mvar was taken as reply. The issue with this was two-fold: - it could be a reply for a different request - there could be DNS replies being sent to the uplink stack leading to Lwt_mvar.put being called, which blocks if there is already a value in the mvar. No, the separate task is a loop reading the mvar, using a Lwt_condition to signal the receive of that ID (potentially discarding if there's no client waiting). The DNS query registers itself (using the ID) in the map with a Lwt_condition, and waits to be notified (or a timeout occurs).
This commit is contained in:
parent
ddfb17c0b2
commit
0e0917f4fe
31
my_dns.ml
31
my_dns.ml
@ -5,11 +5,14 @@ module Transport (R : Mirage_random.S) (C : Mirage_clock.MCLOCK) (Time : Mirage_
|
||||
type io_addr = Ipaddr.V4.t * int
|
||||
type stack = Router.t * (src_port:int -> dst:Ipaddr.V4.t -> dst_port:int -> Cstruct.t -> (unit, [ `Msg of string ]) result Lwt.t) * (Udp_packet.t * Cstruct.t) Lwt_mvar.t
|
||||
|
||||
module IM = Map.Make(Int)
|
||||
|
||||
type t = {
|
||||
protocol : Dns.proto ;
|
||||
nameserver : io_addr ;
|
||||
stack : stack ;
|
||||
timeout_ns : int64 ;
|
||||
mutable requests : Cstruct.t Lwt_condition.t IM.t ;
|
||||
}
|
||||
type context = t
|
||||
|
||||
@ -17,12 +20,26 @@ module Transport (R : Mirage_random.S) (C : Mirage_clock.MCLOCK) (Time : Mirage_
|
||||
let rng = R.generate ?g:None
|
||||
let clock = C.elapsed_ns
|
||||
|
||||
let rec read t =
|
||||
let _, _, answer = t.stack in
|
||||
Lwt_mvar.take answer >>= fun (_, data) ->
|
||||
if Cstruct.length data > 2 then begin
|
||||
match IM.find_opt (Cstruct.BE.get_uint16 data 0) t.requests with
|
||||
| Some cond -> Lwt_condition.broadcast cond data
|
||||
| None -> ()
|
||||
end;
|
||||
read t
|
||||
|
||||
let create ?nameservers ~timeout stack =
|
||||
let protocol, nameserver = match nameservers with
|
||||
| None | Some (_, []) -> invalid_arg "no nameserver found"
|
||||
| Some (proto, ns :: _) -> proto, ns
|
||||
in
|
||||
{ protocol ; nameserver ; stack ; timeout_ns = timeout }
|
||||
let t =
|
||||
{ protocol ; nameserver ; stack ; timeout_ns = timeout ; requests = IM.empty }
|
||||
in
|
||||
Lwt.async (fun () -> read t);
|
||||
t
|
||||
|
||||
let with_timeout timeout_ns f =
|
||||
let timeout = Time.sleep_ns timeout_ns >|= fun () -> Error (`Msg "DNS request timeout") in
|
||||
@ -33,14 +50,18 @@ module Transport (R : Mirage_random.S) (C : Mirage_clock.MCLOCK) (Time : Mirage_
|
||||
let send_recv (ctx : context) buf : (Cstruct.t, [> `Msg of string ]) result Lwt.t =
|
||||
let open Router in
|
||||
let dst, dst_port = ctx.nameserver in
|
||||
let router, send_udp, answer = ctx.stack in
|
||||
let router, send_udp, _ = ctx.stack in
|
||||
let src_port, evict =
|
||||
My_nat.free_udp_port router.nat ~src:router.uplink#my_ip ~dst ~dst_port:53
|
||||
in
|
||||
let id = Cstruct.BE.get_uint16 buf 0 in
|
||||
with_timeout ctx.timeout_ns
|
||||
((send_udp ~src_port ~dst ~dst_port buf >|= Rresult.R.open_error_msg) >>= function
|
||||
| Ok () -> (Lwt_mvar.take answer >|= fun (_, dns_response) -> Ok dns_response)
|
||||
| Error _ as e -> Lwt.return e) >|= fun result ->
|
||||
(let cond = Lwt_condition.create () in
|
||||
ctx.requests <- IM.add id cond ctx.requests;
|
||||
(send_udp ~src_port ~dst ~dst_port buf >|= Rresult.R.open_error_msg) >>= function
|
||||
| Ok () -> Lwt_condition.wait cond >|= fun dns_response -> Ok dns_response
|
||||
| Error _ as e -> Lwt.return e) >|= fun result ->
|
||||
ctx.requests <- IM.remove id ctx.requests;
|
||||
evict ();
|
||||
result
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user