Fix relaxed body canoncalization & introduce further unit tests for both simple and relaxed mode

dkim/canonicalization.py

@@ -48,6 +48,13 @@ def unfold_header_value(content):
     return re.sub(b"\r\n", b"", content)
 
 
+def correct_empty_body(content):
+    if content == b"\r\n":
+        return ""
+    else:
+        return content
+
+
 class Simple:
     """Class that represents the "simple" canonicalization algorithm."""
 
@@ -85,8 +92,8 @@ class Relaxed:
         # Remove all trailing WSP at end of lines.
         # Compress non-line-ending WSP to single space.
         # Ignore all empty lines at the end of the message body.
-        return strip_trailing_lines(
+        return correct_empty_body(strip_trailing_lines(
-            compress_whitespace(strip_trailing_whitespace(body)))
+            compress_whitespace(strip_trailing_whitespace(body))))
 
 
 class CanonicalizationPolicy:

dkim/tests/test_canonicalization.py

@@ -53,6 +53,26 @@ class TestSimpleAlgorithmBody(BaseCanonicalizationTest):
             b'Foo  \tbar    \r\n',
             b'Foo  \tbar    \r\n\r\n')
 
+    def test_adds_crlf(self):
+        self.assertCanonicalForm(
+            b'Foo bar\r\n',
+            b'Foo bar')
+
+    def test_empty_body(self):
+        self.assertCanonicalForm(
+            b'\r\n',
+            b'')
+
+    def test_single_crlf_body(self):
+        self.assertCanonicalForm(
+            b'\r\n',
+            b'\r\n')
+
+    def test_multiple_crlf_body(self):
+        self.assertCanonicalForm(
+            b'\r\n',
+            b'\r\n\r\n')
+
 
 class TestRelaxedAlgorithmHeaders(BaseCanonicalizationTest):
 
@@ -98,6 +118,26 @@ class TestRelaxedAlgorithmBody(BaseCanonicalizationTest):
             b'Foo\r\nbar\r\n',
             b'Foo\r\nbar\r\n\r\n\r\n')
 
+    def test_adds_crlf(self):
+        self.assertCanonicalForm(
+            b'Foo bar\r\n',
+            b'Foo bar')
+
+    def test_empty_body(self):
+        self.assertCanonicalForm(
+            b'',
+            b'')
+
+    def test_single_crlf_body(self):
+        self.assertCanonicalForm(
+            b'',
+            b'\r\n')
+
+    def test_multiple_crlf_body(self):
+        self.assertCanonicalForm(
+            b'',
+            b'\r\n\r\n')
+
 
 class TestCanonicalizationPolicyFromCValue(unittest.TestCase):