core/dnsserver: add Config.UDPDecorateWriterFunc for external plugins (#8257)

* core/dnsserver: add Config.UDPDecorateWriterFunc for external plugins

Allow external plugins to install a dns.DecorateWriter (an existing
miekg/dns Server field CoreDNS never sets) on plain-UDP servers, e.g.
for per-socket write disciplines or overload protection; UDP is the
only transport without a Max*-style concurrency setting.

The field is a factory called once per Server in ServePacket, so each
socket gets its own decorator under multisocket; the *Server argument
is the same value handlers see via ctx.Value(dnsserver.Key{}). Follows
the Config.HTTPRequestValidateFunc precedent, including server-block
propagation and the ServePacket-time lookup over s.zones. Nil (the
default) changes nothing.

Signed-off-by: Ryan Brewster <rpb@anthropic.com>

* core/dnsserver: select UDP writer decorator deterministically in NewServer

Review feedback: s.zones is a map, so choosing the decorator in
ServePacket was non-deterministic when several server blocks share a
listener. Select it in NewServer from the group slice instead (stable
config order, last one set wins) and document the resolution order.

Signed-off-by: Ryan Brewster <rpb@anthropic.com>

* core/dnsserver: group decorator selection with the other last-writer-wins fields

Pure move: place it next to ProxyProtoConnPolicy, the existing field
resolved the same way at the same point.

Signed-off-by: Ryan Brewster <rpb@anthropic.com>

---------

Signed-off-by: Ryan Brewster <rpb@anthropic.com>

core/dnsserver/config.go

@@ -11,6 +11,7 @@ import (
 	"github.com/coredns/coredns/plugin"
 	"github.com/coredns/coredns/request"
 
+	"github.com/miekg/dns"
 	"github.com/pires/go-proxyproto"
 )
 
@@ -49,6 +50,15 @@ type Config struct {
 	// may depend on it.
 	HTTPRequestValidateFunc func(*http.Request) bool
 
+	// If this function is not nil it is called once per Server in ServePacket
+	// (so each UDP listening socket gets its own decorator under multisocket)
+	// and its result is installed as the underlying dns.Server's
+	// DecorateWriter. Plain dns:// UDP listeners only. When several server
+	// blocks sharing a listener set it, the last one in config order wins.
+	// Although this isn't referenced in-tree, external plugins may depend
+	// on it.
+	UDPDecorateWriterFunc func(*Server) dns.DecorateWriter
+
 	// FilterFuncs is used to further filter access
 	// to this handler. E.g. to limit access to a reverse zone
 	// on a non-octet boundary, i.e. /17

core/dnsserver/register.go

@@ -279,6 +279,10 @@ func propagateConfigParams(configs []*Config) {
 		// Propagate HTTPRequestValidateFunc so that custom path validators work in
 		// multi-transport blocks. Otherwise HTTPS 404s on non-"/dns-query" paths.
 		c.HTTPRequestValidateFunc = c.firstConfigInBlock.HTTPRequestValidateFunc
+
+		// Propagate UDPDecorateWriterFunc so a decorator configured once in a
+		// server block applies to the block's UDP listener(s).
+		c.UDPDecorateWriterFunc = c.firstConfigInBlock.UDPDecorateWriterFunc
 	}
 }
 

core/dnsserver/server.go

@@ -55,6 +55,11 @@ type Server struct {
 
 	tsigSecret map[string]string
 
+	// udpDecorateWriterFunc is selected in NewServer from the group configs in
+	// stable order (last one set wins), so the choice is deterministic when
+	// several server blocks share a listener. See Config.UDPDecorateWriterFunc.
+	udpDecorateWriterFunc func(*Server) dns.DecorateWriter
+
 	// Ensure Stop is idempotent when invoked concurrently (e.g., during reload and SIGTERM).
 	stopOnce sync.Once
 	stopErr  error
@@ -138,6 +143,9 @@ func NewServer(addr string, group []*Config) (*Server, error) {
 		if site.ProxyProtoUDPSessionTrackingMaxSessions > 0 {
 			s.udpSessionTrackingMaxSessions = site.ProxyProtoUDPSessionTrackingMaxSessions
 		}
+		if site.UDPDecorateWriterFunc != nil {
+			s.udpDecorateWriterFunc = site.UDPDecorateWriterFunc
+		}
 	}
 
 	if !s.debug {
@@ -179,12 +187,17 @@ func (s *Server) Serve(l net.Listener) error {
 // ServePacket starts the server with an existing packetconn. It blocks until the server stops.
 // This implements caddy.UDPServer interface.
 func (s *Server) ServePacket(p net.PacketConn) error {
+	// Use a custom writer decorator if one was configured.
+	var dw dns.DecorateWriter
+	if s.udpDecorateWriterFunc != nil {
+		dw = s.udpDecorateWriterFunc(s)
+	}
 	s.m.Lock()
 	s.server[udp] = &dns.Server{PacketConn: p, Net: "udp", Handler: dns.HandlerFunc(func(w dns.ResponseWriter, r *dns.Msg) {
 		ctx := context.WithValue(context.Background(), Key{}, s)
 		ctx = context.WithValue(ctx, LoopKey{}, 0)
 		s.ServeDNS(ctx, w, r)
-	}), TsigSecret: s.tsigSecret}
+	}), TsigSecret: s.tsigSecret, DecorateWriter: dw}
 	s.m.Unlock()
 
 	return s.server[udp].ActivateAndServe()

core/dnsserver/server_test.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"net"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -180,3 +181,61 @@ func BenchmarkCoreServeDNS(b *testing.B) {
 		s.ServeDNS(ctx, w, m)
 	}
 }
+
+// recordingWriter counts the packed frames the decorated writer receives
+// before forwarding them to the real writer. The decorator mints a fresh
+// wrapper per packet; the frames counter is shared across them.
+type recordingWriter struct {
+	dns.Writer
+	frames *atomic.Int64
+}
+
+func (rw *recordingWriter) Write(b []byte) (int, error) {
+	rw.frames.Add(1)
+	return rw.Writer.Write(b)
+}
+
+func TestUDPDecorateWriterFunc(t *testing.T) {
+	cfg := testConfig("dns", test.ErrorHandler())
+
+	frames := new(atomic.Int64)
+	var gotServer atomic.Pointer[Server]
+	var calls atomic.Int64
+	cfg.UDPDecorateWriterFunc = func(srv *Server) dns.DecorateWriter {
+		calls.Add(1)
+		gotServer.Store(srv)
+		return func(w dns.Writer) dns.Writer {
+			return &recordingWriter{Writer: w, frames: frames}
+		}
+	}
+
+	s, err := NewServer("127.0.0.1:0", []*Config{cfg})
+	if err != nil {
+		t.Fatalf("NewServer failed: %v", err)
+	}
+
+	pc, err := net.ListenPacket("udp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("ListenPacket failed: %v", err)
+	}
+	defer pc.Close()
+
+	go s.ServePacket(pc)
+	defer s.Stop()
+
+	m := new(dns.Msg)
+	m.SetQuestion("example.com.", dns.TypeA)
+	if _, err := dns.Exchange(m, pc.LocalAddr().String()); err != nil {
+		t.Fatalf("dns.Exchange failed: %v", err)
+	}
+
+	if n := calls.Load(); n != 1 {
+		t.Errorf("expected UDPDecorateWriterFunc to be called once per socket, got %d", n)
+	}
+	if gotServer.Load() != s {
+		t.Errorf("expected UDPDecorateWriterFunc to receive the serving *Server")
+	}
+	if frames.Load() == 0 {
+		t.Errorf("expected the decorated writer to observe the response write")
+	}
+}