git » spf » commit 8f601a6

diff --git a/spf.go b/spf.go
index 2adc3f4..97d2208 100644
--- a/spf.go
+++ b/spf.go
@@ -357,13 +357,6 @@ func (r *resolution) Check(domain string) (Result, error) {
 			continue
 		}
 
-		// Limit the number of resolutions.
-		// https://tools.ietf.org/html/rfc7208#section-4.6.4
-		if r.count > r.maxcount {
-			r.trace("lookup limit reached")
-			return PermError, ErrLookupLimitReached
-		}
-
 		if r.voidcount > r.maxvoidcount {
 			r.trace("void lookup limit reached")
 			return PermError, ErrVoidLookupLimitReached
@@ -485,6 +478,19 @@ func isNotFound(err error) bool {
 	return ok && derr.IsNotFound
 }
 
+// Account for a DNS lookup, and return an error if the limit was already
+// reached. Fields which do a DNS lookup must call this before doing it, so we
+// never go over the limit.
+// https://tools.ietf.org/html/rfc7208#section-4.6.4
+func (r *resolution) countLookup() error {
+	if r.count >= r.maxcount {
+		r.trace("lookup limit reached")
+		return ErrLookupLimitReached
+	}
+	r.count++
+	return nil
+}
+
 // Check if the given DNS error is a "void lookup" (0 answers, or nxdomain),
 // and if so increment the void lookup counter.
 func (r *resolution) checkVoidLookup(nanswers int, err error) {
@@ -549,8 +555,10 @@ func (r *resolution) ptrField(res Result, field, domain string) (bool, Result, e
 	}
 
 	if r.ipNames == nil {
+		if err := r.countLookup(); err != nil {
+			return true, PermError, err
+		}
 		r.ipNames = []string{}
-		r.count++
 		ns, err := r.resolver.LookupAddr(r.ctx, r.ip.String())
 		r.checkVoidLookup(len(ns), err)
 		if err != nil {
@@ -613,7 +621,9 @@ func (r *resolution) existsField(res Result, field, domain string) (bool, Result
 		return true, PermError, ErrInvalidDomain
 	}
 
-	r.count++
+	if err := r.countLookup(); err != nil {
+		return true, PermError, err
+	}
 	ips, err := r.resolver.LookupIPAddr(r.ctx, eDomain)
 	r.checkVoidLookup(len(ips), err)
 	if err != nil {
@@ -642,7 +652,9 @@ func (r *resolution) includeField(res Result, field, domain string) (bool, Resul
 	if err != nil {
 		return true, PermError, ErrInvalidMacro
 	}
-	r.count++
+	if err := r.countLookup(); err != nil {
+		return true, PermError, err
+	}
 	ir, err := r.Check(incdomain)
 	switch ir {
 	case Pass:
@@ -743,7 +755,9 @@ func (r *resolution) aField(res Result, field, domain string) (bool, Result, err
 		return true, PermError, ErrInvalidMacro
 	}
 
-	r.count++
+	if err := r.countLookup(); err != nil {
+		return true, PermError, err
+	}
 	ips, err := r.resolver.LookupIPAddr(r.ctx, aDomain)
 	r.checkVoidLookup(len(ips), err)
 	if err != nil {
@@ -776,7 +790,9 @@ func (r *resolution) mxField(res Result, field, domain string) (bool, Result, er
 		return true, PermError, ErrInvalidMacro
 	}
 
-	r.count++
+	if err := r.countLookup(); err != nil {
+		return true, PermError, err
+	}
 	mxs, err := r.resolver.LookupMX(r.ctx, mxDomain)
 	r.checkVoidLookup(len(mxs), err)
 
@@ -839,7 +855,9 @@ func (r *resolution) redirectField(field, domain string) (Result, error) {
 	}
 
 	// https://tools.ietf.org/html/rfc7208#section-6.1
-	r.count++
+	if err := r.countLookup(); err != nil {
+		return PermError, err
+	}
 	result, err := r.Check(rDomain)
 	if result == None {
 		result = PermError
diff --git a/spf_test.go b/spf_test.go
index 181e2a8..b860ffe 100644
--- a/spf_test.go
+++ b/spf_test.go
@@ -3,6 +3,7 @@ package spf
 import (
 	"context"
 	"net"
+	"strings"
 	"testing"
 
 	"blitiri.com.ar/go/spf/internal/dnstest"
@@ -195,6 +196,55 @@ func TestRecursionLimit(t *testing.T) {
 	}
 }
 
+func TestLookupLimit(t *testing.T) {
+	dns := NewDefaultResolver()
+	defaultTrace = t.Logf
+
+	// "a:d1110" never matches the IP we check, so we use it to consume
+	// lookups; each occurrence counts as one.
+	// https://tools.ietf.org/html/rfc7208#section-4.6.4
+	ten := strings.Repeat("a:d1110 ", 10)
+
+	cases := []struct {
+		txt string
+		res Result
+		err error
+	}{
+		// Ten lookups are allowed, and the fields after them are evaluated.
+		{"v=spf1 " + ten + "-all", Fail, ErrMatchedAll},
+		{"v=spf1 " + ten + "ip4:1.1.1.1", Pass, ErrMatchedIP},
+
+		// The 11th lookup is not performed, even if it would match.
+		{"v=spf1 " + ten + "a:d1111 -all", PermError, ErrLookupLimitReached},
+		{"v=spf1 " + ten + "mx:d1110 -all", PermError, ErrLookupLimitReached},
+		{"v=spf1 " + ten + "ptr:d1111 -all", PermError, ErrLookupLimitReached},
+		{"v=spf1 " + ten + "exists:d1111 -all", PermError, ErrLookupLimitReached},
+		{"v=spf1 " + ten + "include:alldomain -all", PermError, ErrLookupLimitReached},
+		{"v=spf1 " + ten + "redirect=alldomain", PermError, ErrLookupLimitReached},
+
+		// The 11th lookup is the last field, so there is nothing after it to
+		// catch the excess.
+		{"v=spf1 " + ten + "a:d1111", PermError, ErrLookupLimitReached},
+	}
+
+	dns.Ip["d1111"] = []net.IP{ip1111}
+	dns.Ip["d1110"] = []net.IP{ip1110}
+	dns.Mx["d1110"] = []*net.MX{mx("d1111", 5)}
+	dns.Addr["1.1.1.1"] = []string{"d1111."}
+	dns.Txt["alldomain"] = []string{"v=spf1 +all"}
+
+	for _, c := range cases {
+		dns.Txt["domain"] = []string{c.txt}
+		res, err := CheckHost(ip1111, "domain")
+		if res != c.res {
+			t.Errorf("%q: expected %q, got %q", c.txt, c.res, res)
+		}
+		if err != c.err {
+			t.Errorf("%q: expected error [%v], got [%v]", c.txt, c.err, err)
+		}
+	}
+}
+
 func TestRedirect(t *testing.T) {
 	dns := NewDefaultResolver()
 	dns.Txt["domain"] = []string{"v=spf1 redirect=domain2"}