git » chasquid » commit b153955

diff --git a/internal/dkim/header.go b/internal/dkim/header.go
index 680c693..79fb6c4 100644
--- a/internal/dkim/header.go
+++ b/internal/dkim/header.go
@@ -100,6 +100,21 @@ func (sig *dkimSignature) canonicalizationFromString(s string) error {
 	return nil
 }
 
+// Maximum number of entries in the h= tag that we accept.
+//
+// The RFC does not set a limit, but h= is attacker-controlled, so we cap it
+// to bound the work done when verifying. Signatures over the limit are
+// considered invalid (instead of truncating h=, which would also make them
+// fail, but in a more obscure way).
+//
+// 512 is chosen to be well above what legitimate signers use:
+//   - Signers typically use a few dozen entries at most (ours uses ~25).
+//   - OpenDKIM truncates h= to 4096 bytes, so signatures with a longer h=
+//     already fail verification there. With realistic header names (~10
+//     bytes each), 512 entries is above that.
+//   - Exim limits the number of headers it considers when verifying to 512.
+const maxHTagEntries = 512
+
 func (sig *dkimSignature) checkRequiredTags() error {
 	// https://datatracker.ietf.org/doc/html/rfc6376#section-6.1.1
 	if sig.a == "" {
@@ -121,6 +136,11 @@ func (sig *dkimSignature) checkRequiredTags() error {
 		return fmt.Errorf("%w: s=", errMissingRequiredTag)
 	}
 
+	if len(sig.h) > maxHTagEntries {
+		return fmt.Errorf("%w: h= has too many entries (%d > %d)",
+			errInvalidTag, len(sig.h), maxHTagEntries)
+	}
+
 	// h= must contain From.
 	var isFrom = func(s string) bool { return strings.EqualFold(s, "from") }
 	if !slices.ContainsFunc(sig.h, isFrom) {
diff --git a/internal/dkim/header_test.go b/internal/dkim/header_test.go
index dead777..4b91da8 100644
--- a/internal/dkim/header_test.go
+++ b/internal/dkim/header_test.go
@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"strconv"
+	"strings"
 	"testing"
 	"time"
 
@@ -189,6 +190,32 @@ func TestSignatureFromHeader(t *testing.T) {
 	}
 }
 
+func TestSignatureFromHeaderHTagLimit(t *testing.T) {
+	sigWithH := func(n int) string {
+		return "v=1; a=rsa-sha256; d=example.com; s=selector;" +
+			" b=aG9sYQ==; bh=aG9sYQ==;" +
+			" h=" + strings.Repeat("x:", n-1) + "from"
+	}
+
+	// At the limit is fine.
+	sig, err := dkimSignatureFromHeader(sigWithH(maxHTagEntries))
+	if err != nil {
+		t.Fatalf("h= with %d entries: unexpected error: %v",
+			maxHTagEntries, err)
+	}
+	if len(sig.h) != maxHTagEntries {
+		t.Errorf("h= with %d entries: got %d entries",
+			maxHTagEntries, len(sig.h))
+	}
+
+	// Over the limit is not.
+	sig, err = dkimSignatureFromHeader(sigWithH(maxHTagEntries + 1))
+	if sig != nil || !errors.Is(err, errInvalidTag) {
+		t.Errorf("h= with %d entries: wanted nil / %v, got %v / %v",
+			maxHTagEntries+1, errInvalidTag, sig, err)
+	}
+}
+
 func TestCanonicalizationFromString(t *testing.T) {
 	cases := []struct {
 		in     string
diff --git a/internal/dkim/sign.go b/internal/dkim/sign.go
index ac9ce62..1d89394 100644
--- a/internal/dkim/sign.go
+++ b/internal/dkim/sign.go
@@ -75,7 +75,9 @@ func (s *Signer) Sign(ctx context.Context, message string) (string, error) {
 	hsForHeader := []string{}
 	for _, h := range headersToSign {
 		// Append the header as many times as it appears in the message.
-		for i := 0; i < len(headers.FindAll(h)); i++ {
+		// Note the range expression is evaluated only once; evaluating
+		// FindAll on every iteration would make this quadratic.
+		for range len(headers.FindAll(h)) {
 			hsForHeader = append(hsForHeader, h)
 		}
 	}
diff --git a/internal/dkim/sign_test.go b/internal/dkim/sign_test.go
index 07313e1..7184a24 100644
--- a/internal/dkim/sign_test.go
+++ b/internal/dkim/sign_test.go
@@ -129,6 +129,42 @@ func TestSignEd25519(t *testing.T) {
 	}
 }
 
+func TestSignManyHeaders(t *testing.T) {
+	// Regression test for signing being quadratic on the number of instances
+	// of a header. This takes milliseconds, while the quadratic
+	// implementation took minutes with this input. We don't check the
+	// timing, so a regression will show up as this test being very slow.
+	const n = 50_000
+
+	_, priv, err := ed25519.GenerateKey(rand.Reader)
+	if err != nil {
+		t.Fatalf("ed25519.GenerateKey: %v", err)
+	}
+	s := &Signer{
+		Domain:   "example.com",
+		Selector: "test",
+		Signer:   priv,
+	}
+
+	message := "From: from@example.com\r\n" +
+		strings.Repeat("To: to@example.com\r\n", n) +
+		"\r\nHi.\r\n"
+
+	sig, err := s.Sign(t.Context(), message)
+	if err != nil {
+		t.Fatalf("Sign: %v", err)
+	}
+
+	// All instances must be in h=, plus the extra one we always add.
+	hTag := regexp.MustCompile(`h=([^;]*);`).FindStringSubmatch(sig)
+	if hTag == nil {
+		t.Fatalf("h= not found in signature: %q", sig)
+	}
+	if got := strings.Count(hTag[1], "to"); got != n+1 {
+		t.Errorf("wanted %d instances of 'to' in h=, got %d", n+1, got)
+	}
+}
+
 func addSig(sig, message string) string {
 	return "DKIM-Signature: " +
 		strings.ReplaceAll(sig, "\r\n", "\r\n\t") +
diff --git a/internal/dkim/verify.go b/internal/dkim/verify.go
index 00b5107..b8b5d6c 100644
--- a/internal/dkim/verify.go
+++ b/internal/dkim/verify.go
@@ -285,24 +285,28 @@ func headersToInclude(sigH header, hTag []string, headers headers) []header {
 	//    https://datatracker.ietf.org/doc/html/rfc6376#section-3.7
 	//  - Headers may be missing, and that's allowed.
 	//    https://datatracker.ietf.org/doc/html/rfc6376#section-5.4
-	seen := map[string]int{}
+	//
+	// Both h= and the headers are attacker-controlled, so this must be
+	// linear on their sizes: we index the headers by (lowercased) name once,
+	// with the last instance first, and then for every entry in h= we pick
+	// (and consume) the next available instance.
+	byName := map[string][]header{}
+	for _, h := range slices.Backward(headers) {
+		name := strings.ToLower(h.Name)
+		byName[name] = append(byName[name], h)
+	}
+
 	include := []header{}
 	for _, h := range hTag {
-		all := headers.FindAll(h)
-		slices.Reverse(all)
-
-		// We keep track of the last instance of each header that we
-		// included, and find the next one every time it appears in h=.
-		// We have to be careful because the header itself may not be present,
-		// or we may be asked to include it more times than it appears.
+		// The header may not be present, or we may be asked to include it
+		// more times than it appears.
 		lh := strings.ToLower(h)
-		i := seen[lh]
-		if i >= len(all) {
+		available := byName[lh]
+		if len(available) == 0 {
 			continue
 		}
-		seen[lh]++
-
-		selected := all[i]
+		selected := available[0]
+		byName[lh] = available[1:]
 
 		if selected == sigH {
 			continue
diff --git a/internal/dkim/verify_test.go b/internal/dkim/verify_test.go
index 4bf718b..01063f8 100644
--- a/internal/dkim/verify_test.go
+++ b/internal/dkim/verify_test.go
@@ -5,6 +5,7 @@ import (
 	"crypto"
 	"encoding/base64"
 	"errors"
+	"fmt"
 	"net"
 	"strings"
 	"testing"
@@ -341,6 +342,26 @@ func TestHeadersToInclude(t *testing.T) {
 				},
 			},
 		},
+
+		// Check that header names are compared case-insensitively, and that
+		// headers that are not present are skipped.
+		{
+			sigH: header{
+				Name:  "DKIM-Signature",
+				Value: "v=1; a=rsa-sha256; s=brisbane; d=example.com;",
+			},
+			hTag: []string{"from", "X-Missing", "FROM", "subject"},
+			headers: headers{
+				{Name: "FROM", Value: "from1"},
+				{Name: "Subject", Value: "subject1"},
+				{Name: "From", Value: "from2"},
+			},
+			want: []header{
+				{Name: "From", Value: "from2"},
+				{Name: "FROM", Value: "from1"},
+				{Name: "Subject", Value: "subject1"},
+			},
+		},
 	}
 
 	for _, c := range cases {
@@ -352,6 +373,33 @@ func TestHeadersToInclude(t *testing.T) {
 	}
 }
 
+func TestHeadersToIncludeLarge(t *testing.T) {
+	// Regression test for a DoS where headersToInclude was O(len(h=) *
+	// len(headers)), which with large (attacker-controlled) inputs could take
+	// hours. This takes milliseconds, while the quadratic implementation
+	// took minutes with this input. We don't check the timing, so a
+	// regression will show up as this test being very slow.
+	const n = 100_000
+
+	sigH := header{Name: "DKIM-Signature", Value: "v=1;"}
+	hTag := make([]string, n)
+	hdrs := make(headers, n)
+	for i := range n {
+		hTag[i] = "a"
+		hdrs[i] = header{Name: "a", Value: fmt.Sprint(i)}
+	}
+
+	got := headersToInclude(sigH, hTag, hdrs)
+	if len(got) != n {
+		t.Fatalf("wanted %d headers, got %d", n, len(got))
+	}
+	// Last instance first.
+	if got[0].Value != fmt.Sprint(n-1) || got[n-1].Value != "0" {
+		t.Errorf("unexpected order: first %q, last %q",
+			got[0].Value, got[n-1].Value)
+	}
+}
+
 func TestAuthenticationResults(t *testing.T) {
 	resBrisbane := &OneResult{
 		Domain:   "football.example.com",