fix: binary mode for siphon and AtomicFileReplacement

Signed-off-by: Martin Styk <mart.styk@gmail.com>

Author: StykMartin

Common/bkr/common/helpers.py

@@ -135,9 +135,10 @@ class AtomicFileReplacement(object):
     and replace_dest can also be called directly if needed
     """
 
-    def __init__(self, dest_path, mode=0o644):
+    def __init__(self, dest_path, mode=0o644, binary=False):
         self.dest_path = dest_path
         self.mode = mode
+        self.binary = binary
         self._temp_info = None
 
     @property
@@ -152,7 +153,7 @@ def create_temp(self):
         dirname, basename = os.path.split(self.dest_path)
         fd, temp_path = tempfile.mkstemp(prefix='.' + basename, dir=dirname)
         try:
-            f = os.fdopen(fd, 'w')
+            f = os.fdopen(fd, 'wb' if self.binary else 'w')
         except:
             os.unlink(temp_path)
             raise
@@ -247,7 +248,8 @@ def siphon(src, dest):
             break
 
         if six.PY3 and isinstance(chunk, bytes):
-            chunk = chunk.decode('utf-8')
+            if not hasattr(dest, 'mode') or 'b' not in dest.mode:
+                chunk = chunk.decode('utf-8')
 
         dest.write(chunk)
 

Common/bkr/common/test_helpers.py

@@ -3,8 +3,14 @@
 # the Free Software Foundation; either version 2 of the License, or
 # (at your option) any later version.
 
+import os
+import shutil
+import tempfile
 import unittest
 
+import six
+
+from bkr.common.helpers import AtomicFileReplacement, siphon
 from bkr.common.helpers_six import parse_content_type
 
 
@@ -16,3 +22,48 @@ def test_ok(self):
 
     def test_empty(self):
         self.assertEqual(parse_content_type(''), '')
+
+
+class TestSiphon(unittest.TestCase):
+
+    def test_siphon_text_to_text_file(self):
+        src = six.StringIO(u'hello world')
+        dest = six.StringIO()
+        siphon(src, dest)
+        self.assertEqual(dest.getvalue(), u'hello world')
+
+    def test_siphon_binary_to_binary_file(self):
+        binary_data = b'\x00\x01\xcd\xfe\xff' * 1000
+        src = six.BytesIO(binary_data)
+        tmp = tempfile.NamedTemporaryFile(mode='wb', delete=False)
+        try:
+            siphon(src, tmp)
+            tmp.close()
+            with open(tmp.name, 'rb') as f:
+                self.assertEqual(f.read(), binary_data)
+        finally:
+            os.unlink(tmp.name)
+
+
+class TestAtomicFileReplacement(unittest.TestCase):
+
+    def setUp(self):
+        self.tmpdir = tempfile.mkdtemp()
+
+    def tearDown(self):
+        shutil.rmtree(self.tmpdir)
+
+    def test_text_mode(self):
+        dest_path = os.path.join(self.tmpdir, 'textfile')
+        with AtomicFileReplacement(dest_path) as f:
+            f.write(u'hello')
+        with open(dest_path, 'r') as f:
+            self.assertEqual(f.read(), 'hello')
+
+    def test_binary_mode(self):
+        binary_data = b'\x00\x01\xcd\xfe\xff'
+        dest_path = os.path.join(self.tmpdir, 'binfile')
+        with AtomicFileReplacement(dest_path, binary=True) as f:
+            f.write(binary_data)
+        with open(dest_path, 'rb') as f:
+            self.assertEqual(f.read(), binary_data)

LabController/src/bkr/labcontroller/netboot.py

@@ -110,7 +110,7 @@ def copy_default_loader_images():
 def fetch_bootloader_image(fqdn, fqdn_dir, distro_tree_id, image_url):
     timeout = get_conf().get('IMAGE_FETCH_TIMEOUT')
     logger.debug('Fetching bootloader image %s for %s', image_url, fqdn)
-    with atomically_replaced_file(os.path.join(fqdn_dir, 'image')) as dest:
+    with atomically_replaced_file(os.path.join(fqdn_dir, 'image'), binary=True) as dest:
         try:
             siphon(urllib.request.urlopen(image_url, timeout=timeout), dest)
         except Exception as e:
@@ -146,13 +146,13 @@ def fetch_images(distro_tree_id, kernel_url, initrd_url, fqdn):
 
     timeout = get_conf().get('IMAGE_FETCH_TIMEOUT')
     logger.debug('Fetching kernel %s for %s', kernel_url, fqdn)
-    with atomically_replaced_file(os.path.join(images_dir, 'kernel')) as dest:
+    with atomically_replaced_file(os.path.join(images_dir, 'kernel'), binary=True) as dest:
         try:
             siphon(urllib.request.urlopen(kernel_url, timeout=timeout), dest)
         except Exception as e:
             raise ImageFetchingError(kernel_url, distro_tree_id, e)
     logger.debug('Fetching initrd %s for %s', initrd_url, fqdn)
-    with atomically_replaced_file(os.path.join(images_dir, 'initrd')) as dest:
+    with atomically_replaced_file(os.path.join(images_dir, 'initrd'), binary=True) as dest:
         try:
             siphon(urllib.request.urlopen(initrd_url, timeout=timeout), dest)
         except Exception as e:

LabController/src/bkr/labcontroller/pxemenu.py

@@ -88,7 +88,7 @@ def _get_images(tftp_root, distro_tree_id, url, images):
             else:
                 image_url = urllib.parse.urljoin(url, path)
                 print('Fetching %s %s for distro tree %s' % (image_type, image_url, distro_tree_id))
-                with atomically_replaced_file(dest_path) as dest:
+                with atomically_replaced_file(dest_path, binary=True) as dest:
                     siphon(urllib.request.urlopen(image_url), dest)