diff --git a/proxy/config/config.go b/proxy/config/config.go index 99387632..a0f2c671 100644 --- a/proxy/config/config.go +++ b/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 } diff --git a/proxy/go.mod b/proxy/go.mod index b95db1a5..5286608e 100644 --- a/proxy/go.mod +++ b/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 diff --git a/proxy/go.sum b/proxy/go.sum index a2b9f651..15bce604 100644 --- a/proxy/go.sum +++ b/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= diff --git a/proxy/internal/ratelimit/ratelimit.go b/proxy/internal/ratelimit/ratelimit.go index 4a9a9038..a7a516df 100644 --- a/proxy/internal/ratelimit/ratelimit.go +++ b/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,22 +18,31 @@ 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. @@ -41,10 +50,16 @@ 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 } diff --git a/proxy/internal/ratelimit/ratelimit_bench_test.go b/proxy/internal/ratelimit/ratelimit_bench_test.go index 13fa4606..549e955b 100644 --- a/proxy/internal/ratelimit/ratelimit_bench_test.go +++ b/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") diff --git a/proxy/internal/ratelimit/ratelimit_test.go b/proxy/internal/ratelimit/ratelimit_test.go index 0004118f..e72e6571 100644 --- a/proxy/internal/ratelimit/ratelimit_test.go +++ b/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")) +} diff --git a/proxy/server/ratelimit_response_test.go b/proxy/server/ratelimit_response_test.go index 94307937..dc058b93 100644 --- a/proxy/server/ratelimit_response_test.go +++ b/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) diff --git a/proxy/server/server.go b/proxy/server/server.go index 78544985..8bc34d36 100644 --- a/proxy/server/server.go +++ b/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.