From 47c4e9fcb7176b46b9370cc061da96230897d08c Mon Sep 17 00:00:00 2001
From: Andres Erbsen <andreser@yahoo-inc.com>
Date: Tue, 18 Aug 2015 12:09:42 -0700
Subject: [PATCH] allow mocking

---
 dkim.go      | 32 +++++++++++++++++++++-----------
 dkim_test.go | 14 +++++++-------
 pubKeyRep.go | 12 +-----------
 3 files changed, 29 insertions(+), 29 deletions(-)

diff --git a/dkim.go b/dkim.go
index 411b573..1296181 100644
--- a/dkim.go
+++ b/dkim.go
@@ -13,6 +13,7 @@ import (
 	"encoding/base64"
 	"encoding/pem"
 	"hash"
+	"net"
 	"regexp"
 	"strings"
 	"time"
@@ -25,8 +26,6 @@ const (
 	MaxHeaderLineLength = 70
 )
 
-type verifyOutput int
-
 // sigOptions represents signing options
 type SigOptions struct {
 
@@ -181,26 +180,36 @@ func Sign(email []byte, options SigOptions) ([]byte, error) {
 	return append([]byte(dHeader), email...), nil
 }
 
-func Verify(email []byte) (dkimHeader *DKIMHeader, err error) {
-	// parse email
+func Verify(email []byte, LookupTXT func(string) ([]string, error), now func() time.Time) (dkimHeader *DKIMHeader, err error) {
+	if LookupTXT == nil {
+		LookupTXT = net.LookupTXT
+	}
+	if now == nil {
+		now = time.Now
+	}
+
 	dkimHeader, err = newDkimHeaderFromEmail(email)
 	if err != nil {
 		return
 	}
 
-	// we do not set query method because if it's others, validation failed earlier
-	pubKey, err := newPubKeyFromDnsTxt(dkimHeader.Selector, dkimHeader.Domain)
+	txt, err := LookupTXT(dkimHeader.Selector + "._domainkey." + dkimHeader.Domain)
+	if err != nil {
+		return nil, err
+	}
+	pubKey, err := newPubKeyFromDnsTxt(txt)
 	if err != nil {
 		return nil, err
 	}
 
-	// Normalize
 	headers, body, err := canonicalize(email, dkimHeader.MessageCanonicalization, dkimHeader.Headers)
 	if err != nil {
 		return nil, err
 	}
 	sigHash := strings.Split(dkimHeader.Algorithm, "-")
-	// check if hash algo are compatible
+	if len(sigHash) < 2 {
+		return nil, ErrVerifyInappropriateHashAlgo
+	}
 	compatible := false
 	for _, algo := range pubKey.HashAlgo {
 		if sigHash[1] == algo {
@@ -212,9 +221,10 @@ func Verify(email []byte) (dkimHeader *DKIMHeader, err error) {
 		return nil, ErrVerifyInappropriateHashAlgo
 	}
 
-	// expired ?
-	if !dkimHeader.SignatureExpiration.IsZero() && dkimHeader.SignatureExpiration.Second() < time.Now().Second() {
-		return nil, ErrVerifySignatureHasExpired
+	if !dkimHeader.SignatureExpiration.IsZero() {
+		if dkimHeader.SignatureExpiration.Before(now()) {
+			return nil, ErrVerifySignatureHasExpired
+		}
 	}
 
 	bodyHash, err := getBodyHash(body, sigHash[1], dkimHeader.BodyLength)
diff --git a/dkim_test.go b/dkim_test.go
index be4436d..8bb9c39 100644
--- a/dkim_test.go
+++ b/dkim_test.go
@@ -333,38 +333,38 @@ func Test_Sign(t *testing.T) {
 func Test_Verify(t *testing.T) {
 	// no DKIM header
 	email := []byte(emailBase)
-	_, err := Verify(email)
+	_, err := Verify(email, nil, nil)
 	assert.Equal(t, ErrDkimHeaderNotFound, err)
 
 	// No From
 	email = []byte(signedNoFrom)
-	_, err = Verify(email)
+	_, err = Verify(email, nil, nil)
 	assert.Equal(t, ErrVerifyBodyHash, err)
 
 	// missing mandatory 'a' flag
 	email = []byte(signedMissingFlag)
-	_, err = Verify(email)
+	_, err = Verify(email, nil, nil)
 	assert.Error(t, err)
 	assert.Equal(t, ErrDkimHeaderMissingRequiredTag, err)
 
 	// missing bad algo
 	email = []byte(signedBadAlgo)
-	_, err = Verify(email)
+	_, err = Verify(email, nil, nil)
 	assert.Equal(t, ErrSignBadAlgo, err)
 
 	// relaxed
 	email = []byte(signedRelaxedRelaxedLength)
-	_, err = Verify(email)
+	_, err = Verify(email, nil, nil)
 	assert.Equal(t, ErrTesting, err)
 
 	// simple
 	email = []byte(signedSimpleSimpleLength)
-	_, err = Verify(email)
+	_, err = Verify(email, nil, nil)
 	assert.Equal(t, ErrTesting, err)
 
 	// gmail
 	email = []byte(fromGmail)
-	_, err = Verify(email)
+	_, err = Verify(email, nil, nil)
 	assert.NoError(t, err)
 
 }
diff --git a/pubKeyRep.go b/pubKeyRep.go
index 6ec248b..5c2eecd 100644
--- a/pubKeyRep.go
+++ b/pubKeyRep.go
@@ -4,7 +4,6 @@ import (
 	"crypto/rsa"
 	"crypto/x509"
 	"encoding/base64"
-	"net"
 	"strings"
 )
 
@@ -20,16 +19,7 @@ type pubKeyRep struct {
 	FlagIMustBeD bool // flag i
 }
 
-func newPubKeyFromDnsTxt(selector, domain string) (*pubKeyRep, error) {
-	txt, err := net.LookupTXT(selector + "._domainkey." + domain)
-	if err != nil {
-		if strings.HasSuffix(err.Error(), "no such host") {
-			return nil, ErrVerifyNoKeyForSignature
-		} else {
-			return nil, ErrVerifyKeyUnavailable
-		}
-	}
-
+func newPubKeyFromDnsTxt(txt []string) (*pubKeyRep, error) {
 	// empty record
 	if len(txt) == 0 {
 		return nil, ErrVerifyNoKeyForSignature