Merge pull request #470 from djs55/dns

djs55 · web-flow · commit ee5dd643b233 · 2019-04-03T15:39:29.000+01:00
dns: set the TC bit and truncate if UDP response &gt; 512 bytes
diff --git a/src/hostnet/hostnet_dns.ml b/src/hostnet/hostnet_dns.ml
@@ -7,6 +7,9 @@ let src =
 
 module Log = (val Logs.src_log src : Logs.LOG)
 
+(* Maximum size of a UDP DNS response before we must truncate *)
+let max_udp_response = 512
+
 module Config = struct
   type t = [
     | `Upstream of Dns_forward.Config.t
@@ -326,17 +329,35 @@ struct
       Log.info (fun f -> f "Will use the host's DNS resolver");
       Lwt.return { local_ip; builtin_names; resolver = Host }
 
+  let search f low high =
+    if not(f low)
+    then None (* none of the elements satisfy the predicate *)
+    else
+      let rec loop low high =
+        if low = high
+        then Some low
+        else
+          let mid = (low + high + 1) / 2 in
+          (* since low <> high, mid <> low but it might be mid = high *)
+          if f mid
+          then loop mid high
+          else
+            if mid = high
+            then Some low
+            else loop low mid in
+      loop low high
+
   let answer t is_tcp buf =
     let open Dns.Packet in
     let len = Cstruct.len buf in
     match Dns.Protocol.Server.parse (Cstruct.sub buf 0 len) with
     | None ->
       Lwt.return (Error (`Msg "failed to parse DNS packet"))
     | Some ({ questions = [ question ]; _ } as request) ->
-      let reply answers =
+      let reply ~tc answers =
         let id = request.id in
         let detail =
-          { request.detail with Dns.Packet.qr = Dns.Packet.Response; ra = true }
+          { request.detail with Dns.Packet.qr = Dns.Packet.Response; ra = true; tc }
         in
         let questions = request.questions in
         let authorities = [] and additionals = [] in
@@ -354,31 +375,70 @@ struct
         { Dns.Packet.id; detail; questions; answers; authorities;
           additionals }
       in
+      let marshal_reply answers =
+        let buf = marshal @@ reply ~tc:false answers in
+        if is_tcp
+        then Some buf (* No need to truncate for TCP *)
+        else begin
+          (* If the packet is too big then set the TC bit and truncate by dropping answers *)
+          let take n from =
+            let rec loop n from acc = match n, from with
+              | 0, _ -> acc
+              | _, [] -> acc
+              | n, x :: xs -> loop (n - 1) xs (x :: acc) in
+            List.rev @@ loop n from [] in
+          if Cstruct.len buf > max_udp_response then begin
+            match search (fun num ->
+              (* use only the first 'num' answers *)
+              Cstruct.len (marshal @@ reply ~tc:true (take num answers)) <= max_udp_response
+            ) 0 (List.length answers) with
+            | None -> None
+            | Some num -> Some (marshal @@ reply ~tc:true (take num answers))
+          end
+          else Some buf
+        end in
       begin
         (* Consider the builtins (from the command-line) to have higher priority
            than the addresses in the /etc/hosts file. *)
         match try_builtins t.builtin_names question with
         | `Does_not_exist ->
-          Lwt.return (Ok (marshal nxdomain))
+          Lwt.return (Ok (Some (marshal nxdomain)))
         | `Answers answers ->
-          Lwt.return (Ok (marshal @@ reply answers))
+          Lwt.return (Ok (marshal_reply answers))
         | `Dont_know ->
           match try_etc_hosts question with
           | Some answers ->
-            Lwt.return (Ok (marshal @@ reply answers))
+            Lwt.return (Ok (marshal_reply answers))
           | None ->
             match is_tcp, t.resolver with
             | true, Upstream { dns_tcp_resolver; _ } ->
-              Dns_tcp_resolver.answer buf dns_tcp_resolver
+              begin
+                Dns_tcp_resolver.answer buf dns_tcp_resolver
+                >>= function
+                | Error e -> Lwt.return (Error e)
+                | Ok buf -> Lwt.return (Ok (Some buf))
+              end
             | false, Upstream { dns_udp_resolver; _ } ->
-              Dns_udp_resolver.answer buf dns_udp_resolver
+              begin
+                Dns_udp_resolver.answer buf dns_udp_resolver
+                >>= function
+                | Error e -> Lwt.return (Error e)
+                | Ok buf ->
+                  (* We need to parse and re-marshal so we can set the TC bit and truncate *)
+                  begin match Dns.Protocol.Server.parse buf with
+                  | None ->
+                    Lwt.return (Error (`Msg "Failed to unmarshal DNS response from upstream"))
+                  | Some { answers; _ } ->
+                    Lwt.return (Ok (marshal_reply answers))
+                  end
+              end
             | _, Host ->
               D.resolve question
               >>= function
               | [] ->
-                Lwt.return (Ok (marshal nxdomain))
+                Lwt.return (Ok (Some (marshal nxdomain)))
               | answers ->
-                Lwt.return (Ok (marshal @@ reply answers))
+                Lwt.return (Ok (marshal_reply answers))
       end
     | _ ->
       Lwt.return (Error (`Msg "DNS packet had multiple questions"))
@@ -395,7 +455,10 @@ struct
     | Error (`Msg m) ->
       Log.warn (fun f -> f "%s lookup failed: %s" (describe buf) m);
       Lwt.return (Ok ())
-    | Ok buffer ->
+    | Ok None ->
+      Log.err (fun f -> f "%s unable to marshal response" (describe buf));
+      Lwt.return (Ok ())
+    | Ok (Some buffer) ->
       Udp.write ~src_port:53 ~dst:src ~dst_port:src_port udp buffer
 
   let handle_tcp ~t =
@@ -414,7 +477,10 @@ struct
               | Error (`Msg m) ->
                 Log.warn (fun f -> f "%s lookup failed: %s" (describe request) m);
                 Lwt.return_unit
-              | Ok buffer ->
+              | Ok None ->
+                Log.err (fun f -> f "%s unable to marshal response to" (describe request));
+                Lwt.return_unit
+              | Ok (Some buffer) ->
                 Dns_tcp_framing.write packets buffer >>= function
                 | Error (`Msg m) ->
                   Log.warn (fun f ->
diff --git a/src/hostnet_test/slirp_stack.ml b/src/hostnet_test/slirp_stack.ml
@@ -39,6 +39,8 @@ module Dns_policy = struct
 
   let t = ref (IntMap.add 0 google_dns IntMap.empty)
 
+  let clear () = t := (IntMap.add 0 google_dns IntMap.empty)
+
   let config () =
     snd @@ IntMap.max_binding !t
 
diff --git a/src/hostnet_test/suite.ml b/src/hostnet_test/suite.ml
@@ -28,94 +28,6 @@ let test_dhcp_query () =
   in
   run ~pcap:"test_dhcp_query.pcap" t
 
-let set_dns_policy ?builtin_names use_host =
-  Mclock.connect () >|= fun clock ->
-  Dns_policy.remove ~priority:3;
-  Dns_policy.add ~priority:3
-    ~config:(if use_host then `Host else Dns_policy.google_dns);
-  Slirp_stack.Debug.update_dns ?builtin_names clock
-
-let test_dns_query server use_host () =
-  let t _ stack =
-    set_dns_policy use_host >>= fun () ->
-    let resolver = DNS.create stack.Client.t in
-    DNS.gethostbyname ~server resolver "www.google.com" >|= function
-    | (_ :: _) as ips ->
-      Log.info (fun f -> f "www.google.com has IPs: %a" pp_ips ips);
-    | _ ->
-      Log.err (fun f -> f "Failed to lookup www.google.com");
-      failwith "Failed to lookup www.google.com"
-  in
-  run ~pcap:"test_dns_query.pcap" t
-
-let test_builtin_dns_query server use_host () =
-  let name = "experimental.host.name.localhost" in
-  let t _ stack =
-    set_dns_policy ~builtin_names:[ Dns.Name.of_string name, Ipaddr.V4 (Ipaddr.V4.localhost) ] use_host
-    >>= fun () ->
-    let resolver = DNS.create stack.Client.t in
-    DNS.gethostbyname ~server resolver name >>= function
-    | (_ :: _) as ips ->
-      Log.info (fun f -> f "%s has IPs: %a" name pp_ips ips);
-      Lwt.return ()
-    | _ ->
-      Log.err (fun f -> f "Failed to lookup %s" name);
-      failwith ("Failed to lookup " ^ name)
-  in
-  run ~pcap:"test_builtin_dns_query.pcap" t
-
-let test_etc_hosts_query server use_host () =
-  let test_name = "vpnkit.is.cool.yes.really" in
-  let t _ stack =
-    set_dns_policy use_host >>= fun () ->
-    let resolver = DNS.create stack.Client.t in
-    DNS.gethostbyname ~server resolver test_name >>= function
-    | (_ :: _) as ips ->
-      Log.err (fun f ->
-          f "This test relies on the name %s not existing but it really \
-             has IPs: %a" test_name pp_ips ips);
-      Fmt.kstrf failwith "Test name %s really does exist" test_name
-    | _ ->
-      Hosts.etc_hosts := [
-        test_name, Ipaddr.V4 (Ipaddr.V4.localhost);
-      ];
-      DNS.gethostbyname ~server resolver test_name >|= function
-      | (_ :: _) as ips ->
-        Log.info (fun f -> f "Name %s has IPs: %a" test_name pp_ips ips);
-        Hosts.etc_hosts := []
-      | _ ->
-        Log.err (fun f -> f "Failed to lookup name from /etc/hosts");
-        Hosts.etc_hosts := [];
-        failwith "failed to lookup name from /etc/hosts"
-  in
-  run ~pcap:"test_etc_hosts_query.pcap" t
-
-let test_etc_hosts_priority server use_host () =
-  let name = "builtins.should.be.higher.priority" in
-  let builtin_ip = Ipaddr.of_string_exn "127.0.0.1" in
-  let hosts_ip = Ipaddr.of_string_exn "127.0.0.2" in
-  let t _ stack =
-    set_dns_policy ~builtin_names:[ Dns.Name.of_string name, builtin_ip ] use_host
-    >>= fun () ->
-    Hosts.etc_hosts := [
-      name, hosts_ip;
-    ];
-    let resolver = DNS.create stack.Client.t in
-    DNS.gethostbyname ~server resolver name >>= function
-    | [ ip ] ->
-      Log.info (fun f -> f "%s has single IP: %a" name Ipaddr.pp_hum ip);
-      if Ipaddr.compare ip builtin_ip = 0
-      then Lwt.return ()
-      else failwith ("Builtin DNS names should have higher priority than /etc/hosts")
-    | (_ :: _) as ips ->
-      Log.info (fun f -> f "%s has IPs: %a" name pp_ips ips);
-      failwith ("Duplicate DNS names resolved for " ^ name);
-    | _ ->
-      Log.err (fun f -> f "Failed to lookup %s" name);
-      failwith ("Failed to lookup " ^ name)
-  in
-  run ~pcap:"test_etc_hosts_priority.pcap" t
-
 let test_max_connections () =
   let t _ stack =
     Lwt.finalize (fun () ->
@@ -373,21 +285,6 @@ let test_dhcp = [
   ["check that the DHCP server works", `Quick, test_dhcp_query];
 ]
 
-let test_dns use_host =
-  let prefix = if use_host then "Host resolver" else "DNS forwarder" in [
-    prefix ^ ": lookup ",
-    ["", `Quick, test_dns_query primary_dns_ip use_host];
-
-    prefix ^ ": builtins",
-    [ "", `Quick, test_builtin_dns_query primary_dns_ip use_host ];
-
-    prefix ^ ": _etc_hosts",
-    [ "", `Quick, test_etc_hosts_query primary_dns_ip use_host ];
-
-    prefix ^ ": _etc_hosts_priority",
-    [ "", `Quick, test_etc_hosts_priority primary_dns_ip use_host ];
-  ]
-
 let test_tcp = [
   "HTTP GET", [ "HTTP GET http://www.google.com/", `Quick, test_http_fetch ];
 
@@ -412,7 +309,7 @@ let test_tcp = [
 
 let tests =
   Hosts_test.tests @ Forwarding.tests @ test_dhcp
-  @ (test_dns true) @ (test_dns false)
+  @ Test_dns.suite
   @ test_tcp @ Test_nat.tests @ Test_http.tests @ Test_http.Exclude.tests
   @ Test_half_close.tests @ Test_ping.tests
   @ Test_bridge.tests @ Test_forward_protocol.suite
diff --git a/src/hostnet_test/test_dns.ml b/src/hostnet_test/test_dns.ml
