more config

internal/resolver/pool.go

@@ -255,8 +255,10 @@ func (u *Upstream) Stats() Stats {
 // from, the upstreams, and the address index used for per-client upstream
 // overrides. Readers take the pointer once and are then immune to reloads.
 type poolState struct {
-	cfg config.Upstream
-	ups []*Upstream
+	cfg      config.Upstream
+	ups      []*Upstream
+	defaults []*Upstream
+	routes   map[string][]*Upstream
 	// byAddr indexes the upstreams by every address they answer to, so an ACL
 	// group naming either the configured Addr or the resolved one works.
 	byAddr map[string][]*Upstream
@@ -344,6 +346,22 @@ func buildState(cfg config.Upstream, carry *poolState, lookup bool) *poolState {
 		}
 		st.ups = append(st.ups, up)
 	}
+	st.defaults = append([]*Upstream(nil), st.ups...)
+	st.routes = make(map[string][]*Upstream)
+	for _, z := range cfg.ForwardZones {
+		zone := strings.ToLower(dns.Fqdn(z.Zone))
+		for _, s := range z.Servers {
+			addr, resolved := resolveAddr(s, bootstraps, lookup)
+			up := &Upstream{cfg: s, addr: addr, resolved: resolved}
+			if old := prev[s.Addr+"/"+s.Proto]; old != nil {
+				up.stats = old.stats
+			} else {
+				up.stats = newUpstreamStats()
+			}
+			st.ups = append(st.ups, up)
+			st.routes[zone] = append(st.routes[zone], up)
+		}
+	}
 	st.byAddr = make(map[string][]*Upstream, len(st.ups)*2)
 	for _, up := range st.ups {
 		for _, a := range append([]string{up.cfg.Addr, up.addr}, up.resolved...) {
@@ -392,6 +410,9 @@ func (p *Pool) Config() config.Upstream {
 // All returns every upstream.
 func (p *Pool) All() []*Upstream { return p.state.Load().all() }
 
+// Defaults returns the servers eligible for queries without a forwarding route.
+func (p *Pool) Defaults() []*Upstream { return p.state.Load().defaults }
+
 // Get returns an upstream by configured or resolved address, or nil.
 func (p *Pool) Get(addr string) *Upstream {
 	for _, u := range p.All() {
@@ -457,8 +478,26 @@ func (p *Pool) clientFor(u *Upstream) *dns.Client {
 // attempt.
 func (p *Pool) Exchange(ctx context.Context, q *dns.Msg, want []string, timeout time.Duration) (*dns.Msg, string, error) {
 	st := p.state.Load()
-	candidates := st.all()
-	if len(want) > 0 {
+	candidates := st.defaults
+	// Walk complete labels, most specific first. A matched route never falls
+	// back to public resolvers, including when every target is unavailable.
+	matched := false
+	if len(q.Question) > 0 {
+		name := strings.ToLower(dns.Fqdn(q.Question[0].Name))
+		for name != "" {
+			if route, ok := st.routes[name]; ok {
+				candidates = route
+				matched = true
+				break
+			}
+			dot := strings.IndexByte(name, '.')
+			if dot < 0 {
+				break
+			}
+			name = name[dot+1:]
+		}
+	}
+	if !matched && len(want) > 0 {
 		candidates = p.restrict(st, want)
 	}
 	if len(candidates) == 0 {

internal/resolver/resolver_test.go

@@ -656,3 +656,42 @@ func TestHealthCheckProbesEveryUpstream(t *testing.T) {
 		t.Error("the dead upstream was not taken out of rotation")
 	}
 }
+
+func TestForwardZones(t *testing.T) {
+	public := newFakeUpstream(t, fakeOpt{ip: "203.0.113.1"})
+	firewall := newFakeUpstream(t, fakeOpt{ip: "10.10.11.117", truncateUDP: true})
+	specific := newFakeUpstream(t, fakeOpt{ip: "10.10.11.118"})
+	cfg := poolCfg(udpServer(public.addr))
+	cfg.ForwardZones = []config.ForwardZone{
+		{Zone: "local", Servers: []config.UpstreamServer{udpServer(firewall.addr)}},
+		{Zone: "lab.local", Servers: []config.UpstreamServer{udpServer(specific.addr)}},
+	}
+	p := NewPool(cfg)
+	for _, tc := range []struct{ name, addr string }{
+		{"fastpi.LoCaL.", firewall.addr}, {"local.", firewall.addr},
+		{"host.lab.local.", specific.addr}, {"notlocal.", public.addr},
+		{"local.example.", public.addr},
+	} {
+		_, addr, err := p.Exchange(testCtx(t), query(tc.name), []string{public.addr}, time.Second)
+		if err != nil || addr != tc.addr {
+			t.Fatalf("%s: addr=%s err=%v, want %s", tc.name, addr, err, tc.addr)
+		}
+	}
+	if firewall.tcpHits.Load() == 0 {
+		t.Fatal("forwarded truncated UDP answer did not retry TCP")
+	}
+	before := public.udpHits.Load()
+	cfg.ForwardZones[0].Servers = []config.UpstreamServer{udpServer("127.0.0.1:1")}
+	p.Rebind(cfg)
+	if _, _, err := p.Exchange(testCtx(t), query("fastpi.local."), nil, 50*time.Millisecond); err == nil {
+		t.Fatal("dead route succeeded")
+	}
+	if public.udpHits.Load() != before {
+		t.Fatal("private query leaked to public upstream")
+	}
+	cfg.ForwardZones = nil
+	p.Rebind(cfg)
+	if _, addr, err := p.Exchange(testCtx(t), query("fastpi.local."), nil, time.Second); err != nil || addr != public.addr {
+		t.Fatalf("removed route: %s %v", addr, err)
+	}
+}