Pull request 421: AGDNS-3523-add-ratelimit

Merge in GO/dnsproxy from AGDNS-3523-add-ratelimit to master

Squashed commit of the following:

commit 5d51eb1
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Fri Feb 13 12:39:20 2026 +0700

    handler: rename to middleware

commit d281b05
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Fri Feb 13 12:35:08 2026 +0700

    handler: migrate middleware

commit bcb660f
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Thu Feb 12 10:55:26 2026 +0700

    proxy: imp code

commit 53fc44c
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Wed Feb 11 13:13:37 2026 +0700

    ratelimit: imp tests

commit ea3dd77
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Wed Feb 11 12:56:26 2026 +0700

    ratelimit: imp code

commit 4460e5e
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Wed Feb 11 12:41:27 2026 +0700

    all: imp

commit 5411a3f
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Tue Feb 10 10:20:16 2026 +0700

    all: ratelimit package

Author: Mizzick

internal/cmd/proxy.go

@@ -13,7 +13,7 @@ import (
 	"time"
 
 	"github.com/AdguardTeam/dnsproxy/internal/dnsmsg"
-	"github.com/AdguardTeam/dnsproxy/internal/handler"
+	"github.com/AdguardTeam/dnsproxy/internal/middleware"
 	proxynetutil "github.com/AdguardTeam/dnsproxy/internal/netutil"
 	"github.com/AdguardTeam/dnsproxy/proxy"
 	"github.com/AdguardTeam/dnsproxy/upstream"
@@ -39,13 +39,13 @@ func createProxyConfig(
 		return nil, err
 	}
 
-	hosts, err := handler.ReadHosts(ctx, l, hostsFiles)
+	hosts, err := middleware.ReadHosts(ctx, l, hostsFiles)
 	if err != nil {
 		return nil, fmt.Errorf("reading hosts files: %w", err)
 	}
 
-	reqHdlr := handler.NewDefault(&handler.DefaultConfig{
-		Logger: l.With(slogutil.KeyPrefix, "default_handler"),
+	preMw := middleware.New(&middleware.Config{
+		Logger: l.With(slogutil.KeyPrefix, "pre_handler_mw"),
 		// TODO(e.burkov):  Use the configured message constructor.
 		MessageConstructor: dnsmsg.DefaultMessageConstructor{},
 		HaltIPv6:           conf.IPv6Disabled,
@@ -81,7 +81,7 @@ func createProxyConfig(
 		MaxGoroutines:          conf.MaxGoRoutines,
 		UsePrivateRDNS:         conf.UsePrivateRDNS,
 		PrivateSubnets:         netutil.SubnetSetFunc(netutil.IsLocallyServed),
-		RequestHandler:         reqHdlr,
+		RequestHandler:         preMw.Wrap(proxy.DefaultHandler{}),
 		PendingRequests: &proxy.PendingRequestsConfig{
 			Enabled: conf.PendingRequestsEnabled,
 		},

internal/handler/default.go

@@ -1,4 +1,4 @@
-package handler
+package middleware
 
 import (
 	"net/netip"

internal/handler/handler.go

@@ -1,4 +1,4 @@
-package handler
+package middleware
 
 import (
 	"context"
@@ -91,43 +91,46 @@ func readHostsFile(ctx context.Context, strg *hostsfile.DefaultStorage, path str
 
 // resolveFromHosts resolves the DNS query from the hosts file.  It fills the
 // response with the A, AAAA, and PTR records from the hosts file.
-func (h *Default) resolveFromHosts(ctx context.Context, req *dns.Msg) (resp *dns.Msg) {
+func (mw *Default) resolveFromHosts(
+	ctx context.Context,
+	req *dns.Msg,
+) (resp *dns.Msg) {
 	var addrs []netip.Addr
 	var ptrs []string
 
 	q := req.Question[0]
 	name := strings.TrimSuffix(q.Name, ".")
 	switch q.Qtype {
 	case dns.TypeA:
-		addrs = slices.Clone(h.hosts.ByName(name))
+		addrs = slices.Clone(mw.hosts.ByName(name))
 		addrs = slices.DeleteFunc(addrs, netip.Addr.Is6)
 	case dns.TypeAAAA:
-		addrs = slices.Clone(h.hosts.ByName(name))
+		addrs = slices.Clone(mw.hosts.ByName(name))
 		addrs = slices.DeleteFunc(addrs, netip.Addr.Is4)
 	case dns.TypePTR:
 		addr, err := netutil.IPFromReversedAddr(name)
 		if err != nil {
-			h.logger.DebugContext(ctx, "failed parsing ptr", slogutil.KeyError, err)
+			mw.logger.DebugContext(ctx, "failed parsing ptr", slogutil.KeyError, err)
 
 			return nil
 		}
 
-		ptrs = h.hosts.ByAddr(addr)
+		ptrs = mw.hosts.ByAddr(addr)
 	default:
 		return nil
 	}
 
 	switch {
 	case len(addrs) > 0:
-		resp = h.messages.NewIPResponse(req, addrs)
+		resp = mw.messages.NewIPResponse(req, addrs)
 	case len(ptrs) > 0:
-		resp = h.messages.NewCompressedResponse(req, dns.RcodeSuccess)
+		resp = mw.messages.NewCompressedResponse(req, dns.RcodeSuccess)
 		name = req.Question[0].Name
 		for _, ptr := range ptrs {
-			resp.Answer = append(resp.Answer, h.messages.NewPTRAnswer(name, dns.Fqdn(ptr)))
+			resp.Answer = append(resp.Answer, mw.messages.NewPTRAnswer(name, dns.Fqdn(ptr)))
 		}
 	default:
-		h.logger.DebugContext(ctx, "no hosts records found", "name", name, "qtype", q.Qtype)
+		mw.logger.DebugContext(ctx, "no hosts records found", "name", name, "qtype", q.Qtype)
 	}
 
 	return resp

internal/handler/constructor.go internal/middleware/constructor.gointernal/handler/constructor.go renamed to internal/middleware/constructor.go

@@ -1,4 +1,4 @@
-package handler
+package middleware
 
 import (
 	"context"
@@ -8,15 +8,15 @@ import (
 
 // haltAAAA halts the processing of AAAA requests if IPv6 is disabled.  req must
 // not be nil.
-func (h *Default) haltAAAA(ctx context.Context, req *dns.Msg) (resp *dns.Msg) {
-	if h.isIPv6Halted && req.Question[0].Qtype == dns.TypeAAAA {
-		h.logger.DebugContext(
+func (mw *Default) haltAAAA(ctx context.Context, req *dns.Msg) (resp *dns.Msg) {
+	if mw.haltIPv6 && req.Question[0].Qtype == dns.TypeAAAA {
+		mw.logger.DebugContext(
 			ctx,
 			"ipv6 is disabled; replying with empty response",
 			"req", req.Question[0].Name,
 		)
 
-		return h.messages.NewMsgNODATA(req)
+		return mw.messages.NewMsgNODATA(req)
 	}
 
 	return nil

internal/handler/hosts.go internal/middleware/hosts.gointernal/handler/hosts.go renamed to internal/middleware/hosts.go

@@ -0,0 +1,80 @@
+// Package middleware provides some customizable DNS request handling logic used
+// in the proxy.
+package middleware
+
+import (
+	"context"
+	"log/slog"
+
+	"github.com/AdguardTeam/dnsproxy/proxy"
+	"github.com/AdguardTeam/golibs/hostsfile"
+)
+
+// Config is the configuration for [Default].
+type Config struct {
+	// HostsFiles is the index containing the records of the hosts files.  It
+	// must not be nil.
+	HostsFiles hostsfile.Storage
+
+	// Logger is the logger.  It must not be nil.
+	Logger *slog.Logger
+
+	// MessageConstructor constructs DNS messages.  It must not be nil.
+	MessageConstructor proxy.MessageConstructor
+
+	// HaltIPv6 halts the processing of AAAA requests and makes the handler
+	// reply with NODATA to them, if true.
+	HaltIPv6 bool
+}
+
+// Default implements [proxy.Middleware] with default DNS request handling
+// logic.
+type Default struct {
+	hosts    hostsfile.Storage
+	logger   *slog.Logger
+	messages messageConstructor
+	haltIPv6 bool
+}
+
+// New creates a new [*Default].
+func New(conf *Config) (mw *Default) {
+	mc, ok := conf.MessageConstructor.(messageConstructor)
+	if !ok {
+		mc = defaultConstructor{
+			MessageConstructor: conf.MessageConstructor,
+		}
+	}
+
+	return &Default{
+		hosts:    conf.HostsFiles,
+		logger:   conf.Logger,
+		messages: mc,
+		haltIPv6: conf.HaltIPv6,
+	}
+}
+
+// type check
+var _ proxy.Middleware = (*Default)(nil)
+
+// Wrap implements the [proxy.Middleware] interface for *Default.  It validates
+// and resolves the DNS request within proxyCtx.  It only calls h if the request
+// isn't handled by any of the internal handlers.
+func (mw *Default) Wrap(h proxy.Handler) (wrapped proxy.Handler) {
+	f := func(p *proxy.Proxy, proxyCtx *proxy.DNSContext) (err error) {
+		ctx := context.TODO()
+
+		mw.logger.DebugContext(ctx, "handling request", "req", &proxyCtx.Req.Question[0])
+
+		if proxyCtx.Res = mw.haltAAAA(ctx, proxyCtx.Req); proxyCtx.Res != nil {
+			return nil
+		}
+
+		if proxyCtx.Res = mw.resolveFromHosts(ctx, proxyCtx.Req); proxyCtx.Res != nil {
+			return nil
+		}
+
+		return h.ServeDNS(p, proxyCtx)
+	}
+
+	return proxy.HandlerFunc(f)
+}

internal/handler/ipv6halt.go internal/middleware/ipv6halt.gointernal/handler/ipv6halt.go renamed to internal/middleware/ipv6halt.go

@@ -1,4 +1,4 @@
-package handler
+package middleware
 
 import (
 	"net"
@@ -52,23 +52,23 @@ func TestDefault_haltAAAA(t *testing.T) {
 	t.Run("disabled", func(t *testing.T) {
 		t.Parallel()
 
-		hdlr := NewDefault(&DefaultConfig{
-			Logger:             slogutil.NewDiscardLogger(),
+		mw := New(&Config{
+			Logger:             testLogger,
 			MessageConstructor: messages,
 			HaltIPv6:           false,
 		})
 
 		ctx := testutil.ContextWithTimeout(t, defaultTimeout)
 
-		assert.Nil(t, hdlr.haltAAAA(ctx, reqA))
-		assert.Nil(t, hdlr.haltAAAA(ctx, reqAAAA))
+		assert.Nil(t, mw.haltAAAA(ctx, reqA))
+		assert.Nil(t, mw.haltAAAA(ctx, reqAAAA))
 	})
 
 	t.Run("enabled", func(t *testing.T) {
 		t.Parallel()
 
-		hdlr := NewDefault(&DefaultConfig{
-			Logger:             slogutil.NewDiscardLogger(),
+		hdlr := New(&Config{
+			Logger:             testLogger,
 			MessageConstructor: messages,
 			HaltIPv6:           true,
 		})
@@ -94,9 +94,9 @@ func TestDefault_resolveFromHosts(t *testing.T) {
 	strg, err := ReadHosts(ctx, testLogger, []string{absPath, relPath})
 	require.NoError(t, err)
 
-	hdlr := NewDefault(&DefaultConfig{
+	hdlr := New(&Config{
 		MessageConstructor: messages,
-		Logger:             slogutil.NewDiscardLogger(),
+		Logger:             testLogger,
 		HostsFiles:         strg,
 		HaltIPv6:           true,
 	})

internal/middleware/middleware.go

@@ -61,8 +61,9 @@ type Config struct {
 	BeforeRequestHandler BeforeRequestHandler
 
 	// RequestHandler is an optional custom handler for DNS requests.  It's used
-	// instead of DefaultRequestHandler if set.
-	RequestHandler RequestHandler
+	// instead of DefaultHandler if set.  In case of [ErrDrop] error returned
+	// from this handler, the proxy will not send any response to the client.
+	RequestHandler Handler
 
 	// UpstreamConfig is a general set of DNS servers to forward requests to.
 	UpstreamConfig *UpstreamConfig