Pull request 300: AG-27616-ratelimit-whitelist-netip-addr

Squashed commit of the following:

commit 6b8f86f
Merge: 8cc3b79 280cca3
Author: Stanislav Chzhen <[email protected]>
Date:   Tue Nov 28 18:59:53 2023 +0300

    Merge branch 'master' into AG-27616-ratelimit-whitelist-netip-addr

commit 8cc3b79
Author: Stanislav Chzhen <[email protected]>
Date:   Fri Nov 24 16:17:48 2023 +0300

    all: imp logs

commit 53b095f
Author: Stanislav Chzhen <[email protected]>
Date:   Thu Nov 23 15:46:08 2023 +0300

    proxy: imp err msg

commit d5d5d8f
Author: Stanislav Chzhen <[email protected]>
Date:   Wed Nov 22 17:08:45 2023 +0300

    proxy: imp docs

commit f2d1c21
Author: Stanislav Chzhen <[email protected]>
Date:   Tue Nov 21 19:57:33 2023 +0300

    proxy: imp code

commit 5deb3ec
Author: Stanislav Chzhen <[email protected]>
Date:   Tue Nov 21 19:26:45 2023 +0300

    all: imp code

commit 5f36d90
Author: Stanislav Chzhen <[email protected]>
Date:   Tue Nov 21 17:49:54 2023 +0300

    all: add todo

commit e63df13
Author: Stanislav Chzhen <[email protected]>
Date:   Mon Nov 20 20:34:14 2023 +0300

    all: dnscontext addr netip addrport

commit b6588c0
Author: Stanislav Chzhen <[email protected]>
Date:   Thu Nov 16 17:46:29 2023 +0300

    proxy: ratelimit whitelsit netip addr

Author: schzhn

internal/netutil/udp.go

@@ -16,6 +16,8 @@ func UDPSetOptions(c *net.UDPConn) (err error) {
 // UDPRead reads the message from conn using buf and receives a control-message
 // payload of size udpOOBSize from it.  It returns the number of bytes copied
 // into buf and the source address of the message.
+//
+// TODO(s.chzhen):  Consider using netip.Addr.
 func UDPRead(
 	conn *net.UDPConn,
 	buf []byte,
@@ -25,6 +27,8 @@ func UDPRead(
 }
 
 // UDPWrite writes the data to the remoteAddr using conn.
+//
+// TODO(s.chzhen):  Consider using netip.Addr.
 func UDPWrite(
 	data []byte,
 	conn *net.UDPConn,

main.go

@@ -23,7 +23,6 @@ import (
 	"github.com/AdguardTeam/dnsproxy/upstream"
 	"github.com/AdguardTeam/golibs/log"
 	"github.com/AdguardTeam/golibs/mathutil"
-	"github.com/AdguardTeam/golibs/netutil"
 	"github.com/AdguardTeam/golibs/timeutil"
 	"github.com/ameshkov/dnscrypt/v2"
 	goFlags "github.com/jessevdk/go-flags"
@@ -153,7 +152,7 @@ type Options struct {
 
 	// RatelimitSubnetLenIPv6 is a subnet length for IPv6 addresses used for
 	// rate limiting requests
-	RatelimitSubnetLenIPv6 int `yaml:"ratelimit-subnet-len-ipv6" long:"ratelimit-subnet-len-ipv6" description:"Ratelimit subnet length for IPv6." default:"64"`
+	RatelimitSubnetLenIPv6 int `yaml:"ratelimit-subnet-len-ipv6" long:"ratelimit-subnet-len-ipv6" description:"Ratelimit subnet length for IPv6." default:"56"`
 
 	// If true, refuse ANY requests
 	RefuseAny bool `yaml:"refuse-any" long:"refuse-any" description:"If specified, refuse ANY requests" optional:"yes" optional-value:"true"`
@@ -336,8 +335,8 @@ func runPprof(options *Options) {
 // createProxyConfig creates proxy.Config from the command line arguments
 func createProxyConfig(options *Options) (conf proxy.Config) {
 	conf = proxy.Config{
-		RatelimitSubnetMaskIPv4: net.CIDRMask(options.RatelimitSubnetLenIPv4, netutil.IPv4BitLen),
-		RatelimitSubnetMaskIPv6: net.CIDRMask(options.RatelimitSubnetLenIPv6, netutil.IPv6BitLen),
+		RatelimitSubnetLenIPv4: options.RatelimitSubnetLenIPv4,
+		RatelimitSubnetLenIPv6: options.RatelimitSubnetLenIPv6,
 
 		Ratelimit:       options.Ratelimit,
 		CacheEnabled:    options.Cache,

proxy/cache_test.go

@@ -2,6 +2,7 @@ package proxy
 
 import (
 	"net"
+	"net/netip"
 	"strings"
 	"sync"
 	"testing"
@@ -324,7 +325,7 @@ func TestCacheExpirationWithTTLOverride(t *testing.T) {
 
 	t.Run("replace_min", func(t *testing.T) {
 		d.Req = createHostTestMessage("host")
-		d.Addr = &net.TCPAddr{}
+		d.Addr = netip.AddrPort{}
 
 		u.ans = []dns.RR{&dns.A{
 			Hdr: dns.RR_Header{
@@ -348,7 +349,7 @@ func TestCacheExpirationWithTTLOverride(t *testing.T) {
 
 	t.Run("replace_max", func(t *testing.T) {
 		d.Req = createHostTestMessage("host2")
-		d.Addr = &net.TCPAddr{}
+		d.Addr = netip.AddrPort{}
 
 		u.ans = []dns.RR{&dns.A{
 			Hdr: dns.RR_Header{

proxy/config.go

@@ -65,18 +65,26 @@ type Config struct {
 
 	// Rate-limiting and anti-DNS amplification measures
 	// --
+	//
+	// TODO(s.chzhen):  Extract ratelimit settings to a separate structure.
 
-	// RatelimitSubnetMaskIPv4 is a subnet mask for IPv4 addresses used for
+	// RatelimitSubnetLenIPv4 is a subnet length for IPv4 addresses used for
 	// rate limiting requests.
-	RatelimitSubnetMaskIPv4 net.IPMask
+	RatelimitSubnetLenIPv4 int
 
-	// RatelimitSubnetMaskIPv6 is a subnet mask for IPv6 addresses used for
+	// RatelimitSubnetLenIPv6 is a subnet length for IPv6 addresses used for
 	// rate limiting requests.
-	RatelimitSubnetMaskIPv6 net.IPMask
+	RatelimitSubnetLenIPv6 int
+
+	// Ratelimit is a maximum number of requests per second from a given IP (0
+	// to disable).
+	Ratelimit int
+
+	// RatelimitWhitelist is a list of IP addresses excluded from rate limiting.
+	RatelimitWhitelist []netip.Addr
 
-	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
+	// RefuseAny makes proxy refuse the requests of type ANY.
+	RefuseAny bool
 
 	// TrustedProxies is the list of IP addresses and CIDR networks to
 	// detect proxy servers addresses the DoH requests from which should be
@@ -240,22 +248,27 @@ func (p *Proxy) validateRatelimit() (err error) {
 		return nil
 	}
 
-	if p.RatelimitSubnetMaskIPv4 == nil {
-		return errors.Error("ipv4 subnet mask is nil")
+	err = checkInclusion(p.RatelimitSubnetLenIPv4, 0, netutil.IPv4BitLen)
+	if err != nil {
+		return fmt.Errorf("ratelimit subnet len ipv4 is invalid: %w", err)
 	}
 
-	_, bits := p.RatelimitSubnetMaskIPv4.Size()
-	if bits != netutil.IPv4BitLen {
-		return fmt.Errorf("ipv4 subnet mask must contain %d bits, got %d", netutil.IPv4BitLen, bits)
+	err = checkInclusion(p.RatelimitSubnetLenIPv6, 0, netutil.IPv6BitLen)
+	if err != nil {
+		return fmt.Errorf("ratelimit subnet len ipv6 is invalid: %w", err)
 	}
 
-	if p.RatelimitSubnetMaskIPv6 == nil {
-		return errors.Error("ipv6 subnet is nil")
-	}
+	return nil
+}
 
-	_, bits = p.RatelimitSubnetMaskIPv6.Size()
-	if bits != netutil.IPv6BitLen {
-		return fmt.Errorf("ipv6 subnet mask must contain %d bits, got %d", netutil.IPv6BitLen, bits)
+// checkInclusion returns an error if a n is not in the inclusive range between
+// minN and maxN.
+func checkInclusion(n, minN, maxN int) (err error) {
+	switch {
+	case n < minN:
+		return fmt.Errorf("value %d less than min %d", n, minN)
+	case n > maxN:
+		return fmt.Errorf("value %d greater than max %d", n, maxN)
 	}
 
 	return nil
@@ -268,14 +281,11 @@ func (p *Proxy) logConfigInfo() {
 	}
 
 	if p.Ratelimit > 0 {
-		sizeV4, _ := p.RatelimitSubnetMaskIPv4.Size()
-		sizeV6, _ := p.RatelimitSubnetMaskIPv6.Size()
-
 		log.Info(
 			"Ratelimit is enabled and set to %d rps, IPv4 subnet mask len %d, IPv6 subnet mask len %d",
 			p.Ratelimit,
-			sizeV4,
-			sizeV6,
+			p.RatelimitSubnetLenIPv4,
+			p.RatelimitSubnetLenIPv6,
 		)
 	}
 

proxy/dns64_test.go

@@ -156,10 +156,7 @@ func TestProxy_Resolve_dns64(t *testing.T) {
 	require.NoError(t, err)
 	ptrGlobDomain = dns.Fqdn(ptrGlobDomain)
 
-	cliIP := &net.TCPAddr{
-		IP:   net.IP{192, 168, 1, 1},
-		Port: 1234,
-	}
+	cliAddrPort := netip.MustParseAddrPort("192.168.1.1:1234")
 
 	const (
 		sectionAnswer = iota
@@ -354,7 +351,7 @@ func TestProxy_Resolve_dns64(t *testing.T) {
 			req := (&dns.Msg{}).SetQuestion(tc.qname, tc.qtype)
 			dctx := &DNSContext{
 				Req:  req,
-				Addr: cliIP,
+				Addr: cliAddrPort,
 			}
 
 			err = p.Resolve(dctx)

proxy/dnscontext.go

@@ -3,6 +3,7 @@ package proxy
 import (
 	"net"
 	"net/http"
+	"net/netip"
 	"time"
 
 	"github.com/AdguardTeam/dnsproxy/upstream"
@@ -27,9 +28,7 @@ type DNSContext struct {
 	QUICStream quic.Stream
 
 	// Addr is the address of the client.
-	//
-	// TODO(s.chzhen):  Use [netip.AddrPort].
-	Addr net.Addr
+	Addr netip.AddrPort
 
 	// Upstream is the upstream that resolved the request.  In case of cached
 	// response it's nil.

proxy/proxy.go

@@ -24,6 +24,7 @@ import (
 	gocache "github.com/patrickmn/go-cache"
 	"github.com/quic-go/quic-go"
 	"github.com/quic-go/quic-go/http3"
+	"golang.org/x/exp/slices"
 )
 
 const (
@@ -183,6 +184,7 @@ type Proxy struct {
 
 // Init populates fields of p but does not start listeners.
 func (p *Proxy) Init() (err error) {
+	// TODO(s.chzhen):  Consider moving to [Proxy.validateConfig].
 	err = p.validateBasicAuth()
 	if err != nil {
 		return fmt.Errorf("basic auth: %w", err)
@@ -233,6 +235,9 @@ func (p *Proxy) Init() (err error) {
 		return fmt.Errorf("setting up DNS64: %w", err)
 	}
 
+	p.RatelimitWhitelist = slices.Clone(p.RatelimitWhitelist)
+	slices.SortFunc(p.RatelimitWhitelist, netip.Addr.Compare)
+
 	return nil
 }
 
@@ -512,10 +517,9 @@ func (p *Proxy) selectUpstreams(d *DNSContext) (upstreams []upstream.Upstream) {
 		return nil
 	}
 
-	ip, _ := netutil.IPAndPortFromAddr(d.Addr)
 	// TODO(e.burkov):  Detect against the actual configured subnet set.
 	// Perhaps, even much earlier.
-	if !netutil.IsLocallyServed(ip) {
+	if !netutil.IsLocallyServedAddr(d.Addr.Addr()) {
 		return nil
 	}
 
@@ -708,10 +712,7 @@ func (dctx *DNSContext) processECS(cliIP net.IP) {
 
 	// Set ECS.
 	if cliIP == nil {
-		cliIP, _ = netutil.IPAndPortFromAddr(dctx.Addr)
-		if cliIP == nil {
-			return
-		}
+		cliIP = dctx.Addr.Addr().AsSlice()
 	}
 
 	if !netutil.IsSpecialPurpose(cliIP) {

proxy/proxy_test.go

@@ -852,10 +852,8 @@ func TestProxy_ReplyFromUpstream_badResponse(t *testing.T) {
 			0,
 			false,
 		),
-		Req: createHostTestMessage("host"),
-		Addr: &net.TCPAddr{
-			IP: net.IP{1, 2, 3, 0},
-		},
+		Req:  createHostTestMessage("host"),
+		Addr: netip.MustParseAddrPort("1.2.3.0:1234"),
 	}
 
 	var err error
@@ -893,7 +891,7 @@ func TestExchangeCustomUpstreamConfig(t *testing.T) {
 			false,
 		),
 		Req:  createHostTestMessage("host"),
-		Addr: &net.TCPAddr{IP: net.IP{1, 2, 3, 0}},
+		Addr: netip.MustParseAddrPort("1.2.3.0:1234"),
 	}
 
 	err = prx.Resolve(&d)
@@ -945,7 +943,7 @@ func TestExchangeCustomUpstreamConfigCache(t *testing.T) {
 	d := DNSContext{
 		CustomUpstreamConfig: customUpstreamConfig,
 		Req:                  createHostTestMessage("host"),
-		Addr:                 &net.TCPAddr{IP: net.IP{1, 2, 3, 0}},
+		Addr:                 netip.MustParseAddrPort("1.2.3.0:1234"),
 	}
 
 	err = prx.Resolve(&d)
@@ -1036,7 +1034,7 @@ func TestECSProxy(t *testing.T) {
 	t.Run("cache_subnet", func(t *testing.T) {
 		d := DNSContext{
 			Req:  createHostTestMessage("host"),
-			Addr: &net.TCPAddr{IP: ip1230},
+			Addr: netip.MustParseAddrPort("1.2.3.0:1234"),
 		}
 
 		err = prx.Resolve(&d)
@@ -1049,7 +1047,7 @@ func TestECSProxy(t *testing.T) {
 	t.Run("serve_subnet_cache", func(t *testing.T) {
 		d := DNSContext{
 			Req:  createHostTestMessage("host"),
-			Addr: &net.TCPAddr{IP: net.IP{1, 2, 3, 1}},
+			Addr: netip.MustParseAddrPort("1.2.3.1:1234"),
 		}
 		u.ans, u.ecsIP, u.ecsReqIP = nil, nil, nil
 
@@ -1063,7 +1061,7 @@ func TestECSProxy(t *testing.T) {
 	t.Run("another_subnet", func(t *testing.T) {
 		d := DNSContext{
 			Req:  createHostTestMessage("host"),
-			Addr: &net.TCPAddr{IP: ip2230},
+			Addr: netip.MustParseAddrPort("2.2.3.0:1234"),
 		}
 		u.ans = []dns.RR{&dns.A{
 			Hdr: dns.RR_Header{Rrtype: dns.TypeA, Name: "host.", Ttl: 60},
@@ -1081,7 +1079,7 @@ func TestECSProxy(t *testing.T) {
 	t.Run("cache_general", func(t *testing.T) {
 		d := DNSContext{
 			Req:  createHostTestMessage("host"),
-			Addr: &net.TCPAddr{IP: net.IP{127, 0, 0, 1}},
+			Addr: netip.MustParseAddrPort("127.0.0.1:1234"),
 		}
 		u.ans = []dns.RR{&dns.A{
 			Hdr: dns.RR_Header{Rrtype: dns.TypeA, Name: "host.", Ttl: 60},
@@ -1099,7 +1097,7 @@ func TestECSProxy(t *testing.T) {
 	t.Run("serve_general_cache", func(t *testing.T) {
 		d := DNSContext{
 			Req:  createHostTestMessage("host"),
-			Addr: &net.TCPAddr{IP: net.IP{127, 0, 0, 2}},
+			Addr: netip.MustParseAddrPort("127.0.0.2:1234"),
 		}
 		u.ans, u.ecsIP, u.ecsReqIP = nil, nil, nil
 
@@ -1138,7 +1136,7 @@ func TestECSProxyCacheMinMaxTTL(t *testing.T) {
 	// first request
 	d := DNSContext{
 		Req:  createHostTestMessage("host"),
-		Addr: &net.TCPAddr{IP: clientIP},
+		Addr: netip.MustParseAddrPort("1.2.3.0:1234"),
 	}
 	err = prx.Resolve(&d)
 	require.NoError(t, err)
@@ -1156,9 +1154,7 @@ func TestECSProxyCacheMinMaxTTL(t *testing.T) {
 	// 2nd request
 	clientIP = net.IP{1, 2, 4, 0}
 	d.Req = createHostTestMessage("host")
-	d.Addr = &net.TCPAddr{
-		IP: clientIP,
-	}
+	d.Addr = netip.MustParseAddrPort("1.2.4.0:1234")
 	u.ans = []dns.RR{&dns.A{
 		Hdr: dns.RR_Header{
 			Rrtype: dns.TypeA,
@@ -1242,8 +1238,8 @@ func createTestProxy(t *testing.T, tlsConfig *tls.Config) *Proxy {
 
 	p.TrustedProxies = []string{"0.0.0.0/0", "::0/0"}
 
-	p.RatelimitSubnetMaskIPv4 = net.CIDRMask(24, 32)
-	p.RatelimitSubnetMaskIPv6 = net.CIDRMask(64, 128)
+	p.RatelimitSubnetLenIPv4 = 24
+	p.RatelimitSubnetLenIPv6 = 64
 
 	return &p
 }