From a36f346a634b08293d7222c8a4e4808b830e45af Mon Sep 17 00:00:00 2001
From: Stephen Soltesz <stephen.soltesz@gmail.com>
Date: Tue, 13 Feb 2024 18:21:59 -0500
Subject: [PATCH] Add IPv4 only annotator (#64)

* Add IPv4 only annotator
* Update coveralls to travis-pro
---
 .travis.yml              |  2 +-
 asnannotator/asn.go      | 38 +++++++++++++++++++++++++++++---------
 asnannotator/asn_test.go | 38 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 68 insertions(+), 10 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 5e4fa00..8b38adc 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -14,4 +14,4 @@ script:
 - go build ./...
 - go test ./... -cover=1 -coverprofile=_c.cov
 - go test ./... -race
-- $GOPATH/bin/goveralls -service=travis-ci -coverprofile=_c.cov
+- $GOPATH/bin/goveralls -service=travis-pro -coverprofile=_c.cov
diff --git a/asnannotator/asn.go b/asnannotator/asn.go
index ed5c113..83f2489 100644
--- a/asnannotator/asn.go
+++ b/asnannotator/asn.go
@@ -36,6 +36,18 @@ type asnAnnotator struct {
 	asnames    ipinfo.ASNames
 }
 
+// NewIPv4 makes a new IPv4-only Annotator that uses IP addresses to lookup ASN metadata for
+// that IP based on the current copy of RouteViews data stored in the given providers.
+func NewIPv4(ctx context.Context, as4 content.Provider) ASNAnnotator {
+	a := &asnAnnotator{
+		as4: as4,
+	}
+	var err error
+	a.asn4, err = load(ctx, as4, nil)
+	rtx.Must(err, "Could not load Routeviews IPv4 ASN db")
+	return a
+}
+
 // New makes a new Annotator that uses IP addresses to lookup ASN metadata for
 // that IP based on the current copy of RouteViews data stored in the given providers.
 func New(ctx context.Context, as4 content.Provider, as6 content.Provider, asnamedata content.Provider, localIPs []net.IP) ASNAnnotator {
@@ -97,6 +109,10 @@ func (a *asnAnnotator) annotateIPHoldingLock(src string) *annotator.Network {
 		metrics.ASNSearches.WithLabelValues("ipv4-success").Inc()
 		return ann
 	}
+	if a.asn6 == nil {
+		ann.Missing = true
+		return ann
+	}
 
 	ipnet, err = a.asn6.Search(src)
 	if err != nil {
@@ -126,15 +142,19 @@ func (a *asnAnnotator) Reload(ctx context.Context) {
 		log.Println("Could not reload v4 routeviews:", err)
 		return
 	}
-	new6, err := load(ctx, a.as6, a.asn6)
-	if err != nil {
-		log.Println("Could not reload v6 routeviews:", err)
-		return
-	}
-	newnames, err := loadNames(ctx, a.asnamedata, a.asnames)
-	if err != nil {
-		log.Println("Could not reload asnames from ipinfo:", err)
-		return
+	var new6 routeview.Index
+	var newnames ipinfo.ASNames
+	if a.as6 != nil {
+		new6, err = load(ctx, a.as6, a.asn6)
+		if err != nil {
+			log.Println("Could not reload v6 routeviews:", err)
+			return
+		}
+		newnames, err = loadNames(ctx, a.asnamedata, a.asnames)
+		if err != nil {
+			log.Println("Could not reload asnames from ipinfo:", err)
+			return
+		}
 	}
 	// Don't acquire the lock until after the data is in RAM.
 	a.m.Lock()
diff --git a/asnannotator/asn_test.go b/asnannotator/asn_test.go
index 946d12f..70894bd 100644
--- a/asnannotator/asn_test.go
+++ b/asnannotator/asn_test.go
@@ -305,3 +305,41 @@ func TestNewFake(t *testing.T) {
 		t.Error("Should have had a missing return value, not", n3)
 	}
 }
+
+func Test_IPv4Annotator_AnnotateIP(t *testing.T) {
+	setUp()
+	tests := []struct {
+		name string
+		addr string
+		want annotator.Network
+	}{
+		{
+			name: "success-ipv4",
+			addr: "1.0.0.1",
+			want: annotator.Network{
+				CIDR:     "1.0.0.0/24",
+				ASNumber: 13335,
+				Systems: []annotator.System{
+					{ASNs: []uint32{13335}},
+				},
+			},
+		},
+		{
+			name: "success-ipv6",
+			addr: "2001:200::1",
+			want: annotator.Network{
+				Missing: true,
+			},
+		},
+	}
+	ctx := context.Background()
+	a := NewIPv4(ctx, local4Rawfile)
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := a.AnnotateIP(tt.addr)
+			if diff := deep.Equal(*got, tt.want); diff != nil {
+				t.Error("AnnotateIP() wrong value; got!=want", diff)
+			}
+		})
+	}
+}