go-dns/update.go

package godns

import (
	"context"
	"fmt"
	"time"

	"github.com/miekg/dns"
)

// Action is one RFC 2136 UPDATE operation.
type Action int

const (
	// ActionAdd inserts a record.
	ActionAdd Action = iota
	// ActionDelete removes a matching RR (name, type, and rdata).
	ActionDelete
)

// Change is a typed add or delete of one record.
type Change struct {
	Action Action
	Record Record
}

// Add builds an insert change.
func Add(rec Record) Change {
	return Change{Action: ActionAdd, Record: rec}
}

// Delete builds a delete-by-rr change (name, type, and rdata).
func Delete(rec Record) Change {
	return Change{Action: ActionDelete, Record: rec}
}

// UpdateResult is the nameserver's response to an UPDATE.
type UpdateResult struct {
	Rcode string
}

// Update applies RFC 2136 changes. It tries UDP first and retries on TCP when
// the UDP exchange fails or the response is truncated.
func (c *Client) Update(ctx context.Context, changes ...Change) (*UpdateResult, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if len(changes) == 0 {
		return &UpdateResult{Rcode: dns.RcodeToString[dns.RcodeSuccess]}, nil
	}

	msg, err := c.updateMessage(changes)
	if err != nil {
		return nil, err
	}

	client := &dns.Client{
		Net:        "udp",
		TsigSecret: map[string]string{dns.Fqdn(c.tsig.Name): c.tsig.Secret},
		Timeout:    timeoutFrom(ctx),
	}
	udpMsg := msg.Copy()
	resp, _, err := client.ExchangeContext(ctx, udpMsg, c.server)
	if err != nil || (resp != nil && resp.Truncated) {
		client.Net = "tcp"
		resp, _, err = client.ExchangeContext(ctx, msg.Copy(), c.server)
	}
	if err != nil {
		return nil, fmt.Errorf("UPDATE to %s: %w", c.server, err)
	}
	if resp.Rcode != dns.RcodeSuccess {
		return nil, fmt.Errorf("UPDATE rcode %s", dns.RcodeToString[resp.Rcode])
	}
	return &UpdateResult{Rcode: dns.RcodeToString[resp.Rcode]}, nil
}

// updateMessage builds a TSIG-signed RFC 2136 UPDATE message.
func (c *Client) updateMessage(changes []Change) (*dns.Msg, error) {
	zone := dns.Fqdn(c.zone)
	msg := new(dns.Msg)
	msg.SetUpdate(zone)
	for _, change := range changes {
		rr, err := recordToRR(change.Record)
		if err != nil {
			return nil, err
		}
		switch change.Action {
		case ActionAdd:
			msg.Insert([]dns.RR{rr})
		case ActionDelete:
			msg.Remove([]dns.RR{rr})
		default:
			return nil, fmt.Errorf("unsupported update action %d", change.Action)
		}
	}
	msg.SetTsig(dns.Fqdn(c.tsig.Name), tsigAlgorithm(c.tsig.Algorithm), 300, time.Now().Unix())
	return msg, nil
}