From d4f2f06c81b2cbee46fbfa482a65e1a20895bd9f Mon Sep 17 00:00:00 2001
From: Stephen Soltesz <stephen.soltesz@gmail.com>
Date: Wed, 10 Jul 2024 17:31:43 -0400
Subject: [PATCH] Apply provided site probability to Heartbeat registration
 (#23)

* Add Probability parameter
* Apply provided site probability to Heartbeat registration
* Add test cases to cover probability parameter
---
 .github/workflows/coverage.yml     |  2 +-
 handler/handler.go                 | 15 +++++++++
 handler/handler_test.go            | 52 +++++++++++++++++++++++++++++-
 internal/register/register.go      |  3 +-
 internal/register/register_test.go |  1 +
 5 files changed, 70 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml
index 4f611b7..6f1119b 100644
--- a/.github/workflows/coverage.yml
+++ b/.github/workflows/coverage.yml
@@ -13,7 +13,7 @@ jobs:
     - uses: actions/setup-go@v5
       with:
         go-version: "1.20"
-    - run: go test -coverprofile=profile.cov ./...
+    - run: go test -v -coverprofile=profile.cov ./...
     - name: Send coverage
       uses: shogo82148/actions-goveralls@v1
       with:
diff --git a/handler/handler.go b/handler/handler.go
index 4eb6e72..c37c1c7 100644
--- a/handler/handler.go
+++ b/handler/handler.go
@@ -198,6 +198,9 @@ func (s *Server) Register(rw http.ResponseWriter, req *http.Request) {
 	}
 	param.Geo = record
 	param.Network = s.ASN.AnnotateIP(param.IPv4)
+	// Override site probability with user-provided parameter.
+	// TODO(soltesz): include M-Lab override option
+	param.Probability = getProbability(req)
 	r := register.CreateRegisterResponse(param)
 
 	// Register the hostname under the organization zone.
@@ -366,3 +369,15 @@ func getClientIP(req *http.Request) string {
 	hip, _, _ := net.SplitHostPort(req.RemoteAddr)
 	return hip
 }
+
+func getProbability(req *http.Request) float64 {
+	prob := req.URL.Query().Get("probability")
+	if prob == "" {
+		return 1.0
+	}
+	p, err := strconv.ParseFloat(prob, 64)
+	if err != nil {
+		return 1.0
+	}
+	return p
+}
diff --git a/handler/handler_test.go b/handler/handler_test.go
index ca41987..c01b0e1 100644
--- a/handler/handler_test.go
+++ b/handler/handler_test.go
@@ -256,7 +256,57 @@ func TestServer_Register(t *testing.T) {
 	}{
 		{
 			name:   "success",
-			params: "?service=foo&organization=bar&iata=lga&ipv4=192.168.0.1",
+			params: "?service=foo&organization=bar&iata=lga&ipv4=192.168.0.1&probability=1.0",
+			Iata: &fakeIataFinder{
+				findRow: iata.Row{
+					IATA:      "lga",
+					Latitude:  -10,
+					Longitude: -10,
+				},
+			},
+			Maxmind: &fakeMaxmind{
+				// NOTE: this riduculous declaration is needed due to anonymous structs in the geoip2 package.
+				city: &geoip2.City{
+					Country: struct {
+						GeoNameID         uint              `maxminddb:"geoname_id"`
+						IsInEuropeanUnion bool              `maxminddb:"is_in_european_union"`
+						IsoCode           string            `maxminddb:"iso_code"`
+						Names             map[string]string `maxminddb:"names"`
+					}{
+						IsoCode: "US",
+					},
+					Subdivisions: []struct {
+						GeoNameID uint              `maxminddb:"geoname_id"`
+						IsoCode   string            `maxminddb:"iso_code"`
+						Names     map[string]string `maxminddb:"names"`
+					}{
+						{IsoCode: "NY", Names: map[string]string{"en": "New York"}},
+						{IsoCode: "ZZ", Names: map[string]string{"en": "fake thing"}},
+					},
+					Location: struct {
+						AccuracyRadius uint16  `maxminddb:"accuracy_radius"`
+						Latitude       float64 `maxminddb:"latitude"`
+						Longitude      float64 `maxminddb:"longitude"`
+						MetroCode      uint    `maxminddb:"metro_code"`
+						TimeZone       string  `maxminddb:"time_zone"`
+					}{
+						Latitude:  41,
+						Longitude: -73,
+					},
+				},
+			},
+			ASN: &fakeAsn{
+				ann: &annotator.Network{
+					ASNumber: 12345,
+				},
+			},
+			DNS:      &fakeDNS{},
+			wantName: "foo-lga12345-c0a80001.bar.sandbox.measurement-lab.org",
+			wantCode: http.StatusOK,
+		},
+		{
+			name:   "success-probability-invalid",
+			params: "?service=foo&organization=bar&iata=lga&ipv4=192.168.0.1&probability=invalid",
 			Iata: &fakeIataFinder{
 				findRow: iata.Row{
 					IATA:      "lga",
diff --git a/internal/register/register.go b/internal/register/register.go
index f82929f..a366b12 100644
--- a/internal/register/register.go
+++ b/internal/register/register.go
@@ -27,6 +27,7 @@ type Params struct {
 	Geo     *geoip2.City
 	Metro   iata.Row
 	Network *annotator.Network
+	Probability float64
 }
 
 // OrgZone generates the organization zone name based the organization and project.
@@ -101,7 +102,7 @@ func CreateRegisterResponse(p *Params) v0.RegisterResponse {
 				Machine:       machine,
 				Metro:         site[:3],
 				Project:       p.Project,
-				Probability:   1,
+				Probability:   p.Probability,
 				Site:          site,
 				Type:          "unknown", // should be overridden by node.
 				Uplink:        "unknown", // should be overridden by node.
diff --git a/internal/register/register_test.go b/internal/register/register_test.go
index b04acdd..555798b 100644
--- a/internal/register/register_test.go
+++ b/internal/register/register_test.go
@@ -62,6 +62,7 @@ func TestCreateRegisterResponse(t *testing.T) {
 				Network: &annotator.Network{
 					ASNumber: 12345,
 				},
+				Probability: 1.0,
 			},
 			want: v0.RegisterResponse{
 				Registration: &v0.Registration{