From 6a2e0ad8174de86ec9540706d29a9b55fbebb78e Mon Sep 17 00:00:00 2001
From: "fox.cpp" <fox.cpp@disroot.org>
Date: Wed, 12 May 2021 06:11:26 +0300
Subject: [PATCH] check/spf: Use context.Context and custom resolver interface

Closes #354.
---
 go.mod                    |  2 +-
 go.sum                    |  4 ++--
 internal/check/spf/spf.go | 12 ++++++++----
 3 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/go.mod b/go.mod
index 0b1d346..7ee24cc 100644
--- a/go.mod
+++ b/go.mod
@@ -3,7 +3,7 @@ module github.com/foxcpp/maddy
 go 1.14
 
 require (
-	blitiri.com.ar/go/spf v1.1.1
+	blitiri.com.ar/go/spf v1.2.0
 	github.com/GehirnInc/crypt v0.0.0-20200316065508-bb7000b8a962
 	github.com/cpuguy83/go-md2man/v2 v2.0.0 // indirect
 	github.com/emersion/go-imap v1.0.6
diff --git a/go.sum b/go.sum
index c03469e..cd0f35a 100644
--- a/go.sum
+++ b/go.sum
@@ -1,5 +1,5 @@
-blitiri.com.ar/go/spf v1.1.1 h1:H5MKnEe5feN4NjtPDK/vFkkS0fI+ecTIsOfLNCR+6yI=
-blitiri.com.ar/go/spf v1.1.1/go.mod h1:HLmgHxdrsqbBgi5omEopdAKm18PypvUKJGkF/j7BO0w=
+blitiri.com.ar/go/spf v1.2.0 h1:aPpeEVKz5Ue4xb4SEt4AzScCSyES7/pol6znzZGle3A=
+blitiri.com.ar/go/spf v1.2.0/go.mod h1:HLmgHxdrsqbBgi5omEopdAKm18PypvUKJGkF/j7BO0w=
 cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
 cloud.google.com/go v0.34.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
diff --git a/internal/check/spf/spf.go b/internal/check/spf/spf.go
index 261eb3b..654e519 100644
--- a/internal/check/spf/spf.go
+++ b/internal/check/spf/spf.go
@@ -56,13 +56,15 @@ type Check struct {
 	permerrAction  modconfig.FailAction
 	temperrAction  modconfig.FailAction
 
-	log log.Logger
+	log      log.Logger
+	resolver dns.Resolver
 }
 
 func New(_, instName string, _, _ []string) (module.Module, error) {
 	return &Check{
 		instName: instName,
 		log:      log.Logger{Name: modName},
+		resolver: dns.DefaultResolver(),
 	}, nil
 }
 
@@ -243,7 +245,7 @@ func (s *state) relyOnDMARC(ctx context.Context, hdr textproto.Header) bool {
 		return false
 	}
 
-	policyDomain, record, err := maddydmarc.FetchRecord(ctx, dns.DefaultResolver(), fromDomain)
+	policyDomain, record, err := maddydmarc.FetchRecord(ctx, s.c.resolver, fromDomain)
 	if err != nil {
 		s.log.Error("DMARC fetch", err, "from_domain", fromDomain)
 		return false
@@ -323,7 +325,8 @@ func (s *state) CheckConnection(ctx context.Context) module.CheckResult {
 
 	if s.c.enforceEarly {
 		res, err := spf.CheckHostWithSender(ip.IP,
-			dns.FQDN(s.msgMeta.Conn.Hostname), mailFrom)
+			dns.FQDN(s.msgMeta.Conn.Hostname), mailFrom,
+			spf.WithContext(ctx), spf.WithResolver(s.c.resolver))
 		s.log.Debugf("result: %s (%v)", res, err)
 		return s.spfResult(res, err)
 	}
@@ -344,7 +347,8 @@ func (s *state) CheckConnection(ctx context.Context) module.CheckResult {
 
 		defer trace.StartRegion(ctx, "check.spf/CheckConnection (Async)").End()
 
-		res, err := spf.CheckHostWithSender(ip.IP, dns.FQDN(s.msgMeta.Conn.Hostname), mailFrom)
+		res, err := spf.CheckHostWithSender(ip.IP, dns.FQDN(s.msgMeta.Conn.Hostname), mailFrom,
+			spf.WithContext(ctx), spf.WithResolver(s.c.resolver))
 		s.log.Debugf("result: %s (%v)", res, err)
 		s.spfFetch <- spfRes{res, err}
 	}()