New fastest-tcp group to probe TCP connections (#155)

* Implement fastest-tcp group

* Add success-ttl option

* docs

* Change option to success-ttl-min

* Fix option in doc

cmd/routedns/config.go

@@ -111,6 +111,11 @@ type group struct {
 	Prefix6       uint8  // Prefix bits to identify IPv6 client
 	LimitResolver string `toml:"limit-resolver"` // Resolver to use when rate-limit exceeded
 
+	// Fastest-TCP probe options
+	Port          int
+	WaitAll       bool   `toml:"wait-all"`        // Wait for all probes to return and respond with a sorted list. Generally slower
+	SuccessTTLMin uint32 `toml:"success-ttl-min"` // Set the TTL of records that were probed successfully
+
 	// Response Collapse options
 	NullRCode int `toml:"null-rcode"` // Response code if after collapsing, no answers are left
 }

cmd/routedns/example-config/fastest-tcp.toml

@@ -0,0 +1,21 @@
+# Uses TCP connection probes to determine which of the response IPs
+# is fastests. Only the fastest IP is then cached.
+
+[listeners.local-udp]
+address = "127.0.0.1:53"
+protocol = "udp"
+resolver = "fastest-cached"
+
+[groups.fastest-cached]
+type = "cache"
+resolvers = ["tcp-probe"]
+
+[groups.tcp-probe]
+type = "fastest-tcp"
+port = 443
+success-ttl-min = 1800 # Cache successful lookups for a min of 30min
+resolvers = ["cloudflare-dot"]
+
+[resolvers.cloudflare-dot]
+address = "1.1.1.1:853"
+protocol = "dot"

cmd/routedns/main.go

@@ -390,6 +390,16 @@ func instantiateGroup(id string, g group, resolvers map[string]rdns.Resolver) er
 			MaxTTL: g.TTLMax,
 		}
 		resolvers[id] = rdns.NewTTLModifier(id, gr[0], opt)
+	case "fastest-tcp":
+		if len(gr) != 1 {
+			return fmt.Errorf("type fastest-tcp only supports one resolver in '%s'", id)
+		}
+		opt := rdns.FastestTCPOptions{
+			Port:          g.Port,
+			WaitAll:       g.WaitAll,
+			SuccessTTLMin: g.SuccessTTLMin,
+		}
+		resolvers[id] = rdns.NewFastestTCP(id, gr[0], opt)
 	case "ecs-modifier":
 		if len(gr) != 1 {
 			return fmt.Errorf("type ecs-modifier only supports one resolver in '%s'", id)

doc/configuration.md

@@ -32,6 +32,7 @@
   - [Response Collapse](#Response-Collapse)
   - [Router](#Router)
   - [Rate Limiter](#Rate-Limiter)
+  - [Fastest TCP Probe](#Fastest TCP Probe)
 - [Resolvers](#Resolvers)
   - [Plain DNS](#Plain-DNS-Resolver)
   - [DNS-over-TLS](#DNS-over-TLS-Resolver)
@@ -1124,6 +1125,39 @@ rcode = 5 # REFUSED
 
 Example config files: [rate-limiter.toml](../cmd/routedns/example-config/rate-limiter.toml)
 
+### Fastest TCP Probe
+
+The `fastest-tcp` element will first perform a lookup, then send TCP probes to all A or AAAA records in the response. It can then either return just the A/AAAA record for the fastest response, or all A/AAAA sorted by response time (fastest first). Since probing multiple servers can be slow, it is typically used behind a [cache](#Cache) to avoid making too many probes repeatedly. Each instance can only probe one port and if different ports are to be probed depending on the query name, a router should be used in front of it as well.
+
+#### Configuration
+
+A Fastest TCP Probe element is instantiated with `type = "fastest-tcp"` in the groups section of the configuration.
+
+Options:
+
+- `resolvers` - Array of upstream resolvers, only one is supported.
+- `port` - TCP port number to probe. Default: `443`.
+- `wait-all` - Instead of just returning the fastest response, wait for all probes and return them sorted by response time (fastest first). This will generally be slower as the slowest TCP probe determines the query response time. Default: `false`
+- `success-ttl-min` - Minimum TTL of successful probes (in seconds). Default: 0. Similar to the `ttl-min` option of [TTL Modifier](#TTL-modifier). Typically used to cache the response for longer given how resource-intensive and slow probing can be.
+
+Examples:
+
+TCP probe for the HTTPS port. Successful probes are cached for 30min.
+
+```toml
+[groups.fastest-cached]
+type = "cache"
+resolvers = ["tcp-probe"]
+
+[groups.tcp-probe]
+type = "fastest-tcp"
+port = 443
+success-ttl-min = 1800
+resolvers = ["cloudflare-dot"]
+```
+
+Example config files: [fastest-tcp.toml](../cmd/routedns/example-config/fastest-tcp.toml)
+
 ## Resolvers
 
 Resolvers forward queries to other DNS servers over the network and typically represent the end of one or many processing pipelines. Resolvers encode every query that is passed from listeners, modifiers, routers etc and send them to a DNS server without further processing. Like with other elements in the pipeline, resolvers requires a unique identifier to reference them from other elements. The following protocols are supported:

fastest-tcp.go

@@ -0,0 +1,193 @@
+package rdns
+
+import (
+	"context"
+	"errors"
+	"net"
+	"strconv"
+	"time"
+
+	"github.com/miekg/dns"
+	"github.com/sirupsen/logrus"
+)
+
+// FastestTCP first resolves the query with the upstream resolver, then
+// performs TCP connection tests with the response IPs to determine which
+// IP responds the fastest. This IP is then returned in the response.
+// This should be used in combination with a Cache to avoid the TCP
+// connection overhead on every query.
+type FastestTCP struct {
+	id       string
+	resolver Resolver
+	opt      FastestTCPOptions
+	port     string
+}
+
+var _ Resolver = &FastestTCP{}
+
+// FastestTCPOptions contain settings for a resolver that filters responses
+// based on TCP connection probes.
+type FastestTCPOptions struct {
+	// Port number to use for TCP probes, default 443
+	Port int
+
+	// Wait for all connection probes and sort the responses based on time
+	// (fastest first). This is generally slower than just waiting for the
+	// fastest, since the response time is determined by the slowest probe.
+	WaitAll bool
+
+	// TTL set on all RRs when TCP probing was successful. Can be used to
+	// ensure these are kept for longer in a cache and improve performance.
+	SuccessTTLMin uint32
+}
+
+// NewFastestTCP returns a new instance of a TCP probe resolver.
+func NewFastestTCP(id string, resolver Resolver, opt FastestTCPOptions) *FastestTCP {
+	port := strconv.Itoa(opt.Port)
+	if port == "0" {
+		port = "443"
+	}
+	return &FastestTCP{
+		id:       id,
+		resolver: resolver,
+		opt:      opt,
+		port:     port,
+	}
+}
+
+// Resolve a DNS query using a random resolver.
+func (r *FastestTCP) Resolve(q *dns.Msg, ci ClientInfo) (*dns.Msg, error) {
+	log := logger(r.id, q, ci)
+	a, err := r.resolver.Resolve(q, ci)
+	if err != nil {
+		return a, err
+	}
+	question := q.Question[0]
+
+	// Don't need to do anything if the query wasn't for an IP
+	if question.Qtype != dns.TypeA && question.Qtype != dns.TypeAAAA {
+		return a, nil
+	}
+
+	// Extract the IP responses
+	var ipRRs []dns.RR
+	for _, rr := range a.Answer {
+		if rr.Header().Rrtype == question.Qtype {
+			ipRRs = append(ipRRs, rr)
+		}
+	}
+
+	// If there's only one IP in the response, nothing to probe
+	if len(ipRRs) < 2 {
+		return a, nil
+	}
+
+	// Send TCP probes to all, if anything returns an error, just return
+	// the original response rather than trying to be clever and pick one.
+	log = log.WithField("port", r.port)
+	if r.opt.WaitAll {
+		rrs, err := r.probeAll(log, ipRRs)
+		if err != nil {
+			log.WithError(err).Debug("tcp probe failed")
+			return a, nil
+		}
+		r.setTTL(rrs...)
+		a.Answer = rrs
+		return a, nil
+	} else {
+		first, err := r.probeFastest(log, ipRRs)
+		if err != nil {
+			log.WithError(err).Debug("tcp probe failed")
+			return a, nil
+		}
+		r.setTTL(first)
+		a.Answer = []dns.RR{first}
+		return a, nil
+	}
+}
+
+// Sets the TTL of the given RRs if the option was provided
+func (r *FastestTCP) setTTL(rrs ...dns.RR) {
+	for _, rr := range rrs {
+		h := rr.Header()
+		if h.Ttl < r.opt.SuccessTTLMin {
+			h.Ttl = r.opt.SuccessTTLMin
+		}
+	}
+}
+
+func (r *FastestTCP) String() string {
+	return r.id
+}
+
+// Probes all IPs and returns only the RR with the fastest responding IP.
+// Waits for the first one that comes back. Returns an error if the fastest response
+// is an error.
+func (r *FastestTCP) probeFastest(log logrus.FieldLogger, rrs []dns.RR) (dns.RR, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+	resultCh := r.probe(ctx, log, rrs)
+	select {
+	case res := <-resultCh:
+		return res.rr, res.err
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	}
+}
+
+// Probes all IPs and returns them in the order of response time, fastest first. Returns
+// an error if any of the probes fail or if the probe times out.
+func (r *FastestTCP) probeAll(log logrus.FieldLogger, rrs []dns.RR) ([]dns.RR, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+	resultCh := r.probe(ctx, log, rrs)
+	results := make([]dns.RR, 0, len(rrs))
+	for i := 0; i < len(rrs); i++ {
+		select {
+		case res := <-resultCh:
+			if res.err != nil {
+				return nil, res.err
+			}
+			results = append(results, res.rr)
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		}
+	}
+	return results, nil
+}
+
+type tcpProbeResult struct {
+	rr  dns.RR
+	err error
+}
+
+// Probes all IPs and returns a channel with responses in the order they succeed or fail.
+func (r *FastestTCP) probe(ctx context.Context, log logrus.FieldLogger, rrs []dns.RR) <-chan tcpProbeResult {
+	resultCh := make(chan tcpProbeResult)
+	for _, rr := range rrs {
+		var d net.Dialer
+		go func(rr dns.RR) {
+			var network, ip string
+			switch record := rr.(type) {
+			case *dns.A:
+				network, ip = "tcp4", record.A.String()
+			case *dns.AAAA:
+				network, ip = "tcp6", record.AAAA.String()
+			default:
+				resultCh <- tcpProbeResult{err: errors.New("unexpected resource type")}
+				return
+			}
+			start := time.Now()
+			log.WithField("ip", ip).Debug("sending tcp probe")
+			c, err := d.DialContext(ctx, network, net.JoinHostPort(ip, r.port))
+			if err != nil {
+				resultCh <- tcpProbeResult{err: err}
+				return
+			}
+			log.WithField("ip", ip).WithField("response-time", time.Since(start)).Debug("tcp probe finished")
+			defer c.Close()
+			resultCh <- tcpProbeResult{rr: rr}
+		}(rr)
+	}
+	return resultCh
+}