diff --git a/home/clients.go b/home/clients.go index ad246171..69534dba 100644 --- a/home/clients.go +++ b/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 diff --git a/home/clients_test.go b/home/clients_test.go index c4d1837e..f535d69f 100644 --- a/home/clients_test.go +++ b/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") +}