fix(proxy): cap rate-limit bucket stores, bucket only existing profiles, key IPv6 by /64

proxy/config/config.go

@@ -71,6 +71,8 @@ type RateLimitConfig struct {
 	PerProfileRate     int
 	PerProfileBurst    int
 	PerProfileResponse string // "drop" or "refuse" (default)
+	MaxBuckets         int    // cap per bucket store
+	IPv6PrefixLen      int    // prefix length IPv6 clients are grouped by
 }
 
 // MetricsConfig holds Prometheus metrics server settings.
@@ -451,6 +453,8 @@ func loadRateLimitConfig() *RateLimitConfig {
 		PerProfileRate:     600,
 		PerProfileBurst:    1000,
 		PerProfileResponse: RateLimitResponseRefuse,
+		MaxBuckets:         100_000,
+		IPv6PrefixLen:      64,
 	}
 	if v := strings.ToLower(strings.TrimSpace(os.Getenv("RATELIMIT_PER_IP_RESPONSE"))); v == RateLimitResponseDrop || v == RateLimitResponseRefuse {
 		cfg.PerIPResponse = v
@@ -470,6 +474,12 @@ func loadRateLimitConfig() *RateLimitConfig {
 	if v, err := strconv.Atoi(os.Getenv("RATELIMIT_PER_PROFILE_BURST")); err == nil && v > 0 {
 		cfg.PerProfileBurst = v
 	}
+	if v, err := strconv.Atoi(os.Getenv("RATELIMIT_MAX_BUCKETS")); err == nil && v > 0 {
+		cfg.MaxBuckets = v
+	}
+	if v, err := strconv.Atoi(os.Getenv("RATELIMIT_IPV6_PREFIX")); err == nil && v >= 1 && v <= 128 {
+		cfg.IPv6PrefixLen = v
+	}
 	return cfg
 }
 

proxy/go.mod

@@ -7,6 +7,7 @@ require (
 	github.com/AdguardTeam/golibs v0.35.2
 	github.com/Shopify/toxiproxy/v2 v2.12.0
 	github.com/getsentry/sentry-go/zerolog v0.31.1
+	github.com/hashicorp/golang-lru/v2 v2.0.7
 	github.com/ivpn/dns/libs v0.0.0
 	github.com/miekg/dns v1.1.68
 	github.com/oschwald/geoip2-golang v1.13.0

proxy/go.sum

@@ -95,6 +95,8 @@ github.com/hashicorp/errwrap v1.1.0 h1:OxrOeh75EUXMY8TBjag2fzXGZ40LB6IKw45YeGUDY
 github.com/hashicorp/errwrap v1.1.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4=
 github.com/hashicorp/go-multierror v1.1.1 h1:H5DkEtf6CXdFp0N0Em5UCwQpXMWke8IA0+lD48awMYo=
 github.com/hashicorp/go-multierror v1.1.1/go.mod h1:iw975J/qwKPdAO1clOe2L8331t/9/fmwbPZ6JB6eMoM=
+github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k=
+github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
 github.com/klauspost/compress v1.18.5 h1:/h1gH5Ce+VWNLSWqPzOVn6XBO+vJbCNGvjoaGBFW2IE=
 github.com/klauspost/compress v1.18.5/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ=
 github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=

proxy/internal/ratelimit/ratelimit.go

@@ -4,7 +4,7 @@ import (
 	"net/netip"
 	"time"
 
-	gocache "github.com/patrickmn/go-cache"
+	"github.com/hashicorp/golang-lru/v2/expirable"
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
 	"golang.org/x/time/rate"
@@ -18,33 +18,48 @@ type Config struct {
 	PerProfileEnabled bool
 	PerProfileRate    int
 	PerProfileBurst   int
+	// MaxBuckets caps each bucket store; 0 or negative selects defaultMaxBuckets.
+	MaxBuckets int
+	// IPv6PrefixLen is the prefix length IPv6 client addresses are grouped by;
+	// values outside 1..128 select defaultIPv6PrefixLen.
+	IPv6PrefixLen int
 }
 
 // RateLimiter enforces per-IP and per-profile rate limits using token buckets.
 type RateLimiter struct {
 	cfg            Config
-	ipBuckets      *gocache.Cache
-	profileBuckets *gocache.Cache
+	ipBuckets      *expirable.LRU[string, *rate.Limiter]
+	profileBuckets *expirable.LRU[string, *rate.Limiter]
 	metrics        Metrics
 	sampledLogger  zerolog.Logger
 }
 
 const (
-	bucketExpiry  = 1 * time.Hour
-	bucketCleanup = 5 * time.Minute
-	layerIP       = "ip"
-	layerProfile  = "profile"
+	bucketExpiry = 1 * time.Hour
+	layerIP      = "ip"
+	layerProfile = "profile"
+
+	defaultMaxBuckets = 100_000
+	// RFC 4291 §2.5.4: /64 is the standard interface-identifier boundary, so
+	// one subscriber allocation maps to one bucket.
+	defaultIPv6PrefixLen = 64
 )
 
 // New creates a RateLimiter. Pass nil for m to disable metrics recording.
 func New(cfg Config, m Metrics) *RateLimiter {
 	if m == nil {
 		m = noopMetrics{}
 	}
+	if cfg.MaxBuckets <= 0 {
+		cfg.MaxBuckets = defaultMaxBuckets
+	}
+	if cfg.IPv6PrefixLen < 1 || cfg.IPv6PrefixLen > 128 {
+		cfg.IPv6PrefixLen = defaultIPv6PrefixLen
+	}
 	return &RateLimiter{
 		cfg:            cfg,
-		ipBuckets:      gocache.New(bucketExpiry, bucketCleanup),
-		profileBuckets: gocache.New(bucketExpiry, bucketCleanup),
+		ipBuckets:      expirable.NewLRU[string, *rate.Limiter](cfg.MaxBuckets, nil, bucketExpiry),
+		profileBuckets: expirable.NewLRU[string, *rate.Limiter](cfg.MaxBuckets, nil, bucketExpiry),
 		metrics:        m,
 		sampledLogger: log.Logger.Sample(&zerolog.BurstSampler{
 			Burst:       5,
@@ -59,7 +74,22 @@ func (rl *RateLimiter) CheckIP(addr netip.Addr, proto string) bool {
 	if !rl.cfg.PerIPEnabled {
 		return true
 	}
-	return rl.check(rl.ipBuckets, addr.String(), rl.cfg.PerIPRate, rl.cfg.PerIPBurst, layerIP, proto)
+	return rl.check(rl.ipBuckets, rl.ipKey(addr), rl.cfg.PerIPRate, rl.cfg.PerIPBurst, layerIP, proto)
+}
+
+// ipKey groups IPv6 clients by prefix: a subscriber typically holds an entire
+// prefix delegation, so keying on the full address would give one client a
+// bucket per address. 4-mapped-6 addresses are unmapped so they key as IPv4.
+func (rl *RateLimiter) ipKey(addr netip.Addr) string {
+	addr = addr.Unmap()
+	if !addr.Is6() {
+		return addr.String()
+	}
+	prefix, err := addr.Prefix(rl.cfg.IPv6PrefixLen)
+	if err != nil {
+		return addr.String()
+	}
+	return prefix.Addr().String()
 }
 
 // CheckProfile returns true if the query for profileID should be allowed (Layer 2).
@@ -70,10 +100,9 @@ func (rl *RateLimiter) CheckProfile(profileID string, proto string) bool {
 	return rl.check(rl.profileBuckets, profileID, rl.cfg.PerProfileRate, rl.cfg.PerProfileBurst, layerProfile, proto)
 }
 
-func (rl *RateLimiter) check(store *gocache.Cache, key string, rps, burst int, layer, proto string) bool {
-	v, found := store.Get(key)
+func (rl *RateLimiter) check(store *expirable.LRU[string, *rate.Limiter], key string, rps, burst int, layer, proto string) bool {
+	limiter, found := store.Get(key)
 	if found {
-		limiter := v.(*rate.Limiter)
 		if limiter.Allow() {
 			return true
 		}
@@ -85,8 +114,8 @@ func (rl *RateLimiter) check(store *gocache.Cache, key string, rps, burst int, l
 		return false
 	}
 
-	limiter := rate.NewLimiter(rate.Limit(rps), burst)
+	limiter = rate.NewLimiter(rate.Limit(rps), burst)
 	limiter.Allow() // consume first token
-	store.Set(key, limiter, gocache.DefaultExpiration)
+	store.Add(key, limiter)
 	return true
 }

proxy/internal/ratelimit/ratelimit_bench_test.go

@@ -49,6 +49,16 @@ func BenchmarkCheckIP_ManyIPs(b *testing.B) {
 	}
 }
 
+// BenchmarkCheckProfile_Churn exercises the store-full eviction path with a
+// distinct key per iteration.
+func BenchmarkCheckProfile_Churn(b *testing.B) {
+	rl := benchLimiter()
+	b.ResetTimer()
+	for i := range b.N {
+		rl.CheckProfile(fmt.Sprintf("profile-%d", i), "tls")
+	}
+}
+
 func BenchmarkCheckIP_Parallel(b *testing.B) {
 	rl := benchLimiter()
 	addr := netip.MustParseAddr("192.0.2.1")

proxy/internal/ratelimit/ratelimit_test.go

@@ -175,3 +175,57 @@ func TestCheckIP_ManyIPs(t *testing.T) {
 		assert.True(t, rl.CheckIP(addr, "udp"))
 	}
 }
+
+func TestMaxBucketsEvictsOldest(t *testing.T) {
+	rl, _ := newTestLimiter(Config{PerProfileEnabled: true, PerProfileRate: 1, PerProfileBurst: 1, MaxBuckets: 3})
+
+	// prof1's single token is consumed; a second call would be rejected.
+	assert.True(t, rl.CheckProfile("prof1", "udp"))
+
+	// Fill the store past its cap; prof1 is the least recently used entry.
+	assert.True(t, rl.CheckProfile("prof2", "udp"))
+	assert.True(t, rl.CheckProfile("prof3", "udp"))
+	assert.True(t, rl.CheckProfile("prof4", "udp"))
+
+	// prof1 was evicted, so it gets a fresh bucket and passes again.
+	assert.True(t, rl.CheckProfile("prof1", "udp"))
+
+	// prof3 and prof4 are still tracked: their tokens are spent.
+	assert.False(t, rl.CheckProfile("prof3", "udp"))
+	assert.False(t, rl.CheckProfile("prof4", "udp"))
+}
+
+func TestCheckIP_IPv6SamePrefixSharesBucket(t *testing.T) {
+	rl, m := newTestLimiter(Config{PerIPEnabled: true, PerIPRate: 3, PerIPBurst: 3})
+
+	// Three addresses within the same /64 draw from one bucket.
+	assert.True(t, rl.CheckIP(netip.MustParseAddr("2001:db8:1:1::1"), "udp"))
+	assert.True(t, rl.CheckIP(netip.MustParseAddr("2001:db8:1:1::2"), "udp"))
+	assert.True(t, rl.CheckIP(netip.MustParseAddr("2001:db8:1:1:ffff::3"), "udp"))
+	assert.False(t, rl.CheckIP(netip.MustParseAddr("2001:db8:1:1::4"), "udp"))
+	assert.Equal(t, 1, m.count("ip", "udp"))
+
+	// A different /64 is unaffected.
+	assert.True(t, rl.CheckIP(netip.MustParseAddr("2001:db8:1:2::1"), "udp"))
+}
+
+func TestCheckIP_IPv6PrefixLenConfigurable(t *testing.T) {
+	rl, _ := newTestLimiter(Config{PerIPEnabled: true, PerIPRate: 1, PerIPBurst: 1, IPv6PrefixLen: 56})
+
+	assert.True(t, rl.CheckIP(netip.MustParseAddr("2001:db8:1:100::1"), "udp"))
+	// Same /56, different /64: shares the bucket.
+	assert.False(t, rl.CheckIP(netip.MustParseAddr("2001:db8:1:1ff::1"), "udp"))
+	// Different /56: own bucket.
+	assert.True(t, rl.CheckIP(netip.MustParseAddr("2001:db8:1:200::1"), "udp"))
+}
+
+func TestCheckIP_IPv4MappedTreatedAsIPv4(t *testing.T) {
+	rl, _ := newTestLimiter(Config{PerIPEnabled: true, PerIPRate: 1, PerIPBurst: 1})
+
+	// A 4-mapped-6 address shares its bucket with the plain IPv4 form.
+	assert.True(t, rl.CheckIP(netip.MustParseAddr("::ffff:192.0.2.1"), "udp"))
+	assert.False(t, rl.CheckIP(netip.MustParseAddr("192.0.2.1"), "udp"))
+
+	// Other IPv4 addresses keep independent buckets.
+	assert.True(t, rl.CheckIP(netip.MustParseAddr("192.0.2.2"), "udp"))
+}

proxy/server/ratelimit_response_test.go

@@ -2,14 +2,22 @@ package server
 
 import (
 	"errors"
+	"net/http"
 	"net/netip"
+	"net/url"
 	"testing"
+	"time"
 
 	"github.com/AdguardTeam/dnsproxy/proxy"
+	"github.com/ivpn/dns/libs/logging"
 	"github.com/ivpn/dns/proxy/config"
 	"github.com/ivpn/dns/proxy/internal/ratelimit"
+	"github.com/ivpn/dns/proxy/mocks"
+	"github.com/ivpn/dns/proxy/model"
 	"github.com/miekg/dns"
+	gocache "github.com/patrickmn/go-cache"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 )
 
@@ -93,6 +101,89 @@ func TestHandleBefore_IPRateLimit_Refuse(t *testing.T) {
 	assert.ErrorIs(t, err, errRateLimitedIP)
 }
 
+// newProfileRateLimitServer builds a Server that reaches the per-profile
+// rate-limit layer: per-IP limiting is disabled and the profile limiter is
+// rate=1, burst=1 so the second call for the same profile is rejected.
+func newProfileRateLimitServer(c *mocks.Cache, mem *mocks.MemoryCache) *Server {
+	return &Server{
+		Config: &config.Config{
+			Server:   &config.ServerConfig{},
+			Upstream: &config.UpstreamConfig{Default: "default"},
+			RateLimit: &config.RateLimitConfig{
+				PerProfileEnabled:  true,
+				PerProfileRate:     1,
+				PerProfileBurst:    1,
+				PerProfileResponse: config.RateLimitResponseRefuse,
+			},
+		},
+		Cache:                c,
+		InMemoryCache:        mem,
+		ProfileSettingsCache: gocache.New(time.Minute, time.Minute),
+		LoggerFactory:        logging.NewDefaultFactory(),
+		RateLimiter: ratelimit.New(ratelimit.Config{
+			PerProfileEnabled: true,
+			PerProfileRate:    1,
+			PerProfileBurst:   1,
+		}, nil),
+		Metrics: noopMetrics{},
+	}
+}
+
+// newDoHDNSContext carries profileID via the DoH path, the simplest route
+// through clientIDFromDNSContext in tests.
+func newDoHDNSContext(profileID string) *proxy.DNSContext {
+	req := new(dns.Msg)
+	req.SetQuestion(dns.Fqdn("example.com"), dns.TypeA)
+	return &proxy.DNSContext{
+		Req:         req,
+		Addr:        netip.MustParseAddrPort("192.0.2.1:443"),
+		Proto:       proxy.ProtoHTTPS,
+		HTTPRequest: &http.Request{URL: &url.URL{Path: "/dns-query/" + profileID}},
+	}
+}
+
+func TestHandleBefore_UnknownProfileNeverProfileRateLimited(t *testing.T) {
+	c := mocks.NewCache(t)
+	c.EXPECT().GetProfileSettingsBatch(mock.Anything, "unknownprofile1").
+		Return(&model.ProfileSettings{PrivacyErr: errors.New("no [privacy] settings found for profile")}, nil)
+	s := newProfileRateLimitServer(c, mocks.NewMemoryCache(t))
+
+	// Far past the burst of 1: every call must fail on existence, and the
+	// profile rate limiter must never fire for a profile that does not exist.
+	for i := range 5 {
+		err := s.HandleBefore(nil, newDoHDNSContext("unknownprofile1"))
+		require.ErrorIs(t, err, errProfileIdNotFound, "call %d", i)
+		require.NotErrorIs(t, err, errRateLimitedProfile, "call %d", i)
+	}
+}
+
+func TestHandleBefore_ValidProfileRateLimit_Refuse(t *testing.T) {
+	mem := mocks.NewMemoryCache(t)
+	mem.EXPECT().SetRequestCtx(mock.Anything, mock.Anything).Return(nil)
+	s := newProfileRateLimitServer(mocks.NewCache(t), mem)
+
+	fetchErr := errors.New("settings unavailable")
+	s.ProfileSettingsCache.Set("validprofile1", &model.ProfileSettings{
+		Privacy:                map[string]string{},
+		LogsErr:                fetchErr,
+		DNSSECErr:              fetchErr,
+		RebindingProtectionErr: fetchErr,
+		AdvancedErr:            fetchErr,
+	}, gocache.DefaultExpiration)
+
+	// First request consumes the single token.
+	require.NoError(t, s.HandleBefore(nil, newDoHDNSContext("validprofile1")))
+
+	// Second request must be refused by the profile layer.
+	err := s.HandleBefore(nil, newDoHDNSContext("validprofile1"))
+	require.ErrorIs(t, err, errRateLimitedProfile)
+
+	var befErr *proxy.BeforeRequestError
+	require.True(t, errors.As(err, &befErr))
+	require.NotNil(t, befErr.Response)
+	assert.Equal(t, dns.RcodeRefused, befErr.Response.Rcode)
+}
+
 func TestRefusedResponse(t *testing.T) {
 	s := &Server{}
 	req := new(dns.Msg)

proxy/server/server.go

@@ -91,6 +91,8 @@ func NewServer(serverConfig *config.Config, collectorChannels map[string]channel
 		PerProfileEnabled: serverConfig.RateLimit.PerProfileEnabled,
 		PerProfileRate:    serverConfig.RateLimit.PerProfileRate,
 		PerProfileBurst:   serverConfig.RateLimit.PerProfileBurst,
+		MaxBuckets:        serverConfig.RateLimit.MaxBuckets,
+		IPv6PrefixLen:     serverConfig.RateLimit.IPv6PrefixLen,
 	}, metrics.NewRateLimitMetrics(prometheus.DefaultRegisterer))
 
 	server := &Server{
@@ -185,17 +187,6 @@ func (s *Server) HandleBefore(p *proxy.Proxy, dctx *proxy.DNSContext) (err error
 		systemLogger.Warn().Err(errProfileIdNotProvided).Msg(errProfileIdNotProvided.Error())
 		return errProfileIdNotProvided
 	} else {
-		// Layer 2: per-profile rate limit (after profile extraction, before Redis).
-		if !s.RateLimiter.CheckProfile(profileId, string(dctx.Proto)) {
-			if s.Config.RateLimit.PerProfileResponse == config.RateLimitResponseRefuse {
-				return &proxy.BeforeRequestError{
-					Err:      errRateLimitedProfile,
-					Response: s.refusedResponse(dctx.Req),
-				}
-			}
-			return errRateLimitedProfile
-		}
-
 		// Try in-memory profile settings cache first.
 		var settings *model.ProfileSettings
 		if cached, ok := s.ProfileSettingsCache.Get(profileId); ok {
@@ -221,6 +212,18 @@ func (s *Server) HandleBefore(p *proxy.Proxy, dctx *proxy.DNSContext) (err error
 			systemLogger.Debug().Err(settings.PrivacyErr).Msg(errProfileIdNotFound.Error())
 			return errProfileIdNotFound
 		}
+
+		// Layer 2: per-profile rate limit. Runs after the existence check so
+		// buckets are only created for profiles that exist.
+		if !s.RateLimiter.CheckProfile(profileId, string(dctx.Proto)) {
+			if s.Config.RateLimit.PerProfileResponse == config.RateLimitResponseRefuse {
+				return &proxy.BeforeRequestError{
+					Err:      errRateLimitedProfile,
+					Response: s.refusedResponse(dctx.Req),
+				}
+			}
+			return errRateLimitedProfile
+		}
 		prvSettings := settings.Privacy
 
 		// Logs settings: default to enabled if unavailable.