Skip to content

Commit 7b2e2e4

Browse files
MizzickMizzick
committed
Pull request: AGDNS-3523-use-ratelimit
Merge in GO/dnsproxy from AGDNS-3523-use-ratelimit to master Squashed commit of the following: commit a0c0d52 Merge: f18c897 68457db Author: Dimitry Kolyshev <dkolyshev@adguard.com> Date: Wed Feb 18 09:36:55 2026 +0700 Merge remote-tracking branch 'origin/master' into AGDNS-3523-use-ratelimit # Conflicts: # proxy/serverquic_internal_test.go commit f18c897 Author: Dimitry Kolyshev <dkolyshev@adguard.com> Date: Wed Feb 18 09:31:02 2026 +0700 proxy: imp code commit 13774cd Author: Dimitry Kolyshev <dkolyshev@adguard.com> Date: Tue Feb 17 09:04:46 2026 +0700 proxy: mv req validation commit 8ab8031 Author: Dimitry Kolyshev <dkolyshev@adguard.com> Date: Tue Feb 17 08:51:55 2026 +0700 cmd: imp code commit bfa1c9e Author: Dimitry Kolyshev <dkolyshev@adguard.com> Date: Mon Feb 16 13:00:42 2026 +0700 ratelimit: imp commit f7ae7e7 Author: Dimitry Kolyshev <dkolyshev@adguard.com> Date: Mon Feb 16 11:40:45 2026 +0700 proxy: imp tests commit 29ed832 Author: Dimitry Kolyshev <dkolyshev@adguard.com> Date: Mon Feb 16 11:30:20 2026 +0700 all: use ratelimit mw
1 parent 68457db commit 7b2e2e4

28 files changed

Lines changed: 322 additions & 683 deletions

‎internal/cmd/config.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -119,15 +119,15 @@ type configuration struct {
119119
CacheSizeBytes int `yaml:"cache-size"`
120120

121121
// Ratelimit is the maximum number of requests per second.
122-
Ratelimit int `yaml:"ratelimit"`
122+
Ratelimit uint `yaml:"ratelimit"`
123123

124124
// RatelimitSubnetLenIPv4 is a subnet length for IPv4 addresses used for
125125
// rate limiting requests.
126-
RatelimitSubnetLenIPv4 int `yaml:"ratelimit-subnet-len-ipv4"`
126+
RatelimitSubnetLenIPv4 uint `yaml:"ratelimit-subnet-len-ipv4"`
127127

128128
// RatelimitSubnetLenIPv6 is a subnet length for IPv6 addresses used for
129129
// rate limiting requests.
130-
RatelimitSubnetLenIPv6 int `yaml:"ratelimit-subnet-len-ipv6"`
130+
RatelimitSubnetLenIPv6 uint `yaml:"ratelimit-subnet-len-ipv6"`
131131

132132
// UDPBufferSize is the size of the UDP buffer in bytes. A value <= 0 will
133133
// use the system default.

‎internal/cmd/proxy.go‎

Lines changed: 28 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import (
1616
"github.com/AdguardTeam/dnsproxy/internal/middleware"
1717
proxynetutil "github.com/AdguardTeam/dnsproxy/internal/netutil"
1818
"github.com/AdguardTeam/dnsproxy/proxy"
19+
"github.com/AdguardTeam/dnsproxy/ratelimit"
1920
"github.com/AdguardTeam/dnsproxy/upstream"
2021
"github.com/AdguardTeam/golibs/errors"
2122
"github.com/AdguardTeam/golibs/logutil/slogutil"
@@ -52,13 +53,13 @@ func createProxyConfig(
5253
HostsFiles: hosts,
5354
})
5455

55-
proxyConf = &proxy.Config{
56-
Logger: l.With(slogutil.KeyPrefix, proxy.LogPrefix),
57-
58-
RatelimitSubnetLenIPv4: conf.RatelimitSubnetLenIPv4,
59-
RatelimitSubnetLenIPv6: conf.RatelimitSubnetLenIPv6,
56+
ratelimitMw, err := conf.newRatelimitMw(l)
57+
if err != nil {
58+
return nil, fmt.Errorf("ratelimit mw: %w", err)
59+
}
6060

61-
Ratelimit: conf.Ratelimit,
61+
proxyConf = &proxy.Config{
62+
Logger: l.With(slogutil.KeyPrefix, proxy.LogPrefix),
6263
CacheEnabled: conf.Cache,
6364
CacheSizeBytes: conf.CacheSizeBytes,
6465
CacheMinTTL: conf.CacheMinTTL,
@@ -81,7 +82,7 @@ func createProxyConfig(
8182
MaxGoroutines: conf.MaxGoRoutines,
8283
UsePrivateRDNS: conf.UsePrivateRDNS,
8384
PrivateSubnets: netutil.SubnetSetFunc(netutil.IsLocallyServed),
84-
RequestHandler: preMw.Wrap(proxy.DefaultHandler{}),
85+
RequestHandler: ratelimitMw.Wrap(preMw.Wrap(proxy.DefaultHandler{})),
8586
PendingRequests: &proxy.PendingRequestsConfig{
8687
Enabled: conf.PendingRequestsEnabled,
8788
},
@@ -120,6 +121,26 @@ func isEmpty(uc *proxy.UpstreamConfig) (ok bool) {
120121
len(uc.SpecifiedDomainUpstreams) == 0
121122
}
122123

124+
// newRatelimitMw returns the ratelimit middleware. In case of invalid
125+
// ratelimit configuration returns an error. l must not be nil.
126+
func (conf *configuration) newRatelimitMw(l *slog.Logger) (mw proxy.Middleware, err error) {
127+
if conf.Ratelimit == 0 {
128+
return proxy.MiddlewareFunc(proxy.PassThrough), nil
129+
}
130+
131+
rlConf := &ratelimit.Config{
132+
Logger: l.With(slogutil.KeyPrefix, "ratelimit"),
133+
Ratelimit: conf.Ratelimit,
134+
SubnetLenIPv4: conf.RatelimitSubnetLenIPv4,
135+
SubnetLenIPv6: conf.RatelimitSubnetLenIPv6,
136+
}
137+
if err = rlConf.Validate(); err != nil {
138+
return nil, fmt.Errorf("invalid configuration: %w", err)
139+
}
140+
141+
return ratelimit.NewMiddleware(rlConf), nil
142+
}
143+
123144
// defaultLocalTimeout is the default timeout for local operations.
124145
const defaultLocalTimeout = 1 * time.Second
125146

‎proxy/beforerequest_internal_test.go‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@ import (
99
"github.com/AdguardTeam/dnsproxy/internal/dnsproxytest"
1010
"github.com/AdguardTeam/dnsproxy/upstream"
1111
"github.com/AdguardTeam/golibs/errors"
12-
"github.com/AdguardTeam/golibs/logutil/slogutil"
1312
"github.com/AdguardTeam/golibs/netutil"
1413
"github.com/AdguardTeam/golibs/testutil/servicetest"
1514
"github.com/miekg/dns"
@@ -53,7 +52,7 @@ func TestProxy_HandleDNSRequest_beforeRequestHandler(t *testing.T) {
5352
errorResponse := (&dns.Msg{}).SetReply(errorRequest)
5453

5554
p := mustNew(t, &Config{
56-
Logger: slogutil.NewDiscardLogger(),
55+
Logger: testLogger,
5756
TCPListenAddr: []*net.TCPAddr{net.TCPAddrFromAddrPort(localhostAnyPort)},
5857
UpstreamConfig: &UpstreamConfig{
5958
Upstreams: []upstream.Upstream{&dnsproxytest.Upstream{

‎proxy/bogusnxdomain_internal_test.go‎

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@ import (
66
"testing"
77

88
"github.com/AdguardTeam/dnsproxy/upstream"
9-
"github.com/AdguardTeam/golibs/logutil/slogutil"
109
"github.com/AdguardTeam/golibs/testutil/servicetest"
1110
"github.com/miekg/dns"
1211
"github.com/stretchr/testify/assert"
@@ -15,14 +14,12 @@ import (
1514

1615
func TestProxy_IsBogusNXDomain(t *testing.T) {
1716
prx := mustNew(t, &Config{
18-
Logger: slogutil.NewDiscardLogger(),
19-
UDPListenAddr: []*net.UDPAddr{net.UDPAddrFromAddrPort(localhostAnyPort)},
20-
TCPListenAddr: []*net.TCPAddr{net.TCPAddrFromAddrPort(localhostAnyPort)},
21-
UpstreamConfig: newTestUpstreamConfig(t, defaultTimeout, testDefaultUpstreamAddr),
22-
TrustedProxies: defaultTrustedProxies,
23-
RatelimitSubnetLenIPv4: 24,
24-
RatelimitSubnetLenIPv6: 64,
25-
CacheEnabled: true,
17+
Logger: testLogger,
18+
UDPListenAddr: []*net.UDPAddr{net.UDPAddrFromAddrPort(localhostAnyPort)},
19+
TCPListenAddr: []*net.TCPAddr{net.TCPAddrFromAddrPort(localhostAnyPort)},
20+
UpstreamConfig: newTestUpstreamConfig(t, defaultTimeout, testDefaultUpstreamAddr),
21+
TrustedProxies: defaultTrustedProxies,
22+
CacheEnabled: true,
2623
BogusNXDomain: []netip.Prefix{
2724
netip.MustParsePrefix("4.3.2.1/24"),
2825
netip.MustParsePrefix("1.2.3.4/8"),

‎proxy/cache_internal_test.go‎

Lines changed: 31 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@ import (
1111

1212
"github.com/AdguardTeam/dnsproxy/internal/dnsproxytest"
1313
"github.com/AdguardTeam/dnsproxy/upstream"
14-
"github.com/AdguardTeam/golibs/logutil/slogutil"
1514
"github.com/AdguardTeam/golibs/netutil"
1615
"github.com/AdguardTeam/golibs/testutil"
1716
"github.com/AdguardTeam/golibs/testutil/servicetest"
@@ -53,14 +52,12 @@ func newTestCache(tb testing.TB, conf *cacheConfig) (c *cache) {
5352

5453
func TestServeCached(t *testing.T) {
5554
dnsProxy := mustNew(t, &Config{
56-
Logger: slogutil.NewDiscardLogger(),
57-
UDPListenAddr: []*net.UDPAddr{net.UDPAddrFromAddrPort(localhostAnyPort)},
58-
TCPListenAddr: []*net.TCPAddr{net.TCPAddrFromAddrPort(localhostAnyPort)},
59-
UpstreamConfig: newTestUpstreamConfig(t, defaultTimeout, testDefaultUpstreamAddr),
60-
TrustedProxies: defaultTrustedProxies,
61-
RatelimitSubnetLenIPv4: 24,
62-
RatelimitSubnetLenIPv6: 64,
63-
CacheEnabled: true,
55+
Logger: testLogger,
56+
UDPListenAddr: []*net.UDPAddr{net.UDPAddrFromAddrPort(localhostAnyPort)},
57+
TCPListenAddr: []*net.TCPAddr{net.TCPAddrFromAddrPort(localhostAnyPort)},
58+
UpstreamConfig: newTestUpstreamConfig(t, defaultTimeout, testDefaultUpstreamAddr),
59+
TrustedProxies: defaultTrustedProxies,
60+
CacheEnabled: true,
6461
})
6562

6663
// Start listening.
@@ -75,7 +72,7 @@ func TestServeCached(t *testing.T) {
7572
}).SetQuestion("google.com.", dns.TypeA)
7673
reply.SetEdns0(defaultUDPBufSize, false)
7774

78-
dnsProxy.cache.set(reply, upstreamWithAddr, slogutil.NewDiscardLogger())
75+
dnsProxy.cache.set(reply, upstreamWithAddr, testLogger)
7976

8077
// Create a DNS-over-UDP client connection.
8178
addr := dnsProxy.Addr(ProtoUDP)
@@ -199,7 +196,7 @@ func TestCacheDO(t *testing.T) {
199196
reply.SetEdns0(4096, true)
200197

201198
// Store in cache.
202-
testCache.set(reply, upstreamWithAddr, slogutil.NewDiscardLogger())
199+
testCache.set(reply, upstreamWithAddr, testLogger)
203200

204201
// Make a request.
205202
request := (&dns.Msg{}).SetQuestion("google.com.", dns.TypeA)
@@ -230,8 +227,6 @@ func TestCacheDO(t *testing.T) {
230227
}
231228

232229
func TestCacheCNAME(t *testing.T) {
233-
l := slogutil.NewDiscardLogger()
234-
235230
testCache := newTestCache(t, nil)
236231

237232
// Fill the cache
@@ -241,7 +236,7 @@ func TestCacheCNAME(t *testing.T) {
241236
},
242237
Answer: []dns.RR{newRR(t, "google.com.", dns.TypeCNAME, 3600, "test.google.com.")},
243238
}).SetQuestion("google.com.", dns.TypeA)
244-
testCache.set(reply, upstreamWithAddr, l)
239+
testCache.set(reply, upstreamWithAddr, testLogger)
245240

246241
// Create a DNS request.
247242
request := (&dns.Msg{}).SetQuestion("google.com.", dns.TypeA)
@@ -254,7 +249,7 @@ func TestCacheCNAME(t *testing.T) {
254249

255250
// Now fill the cache with a cacheable CNAME response.
256251
reply.Answer = append(reply.Answer, newRR(t, "google.com.", dns.TypeA, 3600, net.IP{8, 8, 8, 8}))
257-
testCache.set(reply, upstreamWithAddr, l)
252+
testCache.set(reply, upstreamWithAddr, testLogger)
258253

259254
// We are testing that a proper CNAME response gets cached
260255
t.Run("cnames_exist", func(t *testing.T) {
@@ -277,7 +272,7 @@ func TestCache_uncacheable(t *testing.T) {
277272
reply := (&dns.Msg{}).SetRcode(request, dns.RcodeBadAlg)
278273

279274
// We are testing that SERVFAIL responses aren't cached
280-
testCache.set(reply, upstreamWithAddr, slogutil.NewDiscardLogger())
275+
testCache.set(reply, upstreamWithAddr, testLogger)
281276

282277
r, expired, _ := testCache.get(request)
283278
assert.Nil(t, r)
@@ -318,20 +313,16 @@ func TestCacheExpiration(t *testing.T) {
318313
t.Parallel()
319314

320315
dnsProxy := mustNew(t, &Config{
321-
Logger: slogutil.NewDiscardLogger(),
322-
UDPListenAddr: []*net.UDPAddr{net.UDPAddrFromAddrPort(localhostAnyPort)},
323-
TCPListenAddr: []*net.TCPAddr{net.TCPAddrFromAddrPort(localhostAnyPort)},
324-
UpstreamConfig: newTestUpstreamConfig(t, defaultTimeout, testDefaultUpstreamAddr),
325-
TrustedProxies: defaultTrustedProxies,
326-
RatelimitSubnetLenIPv4: 24,
327-
RatelimitSubnetLenIPv6: 64,
328-
CacheEnabled: true,
316+
Logger: testLogger,
317+
UDPListenAddr: []*net.UDPAddr{net.UDPAddrFromAddrPort(localhostAnyPort)},
318+
TCPListenAddr: []*net.TCPAddr{net.TCPAddrFromAddrPort(localhostAnyPort)},
319+
UpstreamConfig: newTestUpstreamConfig(t, defaultTimeout, testDefaultUpstreamAddr),
320+
TrustedProxies: defaultTrustedProxies,
321+
CacheEnabled: true,
329322
})
330323

331324
servicetest.RequireRun(t, dnsProxy, testTimeout)
332325

333-
l := slogutil.NewDiscardLogger()
334-
335326
// Create dns messages with TTL of 1 second.
336327
rrs := []dns.RR{
337328
newRR(t, "youtube.com.", dns.TypeA, 1, net.IP{173, 194, 221, 198}),
@@ -346,7 +337,7 @@ func TestCacheExpiration(t *testing.T) {
346337
},
347338
Answer: []dns.RR{dns.Copy(rr)},
348339
}).SetQuestion(rr.Header().Name, dns.TypeA)
349-
dnsProxy.cache.set(rep, upstreamWithAddr, l)
340+
dnsProxy.cache.set(rep, upstreamWithAddr, testLogger)
350341
replies[i] = rep
351342
}
352343

@@ -375,18 +366,16 @@ func TestCacheExpirationWithTTLOverride(t *testing.T) {
375366
u := testUpstream{}
376367

377368
dnsProxy := mustNew(t, &Config{
378-
Logger: slogutil.NewDiscardLogger(),
369+
Logger: testLogger,
379370
UDPListenAddr: []*net.UDPAddr{net.UDPAddrFromAddrPort(localhostAnyPort)},
380371
TCPListenAddr: []*net.TCPAddr{net.TCPAddrFromAddrPort(localhostAnyPort)},
381372
UpstreamConfig: &UpstreamConfig{
382373
Upstreams: []upstream.Upstream{&u},
383374
},
384-
TrustedProxies: defaultTrustedProxies,
385-
RatelimitSubnetLenIPv4: 24,
386-
RatelimitSubnetLenIPv6: 64,
387-
CacheEnabled: true,
388-
CacheMinTTL: 20,
389-
CacheMaxTTL: 40,
375+
TrustedProxies: defaultTrustedProxies,
376+
CacheEnabled: true,
377+
CacheMinTTL: 20,
378+
CacheMaxTTL: 40,
390379
})
391380

392381
servicetest.RequireRun(t, dnsProxy, testTimeout)
@@ -556,8 +545,6 @@ func TestCache(t *testing.T) {
556545
}
557546

558547
func (tests testCases) run(t *testing.T) {
559-
l := slogutil.NewDiscardLogger()
560-
561548
testCache := newTestCache(t, nil)
562549

563550
for _, res := range tests.cache {
@@ -567,7 +554,7 @@ func (tests testCases) run(t *testing.T) {
567554
},
568555
Answer: res.a,
569556
}).SetQuestion(res.q, res.t)
570-
testCache.set(reply, upstreamWithAddr, l)
557+
testCache.set(reply, upstreamWithAddr, testLogger)
571558
}
572559

573560
for _, tc := range tests.cases {
@@ -590,7 +577,7 @@ func (tests testCases) run(t *testing.T) {
590577
Answer: tc.a,
591578
}).SetQuestion(tc.q, tc.t)
592579

593-
testCache.set(reply, upstreamWithAddr, l)
580+
testCache.set(reply, upstreamWithAddr, testLogger)
594581

595582
requireEqualMsgs(t, ci.m, reply)
596583
}
@@ -637,7 +624,7 @@ func setAndGetCache(t *testing.T, c *cache, g *sync.WaitGroup, host, ip string)
637624
Answer: []dns.RR{newRR(t, host, dns.TypeA, 1, ipAddr)},
638625
}).SetQuestion(host, dns.TypeA)
639626

640-
c.set(dnsMsg, upstreamWithAddr, slogutil.NewDiscardLogger())
627+
c.set(dnsMsg, upstreamWithAddr, testLogger)
641628

642629
for range 2 {
643630
ci, expired, key := c.get(dnsMsg)
@@ -663,7 +650,6 @@ func TestCache_getWithSubnet(t *testing.T) {
663650
req := (&dns.Msg{}).SetQuestion(testFQDN, dns.TypeA)
664651
mask16 := net.CIDRMask(16, netutil.IPv4BitLen)
665652
mask24 := net.CIDRMask(24, netutil.IPv4BitLen)
666-
l := slogutil.NewDiscardLogger()
667653

668654
c := newTestCache(t, &cacheConfig{withECS: true})
669655

@@ -677,7 +663,7 @@ func TestCache_getWithSubnet(t *testing.T) {
677663
resp := (&dns.Msg{
678664
Answer: []dns.RR{newRR(t, testFQDN, dns.TypeA, 1, net.IP{1, 1, 1, 1})},
679665
}).SetReply(req)
680-
c.setWithSubnet(resp, upstreamWithAddr, &net.IPNet{IP: ip1234, Mask: mask16}, slogutil.NewDiscardLogger())
666+
c.setWithSubnet(resp, upstreamWithAddr, &net.IPNet{IP: ip1234, Mask: mask16}, testLogger)
681667

682668
t.Run("different_ip", func(t *testing.T) {
683669
ci, expired, key := c.getWithSubnet(req, &net.IPNet{IP: ip2234, Mask: mask24})
@@ -690,13 +676,13 @@ func TestCache_getWithSubnet(t *testing.T) {
690676
resp = (&dns.Msg{
691677
Answer: []dns.RR{newRR(t, testFQDN, dns.TypeA, 1, net.IP{2, 2, 2, 2})},
692678
}).SetReply(req)
693-
c.setWithSubnet(resp, upstreamWithAddr, &net.IPNet{IP: ip2234, Mask: mask16}, l)
679+
c.setWithSubnet(resp, upstreamWithAddr, &net.IPNet{IP: ip2234, Mask: mask16}, testLogger)
694680

695681
// Add a response entry without subnet.
696682
resp = (&dns.Msg{
697683
Answer: []dns.RR{newRR(t, testFQDN, dns.TypeA, 1, net.IP{3, 3, 3, 3})},
698684
}).SetReply(req)
699-
c.setWithSubnet(resp, upstreamWithAddr, &net.IPNet{IP: nil, Mask: nil}, l)
685+
c.setWithSubnet(resp, upstreamWithAddr, &net.IPNet{IP: nil, Mask: nil}, testLogger)
700686

701687
t.Run("with_subnet_1", func(t *testing.T) {
702688
ci, expired, key := c.getWithSubnet(req, &net.IPNet{IP: ip1234, Mask: mask24})
@@ -766,7 +752,7 @@ func TestCache_getWithSubnet_mask(t *testing.T) {
766752
resp,
767753
upstreamWithAddr,
768754
&net.IPNet{IP: cachedIP, Mask: cidrMask},
769-
slogutil.NewDiscardLogger(),
755+
testLogger,
770756
)
771757

772758
t.Run("mask_matched", func(t *testing.T) {
@@ -994,7 +980,7 @@ func TestCache_IsCacheable_negative(t *testing.T) {
994980

995981
for _, tc := range testCases {
996982
t.Run(tc.name, func(t *testing.T) {
997-
assert.Equal(t, tc.wantTTL, cacheTTL(tc.req, slogutil.NewDiscardLogger()))
983+
assert.Equal(t, tc.wantTTL, cacheTTL(tc.req, testLogger))
998984
})
999985
}
1000986
}

0 commit comments

Comments
 (0)