Resolve the country once per check in allow_match any mode

proxy/internal/restrict/restrict.go

@@ -261,70 +261,61 @@ func (f *Filter) Check(addr netip.Addr, geo GeoResolver) Verdict {
 // Blocklists remain a hard-deny gate evaluated first and are independent of the
 // allow-combine mode, so a blocklist match (or unverifiable country block) still
 // denies. CrowdSec runs last, as in the default path.
+//
+// The country is resolved at most once and shared by both the blocklist and the
+// allowlist, matching what the all-mode path does. Splitting the two checks into
+// separate helpers cost a second geo lookup per connection whenever both country
+// lists were configured.
 func (f *Filter) checkAny(addr netip.Addr, geo GeoResolver) Verdict {
-	if v := f.checkBlocked(addr, geo); v != Allow {
-		return v
-	}
-	if v := f.checkAllowedAny(addr, geo); v != Allow {
-		return v
-	}
-	return f.checkCrowdSec(addr)
-}
-
-// checkBlocked is the hard-deny gate: it denies on any blocklist match,
-// regardless of the allow-combine mode. A configured country blocklist with an
-// unavailable geo lookup fails closed.
-func (f *Filter) checkBlocked(addr netip.Addr, geo GeoResolver) Verdict {
 	for _, prefix := range f.BlockedCIDRs {
 		if prefix.Contains(addr) {
 			return DenyCIDR
 		}
 	}
 
-	if len(f.BlockedCountries) == 0 {
-		return Allow
-	}
-	if geo == nil || !geo.Available() {
-		return DenyGeoUnavailable
-	}
-	if code := geo.LookupAddr(addr).CountryCode; code != "" && slices.Contains(f.BlockedCountries, code) {
-		return DenyCountry
-	}
-	return Allow
-}
-
-// checkAllowedAny admits the address if it matches any active allowlist. The
-// CIDR allowlist is evaluated first so a match admits without a geo lookup;
-// only when it does not match is the country allowlist consulted, where an
-// unavailable geo lookup fails closed.
-func (f *Filter) checkAllowedAny(addr netip.Addr, geo GeoResolver) Verdict {
 	cidrActive := len(f.AllowedCIDRs) > 0
-	countryActive := len(f.AllowedCountries) > 0
+	cidrAllowed := false
-	if !cidrActive && !countryActive {
-		return Allow
-	}
-
 	if cidrActive {
 		for _, prefix := range f.AllowedCIDRs {
 			if prefix.Contains(addr) {
-				return Allow
+				cidrAllowed = true
+				break
 			}
 		}
 	}
 
-	if countryActive {
+	countryActive := len(f.AllowedCountries) > 0
+	// The blocklist is a hard gate, so it needs the country even when a CIDR
+	// allowlist already admitted the address. The allowlist needs it only when
+	// the CIDR list did not admit it, which is why a matching allowed CIDR
+	// still skips the lookup when no country blocklist is configured.
+	needCountry := len(f.BlockedCountries) > 0 || (countryActive && !cidrAllowed)
+
+	country := ""
+	if needCountry {
 		if geo == nil || !geo.Available() {
 			return DenyGeoUnavailable
 		}
-		if code := geo.LookupAddr(addr).CountryCode; code != "" && slices.Contains(f.AllowedCountries, code) {
+		country = geo.LookupAddr(addr).CountryCode
-			return Allow
-		}
 	}
 
-	if cidrActive {
+	if country != "" && slices.Contains(f.BlockedCountries, country) {
-		return DenyCIDR
+		return DenyCountry
 	}
-	return DenyCountry
+
+	allowed := (!cidrActive && !countryActive) ||
+		cidrAllowed ||
+		(countryActive && country != "" && slices.Contains(f.AllowedCountries, country))
+	if !allowed {
+		// Both allowlists missing is reported against the CIDR list, the one
+		// checked first, so the reason stays stable for existing access logs.
+		if cidrActive {
+			return DenyCIDR
+		}
+		return DenyCountry
+	}
+
+	return f.checkCrowdSec(addr)
 }
 
 func (f *Filter) checkCIDR(addr netip.Addr) Verdict {

proxy/internal/restrict/restrict_test.go

@@ -5,6 +5,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 
 	"github.com/netbirdio/netbird/proxy/internal/geolocation"
 )
@@ -733,3 +734,81 @@ func TestFilter_HasRestrictions_CrowdSec(t *testing.T) {
 	f2 := ParseFilter(FilterConfig{CrowdSec: nil, CrowdSecMode: CrowdSecEnforce})
 	assert.True(t, f2.HasRestrictions())
 }
+
+// countingGeo records how many times an address was resolved.
+type countingGeo struct {
+	countries map[string]string
+	lookups   int
+}
+
+func (c *countingGeo) LookupAddr(addr netip.Addr) geolocation.Result {
+	c.lookups++
+	return geolocation.Result{CountryCode: c.countries[addr.String()]}
+}
+
+func (c *countingGeo) Available() bool { return true }
+
+// The geo lookup is the expensive part of the check and runs per connection, so
+// "any" mode must resolve the country once and share it between the blocklist
+// and the allowlist, the way "all" mode does.
+func TestCheck_AnyResolvesCountryOnce(t *testing.T) {
+	tests := []struct {
+		name        string
+		ip          string
+		want        Verdict
+		wantLookups int
+	}{
+		{"blocked and allowed lists both active", "203.0.113.1", Allow, 1},
+		{"blocked country denies", "198.51.100.1", DenyCountry, 1},
+		{"neither allowlist matches", "192.0.2.1", DenyCIDR, 1},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			geo := &countingGeo{countries: map[string]string{
+				"203.0.113.1":  "US",
+				"198.51.100.1": "CN",
+				"192.0.2.1":    "FR",
+			}}
+			f := ParseFilter(FilterConfig{
+				AllowedCIDRs:     []string{"10.0.0.0/8"},
+				AllowedCountries: []string{"US"},
+				BlockedCountries: []string{"CN"},
+				AllowMatch:       AllowMatchAny,
+			})
+			require.NotNil(t, f)
+
+			assert.Equal(t, tt.want, f.Check(netip.MustParseAddr(tt.ip), geo))
+			assert.Equal(t, tt.wantLookups, geo.lookups, "the country must be resolved at most once per check")
+		})
+	}
+}
+
+// A matching allowed CIDR short-circuits the allowlist, so with no country
+// blocklist configured there is nothing left to resolve.
+func TestCheck_AnySkipsLookupWhenCIDRAdmits(t *testing.T) {
+	geo := &countingGeo{countries: map[string]string{"10.1.2.3": "US"}}
+	f := ParseFilter(FilterConfig{
+		AllowedCIDRs:     []string{"10.0.0.0/8"},
+		AllowedCountries: []string{"DE"},
+		AllowMatch:       AllowMatchAny,
+	})
+	require.NotNil(t, f)
+
+	assert.Equal(t, Allow, f.Check(netip.MustParseAddr("10.1.2.3"), geo))
+	assert.Zero(t, geo.lookups, "an admitted CIDR needs no geo lookup")
+}
+
+// The blocklist is a hard gate, so it is consulted even when a CIDR allowlist
+// already admitted the address.
+func TestCheck_AnyBlocklistOutranksAllowedCIDR(t *testing.T) {
+	geo := &countingGeo{countries: map[string]string{"10.1.2.3": "CN"}}
+	f := ParseFilter(FilterConfig{
+		AllowedCIDRs:     []string{"10.0.0.0/8"},
+		BlockedCountries: []string{"CN"},
+		AllowMatch:       AllowMatchAny,
+	})
+	require.NotNil(t, f)
+
+	assert.Equal(t, DenyCountry, f.Check(netip.MustParseAddr("10.1.2.3"), geo))
+	assert.Equal(t, 1, geo.lookups)
+}