DDI/agent/role/dns/soa.go

// SPDX-License-Identifier: Apache-2.0

package dns

import (
	"context"
	"crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"strings"
	"time"

	"golang.org/x/net/dns/dnsmessage"
)

// queryTimeout bounds one SOA query (UDP and, if truncated, TCP).
const queryTimeout = 3 * time.Second

type dialFunc func(ctx context.Context, network, address string) (net.Conn, error)

func defaultDial(ctx context.Context, network, address string) (net.Conn, error) {
	var d net.Dialer
	return d.DialContext(ctx, network, address)
}

// querySOA asks server for the SOA record of zone and returns its serial.
// The answer must be authoritative, so that a zone BIND failed to load
// (and that might be answered from elsewhere) does not count. A truncated
// UDP answer is repeated over TCP.
func querySOA(ctx context.Context, dial dialFunc, server, zone string) (uint32, error) {
	ctx, cancel := context.WithTimeout(ctx, queryTimeout)
	defer cancel()
	name, err := dnsmessage.NewName(strings.TrimSuffix(zone, ".") + ".")
	if err != nil {
		return 0, fmt.Errorf("invalid zone name %q: %w", zone, err)
	}
	var idb [2]byte
	if _, err := rand.Read(idb[:]); err != nil {
		return 0, err
	}
	id := binary.BigEndian.Uint16(idb[:])
	b := dnsmessage.NewBuilder(nil, dnsmessage.Header{ID: id})
	b.EnableCompression()
	if err := b.StartQuestions(); err != nil {
		return 0, err
	}
	q := dnsmessage.Question{Name: name, Type: dnsmessage.TypeSOA, Class: dnsmessage.ClassINET}
	if err := b.Question(q); err != nil {
		return 0, err
	}
	msg, err := b.Finish()
	if err != nil {
		return 0, err
	}

	resp, err := exchangeUDP(ctx, dial, server, msg, id)
	if err != nil {
		return 0, err
	}
	serial, truncated, err := parseSOA(resp, id, q)
	if truncated {
		if resp, err = exchangeTCP(ctx, dial, server, msg); err != nil {
			return 0, err
		}
		serial, _, err = parseSOA(resp, id, q)
	}
	return serial, err
}

func exchangeUDP(ctx context.Context, dial dialFunc, server string, msg []byte, id uint16) ([]byte, error) {
	conn, err := dial(ctx, "udp", server)
	if err != nil {
		return nil, err
	}
	defer func() { _ = conn.Close() }()
	if dl, ok := ctx.Deadline(); ok {
		_ = conn.SetDeadline(dl)
	}
	if _, err := conn.Write(msg); err != nil {
		return nil, err
	}
	buf := make([]byte, 65535)
	for {
		n, err := conn.Read(buf)
		if err != nil {
			return nil, err
		}
		// Ignore stray datagrams with another ID.
		if n >= 2 && binary.BigEndian.Uint16(buf[:2]) == id {
			return buf[:n], nil
		}
	}
}

func exchangeTCP(ctx context.Context, dial dialFunc, server string, msg []byte) ([]byte, error) {
	conn, err := dial(ctx, "tcp", server)
	if err != nil {
		return nil, err
	}
	defer func() { _ = conn.Close() }()
	if dl, ok := ctx.Deadline(); ok {
		_ = conn.SetDeadline(dl)
	}
	framed := make([]byte, 2+len(msg))
	binary.BigEndian.PutUint16(framed, uint16(len(msg))) //nolint:gosec // a query is far below 64 KiB
	copy(framed[2:], msg)
	if _, err := conn.Write(framed); err != nil {
		return nil, err
	}
	var l [2]byte
	if _, err := io.ReadFull(conn, l[:]); err != nil {
		return nil, err
	}
	resp := make([]byte, binary.BigEndian.Uint16(l[:]))
	if _, err := io.ReadFull(conn, resp); err != nil {
		return nil, err
	}
	return resp, nil
}

// parseSOA extracts the SOA serial for q from a response. truncated is
// true if the server set the TC bit (the caller retries over TCP).
func parseSOA(resp []byte, id uint16, q dnsmessage.Question) (serial uint32, truncated bool, err error) {
	var p dnsmessage.Parser
	h, err := p.Start(resp)
	if err != nil {
		return 0, false, fmt.Errorf("invalid DNS response: %w", err)
	}
	switch {
	case h.ID != id || !h.Response:
		return 0, false, errors.New("unexpected DNS response")
	case h.Truncated:
		return 0, true, nil
	case h.RCode != dnsmessage.RCodeSuccess:
		return 0, false, fmt.Errorf("server answered %s", rcodeString(h.RCode))
	case !h.Authoritative:
		return 0, false, errors.New("answer is not authoritative (zone not loaded?)")
	}
	qs, err := p.AllQuestions()
	if err != nil {
		return 0, false, fmt.Errorf("invalid DNS response: %w", err)
	}
	if len(qs) != 1 || qs[0].Type != q.Type || !strings.EqualFold(qs[0].Name.String(), q.Name.String()) {
		return 0, false, errors.New("DNS response is for another question")
	}
	for {
		ah, err := p.AnswerHeader()
		if errors.Is(err, dnsmessage.ErrSectionDone) {
			return 0, false, errors.New("no SOA record in the answer")
		}
		if err != nil {
			return 0, false, fmt.Errorf("invalid DNS response: %w", err)
		}
		if ah.Type != dnsmessage.TypeSOA || !strings.EqualFold(ah.Name.String(), q.Name.String()) {
			if err := p.SkipAnswer(); err != nil {
				return 0, false, fmt.Errorf("invalid DNS response: %w", err)
			}
			continue
		}
		soa, err := p.SOAResource()
		if err != nil {
			return 0, false, fmt.Errorf("invalid SOA record: %w", err)
		}
		return soa.Serial, false, nil
	}
}

func rcodeString(rc dnsmessage.RCode) string {
	s := rc.String()
	return strings.TrimPrefix(s, "RCode")
}