From db0ef5cf5c185411caeb4d1089da49a065844dc5 Mon Sep 17 00:00:00 2001
From: Alvaro Saurin <alvaro.saurin@gmail.com>
Date: Wed, 13 May 2015 10:29:11 +0200
Subject: [PATCH] Remove duplicates in answers

---
 nameserver/zone_lookup.go | 63 ++++++++++++++++++++++++++++++++-------
 1 file changed, 52 insertions(+), 11 deletions(-)

diff --git a/nameserver/zone_lookup.go b/nameserver/zone_lookup.go
index 6eac008e..b38971d3 100644
--- a/nameserver/zone_lookup.go
+++ b/nameserver/zone_lookup.go
@@ -6,18 +6,52 @@ import (
 	"time"
 )
 
+type uniqZoneRecordKey struct {
+	name string
+	ipv4 IPv4
+}
+
+// A group of ZoneRecords where there are no duplicates (according to the name & IPv4)
+type uniqZoneRecords map[uniqZoneRecordKey]ZoneRecord
+
+func newUniqZoneRecords() uniqZoneRecords {
+	return make(uniqZoneRecords, 0)
+}
+
+// Add a new ZoneRecord to the group
+func (uzr *uniqZoneRecords) add(zr ZoneRecord) {
+	key := uniqZoneRecordKey{zr.Name(), ipToIPv4(zr.IP())}
+	(*uzr)[key] = zr
+}
+
+// Return the group as an slice
+func (uzr *uniqZoneRecords) toSlice() []ZoneRecord {
+	res := make([]ZoneRecord, len(*uzr))
+	i := 0
+	for _, r := range *uzr {
+		res[i] = r
+		i++
+	}
+	return res
+}
+
+//////////////////////////////////////////////////////////////////////////////
+
+// Lookup in the database for locally-introduced information
 func (zone *zoneDb) lookup(target string, lfun func(ns *namesSet) []*recordEntry) (res []ZoneRecord, err error) {
+	uniq := newUniqZoneRecords()
 	for identName, nameset := range zone.idents {
 		if identName != defaultRemoteIdent {
 			for _, ze := range lfun(nameset) {
-				res = append(res, ze)
+				uniq.add(ze)
 			}
 		}
 	}
-	if len(res) == 0 {
-		err = LookupError(target)
+	if len(uniq) == 0 {
+		return nil, LookupError(target)
+	} else {
+		return uniq.toSlice(), nil
 	}
-	return
 }
 
 // Perform a lookup for a name in the zone
@@ -63,16 +97,19 @@ func (zone *zoneDb) domainLookup(target string, lfun ZoneLookupFunc) (res []Zone
 	res = make([]ZoneRecord, len(lanswers))
 	zone.mx.Lock()
 	now := zone.clock.Now()
-	for i, answer := range lanswers {
-		res[i], err = zone.getNamesSet(defaultRemoteIdent).addIPToName(answer, now)
+	uniq := newUniqZoneRecords()
+	for _, answer := range lanswers {
+		r, err := zone.getNamesSet(defaultRemoteIdent).addIPToName(answer, now)
 		if err != nil {
 			zone.mx.Unlock()
 			Warning.Printf("[zonedb] '%s' insertion for %s failed: %s", answer, target, err)
 			return nil, err
 		}
+		uniq.add(r)
 	}
 	zone.mx.Unlock()
-	return
+
+	return uniq.toSlice(), nil
 }
 
 // Perform a lookup for a name in the zone
@@ -83,6 +120,7 @@ func (zone *zoneDb) DomainLookupName(name string) (res []ZoneRecord, err error)
 
 	zone.mx.RLock()
 	now := zone.clock.Now()
+	uniq := newUniqZoneRecords()
 	for identName, nameset := range zone.idents {
 		for _, ze := range nameset.getEntriesForName(name) {
 			// filter the entries with expired TTL
@@ -91,7 +129,7 @@ func (zone *zoneDb) DomainLookupName(name string) (res []ZoneRecord, err error)
 				Debug.Printf("[zonedb] '%s': expired entry '%s' ignored: removing", name, ze)
 				nameset.deleteName(name)
 			} else {
-				res = append(res, ze)
+				uniq.add(ze)
 			}
 		}
 		if identName != defaultRemoteIdent {
@@ -100,8 +138,9 @@ func (zone *zoneDb) DomainLookupName(name string) (res []ZoneRecord, err error)
 	}
 	zone.mx.RUnlock()
 
-	if len(res) > 0 {
+	if len(uniq) > 0 {
 		Debug.Printf("[zonedb] '%s' resolved in local database", name)
+		res = uniq.toSlice()
 	} else {
 		res, err = zone.domainLookup(name, zone.mdnsCli.LookupName)
 	}
@@ -126,6 +165,7 @@ func (zone *zoneDb) DomainLookupInaddr(inaddr string) (res []ZoneRecord, err err
 
 	zone.mx.RLock()
 	now := zone.clock.Now()
+	uniq := newUniqZoneRecords()
 	for identName, nameset := range zone.idents {
 		for _, ze := range nameset.getEntriesForIP(revIPv4) {
 			// filter the entries with expired TTL
@@ -134,7 +174,7 @@ func (zone *zoneDb) DomainLookupInaddr(inaddr string) (res []ZoneRecord, err err
 				Debug.Printf("[zonedb] '%s': expired entry '%s' ignored: removing", revIPv4, ze)
 				nameset.deleteIP(revIPv4)
 			} else {
-				res = append(res, ZoneRecord(ze))
+				uniq.add(ze)
 				if identName != defaultRemoteIdent {
 					nameset.touchName(ze.Name(), now)
 				}
@@ -143,8 +183,9 @@ func (zone *zoneDb) DomainLookupInaddr(inaddr string) (res []ZoneRecord, err err
 	}
 	zone.mx.RUnlock()
 
-	if len(res) > 0 {
+	if len(uniq) > 0 {
 		Debug.Printf("[zonedb] '%s' resolved in local database", inaddr)
+		res = uniq.toSlice()
 	} else {
 		res, err = zone.domainLookup(inaddr, zone.mdnsCli.LookupInaddr)
 	}