Merge pull request #173 from mailgun/thrawn/rand-mx-hosts

mxresolv.Lookup() now randomizes the mx host records when cached

go.mod

@@ -78,7 +78,7 @@ require (
 	go.uber.org/zap v1.21.0 // indirect
 	golang.org/x/mod v0.8.0 // indirect
 	golang.org/x/sys v0.8.0 // indirect
-	golang.org/x/text v0.8.0 // indirect
+	golang.org/x/text v0.9.0 // indirect
 	golang.org/x/tools v0.6.0 // indirect
 	google.golang.org/genproto v0.0.0-20230306155012-7f2fa6fef1f4 // indirect
 	google.golang.org/protobuf v1.30.0 // indirect

go.sum

@@ -599,8 +599,8 @@ golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.5/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
-golang.org/x/text v0.8.0 h1:57P1ETyNKtuIjB4SRd15iJxuhj8Gc416Y78H3qgMh68=
-golang.org/x/text v0.8.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
+golang.org/x/text v0.9.0 h1:2sjJmO8cDvYveuX97RDLsxlyUxLl+GHoLxBiRdHllBE=
+golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
 golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=

mxresolv/mxresolv.go

@@ -45,10 +45,14 @@ func init() {
 //
 // It uses an LRU cache with a timeout to reduce the number of network requests.
 func Lookup(ctx context.Context, hostname string) (retMxHosts []string, retImplicit bool, reterr error) {
-	if cachedVal, ok := lookupResultCache.Get(hostname); ok {
-		lookupResult := cachedVal.(lookupResult)
-		return lookupResult.mxHosts, lookupResult.implicit, lookupResult.err
+	if obj, ok := lookupResultCache.Get(hostname); ok {
+		cached := obj.(lookupResult)
+		if len(cached.mxRecords) != 0 {
+			return shuffleMXRecords(cached.mxRecords), cached.implicit, cached.err
+		}
+		return cached.mxHosts, cached.implicit, cached.err
 	}
+
 	asciiHostname, err := ensureASCII(hostname)
 	if err != nil {
 		return nil, false, errors.Wrap(err, "invalid hostname")
@@ -62,23 +66,23 @@ func Lookup(ctx context.Context, hostname string) (retMxHosts []string, retImpli
 		var netDNSError *net.DNSError
 		if errors.As(err, &netDNSError) && netDNSError.Err == "no such host" {
 			if _, err := DefaultResolver.LookupIPAddr(ctx, asciiHostname); err != nil {
-				return cacheAndReturn(hostname, nil, false, errors.WithStack(err))
+				return cacheAndReturn(hostname, nil, nil, false, errors.WithStack(err))
 			}
-			return cacheAndReturn(hostname, []string{asciiHostname}, true, nil)
+			return cacheAndReturn(hostname, []string{asciiHostname}, nil, true, nil)
 		}
 		if mxRecords == nil {
-			return cacheAndReturn(hostname, nil, false, errors.WithStack(err))
+			return cacheAndReturn(hostname, nil, nil, false, errors.WithStack(err))
 		}
 	}
 	// Check for "Null MX" record (https://tools.ietf.org/html/rfc7505).
 	if len(mxRecords) == 1 {
 		if mxRecords[0].Host == "." {
-			return cacheAndReturn(hostname, nil, false, errNullMXRecord)
+			return cacheAndReturn(hostname, nil, nil, false, errNullMXRecord)
 		}
 		// 0.0.0.0 is not really a "Null MX" record, but some people apparently
 		// have never heard of RFC7505 and configure it this way.
 		if strings.HasPrefix(mxRecords[0].Host, "0.0.0.0") {
-			return cacheAndReturn(hostname, nil, false, errNullMXRecord)
+			return cacheAndReturn(hostname, nil, nil, false, errNullMXRecord)
 		}
 	}
 	// Normalize returned hostnames: drop trailing '.' and lowercase.
@@ -94,6 +98,14 @@ func Lookup(ctx context.Context, hostname string) (retMxHosts []string, retImpli
 		return mxRecords[i].Pref < mxRecords[j].Pref ||
 			(mxRecords[i].Pref == mxRecords[j].Pref && mxRecords[i].Host < mxRecords[j].Host)
 	})
+	mxHosts := shuffleMXRecords(mxRecords)
+	if len(mxHosts) == 0 {
+		return cacheAndReturn(hostname, nil, nil, false, errNoValidMXHosts)
+	}
+	return cacheAndReturn(hostname, mxHosts, mxRecords, false, nil)
+}
+
+func shuffleMXRecords(mxRecords []*net.MX) []string {
 	// Shuffle records within preference groups unless disabled in tests.
 	if shuffle {
 		mxRecordCount := len(mxRecords)
@@ -119,10 +131,7 @@ func Lookup(ctx context.Context, hostname string) (retMxHosts []string, retImpli
 		}
 		mxHosts = append(mxHosts, mxRecord.Host)
 	}
-	if len(mxHosts) == 0 {
-		return cacheAndReturn(hostname, nil, false, errNoValidMXHosts)
-	}
-	return cacheAndReturn(hostname, mxHosts, false, nil)
+	return mxHosts
 }
 
 func ensureASCII(hostname string) (string, error) {
@@ -146,13 +155,14 @@ func isASCII(s string) bool {
 }
 
 type lookupResult struct {
-	mxHosts  []string
-	implicit bool
-	err      error
+	mxRecords []*net.MX
+	mxHosts   []string
+	implicit  bool
+	err       error
 }
 
-func cacheAndReturn(hostname string, mxHosts []string, implicit bool, err error) (retMxHosts []string, retImplicit bool, reterr error) {
-	lookupResultCache.AddWithTTL(hostname, lookupResult{mxHosts, implicit, err}, cacheTTL)
+func cacheAndReturn(hostname string, mxHosts []string, mxRecords []*net.MX, implicit bool, err error) (retMxHosts []string, retImplicit bool, reterr error) {
+	lookupResultCache.AddWithTTL(hostname, lookupResult{mxHosts: mxHosts, mxRecords: mxRecords, implicit: implicit, err: err}, cacheTTL)
 	return mxHosts, implicit, err
 }
 

mxresolv/mxresolv_test.go

@@ -2,6 +2,7 @@ package mxresolv
 
 import (
 	"context"
+	"math"
 	"regexp"
 	"sort"
 	"testing"
@@ -178,6 +179,48 @@ func TestLookupShuffle(t *testing.T) {
 	assert.Equal(t, shuffle1[4:], shuffle2[4:])
 }
 
+func TestDistribution(t *testing.T) {
+	dist := make(map[string]int, 3)
+	for i := 0; i < 1000; i++ {
+		s, _, _ := Lookup(context.Background(), "test-mx.definbox.com")
+		_, ok := dist[s[0]]
+		if ok {
+			dist[s[0]] += 1
+		} else {
+			dist[s[0]] = 0
+		}
+	}
+
+	// Calculate the mean of the distribution
+	var sum int
+	for _, value := range dist {
+		sum += value
+	}
+	mean := float64(sum) / float64(len(dist))
+
+	// Calculate the sum of squared differences
+	var squaredDifferences float64
+	for _, value := range dist {
+		diff := float64(value) - mean
+		squaredDifferences += diff * diff
+	}
+
+	// Calculate the variance and standard deviation
+	variance := squaredDifferences / float64(len(dist))
+	stdDev := math.Sqrt(variance)
+
+	// The distribution of random hosts chosen should not exceed 30
+	assert.False(t, stdDev > 30.0, "Standard deviation is greater than 30: %.2f", stdDev)
+
+	// For example this is what a standard distribution looks like when 3 hosts have the same MX priority
+	// spew.Dump(dist)
+	// (map[string]int) (len=3) {
+	// 	(string) (len=16) "mxa.definbox.com": (int) 324,
+	//	(string) (len=16) "mxe.definbox.com": (int) 359,
+	//	(string) (len=16) "mxi.definbox.com": (int) 314
+	// }
+}
+
 func disableShuffle() func() {
 	shuffle = false
 	return func() {