Add support for site annotations to the site package (#285)

* Add support for site annotations to the site package. * Remove some debug strings. * Add test cases for IPv6. * Remove double import. * Typo in a comment. * Address review comments.
m-lab · Apr 27, 2021 · 6f5ace9 · 6f5ace9
1 parent 8e4d70d
commit 6f5ace9
Show file tree

Hide file tree

Showing 3 changed files with 85 additions and 70 deletions.
diff --git a/site/site.go b/site/site.go
@@ -5,7 +5,9 @@ import (
 	"context"
 	"encoding/json"
 	"flag"
+	"fmt"
 	"log"
+	"net"
 	"time"
 
 	"github.com/m-lab/go/content"
@@ -29,9 +31,9 @@ func init() {
 }
 
 // Annotate adds site annotation for a site/machine
-func Annotate(site, machine string, server *uuid.ServerAnnotations) {
+func Annotate(ip string, server *uuid.ServerAnnotations) {
 	if globalAnnotator != nil {
-		globalAnnotator.Annotate(site, machine, server)
+		globalAnnotator.Annotate(ip, server)
 	}
 }
 
@@ -40,7 +42,7 @@ func LoadFrom(ctx context.Context, js content.Provider, retiredJS content.Provid
 	globalAnnotator = &annotator{
 		siteinfoSource:        js,
 		siteinfoRetiredSource: retiredJS,
-		sites:                 make(map[string]uuid.ServerAnnotations, 200),
+		networks:              make(map[string]uuid.ServerAnnotations, 400),
 	}
 	err := globalAnnotator.load(ctx)
 	log.Println(len(globalAnnotator.sites), "sites loaded")
@@ -85,7 +87,8 @@ type annotator struct {
 	siteinfoRetiredSource content.Provider
 	// Each site has a single ServerAnnotations struct, which
 	// is later customized for each machine.
-	sites map[string]uuid.ServerAnnotations
+	sites    map[string]uuid.ServerAnnotations
+	networks map[string]uuid.ServerAnnotations
 }
 
 // missing is used if annotation is requested for a non-existant server.
@@ -98,22 +101,33 @@ var missing = uuid.ServerAnnotations{
 	},
 }
 
-// Annotate annotates the server with the approprate annotations.
-func (sa *annotator) Annotate(site, machine string, server *uuid.ServerAnnotations) {
+// Annotate annotates the server with the appropriate annotations.
+func (sa *annotator) Annotate(ip string, server *uuid.ServerAnnotations) {
 	if server == nil {
 		return
 	}
 
-	server.Machine = machine
-	server.Site = site
-	s, ok := sa.sites[site]
-	if !ok {
-		server.Geo = missing.Geo
-		server.Network = missing.Network
+	parsedIP := net.ParseIP(ip)
+	if parsedIP == nil {
 		return
 	}
-	server.Geo = s.Geo
-	server.Network = s.Network
+
+	// Find CIDR corresponding to the provided ip.
+	// All of our subnets are /26 if IPv4, /64 if IPv6.
+	var cidr string
+	if parsedIP.To4() == nil {
+		mask := net.CIDRMask(64, 128)
+		cidr = fmt.Sprintf("%s/64", parsedIP.Mask(mask))
+	} else {
+		mask := net.CIDRMask(26, 32)
+		cidr = fmt.Sprintf("%s/26", parsedIP.Mask(mask))
+	}
+
+	if ann, ok := sa.networks[cidr]; ok {
+		*server = ann
+	} else {
+		*server = missing
+	}
 }
 
 // load loads siteinfo dataset and returns them.
@@ -151,7 +165,34 @@ func (sa *annotator) load(ctx context.Context) error {
 	for _, ann := range s {
 		// Machine should always be empty, filled in later.
 		ann.Annotation.Machine = ""
-		sa.sites[ann.Annotation.Site] = ann.Annotation // Copy out of array.
+
+		// Make a map of CIDR -> Annotation.
+		// Verify that the CIDRs are valid by trying to parse them.
+		// If either the IPv4 or IPv6 CIDRs are wrong, the entry is
+		// discarded. The IPv6 CIDR can be empty in some cases.
+		if ann.Network.IPv4 == "" {
+			continue
+		}
+		_, _, err := net.ParseCIDR(ann.Network.IPv4)
+		if err != nil {
+			log.Printf("Found incorrect IPv4 in siteinfo: %s\n",
+				ann.Network.IPv4)
+			continue
+		}
+
+		// Check the IPv6 CIDR only if not empty.
+		if ann.Network.IPv6 != "" {
+			_, _, err = net.ParseCIDR(ann.Network.IPv6)
+			if err != nil {
+				log.Printf("Found incorrect IPv6 in siteinfo: %s\n",
+					ann.Network.IPv6)
+				continue
+			}
+			sa.networks[ann.Network.IPv6] = ann.Annotation
+		}
+
+		sa.networks[ann.Network.IPv4] = ann.Annotation
 	}
+
 	return nil
 }
diff --git a/site/site_test.go b/site/site_test.go
@@ -9,7 +9,7 @@ import (
 	"time"
 
 	"github.com/go-test/deep"
-	"github.com/m-lab/etl/site"
+	"github.com/m-lab/annotation-service/site"
 	"github.com/m-lab/go/content"
 	"github.com/m-lab/go/flagx"
 	"github.com/m-lab/go/osx"
@@ -52,9 +52,8 @@ func TestBasic(t *testing.T) {
 	setUp()
 	ctx := context.Background()
 	site.LoadFrom(ctx, localRawfile, retiredFile)
-	missingServerAnn := annotator.ServerAnnotations{
-		Machine: "foo",
-		Site:    "bar",
+
+	var missingServerAnn = annotator.ServerAnnotations{
 		Geo: &annotator.Geolocation{
 			Missing: true,
 		},
@@ -64,7 +63,7 @@ func TestBasic(t *testing.T) {
 	}
 
 	defaultServerAnn := annotator.ServerAnnotations{
-		Machine: "mlab1",
+		Machine: "",
 		Site:    "lga03",
 		Geo: &annotator.Geolocation{
 			ContinentCode: "NA",
@@ -83,7 +82,7 @@ func TestBasic(t *testing.T) {
 	}
 
 	retiredServerann := annotator.ServerAnnotations{
-		Machine: "mlab1",
+		Machine: "",
 		Site:    "acc01",
 		Geo: &annotator.Geolocation{
 			ContinentCode: "AF",
@@ -102,33 +101,40 @@ func TestBasic(t *testing.T) {
 	}
 
 	tests := []struct {
-		name          string
-		site, machine string
-		want          annotator.ServerAnnotations
+		name string
+		ip   string
+		want annotator.ServerAnnotations
 	}{
 		{
-			name:    "success",
-			site:    "lga03",
-			machine: "mlab1",
-			want:    defaultServerAnn,
+			name: "success",
+			ip:   "64.86.148.130",
+			want: defaultServerAnn,
+		},
+		{
+			name: "success-ipv6",
+			ip:   "2001:5a0:4300::1",
+			want: defaultServerAnn,
+		},
+		{
+			name: "success-retired-site",
+			ip:   "196.201.2.192",
+			want: retiredServerann,
 		},
 		{
-			name:    "success-retired-site",
-			site:    "acc01",
-			machine: "mlab1",
-			want:    retiredServerann,
+			name: "missing",
+			ip:   "0.0.0.0",
+			want: missingServerAnn,
 		},
 		{
-			name:    "missing",
-			site:    "bar",
-			machine: "foo",
-			want:    missingServerAnn,
+			name: "missing-ipv6",
+			ip:   "::1",
+			want: missingServerAnn,
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			ann := annotator.ServerAnnotations{}
-			site.Annotate(tt.site, tt.machine, &ann)
+			site.Annotate(tt.ip, &ann)
 			if diff := deep.Equal(ann, tt.want); diff != nil {
 				t.Errorf("Annotate() failed; %s", strings.Join(diff, "\n"))
 			}
@@ -155,7 +161,7 @@ func TestNilServer(t *testing.T) {
 		t.Error(err)
 	}
 	// Should not panic!  Nothing else to check.
-	site.Annotate("lga03", "mlab1", nil)
+	site.Annotate("64.86.148.128", nil)
 }
 
 func TestCorrupt(t *testing.T) {

diff --git a/site/testdata/annotations.json b/site/testdata/annotations.json
@@ -28,38 +28,6 @@
          "IPv6": "2001:5a0:4300::/64"
       }
    },
-   {
-      "Annotation": {
-         "Geo": {
-            "City": "New York"
-         },
-         "Network": {
-            "ASName": "TATA COMMUNICATIONS (AMERICA) INC"
-         },
-         "Site": "six02"
-      },
-      "Name": "six02",
-      "Network": {
-         "IPv4": "",
-         "IPv6": "2001:5a0:4300::/64"
-      }
-   },
-   {
-      "Annotation": {
-         "Geo": {
-            "City": "New York"
-         },
-         "Network": {
-            "ASName": "TATA COMMUNICATIONS (AMERICA) INC"
-         },
-         "Site": "six01"
-      },
-      "Name": "six01",
-      "Network": {
-         "IPv4": "64.86.148.128/26",
-         "IPv6": ""
-      }
-   },
    {
       "Annotation": {
          "Geo": {