Merge: - whois: couldn't set info on existing auto-clients

Close #1059

* commit 'c299753b6733fd7841ed2acfc8f383df99964da4':
  - whois: couldn't set info on existing auto-clients

home/clients.go

@@ -73,9 +73,9 @@ type ClientHost struct {
 }
 
 type clientsContainer struct {
-	list    map[string]*Client    // name -> client
-	ipIndex map[string]*Client    // IP -> client
-	ipHost  map[string]ClientHost // IP -> Hostname
+	list    map[string]*Client     // name -> client
+	ipIndex map[string]*Client     // IP -> client
+	ipHost  map[string]*ClientHost // IP -> Hostname
 	lock    sync.Mutex
 }
 
@@ -87,7 +87,7 @@ func (clients *clientsContainer) Init() {
 	}
 	clients.list = make(map[string]*Client)
 	clients.ipIndex = make(map[string]*Client)
-	clients.ipHost = make(map[string]ClientHost)
+	clients.ipHost = make(map[string]*ClientHost)
 
 	go clients.periodicUpdate()
 }
@@ -303,7 +303,7 @@ func (clients *clientsContainer) SetWhoisInfo(ip string, info [][]string) {
 		return
 	}
 
-	ch = ClientHost{
+	ch = &ClientHost{
 		Source: ClientSourceWHOIS,
 	}
 	ch.WhoisInfo = info
@@ -324,16 +324,18 @@ func (clients *clientsContainer) AddHost(ip, host string, source clientSource) (
 		return false, nil
 	}
 
-	// check index
-	c, ok := clients.ipHost[ip]
-	if ok && c.Source > source {
+	// check auto-clients index
+	ch, ok := clients.ipHost[ip]
+	if ok && ch.Source > source {
 		return false, nil
-	}
-
-	clients.ipHost[ip] = ClientHost{
-		Host:      host,
-		Source:    source,
-		WhoisInfo: c.WhoisInfo,
+	} else if ok {
+		ch.Source = source
+	} else {
+		ch = &ClientHost{
+			Host:   host,
+			Source: source,
+		}
+		clients.ipHost[ip] = ch
 	}
 	log.Tracef("'%s' -> '%s' [%d]", ip, host, len(clients.ipHost))
 	return true, nil

home/clients_test.go

@@ -135,3 +135,29 @@ func TestClients(t *testing.T) {
 	// get
 	assert.True(t, clients.Exists("1.1.1.1", ClientSourceHostsFile))
 }
+
+func TestClientsWhois(t *testing.T) {
+	var c Client
+	clients := clientsContainer{}
+	clients.Init()
+
+	whois := [][]string{{"orgname", "orgname-val"}, {"country", "country-val"}}
+	// set whois info on new client
+	clients.SetWhoisInfo("1.1.1.255", whois)
+	assert.True(t, clients.ipHost["1.1.1.255"].WhoisInfo[0][1] == "orgname-val")
+
+	// set whois info on existing auto-client
+	_, _ = clients.AddHost("1.1.1.1", "host", ClientSourceRDNS)
+	clients.SetWhoisInfo("1.1.1.1", whois)
+	assert.True(t, clients.ipHost["1.1.1.1"].WhoisInfo[0][1] == "orgname-val")
+
+	// set whois info on existing client
+	c = Client{
+		IP:   "1.1.1.2",
+		Name: "client1",
+	}
+	_, _ = clients.Add(c)
+	clients.SetWhoisInfo("1.1.1.2", whois)
+	assert.True(t, clients.ipIndex["1.1.1.2"].WhoisInfo[0][1] == "orgname-val")
+	_ = clients.Del("client1")
+}