add support for country wide checks

netbirdio · Jan 31, 2024 · d2d5f2b · d2d5f2b
1 parent 00c7c00
commit d2d5f2b
Show file tree

Hide file tree

Showing 3 changed files with 101 additions and 9 deletions.
diff --git a/management/server/http/posture_checks_handler.go b/management/server/http/posture_checks_handler.go
@@ -3,6 +3,7 @@ package http
 import (
 	"encoding/json"
 	"net/http"
+	"regexp"
 
 	"github.com/gorilla/mux"
 	"github.com/rs/xid"
@@ -247,8 +248,10 @@ func validatePostureChecksUpdate(req api.PostureCheckUpdate) error {
 			if loc.CountryCode == "" {
 				return status.Errorf(status.InvalidArgument, "country code for geolocation check shouldn't be empty")
 			}
-			if loc.CityName == "" {
-				return status.Errorf(status.InvalidArgument, "city name for geolocation check shouldn't be empty")
+
+			countryCodeRegex := regexp.MustCompile("^[a-zA-Z]{2}$")
+			if !countryCodeRegex.MatchString(loc.CountryCode) {
+				return status.Errorf(status.InvalidArgument, "country code must be 2 letters (ISO 3166-1 alpha-2 format)")
 			}
 		}
 

diff --git a/management/server/posture/geo_location.go b/management/server/posture/geo_location.go
@@ -31,16 +31,17 @@ type GeoLocationCheck struct {
 
 func (g *GeoLocationCheck) Check(peer nbpeer.Peer) (bool, error) {
 	for _, loc := range g.Locations {
-		if loc.CountryCode == peer.Meta.Location.CountryCode && loc.CityName == peer.Meta.Location.CityName {
-			switch g.Action {
-			case GeoLocationActionDeny:
-				return false, nil
-			case GeoLocationActionAllow:
-				return true, nil
+		if loc.CountryCode == peer.Meta.Location.CountryCode {
+			if loc.CityName == "" || loc.CityName == peer.Meta.Location.CityName {
+				switch g.Action {
+				case GeoLocationActionDeny:
+					return false, nil
+				case GeoLocationActionAllow:
+					return true, nil
+				}
 			}
 		}
 	}
-
 	// At this point, no location in the list matches the peer's location
 	// For action deny and no location match, allow the peer
 	if g.Action == GeoLocationActionDeny {

diff --git a/management/server/posture/geo_location_test.go b/management/server/posture/geo_location_test.go
@@ -42,6 +42,27 @@ func TestGeoLocationCheck_Check(t *testing.T) {
 			wantErr: false,
 			isValid: true,
 		},
+		{
+			name: "Peer location matches the location in the allow country only",
+			input: peer.Peer{
+				Meta: peer.PeerSystemMeta{
+					Location: Location{
+						CountryCode: "DE",
+						CityName:    "Berlin",
+					},
+				},
+			},
+			check: GeoLocationCheck{
+				Locations: []Location{
+					{
+						CountryCode: "DE",
+					},
+				},
+				Action: GeoLocationActionAllow,
+			},
+			wantErr: false,
+			isValid: true,
+		},
 		{
 			name: "Peer location doesn't match the location in the allow sets",
 			input: peer.Peer{
@@ -68,6 +89,27 @@ func TestGeoLocationCheck_Check(t *testing.T) {
 			wantErr: false,
 			isValid: false,
 		},
+		{
+			name: "Peer location doesn't match the location in the allow country only",
+			input: peer.Peer{
+				Meta: peer.PeerSystemMeta{
+					Location: Location{
+						CountryCode: "DE",
+						CityName:    "Frankfurt am Main",
+					},
+				},
+			},
+			check: GeoLocationCheck{
+				Locations: []Location{
+					{
+						CountryCode: "US",
+					},
+				},
+				Action: GeoLocationActionAllow,
+			},
+			wantErr: false,
+			isValid: false,
+		},
 		{
 			name: "Peer location matches the location in the deny sets",
 			input: peer.Peer{
@@ -94,6 +136,30 @@ func TestGeoLocationCheck_Check(t *testing.T) {
 			wantErr: false,
 			isValid: false,
 		},
+		{
+			name: "Peer location matches the location in the deny country only",
+			input: peer.Peer{
+				Meta: peer.PeerSystemMeta{
+					Location: Location{
+						CountryCode: "DE",
+						CityName:    "Berlin",
+					},
+				},
+			},
+			check: GeoLocationCheck{
+				Locations: []Location{
+					{
+						CountryCode: "DE",
+					},
+					{
+						CountryCode: "US",
+					},
+				},
+				Action: GeoLocationActionDeny,
+			},
+			wantErr: false,
+			isValid: false,
+		},
 		{
 			name: "Peer location doesn't match the location in the deny sets",
 			input: peer.Peer{
@@ -120,6 +186,28 @@ func TestGeoLocationCheck_Check(t *testing.T) {
 			wantErr: false,
 			isValid: true,
 		},
+		{
+			name: "Peer location doesn't match the location in the deny country only",
+			input: peer.Peer{
+				Meta: peer.PeerSystemMeta{
+					Location: Location{
+						CountryCode: "DE",
+						CityName:    "Frankfurt am Main",
+					},
+				},
+			},
+			check: GeoLocationCheck{
+				Locations: []Location{
+					{
+						CountryCode: "US",
+						CityName:    "Los Angeles",
+					},
+				},
+				Action: GeoLocationActionDeny,
+			},
+			wantErr: false,
+			isValid: true,
+		},
 	}
 
 	for _, tt := range tests {