From 2d106be341ae5b13d46712d0202dbfa7ff424753 Mon Sep 17 00:00:00 2001 From: rpb-ant Date: Thu, 9 Jul 2026 13:15:13 -0400 Subject: [PATCH] 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 * 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 * 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 --------- Signed-off-by: Ryan Brewster --- core/dnsserver/config.go | 10 ++++++ core/dnsserver/register.go | 4 +++ core/dnsserver/server.go | 15 ++++++++- core/dnsserver/server_test.go | 59 +++++++++++++++++++++++++++++++++++ 4 files changed, 87 insertions(+), 1 deletion(-) diff --git a/core/dnsserver/config.go b/core/dnsserver/config.go index c46b2fb77..130baf2ee 100644 --- a/core/dnsserver/config.go +++ b/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 diff --git a/core/dnsserver/register.go b/core/dnsserver/register.go index 55d311235..015bc1c0f 100644 --- a/core/dnsserver/register.go +++ b/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 } } diff --git a/core/dnsserver/server.go b/core/dnsserver/server.go index 942520ace..677455b83 100644 --- a/core/dnsserver/server.go +++ b/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() diff --git a/core/dnsserver/server_test.go b/core/dnsserver/server_test.go index 36395fcac..cfb50ed0c 100644 --- a/core/dnsserver/server_test.go +++ b/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") + } +}