dmarc, dkim: allow setting lookup DNS TXT function

References: #10

dkim/query.go

@@ -62,14 +62,22 @@ const (
 	QueryMethodDNSTXT QueryMethod = "dns/txt"
 )
 
-type queryFunc func(domain, selector string) (*queryResult, error)
+type txtLookupFunc func(domain string) ([]string, error)
+type queryFunc func(domain, selector string, txtLookup txtLookupFunc) (*queryResult, error)
 
 var queryMethods = map[QueryMethod]queryFunc{
 	QueryMethodDNSTXT: queryDNSTXT,
 }
 
-func queryDNSTXT(domain, selector string) (*queryResult, error) {
-	txts, err := net.LookupTXT(selector + "._domainkey." + domain)
+func queryDNSTXT(domain, selector string, txtLookup txtLookupFunc) (*queryResult, error) {
+	var txts []string
+	var err error
+	if txtLookup != nil {
+		txts, err = txtLookup(selector + "._domainkey." + domain)
+	} else {
+		txts, err = net.LookupTXT(selector + "._domainkey." + domain)
+	}
+
 	if netErr, ok := err.(net.Error); ok && netErr.Temporary() {
 		return nil, tempFailError("key unavailable: " + err.Error())
 	} else if err != nil {

dkim/query_test.go

@@ -16,7 +16,7 @@ func init() {
 	queryMethods["dns/txt"] = queryTest
 }
 
-func queryTest(domain, selector string) (*queryResult, error) {
+func queryTest(domain, selector string, txtLookup txtLookupFunc) (*queryResult, error) {
 	record := selector + "._domainkey." + domain
 	switch record {
 	case "brisbane._domainkey.example.com", "brisbane._domainkey.example.org", "test._domainkey.football.example.com":

dkim/verify.go

@@ -86,11 +86,21 @@ type signature struct {
 	v string
 }
 
+// VerifyOptions allows to customize the default signature verification behavior
+// LookupTXT returns the DNS TXT records for the given domain name. If nil, net.LookupTXT is used
+type VerifyOptions struct {
+	LookupTXT func(domain string) ([]string, error)
+}
+
 // Verify checks if a message's signatures are valid. It returns one
 // verification per signature.
 //
 // There is no guarantee that the reader will be completely consumed.
 func Verify(r io.Reader) ([]*Verification, error) {
+	return VerifyWithOptions(r, nil)
+}
+
+func VerifyWithOptions(r io.Reader, options *VerifyOptions) ([]*Verification, error) {
 	// TODO: be able to specify options such as the max number of signatures to
 	// check
 
@@ -111,11 +121,11 @@ func Verify(r io.Reader) ([]*Verification, error) {
 	}
 
 	if len(signatures) != 1 {
-		return parallelVerify(bufr, h, signatures)
+		return parallelVerify(bufr, h, signatures, options)
 	}
 
 	// If there is only one signature - just verify it.
-	v, err := verify(h, bufr, h[signatures[0].i], signatures[0].v)
+	v, err := verify(h, bufr, h[signatures[0].i], signatures[0].v, options)
 	if err != nil && !IsTempFail(err) && !IsPermFail(err) && !isFail(err) {
 		return nil, err
 	}
@@ -124,7 +134,7 @@ func Verify(r io.Reader) ([]*Verification, error) {
 	return []*Verification{v}, nil
 }
 
-func parallelVerify(r io.Reader, h header, signatures []*signature) ([]*Verification, error) {
+func parallelVerify(r io.Reader, h header, signatures []*signature, options *VerifyOptions) ([]*Verification, error) {
 	pipeWriters := make([]*io.PipeWriter, len(signatures))
 	// We can't pass pipeWriter to io.MultiWriter directly,
 	// we need a slice of io.Writer, but we also need *io.PipeWriter
@@ -143,7 +153,7 @@ func parallelVerify(r io.Reader, h header, signatures []*signature) ([]*Verifica
 		pipeWriters[i] = pw
 
 		go func() {
-			v, err := verify(h, pr, h[sig.i], sig.v)
+			v, err := verify(h, pr, h[sig.i], sig.v, options)
 
 			// Make sure we consume the whole reader, otherwise io.Copy on
 			// other side can block forever.
@@ -177,7 +187,7 @@ func parallelVerify(r io.Reader, h header, signatures []*signature) ([]*Verifica
 	return verifications, nil
 }
 
-func verify(h header, r io.Reader, sigField, sigValue string) (*Verification, error) {
+func verify(h header, r io.Reader, sigField, sigValue string, options *VerifyOptions) (*Verification, error) {
 	verif := new(Verification)
 
 	params, err := parseHeaderParams(sigValue)
@@ -246,7 +256,11 @@ func verify(h header, r io.Reader, sigField, sigValue string) (*Verification, er
 	var res *queryResult
 	for _, method := range methods {
 		if query, ok := queryMethods[QueryMethod(method)]; ok {
-			res, err = query(verif.Domain, stripWhitespace(params["s"]))
+			if options != nil {
+				res, err = query(verif.Domain, stripWhitespace(params["s"]), options.LookupTXT)
+			} else {
+				res, err = query(verif.Domain, stripWhitespace(params["s"]), nil)
+			}
 			break
 		}
 	}

dkim/verify_test.go

@@ -3,6 +3,7 @@ package dkim
 import (
 	"errors"
 	"io"
+	"net"
 	"reflect"
 	"strings"
 	"testing"
@@ -84,6 +85,36 @@ func TestVerify(t *testing.T) {
 	}
 }
 
+func TestVerifyWithOption(t *testing.T) {
+	r := newMailStringReader(verifiedMailString)
+	option := VerifyOptions{}
+	verifications, err := VerifyWithOptions(r, &option)
+	if err != nil {
+		t.Fatalf("Expected no error while verifying signature, got: %v", err)
+	} else if len(verifications) != 1 {
+		t.Fatalf("Expected exactly one verification, got %v", len(verifications))
+	}
+
+	v := verifications[0]
+	if !reflect.DeepEqual(testVerification, v) {
+		t.Errorf("Expected verification to be \n%+v\n but got \n%+v", testVerification, v)
+	}
+
+	r = newMailStringReader(verifiedMailString)
+	option = VerifyOptions{LookupTXT: net.LookupTXT}
+	verifications, err = VerifyWithOptions(r, &option)
+	if err != nil {
+		t.Fatalf("Expected no error while verifying signature, got: %v", err)
+	} else if len(verifications) != 1 {
+		t.Fatalf("Expected exactly one verification, got %v", len(verifications))
+	}
+
+	v = verifications[0]
+	if !reflect.DeepEqual(testVerification, v) {
+		t.Errorf("Expected verification to be \n%+v\n but got \n%+v", testVerification, v)
+	}
+}
+
 const verifiedEd25519MailString = `DKIM-Signature: v=1; a=ed25519-sha256; c=relaxed/relaxed;
  d=football.example.com; [email protected];
  q=dns/txt; s=brisbane; t=1528637909; h=from : to :
@@ -137,7 +168,7 @@ func TestVerify_ed25519(t *testing.T) {
 
 // errorReader reads from r and then returns an arbitrary error.
 type errorReader struct {
-	r io.Reader
+	r   io.Reader
 	err error
 }
 
@@ -159,7 +190,7 @@ func TestVerify_invalid(t *testing.T) {
 	expectedErr := errors.New("expected test error")
 
 	r = &errorReader{
-		r: newMailStringReader(verifiedEd25519MailString),
+		r:   newMailStringReader(verifiedEd25519MailString),
 		err: expectedErr,
 	}
 	_, err = Verify(r)

dmarc/lookup.go

@@ -24,9 +24,25 @@ func IsTempFail(err error) bool {
 
 var ErrNoPolicy = errors.New("dmarc: no policy found for domain")
 
+// LookupOptions allows to customize the default signature verification behavior
+// LookupTXT returns the DNS TXT records for the given domain name. If nil, net.LookupTXT is used
+type LookupOptions struct {
+	LookupTXT func(domain string) ([]string, error)
+}
+
 // Lookup queries a DMARC record for a specified domain.
 func Lookup(domain string) (*Record, error) {
-	txts, err := net.LookupTXT("_dmarc." + domain)
+	return LookupWithOptions(domain, nil)
+}
+
+func LookupWithOptions(domain string, options *LookupOptions) (*Record, error) {
+	var txts []string
+	var err error
+	if options != nil && options.LookupTXT != nil {
+		txts, err = options.LookupTXT("_dmarc." + domain)
+	} else {
+		txts, err = net.LookupTXT("_dmarc." + domain)
+	}
 	if netErr, ok := err.(net.Error); ok && netErr.Temporary() {
 		return nil, tempFailError("TXT record unavailable: " + err.Error())
 	} else if err != nil {