forward.go

// Package forward implements a forwarding proxy. It caches an upstream net.Conn for some time, so if the same
// client returns the upstream's Conn will be precached. Depending on how you benchmark this looks to be
// 50% faster than just opening a new connection for every client. It works with UDP and TCP and uses
// inband healthchecking.
package forward

import (
	"context"
	"crypto/tls"
	"errors"
	"net"
	"sync/atomic"
	"time"

	"github.com/inverse-inc/wireguard-go/dns/plugin"
	"github.com/inverse-inc/wireguard-go/dns/plugin/debug"
	"github.com/inverse-inc/wireguard-go/dns/plugin/dnstap"
	clog "github.com/inverse-inc/wireguard-go/dns/plugin/pkg/log"
	"github.com/inverse-inc/wireguard-go/dns/request"

	"github.com/miekg/dns"
	ot "github.com/opentracing/opentracing-go"
)

var log = clog.NewWithPlugin("forward")

// Forward represents a plugin instance that can proxy requests to another (DNS) server. It has a list
// of proxies each representing one upstream proxy.
type Forward struct {
	concurrent int64 // atomic counters need to be first in struct for proper alignment

	proxies    []*Proxy
	p          Policy
	hcInterval time.Duration

	from    string
	ignored []string

	tlsConfig     *tls.Config
	tlsServerName string
	maxfails      uint32
	expire        time.Duration
	maxConcurrent int64

	sourceNetwork net.IPNet

	opts options // also here for testing

	// ErrLimitExceeded indicates that a query was rejected because the number of concurrent queries has exceeded
	// the maximum allowed (maxConcurrent)
	ErrLimitExceeded error

	tapPlugin *dnstap.Dnstap // when the dnstap plugin is loaded, we use to this to send messages out.

	// Next plugin.Handler
}

type Forwards struct {
	Forward []*Forward

	Next plugin.Handler
}

// New returns a new Forward.
func New() *Forward {
	f := &Forward{maxfails: 2, tlsConfig: new(tls.Config), expire: defaultExpire, p: new(random), from: ".", hcInterval: hcInterval, opts: options{forceTCP: false, preferUDP: false, hcRecursionDesired: true}, sourceNetwork: net.IPNet{IP: []byte{0, 0, 0, 0}, Mask: []byte{0, 0, 0, 0}}}
	return f
}

// SetProxy appends p to the proxy list and starts healthchecking.
func (f *Forward) SetProxy(p *Proxy) {
	f.proxies = append(f.proxies, p)
	p.start(f.hcInterval)
}

// Len returns the number of configured proxies.
func (f *Forward) Len() int { return len(f.proxies) }

// Name implements plugin.Handler.
func (fs Forwards) Name() string { return "forward" }

// ServeDNS implements plugin.Handler.
func (fs Forwards) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {

	for _, f := range fs.Forward {
		state := request.Request{W: w, Req: r}
		if !f.match(state) {
			continue
		}
		if f.maxConcurrent > 0 {
			count := atomic.AddInt64(&(f.concurrent), 1)
			defer atomic.AddInt64(&(f.concurrent), -1)
			if count > f.maxConcurrent {
				MaxConcurrentRejectCount.Add(1)
				return dns.RcodeServerFailure, f.ErrLimitExceeded
			}
		}

		fails := 0
		var span, child ot.Span
		var upstreamErr error
		span = ot.SpanFromContext(ctx)
		i := 0
		list := f.List()
		deadline := time.Now().Add(defaultTimeout)
		start := time.Now()
		for time.Now().Before(deadline) {
			if i >= len(list) {
				// reached the end of list, reset to begin
				i = 0
				fails = 0
			}

			proxy := list[i]
			i++
			if proxy.Down(f.maxfails) {
				fails++
				if fails < len(f.proxies) {
					continue
				}
				// All upstream proxies are dead, assume healthcheck is completely broken and randomly
				// select an upstream to connect to.
				r := new(random)
				proxy = r.List(f.proxies)[0]

				HealthcheckBrokenCount.Add(1)
			}

			if span != nil {
				child = span.Tracer().StartSpan("connect", ot.ChildOf(span.Context()))
				ctx = ot.ContextWithSpan(ctx, child)
			}

			var (
				ret *dns.Msg
				err error
			)
			opts := f.opts
			for {
				ret, err = proxy.Connect(ctx, state, opts)
				if err == ErrCachedClosed { // Remote side closed conn, can only happen with TCP.
					continue
				}
				// Retry with TCP if truncated and prefer_udp configured.
				if ret != nil && ret.Truncated && !opts.forceTCP && opts.preferUDP {
					opts.forceTCP = true
					continue
				}
				break
			}

			if child != nil {
				child.Finish()
			}

			if f.tapPlugin != nil {
				toDnstap(f, proxy.addr, state, opts, ret, start)
			}

			upstreamErr = err

			if err != nil {
				// Kick off health check to see if *our* upstream is broken.
				if f.maxfails != 0 {
					proxy.Healthcheck()
				}

				if fails < len(f.proxies) {
					continue
				}
				break
			}

			// Check if the reply is correct; if not return FormErr.
			if !state.Match(ret) {
				debug.Hexdumpf(ret, "Wrong reply for id: %d, %s %d", ret.Id, state.QName(), state.QType())

				formerr := new(dns.Msg)
				formerr.SetRcode(state.Req, dns.RcodeFormatError)
				w.WriteMsg(formerr)
				return 0, nil
			}

			w.WriteMsg(ret)
			return 0, nil
		}

		if upstreamErr != nil {
			return dns.RcodeServerFailure, upstreamErr
		}

		return dns.RcodeServerFailure, ErrNoHealthy
	}
	return plugin.NextOrFailure(fs.Name(), fs.Next, ctx, w, r)
}

func (f *Forward) match(state request.Request) bool {
	if !plugin.Name(f.from).Matches(state.Name()) || !f.isAllowedDomain(state.Name()) {
		return false
	}

	network := f.sourceNetwork

	if network.Contains(net.ParseIP(state.IP())) {
		return true
	}
	return false
}

func (f *Forward) isAllowedDomain(name string) bool {
	if dns.Name(name) == dns.Name(f.from) {
		return true
	}

	for _, ignore := range f.ignored {
		if plugin.Name(ignore).Matches(name) {
			return false
		}
	}
	return true
}

// ForceTCP returns if TCP is forced to be used even when the request comes in over UDP.
func (f *Forward) ForceTCP() bool { return f.opts.forceTCP }

// PreferUDP returns if UDP is preferred to be used even when the request comes in over TCP.
func (f *Forward) PreferUDP() bool { return f.opts.preferUDP }

// List returns a set of proxies to be used for this client depending on the policy in f.
func (f *Forward) List() []*Proxy { return f.p.List(f.proxies) }

var (
	// ErrNoHealthy means no healthy proxies left.
	ErrNoHealthy = errors.New("no healthy proxies")
	// ErrNoForward means no forwarder defined.
	ErrNoForward = errors.New("no forwarder defined")
	// ErrCachedClosed means cached connection was closed by peer.
	ErrCachedClosed = errors.New("cached connection was closed by peer")
)

// options holds various options that can be set.
type options struct {
	forceTCP           bool
	preferUDP          bool
	hcRecursionDesired bool
}

const defaultTimeout = 5 * time.Second