unikernel/duniverse/ocaml-dns/resolver/dns_block.ml

@@ -0,0 +1,100 @@
+open Dns
+
+let src = Logs.Src.create "dns_block" ~doc:"DNS block"
+module Log = (val Logs.src_log src : Logs.LOG)
+
+let nameserver =
+  let lh = Domain_name.of_string_exn "localhost"
+  and bl = Domain_name.of_string_exn "blocked"
+  in
+  fun ns -> Domain_name.equal ns lh || Domain_name.equal ns bl
+
+let ipv4 =
+  let lh = Ipaddr.V4.(Set.singleton localhost)
+  and any = Ipaddr.V4.(Set.singleton any)
+  in
+  fun ipv4s -> Ipaddr.V4.Set.equal ipv4s lh || Ipaddr.V4.Set.equal ipv4s any
+
+let ipv6 =
+  let lh = Ipaddr.V6.(Set.singleton localhost)
+  and un = Ipaddr.V6.(Set.singleton unspecified)
+  in
+  fun ipv6s -> Ipaddr.V6.Set.equal ipv6s lh || Ipaddr.V6.Set.equal ipv6s un
+
+let likely reply =
+  (* HACK! We assume blocked domains have a certain shape. *)
+  let blocked_soa auth =
+    Domain_name.Map.cardinal auth > 0 &&
+    Domain_name.Map.for_all (fun _domain rr ->
+        match Rr_map.find Rr_map.Soa rr with
+        | None -> false
+        | Some soa -> nameserver soa.nameserver)
+      auth
+  in
+  match reply.Packet.data with
+  | `Answer (answ, _auth) ->
+    Domain_name.Map.for_all
+      (fun _domain rr ->
+         Rr_map.for_all
+           (function
+             | Rr_map.B (Rr_map.A, (_, ips)) -> ipv4 ips
+             | Rr_map.B (Rr_map.Aaaa, (_, ips)) -> ipv6 ips
+             | _ -> false)
+           rr)
+      answ
+  | `Rcode_error (Rcode.NXDomain, _, Some (_answ, auth)) -> blocked_soa auth
+  | _ -> false
+
+let reason reply =
+  let find_soa_hostmaster rr =
+    match Rr_map.find Rr_map.Soa rr with
+    | None -> None
+    | Some soa ->
+      if nameserver soa.Soa.nameserver then
+        Some (Domain_name.to_string soa.Soa.hostmaster)
+      else
+        None
+  in
+  let find_soa_hostmaster_in_domain_map map =
+    Domain_name.Map.fold (fun _domain rr acc ->
+        match acc, find_soa_hostmaster rr with
+        | None, x -> x
+        | Some x, None -> Some x
+        | Some x, Some y ->
+          if not (String.equal x y) then
+            Log.info (fun m -> m "finding blocklist resulted in %S and %S, using the first" x y);
+          Some x) map None
+  in
+  let find_soa_hostmaster_in_reply answer authority =
+    match find_soa_hostmaster_in_domain_map answer, find_soa_hostmaster_in_domain_map authority with
+    | None, x -> x
+    | Some x, None -> Some x
+    | Some x, Some y ->
+      if not (String.equal x y) then
+        Log.info (fun m -> m "finding blocklist resulted in %S (answer) and %S (authority), using the first" x y);
+      Some x
+  in
+  let r =
+    match reply.Packet.data with
+    | `Answer (answ, auth) -> find_soa_hostmaster_in_reply answ auth
+    | `Rcode_error (Rcode.NXDomain, _, Some (answ, auth)) ->
+      find_soa_hostmaster_in_reply answ auth
+    | _ -> None
+  in
+  Option.map (fun reason -> "appears in blocklist " ^ reason) r
+
+let edns reply =
+  if likely reply then
+    (* After guessing that a domain is blocked we add [`Filtered] extended error
+       code and emit a [`Blocked] metrics event. *)
+    let reason = reason reply in
+    match reply.edns with
+    | None ->
+      Some (Edns.create ~extended_error:(`Blocked, reason) ())
+    | Some ({ Edns.extensions = []; extended_rcode; version; dnssec_ok; payload_size }) ->
+      Some (Edns.create ~extended_error:(`Blocked, reason) ~extended_rcode ~version ~dnssec_ok ~payload_size ())
+    | Some edns ->
+      Log.warn (fun m -> m "don't know how to extend edns to add extended error; not doing anything:@ %a" Edns.pp edns);
+      Some edns
+  else
+    None