feat: global recursion upstream handler

wintbiit · Dec 12, 2023 · 375a4ca · 375a4ca
1 parent 94bde68
commit 375a4ca
Show file tree

Hide file tree

Showing 11 changed files with 34 additions and 108 deletions.
diff --git a/resolver/A.go b/resolver/A.go
@@ -1,8 +1,6 @@
 package resolver
 
 import (
-	"fmt"
-
 	"github.com/miekg/dns"
 	"github.com/wintbiit/ninedns/model"
 )
@@ -13,19 +11,10 @@ func init() {
 	resolvers[dns.TypeA] = &A{}
 }
 
-func (_ *A) Resolve(s model.RecordProvider, r *dns.Msg, name string) ([]dns.RR, error) {
+func (_ *A) Resolve(s model.RecordProvider, name string) ([]dns.RR, error) {
 	record := s.FindRecord(name, dns.TypeA)
 	if record == nil {
-		if !s.Recursion() {
-			return nil, fmt.Errorf("no record found for question: %+v", name)
-		}
-
-		resp, err := s.Exchange(r)
-		if err != nil {
-			return nil, err
-		}
-
-		return resp.Answer, nil
+		return nil, nil
 	}
 
 	rr := &dns.A{

diff --git a/resolver/AAAA.go b/resolver/AAAA.go
@@ -1,8 +1,6 @@
 package resolver
 
 import (
-	"fmt"
-
 	"github.com/miekg/dns"
 	"github.com/wintbiit/ninedns/model"
 )
@@ -13,19 +11,10 @@ func init() {
 	resolvers[dns.TypeAAAA] = &AAAA{}
 }
 
-func (_ *AAAA) Resolve(s model.RecordProvider, r *dns.Msg, name string) ([]dns.RR, error) {
+func (_ *AAAA) Resolve(s model.RecordProvider, name string) ([]dns.RR, error) {
 	record := s.FindRecord(name, dns.TypeAAAA)
 	if record == nil {
-		if !s.Recursion() {
-			return nil, fmt.Errorf("no record found for question: %+v", name)
-		}
-
-		resp, err := s.Exchange(r)
-		if err != nil {
-			return nil, err
-		}
-
-		return resp.Answer, nil
+		return nil, nil
 	}
 
 	rr := &dns.AAAA{

diff --git a/resolver/CNAME.go b/resolver/CNAME.go
@@ -1,8 +1,6 @@
 package resolver
 
 import (
-	"fmt"
-
 	"github.com/miekg/dns"
 	"github.com/wintbiit/ninedns/model"
 )
@@ -13,18 +11,10 @@ func init() {
 	resolvers[dns.TypeCNAME] = &CNAME{}
 }
 
-func (_ *CNAME) Resolve(s model.RecordProvider, r *dns.Msg, name string) ([]dns.RR, error) {
+func (_ *CNAME) Resolve(s model.RecordProvider, name string) ([]dns.RR, error) {
 	record := s.FindRecord(name, dns.TypeCNAME)
 	if record == nil {
-		if !s.Recursion() {
-			return nil, fmt.Errorf("no record found for question: %+v", name)
-		}
-		resp, err := s.Exchange(r)
-		if err != nil {
-			return nil, err
-		}
-
-		return resp.Answer, nil
+		return nil, nil
 	}
 
 	cname := record.Value.String()

diff --git a/resolver/MX.go b/resolver/MX.go
@@ -1,8 +1,6 @@
 package resolver
 
 import (
-	"fmt"
-
 	"github.com/miekg/dns"
 	"github.com/wintbiit/ninedns/model"
 )
@@ -13,19 +11,10 @@ func init() {
 	resolvers[dns.TypeMX] = &MX{}
 }
 
-func (_ *MX) Resolve(s model.RecordProvider, r *dns.Msg, name string) ([]dns.RR, error) {
+func (_ *MX) Resolve(s model.RecordProvider, name string) ([]dns.RR, error) {
 	record := s.FindRecord(name, dns.TypeMX)
 	if record == nil {
-		if !s.Recursion() {
-			return nil, fmt.Errorf("no record found for question: %+v", name)
-		}
-
-		resp, err := s.Exchange(r)
-		if err != nil {
-			return nil, err
-		}
-
-		return resp.Answer, nil
+		return nil, nil
 	}
 
 	mx, err := record.Value.MX()

diff --git a/resolver/NS.go b/resolver/NS.go
@@ -1,8 +1,6 @@
 package resolver
 
 import (
-	"fmt"
-
 	"github.com/miekg/dns"
 	"github.com/wintbiit/ninedns/model"
 )
@@ -13,19 +11,10 @@ func init() {
 	resolvers[dns.TypeNS] = &NS{}
 }
 
-func (_ *NS) Resolve(s model.RecordProvider, r *dns.Msg, name string) ([]dns.RR, error) {
+func (_ *NS) Resolve(s model.RecordProvider, name string) ([]dns.RR, error) {
 	record := s.FindRecord(name, dns.TypeNS)
 	if record == nil {
-		if !s.Recursion() {
-			return nil, fmt.Errorf("no record found for question: %+v", name)
-		}
-
-		resp, err := s.Exchange(r)
-		if err != nil {
-			return nil, err
-		}
-
-		return resp.Answer, nil
+		return nil, nil
 	}
 
 	rr := &dns.NS{

diff --git a/resolver/SOA.go b/resolver/SOA.go
@@ -1,8 +1,6 @@
 package resolver
 
 import (
-	"fmt"
-
 	"github.com/miekg/dns"
 	"github.com/wintbiit/ninedns/model"
 )
@@ -13,20 +11,11 @@ func init() {
 	resolvers[dns.TypeSOA] = &SOA{}
 }
 
-func (_ *SOA) Resolve(s model.RecordProvider, r *dns.Msg, name string) ([]dns.RR, error) {
+func (_ *SOA) Resolve(s model.RecordProvider, name string) ([]dns.RR, error) {
 	record := s.FindRecord(name, dns.TypeSOA)
 
 	if record == nil {
-		if !s.Recursion() {
-			return nil, fmt.Errorf("no record found for question: %+v", name)
-		}
-
-		resp, err := s.Exchange(r)
-		if err != nil {
-			return nil, err
-		}
-
-		return resp.Answer, nil
+		return nil, nil
 	}
 
 	soa, err := record.Value.SOA()

diff --git a/resolver/SRV.go b/resolver/SRV.go
@@ -1,8 +1,6 @@
 package resolver
 
 import (
-	"fmt"
-
 	"github.com/miekg/dns"
 	"github.com/wintbiit/ninedns/model"
 )
@@ -13,20 +11,11 @@ func init() {
 	resolvers[dns.TypeSRV] = &SRV{}
 }
 
-func (_ *SRV) Resolve(s model.RecordProvider, r *dns.Msg, name string) ([]dns.RR, error) {
+func (_ *SRV) Resolve(s model.RecordProvider, name string) ([]dns.RR, error) {
 	record := s.FindRecord(name, dns.TypeSRV)
 
 	if record == nil {
-		if !s.Recursion() {
-			return nil, fmt.Errorf("no record found for question: %+v", name)
-		}
-
-		resp, err := s.Exchange(r)
-		if err != nil {
-			return nil, err
-		}
-
-		return resp.Answer, nil
+		return nil, nil
 	}
 
 	srv, err := record.Value.SRV()

diff --git a/resolver/TXT.go b/resolver/TXT.go
@@ -11,20 +11,11 @@ func init() {
 	resolvers[dns.TypeTXT] = &TXT{}
 }
 
-func (_ *TXT) Resolve(s model.RecordProvider, r *dns.Msg, name string) ([]dns.RR, error) {
+func (_ *TXT) Resolve(s model.RecordProvider, name string) ([]dns.RR, error) {
 	records := s.FindRecords(name, dns.TypeTXT)
 
 	if records == nil {
-		if !s.Recursion() {
-			return nil, nil
-		}
-
-		resp, err := s.Exchange(r)
-		if err != nil {
-			return nil, err
-		}
-
-		return resp.Answer, nil
+		return nil, nil
 	}
 
 	rrs := make([]dns.RR, len(records))

diff --git a/resolver/resolver.go b/resolver/resolver.go
@@ -6,16 +6,16 @@ import (
 )
 
 type Resolver interface {
-	Resolve(model.RecordProvider, *dns.Msg, string) ([]dns.RR, error)
+	Resolve(model.RecordProvider, string) ([]dns.RR, error)
 }
 
 var resolvers = make(map[uint16]Resolver)
 
-func Resolve(typ uint16, p model.RecordProvider, r *dns.Msg, name string) ([]dns.RR, error) {
+func Resolve(typ uint16, p model.RecordProvider, name string) ([]dns.RR, error) {
 	resolver, ok := resolvers[typ]
 	if !ok {
 		return nil, nil
 	}
 
-	return resolver.Resolve(p, r, name)
+	return resolver.Resolve(p, name)
 }
diff --git a/server/ruleset.go b/server/ruleset.go
@@ -55,7 +55,7 @@ func (s *Server) newRuleSet(name string, rule model.Rule) *RuleSet {
 	return ruleSet
 }
 
-func (s *RuleSet) query(r, m *dns.Msg) {
+func (s *RuleSet) resolve(r, m *dns.Msg) {
 	for _, q := range r.Question {
 		// 1. Try CNAME
 		name := q.Name
@@ -74,11 +74,22 @@ func (s *RuleSet) query(r, m *dns.Msg) {
 }
 
 func (s *RuleSet) question(q *dns.Question, r, m *dns.Msg, name string) error {
-	records, err := resolver.Resolve(q.Qtype, s, r, name)
+	records, err := resolver.Resolve(q.Qtype, s, name)
 	if err != nil {
 		return err
 	}
 
+	if len(records) == 0 && s.Recursion() {
+		s.l.Debugf("Question %s not found, try upstream", q.String())
+		resp, err := s.Exchange(r)
+		if err != nil {
+			return err
+		}
+
+		m.Answer = resp.Answer
+		return nil
+	}
+
 	for _, record := range records {
 		record.Header().Name = name
 		m.Answer = append(m.Answer, record)
@@ -132,7 +143,7 @@ func (s *RuleSet) FindRecords(name string, quesType uint16) []model.Record {
 	}
 	records, err := s.cacheClient.FindRecords(name, model.ReadRecordType(quesType).String(), s.Name)
 	if err != nil {
-		s.l.Errorf("Failed to query records: %s", err)
+		s.l.Errorf("Failed to resolve records: %s", err)
 		return nil
 	}
 

diff --git a/server/server.go b/server/server.go
@@ -126,7 +126,7 @@ func (s *Server) handle(w dns.ResponseWriter, r *dns.Msg) {
 
 	s.l.Debugf("Found rule for %s: %+v", remoteAddr, handler)
 
-	handler.query(r, m)
+	handler.resolve(r, m)
 
 	if err := w.WriteMsg(m); err != nil {
 		s.l.Errorf("Failed to write response: %s", err)