lib/dns: make the TCP forwarders as complementary of UDP

The TCP forwarders only active when client send the DNS request as TCP.
When the server receive that request it should also forward the request
as TCP not as UDP to prevent the truncated response.

Another use case for TCP is when the response is truncated, the client
will send the query back through TCP connection.  The server should
forward this request using TCP instead of UDP.

3 files changed, 60 insertions, 28 deletions

diff --git a/lib/dns/server.go b/lib/dns/server.go
index 43640e0e..ea6cb257 100644
--- a/lib/dns/server.go
+++ b/lib/dns/server.go
@@ -89,6 +89,7 @@ type Server struct {
 
 	requestq chan *request
 	primaryq chan *request
+	tcpq     chan *request
 
 	fwLocker   sync.Mutex
 	fwStoppers []chan bool
@@ -108,6 +109,7 @@ func NewServer(opts *ServerOptions) (srv *Server, err error) {
 		opts:     opts,
 		requestq: make(chan *request, 512),
 		primaryq: make(chan *request, 512),
+		tcpq:     make(chan *request, 512),
 	}
 
 	udpAddr := opts.getUDPAddress()
@@ -614,7 +616,9 @@ func (srv *Server) serveTCPClient(cl *TCPClient, kind connType) {
 
 		req.message, err = cl.recv()
 		if err != nil {
-			log.Printf("serveTCPClient: %s: %s", connTypeNames[kind], err)
+			if !errors.Is(err, io.EOF) {
+				log.Printf("serveTCPClient: %s: %s", connTypeNames[kind], err)
+			}
 			break
 		}
 
@@ -682,7 +686,11 @@ func (srv *Server) processRequest() {
 		if ans == nil || an == nil {
 			switch {
 			case srv.hasForwarders():
-				srv.primaryq <- req
+				if req.kind == connTypeTCP {
+					srv.tcpq <- req
+				} else {
+					srv.primaryq <- req
+				}
 			default:
 				if debug.Value >= 1 {
 					fmt.Printf("dns: * %s %d:%s\n",
@@ -704,7 +712,11 @@ func (srv *Server) processRequest() {
 						req.message.Header.ID,
 						req.message.Question.String())
 				}
-				srv.primaryq <- req
+				if req.kind == connTypeTCP {
+					srv.tcpq <- req
+				} else {
+					srv.primaryq <- req
+				}
 
 			default:
 				if debug.Value >= 1 {
@@ -779,26 +791,26 @@ func (srv *Server) startAllForwarders() {
 	for x := 0; x < len(srv.opts.primaryUDP); x++ {
 		tag := fmt.Sprintf("UDP-%d-%s", x, asPrimary)
 		nameserver := srv.opts.primaryUDP[x].String()
-		go srv.runUDPForwarder(tag, nameserver, srv.primaryq)
+		go srv.runUDPForwarder(tag, nameserver)
 	}
 	for x := 0; x < len(srv.opts.primaryTCP); x++ {
 		tag := fmt.Sprintf("TCP-%d-%s", x, asPrimary)
 		nameserver := srv.opts.primaryTCP[x].String()
-		go srv.runTCPForwarder(tag, nameserver, srv.primaryq)
+		go srv.runTCPForwarder(tag, nameserver)
 	}
 	for x := 0; x < len(srv.opts.primaryDoh); x++ {
 		tag := fmt.Sprintf("DoH-%d-%s", x, asPrimary)
 		nameserver := srv.opts.primaryDoh[x]
-		go srv.runDohForwarder(tag, nameserver, srv.primaryq)
+		go srv.runDohForwarder(tag, nameserver)
 	}
 	for x := 0; x < len(srv.opts.primaryDot); x++ {
 		tag := fmt.Sprintf("DoT-%d-%s", x, asPrimary)
 		nameserver := srv.opts.primaryDot[x]
-		go srv.runTLSForwarder(tag, nameserver, srv.primaryq)
+		go srv.runTLSForwarder(tag, nameserver)
 	}
 }
 
-func (srv *Server) runDohForwarder(tag, nameserver string, primaryq <-chan *request) {
+func (srv *Server) runDohForwarder(tag, nameserver string) {
 	stopper := srv.newStopper()
 
 	defer func() {
@@ -828,7 +840,7 @@ func (srv *Server) runDohForwarder(tag, nameserver string, primaryq <-chan *requ
 		ticker := time.NewTicker(aliveInterval)
 		for isRunning {
 			select {
-			case req, ok := <-primaryq:
+			case req, ok := <-srv.primaryq:
 				if !ok {
 					log.Println("dns: primary queue has been closed")
 					srv.stopForwarder(forwarder)
@@ -863,7 +875,7 @@ func (srv *Server) runDohForwarder(tag, nameserver string, primaryq <-chan *requ
 	}
 }
 
-func (srv *Server) runTLSForwarder(tag, nameserver string, primaryq <-chan *request) {
+func (srv *Server) runTLSForwarder(tag, nameserver string) {
 	stopper := srv.newStopper()
 
 	defer func() {
@@ -893,7 +905,7 @@ func (srv *Server) runTLSForwarder(tag, nameserver string, primaryq <-chan *requ
 		ticker := time.NewTicker(aliveInterval)
 		for isRunning {
 			select {
-			case req, ok := <-primaryq:
+			case req, ok := <-srv.primaryq:
 				if !ok {
 					log.Println("dns: primary queue has been closed")
 					srv.stopForwarder(forwarder)
@@ -929,7 +941,7 @@ func (srv *Server) runTLSForwarder(tag, nameserver string, primaryq <-chan *requ
 	}
 }
 
-func (srv *Server) runTCPForwarder(tag, nameserver string, primaryq <-chan *request) {
+func (srv *Server) runTCPForwarder(tag, nameserver string) {
 	stopper := srv.newStopper()
 
 	log.Printf("dns: starting forwarder %s for %s", tag, nameserver)
@@ -944,7 +956,7 @@ func (srv *Server) runTCPForwarder(tag, nameserver string, primaryq <-chan *requ
 	ticker := time.NewTicker(aliveInterval)
 	for {
 		select {
-		case req, ok := <-primaryq:
+		case req, ok := <-srv.tcpq:
 			if !ok {
 				log.Println("dns: primary queue has been closed")
 				return
@@ -981,10 +993,10 @@ func (srv *Server) runTCPForwarder(tag, nameserver string, primaryq <-chan *requ
 }
 
 //
-// runUDPForwarder create a UDP client that consume request from forward queue
-// and forward it to parent server at "nameserver".
+// runUDPForwarder create a UDP client that consume request from queue
+// and forward it to parent name server.
 //
-func (srv *Server) runUDPForwarder(tag, nameserver string, primaryq <-chan *request) {
+func (srv *Server) runUDPForwarder(tag, nameserver string) {
 	stopper := srv.newStopper()
 
 	defer func() {
@@ -1016,7 +1028,7 @@ func (srv *Server) runUDPForwarder(tag, nameserver string, primaryq <-chan *requ
 		ticker := time.NewTicker(aliveInterval)
 		for isRunning {
 			select {
-			case req, ok := <-primaryq:
+			case req, ok := <-srv.primaryq:
 				if !ok {
 					log.Println("dns: primary queue has been closed")
 					srv.stopForwarder(forwarder)
diff --git a/lib/dns/server_options.go b/lib/dns/server_options.go
index b038acf5..b8ccada9 100644
--- a/lib/dns/server_options.go
+++ b/lib/dns/server_options.go
@@ -46,7 +46,7 @@ type ServerOptions struct {
 	// The name server use the URI format,
 	//
 	//	nameserver  = [ scheme "://" ] ( ip-address / domain-name ) [:port]
-	//	scheme      = ( "udp" / "tcp" / "https" )
+	//	scheme      = ( "udp" / "https" )
 	//	ip-address  = ( ip4 / ip6 )
 	//	domain-name = ; fully qualified domain name
 	//
@@ -56,7 +56,6 @@ type ServerOptions struct {
 	// Example,
 	//
 	//	udp://35.240.172.103
-	//	tcp://35.240.172.103:5353
 	//	https://35.240.172.103:853 (DNS over TLS)
 	//	https://kilabit.info/dns-query (DNS over HTTPS)
 	//
@@ -147,7 +146,7 @@ func (opts *ServerOptions) init() (err error) {
 
 	opts.initNameServers()
 
-	if len(opts.primaryUDP) == 0 && len(opts.primaryTCP) == 0 && len(opts.primaryDoh) == 0 && len(opts.primaryDot) == 0 {
+	if len(opts.primaryUDP) == 0 && len(opts.primaryDoh) == 0 && len(opts.primaryDot) == 0 {
 		return fmt.Errorf("dns: no valid name servers")
 	}
 
@@ -201,10 +200,9 @@ func (opts *ServerOptions) parseNameServers(nameServers []string) {
 		case "tcp":
 			tcpAddr, err := libnet.ParseTCPAddr(dnsURL.Host, DefaultPort)
 			if err != nil {
-				log.Printf("dns: invalid TCP IP address %q", dnsURL.Host)
+				log.Printf("dns: invalid IP address %q", dnsURL.Host)
 				continue
 			}
-
 			opts.primaryTCP = append(opts.primaryTCP, tcpAddr)
 
 		case "https":
@@ -222,11 +220,17 @@ func (opts *ServerOptions) parseNameServers(nameServers []string) {
 
 			udpAddr, err := libnet.ParseUDPAddr(ns, DefaultPort)
 			if err != nil {
-				log.Printf("dns: invalid UDP IP address %q", ns)
+				log.Printf("dns: invalid IP address %q", ns)
 				continue
 			}
-
 			opts.primaryUDP = append(opts.primaryUDP, udpAddr)
+
+			tcpAddr, err := libnet.ParseTCPAddr(ns, DefaultPort)
+			if err != nil {
+				log.Printf("dns: invalid IP address %q", ns)
+				continue
+			}
+			opts.primaryTCP = append(opts.primaryTCP, tcpAddr)
 		}
 	}
 }
diff --git a/lib/dns/server_options_test.go b/lib/dns/server_options_test.go
index 39ff200b..5f899c54 100644
--- a/lib/dns/server_options_test.go
+++ b/lib/dns/server_options_test.go
@@ -68,6 +68,12 @@ func TestServerOptionsInit(t *testing.T) {
 					Port: 53,
 				},
 			},
+			primaryTCP: []net.Addr{
+				&net.TCPAddr{
+					IP:   net.ParseIP("127.0.0.1"),
+					Port: 53,
+				},
+			},
 		},
 	}}
 
@@ -120,6 +126,10 @@ func TestServerOptionsParseNameServers(t *testing.T) {
 			IP:   ip,
 			Port: 53,
 		}},
+		expTCPServers: []net.Addr{&net.TCPAddr{
+			IP:   ip,
+			Port: 53,
+		}},
 	}, {
 		desc: "With valid name servers",
 		nameServers: []string{
@@ -131,10 +141,16 @@ func TestServerOptionsParseNameServers(t *testing.T) {
 			IP:   ip,
 			Port: 53,
 		}},
-		expTCPServers: []net.Addr{&net.TCPAddr{
-			IP:   ip,
-			Port: 5353,
-		}},
+		expTCPServers: []net.Addr{
+			&net.TCPAddr{
+				IP:   ip,
+				Port: 53,
+			},
+			&net.TCPAddr{
+				IP:   ip,
+				Port: 5353,
+			},
+		},
 		expDoHServers: []string{
 			"https://localhost/dns-query",
 		},