From 749354b3bb013ab9979802ef5cefca6f2ef5c629 Mon Sep 17 00:00:00 2001
From: Daniel McCarney <cpu@letsencrypt.org>
Date: Tue, 27 Aug 2019 13:36:20 -0400
Subject: [PATCH] dns: add support for mocking SERVFAIL responses. (#10)

---
 challenge-servers.go | 15 +++++++++------
 dns.go               |  7 +++++++
 mockdns.go           | 27 +++++++++++++++++++++++++++
 3 files changed, 43 insertions(+), 6 deletions(-)

diff --git a/challenge-servers.go b/challenge-servers.go
index c069432..7318f7a 100644
--- a/challenge-servers.go
+++ b/challenge-servers.go
@@ -76,6 +76,8 @@ type mockDNSData struct {
 	caaRecords map[string][]MockCAAPolicy
 	// A map of host to CNAME records.
 	cnameRecords map[string]string
+	// A map of hostnames that should receive a SERVFAIL response for all queries.
+	servFailRecords map[string]bool
 }
 
 // MockCAAPolicy holds a tag and a value for a CAA record. See
@@ -133,12 +135,13 @@ func New(config Config) (*ChallSrv, error) {
 		tlsALPNOne:     make(map[string]string),
 		redirects:      make(map[string]string),
 		dnsMocks: mockDNSData{
-			defaultIPv4:  defaultIPv4,
-			defaultIPv6:  defaultIPv6,
-			aRecords:     make(map[string][]string),
-			aaaaRecords:  make(map[string][]string),
-			caaRecords:   make(map[string][]MockCAAPolicy),
-			cnameRecords: make(map[string]string),
+			defaultIPv4:     defaultIPv4,
+			defaultIPv6:     defaultIPv6,
+			aRecords:        make(map[string][]string),
+			aaaaRecords:     make(map[string][]string),
+			caaRecords:      make(map[string][]MockCAAPolicy),
+			cnameRecords:    make(map[string]string),
+			servFailRecords: make(map[string]bool),
 		},
 	}
 
diff --git a/dns.go b/dns.go
index d2b71fb..5b36e93 100644
--- a/dns.go
+++ b/dns.go
@@ -170,6 +170,13 @@ func (s *ChallSrv) dnsHandler(w dns.ResponseWriter, r *dns.Msg) {
 			Question: q,
 		})
 
+		// If there is a ServFail mock set then ignore the question and set the
+		// SERVFAIL rcode and continue.
+		if s.GetDNSServFailRecord(q.Name) {
+			m.SetRcode(r, dns.RcodeServerFailure)
+			continue
+		}
+
 		// If a CNAME exists for the question include the CNAME record and modify
 		// the question to instead lookup based on that CNAME's target
 		if cname := s.GetDNSCNAMERecord(q.Name); cname != "" {
diff --git a/mockdns.go b/mockdns.go
index e683bbf..b192c47 100644
--- a/mockdns.go
+++ b/mockdns.go
@@ -145,3 +145,30 @@ func (s *ChallSrv) GetDNSCAARecord(host string) []MockCAAPolicy {
 	host = dns.Fqdn(host)
 	return s.dnsMocks.caaRecords[host]
 }
+
+// AddDNSServFailRecord configures the chall srv to return SERVFAIL responses
+// for all queries for the given host.
+func (s *ChallSrv) AddDNSServFailRecord(host string) {
+	s.challMu.Lock()
+	defer s.challMu.Unlock()
+	host = dns.Fqdn(host)
+	s.dnsMocks.servFailRecords[host] = true
+}
+
+// DeleteDNSServFailRecord configures the chall srv to no longer return SERVFAIL
+// responses for all queries for the given host.
+func (s *ChallSrv) DeleteDNSServFailRecord(host string) {
+	s.challMu.Lock()
+	defer s.challMu.Unlock()
+	host = dns.Fqdn(host)
+	delete(s.dnsMocks.servFailRecords, host)
+}
+
+// GetDNSServFailRecord returns true when the chall srv has been configured with
+// AddDNSServFailRecord to return SERVFAIL for all queries to the given host.
+func (s *ChallSrv) GetDNSServFailRecord(host string) bool {
+	s.challMu.RLock()
+	defer s.challMu.RUnlock()
+	host = dns.Fqdn(host)
+	return s.dnsMocks.servFailRecords[host]
+}