#19063: fix set_payload handling of non-ASCII string input.

This version of the fix raises an error instead of accepting the invalid
input (ie: if a non-ASCII string is used but no charset is specified).

2 files changed, 29 insertions, 22 deletions

diff --git a/Lib/email/charset.py b/Lib/email/charset.py
index 892bab5..51f24b2 100644
--- a/Lib/email/charset.py
+++ b/Lib/email/charset.py
@@ -378,18 +378,19 @@ class Charset:
             return None
 
     def body_encode(self, string):
-        """Body-encode a string by converting it first to bytes.
+        """Body-encode a string, converting it first to bytes if needed.
 
         The type of encoding (base64 or quoted-printable) will be based on
-        self.body_encoding.  If body_encoding is None, we assume the
-        output charset is a 7bit encoding, so re-encoding the decoded
-        string using the ascii codec produces the correct string version
-        of the content.
+        self.body_encoding.  If body_encoding is None, we perform no CTE
+        encoding (the CTE will be either 7bit or 8bit), we just encode the
+        binary representation to ascii using the surrogateescape error handler,
+        which will enable the Generators to produce the correct output.
         """
-        # 7bit/8bit encodings return the string unchanged (module conversions)
+        if not string:
+            return string
+        if isinstance(string, str):
+            string = string.encode(self.output_charset)
         if self.body_encoding is BASE64:
-            if isinstance(string, str):
-                string = string.encode(self.output_charset)
             return email.base64mime.body_encode(string)
         elif self.body_encoding is QP:
             # quopromime.body_encode takes a string, but operates on it as if
@@ -398,15 +399,7 @@ class Charset:
             # character set, then, we must turn it into pseudo bytes via the
             # latin1 charset, which will encode any byte as a single code point
             # between 0 and 255, which is what body_encode is expecting.
-            #
-            # Note that this clause doesn't handle the case of a _payload that
-            # is already bytes.  It never did, and the semantics of _payload
-            # being bytes has never been nailed down, so fixing that is a
-            # longer term TODO.
-            if isinstance(string, str):
-                string = string.encode(self.output_charset).decode('latin1')
+            string = string.decode('latin1')
             return email.quoprimime.body_encode(string)
         else:
-            if isinstance(string, str):
-                string = string.encode(self.output_charset).decode('ascii')
-            return string
+            return string.decode('ascii', 'surrogateescape')
diff --git a/Lib/email/message.py b/Lib/email/message.py
index ce673b0..9d295fc 100644
--- a/Lib/email/message.py
+++ b/Lib/email/message.py
@@ -301,9 +301,23 @@ class Message:
         Optional charset sets the message's default character set.  See
         set_charset() for details.
         """
-        if isinstance(payload, bytes):
-            payload = payload.decode('ascii', 'surrogateescape')
-        self._payload = payload
+        if hasattr(payload, 'encode'):
+            if charset is None:
+                try:
+                    payload.encode('ascii', 'surrogateescape')
+                except UnicodeError:
+                    raise TypeError("charset argument must be specified"
+                        " when non-ASCII characters are used in the"
+                        " payload") from None
+                self._payload = payload
+                return
+            if not isinstance(charset, Charset):
+                charset = Charset(charset)
+            payload = payload.encode(charset.output_charset)
+        if hasattr(payload, 'decode'):
+            self._payload = payload.decode('ascii', 'surrogateescape')
+        else:
+            self._payload = payload
         if charset is not None:
             self.set_charset(charset)
 
@@ -342,7 +356,7 @@ class Message:
             try:
                 cte(self)
             except TypeError:
-                self._payload = charset.body_encode(self._payload)
+                self._payload = charset.body_encode(self.get_payload(decode=True))
                 self.add_header('Content-Transfer-Encoding', cte)
 
     def get_charset(self):