Add util.message_from_string to wrap email.message_from_string.

This is to work-around the fact that email.message_from_string uses
cStringIO in Python 2.6, which can't handle Unicode.

3 files changed, 15 insertions, 3 deletions

diff --git a/cloudinit/user_data.py b/cloudinit/user_data.py
index 663a9048..eb3c7336 100644
--- a/cloudinit/user_data.py
+++ b/cloudinit/user_data.py
@@ -22,8 +22,6 @@
 
 import os
 
-import email
-
 from email.mime.base import MIMEBase
 from email.mime.multipart import MIMEMultipart
 from email.mime.nonmultipart import MIMENonMultipart
@@ -338,7 +336,7 @@ def convert_string(raw_data, headers=None):
         headers = {}
     data = util.decode_binary(util.decomp_gzip(raw_data))
     if "mime-version:" in data[0:4096].lower():
-        msg = email.message_from_string(data)
+        msg = util.message_from_string(data)
         for (key, val) in headers.items():
             _replace_header(msg, key, val)
     else:
diff --git a/cloudinit/util.py b/cloudinit/util.py
index b6065410..971c1c2d 100644
--- a/cloudinit/util.py
+++ b/cloudinit/util.py
@@ -23,6 +23,7 @@
 import contextlib
 import copy as obj_copy
 import ctypes
+import email
 import errno
 import glob
 import grp
@@ -2187,3 +2188,9 @@ def read_dmi_data(key):
     LOG.warn("did not find either path %s or dmidecode command",
              DMI_SYS_PATH)
     return None
+
+
+def message_from_string(string):
+    if sys.version_info[:2] < (2, 7):
+        return email.message_from_file(six.StringIO(string))
+    return email.message_from_string(string)
diff --git a/tests/unittests/test_util.py b/tests/unittests/test_util.py
index 7da1f755..1619b5d2 100644
--- a/tests/unittests/test_util.py
+++ b/tests/unittests/test_util.py
@@ -452,4 +452,11 @@ class TestMultiLog(helpers.FilesystemMockingTestCase):
         util.multi_log('message', log=log, log_level=log_level)
         self.assertEqual((log_level, mock.ANY), log.log.call_args[0])
 
+
+class TestMessageFromString(helpers.TestCase):
+
+    def test_unicode_not_messed_up(self):
+        roundtripped = util.message_from_string(u'\n').as_string()
+        self.assertNotIn('\x00', roundtripped)
+
 # vi: ts=4 expandtab