chore: fix our dns server implementation

This PR does those things:
* No longer shuffles dns servers for each request.
* Sets a context timeout of 4.5 seconds.
* Correctly returns a proper error from the root layer.

Signed-off-by: Dmitriy Matrenichev <dmitry.matrenichev@siderolabs.com>

internal/pkg/dns/dns.go

@@ -10,7 +10,6 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"math/rand/v2"
 	"net"
 	"net/netip"
 	"slices"
@@ -50,11 +49,42 @@ func NewCache(next plugin.Handler, l *zap.Logger) *Cache {
 
 // ServeDNS implements [dns.Handler].
 func (c *Cache) ServeDNS(wr dns.ResponseWriter, msg *dns.Msg) {
-	_, err := c.cache.ServeDNS(context.Background(), request.NewScrubWriter(msg, wr), msg)
+	wr = request.NewScrubWriter(msg, wr)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 4500*time.Millisecond)
+	defer cancel()
+
+	code, err := c.cache.ServeDNS(ctx, wr, msg)
 	if err != nil {
 		// we should probably call newProxy.Healthcheck() if there are too many errors
 		c.logger.Warn("error serving dns request", zap.Error(err))
 	}
+
+	if clientWrite(code) {
+		return
+	}
+
+	// Something went wrong
+	state := request.Request{W: wr, Req: msg}
+
+	answer := new(dns.Msg)
+	answer.SetRcode(msg, code)
+	state.SizeAndDo(answer)
+
+	err = wr.WriteMsg(answer)
+	if err != nil {
+		c.logger.Warn("error writing dns response", zap.Error(err))
+	}
+}
+
+// clientWrite returns true if the response has been written to the client.
+func clientWrite(rcode int) bool {
+	switch rcode {
+	case dns.RcodeServerFailure, dns.RcodeRefused, dns.RcodeFormatError, dns.RcodeNotImplemented:
+		return false
+	default:
+		return true
+	}
 }
 
 // Handler is a dns proxy selector.
@@ -87,28 +117,16 @@ func (h *Handler) ServeDNS(ctx context.Context, wrt dns.ResponseWriter, msg *dns
 
 	h.logger.Debug("dns request", zap.Stringer("data", msg))
 
-	upstreams := slices.Clone(h.dests)
-
-	if len(upstreams) == 0 {
-		emptyProxyErr := new(dns.Msg).SetRcode(req.Req, dns.RcodeServerFailure)
-
-		err := wrt.WriteMsg(emptyProxyErr)
-		if err != nil {
-			// We can't do much here, but at least log the error.
-			h.logger.Warn("failed to write 'no destination available' error dns response", zap.Error(err))
-		}
-
+	if len(h.dests) == 0 {
 		return dns.RcodeServerFailure, errors.New("no destination available")
 	}
 
-	rand.Shuffle(len(upstreams), func(i, j int) { upstreams[i], upstreams[j] = upstreams[j], upstreams[i] })
-
 	var (
 		resp *dns.Msg
 		err  error
 	)
 
-	for _, ups := range upstreams {
+	for _, ups := range h.dests {
 		opts := proxy.Options{}
 
 		for {
@@ -126,26 +144,20 @@ func (h *Handler) ServeDNS(ctx context.Context, wrt dns.ResponseWriter, msg *dns
 			break
 		}
 
-		if err == nil {
+		if ctx.Err() != nil || err == nil {
 			break
 		}
 
 		continue
 	}
 
-	if err != nil {
+	if ctx.Err() != nil {
+		return dns.RcodeServerFailure, ctx.Err()
+	} else if err != nil {
 		return dns.RcodeServerFailure, err
 	}
 
 	if !req.Match(resp) {
-		resp = new(dns.Msg).SetRcode(req.Req, dns.RcodeFormatError)
-
-		err = wrt.WriteMsg(resp)
-		if err != nil {
-			// We can't do much here, but at least log the error.
-			h.logger.Warn("failed to write non-matched response", zap.Error(err))
-		}
-
 		h.logger.Warn("dns response didn't match", zap.Stringer("data", resp))
 
 		return dns.RcodeFormatError, nil