From 1a9843735eb898652a3a182479c1bd44484742d4 Mon Sep 17 00:00:00 2001
From: Stanislav Chzhen <s.chzhen@adguard.com>
Date: Mon, 30 Oct 2023 19:10:50 +0300
Subject: [PATCH] all: ratelimit subnet len

---
 main.go            |  6 ++++++
 proxy/config.go    |  3 +++
 proxy/ratelimit.go | 10 ++++++++++
 3 files changed, 19 insertions(+)

diff --git a/main.go b/main.go
index 654631c49..11f6951e3 100644
--- a/main.go
+++ b/main.go
@@ -140,6 +140,9 @@ type Options struct {
 	// Anti-DNS amplification measures
 	// --
 
+	RatelimitSubnetLenV4 int `yaml:"ratelimit-subnet-len-v4" long:"ratelimit-subnet-len-v4"`
+	RatelimitSubnetLenV6 int `yaml:"ratelimit-subnet-len-v6" long:"ratelimit-subnet-len-v6"`
+
 	// Ratelimit value
 	Ratelimit int `yaml:"ratelimit" short:"r" long:"ratelimit" description:"Ratelimit (requests per second)"`
 
@@ -320,6 +323,9 @@ func runPprof(options *Options) {
 func createProxyConfig(options *Options) proxy.Config {
 	// Create the config
 	config := proxy.Config{
+		RatelimitSubnetLenV4: options.RatelimitSubnetLenV4,
+		RatelimitSubnetLenV6: options.RatelimitSubnetLenV6,
+
 		Ratelimit:       options.Ratelimit,
 		CacheEnabled:    options.Cache,
 		CacheSizeBytes:  options.CacheSizeBytes,
diff --git a/proxy/config.go b/proxy/config.go
index 45a18ee36..80b63cad0 100644
--- a/proxy/config.go
+++ b/proxy/config.go
@@ -64,6 +64,9 @@ type Config struct {
 	// Rate-limiting and anti-DNS amplification measures
 	// --
 
+	RatelimitSubnetLenV4 int
+	RatelimitSubnetLenV6 int
+
 	Ratelimit          int      // max number of requests per second from a given IP (0 to disable)
 	RatelimitWhitelist []string // a list of whitelisted client IP addresses
 	RefuseAny          bool     // if true, refuse ANY requests
diff --git a/proxy/ratelimit.go b/proxy/ratelimit.go
index 0a73cc1f1..45061b5a7 100644
--- a/proxy/ratelimit.go
+++ b/proxy/ratelimit.go
@@ -53,6 +53,16 @@ func (p *Proxy) isRatelimited(addr net.Addr) (ok bool) {
 		}
 	}
 
+	var mask net.IPMask
+	if len(ip) == net.IPv4len {
+		mask = net.CIDRMask(p.RatelimitSubnetLenV4, 32)
+	} else {
+		mask = net.CIDRMask(p.RatelimitSubnetLenV6, 128)
+	}
+
+	ip = ip.Mask(mask)
+	ipStr = ip.String()
+
 	value := p.limiterForIP(ipStr)
 	rl, ok := value.(*rate.RateLimiter)
 	if !ok {