git » chasquid » commit 2601754

diff --git a/internal/courier/smtp.go b/internal/courier/smtp.go
index 3df9a5c..4c779af 100644
--- a/internal/courier/smtp.go
+++ b/internal/courier/smtp.go
@@ -26,8 +26,10 @@ var (
 	smtpPort = flag.String("testing__outgoing_smtp_port", "25",
 		"port to use for outgoing SMTP connections, ONLY FOR TESTING")
 
-	// Fake MX records, used for testing only.
-	fakeMX = map[string][]string{}
+	// Allow overriding of net.LookupMX for testing purposes.
+	// TODO: replace this with proper lookup interception once it is supported
+	// by Go.
+	netLookupMX = net.LookupMX
 )
 
 // Exported variables.
@@ -198,10 +200,6 @@ retry:
 }
 
 func lookupMXs(tr *trace.Trace, domain string) ([]string, error) {
-	if v, ok := fakeMX[domain]; ok {
-		return v, nil
-	}
-
 	domain, err := idna.ToASCII(domain)
 	if err != nil {
 		return nil, err
@@ -209,7 +207,7 @@ func lookupMXs(tr *trace.Trace, domain string) ([]string, error) {
 
 	mxs := []string{}
 
-	mxRecords, err := net.LookupMX(domain)
+	mxRecords, err := netLookupMX(domain)
 	if err != nil {
 		// There was an error. It could be that the domain has no MX, in which
 		// case we have to fall back to A, or a bigger problem.
diff --git a/internal/courier/smtp_test.go b/internal/courier/smtp_test.go
index c51d691..b12d7be 100644
--- a/internal/courier/smtp_test.go
+++ b/internal/courier/smtp_test.go
@@ -2,15 +2,32 @@ package courier
 
 import (
 	"bufio"
+	"fmt"
 	"net"
 	"net/textproto"
+	"strings"
 	"testing"
 	"time"
 
 	"blitiri.com.ar/go/chasquid/internal/domaininfo"
 	"blitiri.com.ar/go/chasquid/internal/testlib"
+	"blitiri.com.ar/go/chasquid/internal/trace"
 )
 
+// This domain will cause idna.ToASCII to fail.
+var invalidDomain = "test " + strings.Repeat("x", 65536) + "\uff00"
+
+// Override the netLookupMX function, to return controlled results for
+// testing.
+var testMX = map[string][]*net.MX{}
+var testMXErr = map[string]error{}
+
+func init() {
+	netLookupMX = func(name string) ([]*net.MX, error) {
+		return testMX[name], testMXErr[name]
+	}
+}
+
 func newSMTP(t *testing.T) (*SMTP, string) {
 	dir := testlib.MustTempDir(t)
 	dinfo, err := domaininfo.New(dir)
@@ -88,7 +105,7 @@ func TestSMTP(t *testing.T) {
 	// lookup whick makes the test more hermetic. This is a hack, ideally we
 	// would be able to override the default resolver, but Go does not
 	// implement that yet.
-	fakeMX["to"] = []string{":::", host}
+	testMX["to"] = []*net.MX{{":::", 10}, {host, 20}}
 	*smtpPort = port
 
 	s, tmpDir := newSMTP(t)
@@ -149,7 +166,7 @@ func TestSMTPErrors(t *testing.T) {
 		addr := fakeServer(t, rs)
 		host, port, _ := net.SplitHostPort(addr)
 
-		fakeMX["to"] = []string{host}
+		testMX["to"] = []*net.MX{{Host: host, Pref: 10}}
 		*smtpPort = port
 
 		s, tmpDir := newSMTP(t)
@@ -163,7 +180,7 @@ func TestSMTPErrors(t *testing.T) {
 }
 
 func TestNoMXServer(t *testing.T) {
-	fakeMX["to"] = []string{}
+	testMX["to"] = []*net.MX{}
 
 	s, tmpDir := newSMTP(t)
 	defer testlib.RemoveIfOk(t, tmpDir)
@@ -177,4 +194,71 @@ func TestNoMXServer(t *testing.T) {
 	t.Logf("got permanent failure, as expected: %v", err)
 }
 
+func TestTooManyMX(t *testing.T) {
+	tr := trace.New("test", "test")
+	testMX["domain"] = []*net.MX{
+		{Host: "h1", Pref: 10}, {Host: "h2", Pref: 20},
+		{Host: "h3", Pref: 30}, {Host: "h4", Pref: 40},
+		{Host: "h5", Pref: 50}, {Host: "h5", Pref: 60},
+	}
+	mxs, err := lookupMXs(tr, "domain")
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if len(mxs) != 5 {
+		t.Errorf("expected len(mxs) == 5, got: %v", mxs)
+	}
+}
+
+func TestFallbackToA(t *testing.T) {
+	tr := trace.New("test", "test")
+	testMX["domain"] = nil
+	testMXErr["domain"] = &net.DNSError{
+		Err:         "no such host (test)",
+		IsTemporary: false,
+	}
+
+	mxs, err := lookupMXs(tr, "domain")
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if !(len(mxs) == 1 && mxs[0] == "domain") {
+		t.Errorf("expected mxs == [domain], got: %v", mxs)
+	}
+}
+
+func TestTemporaryDNSerror(t *testing.T) {
+	tr := trace.New("test", "test")
+	testMX["domain"] = nil
+	testMXErr["domain"] = &net.DNSError{
+		Err:         "temp error (test)",
+		IsTemporary: true,
+	}
+
+	mxs, err := lookupMXs(tr, "domain")
+	if !(mxs == nil && err == testMXErr["domain"]) {
+		t.Errorf("expected mxs == nil, err == test error, got: %v, %v", mxs, err)
+	}
+}
+
+func TestMXLookupError(t *testing.T) {
+	tr := trace.New("test", "test")
+	testMX["domain"] = nil
+	testMXErr["domain"] = fmt.Errorf("test error")
+
+	mxs, err := lookupMXs(tr, "domain")
+	if !(mxs == nil && err == testMXErr["domain"]) {
+		t.Errorf("expected mxs == nil, err == test error, got: %v, %v", mxs, err)
+	}
+}
+
+func TestLookupInvalidDomain(t *testing.T) {
+	tr := trace.New("test", "test")
+
+	mxs, err := lookupMXs(tr, invalidDomain)
+	if !(mxs == nil && err != nil) {
+		t.Errorf("expected err != nil, got: %v, %v", mxs, err)
+	}
+}
+
 // TODO: Test STARTTLS negotiation.