net: support all PacketConn and Conn returned by Resolver.Dial

Allow the Resolver.Dial func to return instances of Conn other than
*TCPConn and *UDPConn. If the Conn is also a PacketConn, assume DNS
messages transmitted over the Conn adhere to section 4.2.1. "UDP usage".
Otherwise, follow section 4.2.2. "TCP usage".

Provides a hook mechanism so that DNS queries generated by the net
package may be answered or modified before being sent to over the
network.

Updates #19910

Change-Id: Ib089a28ad4a1848bbeaf624ae889f1e82d56655b
Reviewed-on: https://go-review.googlesource.com/45153
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>

1 files changed, 66 insertions, 12 deletions

diff --git a/src/net/dnsclient_unix_test.go b/src/net/dnsclient_unix_test.go
index d0ac4302b1..73b628c1b5 100644
--- a/src/net/dnsclient_unix_test.go
+++ b/src/net/dnsclient_unix_test.go
@@ -8,6 +8,7 @@ package net
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"internal/poll"
 	"io/ioutil"
@@ -43,11 +44,14 @@ var dnsTransportFallbackTests = []struct {
 
 func TestDNSTransportFallback(t *testing.T) {
 	fake := fakeDNSServer{
-		rh: func(n, _ string, _ *dnsMsg, _ time.Time) (*dnsMsg, error) {
+		rh: func(n, _ string, q *dnsMsg, _ time.Time) (*dnsMsg, error) {
 			r := &dnsMsg{
 				dnsMsgHdr: dnsMsgHdr{
-					rcode: dnsRcodeSuccess,
+					id:       q.id,
+					response: true,
+					rcode:    dnsRcodeSuccess,
 				},
+				question: q.question,
 			}
 			if n == "udp" {
 				r.truncated = true
@@ -98,8 +102,10 @@ func TestSpecialDomainName(t *testing.T) {
 	fake := fakeDNSServer{func(_, _ string, q *dnsMsg, _ time.Time) (*dnsMsg, error) {
 		r := &dnsMsg{
 			dnsMsgHdr: dnsMsgHdr{
-				id: q.id,
+				id:       q.id,
+				response: true,
 			},
+			question: q.question,
 		}
 
 		switch q.question[0].Name {
@@ -612,8 +618,10 @@ func TestErrorForOriginalNameWhenSearching(t *testing.T) {
 	fake := fakeDNSServer{func(_, _ string, q *dnsMsg, _ time.Time) (*dnsMsg, error) {
 		r := &dnsMsg{
 			dnsMsgHdr: dnsMsgHdr{
-				id: q.id,
+				id:       q.id,
+				response: true,
 			},
+			question: q.question,
 		}
 
 		switch q.question[0].Name {
@@ -751,7 +759,7 @@ type fakeDNSServer struct {
 }
 
 func (server *fakeDNSServer) DialContext(_ context.Context, n, s string) (Conn, error) {
-	return &fakeDNSConn{nil, server, n, s, time.Time{}}, nil
+	return &fakeDNSConn{nil, server, n, s, nil, time.Time{}}, nil
 }
 
 type fakeDNSConn struct {
@@ -759,6 +767,7 @@ type fakeDNSConn struct {
 	server *fakeDNSServer
 	n      string
 	s      string
+	q      *dnsMsg
 	t      time.Time
 }
 
@@ -766,15 +775,45 @@ func (f *fakeDNSConn) Close() error {
 	return nil
 }
 
+func (f *fakeDNSConn) Read(b []byte) (int, error) {
+	resp, err := f.server.rh(f.n, f.s, f.q, f.t)
+	if err != nil {
+		return 0, err
+	}
+
+	bb, ok := resp.Pack()
+	if !ok {
+		return 0, errors.New("cannot marshal DNS message")
+	}
+	if len(b) < len(bb) {
+		return 0, errors.New("read would fragment DNS message")
+	}
+
+	copy(b, bb)
+	return len(bb), nil
+}
+
+func (f *fakeDNSConn) ReadFrom(b []byte) (int, Addr, error) {
+	return 0, nil, nil
+}
+
+func (f *fakeDNSConn) Write(b []byte) (int, error) {
+	f.q = new(dnsMsg)
+	if !f.q.Unpack(b) {
+		return 0, errors.New("cannot unmarshal DNS message")
+	}
+	return len(b), nil
+}
+
+func (f *fakeDNSConn) WriteTo(b []byte, addr Addr) (int, error) {
+	return 0, nil
+}
+
 func (f *fakeDNSConn) SetDeadline(t time.Time) error {
 	f.t = t
 	return nil
 }
 
-func (f *fakeDNSConn) dnsRoundTrip(q *dnsMsg) (*dnsMsg, error) {
-	return f.server.rh(f.n, f.s, q, f.t)
-}
-
 // UDP round-tripper algorithm should ignore invalid DNS responses (issue 13281).
 func TestIgnoreDNSForgeries(t *testing.T) {
 	c, s := Pipe()
@@ -837,7 +876,8 @@ func TestIgnoreDNSForgeries(t *testing.T) {
 		},
 	}
 
-	resp, err := dnsRoundTripUDP(c, msg)
+	dc := &dnsPacketConn{c}
+	resp, err := dc.dnsRoundTrip(msg)
 	if err != nil {
 		t.Fatalf("dnsRoundTripUDP failed: %v", err)
 	}
@@ -1113,7 +1153,14 @@ func TestStrictErrorsLookupIP(t *testing.T) {
 			case resolveOpError:
 				return nil, &OpError{Op: "write", Err: fmt.Errorf("socket on fire")}
 			case resolveServfail:
-				return &dnsMsg{dnsMsgHdr: dnsMsgHdr{id: q.id, rcode: dnsRcodeServerFailure}}, nil
+				return &dnsMsg{
+					dnsMsgHdr: dnsMsgHdr{
+						id:       q.id,
+						response: true,
+						rcode:    dnsRcodeServerFailure,
+					},
+					question: q.question,
+				}, nil
 			case resolveTimeout:
 				return nil, poll.ErrTimeout
 			default:
@@ -1123,7 +1170,14 @@ func TestStrictErrorsLookupIP(t *testing.T) {
 			switch q.question[0].Name {
 			case searchX, name + ".":
 				// Return NXDOMAIN to utilize the search list.
-				return &dnsMsg{dnsMsgHdr: dnsMsgHdr{id: q.id, rcode: dnsRcodeNameError}}, nil
+				return &dnsMsg{
+					dnsMsgHdr: dnsMsgHdr{
+						id:       q.id,
+						response: true,
+						rcode:    dnsRcodeNameError,
+					},
+					question: q.question,
+				}, nil
 			case searchY:
 				// Return records below.
 			default: