Merge pull request #349 from RyanGibb/add-packet-callback

add a packet callback
mirage · May 17, 2024 · 2c0f350 · 2c0f350
2 parents 9f395ef + 18a51ab
commit 2c0f350
Show file tree

Hide file tree

Showing 4 changed files with 57 additions and 31 deletions.
diff --git a/dns-client.opam b/dns-client.opam
@@ -20,6 +20,8 @@ depends: [
   "domain-name" {>= "0.4.0"}
   "mtime" {>= "1.2.0"}
   "mirage-crypto-rng" {>= "0.11.0"}
+  "fmt" {>= "0.9.0"}
+  "ipaddr" {>= "5.5.0"}
   "alcotest" {with-test}
 ]
 synopsis: "DNS client API"

diff --git a/dns-server.opam b/dns-server.opam
@@ -24,6 +24,7 @@ depends: [
   "dns-tsig" {with-test}
   "base64" {with-test & >= "3.0.0"}
   "metrics"
+  "logs" {>= "0.7.0"}
 ]
 
 build: [

diff --git a/server/dns_server.ml b/server/dns_server.ml
@@ -407,6 +407,8 @@ let safe_decode buf =
     rx_metrics v.Packet.data;
     Ok v
 
+type packet_callback = Packet.Question.t -> Packet.reply option
+
 let handle_question t (name, typ) =
   (* TODO allow/disallowlist of allowed qtypes? what about ANY and UDP? *)
   match typ with
@@ -987,7 +989,7 @@ module Primary = struct
       end
     | _ -> Error ()
 
-  let handle_packet (t, m, l, ns) now ts proto ip _port p key =
+  let handle_packet ?(packet_callback = fun _q -> None) (t, m, l, ns) now ts proto ip _port p key =
     let key = match key with
       | None -> None
       | Some k -> Some (Domain_name.raw k)
@@ -1023,10 +1025,14 @@ module Primary = struct
         | _ -> l, ns, [], None
       in
       let answer =
-        let flags, data, additional = match handle_question t p.question with
-          | Ok (flags, data, additional) -> flags, `Answer data, additional
-          | Error (rcode, data) ->
-            err_flags rcode, `Rcode_error (rcode, Opcode.Query, data), None
+        let flags, data, additional =
+          match packet_callback p.question with
+          | Some reply -> Packet.Flags.singleton `Authoritative, (reply :> Packet.data), None
+          | None ->
+            match handle_question t p.question with
+            | Ok (flags, data, additional) -> flags, `Answer data, additional
+            | Error (rcode, data) ->
+              err_flags rcode, `Rcode_error (rcode, Opcode.Query, data), None
         in
         Packet.create ?additional (fst p.header, flags) p.question data
       in
@@ -1091,7 +1097,7 @@ module Primary = struct
       Log.err (fun m -> m "ignoring unsolicited %a" Packet.pp_data p);
       (t, m, l, ns), None, [], None
 
-  let handle_buf t now ts proto ip port buf =
+  let handle_buf ?(packet_callback = fun _q -> None) t now ts proto ip port buf =
     match
       let* res = safe_decode buf in
       Log.debug (fun m -> m "from %a received:@[%a@]" Ipaddr.pp ip
@@ -1109,7 +1115,7 @@ module Primary = struct
     | Ok p ->
       let handle_inner tsig_size keyname =
         let t, answer, out, notify =
-          handle_packet t now ts proto ip port p keyname
+          handle_packet ~packet_callback t now ts proto ip port p keyname
         in
         let answer = match answer with
           | Some answer ->
@@ -1661,17 +1667,21 @@ module Secondary = struct
                    Packet.Question.pp (Domain_name.raw zone, typ));
       Error Rcode.Refused
 
-  let handle_packet (t, zones) now ts ip p keyname =
+  let handle_packet ?(packet_callback = fun _q -> None) (t, zones) now ts ip p keyname =
     let keyname = match keyname with
       | None -> None
       | Some k -> Some (Domain_name.raw k)
     in
     match p.Packet.data with
     | `Query ->
-      let flags, data, additional = match handle_question t p.question with
-        | Ok (flags, data, additional) -> flags, `Answer data, additional
-        | Error (rcode, data) ->
-          err_flags rcode, `Rcode_error (rcode, Opcode.Query, data), None
+      let flags, data, additional =
+        match packet_callback p.question with
+        | Some reply -> Packet.Flags.singleton `Authoritative, (reply :> Packet.data), None
+        | None ->
+          match handle_question t p.question with
+          | Ok (flags, data, additional) -> flags, `Answer data, additional
+          | Error (rcode, data) ->
+            err_flags rcode, `Rcode_error (rcode, Opcode.Query, data), None
       in
       let answer =
         Packet.create ?additional (fst p.header, flags) p.question data
@@ -1850,7 +1860,7 @@ module Secondary = struct
         | Some (Processing_axfr (_, _, mac, _, _), _, _) -> Some mac
         | _ -> None
 
-  let handle_buf t now ts proto ip buf =
+  let handle_buf ?(packet_callback = fun _q -> None) t now ts proto ip buf =
     match
       let* res = safe_decode buf in
       Log.debug (fun m -> m "received a packet from %a: %a" Ipaddr.pp ip
@@ -1862,7 +1872,7 @@ module Secondary = struct
       t, Packet.raw_error buf rcode, None
     | Ok p ->
       let handle_inner keyname =
-        let t, answer, out = handle_packet t now ts ip p keyname in
+        let t, answer, out = handle_packet ~packet_callback t now ts ip p keyname in
         let answer = match answer with
           | Some answer ->
             let max_size, edns = Edns.reply p.edns in

diff --git a/server/dns_server.mli b/server/dns_server.mli
@@ -109,6 +109,9 @@ val handle_tsig : ?mac:Cstruct.t -> t -> Ptime.t -> Packet.t ->
 (** [handle_tsig ~mac t now packet buffer] verifies the tsig
     signature if present, returning the keyname, tsig, mac, and used key. *)
 
+type packet_callback = Packet.Question.t -> Packet.reply option
+(** [packet_callback question] either returns a reply to a DNS question [Some reply] or [None]. *)
+
 module Primary : sig
 
   type s
@@ -140,24 +143,32 @@ module Primary : sig
      data] creates a primary server. If [unauthenticated_zone_transfer] is
      provided and [true] (defaults to [false]), anyone can transfer the zones. *)
 
-  val handle_packet : s -> Ptime.t -> int64 -> proto -> Ipaddr.t -> int ->
-    Packet.t -> 'a Domain_name.t option ->
+  val handle_packet : ?packet_callback:packet_callback -> s -> Ptime.t -> int64
+    -> proto -> Ipaddr.t -> int -> Packet.t -> 'a Domain_name.t option ->
     s * Packet.t option * (Ipaddr.t * Cstruct.t list) list *
     [> `Notify of Soa.t option | `Keep ] option
-  (** [handle_packet s now ts src src_port proto key packet] handles the given
-     [packet], returning new state, an answer, and potentially notify packets to
-     secondary name servers. *)
-
-  val handle_buf : s -> Ptime.t -> int64 -> proto ->
-    Ipaddr.t -> int -> Cstruct.t ->
+  (** [handle_packet ~packet_callback s now ts src src_port proto key packet]
+      handles the given [packet], returning new state, an answer, and
+      potentially notify packets to secondary name servers. If [packet_callback]
+      is specified, it is called for each incoming query. If it returns
+      [Some reply], this reply is used instead of the usual lookup in the
+      zone data. It can be used for custom query processing, such as for load
+      balancing or transporting data. *)
+
+  val handle_buf : ?packet_callback:packet_callback -> s -> Ptime.t -> int64
+    -> proto -> Ipaddr.t -> int -> Cstruct.t ->
     s * Cstruct.t list * (Ipaddr.t * Cstruct.t list) list *
     [ `Notify of Soa.t option | `Signed_notify of Soa.t option | `Keep ] option *
     [ `raw ] Domain_name.t option
-  (** [handle_buf s now ts proto src src_port buffer] decodes the [buffer],
-     processes the DNS frame using {!handle_packet}, and encodes the reply.
-     The result is a new state, potentially a list of answers to the requestor,
-     a list of notifications to send out, information whether a notify (or
-     signed notify) was received, and the hmac key used for authentication. *)
+  (** [handle_buf ~packet_callback s now ts proto src src_port buffer] decodes
+      the [buffer], processes the DNS frame using {!handle_packet}, and encodes
+      the reply. The result is a new state, potentially a list of answers to the
+      requestor, a list of notifications to send out, information whether a
+      notify (or signed notify) was received, and the hmac key used for
+      authentication. If [packet_callback] is specified, it is called for each
+      incoming query. If it returns [Some reply], this reply is used instead of
+      the usual lookup in the zone data. This can be used for custom query
+      processing, such as for load balancing or transporting data. *)
 
   val closed : s -> Ipaddr.t -> s
   (** [closed s ip] marks the connection to [ip] closed. *)
@@ -192,14 +203,16 @@ module Secondary : sig
   (** [create ~primary ~tsig_verify ~tsig_sign ~rng keys] creates a secondary
      DNS server state. *)
 
-  val handle_packet : s -> Ptime.t -> int64 -> Ipaddr.t ->
-    Packet.t -> 'a Domain_name.t option ->
+  val handle_packet : ?packet_callback:packet_callback -> s -> Ptime.t -> int64 ->
+    Ipaddr.t -> Packet.t -> 'a Domain_name.t option ->
     s * Packet.t option * (Ipaddr.t * Cstruct.t) option
   (** [handle_packet s now ts ip proto key t] handles the incoming packet. *)
 
-  val handle_buf : s -> Ptime.t -> int64 -> proto -> Ipaddr.t -> Cstruct.t ->
+
+  val handle_buf : ?packet_callback:packet_callback -> s -> Ptime.t -> int64 ->
+    proto -> Ipaddr.t -> Cstruct.t ->
     s * Cstruct.t option * (Ipaddr.t * Cstruct.t) option
-  (** [handle_buf s now ts proto src buf] decodes [buf], processes with
+  (** [handle_buf ~packet_callback s now ts proto src buf] decodes [buf], processes with
       {!handle_packet}, and encodes the results. *)
 
   val timer : s -> Ptime.t -> int64 ->