git » chasquid » commit ce236e9

diff --git a/internal/dkim/message.go b/internal/dkim/message.go
index 5ad3a7a..a8da856 100644
--- a/internal/dkim/message.go
+++ b/internal/dkim/message.go
@@ -32,8 +32,11 @@ var errInvalidHeader = errors.New("invalid header")
 // Does NOT touch whitespace, this is important to preserve the original
 // message and headers, which is required for the signature.
 func parseMessage(message string) (headers, string, error) {
-	headers := make(headers, 0)
 	lines := strings.Split(message, "\r\n")
+
+	// Group the lines of each header: the first line, followed by its
+	// continuation lines (if any).
+	groups := [][]string{}
 	eoh := 0
 	for i, line := range lines {
 		if line == "" {
@@ -43,35 +46,43 @@ func parseMessage(message string) (headers, string, error) {
 
 		if strings.HasPrefix(line, " ") || strings.HasPrefix(line, "\t") {
 			// Continuation of the previous header.
-			if len(headers) == 0 {
+			if len(groups) == 0 {
 				return nil, "", fmt.Errorf(
 					"%w: bad continuation", errInvalidHeader)
 			}
-			headers[len(headers)-1].Value += "\r\n" + line
-			headers[len(headers)-1].Source += "\r\n" + line
+			last := len(groups) - 1
+			groups[last] = append(groups[last], line)
 		} else {
 			// New header.
-			h, err := parseHeader(line)
-			if err != nil {
-				return nil, "", err
-			}
+			groups = append(groups, []string{line})
+		}
+	}
 
-			headers = append(headers, h)
+	headers := make(headers, 0, len(groups))
+	for _, g := range groups {
+		h, err := parseHeader(g)
+		if err != nil {
+			return nil, "", err
 		}
+		headers = append(headers, h)
 	}
 
 	return headers, strings.Join(lines[eoh+1:], "\r\n"), nil
 }
 
-func parseHeader(line string) (header, error) {
-	name, value, found := strings.Cut(line, ":")
-	if !found {
+// Parse a header, given its lines (the first one, followed by the
+// continuation lines, if any).
+func parseHeader(lines []string) (header, error) {
+	// The colon must be in the first line.
+	if !strings.Contains(lines[0], ":") {
 		return header{}, fmt.Errorf("%w: no colon", errInvalidHeader)
 	}
 
+	source := strings.Join(lines, "\r\n")
+	name, value, _ := strings.Cut(source, ":")
 	return header{
 		Name:   name,
 		Value:  value,
-		Source: line,
+		Source: source,
 	}, nil
 }
diff --git a/internal/dkim/message_test.go b/internal/dkim/message_test.go
index 0e28b7d..9006ea4 100644
--- a/internal/dkim/message_test.go
+++ b/internal/dkim/message_test.go
@@ -1,6 +1,7 @@
 package dkim
 
 import (
+	"strings"
 	"testing"
 
 	"blitiri.com.ar/go/chasquid/internal/normalize"
@@ -34,6 +35,34 @@ body`),
 			},
 			body: "body",
 		},
+		{
+			// Multiple continuation lines, with spaces and tabs, in a header
+			// that is not the last one. Colons in continuation lines are
+			// part of the value.
+			message: normalize.StringToCRLF(`From: a@b
+Multi: one
+ two:three
+	four
+To: c@d
+
+body`),
+			headers: headers{
+				header{Name: "From", Value: " a@b",
+					Source: "From: a@b"},
+				header{Name: "Multi",
+					Value:  " one\r\n two:three\r\n\tfour",
+					Source: "Multi: one\r\n two:three\r\n\tfour"},
+				header{Name: "To", Value: " c@d",
+					Source: "To: c@d"},
+			},
+			body: "body",
+		},
+		{
+			// Empty message.
+			message: "",
+			headers: headers{},
+			body:    "",
+		},
 	}
 
 	for i, c := range cases {
@@ -68,6 +97,12 @@ func TestParseMessageWithErrors(t *testing.T) {
 			message: "No colon",
 			err:     errInvalidHeader,
 		},
+		{
+			// Header without ':' in its first line, even if there's one in
+			// a continuation line.
+			message: "No colon\r\n but: here",
+			err:     errInvalidHeader,
+		},
 	}
 
 	for i, c := range cases {
@@ -79,6 +114,54 @@ func TestParseMessageWithErrors(t *testing.T) {
 	}
 }
 
+func TestParseMessageWithoutBlankLine(t *testing.T) {
+	// If there's no blank line, all the lines are headers, including the
+	// continuation lines of the last one.
+	hs, _, err := parseMessage("From: a@b\r\nTo: c@d\r\n e@f")
+	if err != nil {
+		t.Fatalf("parseMessage error: %v", err)
+	}
+	want := headers{
+		{Name: "From", Value: " a@b", Source: "From: a@b"},
+		{Name: "To", Value: " c@d\r\n e@f", Source: "To: c@d\r\n e@f"},
+	}
+	if diff := cmp.Diff(want, hs); diff != "" {
+		t.Errorf("parseMessage headers mismatch (-want +got):\n%s", diff)
+	}
+}
+
+func TestParseMessageManyContinuations(t *testing.T) {
+	// Regression test for parsing being quadratic on the number of
+	// continuation lines 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 = 1_000_000
+	value := strings.Repeat("\r\n a", n)
+	message := "From: a@b\r\nLong:" + value + "\r\n\r\nbody"
+
+	headers, body, err := parseMessage(message)
+	if err != nil {
+		t.Fatalf("parseMessage error: %v", err)
+	}
+	if len(headers) != 2 {
+		t.Fatalf("wanted 2 headers, got %d", len(headers))
+	}
+	if headers[1].Name != "Long" {
+		t.Fatalf("unexpected header name: %q", headers[1].Name)
+	}
+	if headers[1].Value != value {
+		t.Errorf("unexpected value (len %d, want len %d)",
+			len(headers[1].Value), len(value))
+	}
+	if headers[1].Source != "Long:"+value {
+		t.Errorf("unexpected source (len %d, want len %d)",
+			len(headers[1].Source), len("Long:"+value))
+	}
+	if body != "body" {
+		t.Errorf("unexpected body: %q", body)
+	}
+}
+
 func TestHeadersFindAll(t *testing.T) {
 	hs := headers{
 		{Name: "From", Value: "a@b", Source: "From: a@b"},