From 816fefe10e5db7c282b98004ec09c3e84e20833c Mon Sep 17 00:00:00 2001
From: Simon Ser <contact@emersion.fr>
Date: Tue, 24 Nov 2020 12:05:22 +0100
Subject: [PATCH] dkim: fix CRLF split between two Write calls

Closes: https://github.com/emersion/go-msgauth/issues/38
---
 dkim/canonical.go      | 41 +++++++++++++++++++++++++++--------------
 dkim/canonical_test.go | 21 +++++++++++++++++++++
 2 files changed, 48 insertions(+), 14 deletions(-)

diff --git a/dkim/canonical.go b/dkim/canonical.go
index 440fe15..acef8a1 100644
--- a/dkim/canonical.go
+++ b/dkim/canonical.go
@@ -26,14 +26,25 @@ var canonicalizers = map[Canonicalization]canonicalizer{
 	CanonicalizationRelaxed: new(relaxedCanonicalizer),
 }
 
-// Fix any \n without a matching \r
-func fixCRLF(b []byte) []byte {
+// crlfFixer fixes any lone LF without a preceding CR.
+type crlfFixer struct {
+	cr bool
+}
+
+func (cf *crlfFixer) Fix(b []byte) []byte {
 	res := make([]byte, 0, len(b))
-	for i := range b {
-		if b[i] == '\n' && (i == 0 || b[i-1] != '\r') {
-			res = append(res, '\r')
+	for _, ch := range b {
+		prevCR := cf.cr
+		cf.cr = false
+		switch ch {
+		case '\r':
+			cf.cr = true
+		case '\n':
+			if !prevCR {
+				res = append(res, '\r')
+			}
 		}
-		res = append(res, b[i])
+		res = append(res, ch)
 	}
 	return res
 }
@@ -45,15 +56,16 @@ func (c *simpleCanonicalizer) CanonicalizeHeader(s string) string {
 }
 
 type simpleBodyCanonicalizer struct {
-	w       io.Writer
-	crlfBuf []byte
+	w         io.Writer
+	crlfBuf   []byte
+	crlfFixer crlfFixer
 }
 
 func (c *simpleBodyCanonicalizer) Write(b []byte) (int, error) {
 	written := len(b)
 	b = append(c.crlfBuf, b...)
 
-	b = fixCRLF(b)
+	b = c.crlfFixer.Fix(b)
 
 	end := len(b)
 	// If it ends with \r, maybe the next write will begin with \n
@@ -116,16 +128,17 @@ func (c *relaxedCanonicalizer) CanonicalizeHeader(s string) string {
 }
 
 type relaxedBodyCanonicalizer struct {
-	w       io.Writer
-	crlfBuf []byte
-	wspBuf  []byte
-	written bool
+	w         io.Writer
+	crlfBuf   []byte
+	wspBuf    []byte
+	written   bool
+	crlfFixer crlfFixer
 }
 
 func (c *relaxedBodyCanonicalizer) Write(b []byte) (int, error) {
 	written := len(b)
 
-	b = fixCRLF(b)
+	b = c.crlfFixer.Fix(b)
 
 	canonical := make([]byte, 0, len(b))
 	for _, ch := range b {
diff --git a/dkim/canonical_test.go b/dkim/canonical_test.go
index 24112ac..e9ad9ba 100644
--- a/dkim/canonical_test.go
+++ b/dkim/canonical_test.go
@@ -168,3 +168,24 @@ func TestRelaxedCanonicalizer_CanonicalBody(t *testing.T) {
 		}
 	}
 }
+
+func TestRelaxedCanonicalizer_CanonicalBody_splitCRLF(t *testing.T) {
+	want := "line 1\r\nline 2\r\n"
+	writes := [][]byte{
+		[]byte("line 1\r"),
+		[]byte("\nline 2"),
+	}
+
+	var b bytes.Buffer
+	wc := new(relaxedCanonicalizer).CanonicalizeBody(&b)
+	for _, b := range writes {
+		if _, err := wc.Write(b); err != nil {
+			t.Errorf("Expected no error while writing to relaxed body canonicalizer, got: %v", err)
+		}
+	}
+	if err := wc.Close(); err != nil {
+		t.Errorf("Expected no error while closing relaxed body canonicalizer, got: %v", err)
+	} else if s := b.String(); s != want {
+		t.Errorf("Expected canonical body to be %q, but got %q", want, s)
+	}
+}