Updated and added unit testing for the ASNCache

requests/asncache.go

@@ -72,19 +72,15 @@ func (c *ASNCache) Update(req *ASNRequest) {
 	c.Lock()
 	defer c.Unlock()
 
-	if _, found := c.cache[req.ASN]; !found {
+	as, found := c.cache[req.ASN]
+	if !found {
 		c.cache[req.ASN] = req
 		if req.Netblocks == nil {
 			req.Netblocks = stringset.New(req.Prefix)
 		}
 		return
 	}
-
-	as := c.cache[req.ASN]
 	// This is additional information for an ASN entry
-	if as.Prefix == "" && req.Prefix != "" {
-		as.Prefix = req.Prefix
-	}
 	if as.CC == "" && req.CC != "" {
 		as.CC = req.CC
 	}
@@ -94,12 +90,12 @@ func (c *ASNCache) Update(req *ASNRequest) {
 	if as.AllocationDate.IsZero() && !req.AllocationDate.IsZero() {
 		as.AllocationDate = req.AllocationDate
 	}
-	if as.Description == "" && req.Description != "" {
+	if len(as.Description) < len(req.Description) {
 		as.Description = req.Description
 	}
-	if req.Netblocks == nil {
-		as.Netblocks.Union(stringset.New(req.Prefix))
-	} else {
+
+	as.Netblocks.Insert(req.Prefix)
+	if req.Netblocks != nil {
 		as.Netblocks.Union(req.Netblocks)
 	}
 }
@@ -140,12 +136,16 @@ func (c *ASNCache) AddrSearch(addr string) *ASNRequest {
 		}
 	}
 
+	prefix := entry.IPNet.String()
+	netblocks := stringset.New(prefix)
+	netblocks.Union(entry.Data.Netblocks)
+
 	return &ASNRequest{
 		Address:     addr,
 		ASN:         entry.Data.ASN,
 		CC:          entry.Data.CC,
-		Prefix:      entry.IPNet.String(),
-		Netblocks:   stringset.New(entry.IPNet.String()),
+		Prefix:      prefix,
+		Netblocks:   netblocks,
 		Description: entry.Data.Description,
 		Tag:         RIR,
 		Source:      "RIR",

requests/asncache_test.go

@@ -0,0 +1,121 @@
+// Copyright 2021 Jeff Foley. All rights reserved.
+// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
+
+package requests
+
+import (
+	"net"
+	"testing"
+	"time"
+
+	"github.com/caffix/stringset"
+)
+
+func TestEmpty(t *testing.T) {
+	cache := NewASNCache()
+
+	if entry := cache.ASNSearch(0); entry != nil {
+		t.Errorf("ASNSearch returned a non-nil value when empty: %v", entry)
+	}
+	if entry := cache.AddrSearch("72.237.4.113"); entry != nil {
+		t.Errorf("AddrSearch returned a non-nil value when empty: %v", entry)
+	}
+}
+
+func TestUpdate(t *testing.T) {
+	cache := NewASNCache()
+
+	if entry := cache.AddrSearch("72.237.4.113"); entry != nil {
+		t.Errorf("AddrSearch returned a non-nil value when empty: %v", entry)
+	}
+
+	cache.Update(&ASNRequest{
+		Address: "72.237.4.113",
+		ASN:     26808,
+		Prefix:  "72.237.4.0/24",
+		Tag:     RIR,
+		Source:  "RIR",
+	})
+
+	if entry := cache.AddrSearch("72.237.4.113"); entry == nil {
+		t.Errorf("AddrSearch returned nil after updated with the correct data")
+	}
+
+	cache.Update(&ASNRequest{
+		Address:        "72.237.4.113",
+		ASN:            26808,
+		Prefix:         "72.237.4.0/24",
+		CC:             "US",
+		Registry:       "ARIN",
+		AllocationDate: time.Now(),
+		Description:    "UTICA-COLLEGE",
+		Netblocks:      stringset.New("72.237.4.0/24", "8.24.68.0/23"),
+		Tag:            RIR,
+		Source:         "RIR",
+	})
+
+	if entry := cache.AddrSearch("72.237.4.113"); entry == nil || entry.CC != "US" || entry.Description != "UTICA-COLLEGE" {
+		t.Errorf("Update failed to enhance the ASN entry with the more detailed data")
+	}
+	if entry := cache.AddrSearch("8.24.68.1"); entry == nil || entry.ASN != 26808 {
+		t.Errorf("Update failed to add the new netblock to the ASN")
+	}
+}
+
+func TestASNSearch(t *testing.T) {
+	cache := NewASNCache()
+
+	if entry := cache.ASNSearch(26808); entry != nil {
+		t.Errorf("ASNSearch returned a non-nil value when empty: %v", entry)
+	}
+
+	cache.Update(&ASNRequest{
+		Address: "72.237.4.113",
+		ASN:     26808,
+		Prefix:  "72.237.4.0/24",
+		Tag:     RIR,
+		Source:  "RIR",
+	})
+
+	if entry := cache.ASNSearch(26808); entry == nil {
+		t.Errorf("ASNSearch returned nil after updated with the correct data")
+	}
+}
+
+func TestAddrSearch(t *testing.T) {
+	cache := NewASNCache()
+
+	if entry := cache.AddrSearch("127.0.0.1"); entry == nil || entry.ASN != 0 {
+		t.Errorf("AddrSearch returned nil when searching for a reserved network address block")
+	}
+	if entry := cache.AddrSearch("72.237.4.113"); entry != nil {
+		t.Errorf("AddrSearch returned a non-nil value when empty: %v", entry)
+	}
+
+	cache.Update(&ASNRequest{
+		Address:        "72.237.4.113",
+		ASN:            26808,
+		Prefix:         "72.237.4.0/24",
+		CC:             "US",
+		Registry:       "ARIN",
+		AllocationDate: time.Now(),
+		Description:    "UTICA-COLLEGE",
+		Netblocks:      stringset.New("72.237.4.0/24", "8.24.68.0/23"),
+		Tag:            RIR,
+		Source:         "RIR",
+	})
+
+	if entry := cache.AddrSearch("72.237.4.120"); entry == nil {
+		t.Errorf("AddrSearch returned nil after updated with the correct data")
+	}
+
+	addr := "8.24.68.1"
+	ip := net.ParseIP(addr)
+	entry := cache.AddrSearch(addr)
+	if entry == nil {
+		t.Fatalf("AddrSearch returned nil for an address known to the cache")
+	}
+	if _, ipnet, err := net.ParseCIDR(entry.Prefix); err != nil || !ipnet.Contains(ip) {
+		t.Errorf("AddrSearch returned the wrong Prefix value for the provided IP address: %s", entry.Prefix)
+	}
+}