dnsproxy/upstream/plain.go

package upstream

import (
	"context"
	"fmt"
	"io"
	"log/slog"
	"net"
	"net/url"
	"time"

	"github.com/AdguardTeam/dnsproxy/internal/bootstrap"
	"github.com/AdguardTeam/golibs/errors"
	"github.com/AdguardTeam/golibs/logutil/slogutil"
	"github.com/miekg/dns"
)

// network is the semantic type alias of the network to pass to dialing
// functions.  It's either [networkUDP] or [networkTCP].  It may also be used as
// URL scheme for plain upstreams.
type network = string

const (
	// networkUDP is the UDP network.
	networkUDP network = "udp"

	// networkTCP is the TCP network.
	networkTCP network = "tcp"
)

// plainDNS implements the [Upstream] interface for the regular DNS protocol.
type plainDNS struct {
	// addr is the DNS server URL.  Scheme is always "udp" or "tcp".
	addr *url.URL

	// logger is used for exchange logging.  It is never nil.
	logger *slog.Logger

	// getDialer either returns an initialized dial handler or creates a new
	// one.
	getDialer DialerInitializer

	// net is the network of the connections.
	net network

	// timeout is the timeout for DNS requests.
	timeout time.Duration
}

// newPlain returns the plain DNS Upstream.  addr.Scheme should be either "udp"
// or "tcp".
func newPlain(addr *url.URL, opts *Options) (u *plainDNS, err error) {
	switch addr.Scheme {
	case networkUDP, networkTCP:
		// Go on.
	default:
		return nil, fmt.Errorf("unsupported url scheme: %s", addr.Scheme)
	}

	addPort(addr, defaultPortPlain)

	return &plainDNS{
		addr:      addr,
		logger:    opts.Logger,
		getDialer: newDialerInitializer(addr, opts),
		net:       addr.Scheme,
		timeout:   opts.Timeout,
	}, nil
}

// type check
var _ Upstream = &plainDNS{}

// Address implements the [Upstream] interface for *plainDNS.
func (p *plainDNS) Address() string {
	switch p.net {
	case networkUDP:
		return p.addr.Host
	case networkTCP:
		return p.addr.String()
	default:
		panic(fmt.Sprintf("unexpected network: %s", p.net))
	}
}

// dialExchange performs a DNS exchange with the specified dial handler.
// network must be either [networkUDP] or [networkTCP].
func (p *plainDNS) dialExchange(
	ctx context.Context,
	network network,
	dial bootstrap.DialHandler,
	req *dns.Msg,
) (resp *dns.Msg, err error) {
	var cancel context.CancelFunc
	if p.timeout > 0 {
		ctx, cancel = context.WithTimeout(ctx, p.timeout)
		defer cancel()
	}

	addr := p.Address()

	// NOTE: It is important to set a client timeout to prevent the use of the
	// default timeout settings.
	client := &dns.Client{Timeout: p.timeout}

	conn := &dns.Conn{}
	upstreamReq := setRequestForNetwork(req, conn, network)
	defer func() {
		if resp != nil {
			resp.Id = req.Id
		}
	}()

	logBegin(ctx, p.logger, addr, network, upstreamReq)
	defer func() { logFinish(ctx, p.logger, addr, network, err) }()

	// NOTE: Discard context timeout to prevent deadline sharing between dial
	// calls.
	//
	// TODO(f.setrakov): Consider supporting context timeout here.
	conn.Conn, err = dial(context.WithoutCancel(ctx), network, "")
	if err != nil {
		return nil, fmt.Errorf("dialing %s over %s: %w", p.addr.Host, network, err)
	}
	defer func(c net.Conn) { err = errors.WithDeferred(err, c.Close()) }(conn.Conn)

	resp, _, err = client.ExchangeWithConnContext(ctx, upstreamReq, conn)
	if isExpectedConnErr(err) {
		if p.timeout > 0 {
			ctx = context.WithoutCancel(ctx)
			ctx, cancel = context.WithTimeout(ctx, p.timeout)
			defer cancel()
		}

		conn.Conn, err = dial(context.WithoutCancel(ctx), network, "")
		if err != nil {
			return nil, fmt.Errorf("dialing %s over %s again: %w", p.addr.Host, network, err)
		}
		defer func(c net.Conn) { err = errors.WithDeferred(err, c.Close()) }(conn.Conn)

		resp, _, err = client.ExchangeWithConnContext(ctx, upstreamReq, conn)
	}

	if err != nil {
		return resp, fmt.Errorf("exchanging with %s over %s: %w", addr, network, err)
	}

	return resp, validateResponse(upstreamReq, resp)
}

// setRequestForNetwork sets connection options in conn and overrides the
// upstream request, if necessary, depending on network.  If network is
// [networkUDP] and orig has a zero ID, req is a copy of orig with a new ID to
// increase entropy.  network must be either [networkUDP] or [networkTCP].  orig
// and conn must not be nil.
func setRequestForNetwork(orig *dns.Msg, conn *dns.Conn, network network) (req *dns.Msg) {
	req = orig
	if network != networkUDP {
		return req
	}

	conn.UDPSize = dns.MinMsgSize

	if orig.Id == 0 {
		req = orig.Copy()
		req.Id = dns.Id()
	}

	return req
}

// isExpectedConnErr returns true if the error is expected.  In this case,
// we will make a second attempt to process the request.
func isExpectedConnErr(err error) (is bool) {
	var netErr net.Error

	return err != nil && (errors.As(err, &netErr) || errors.Is(err, io.EOF))
}

// Exchange implements the [Upstream] interface for *plainDNS.
func (p *plainDNS) Exchange(ctx context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
	dial, err := p.getDialer()
	if err != nil {
		// Don't wrap the error since it's informative enough as is.
		return nil, err
	}

	addr := p.Address()

	resp, err = p.dialExchange(ctx, p.net, dial, req)
	if p.net != networkUDP {
		// The network is already TCP.
		return resp, err
	}

	if resp == nil {
		// There is likely an error with the upstream.
		return resp, err
	}

	if errors.Is(err, errQuestion) {
		// The upstream responds with malformed messages, so try TCP.
		p.logger.DebugContext(
			ctx,
			"plain response is malformed, using tcp",
			"addr", addr,
			slogutil.KeyError, err,
		)

		return p.dialExchange(ctx, networkTCP, dial, req)
	} else if resp.Truncated {
		// Fallback to TCP on truncated responses.
		p.logger.DebugContext(
			ctx,
			"plain response is truncated, using tcp",
			"question", &req.Question[0],
			"addr", addr,
		)

		return p.dialExchange(ctx, networkTCP, dial, req)
	}

	// There is either no error or the error isn't related to the received
	// message.
	return resp, err
}

// Close implements the [Upstream] interface for *plainDNS.
func (p *plainDNS) Close() (err error) {
	return nil
}