Fix: inability to load duplicate records when loading local HOSTS

Signed-off-by: Fxzx micah <[email protected]>
fxzxmicah · Jan 19, 2024 · a04e81c · a04e81c
1 parent 9ef1e44
commit a04e81c
Show file tree

Hide file tree

Showing 5 changed files with 36 additions and 37 deletions.
diff --git a/component/resolver/resolver.go b/component/resolver/resolver.go
@@ -48,7 +48,7 @@ type Resolver interface {
 // LookupIPv4 with a host, return ipv4 list
 func LookupIPv4(ctx context.Context, host string) ([]net.IP, error) {
 	if node := DefaultHosts.Search(host); node != nil {
-		if ip := node.Data.(net.IP).To4(); ip != nil {
+		if ip := node.Data.([]net.IP)[0].To4(); ip != nil {
 			return []net.IP{ip}, nil
 		}
 	}
@@ -95,7 +95,7 @@ func LookupIPv6(ctx context.Context, host string) ([]net.IP, error) {
 	}
 
 	if node := DefaultHosts.Search(host); node != nil {
-		if ip := node.Data.(net.IP).To16(); ip != nil {
+		if ip := node.Data.([]net.IP)[0].To16(); ip != nil {
 			return []net.IP{ip}, nil
 		}
 	}
@@ -138,7 +138,7 @@ func ResolveIPv6(host string) (net.IP, error) {
 // LookupIPWithResolver same as ResolveIP, but with a resolver
 func LookupIPWithResolver(ctx context.Context, host string, r Resolver) ([]net.IP, error) {
 	if node := DefaultHosts.Search(host); node != nil {
-		return []net.IP{node.Data.(net.IP)}, nil
+		return []net.IP{node.Data.([]net.IP)[0]}, nil
 	}
 
 	if r != nil {

diff --git a/config/config.go b/config/config.go
@@ -514,7 +514,7 @@ func parseHosts(cfg *RawConfig) (*trie.DomainTrie, error) {
 			if ip == nil {
 				return nil, fmt.Errorf("%s is not a valid IP", ipStr)
 			}
-			tree.Insert(strings.ToLower(domain), ip)
+			tree.Insert(strings.ToLower(domain), []net.IP{ip})
 		}
 	}
 

diff --git a/dns/hosts.go b/dns/hosts.go
@@ -20,6 +20,7 @@ func LoadHosts() *trie.DomainTrie {
 	}
 
 	t := trie.New()
+	h := map[string][]net.IP{}
 
 	sc := bufio.NewScanner(f)
 	for sc.Scan() {
@@ -52,25 +53,22 @@ func LoadHosts() *trie.DomainTrie {
 				continue
 			}
 
-			t.Insert(strings.ToLower(name), ip)
+			h[name] = append(h[name], ip)
 		}
 	}
 
+	for name, ips := range h {
+		t.Insert(strings.ToLower(name), ips)
+	}
+
+	defer f.Close()
 	return t
 }
 
 func hostsPath() string {
 	switch runtime.GOOS {
 	case "windows":
-		var sysRoot string
-		for _, e := range os.Environ() {
-			subs := strings.SplitN(e, "=", 2)
-			if subs[0] == "SystemRoot" && len(subs) == 2 {
-				sysRoot = subs[1]
-				break
-			}
-		}
-		return path.Join(sysRoot, "\\System32\\drivers\\etc\\hosts")
+		return path.Join(os.Getenv("SYSTEMROOT"), "\\System32\\drivers\\etc\\hosts")
 	default:
 		return "/etc/hosts"
 	}

diff --git a/dns/middleware.go b/dns/middleware.go
@@ -34,31 +34,33 @@ func withHosts(hosts *trie.DomainTrie) middleware {
 				return next(ctx, r)
 			}
 
-			ip := record.Data.(net.IP)
 			msg := r.Copy()
+			for _, data := range record.Data.([]net.IP) {
+				if v4 := data.To4(); v4 != nil && q.Qtype == D.TypeA {
+					rr := &D.A{}
+					rr.Hdr = D.RR_Header{Name: q.Name, Rrtype: D.TypeA, Class: D.ClassINET, Ttl: 10}
+					rr.A = v4
+
+					msg.Answer = append(msg.Answer, rr)
+				} else if v6 := data.To16(); v6 != nil && q.Qtype == D.TypeAAAA {
+					rr := &D.AAAA{}
+					rr.Hdr = D.RR_Header{Name: q.Name, Rrtype: D.TypeAAAA, Class: D.ClassINET, Ttl: 10}
+					rr.AAAA = v6
+
+					msg.Answer = append(msg.Answer, rr)
+				}
+			}
 
-			if v4 := ip.To4(); v4 != nil && q.Qtype == D.TypeA {
-				rr := &D.A{}
-				rr.Hdr = D.RR_Header{Name: q.Name, Rrtype: D.TypeA, Class: D.ClassINET, Ttl: dnsDefaultTTL}
-				rr.A = v4
-
-				msg.Answer = []D.RR{rr}
-			} else if v6 := ip.To16(); v6 != nil && q.Qtype == D.TypeAAAA {
-				rr := &D.AAAA{}
-				rr.Hdr = D.RR_Header{Name: q.Name, Rrtype: D.TypeAAAA, Class: D.ClassINET, Ttl: dnsDefaultTTL}
-				rr.AAAA = v6
+			if len(msg.Answer) > 0 {
+				ctx.SetType(context.DNSTypeHost)
+				msg.SetRcode(r, D.RcodeSuccess)
+				msg.Authoritative = true
+				msg.RecursionAvailable = true
 
-				msg.Answer = []D.RR{rr}
-			} else {
-				return next(ctx, r)
+				return msg, nil
 			}
 
-			ctx.SetType(context.DNSTypeHost)
-			msg.SetRcode(r, D.RcodeSuccess)
-			msg.Authoritative = true
-			msg.RecursionAvailable = true
-
-			return msg, nil
+			return next(ctx, r)
 		}
 	}
 }

diff --git a/tunnel/tunnel.go b/tunnel/tunnel.go
@@ -141,7 +141,7 @@ func preHandleMetadata(metadata *C.Metadata) error {
 				metadata.DNSMode = C.DNSFakeIP
 			} else if node := resolver.DefaultHosts.Search(host); node != nil {
 				// redir-host should lookup the hosts
-				metadata.DstIP = node.Data.(net.IP)
+				metadata.DstIP = node.Data.([]net.IP)[0]
 			}
 		} else if resolver.IsFakeIP(metadata.DstIP) {
 			return fmt.Errorf("fake DNS record %s missing", metadata.DstIP)
@@ -381,8 +381,7 @@ func match(metadata *C.Metadata) (C.Proxy, C.Rule, error) {
 	var processFound bool
 
 	if node := resolver.DefaultHosts.Search(metadata.Host); node != nil {
-		ip := node.Data.(net.IP)
-		metadata.DstIP = ip
+		metadata.DstIP = node.Data.([]net.IP)[0]
 		resolved = true
 	}