coredns/core/dnsserver/quic.go

package dnsserver

import (
	"crypto/tls"
	"encoding/binary"
	"errors"
	"net"

	"github.com/miekg/dns"
	"github.com/quic-go/quic-go"
)

type DoQWriter struct {
	localAddr  net.Addr
	remoteAddr net.Addr
	stream     *quic.Stream
	conn       *quic.Conn
	Msg        *dns.Msg
	tsigStatus error
}

func (w *DoQWriter) Write(b []byte) (int, error) {
	if w.stream == nil {
		return 0, errors.New("stream is nil")
	}
	b = AddPrefix(b)
	return w.stream.Write(b)
}

func (w *DoQWriter) WriteMsg(m *dns.Msg) error {
	bytes, err := m.Pack()
	if err != nil {
		return err
	}

	_, err = w.Write(bytes)
	if err != nil {
		return err
	}

	return w.Close()
}

// Close sends the STREAM FIN signal.
// The server MUST send the response(s) on the same stream and MUST
// indicate, after the last response, through the STREAM FIN
// mechanism that no further data will be sent on that stream.
// See https://www.rfc-editor.org/info/rfc9250/#section-4.2-7
func (w *DoQWriter) Close() error {
	if w.stream == nil {
		return errors.New("stream is nil")
	}
	return w.stream.Close()
}

// AddPrefix adds a 2-byte prefix with the DNS message length.
func AddPrefix(b []byte) (m []byte) {
	m = make([]byte, 2+len(b))
	binary.BigEndian.PutUint16(m, uint16(len(b))) // #nosec G115 -- DNS message length fits in uint16
	copy(m[2:], b)

	return m
}

// These methods implement the dns.ResponseWriter interface from Go DNS.

func (w *DoQWriter) TsigStatus() error      { return w.tsigStatus }
func (w *DoQWriter) TsigTimersOnly(_b bool) {}
func (w *DoQWriter) Hijack()                {}
func (w *DoQWriter) LocalAddr() net.Addr    { return w.localAddr }
func (w *DoQWriter) RemoteAddr() net.Addr   { return w.remoteAddr }
func (w *DoQWriter) Network() string        { return "" }

// ConnectionState implements the dns.ConnectionStater interface, exposing the
// TLS state of the underlying QUIC connection (e.g. for plugins that need to
// read the SNI ServerName). Mirrors the DoT behavior already provided by the
// miekg/dns response writer.
func (w *DoQWriter) ConnectionState() *tls.ConnectionState {
	if w.conn == nil {
		return nil
	}
	state := w.conn.ConnectionState().TLS
	return &state
}