WIP Mock in-memory DNS server

core/local/local_test.go

@@ -974,23 +974,7 @@ func TestExecutionSchedulerIsRunning(t *testing.T) {
 	assert.NoError(t, <-err)
 }
 
-type mockDNSRunner struct {
-	*js.Runner
-	resolver *testutils.MockResolver
-}
-
-// NewVU initializes a new VU replacing its resolver with a MockResolver.
-func (r *mockDNSRunner) NewVU(id int64, out chan<- stats.SampleContainer) (lib.InitializedVU, error) {
-	initVU, err := r.Runner.NewVU(id, out)
-	if err != nil {
-		return nil, err
-	}
-	(initVU.(*js.VU)).Dialer.Resolver = r.resolver
-	return initVU, nil
-}
-
 func TestDNS(t *testing.T) {
-	t.Parallel()
 	tb := httpmultibin.NewHTTPMultiBin(t)
 	defer tb.Cleanup()
 	sr := tb.Replacer.Replace
@@ -1009,6 +993,11 @@ func TestDNS(t *testing.T) {
 			sleep(0.45);  // not an even multiple of 0.5 to minimize races with asserts
 		}`)
 
+	mockRes := testutils.NewMockResolver(nil)
+	defaultRes := net.DefaultResolver
+	net.DefaultResolver = testutils.NewMockDNSServer(mockRes)
+	defer func() { net.DefaultResolver = defaultRes }()
+
 	t.Run("cache", func(t *testing.T) {
 		testCases := map[string]struct {
 			opts lib.Options
@@ -1033,21 +1022,20 @@ func TestDNS(t *testing.T) {
 		for name, tc := range testCases {
 			tc := tc
 			t.Run(name, func(t *testing.T) {
+				mockRes.Set("myhost", sr("HTTPBIN_IP"))
+				time.AfterFunc(2*time.Second, func() {
+					mockRes.Set("myhost", "127.0.0.254")
+				})
 				logger := logrus.New()
 				logger.SetOutput(testutils.NewTestOutput(t))
 				runner, err := js.New(logger, &loader.SourceData{
 					URL: &url.URL{Path: "/script.js"}, Data: []byte(script),
 				}, nil, lib.RuntimeOptions{})
 				require.NoError(t, err)
 
-				mr := &mockDNSRunner{runner, testutils.NewMockResolver(nil)}
-				ctx, cancel, execScheduler, samples := newTestExecutionScheduler(t, mr, logger, tc.opts)
+				ctx, cancel, execScheduler, samples := newTestExecutionScheduler(t, runner, logger, tc.opts)
 				defer cancel()
 
-				mr.resolver.Set("myhost", sr("HTTPBIN_IP"))
-				time.AfterFunc(1*time.Second, func() {
-					mr.resolver.Set("myhost", "127.0.0.254")
-				})
 				errCh := make(chan error, 1)
 				go func() { errCh <- execScheduler.Run(ctx, ctx, samples) }()
 

go.mod

@@ -13,6 +13,7 @@ require (
 	github.com/andybalholm/brotli v0.0.0-20190704151324-71eb68cc467c
 	github.com/andybalholm/cascadia v1.0.0 // indirect
 	github.com/daaku/go.zipexe v0.0.0-20150329023125-a5fe2436ffcb // indirect
+	github.com/davecgh/go-spew v1.1.1
 	github.com/dlclark/regexp2 v1.2.1-0.20200807145002-74bac81f00cf // indirect
 	github.com/dop251/goja v0.0.0-20200818110326-5574b5dbd2b9
 	github.com/dustin/go-humanize v0.0.0-20171111073723-bb3d318650d4

lib/testutils/resolver.go

@@ -21,9 +21,15 @@
 package testutils
 
 import (
+	"context"
+	"errors"
 	"fmt"
 	"net"
 	"sync"
+	"time"
+
+	"github.com/davecgh/go-spew/spew"
+	"golang.org/x/net/dns/dnsmessage"
 )
 
 // MockResolver implements netext.DNSResolver, and allows changing the host
@@ -49,8 +55,174 @@ func (r *MockResolver) Fetch(host string) (net.IP, error) {
 	return nil, fmt.Errorf("lookup %s: no such host", host)
 }
 
+func (mr *MockResolver) handle(s net.Conn) {
+	for {
+		b := make([]byte, 512)
+		n, err := s.Read(b)
+		if err != nil {
+			return
+		}
+
+		var msg dnsmessage.Message
+		// FIXME: This returns the error `unpacking Question.Name: segment prefix is reserved`
+		// https://github.com/golang/net/blob/62affa334b73ec65ed44a326519ac12c421905e3/dns/dnsmessage/message.go#L1999
+		if err := msg.Unpack(b[:n]); err != nil {
+			spew.Dump(err)
+			return
+		}
+
+		if len(msg.Questions) == 0 {
+			return
+		}
+		q := msg.Questions[0]
+		if q.Type != dnsmessage.TypeA {
+			return
+		}
+
+		ip, err := mr.Fetch(q.Name.String())
+		if err != nil {
+			// TODO: Write NXDOMAIN response?
+			return
+		}
+
+		var ip4 [4]byte
+		copy(ip4[:], ip)
+
+		msg.Header.Response = true
+		msg.Answers = []dnsmessage.Resource{
+			{
+				Header: dnsmessage.ResourceHeader{
+					Name:   q.Name,
+					Type:   q.Type,
+					Class:  q.Class,
+					Length: 4,
+				},
+				Body: &dnsmessage.AResource{
+					A: ip4,
+				},
+			},
+		}
+
+		b, err = msg.Pack()
+		if err != nil {
+			return
+		}
+		s.Write(b)
+	}
+}
+
 func (r *MockResolver) Set(host, ip string) {
 	r.m.Lock()
 	defer r.m.Unlock()
 	r.hosts[host] = net.ParseIP(ip)
 }
+
+type BaseMockResolver struct {
+	*net.Resolver
+	mr *MockResolver
+}
+
+func NewMockDNSServer(mr *MockResolver) *net.Resolver {
+	c, s := net.Pipe()
+
+	go mr.handle(s)
+
+	return &net.Resolver{
+		PreferGo: true,
+		Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
+			return c, nil
+		},
+	}
+}
+
+// This is another mocking approach copied from go/src/net/dnsclient_unix_test.go that additionally
+// mocks net.Conn, but the net.Pipe() approach is much simpler so I'd prefer to get that working.
+type fakeDNSServer struct {
+	rh        func(n, s string, q dnsmessage.Message, t time.Time) (dnsmessage.Message, error)
+	alwaysTCP bool
+}
+
+func (server *fakeDNSServer) DialContext(_ context.Context, n, s string) (net.Conn, error) {
+	if server.alwaysTCP || n == "tcp" || n == "tcp4" || n == "tcp6" {
+		return &fakeDNSConn{tcp: true, server: server, n: n, s: s}, nil
+	}
+	return &fakeDNSPacketConn{fakeDNSConn: fakeDNSConn{tcp: false, server: server, n: n, s: s}}, nil
+}
+
+type fakeDNSConn struct {
+	net.Conn
+	tcp    bool
+	server *fakeDNSServer
+	n      string
+	s      string
+	q      dnsmessage.Message
+	t      time.Time
+	buf    []byte
+}
+
+func (f *fakeDNSConn) Close() error {
+	return nil
+}
+
+func (f *fakeDNSConn) Read(b []byte) (int, error) {
+	if len(f.buf) > 0 {
+		n := copy(b, f.buf)
+		f.buf = f.buf[n:]
+		return n, nil
+	}
+
+	resp, err := f.server.rh(f.n, f.s, f.q, f.t)
+	if err != nil {
+		return 0, err
+	}
+
+	bb := make([]byte, 2, 514)
+	bb, err = resp.AppendPack(bb)
+	if err != nil {
+		return 0, fmt.Errorf("cannot marshal DNS message: %v", err)
+	}
+
+	if f.tcp {
+		l := len(bb) - 2
+		bb[0] = byte(l >> 8)
+		bb[1] = byte(l)
+		f.buf = bb
+		return f.Read(b)
+	}
+
+	bb = bb[2:]
+	if len(b) < len(bb) {
+		return 0, errors.New("read would fragment DNS message")
+	}
+
+	copy(b, bb)
+	return len(bb), nil
+}
+
+func (f *fakeDNSConn) Write(b []byte) (int, error) {
+	if f.tcp && len(b) >= 2 {
+		b = b[2:]
+	}
+	if f.q.Unpack(b) != nil {
+		return 0, fmt.Errorf("cannot unmarshal DNS message fake %s (%d)", f.n, len(b))
+	}
+	return len(b), nil
+}
+
+func (f *fakeDNSConn) SetDeadline(t time.Time) error {
+	f.t = t
+	return nil
+}
+
+type fakeDNSPacketConn struct {
+	net.PacketConn
+	fakeDNSConn
+}
+
+func (f *fakeDNSPacketConn) SetDeadline(t time.Time) error {
+	return f.fakeDNSConn.SetDeadline(t)
+}
+
+func (f *fakeDNSPacketConn) Close() error {
+	return f.fakeDNSConn.Close()
+}