Feature: add PTR domain name resolution

Chore: optimize judgment logic

Signed-off-by: Fxzx micah <[email protected]>

dns/hosts.go

@@ -21,6 +21,7 @@ func LoadHosts() *trie.DomainTrie {
 	t := trie.New()
 	h := map[string][]net.IP{}
 	p := map[string][]string{}
+	n := map[string][]net.IP{}
 
 	sc := bufio.NewScanner(f)
 	for sc.Scan() {
@@ -36,22 +37,23 @@ func LoadHosts() *trie.DomainTrie {
 			continue
 		}
 
-		name := strings.Fields(line)
+		names := strings.Fields(line)
 		// ignore blank lines
-		if len(name) == 0 {
+		if len(names) == 0 {
 			continue
 		}
 
-		ip := net.ParseIP(name[0])
+		ip := net.ParseIP(names[0])
 		// ignore lines that do not start with IP
 		if ip == nil {
 			continue
 		}
 
 		ptr := transIpToPtr(ip)
-		for _, name := range name[1:] {
+		for _, name := range names[1:] {
 			h[name] = append(h[name], ip)
 			p[ptr] = append(p[ptr], name+".")
+			n[ptr] = append(n[ptr], h[name]...)
 		}
 	}
 
@@ -63,6 +65,10 @@ func LoadHosts() *trie.DomainTrie {
 		t.Insert(ptr, names)
 	}
 
+	for ptr, ips := range n {
+		t.Insert(ptr, ips)
+	}
+
 	return t
 }
 

dns/middleware.go

@@ -38,13 +38,15 @@ func withHosts(hosts *trie.DomainTrie) middleware {
 			switch data := record.Data.(type) {
 			case []net.IP:
 				for _, ip := range data {
-					if v4 := ip.To4(); v4 != nil && q.Qtype == D.TypeA {
+					if q.Qtype == D.TypeA && v4 := ip.To4(); v4 != nil {
 						rr := &D.A{}
 						rr.Hdr = D.RR_Header{Name: q.Name, Rrtype: D.TypeA, Class: D.ClassINET, Ttl: 60}
 						rr.A = v4
 
 						msg.Answer = append(msg.Answer, rr)
-					} else if v6 := ip.To16(); v6 != nil && q.Qtype == D.TypeAAAA {
+					}
+
+					if q.Qtype == D.TypeAAAA && v6 := ip.To16(); v6 != nil {
 						rr := &D.AAAA{}
 						rr.Hdr = D.RR_Header{Name: q.Name, Rrtype: D.TypeAAAA, Class: D.ClassINET, Ttl: 60}
 						rr.AAAA = v6