feat: Add option to define which DNS server a DNS monitor queries (#2389)

agent/network_monitor_probe.go

@@ -28,7 +28,7 @@ func networkMonitorProbe(client *http.Client) monitorProbe {
 		case "http":
 			return monitorHTTP(ctx, client, config.Target)
 		case "dns":
-			return monitorDNS(ctx, config.Target)
+			return monitorDNS(ctx, config.Target, config.Server)
 		default:
 			return -1, fmt.Errorf("unknown monitor protocol: %s", config.Protocol)
 		}
@@ -73,19 +73,43 @@ func monitorTCP(ctx context.Context, target string, port uint16) (int64, error)
 	return -1, err
 }
 
-// monitorDNS measures DNS resolution response time in microseconds. Returns -1 and an error on failure.
-func monitorDNS(ctx context.Context, target string) (int64, error) {
+// monitorDNS measures DNS resolution response time in microseconds. If server is
+// non-empty, the lookup is sent to that DNS server (host or host:port, default
+// port 53) instead of the system resolver. Returns -1 and an error on failure.
+func monitorDNS(ctx context.Context, target, server string) (int64, error) {
 	ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
 	defer cancel()
 
+	resolver := net.DefaultResolver
+	if server != "" {
+		resolver = dnsResolverForServer(server)
+	}
+
 	start := time.Now()
-	ips, err := net.DefaultResolver.LookupHost(ctx, target)
+	ips, err := resolver.LookupHost(ctx, target)
 	if err != nil || len(ips) == 0 {
 		return -1, err
 	}
 	return time.Since(start).Microseconds(), nil
 }
 
+// dnsResolverForServer builds a resolver that sends lookups to the given DNS
+// server address instead of the system resolver. server may be a bare host or
+// host:port; when no port is given, the standard DNS port 53 is used.
+func dnsResolverForServer(server string) *net.Resolver {
+	address := server
+	if _, _, err := net.SplitHostPort(server); err != nil {
+		address = net.JoinHostPort(server, "53")
+	}
+	return &net.Resolver{
+		PreferGo: true,
+		Dial: func(ctx context.Context, network, _ string) (net.Conn, error) {
+			var dialer net.Dialer
+			return dialer.DialContext(ctx, network, address)
+		},
+	}
+}
+
 // monitorHTTP measures HTTP GET request response in microseconds. Returns -1 and an error on failure.
 func monitorHTTP(ctx context.Context, client *http.Client, url string) (int64, error) {
 	if client == nil {

agent/network_monitor_test.go

@@ -376,15 +376,79 @@ func tcpMonitorTestResolver(ips []string) *net.Resolver {
 	}}
 }
 
+// udpDNSTestServer starts a UDP server on loopback that answers A queries with the
+// given IPs, and returns its listen address (host:port).
+func udpDNSTestServer(t *testing.T, ips []string) string {
+	t.Helper()
+	conn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP("127.0.0.1")})
+	require.NoError(t, err)
+	t.Cleanup(func() { conn.Close() })
+
+	go func() {
+		buf := make([]byte, 512)
+		for {
+			n, addr, err := conn.ReadFromUDP(buf)
+			if err != nil {
+				return
+			}
+			var msg dnsmessage.Message
+			if err := msg.Unpack(buf[:n]); err != nil {
+				continue
+			}
+			msg.Header.Response = true
+			msg.Header.RecursionAvailable = true
+			for _, question := range msg.Questions {
+				if question.Type != dnsmessage.TypeA {
+					continue
+				}
+				for _, ip := range ips {
+					msg.Answers = append(msg.Answers, dnsmessage.Resource{
+						Header: dnsmessage.ResourceHeader{Name: question.Name, Type: dnsmessage.TypeA, Class: dnsmessage.ClassINET},
+						Body:   &dnsmessage.AResource{A: [4]byte(net.ParseIP(ip).To4())},
+					})
+				}
+			}
+			packet, err := msg.Pack()
+			if err != nil {
+				continue
+			}
+			_, _ = conn.WriteToUDP(packet, addr)
+		}
+	}()
+
+	return conn.LocalAddr().String()
+}
+
 func TestMonitorDNS(t *testing.T) {
 	t.Run("success", func(t *testing.T) {
-		responseUs, err := monitorDNS(context.Background(), "localhost")
+		responseUs, err := monitorDNS(context.Background(), "localhost", "")
 		require.NoError(t, err)
 		assert.GreaterOrEqual(t, responseUs, int64(0))
 	})
 
 	t.Run("lookup failure", func(t *testing.T) {
-		responseUs, err := monitorDNS(context.Background(), "")
+		responseUs, err := monitorDNS(context.Background(), "", "")
+		assert.Equal(t, int64(-1), responseUs)
+		require.Error(t, err)
+	})
+
+	t.Run("custom server", func(t *testing.T) {
+		serverAddr := udpDNSTestServer(t, []string{"192.0.2.10"})
+		responseUs, err := monitorDNS(context.Background(), "example.test.", serverAddr)
+		require.NoError(t, err)
+		assert.GreaterOrEqual(t, responseUs, int64(0))
+	})
+
+	t.Run("custom server without port defaults to 53", func(t *testing.T) {
+		resolver := dnsResolverForServer("127.0.0.1")
+		conn, err := resolver.Dial(context.Background(), "udp", "")
+		require.NoError(t, err)
+		defer conn.Close()
+		assert.Equal(t, "127.0.0.1:53", conn.RemoteAddr().String())
+	})
+
+	t.Run("custom server unreachable", func(t *testing.T) {
+		responseUs, err := monitorDNS(context.Background(), "example.test.", "127.0.0.1:1")
 		assert.Equal(t, int64(-1), responseUs)
 		require.Error(t, err)
 	})
@@ -479,7 +543,7 @@ func TestMonitorResolutionCancellation(t *testing.T) {
 				case "tcp":
 					_, err = monitorTCP(ctx, "monitor-cancellation.invalid.", 80)
 				case "dns":
-					_, err = monitorDNS(ctx, "monitor-cancellation.invalid.")
+					_, err = monitorDNS(ctx, "monitor-cancellation.invalid.", "")
 				case "icmp":
 					_, err = monitorICMP(ctx, "monitor-cancellation.invalid.")
 				}