Add "keep" functionality to leave specified mechanisms unflattened

.github/workflows/check-spf-record.yaml

@@ -15,7 +15,8 @@ jobs:
         shell: bash
         run: |
           initialSPF="v=spf1 include:_spf.google.com ip4:23.178.112.0/24 ip4:66.133.109.36 ip4:64.78.149.164 include:spf.mandrillapp.com include:aspmx.pardot.com include:mail.zendesk.com include:_spf.intacct.com include:mg-spf.greenhouse.io -all"
-          warning=$(go run cmd/spf-flatten/main.go --domain "abetterinternet.org" --initialSPF "${initialSPF}" --logLevel warn)
+          keep="include:_spf.intacct.com"
+          warning=$(go run cmd/spf-flatten/main.go --domain "abetterinternet.org" --initialSPF "${initialSPF}" --keep "${keep}" --logLevel warn)
           if [ -z "${warning}" ]; then
             echo "SPF record for abetterinternet.org has NOT changed"
             exit 0
@@ -26,10 +27,12 @@ jobs:
           fi
 
       - name: Check letsencrypt.org SPF
+        if: ${{ !cancelled() }}
         shell: bash
         run: |
           initialSPF="v=spf1 include:_spf.google.com ip4:23.178.112.0/24 ip4:66.133.109.36 ip4:64.78.149.164 include:spf.mandrillapp.com include:aspmx.pardot.com include:mail.zendesk.com include:shops.shopify.com include:_spf.intacct.com include:mg-spf.greenhouse.io -all"
-          warning=$(go run cmd/spf-flatten/main.go --domain "letsencrypt.org" --initialSPF "${initialSPF}" --logLevel warn)
+          keep="include:_spf.intacct.com"
+          warning=$(go run cmd/spf-flatten/main.go --domain "letsencrypt.org" --initialSPF "${initialSPF}" --keep "${keep}" --logLevel warn)
           if [ -z "${warning}" ]; then
             echo "SPF record for letsencrypt.org has NOT changed"
             exit 0

cmd/spf-flatten/main.go

@@ -16,6 +16,7 @@ type flags struct {
 	logLevel   slog.Level
 	dryrun     bool
 	warn       bool
+	keep       string
 	url        string
 	authEmail  string
 	authKey    string
@@ -28,14 +29,15 @@ func parseFlags() (flags, error) {
 	logLevelF := flag.String("logLevel", "LevelInfo", "")                            // optional
 	dryrunF := flag.Bool("dryrun", true, "")                                         // optional
 	warnF := flag.Bool("warn", true, "")                                             // optional TODO: come up with better name
+	keepF := flag.String("keep", "", "list SPF mechanisms to leave unflattened")     // optional
 	urlF := flag.String("url", "", "API URL for SPF record")                         // optional
 	authEmailF := flag.String("authEmail", "", "API key for X-Auth-Email header")    // optional unless dryrun is false
 	authKeyF := flag.String("authKey", "", "API key for X-Auth-Key header")          // optional unless dryrun is false
 	// TBD: should url, authEmail, and authKey be retrieved from os.Getenv instead of passed as flags?
 	// TBD: possible flags to add in the future: list of mechanisms to ignore or fail on, timeout/max recursions
 	flag.Parse()
 	f := flags{rootDomain: *rootDomainF, initialSPF: *ogRecordF, dryrun: *dryrunF, warn: *warnF,
-		url: *urlF, authEmail: *authEmailF, authKey: *authKeyF}
+		keep: *keepF, url: *urlF, authEmail: *authEmailF, authKey: *authKeyF}
 
 	// Require domain to be nonempty
 	if f.rootDomain == "" {
@@ -71,7 +73,7 @@ func main() {
 	slog.SetDefault(logger)
 
 	/// Flatten SPF record for input domain
-	r := spf.NewRootSPF(inputs.rootDomain, spf.NetLookup{})
+	r := spf.NewRootSPF(inputs.rootDomain, spf.NetLookup{}, inputs.keep)
 	if err = r.FlattenSPF(r.RootDomain, inputs.initialSPF); err != nil {
 		slog.Error("Could not flatten SPF record for initial domain", "error", err)
 		os.Exit(1)

internal/spf/flatten.go

@@ -38,13 +38,19 @@ func writeIPMech(ip net.IP, prefix string) string {
 type RootSPF struct {
 	RootDomain   string
 	AllMechanism string
+	MapKeep      map[string]bool
 	MapIPs       map[string]bool
 	MapNonflat   map[string]bool
 	LookupIF     Lookup
 }
 
-func NewRootSPF(rootDomain string, lookupIF Lookup) RootSPF {
-	return RootSPF{RootDomain: rootDomain, LookupIF: lookupIF, MapIPs: map[string]bool{}, MapNonflat: map[string]bool{}}
+func NewRootSPF(rootDomain string, lookupIF Lookup, keep string) RootSPF {
+	mapKeep := make(map[string]bool)
+	for _, mechanism := range strings.Fields(keep) {
+		mapKeep[mechanism] = true
+	}
+	return RootSPF{RootDomain: rootDomain, LookupIF: lookupIF, MapKeep: mapKeep,
+		MapIPs: map[string]bool{}, MapNonflat: map[string]bool{}}
 }
 
 var allInRecordRegex = regexp.MustCompile(`^.* (\+|-|~|\?)?all$`)
@@ -69,6 +75,11 @@ func (r *RootSPF) FlattenSPF(domain, spfRecord string) error {
 	}
 	containsAll := allInRecordRegex.MatchString(spfRecord)
 	for _, mechanism := range strings.Fields(spfRecord)[1:] {
+		// If mechanism is in string of "keep" mechanisms, add to MapNonflat and don't parse
+		if _, ok := r.MapKeep[mechanism]; ok {
+			r.MapNonflat[mechanism] = true
+			continue
+		}
 		// If not `all`, skip mechanism if fail modifier (- or ~) and ignore modifier otherwise
 		if modifierRegex.MatchString(mechanism[:1]) && !allRegex.MatchString(mechanism) {
 			if mechanism[:1] == "-" || mechanism[:1] == "~" {

internal/spf/flatten_test.go

@@ -61,7 +61,7 @@ func (r *RootSPF) compareExpected(err error, expAll string, expIPs, expNF []stri
 }
 
 func TestParseMechanismAll(t *testing.T) {
-	r := NewRootSPF("myrootdomain", mockLookup{})
+	r := NewRootSPF("myrootdomain", mockLookup{}, "")
 	// Test `all`` mechanism set if domain is root
 	err := r.ParseMechanism("~all", "myrootdomain")
 	if err = r.compareExpected(err, " ~all", []string{}, []string{}); err != nil {
@@ -75,7 +75,7 @@ func TestParseMechanismAll(t *testing.T) {
 }
 
 func TestParseMechanismIP(t *testing.T) {
-	r := NewRootSPF("", mockLookup{})
+	r := NewRootSPF("", mockLookup{}, "")
 	// Test ip mechanisms of the form `ip4:<ipaddr>`, `ip4:<ipaddr>/<prefix-length>, `ip6:<ipaddr>`, and `ip6:<ipaddr>/<prefix-length>`
 	ipMechs := []string{"ip4:abcd", "ip4:8.8.8.8", "ip6:efgh/36", "ip6:2001:0db8:85a3:0000:0000:8a2e:0370:7334", "ip6:11:22::33/128"}
 	for _, mech := range ipMechs {
@@ -87,7 +87,7 @@ func TestParseMechanismIP(t *testing.T) {
 }
 
 func TestParseMechanismA(t *testing.T) {
-	r := NewRootSPF("", mockLookup{})
+	r := NewRootSPF("", mockLookup{}, "")
 	// Test `a` mechanism of the form `a`, `a:<domain>`, `a/<prefix-length>`, and `a:<domain>/<prefix-length`
 	testInputs := [][]string{ //mechanism, prefix, aDomain, currentDomain
 		{"a", "", "mydomain", "mydomain"},
@@ -109,7 +109,7 @@ func TestParseMechanismA(t *testing.T) {
 }
 
 func TestParseMechanismMX(t *testing.T) {
-	r := NewRootSPF("", mockLookup{})
+	r := NewRootSPF("", mockLookup{}, "")
 	// Test `mx` mechanism of the form `mx`, `mx:<domain>`, `mx/<prefix-length>`, and `mx:<domain>/<prefix-length`
 	testInputs := [][]string{ // mechanism, prefix, mxDomain, currentDomain
 		{"mx", "", "anotherdomain", "anotherdomain"},
@@ -133,7 +133,7 @@ func TestParseMechanismMX(t *testing.T) {
 }
 
 func TestParseMechanismNonFlat(t *testing.T) {
-	r := NewRootSPF("", mockLookup{})
+	r := NewRootSPF("", mockLookup{}, "")
 	// Test ptr mechanism of the form `ptr`
 	err := r.ParseMechanism("ptr", "domain")
 	if err = r.compareExpected(err, "", []string{}, []string{"ptr:domain"}); err != nil {
@@ -150,7 +150,7 @@ func TestParseMechanismNonFlat(t *testing.T) {
 }
 
 func TestParseMechanismInclude(t *testing.T) {
-	r := NewRootSPF("", mockLookup{})
+	r := NewRootSPF("", mockLookup{}, "")
 	// Test mechanism of the form `include:<domain>`
 	includeDomain := "mydomain" // SPF record is just `a`, so expect IPs for "mydomain"
 	expIPs := []string{}
@@ -164,7 +164,7 @@ func TestParseMechanismInclude(t *testing.T) {
 }
 
 func TestParseMechanismRedirect(t *testing.T) {
-	r := NewRootSPF("", mockLookup{})
+	r := NewRootSPF("", mockLookup{}, "")
 	// Test mechanism of the form `redirect=<domain>`
 	redirectDomain := "test.com" // SPF record is just ip4:10.10.10.10
 	err := r.ParseMechanism("redirect="+redirectDomain, "notmydomain")
@@ -174,7 +174,7 @@ func TestParseMechanismRedirect(t *testing.T) {
 }
 
 func TestParseMechanismFails(t *testing.T) {
-	r := NewRootSPF("", mockLookup{})
+	r := NewRootSPF("", mockLookup{}, "")
 	// Test that parseMechanism fails on unexpected mechanism or syntax error
 	noMatchRegex := regexp.MustCompile(`^received unexpected SPF mechanism or syntax.*$`)
 	for _, wrongMech := range []string{"redirect:domain", "include=anotherdomain", "ip:0.0.0.0", "1.1.1.1", "", "ip6", "exp:explanation", "notMechanism:hello"} {
@@ -186,7 +186,7 @@ func TestParseMechanismFails(t *testing.T) {
 }
 
 func TestFlattenSPF(t *testing.T) {
-	r := NewRootSPF("", mockLookup{})
+	r := NewRootSPF("", mockLookup{}, "")
 	// Test that SPF record with multiple different entries gets flattened as expected
 	domain, spf := "abcd.net", "v=spf1 ip4:0.0.0.0/24 include:example.com ip6:56:0::0:7 a mx:test.domain exists:somedomain.edu -all"
 	// include:example.com --> (mydomain spf --> a --> mydomain IPs) + exp=exp.example.com
@@ -206,8 +206,35 @@ func TestFlattenSPF(t *testing.T) {
 	}
 }
 
+func TestNoFlattenKeeps(t *testing.T) {
+	// Test that non-keep mechanisms get flattened
+	r := NewRootSPF("", mockLookup{}, "include:example.com")
+	domain, spf := "abcd.net", "v=spf1 ip4:8.8.8.8/24 include:example.com ip6:56:0::0:7 a mx:test.domain exists:somedomain.edu -all"
+	r.RootDomain = domain
+	expIPs := []string{"ip4:8.8.8.8/24", "ip6:56:0::0:7"}
+	expNFs := []string{"exists:somedomain.edu", "include:example.com"}
+	for _, d := range []string{"abcd.net", "info.info"} {
+		for _, ip := range ipLookup[d] {
+			expIPs = append(expIPs, writeIPMech(ip, ""))
+		}
+	}
+	err := r.FlattenSPF(domain, spf)
+	if err = r.compareExpected(err, " -all", expIPs, expNFs); err != nil {
+		t.Fatal(err)
+	}
+	// Test that keep mechanism "include:example.com" not flattened
+	for _, ip := range ipLookup["mydomain"] {
+		if _, ok := r.MapIPs[writeIPMech(ip, "")]; ok {
+			t.Fatal("keep mechanism 'include:example.com' should not have been flattened")
+		}
+	}
+	if _, ok := r.MapNonflat["exp=exp.example.com"]; ok {
+		t.Fatal("keep mechanism 'include:example.com' should not have been flattened")
+	}
+}
+
 func TestFlattenRedirects(t *testing.T) {
-	r := NewRootSPF("", mockLookup{})
+	r := NewRootSPF("", mockLookup{}, "")
 	domain, spf := "somedomain", "v=spf1 ip4:9.9.9.9 redirect=mydomain ip6:2:2:2:2::::"
 	// Test that mechanisms after redirects are ignored
 	expIPs := []string{"ip4:9.9.9.9"}
@@ -234,7 +261,7 @@ func TestFlattenRedirects(t *testing.T) {
 
 func TestFlattenModifiers(t *testing.T) {
 	// Test modifier logic
-	r := NewRootSPF("", mockLookup{})
+	r := NewRootSPF("", mockLookup{}, "")
 	domain, spf := "somedomain", "v=spf1 -include:mydomain +ip4:9.8.7.6/54 ~exists:otherdomain ?mx:test.domain -all"
 	r.RootDomain = domain
 	expIPs := []string{"ip4:9.8.7.6/54"}

internal/spf/record.go

@@ -184,6 +184,9 @@ type PatchRequest struct {
 
 // PATCH the updated, flattened SPF record
 func UpdateSPFRecord(rootDomain, flatSPF, url, authEmail, authKey string) error {
+	if len(flatSPF) > 2048 {
+		return fmt.Errorf("SPF record is too long (got %d > 2048 characters)", len(flatSPF))
+	}
 	// {\n  \"content\": \"<SPF_RECORD>\",\n  \"name\": \"<DOMAIN>\",\n  \"type\": \"TXT\",\n  \"comment\": \"Dynamically updated, flattened SPF record\"}
 	patchReq := PatchRequest{
 		Content: flatSPF,

internal/spf/record_test.go

@@ -48,7 +48,7 @@ func TestGetLogLevel(t *testing.T) {
 }
 
 func TestGetDomainSPF(t *testing.T) {
-	r := NewRootSPF("", mockLookup{})
+	r := NewRootSPF("", mockLookup{}, "")
 	// Check that correctly filters TXT records for SPF record
 	if record, _ := GetDomainSPFRecord("mydomain", r.LookupIF); record != "v=spf1 a ~all" {
 		t.Fatal("Filtering for SPF record failed.")
@@ -61,7 +61,7 @@ func TestGetDomainSPF(t *testing.T) {
 }
 
 func TestCheckDomainSPF(t *testing.T) {
-	r := NewRootSPF("", mockLookup{})
+	r := NewRootSPF("", mockLookup{}, "")
 	// Check that returns error when SPF record does not match expected format
 	err := CheckSPFRecord("mydomain", "spf1 a ~all", r.LookupIF)
 	if !regexp.MustCompile("^.*did not match expected format.*$").MatchString(err.Error()) {
@@ -93,3 +93,13 @@ func TestCompareRecords(t *testing.T) {
 	}
 }
 
+func TestSPFTooLongError(t *testing.T) {
+	spfRecord := "v=spf1 " + strings.Repeat("a", 2042)
+	if err := UpdateSPFRecord("mydomain", spfRecord, "", "", ""); !strings.HasPrefix(err.Error(), "SPF record is too long") {
+		t.Fatalf("expected to fail since SPF record is too long")
+	}
+	spfRecord = "v=spf1 " + strings.Repeat("a", 2041)
+	if err := UpdateSPFRecord("mydomain", spfRecord, "", "", ""); strings.HasPrefix(err.Error(), "SPF record is too long") {
+		t.Fatalf("unexpected `SPF record is too long` error")
+	}
+}