git » dnss » commit af006ed

diff --git a/dnss.go b/dnss.go
index 4dd6cc1..9c7d62a 100644
--- a/dnss.go
+++ b/dnss.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"flag"
+	"strings"
 	"sync"
 	"time"
 
@@ -22,9 +23,16 @@ import (
 var (
 	dnsListenAddr = flag.String("dns_listen_addr", ":53",
 		"address to listen on for DNS")
+
 	dnsUnqualifiedUpstream = flag.String("dns_unqualified_upstream", "",
 		"DNS server to forward unqualified requests to")
 
+	fallbackUpstream = flag.String("fallback_upstream", "8.8.8.8:53",
+		"DNS server to resolve domains in --fallback_domains")
+	fallbackDomains = flag.String("fallback_domains", "dns.google.com.",
+		"Domains we resolve via DNS, using --fallback_upstream"+
+			" (space-separated list)")
+
 	enableDNStoGRPC = flag.Bool("enable_dns_to_grpc", false,
 		"enable DNS-to-GRPC server")
 	grpcUpstream = flag.String("grpc_upstream", "localhost:9953",
@@ -101,6 +109,8 @@ func main() {
 		r := dnstox.NewGRPCResolver(*grpcUpstream, *grpcClientCAFile)
 		cr := dnstox.NewCachingResolver(r)
 		dtg := dnstox.New(*dnsListenAddr, cr, *dnsUnqualifiedUpstream)
+		dtg.SetFallback(
+			*fallbackUpstream, strings.Split(*fallbackDomains, " "))
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
@@ -128,6 +138,8 @@ func main() {
 		r := dnstox.NewHTTPSResolver(*httpsUpstream, *httpsClientCAFile)
 		cr := dnstox.NewCachingResolver(r)
 		dth := dnstox.New(*dnsListenAddr, cr, *dnsUnqualifiedUpstream)
+		dth.SetFallback(
+			*fallbackUpstream, strings.Split(*fallbackDomains, " "))
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
diff --git a/dnstox/dnstox.go b/dnstox/dnstox.go
index fb5f073..52f7238 100644
--- a/dnstox/dnstox.go
+++ b/dnstox/dnstox.go
@@ -47,13 +47,24 @@ type Server struct {
 	Addr        string
 	unqUpstream string
 	resolver    Resolver
+
+	fallbackDomains  map[string]struct{}
+	fallbackUpstream string
 }
 
 func New(addr string, resolver Resolver, unqUpstream string) *Server {
 	return &Server{
-		Addr:        addr,
-		resolver:    resolver,
-		unqUpstream: unqUpstream,
+		Addr:            addr,
+		resolver:        resolver,
+		unqUpstream:     unqUpstream,
+		fallbackDomains: map[string]struct{}{},
+	}
+}
+
+func (s *Server) SetFallback(upstream string, domains []string) {
+	s.fallbackUpstream = upstream
+	for _, d := range domains {
+		s.fallbackDomains[d] = struct{}{}
 	}
 }
 
@@ -67,11 +78,18 @@ func (s *Server) Handler(w dns.ResponseWriter, r *dns.Msg) {
 		tr.LazyPrintf(util.QuestionsToString(r.Question))
 	}
 
+	// We only support single-question queries.
+	if len(r.Question) != 1 {
+		tr.LazyPrintf("len(Q) != 1, failing")
+		dns.HandleFailed(w, r)
+		return
+	}
+
 	// Forward to the unqualified upstream server if:
 	//  - We have one configured.
 	//  - There's only one question in the request, to keep things simple.
 	//  - The question is unqualified (only one '.' in the name).
-	useUnqUpstream := s.unqUpstream != "" && len(r.Question) == 1 &&
+	useUnqUpstream := s.unqUpstream != "" &&
 		strings.Count(r.Question[0].Name, ".") <= 1
 	if useUnqUpstream {
 		u, err := dns.Exchange(r, s.unqUpstream)
@@ -81,10 +99,29 @@ func (s *Server) Handler(w dns.ResponseWriter, r *dns.Msg) {
 				util.TraceAnswer(tr, u)
 			}
 			w.WriteMsg(u)
-			return
 		} else {
 			tr.LazyPrintf("unqualified upstream error: %v", err)
+			dns.HandleFailed(w, r)
+		}
+
+		return
+	}
+
+	// Forward to the fallback server if the domain is on our list.
+	if _, ok := s.fallbackDomains[r.Question[0].Name]; ok {
+		u, err := dns.Exchange(r, s.fallbackUpstream)
+		if err == nil {
+			tr.LazyPrintf("used fallback upstream (%s)", s.fallbackUpstream)
+			if glog.V(3) {
+				util.TraceAnswer(tr, u)
+			}
+			w.WriteMsg(u)
+		} else {
+			tr.LazyPrintf("fallback upstream error: %v", err)
+			dns.HandleFailed(w, r)
 		}
+
+		return
 	}
 
 	// Create our own IDs, in case different users pick the same id and we