Merge pull request #99 from mailgun/maxim/develop

PIP-1627: Add mxresolv package

mxresolv/mxresolv.go

@@ -0,0 +1,145 @@
+package mxresolv
+
+import (
+	"context"
+	"math/rand"
+	"net"
+	"sort"
+	"strings"
+	"unicode"
+
+	"github.com/mailgun/holster/v4/clock"
+	"github.com/mailgun/holster/v4/collections"
+	"github.com/pkg/errors"
+	"golang.org/x/net/idna"
+)
+
+const (
+	cacheSize = 1000
+	cacheTTL  = 10 * clock.Minute
+)
+
+var (
+	errNullMXRecord   = errors.New("domain accepts no mail")
+	errNoValidMXHosts = errors.New("no valid MX hosts")
+
+	lookupResultCache *collections.LRUCache
+
+	// It is modified only in tests to make them deterministic.
+	shuffle = true
+)
+
+func init() {
+	lookupResultCache = collections.NewLRUCache(cacheSize)
+}
+
+// Lookup performs a DNS lookup of MX records for the specified hostname. It
+// returns a prioritised list of MX hostnames, where hostnames with the same
+// priority are shuffled. If the second returned value is true, then the host
+// does not have explicit MX records, and its A record is returned instead.
+//
+// It uses an LRU cache with a timeout to reduce the number of network requests.
+func Lookup(ctx context.Context, hostname string) ([]string, bool, error) {
+	if cachedVal, ok := lookupResultCache.Get(hostname); ok {
+		lookupResult := cachedVal.(lookupResult)
+		return lookupResult.mxHosts, lookupResult.implicit, lookupResult.err
+	}
+	asciiHostname, err := ensureASCII(hostname)
+	if err != nil {
+		return nil, false, errors.Wrap(err, "invalid hostname")
+	}
+	mxRecords, err := net.DefaultResolver.LookupMX(ctx, asciiHostname)
+	if err != nil {
+		var timeouter interface{ Timeout() bool }
+		if errors.As(err, &timeouter) && timeouter.Timeout() {
+			return nil, false, errors.WithStack(err)
+		}
+		var netDNSError *net.DNSError
+		if errors.As(err, &netDNSError) && netDNSError.Err == "no such host" {
+			if _, err := net.DefaultResolver.LookupIPAddr(ctx, asciiHostname); err != nil {
+				return cacheAndReturn(hostname, nil, false, errors.WithStack(err))
+			}
+			return cacheAndReturn(hostname, []string{asciiHostname}, true, nil)
+		}
+		if mxRecords == nil {
+			return cacheAndReturn(hostname, nil, false, errors.WithStack(err))
+		}
+	}
+	// Check for "Null MX" record (https://tools.ietf.org/html/rfc7505).
+	if len(mxRecords) == 1 && mxRecords[0].Host == "." {
+		return cacheAndReturn(hostname, nil, false, errNullMXRecord)
+	}
+	// Normalize returned hostnames: drop trailing '.' and lowercase.
+	for _, mxRecord := range mxRecords {
+		lastCharIndex := len(mxRecord.Host) - 1
+		if mxRecord.Host[lastCharIndex] == '.' {
+			mxRecord.Host = strings.ToLower(mxRecord.Host[:lastCharIndex])
+		}
+	}
+	// Sort records in order of preference and lexicographically within a
+	// preference group. The latter is only to make tests deterministic.
+	sort.Slice(mxRecords, func(i, j int) bool {
+		return mxRecords[i].Pref < mxRecords[j].Pref ||
+			(mxRecords[i].Pref == mxRecords[j].Pref && mxRecords[i].Host < mxRecords[j].Host)
+	})
+	// Shuffle records within preference groups unless disabled in tests.
+	if shuffle {
+		mxRecordCount := len(mxRecords)
+		groupBegin := 0
+		for i := 1; i < mxRecordCount; i++ {
+			if mxRecords[i].Pref != mxRecords[groupBegin].Pref || i == mxRecordCount-1 {
+				groupSlice := mxRecords[groupBegin:i]
+				rand.Shuffle(len(groupSlice), func(i, j int) {
+					groupSlice[i], groupSlice[j] = groupSlice[j], groupSlice[i]
+				})
+				groupBegin = i
+			}
+		}
+	}
+	// Make a hostname list, but skip non-ASCII names, that cause issues.
+	mxHosts := make([]string, 0, len(mxRecords))
+	for _, mxRecord := range mxRecords {
+		if !isASCII(mxRecord.Host) {
+			continue
+		}
+		if mxRecord.Host == "" {
+			continue
+		}
+		mxHosts = append(mxHosts, mxRecord.Host)
+	}
+	if len(mxHosts) == 0 {
+		return cacheAndReturn(hostname, nil, false, errNoValidMXHosts)
+	}
+	return cacheAndReturn(hostname, mxHosts, false, nil)
+}
+
+func ensureASCII(hostname string) (string, error) {
+	if isASCII(hostname) {
+		return hostname, nil
+	}
+	hostname, err := idna.ToASCII(hostname)
+	if err != nil {
+		return "", errors.WithStack(err)
+	}
+	return hostname, nil
+}
+
+func isASCII(s string) bool {
+	for i := 0; i < len(s); i++ {
+		if s[i] > unicode.MaxASCII {
+			return false
+		}
+	}
+	return true
+}
+
+type lookupResult struct {
+	mxHosts  []string
+	implicit bool
+	err      error
+}
+
+func cacheAndReturn(hostname string, mxHosts []string, implicit bool, err error) ([]string, bool, error) {
+	lookupResultCache.AddWithTTL(hostname, lookupResult{mxHosts, implicit, err}, cacheTTL)
+	return mxHosts, implicit, err
+}

mxresolv/mxresolv_test.go

@@ -0,0 +1,157 @@
+package mxresolv
+
+import (
+	"context"
+	"fmt"
+	"regexp"
+	"sort"
+	"testing"
+
+	"github.com/mailgun/holster/v4/clock"
+	"github.com/mailgun/holster/v4/collections"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestLookup(t *testing.T) {
+	defer disableShuffle()()
+	for i, tc := range []struct {
+		desc          string
+		inDomainName  string
+		outMXHosts    []string
+		outImplicitMX bool
+	}{{
+		desc:         "MX record preference is respected",
+		inDomainName: "test-mx.definbox.com",
+		outMXHosts: []string{
+			/* 1 */ "mxa.definbox.com", "mxe.definbox.com", "mxi.definbox.com",
+			/* 2 */ "mxc.definbox.com",
+			/* 3 */ "mxb.definbox.com", "mxd.definbox.com", "mxf.definbox.com", "mxg.definbox.com", "mxh.definbox.com"},
+		outImplicitMX: false,
+	}, {
+		inDomainName:  "test-a.definbox.com",
+		outMXHosts:    []string{"test-a.definbox.com"},
+		outImplicitMX: true,
+	}, {
+		inDomainName:  "test-cname.definbox.com",
+		outMXHosts:    []string{"mxa.ninomail.com", "mxb.ninomail.com"},
+		outImplicitMX: false,
+	}, {
+		desc: "If an MX host returned by the resolver contains non ASCII " +
+			"characters then it is silently dropped from the returned list",
+		inDomainName:  "test-unicode.definbox.com",
+		outMXHosts:    []string{"mxa.definbox.com", "mxb.definbox.com"},
+		outImplicitMX: false,
+	}, {
+		desc:          "Underscore is allowed in domain names",
+		inDomainName:  "test-underscore.definbox.com",
+		outMXHosts:    []string{"foo_bar.definbox.com"},
+		outImplicitMX: false,
+	}, {
+		inDomainName:  "test-яндекс.definbox.com",
+		outMXHosts:    []string{"xn--test---mofb0ab4b8camvcmn8gxd.definbox.com"},
+		outImplicitMX: false,
+	}, {
+		inDomainName:  "xn--test--xweh4bya7b6j.definbox.com",
+		outMXHosts:    []string{"xn--test---mofb0ab4b8camvcmn8gxd.definbox.com"},
+		outImplicitMX: false,
+	}} {
+		fmt.Printf("Test case #%d: %s, %s\n", i, tc.inDomainName, tc.desc)
+		// When
+		ctx, cancel := context.WithTimeout(context.Background(), 3*clock.Second)
+		mxHosts, explictMX, err := Lookup(ctx, tc.inDomainName)
+		cancel()
+		// Then
+		assert.NoError(t, err)
+		assert.Equal(t, tc.outMXHosts, mxHosts)
+		assert.Equal(t, tc.outImplicitMX, explictMX)
+
+		// The second lookup returns the cached result, that only shows on the
+		// coverage report.
+		mxHosts, explictMX, err = Lookup(ctx, tc.inDomainName)
+		assert.NoError(t, err)
+		assert.Equal(t, tc.outMXHosts, mxHosts)
+		assert.Equal(t, tc.outImplicitMX, explictMX)
+	}
+}
+
+func TestLookupError(t *testing.T) {
+	defer disableShuffle()()
+	for i, tc := range []struct {
+		desc       string
+		inHostname string
+		outError   string
+	}{{
+		inHostname: "test-broken.definbox.com",
+		outError:   "lookup test-broken.definbox.com.*: no such host",
+	}, {
+		inHostname: "",
+		outError:   "lookup : no such host",
+	}, {
+		inHostname: "kaboom",
+		outError:   "lookup kaboom.*: no such host",
+	}, {
+		inHostname: "example.com",
+		outError:   "domain accepts no mail",
+	}} {
+		fmt.Printf("Test case #%d: %s, %s\n", i, tc.inHostname, tc.desc)
+		// When
+		ctx, cancel := context.WithTimeout(context.Background(), 3*clock.Second)
+		_, _, err := Lookup(ctx, tc.inHostname)
+		cancel()
+		// Then
+		assert.Regexp(t, regexp.MustCompile(tc.outError), err.Error())
+
+		// The second lookup returns the cached result, that only shows on the
+		// coverage report.
+		_, _, err = Lookup(ctx, tc.inHostname)
+		assert.Regexp(t, regexp.MustCompile(tc.outError), err.Error())
+	}
+}
+
+// Shuffling only does not cross preference group boundaries.
+//
+// Preference groups are:
+//  1: mxa.definbox.com, mxe.definbox.com, mxi.definbox.com
+//  2: mxc.definbox.com
+//  3: mxb.definbox.com, mxd.definbox.com, mxf.definbox.com, mxg.definbox.com, mxh.definbox.com
+//
+// Warning: since the data set is pretty small subsequent shuffles can produce
+// the same result causing the test to fail.
+func TestLookupShuffle(t *testing.T) {
+	// When
+	ctx, cancel := context.WithTimeout(context.Background(), 3*clock.Second)
+	defer cancel()
+	shuffle1, _, err := Lookup(ctx, "test-mx.definbox.com")
+	assert.NoError(t, err)
+	resetCache()
+	shuffle2, _, err := Lookup(ctx, "test-mx.definbox.com")
+	assert.NoError(t, err)
+
+	// Then
+	assert.NotEqual(t, shuffle1[:3], shuffle2[:3])
+	assert.NotEqual(t, shuffle1[4:], shuffle2[4:])
+
+	sort.Strings(shuffle1[:3])
+	sort.Strings(shuffle2[:3])
+	assert.Equal(t, []string{"mxa.definbox.com", "mxe.definbox.com", "mxi.definbox.com"}, shuffle1[:3])
+	assert.Equal(t, shuffle1[:3], shuffle2[:3])
+
+	assert.Equal(t, "mxc.definbox.com", shuffle1[3])
+	assert.Equal(t, shuffle1[3], shuffle2[3])
+
+	sort.Strings(shuffle1[4:])
+	sort.Strings(shuffle2[4:])
+	assert.Equal(t, []string{"mxb.definbox.com", "mxd.definbox.com", "mxf.definbox.com", "mxg.definbox.com", "mxh.definbox.com"}, shuffle1[4:])
+	assert.Equal(t, shuffle1[4:], shuffle2[4:])
+}
+
+func disableShuffle() func() {
+	shuffle = false
+	return func() {
+		shuffle = true
+	}
+}
+
+func resetCache() {
+	lookupResultCache = collections.NewLRUCache(1000)
+}