Skip to content

Commit 319ee6c

Browse files
MizzickMizzick
committed
all: use ratelimit middleware
1 parent b1e1f52 commit 319ee6c

18 files changed

+168
-262
lines changed

‎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: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import (
1616
"github.com/AdguardTeam/dnsproxy/internal/handler"
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"
@@ -44,6 +45,19 @@ func createProxyConfig(
4445
return nil, fmt.Errorf("reading hosts files: %w", err)
4546
}
4647

48+
mwConf := &ratelimit.Config{
49+
Logger: l.With(slogutil.KeyPrefix, "ratelimit"),
50+
AllowlistAddrs: []netip.Addr{},
51+
Ratelimit: conf.Ratelimit,
52+
SubnetLenIPv4: conf.RatelimitSubnetLenIPv4,
53+
SubnetLenIPv6: conf.RatelimitSubnetLenIPv6,
54+
}
55+
if err = mwConf.Validate(); err != nil {
56+
return nil, fmt.Errorf("invalid ratelimit configuration: %w", err)
57+
}
58+
59+
mw := ratelimit.NewMiddleware(mwConf)
60+
4761
reqHdlr := handler.NewDefault(&handler.DefaultConfig{
4862
Logger: l.With(slogutil.KeyPrefix, "default_handler"),
4963
// TODO(e.burkov): Use the configured message constructor.
@@ -55,10 +69,6 @@ func createProxyConfig(
5569
proxyConf = &proxy.Config{
5670
Logger: l.With(slogutil.KeyPrefix, proxy.LogPrefix),
5771

58-
RatelimitSubnetLenIPv4: conf.RatelimitSubnetLenIPv4,
59-
RatelimitSubnetLenIPv6: conf.RatelimitSubnetLenIPv6,
60-
61-
Ratelimit: conf.Ratelimit,
6272
CacheEnabled: conf.Cache,
6373
CacheSizeBytes: conf.CacheSizeBytes,
6474
CacheMinTTL: conf.CacheMinTTL,
@@ -81,6 +91,7 @@ func createProxyConfig(
8191
MaxGoroutines: conf.MaxGoRoutines,
8292
UsePrivateRDNS: conf.UsePrivateRDNS,
8393
PrivateSubnets: netutil.SubnetSetFunc(netutil.IsLocallyServed),
94+
Middleware: mw,
8495
RequestHandler: reqHdlr,
8596
PendingRequests: &proxy.PendingRequestsConfig{
8697
Enabled: conf.PendingRequestsEnabled,

‎proxy/bogusnxdomain_internal_test.go‎

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -15,14 +15,12 @@ import (
1515

1616
func TestProxy_IsBogusNXDomain(t *testing.T) {
1717
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,
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+
CacheEnabled: true,
2624
BogusNXDomain: []netip.Prefix{
2725
netip.MustParsePrefix("4.3.2.1/24"),
2826
netip.MustParsePrefix("1.2.3.4/8"),

‎proxy/cache_internal_test.go‎

Lines changed: 16 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -53,14 +53,12 @@ func newTestCache(tb testing.TB, conf *cacheConfig) (c *cache) {
5353

5454
func TestServeCached(t *testing.T) {
5555
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,
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+
CacheEnabled: true,
6462
})
6563

6664
// Start listening.
@@ -318,14 +316,12 @@ func TestCacheExpiration(t *testing.T) {
318316
t.Parallel()
319317

320318
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,
319+
Logger: slogutil.NewDiscardLogger(),
320+
UDPListenAddr: []*net.UDPAddr{net.UDPAddrFromAddrPort(localhostAnyPort)},
321+
TCPListenAddr: []*net.TCPAddr{net.TCPAddrFromAddrPort(localhostAnyPort)},
322+
UpstreamConfig: newTestUpstreamConfig(t, defaultTimeout, testDefaultUpstreamAddr),
323+
TrustedProxies: defaultTrustedProxies,
324+
CacheEnabled: true,
329325
})
330326

331327
servicetest.RequireRun(t, dnsProxy, testTimeout)
@@ -381,12 +377,10 @@ func TestCacheExpirationWithTTLOverride(t *testing.T) {
381377
UpstreamConfig: &UpstreamConfig{
382378
Upstreams: []upstream.Upstream{&u},
383379
},
384-
TrustedProxies: defaultTrustedProxies,
385-
RatelimitSubnetLenIPv4: 24,
386-
RatelimitSubnetLenIPv6: 64,
387-
CacheEnabled: true,
388-
CacheMinTTL: 20,
389-
CacheMaxTTL: 40,
380+
TrustedProxies: defaultTrustedProxies,
381+
CacheEnabled: true,
382+
CacheMinTTL: 20,
383+
CacheMaxTTL: 40,
390384
})
391385

392386
servicetest.RequireRun(t, dnsProxy, testTimeout)

‎proxy/config.go‎

Lines changed: 0 additions & 75 deletions
Original file line numberDiff line numberDiff line change
@@ -150,34 +150,9 @@ type Config struct {
150150
// default Well-Known Prefix.
151151
DNS64Prefs []netip.Prefix
152152

153-
// RatelimitWhitelist is a list of IP addresses excluded from rate limiting.
154-
//
155-
// TODO(d.kolyshev): !! Remove.
156-
RatelimitWhitelist []netip.Addr
157-
158153
// EDNSAddr is the ECS IP used in request.
159154
EDNSAddr net.IP
160155

161-
// TODO(s.chzhen): Extract ratelimit settings to a separate structure.
162-
163-
// RatelimitSubnetLenIPv4 is a subnet length for IPv4 addresses used for
164-
// rate limiting requests.
165-
//
166-
// TODO(d.kolyshev): !! Remove.
167-
RatelimitSubnetLenIPv4 int
168-
169-
// RatelimitSubnetLenIPv6 is a subnet length for IPv6 addresses used for
170-
// rate limiting requests.
171-
//
172-
// TODO(d.kolyshev): !! Remove.
173-
RatelimitSubnetLenIPv6 int
174-
175-
// Ratelimit is a maximum number of requests per second from a given IP (0
176-
// to disable).
177-
//
178-
// TODO(d.kolyshev): !! Remove.
179-
Ratelimit int
180-
181156
// CacheSizeBytes is the maximum cache size in bytes.
182157
CacheSizeBytes int
183158

@@ -293,11 +268,6 @@ func (p *Proxy) validateConfig() (err error) {
293268
return fmt.Errorf("fallbacks: %w", err)
294269
}
295270

296-
err = p.validateRatelimit()
297-
if err != nil {
298-
return fmt.Errorf("ratelimit: %w", err)
299-
}
300-
301271
switch p.UpstreamMode {
302272
case
303273
"",
@@ -319,57 +289,12 @@ func (p *Proxy) validateConfig() (err error) {
319289
return nil
320290
}
321291

322-
// validateRatelimit validates ratelimit configuration and returns an error if
323-
// it's invalid.
324-
func (p *Proxy) validateRatelimit() (err error) {
325-
if p.Ratelimit == 0 {
326-
return nil
327-
}
328-
329-
err = checkInclusion(p.RatelimitSubnetLenIPv4, 0, netutil.IPv4BitLen)
330-
if err != nil {
331-
return fmt.Errorf("ratelimit subnet len ipv4 is invalid: %w", err)
332-
}
333-
334-
err = checkInclusion(p.RatelimitSubnetLenIPv6, 0, netutil.IPv6BitLen)
335-
if err != nil {
336-
return fmt.Errorf("ratelimit subnet len ipv6 is invalid: %w", err)
337-
}
338-
339-
return nil
340-
}
341-
342-
// checkInclusion returns an error if a n is not in the inclusive range between
343-
// minN and maxN.
344-
func checkInclusion(n, minN, maxN int) (err error) {
345-
switch {
346-
case n < minN:
347-
return fmt.Errorf("value %d less than min %d", n, minN)
348-
case n > maxN:
349-
return fmt.Errorf("value %d greater than max %d", n, maxN)
350-
}
351-
352-
return nil
353-
}
354-
355292
// logConfigInfo logs proxy configuration information.
356293
func (p *Proxy) logConfigInfo() {
357294
if p.CacheMinTTL > 0 || p.CacheMaxTTL > 0 {
358295
p.logger.Info("cache ttl override is enabled", "min", p.CacheMinTTL, "max", p.CacheMaxTTL)
359296
}
360297

361-
if p.Ratelimit > 0 {
362-
p.logger.Info(
363-
"ratelimit is enabled",
364-
"rps",
365-
p.Ratelimit,
366-
"ipv4_subnet_mask_len",
367-
p.RatelimitSubnetLenIPv4,
368-
"ipv6_subnet_mask_len",
369-
p.RatelimitSubnetLenIPv6,
370-
)
371-
}
372-
373298
if p.RefuseAny {
374299
p.logger.Info("server will refuse requests of type any")
375300
}

‎proxy/dns64_internal_test.go‎

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -50,9 +50,7 @@ func TestDNS64Race(t *testing.T) {
5050
PrivateRDNSUpstreamConfig: &UpstreamConfig{
5151
Upstreams: []upstream.Upstream{localUps},
5252
},
53-
TrustedProxies: defaultTrustedProxies,
54-
RatelimitSubnetLenIPv4: 24,
55-
RatelimitSubnetLenIPv6: 64,
53+
TrustedProxies: defaultTrustedProxies,
5654

5755
UseDNS64: true,
5856
UsePrivateRDNS: true,
@@ -363,10 +361,8 @@ func TestProxy_Resolve_dns64(t *testing.T) {
363361
PrivateRDNSUpstreamConfig: &UpstreamConfig{
364362
Upstreams: []upstream.Upstream{localUps},
365363
},
366-
TrustedProxies: defaultTrustedProxies,
367-
RatelimitSubnetLenIPv4: 24,
368-
RatelimitSubnetLenIPv6: 64,
369-
CacheEnabled: true,
364+
TrustedProxies: defaultTrustedProxies,
365+
CacheEnabled: true,
370366

371367
UseDNS64: true,
372368
UsePrivateRDNS: true,

‎proxy/exchange_internal_test.go‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -215,9 +215,7 @@ func TestProxy_Exchange_loadBalance(t *testing.T) {
215215
UpstreamConfig: &UpstreamConfig{
216216
Upstreams: ups,
217217
},
218-
TrustedProxies: defaultTrustedProxies,
219-
RatelimitSubnetLenIPv4: 24,
220-
RatelimitSubnetLenIPv6: 64,
218+
TrustedProxies: defaultTrustedProxies,
221219
})
222220
p.time = tc.clock
223221
p.randSrc = randSrc

‎proxy/handler_internal_test.go‎

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -40,14 +40,12 @@ func TestFilteringHandler(t *testing.T) {
4040

4141
// Prepare the proxy server.
4242
dnsProxy := mustNew(t, &Config{
43-
Logger: slogutil.NewDiscardLogger(),
44-
TrustedProxies: defaultTrustedProxies,
45-
UpstreamConfig: newTestUpstreamConfig(t, defaultTimeout, testDefaultUpstreamAddr),
46-
RequestHandler: reqHandler,
47-
UDPListenAddr: []*net.UDPAddr{net.UDPAddrFromAddrPort(localhostAnyPort)},
48-
TCPListenAddr: []*net.TCPAddr{net.TCPAddrFromAddrPort(localhostAnyPort)},
49-
RatelimitSubnetLenIPv4: 24,
50-
RatelimitSubnetLenIPv6: 64,
43+
Logger: slogutil.NewDiscardLogger(),
44+
TrustedProxies: defaultTrustedProxies,
45+
UpstreamConfig: newTestUpstreamConfig(t, defaultTimeout, testDefaultUpstreamAddr),
46+
RequestHandler: reqHandler,
47+
UDPListenAddr: []*net.UDPAddr{net.UDPAddrFromAddrPort(localhostAnyPort)},
48+
TCPListenAddr: []*net.TCPAddr{net.TCPAddrFromAddrPort(localhostAnyPort)},
5149
})
5250

5351
servicetest.RequireRun(t, dnsProxy, testTimeout)

‎proxy/pending_test.go‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -115,9 +115,6 @@ func TestPendingRequests(t *testing.T) {
115115
RequestHandler: reqHandler,
116116
UDPListenAddr: []*net.UDPAddr{net.UDPAddrFromAddrPort(localhostAnyPort)},
117117
TCPListenAddr: []*net.TCPAddr{net.TCPAddrFromAddrPort(localhostAnyPort)},
118-
RatelimitSubnetLenIPv4: 24,
119-
RatelimitSubnetLenIPv6: 64,
120-
Ratelimit: 0,
121118
CacheSizeBytes: testCacheSize,
122119
CacheEnabled: true,
123120
EnableEDNSClientSubnet: true,

‎proxy/proxy.go‎

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@ import (
1212
"net"
1313
"net/http"
1414
"net/netip"
15-
"slices"
1615
"sync"
1716
"sync/atomic"
1817
"time"
@@ -287,10 +286,6 @@ func New(c *Config) (p *Proxy, err error) {
287286
return nil, fmt.Errorf("setting up DNS64: %w", err)
288287
}
289288

290-
// TODO(e.burkov): Clone all mutable fields of Config.
291-
p.RatelimitWhitelist = slices.Clone(p.RatelimitWhitelist)
292-
slices.SortFunc(p.RatelimitWhitelist, netip.Addr.Compare)
293-
294289
return p, nil
295290
}
296291

0 commit comments

Comments
 (0)