From a7dc7fc6e847b1fe2d20c418568a6d832659c379 Mon Sep 17 00:00:00 2001
From: Myzel394 <50424412+Myzel394@users.noreply.github.com>
Date: Sun, 8 Sep 2024 22:08:00 +0200
Subject: [PATCH] fix(hosts): Fix ip address detection

---
 handlers/hosts/analyzer/resolver.go      | 23 ++++++++++++-----
 handlers/hosts/analyzer/resolver_test.go | 33 ++++++++++++++++++++++++
 handlers/hosts/indexes/resolver.go       |  6 ++---
 3 files changed, 53 insertions(+), 9 deletions(-)

diff --git a/handlers/hosts/analyzer/resolver.go b/handlers/hosts/analyzer/resolver.go
index d59b412..bedbfe5 100644
--- a/handlers/hosts/analyzer/resolver.go
+++ b/handlers/hosts/analyzer/resolver.go
@@ -19,9 +19,9 @@ func createEntry(
 	}
 
 	if ipv4 := ip.To4(); ipv4 != nil {
-		entry.IPv4Address = ipv4
+		entry.IPv4Address = &ipv4
 	} else if ipv6 := ip.To16(); ipv6 != nil {
-		entry.IPv6Address = ipv6
+		entry.IPv6Address = &ipv6
 	}
 
 	return entry
@@ -35,7 +35,7 @@ type hostnameEntry struct {
 func createResolverFromParser(p ast.HostsParser) (indexes.Resolver, []common.LSPError) {
 	errors := make([]common.LSPError, 0)
 	resolver := indexes.Resolver{
-		Entries: make(map[string]indexes.ResolverEntry),
+		Entries: make(map[string]*indexes.ResolverEntry),
 	}
 
 	it := p.Tree.Entries.Iterator()
@@ -64,12 +64,22 @@ func createResolverFromParser(p ast.HostsParser) (indexes.Resolver, []common.LSP
 			)
 
 			for _, hostName := range hostNames {
-				entry := createEntry(
+				newResolv := createEntry(
 					lineNumber,
 					entry.IPAddress.Value.IP,
 				)
 
 				if resolv, found := resolver.Entries[hostName.HostName]; found {
+					if resolv.IPv4Address == nil && newResolv.IPv4Address != nil {
+						resolv.IPv4Address = newResolv.IPv4Address
+						continue
+					}
+
+					if resolv.IPv6Address == nil && newResolv.IPv6Address != nil {
+						resolv.IPv6Address = newResolv.IPv6Address
+						continue
+					}
+
 					errors = append(
 						errors,
 						common.LSPError{
@@ -80,9 +90,10 @@ func createResolverFromParser(p ast.HostsParser) (indexes.Resolver, []common.LSP
 							},
 						},
 					)
-				} else {
-					resolver.Entries[hostName.HostName] = entry
+					continue
 				}
+
+				resolver.Entries[hostName.HostName] = &newResolv
 			}
 		}
 	}
diff --git a/handlers/hosts/analyzer/resolver_test.go b/handlers/hosts/analyzer/resolver_test.go
index 0dd8337..d4deb2f 100644
--- a/handlers/hosts/analyzer/resolver_test.go
+++ b/handlers/hosts/analyzer/resolver_test.go
@@ -115,3 +115,36 @@ func TestResolverEntriesWithComplexOverlapping(
 		t.Errorf("Expected test.com to have no IPv4 address, but got %v", resolver.Entries["test.com"].IPv4Address)
 	}
 }
+
+func TestResolverEntriesWithDoubleHostNameButDifferentIPs(
+	t *testing.T,
+) {
+	input := utils.Dedent(`
+127.0.0.1 hello.com
+::1 hello.com
+`)
+	parser := ast.NewHostsParser()
+	errors := parser.Parse(input)
+
+	if len(errors) != 0 {
+		t.Fatalf("PARER FAILED! Expected no errors, but got %v", errors)
+	}
+
+	resolver, errors := createResolverFromParser(parser)
+
+	if len(errors) != 0 {
+		t.Fatalf("Expected no errors, but got %v", errors)
+	}
+
+	if len(resolver.Entries) != 1 {
+		t.Errorf("Expected 1 entry, but got %v", len(resolver.Entries))
+	}
+
+	if !(resolver.Entries["hello.com"].IPv4Address.String() == "127.0.0.1") {
+		t.Errorf("Expected hello.com's ipv4 address to be 127.0.0.1, but got %v", resolver.Entries["hello.com"].IPv4Address)
+	}
+
+	if !(resolver.Entries["hello.com"].IPv6Address.String() == "::1") {
+		t.Errorf("Expected hello.com's ipv6 address to be ::1, but got %v", resolver.Entries["hello.com"].IPv6Address)
+	}
+}
diff --git a/handlers/hosts/indexes/resolver.go b/handlers/hosts/indexes/resolver.go
index aabb6f4..dc3655c 100644
--- a/handlers/hosts/indexes/resolver.go
+++ b/handlers/hosts/indexes/resolver.go
@@ -5,8 +5,8 @@ import (
 )
 
 type ResolverEntry struct {
-	IPv4Address net.IP
-	IPv6Address net.IP
+	IPv4Address *net.IP
+	IPv6Address *net.IP
 	Line        uint32
 }
 
@@ -19,5 +19,5 @@ func (e ResolverEntry) GetInfo() string {
 }
 
 type Resolver struct {
-	Entries map[string]ResolverEntry
+	Entries map[string]*ResolverEntry
 }