Separate UDP and TCP queries for DNS Server

pkg/tools/dnsutils/dnsconfigs/handler.go

@@ -30,22 +30,21 @@ import (
 	"github.com/networkservicemesh/sdk/pkg/tools/dnsutils"
 	"github.com/networkservicemesh/sdk/pkg/tools/dnsutils/next"
 	"github.com/networkservicemesh/sdk/pkg/tools/dnsutils/searches"
+	"github.com/networkservicemesh/sdk/pkg/tools/log"
 )
 
 type dnsConfigsHandler struct {
 	configs *dnsconfig.Map
 }
 
-func (h *dnsConfigsHandler) ServeDNS(ctx context.Context, rp dns.ResponseWriter, m *dns.Msg) {
+func (h *dnsConfigsHandler) ServeDNS(ctx context.Context, rw dns.ResponseWriter, m *dns.Msg) {
 	dnsIPs := make([]url.URL, 0)
 	searchDomains := make([]string, 0)
 
 	h.configs.Range(func(key string, value []*networkservice.DNSConfig) bool {
 		for _, conf := range value {
 			for _, ip := range conf.DnsServerIps {
-				dnsIPs = append(dnsIPs,
-					url.URL{Scheme: "udp", Host: ip},
-					url.URL{Scheme: "tcp", Host: ip})
+				dnsIPs = append(dnsIPs, url.URL{Scheme: "udp", Host: ip})
 			}
 			searchDomains = append(searchDomains, conf.SearchDomains...)
 		}
@@ -55,7 +54,36 @@ func (h *dnsConfigsHandler) ServeDNS(ctx context.Context, rp dns.ResponseWriter,
 
 	ctx = clienturlctx.WithClientURLs(ctx, dnsIPs)
 	ctx = searches.WithSearchDomains(ctx, searchDomains)
-	next.Handler(ctx).ServeDNS(ctx, rp, m)
+
+	udpRW := &responseWriter{Response: nil}
+	next.Handler(ctx).ServeDNS(ctx, udpRW, m)
+
+	if resp := udpRW.Response; resp != nil {
+		if err := rw.WriteMsg(resp); err != nil {
+			log.FromContext(ctx).WithField("dnsConfigHandler", "ServeDNS").Warnf("got an error during writing the message: %v", err.Error())
+			dns.HandleFailed(rw, resp)
+			return
+		}
+		return
+	}
+
+	for i := range dnsIPs {
+		dnsIPs[i].Scheme = "tcp"
+	}
+
+	tcpRW := &responseWriter{Response: nil}
+	next.Handler(ctx).ServeDNS(ctx, tcpRW, m)
+
+	if resp := tcpRW.Response; resp != nil {
+		if err := rw.WriteMsg(resp); err != nil {
+			log.FromContext(ctx).WithField("dnsConfigHandler", "ServeDNS").Warnf("got an error during writing the message: %v", err.Error())
+			dns.HandleFailed(rw, resp)
+			return
+		}
+		return
+	}
+
+	dns.HandleFailed(rw, m)
 }
 
 // NewDNSHandler creates a new dns handler that stores DNS configs

pkg/tools/dnsutils/dnsconfigs/handler_test.go

@@ -32,6 +32,16 @@ import (
 	"github.com/networkservicemesh/sdk/pkg/tools/dnsutils/searches"
 )
 
+type responseWriter struct {
+	dns.ResponseWriter
+	Response *dns.Msg
+}
+
+func (r *responseWriter) WriteMsg(m *dns.Msg) error {
+	r.Response = m
+	return nil
+}
+
 type checkHandler struct {
 	Domains []string
 	URLs    []string
@@ -77,7 +87,8 @@ func TestDNSConfigs(t *testing.T) {
 		check,
 	)
 
-	handler.ServeDNS(ctx, nil, new(dns.Msg))
+	r := &responseWriter{}
+	handler.ServeDNS(ctx, r, new(dns.Msg))
 
 	domains := check.Domains
 	require.Equal(t, len(domains), 3)
@@ -88,9 +99,9 @@ func TestDNSConfigs(t *testing.T) {
 	urls := check.URLs
 	require.Equal(t, len(urls), 6)
 	require.Contains(t, urls, "udp://7.7.7.7")
-	require.Contains(t, urls, "tcp://7.7.7.7")
 	require.Contains(t, urls, "udp://1.1.1.1")
-	require.Contains(t, urls, "tcp://1.1.1.1")
 	require.Contains(t, urls, "udp://9.9.9.9")
+	require.Contains(t, urls, "tcp://7.7.7.7")
+	require.Contains(t, urls, "tcp://1.1.1.1")
 	require.Contains(t, urls, "tcp://9.9.9.9")
 }

pkg/tools/dnsutils/dnsconfigs/response_writer.go

@@ -0,0 +1,31 @@
+// Copyright (c) 2022 Cisco and/or its affiliates.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at:
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package dnsconfigs
+
+import (
+	"github.com/miekg/dns"
+)
+
+type responseWriter struct {
+	dns.ResponseWriter
+	Response *dns.Msg
+}
+
+func (r *responseWriter) WriteMsg(m *dns.Msg) error {
+	r.Response = m
+	return nil
+}

pkg/tools/dnsutils/dnsconfigs/tcp_dns_server_test.go

@@ -0,0 +1,184 @@
+// Copyright (c) 2022 Cisco and/or its affiliates.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at:
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build !windows
+// +build !windows
+
+package dnsconfigs_test
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/miekg/dns"
+	"github.com/networkservicemesh/api/pkg/api/networkservice"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/goleak"
+
+	"github.com/networkservicemesh/sdk/pkg/tools/dnsconfig"
+	"github.com/networkservicemesh/sdk/pkg/tools/dnsutils"
+	"github.com/networkservicemesh/sdk/pkg/tools/dnsutils/dnsconfigs"
+	"github.com/networkservicemesh/sdk/pkg/tools/dnsutils/fanout"
+	"github.com/networkservicemesh/sdk/pkg/tools/dnsutils/next"
+	"github.com/networkservicemesh/sdk/pkg/tools/dnsutils/noloop"
+	"github.com/networkservicemesh/sdk/pkg/tools/dnsutils/searches"
+)
+
+func requireIPv4Lookup(ctx context.Context, t *testing.T, r *net.Resolver, host, expected string) {
+	addrs, err := r.LookupIP(ctx, "ip4", host)
+	require.NoError(t, err)
+	require.Len(t, addrs, 1)
+	require.Equal(t, expected, addrs[0].String())
+}
+
+type proxyDNSServer struct {
+	ListenOn  string
+	TCPServer *dns.Server
+	UDPServer *dns.Server
+}
+
+func (p *proxyDNSServer) listenAndServe(t *testing.T, tcpHandler, udpHandler dns.Handler) {
+	p.TCPServer = &dns.Server{Addr: p.ListenOn, Net: "tcp", Handler: tcpHandler}
+	p.UDPServer = &dns.Server{Addr: p.ListenOn, Net: "udp", Handler: udpHandler}
+
+	go func() {
+		err := p.TCPServer.ListenAndServe()
+		require.NoError(t, err)
+	}()
+
+	go func() {
+		err := p.UDPServer.ListenAndServe()
+		require.NoError(t, err)
+	}()
+}
+
+func (p *proxyDNSServer) shutdown() error {
+	tcpErr := p.TCPServer.Shutdown()
+	udpErr := p.UDPServer.Shutdown()
+
+	if tcpErr != nil {
+		return tcpErr
+	}
+
+	if udpErr != nil {
+		return udpErr
+	}
+
+	return nil
+}
+
+type tcpHandler struct {
+}
+
+func (h *tcpHandler) ServeDNS(rw dns.ResponseWriter, m *dns.Msg) {
+	time.Sleep(time.Second * 5)
+	dns.HandleFailed(rw, m)
+}
+
+type udpHandler struct {
+}
+
+func (h *udpHandler) ServeDNS(rw dns.ResponseWriter, m *dns.Msg) {
+	if m.Question[0].Name == "my.domain." {
+		dns.HandleFailed(rw, m)
+		return
+	}
+
+	name := dns.Name(m.Question[0].Name).String()
+	rr := new(dns.A)
+	rr.Hdr = dns.RR_Header{Name: name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 3600}
+	rr.A = net.ParseIP("1.1.1.1")
+
+	resp := new(dns.Msg)
+	resp.SetReply(m)
+	resp.Authoritative = true
+	resp.Answer = append(resp.Answer, rr)
+
+	if err := rw.WriteMsg(resp); err != nil {
+		dns.HandleFailed(rw, m)
+	}
+}
+
+func getFreePort() (int, error) {
+	addr, err := net.ResolveTCPAddr("tcp", "localhost:0")
+	if err != nil {
+		return 0, err
+	}
+
+	l, err := net.ListenTCP("tcp", addr)
+	if err != nil {
+		return 0, err
+	}
+	err = l.Close()
+	if err != nil {
+		return 0, err
+	}
+
+	return l.Addr().(*net.TCPAddr).Port, nil
+}
+
+func Test_TCPDNSServerTimeout(t *testing.T) {
+	t.Cleanup(func() { goleak.VerifyNone(t) })
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
+	defer cancel()
+
+	proxyPort, err := getFreePort()
+	require.NoError(t, err)
+	proxyAddr := fmt.Sprintf("127.0.0.1:%d", proxyPort)
+
+	proxy := &proxyDNSServer{ListenOn: proxyAddr}
+	proxy.listenAndServe(t, &tcpHandler{}, &udpHandler{})
+
+	clientPort, err := getFreePort()
+	require.NoError(t, err)
+	clientAddr := fmt.Sprintf("127.0.0.1:%d", clientPort)
+
+	dnsConfigsMap := new(dnsconfig.Map)
+	dnsConfigsMap.Store("1", []*networkservice.DNSConfig{
+		{
+			DnsServerIps:  []string{proxyAddr},
+			SearchDomains: []string{"com"},
+		},
+	})
+
+	clientDNSHandler := next.NewDNSHandler(
+		dnsconfigs.NewDNSHandler(dnsConfigsMap),
+		searches.NewDNSHandler(),
+		noloop.NewDNSHandler(),
+		fanout.NewDNSHandler(),
+	)
+	dnsutils.ListenAndServe(ctx, clientDNSHandler, clientAddr)
+
+	resolver := net.Resolver{
+		PreferGo: true,
+		Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
+			var dialer net.Dialer
+			return dialer.DialContext(ctx, network, clientAddr)
+		},
+	}
+
+	resolveCtx, resolveCancel := context.WithTimeout(context.Background(), time.Second*5)
+	defer resolveCancel()
+
+	requireIPv4Lookup(resolveCtx, t, &resolver, "my.domain", "1.1.1.1")
+
+	err = proxy.shutdown()
+	require.NoError(t, err)
+}

pkg/tools/dnsutils/searches/handler.go

@@ -19,7 +19,6 @@ package searches
 
 import (
 	"context"
-	"time"
 
 	"github.com/miekg/dns"
 
@@ -28,10 +27,6 @@ import (
 	"github.com/networkservicemesh/sdk/pkg/tools/log"
 )
 
-const (
-	timeout = 5 * time.Second
-)
-
 type searchDomainsHandler struct {
 }
 
@@ -41,12 +36,9 @@ func (h *searchDomainsHandler) ServeDNS(ctx context.Context, rw dns.ResponseWrit
 	r := &responseWriter{
 		ResponseWriter: rw,
 		Responses:      make([]*dns.Msg, len(domains)+1),
-		index:          0,
+		Index:          0,
 	}
 
-	ctx, cancel := context.WithTimeout(ctx, timeout)
-	defer cancel()
-
 	next.Handler(ctx).ServeDNS(ctx, r, m)
 
 	for _, d := range SearchDomains(ctx) {

pkg/tools/dnsutils/searches/response_writer.go

@@ -23,11 +23,11 @@ import (
 type responseWriter struct {
 	dns.ResponseWriter
 	Responses []*dns.Msg
-	index     int
+	Index     int
 }
 
 func (r *responseWriter) WriteMsg(m *dns.Msg) error {
-	r.Responses[r.index] = m
-	r.index++
+	r.Responses[r.Index] = m
+	r.Index++
 	return nil
 }