From 330a03ce754a417b12d819fa11b15f7786ed4c25 Mon Sep 17 00:00:00 2001 From: mlsmaycon Date: Tue, 23 Jun 2026 14:32:22 +0200 Subject: [PATCH] Add unit tests for posture checks and peer location resolution logic --- management/server/peer_test.go | 139 ++++++++++++++++++++++ management/server/posture/affects_test.go | 129 ++++++++++++++++++++ management/server/posture/checks.go | 2 +- 3 files changed, 269 insertions(+), 1 deletion(-) create mode 100644 management/server/posture/affects_test.go diff --git a/management/server/peer_test.go b/management/server/peer_test.go index 98cf10acf..6f139e43f 100644 --- a/management/server/peer_test.go +++ b/management/server/peer_test.go @@ -49,6 +49,7 @@ import ( nbdns "github.com/netbirdio/netbird/dns" "github.com/netbirdio/netbird/management/server/activity" + "github.com/netbirdio/netbird/management/server/geolocation" nbpeer "github.com/netbirdio/netbird/management/server/peer" "github.com/netbirdio/netbird/management/server/posture" "github.com/netbirdio/netbird/management/server/store" @@ -2893,3 +2894,141 @@ func TestUpdatePeer_DnsLabelUniqueName(t *testing.T) { require.NoError(t, err, "renaming to unique FQDN should succeed") assert.Equal(t, "api-server", updated.DNSLabel, "DNS label should be first label of FQDN") } + +// fakeGeo is a configurable geolocation.Geolocation implementation for tests. It +// returns a record built from the configured city geoname id, or an error when set. +type fakeGeo struct { + geoNameID uint + isoCode string + cityName string + err error +} + +func (g *fakeGeo) Lookup(net.IP) (*geolocation.Record, error) { + if g.err != nil { + return nil, g.err + } + record := &geolocation.Record{} + record.City.GeonameID = g.geoNameID + record.City.Names.En = g.cityName + record.Country.ISOCode = g.isoCode + return record, nil +} + +func (g *fakeGeo) GetAllCountries() ([]geolocation.Country, error) { return nil, nil } + +func (g *fakeGeo) GetCitiesByCountry(string) ([]geolocation.City, error) { return nil, nil } + +func (g *fakeGeo) Stop() error { return nil } + +func TestResolvePeerLocation(t *testing.T) { + realIP := net.ParseIP("203.0.113.10") + + tests := []struct { + name string + geo geolocation.Geolocation + peer *nbpeer.Peer + realIP net.IP + want *nbpeer.Location + wantNil bool + }{ + { + name: "no geo configured returns nil", + geo: nil, + peer: &nbpeer.Peer{ID: "p1"}, + realIP: realIP, + wantNil: true, + }, + { + name: "nil real IP returns nil", + geo: &fakeGeo{geoNameID: 100}, + peer: &nbpeer.Peer{ID: "p1"}, + realIP: nil, + wantNil: true, + }, + { + name: "lookup error returns nil", + geo: &fakeGeo{err: fmt.Errorf("lookup boom")}, + peer: &nbpeer.Peer{ID: "p1"}, + realIP: realIP, + wantNil: true, + }, + { + name: "same IP and same geoname returns nil", + geo: &fakeGeo{geoNameID: 100, isoCode: "US", cityName: "City A"}, + peer: &nbpeer.Peer{ + ID: "p1", + Location: nbpeer.Location{ + ConnectionIP: realIP, + GeoNameID: 100, + }, + }, + realIP: realIP, + wantNil: true, + }, + { + name: "same IP but changed geoname returns location", + geo: &fakeGeo{geoNameID: 200, isoCode: "US", cityName: "City B"}, + peer: &nbpeer.Peer{ + ID: "p1", + Location: nbpeer.Location{ + ConnectionIP: realIP, + GeoNameID: 100, + }, + }, + realIP: realIP, + want: &nbpeer.Location{ + ConnectionIP: realIP, + CountryCode: "US", + CityName: "City B", + GeoNameID: 200, + }, + }, + { + name: "different IP returns location", + geo: &fakeGeo{geoNameID: 100, isoCode: "US", cityName: "City A"}, + peer: &nbpeer.Peer{ + ID: "p1", + Location: nbpeer.Location{ + ConnectionIP: net.ParseIP("198.51.100.7"), + GeoNameID: 100, + }, + }, + realIP: realIP, + want: &nbpeer.Location{ + ConnectionIP: realIP, + CountryCode: "US", + CityName: "City A", + GeoNameID: 100, + }, + }, + { + name: "no prior location returns location", + geo: &fakeGeo{geoNameID: 100, isoCode: "US", cityName: "City A"}, + peer: &nbpeer.Peer{ID: "p1"}, + realIP: realIP, + want: &nbpeer.Location{ + ConnectionIP: realIP, + CountryCode: "US", + CityName: "City A", + GeoNameID: 100, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + am := &DefaultAccountManager{geo: tt.geo} + got := am.resolvePeerLocation(context.Background(), tt.peer, tt.realIP) + if tt.wantNil { + assert.Nil(t, got, "resolved location should be nil") + return + } + require.NotNil(t, got, "resolved location should not be nil") + assert.True(t, tt.want.ConnectionIP.Equal(got.ConnectionIP), "connection IP should match") + assert.Equal(t, tt.want.CountryCode, got.CountryCode, "country code should match") + assert.Equal(t, tt.want.CityName, got.CityName, "city name should match") + assert.Equal(t, tt.want.GeoNameID, got.GeoNameID, "geoname id should match") + }) + } +} diff --git a/management/server/posture/affects_test.go b/management/server/posture/affects_test.go new file mode 100644 index 000000000..fbf1bf0d0 --- /dev/null +++ b/management/server/posture/affects_test.go @@ -0,0 +1,129 @@ +package posture + +import ( + "net/netip" + "testing" + + "github.com/stretchr/testify/assert" + + nbpeer "github.com/netbirdio/netbird/management/server/peer" +) + +func TestAffectsPosture(t *testing.T) { + processCheck := &Checks{Checks: ChecksDefinition{ProcessCheck: &ProcessCheck{}}} + osCheck := &Checks{Checks: ChecksDefinition{OSVersionCheck: &OSVersionCheck{}}} + nbCheck := &Checks{Checks: ChecksDefinition{NBVersionCheck: &NBVersionCheck{}}} + geoCheck := &Checks{Checks: ChecksDefinition{GeoLocationCheck: &GeoLocationCheck{}}} + + privateRangeCheck := &Checks{Checks: ChecksDefinition{ + PeerNetworkRangeCheck: &PeerNetworkRangeCheck{ + Ranges: []netip.Prefix{netip.MustParsePrefix("10.0.0.0/8")}, + }, + }} + publicRangeCheck := &Checks{Checks: ChecksDefinition{ + PeerNetworkRangeCheck: &PeerNetworkRangeCheck{ + Ranges: []netip.Prefix{netip.MustParsePrefix("203.0.113.0/24")}, + }, + }} + mixedRangeCheck := &Checks{Checks: ChecksDefinition{ + PeerNetworkRangeCheck: &PeerNetworkRangeCheck{ + Ranges: []netip.Prefix{ + netip.MustParsePrefix("203.0.113.0/24"), + netip.MustParsePrefix("192.168.0.0/16"), + }, + }, + }} + + tests := []struct { + name string + diff *nbpeer.MetaDiff + checks []*Checks + want bool + }{ + { + name: "nil diff never affects posture", + diff: nil, + checks: []*Checks{processCheck}, + want: false, + }, + { + name: "process check affected by files change", + diff: &nbpeer.MetaDiff{Files: true}, + checks: []*Checks{processCheck}, + want: true, + }, + { + name: "process check ignores unrelated change", + diff: &nbpeer.MetaDiff{Hostname: true}, + checks: []*Checks{processCheck}, + want: false, + }, + { + name: "os check affected by os version change", + diff: &nbpeer.MetaDiff{OSVersion: true}, + checks: []*Checks{osCheck}, + want: true, + }, + { + name: "nb check affected by wt version change", + diff: &nbpeer.MetaDiff{WtVersion: true}, + checks: []*Checks{nbCheck}, + want: true, + }, + { + name: "geo check affected by location change", + diff: &nbpeer.MetaDiff{LocationChanged: true}, + checks: []*Checks{geoCheck}, + want: true, + }, + { + name: "network range check not affected without network address or location change", + diff: &nbpeer.MetaDiff{Hostname: true}, + checks: []*Checks{privateRangeCheck}, + want: false, + }, + { + name: "private range check affected by network address change", + diff: &nbpeer.MetaDiff{NetworkAddresses: true}, + checks: []*Checks{privateRangeCheck}, + want: true, + }, + { + name: "public range check not affected by network address change alone", + diff: &nbpeer.MetaDiff{NetworkAddresses: true}, + checks: []*Checks{publicRangeCheck}, + want: false, + }, + { + name: "public range check affected by location change alone", + diff: &nbpeer.MetaDiff{LocationChanged: true}, + checks: []*Checks{publicRangeCheck}, + want: true, + }, + { + name: "private range check affected by location change alone", + diff: &nbpeer.MetaDiff{LocationChanged: true}, + checks: []*Checks{privateRangeCheck}, + want: true, + }, + { + name: "public range check affected when location also changed", + diff: &nbpeer.MetaDiff{NetworkAddresses: true, LocationChanged: true}, + checks: []*Checks{publicRangeCheck}, + want: true, + }, + { + name: "mixed ranges affected by network address change due to private range", + diff: &nbpeer.MetaDiff{NetworkAddresses: true}, + checks: []*Checks{mixedRangeCheck}, + want: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := AffectsPosture(tt.diff, tt.checks) + assert.Equal(t, tt.want, got, "AffectsPosture result should match expectation") + }) + } +} diff --git a/management/server/posture/checks.go b/management/server/posture/checks.go index f7e6991d4..b1d49bf76 100644 --- a/management/server/posture/checks.go +++ b/management/server/posture/checks.go @@ -73,7 +73,7 @@ func AffectsPosture(diff *nbpeer.MetaDiff, checks []*Checks) bool { if c.Checks.GeoLocationCheck != nil && diff.LocationChanged { return true } - if c.Checks.PeerNetworkRangeCheck != nil && diff.NetworkAddresses { + if c.Checks.PeerNetworkRangeCheck != nil && (diff.NetworkAddresses || diff.LocationChanged) { if diff.LocationChanged { return true }