fix(resolver): restore fallbackUpstream for rewritten queries (#2223)

Author: maximilize

resolver/conditional_upstream_resolver.go

@@ -94,12 +94,9 @@ func (r *ConditionalUpstreamResolver) Resolve(ctx context.Context, request *mode
 	resolved := false
 	if len(r.mapping) > 0 {
 		resolved, response, err = r.processRequest(ctx, request)
-		if err != nil {
-			return nil, err
-		}
 	}
 
-	if !resolved {
+	if !resolved && err == nil {
 		logger.WithField("next_resolver", Name(r.next)).Trace("go to next resolver")
 		response, err = r.next.Resolve(ctx, request)
 		if err != nil {
@@ -110,6 +107,18 @@ func (r *ConditionalUpstreamResolver) Resolve(ctx context.Context, request *mode
 	// Revert the request
 	request.Req = original
 
+	// The mapped resolver failed or had nothing: ask the rest of the chain,
+	// using the original name (`fallbackUpstream`).
+	if shouldFallbackUpstream(&r.cfg.RewriterConfig, resolved, response, err) {
+		logger.WithField("next_resolver", Name(r.next)).Trace("fallback to next resolver")
+
+		return r.next.Resolve(ctx, request)
+	}
+
+	if err != nil {
+		return nil, err
+	}
+
 	// Revert rewrites in the response
 	if rewritten != nil && response != NoResponse && response != nil && response.Res != nil {
 		revertRewritesInResponse(response.Res, originalNames)

resolver/conditional_upstream_resolver_test.go

@@ -257,6 +257,121 @@ var _ = Describe("ConditionalUpstreamResolver", Label("conditionalResolver"), fu
 			})
 		})
 
+		When("the mapped upstream has no answer and fallbackUpstream is set", func() {
+			BeforeEach(func() {
+				emptyUpstream := NewMockUDPUpstreamServer().WithAnswerFn(func(request *dns.Msg) (response *dns.Msg) {
+					response = new(dns.Msg)
+					response.SetReply(request)
+
+					return response
+				})
+
+				sutConfig.FallbackUpstream = true
+				sutConfig.Mapping.Upstreams["empty.box"] = []config.Upstream{emptyUpstream.Start()}
+				sutConfig.Rewrite = map[string]string{"source.test": "empty.box"}
+			})
+
+			It("should ask the next resolver with the original name", func() {
+				var seen string
+
+				m = &mockResolver{}
+				m.On("Resolve", mock.Anything).Return(&Response{Res: new(dns.Msg)}, nil)
+				m.ResolveFn = func(_ context.Context, req *Request) (*Response, error) {
+					seen = req.Req.Question[0].Name
+					resp, err := util.NewMsgWithAnswer(seen, 250, A, "192.192.192.192")
+					Expect(err).Should(Succeed())
+
+					return &Response{Res: resp, RType: ResponseTypeRESOLVED, Reason: "RESOLVED"}, nil
+				}
+				sut.Next(m)
+
+				Expect(sut.Resolve(ctx, newRequest("www.source.test.", A))).
+					Should(
+						SatisfyAll(
+							BeDNSRecord("www.source.test.", A, "192.192.192.192"),
+							HaveResponseType(ResponseTypeRESOLVED),
+							HaveReturnCode(dns.RcodeSuccess),
+						))
+
+				Expect(seen).Should(Equal("www.source.test."))
+			})
+		})
+
+		When("the mapped upstream fails and fallbackUpstream is set", func() {
+			BeforeEach(func() {
+				deadUpstream := NewMockUDPUpstreamServer().WithAnswerRR("dead.box. 123 IN A 1.2.3.4")
+				upstream := deadUpstream.Start()
+				deadUpstream.Close()
+
+				sutConfig.FallbackUpstream = true
+				sutConfig.Mapping.Upstreams["dead.box"] = []config.Upstream{upstream}
+				sutConfig.Rewrite = map[string]string{"source.test": "dead.box"}
+			})
+
+			It("should ask the next resolver with the original name", func() {
+				var seen string
+
+				m = &mockResolver{}
+				m.On("Resolve", mock.Anything).Return(&Response{Res: new(dns.Msg)}, nil)
+				m.ResolveFn = func(_ context.Context, req *Request) (*Response, error) {
+					seen = req.Req.Question[0].Name
+					resp, err := util.NewMsgWithAnswer(seen, 250, A, "192.192.192.192")
+					Expect(err).Should(Succeed())
+
+					return &Response{Res: resp, RType: ResponseTypeRESOLVED, Reason: "RESOLVED"}, nil
+				}
+				sut.Next(m)
+
+				Expect(sut.Resolve(ctx, newRequest("www.source.test.", A))).
+					Should(HaveResponseType(ResponseTypeRESOLVED))
+
+				Expect(seen).Should(Equal("www.source.test."))
+			})
+		})
+
+		When("the mapped upstream fails and fallbackUpstream is not set", func() {
+			BeforeEach(func() {
+				deadUpstream := NewMockUDPUpstreamServer().WithAnswerRR("dead.box. 123 IN A 1.2.3.4")
+				upstream := deadUpstream.Start()
+				deadUpstream.Close()
+
+				sutConfig.Mapping.Upstreams["dead.box"] = []config.Upstream{upstream}
+				sutConfig.Rewrite = map[string]string{"source.test": "dead.box"}
+			})
+
+			It("should return the error", func() {
+				_, err := sut.Resolve(ctx, newRequest("www.source.test.", A))
+				Expect(err).Should(HaveOccurred())
+				Expect(m.Calls).Should(BeEmpty())
+			})
+		})
+
+		When("the mapped upstream has no answer and fallbackUpstream is not set", func() {
+			BeforeEach(func() {
+				emptyUpstream := NewMockUDPUpstreamServer().WithAnswerFn(func(request *dns.Msg) (response *dns.Msg) {
+					response = new(dns.Msg)
+					response.SetReply(request)
+
+					return response
+				})
+
+				sutConfig.Mapping.Upstreams["empty.box"] = []config.Upstream{emptyUpstream.Start()}
+				sutConfig.Rewrite = map[string]string{"source.test": "empty.box"}
+			})
+
+			It("should return the empty answer and not ask the next resolver", func() {
+				Expect(sut.Resolve(ctx, newRequest("www.source.test.", A))).
+					Should(
+						SatisfyAll(
+							HaveNoAnswer(),
+							HaveResponseType(ResponseTypeCONDITIONAL),
+							HaveReturnCode(dns.RcodeSuccess),
+						))
+
+				Expect(m.Calls).Should(BeEmpty())
+			})
+		})
+
 		When("request does not match rewrite rule but matches conditional mapping", func() {
 			It("should not rewrite and resolve via conditional upstream", func() {
 				// Direct request to fritz.box (no rewrite)

resolver/custom_dns_resolver.go

@@ -230,6 +230,16 @@ func (r *CustomDNSResolver) Resolve(ctx context.Context, request *model.Request)
 	// Revert the request
 	request.Req = original
 
+	// A response we produced ourselves without an answer means the mapping had
+	// nothing for this query: ask the rest of the chain, using the original
+	// name (`fallbackUpstream`).
+	answered := err == nil && response != nil && response.RType == model.ResponseTypeCUSTOMDNS
+	if shouldFallbackUpstream(&r.cfg.RewriterConfig, answered, response, nil) {
+		logger.WithField("next_resolver", Name(r.next)).Trace("fallback to next resolver")
+
+		return r.next.Resolve(ctx, request)
+	}
+
 	if err != nil {
 		return response, err
 	}

resolver/custom_dns_resolver_test.go

@@ -10,6 +10,7 @@ import (
 	. "github.com/0xERR0R/blocky/helpertest"
 	"github.com/0xERR0R/blocky/log"
 	. "github.com/0xERR0R/blocky/model"
+	"github.com/0xERR0R/blocky/util"
 	"github.com/miekg/dns"
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
@@ -592,6 +593,39 @@ var _ = Describe("CustomDNSResolver", func() {
 			sut.Next(m)
 		})
 
+		When("the mapping has no answer of the requested type and fallbackUpstream is set", func() {
+			BeforeEach(func() {
+				cfg.FallbackUpstream = true
+				sut = NewCustomDNSResolver(cfg)
+			})
+
+			It("should ask the next resolver with the original name", func() {
+				var seen string
+
+				m = &mockResolver{}
+				m.On("Resolve", mock.Anything).Return(&Response{Res: new(dns.Msg)}, nil)
+				m.ResolveFn = func(_ context.Context, req *Request) (*Response, error) {
+					seen = req.Req.Question[0].Name
+					resp, err := util.NewMsgWithAnswer(seen, uint(TTL), AAAA, "2001:db8::1")
+					Expect(err).Should(Succeed())
+
+					return &Response{Res: resp, RType: ResponseTypeRESOLVED, Reason: "RESOLVED"}, nil
+				}
+				sut.Next(m)
+
+				// custom.domain has an A record only, so the AAAA query ends up
+				// as an empty CUSTOM DNS answer
+				Expect(sut.Resolve(ctx, newRequest("www.source.test.", AAAA))).
+					Should(
+						SatisfyAll(
+							HaveResponseType(ResponseTypeRESOLVED),
+							HaveReturnCode(dns.RcodeSuccess),
+						))
+
+				Expect(seen).Should(Equal("www.source.test."))
+			})
+		})
+
 		When("request matches rewrite rule", func() {
 			It("should rewrite subdomain and resolve from mapping", func() {
 				// Request for www.source.test should be rewritten to www.custom.domain

resolver/rewrite_helper.go

@@ -3,6 +3,8 @@ package resolver
 import (
 	"strings"
 
+	"github.com/0xERR0R/blocky/config"
+	"github.com/0xERR0R/blocky/model"
 	"github.com/0xERR0R/blocky/util"
 
 	"github.com/miekg/dns"
@@ -64,6 +66,24 @@ func rewriteDomain(domain string, rewriteMap map[string]string) (string, string)
 	return domain, ""
 }
 
+// shouldFallbackUpstream reports whether a rewritten query that ended without an
+// answer should be retried, with its original name, on the rest of the chain.
+// See `fallbackUpstream` in the documentation.
+//
+// answered tells whether the resolver produced the response itself: a response
+// it only passed through from the next resolver must not be retried there.
+func shouldFallbackUpstream(cfg *config.RewriterConfig, answered bool, response *model.Response, err error) bool {
+	if !cfg.FallbackUpstream || len(cfg.Rewrite) == 0 || !answered {
+		return false
+	}
+
+	if err != nil {
+		return true
+	}
+
+	return response != nil && response != NoResponse && response.Res != nil && len(response.Res.Answer) == 0
+}
+
 // revertRewritesInResponse reverts domain rewrites in the DNS response
 func revertRewritesInResponse(response *dns.Msg, originalNames map[string]string) {
 	if len(originalNames) == 0 {