Fix first-request good bot DNS verification (#90)

Author: joecorall

README.md

@@ -90,7 +90,7 @@ services:
             --providers.docker=true
             --providers.docker.network=default
             --experimental.plugins.captcha-protect.modulename=github.com/libops/captcha-protect
-            --experimental.plugins.captcha-protect.version=v1.12.3
+            --experimental.plugins.captcha-protect.version=v1.12.4
         volumes:
             - /var/run/docker.sock:/var/run/docker.sock:z
             - /CHANGEME/TO/A/HOST/PATH/FOR/STATE/FILE:/tmp/state.json:rw

internal/helper/ip.go

@@ -1,13 +1,14 @@
 package helper
 
 import (
+	"context"
 	"net"
 	"strings"
 )
 
 var (
-	lookupAddrFunc = net.LookupAddr
-	lookupIPFunc   = net.LookupIP
+	lookupAddrFunc = net.DefaultResolver.LookupAddr
+	lookupIPFunc   = net.DefaultResolver.LookupIP
 )
 
 func IsIpExcluded(clientIP string, exemptIps []*net.IPNet) bool {
@@ -29,39 +30,48 @@ func ParseCIDR(cidr string) (*net.IPNet, error) {
 	return ipNet, nil
 }
 
-func IsIpGoodBot(clientIP string, goodBots []string) bool {
+func IsIpGoodBot(ctx context.Context, clientIP string, goodBots []string) bool {
 	if len(goodBots) == 0 {
 		return false
 	}
-
-	// lookup the hostname for a given IP
-	hostname, err := lookupAddrFunc(clientIP)
-	if err != nil || len(hostname) == 0 {
+	ip := net.ParseIP(clientIP)
+	if ip == nil {
 		return false
 	}
 
-	// then nslookup that hostname to avoid spoofing
-	resolvedIP, err := lookupIPFunc(hostname[0])
-	if err != nil || len(resolvedIP) == 0 || resolvedIP[0].String() != clientIP {
+	// lookup the hostname for a given IP
+	hostnames, err := lookupAddrFunc(ctx, clientIP)
+	if err != nil || len(hostnames) == 0 {
 		return false
 	}
 
-	// get the sld
-	// will be like 194.114.135.34.bc.googleusercontent.com.
-	// notice the trailing period
-	parts := strings.Split(hostname[0], ".")
-	l := len(parts)
-	if l < 3 {
-		return false
+	for _, hostname := range hostnames {
+		hostname = strings.ToLower(strings.TrimSuffix(hostname, "."))
+		if !matchesGoodBotDomain(hostname, goodBots) {
+			continue
+		}
+
+		// Resolve the PTR hostname forward to prevent forged reverse DNS records.
+		resolvedIPs, err := lookupIPFunc(ctx, "ip", hostname)
+		if err != nil {
+			continue
+		}
+		for _, resolvedIP := range resolvedIPs {
+			if resolvedIP.Equal(ip) {
+				return true
+			}
+		}
 	}
-	tld := parts[l-2]
-	domain := parts[l-3] + "." + tld
 
+	return false
+}
+
+func matchesGoodBotDomain(hostname string, goodBots []string) bool {
 	for _, bot := range goodBots {
-		if domain == bot {
+		bot = strings.ToLower(strings.Trim(strings.TrimSpace(bot), "."))
+		if bot != "" && (hostname == bot || strings.HasSuffix(hostname, "."+bot)) {
 			return true
 		}
 	}
-
 	return false
 }

internal/helper/ip_test.go

@@ -1,8 +1,10 @@
 package helper
 
 import (
+	"context"
 	"errors"
 	"net"
+	"strings"
 	"testing"
 )
 
@@ -46,7 +48,7 @@ func TestIsIpGoodBot(t *testing.T) {
 			clientIP: "1.2.3.4",
 			goodBots: []string{"google.com"},
 			lookupAddrReturn: []string{
-				"host.example.com.",
+				"crawl.google.com.",
 			},
 			lookupAddrErr: nil,
 			lookupIPReturn: []net.IP{
@@ -97,25 +99,42 @@ func TestIsIpGoodBot(t *testing.T) {
 			lookupIPErr: nil,
 			expected:    true,
 		},
+		{
+			name:     "Client IP is not the first forward result",
+			clientIP: "1.2.3.4",
+			goodBots: []string{"example.com"},
+			lookupAddrReturn: []string{
+				"crawler.example.com.",
+			},
+			lookupIPReturn: []net.IP{
+				net.ParseIP("5.6.7.8"),
+				net.ParseIP("1.2.3.4"),
+			},
+			expected: true,
+		},
 	}
 
 	for _, tc := range tests {
-		lookupAddrFunc = func(ip string) ([]string, error) {
+		lookupAddrFunc = func(_ context.Context, ip string) ([]string, error) {
 			if ip != tc.clientIP {
 				t.Errorf("Expected lookupAddr to be called with %q; got %q", tc.clientIP, ip)
 			}
 			return tc.lookupAddrReturn, tc.lookupAddrErr
 		}
 
-		lookupIPFunc = func(host string) ([]net.IP, error) {
-			if len(tc.lookupAddrReturn) == 0 || host != tc.lookupAddrReturn[0] {
-				t.Errorf("Expected lookupIP to be called with %q; got %q", tc.lookupAddrReturn[0], host)
+		lookupIPFunc = func(_ context.Context, network, host string) ([]net.IP, error) {
+			if network != "ip" {
+				t.Errorf("Expected lookupIP network %q; got %q", "ip", network)
+			}
+			expectedHost := strings.TrimSuffix(tc.lookupAddrReturn[0], ".")
+			if host != expectedHost {
+				t.Errorf("Expected lookupIP to be called with %q; got %q", expectedHost, host)
 			}
 			return tc.lookupIPReturn, tc.lookupIPErr
 		}
 
 		t.Run(tc.name, func(t *testing.T) {
-			result := IsIpGoodBot(tc.clientIP, tc.goodBots)
+			result := IsIpGoodBot(context.Background(), tc.clientIP, tc.goodBots)
 			if result != tc.expected {
 				t.Errorf("IsIpGoodBot(%q) = %v; expected %v", tc.clientIP, result, tc.expected)
 			}

main.go

@@ -38,6 +38,7 @@ const (
 	// Default health check settings (disabled by default)
 	DefaultHealthCheckPeriodSeconds    = 0 // How often to check captcha provider health
 	DefaultHealthCheckFailureThreshold = 0 // Number of consecutive health check failures before opening circuit
+	goodBotLookupTimeout               = 2 * time.Second
 )
 
 type circuitState int
@@ -106,6 +107,7 @@ type CaptchaProtect struct {
 	stateDirty         uint64
 	stateSavedDirty    uint64
 	stateFileModTime   time.Time
+	goodBotLookup      func(context.Context, string, []string) bool
 
 	// Circuit breaker fields
 	mu                      sync.RWMutex
@@ -292,6 +294,7 @@ func NewCaptchaProtect(ctx context.Context, next http.Handler, config *Config, n
 		},
 		rateCache:          lru.New(expiration, 1*time.Minute),
 		botCache:           lru.New(expiration, 1*time.Hour),
+		goodBotLookup:      helper.IsIpGoodBot,
 		verifiedCache:      lru.New(expiration, 1*time.Hour),
 		exemptIps:          ips,
 		tmpl:               tmpl,
@@ -1014,7 +1017,9 @@ func (bc *CaptchaProtect) isGoodBot(req *http.Request, clientIP string) bool {
 		}
 	}
 	if !v {
-		v = helper.IsIpGoodBot(clientIP, bc.config.GoodBots)
+		ctx, cancel := context.WithTimeout(req.Context(), goodBotLookupTimeout)
+		defer cancel()
+		v = bc.goodBotLookup(ctx, clientIP, bc.config.GoodBots)
 	}
 	bc.botCache.Set(clientIP, v, lru.DefaultExpiration)
 	return v

main_test.go

@@ -566,6 +566,48 @@ func TestServeHTTP(t *testing.T) {
 	}
 }
 
+func TestServeHTTPAllowsGoodBotOnFirstRequest(t *testing.T) {
+	upstreamCalled := false
+	next := http.HandlerFunc(func(rw http.ResponseWriter, _ *http.Request) {
+		upstreamCalled = true
+		rw.WriteHeader(http.StatusOK)
+	})
+	config := CreateConfig()
+	config.SiteKey = "test-site-key"
+	config.SecretKey = "test-secret-key"
+	config.RateLimit = 0
+	config.ProtectRoutes = []string{"/"}
+	config.GoodBots = []string{"yandex.com"}
+
+	cp, err := NewCaptchaProtect(context.Background(), next, config, "captcha-protect")
+	if err != nil {
+		t.Fatal(err)
+	}
+	lookupCalled := false
+	cp.goodBotLookup = func(ctx context.Context, clientIP string, goodBots []string) bool {
+		lookupCalled = true
+		if _, ok := ctx.Deadline(); !ok {
+			t.Error("expected DNS lookup context to have a deadline")
+		}
+		return clientIP == "5.255.231.189" && len(goodBots) == 1 && goodBots[0] == "yandex.com"
+	}
+
+	req := httptest.NewRequest(http.MethodGet, "http://example.com/", nil)
+	req.RemoteAddr = "5.255.231.189:1234"
+	rw := httptest.NewRecorder()
+	cp.ServeHTTP(rw, req)
+
+	if !lookupCalled {
+		t.Fatal("expected DNS verification on the first request")
+	}
+	if !upstreamCalled {
+		t.Fatal("expected verified bot's first request to reach the upstream handler")
+	}
+	if rw.Code != http.StatusOK {
+		t.Fatalf("expected status %d, got %d", http.StatusOK, rw.Code)
+	}
+}
+
 func TestIsGoodUserAgent(t *testing.T) {
 	tests := []struct {
 		name             string