dnsxl/xlist.go

package dnsxl

import (
	"context"
	"errors"
	"net"
	"time"
)

const (
	// DefaultTimeout is the maximum time allowed in XList.LookupIP. You can
	// use a different timeout by calling LookupIPCtx instead.
	DefaultTimeout = 5 * time.Second

	// DefaultWeight is the weight given to a server by AddServer when an IP
	// is found to be in the server's list.
	DefaultWeight = 1000
)

var (
	NoServersConfigured = errors.New("no servers configured")
)

// LookupFailed is returned from XList.LookupIP if all lookups fail.
type LookupFailed struct {
	Exemplar error
}

func (lf *LookupFailed) Error() string {
	return "all lookups failed; example error: " + lf.Exemplar.Error()
}

// XList implements a DNSxL (blocklist / allowlist) lookup. Create one with New
// and then call AddServers.
type XList struct {
	servers []server
	res     *net.Resolver
}

// New creates a fresh DNSxL lookup object, using the specified resolver. You
// may pass nil to use net.DefaultResolver.
func New(res *net.Resolver) *XList {
	xl := &XList{
		res: res,
	}
	if xl.res == nil {
		xl.res = net.DefaultResolver
	}
	return xl
}

// AddServers adds DNSxL servers to the list using the default weights.
func (xl *XList) AddServers(names ...string) {
	xl.AddWeightedServers(DefaultWeight, 0, names...)
}

// AddWeightedServers adds DNSxL servers to the list using the specified weights.
func (xl *XList) AddWeightedServers(weightInList, weightNotInList int, names ...string) {
	for _, name := range names {
		xl.servers = append(xl.servers, server{
			Addr:            name,
			WeightInList:    weightInList,
			WeightNotInList: weightNotInList,
		})
	}
}

// Results is the set of results of an IP lookup.
type Results []Result

// Weight returns the sum of each server's resulting weight.
func (r Results) Weight() int {
	var sum int
	for i := range r {
		sum += r[i].Weight
	}
	return sum
}

// LookupIP checks the configured servers and returns a set of results. It uses
// DefaultTimeout.
func (xl *XList) LookupIP(ip net.IP) (Results, error) {
	ctx, cancel := context.WithTimeout(context.Background(), DefaultTimeout)
	defer cancel()
	return xl.LookupIPCtx(ctx, ip)
}

// LookupIPCtx checks the given IP address at each of the configured servers,
// and returns a set of results. The individual server lookups are performed
// in parallel and may arrive in any order in the result set.
func (xl *XList) LookupIPCtx(ctx context.Context, ip net.IP) ([]Result, error) {
	if len(xl.servers) == 0 {
		return nil, NoServersConfigured
	}

	reversed := reverseName(ip)
	res := make([]Result, 0, len(xl.servers))
	resch := make(chan Result)

	for i := range xl.servers {
		go func(i int) {
			resch <- xl.servers[i].Lookup(ctx, xl.res, reversed)
		}(i)
	}

	var nerr int
	for _ = range xl.servers {
		r := <-resch
		if r.Error != nil {
			nerr++
		}
		res = append(res, r)
	}
	if nerr == len(xl.servers) {
		return res, &LookupFailed{
			Exemplar: res[0].Error,
		}
	}
	return res, nil
}