wgengine/tsdns: add test and prevent useless updates (#449)

Signed-Off-By: Dmytro Shynkevych <dmytro@tailscale.com>

wgengine/tsdns/tsdns.go

@@ -2,7 +2,7 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
-// Package tsdns provides a Resolver struct capable of resolving
+// Package tsdns provides a Resolver capable of resolving
 // domains on a Tailscale network.
 package tsdns
 
@@ -11,6 +11,7 @@ import (
 	"errors"
 	"strings"
 	"sync"
+	"time"
 
 	dns "golang.org/x/net/dns/dnsmessage"
 	"inet.af/netaddr"
@@ -18,15 +19,17 @@ import (
 	"tailscale.com/wgengine/packet"
 )
 
-// defaultTTL is the TTL in seconds of all responses from Resolver.
-const defaultTTL = 600
+// defaultTTL is the TTL of all responses from Resolver.
+const defaultTTL = 600 * time.Second
 
 var (
 	errMapNotSet      = errors.New("domain map not set")
 	errNoSuchDomain   = errors.New("domain does not exist")
 	errNotImplemented = errors.New("query type not implemented")
 	errNotOurName     = errors.New("not an *.ipn.dev domain")
+	errNotOurQuery    = errors.New("query not for this resolver")
 	errNotQuery       = errors.New("not a DNS query")
+	errSmallBuffer    = errors.New("response buffer too small")
 )
 
 var (
@@ -36,13 +39,20 @@ var (
 
 // Map is all the data Resolver needs to resolve DNS queries.
 type Map struct {
-	// DomainToIP is a mapping of Tailscale domains to their IP addresses.
+	// domainToIP is a mapping of Tailscale domains to their IP addresses.
 	// For example, monitoring.ipn.dev -> 100.64.0.1.
-	DomainToIP map[string]netaddr.IP
+	domainToIP map[string]netaddr.IP
 }
 
-// Resolver is a DNS resolver for domain names of the form *.ipn.dev
-// It is intended
+// NewMap returns a new Map with domain to address mapping given by domainToIP.
+// It takes ownership of the provided map.
+func NewMap(domainToIP map[string]netaddr.IP) *Map {
+	return &Map{
+		domainToIP: domainToIP,
+	}
+}
+
+// Resolver is a DNS resolver for domain names of the form *.ipn.dev.
 type Resolver struct {
 	logf logger.Logf
 
@@ -96,7 +106,7 @@ func (r *Resolver) Resolve(domain string) (netaddr.IP, dns.RCode, error) {
 		r.mu.Unlock()
 		return netaddr.IP{}, dns.RCodeServerFailure, errMapNotSet
 	}
-	addr, found := r.dnsMap.DomainToIP[domain]
+	addr, found := r.dnsMap.domainToIP[domain]
 	r.mu.Unlock()
 
 	if !found {
@@ -156,6 +166,7 @@ func (r *Resolver) makeResponse(resp *response) error {
 }
 
 // marshalAnswer serializes the answer record into an active builder.
+// The caller may continue using the builder following the call.
 func marshalAnswer(resp *response, builder *dns.Builder) error {
 	var answer dns.AResource
 
@@ -168,7 +179,7 @@ func marshalAnswer(resp *response, builder *dns.Builder) error {
 		Name:  resp.Question.Name,
 		Type:  dns.TypeA,
 		Class: dns.ClassINET,
-		TTL:   defaultTTL,
+		TTL:   uint32(defaultTTL / time.Second),
 	}
 	ip := resp.IP.As16()
 	copy(answer.A[:], ip[12:])
@@ -176,44 +187,53 @@ func marshalAnswer(resp *response, builder *dns.Builder) error {
 }
 
 // marshalResponse serializes the DNS response into an active builder.
-func marshalResponse(resp *response, builder *dns.Builder) ([]byte, error) {
+// The caller may continue using the builder following the call.
+func marshalResponse(resp *response, builder *dns.Builder) error {
+	err := builder.StartQuestions()
+	if err != nil {
+		return err
+	}
+
+	err = builder.Question(resp.Question)
+	if err != nil {
+		return err
+	}
+
+	if resp.Header.RCode == dns.RCodeSuccess {
+		err = marshalAnswer(resp, builder)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+// marshalReponsePacket marshals a full DNS packet (including headers)
+// representing resp, which is a response to query, into buf.
+// It returns buf trimmed to the length of the response packet.
+func marshalResponsePacket(query *packet.ParsedPacket, resp *response, buf []byte) ([]byte, error) {
+	udpHeader := query.UDPHeader()
+	udpHeader.ToResponse()
+	offset := udpHeader.Len()
+
 	resp.Header.Response = true
 	resp.Header.Authoritative = true
 	if resp.Header.RecursionDesired {
 		resp.Header.RecursionAvailable = true
 	}
 
-	err := builder.StartQuestions()
-	if err != nil {
-		return nil, err
-	}
-
-	err = builder.Question(resp.Question)
-	if err != nil {
-		return nil, err
-	}
-
-	if resp.Header.RCode == dns.RCodeSuccess {
-		err = marshalAnswer(resp, builder)
-		if err != nil {
-			return nil, err
-		}
-	}
-
-	return builder.Finish()
-}
-
-func marshalResponsePacket(query *packet.ParsedPacket, resp *response, buf []byte) ([]byte, error) {
-	udpHeader := query.UDPHeader()
-	udpHeader.ToResponse()
-	offset := udpHeader.Len()
-
 	// dns.Builder appends to the passed buffer (without reallocation when possible),
 	// so we pass in a zero-length slice starting at the point it should start writing.
 	builder := dns.NewBuilder(buf[offset:offset], resp.Header)
 
+	err := marshalResponse(resp, &builder)
+	if err != nil {
+		return nil, err
+	}
+
 	// rbuf is the response slice with the correct length starting at offset.
-	rbuf, err := marshalResponse(resp, &builder)
+	rbuf, err := builder.Finish()
 	if err != nil {
 		return nil, err
 	}
@@ -235,15 +255,11 @@ func (r *Resolver) Respond(query *packet.ParsedPacket, buf []byte) ([]byte, erro
 
 	// 0. Verify that contract is upheld.
 	if !r.AcceptsPacket(query) {
-		r.logf("[unexpected] tsdns: Respond called on query not for this resolver")
-		resp.Header.RCode = dns.RCodeServerFailure
-		return marshalResponsePacket(query, &resp, buf)
+		return nil, errNotOurQuery
 	}
 	// A DNS response is at least as long as the query
 	if len(buf) < len(query.Buffer()) {
-		r.logf("[unexpected] tsdns: response buffer is too small")
-		resp.Header.RCode = dns.RCodeServerFailure
-		return marshalResponsePacket(query, &resp, buf)
+		return nil, errSmallBuffer
 	}
 
 	// 1. Parse query packet.