osutils: WriteFile: support append+sudo using dd

BUG=chromium:1050646
TEST=CQ passes

Change-Id: I3eb97c0cf35e465a4ffa8166ff3150732c619955
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/2067184
Commit-Queue: Mike Frysinger <vapier@chromium.org>
Tested-by: Mike Frysinger <vapier@chromium.org>
Reviewed-by: Andrew Lassalle <andrewlassalle@chromium.org>
diff --git a/lib/osutils.py b/lib/osutils.py
index c463bba..8814b9c 100644
--- a/lib/osutils.py
+++ b/lib/osutils.py
@@ -140,8 +140,8 @@
     raise ValueError('mode must be one of {"%s"}, not %r' %
                      ('", "'.join(sorted(_VALID_WRITE_MODES)), mode))
 
-  if sudo and ('a' in mode or '+' in mode):
-    raise ValueError('append mode does not work in sudo mode')
+  if sudo and atomic and ('a' in mode or '+' in mode):
+    raise ValueError('append mode does not work in sudo+atomic mode')
 
   if 'b' in mode:
     if encoding is not None or errors is not None:
@@ -168,24 +168,33 @@
   # If the file needs to be written as root and we are not root, write to a temp
   # file, move it and change the permission.
   if sudo and os.getuid() != 0:
-    with tempfile.NamedTemporaryFile(mode=mode, delete=False) as temp:
-      write_path = temp.name
-      temp.writelines(write_wrapper(cros_build_lib.iflatten_instance(content)))
-    os.chmod(write_path, 0o644)
+    if 'a' in mode or '+' in mode:
+      # Use dd to run through sudo & append the output, and write the new data
+      # to it through stdin.
+      cros_build_lib.sudo_run(
+          ['dd', 'conv=notrunc', 'oflag=append', 'status=none',
+           'of=%s' % (path,)], print_cmd=False, input=content)
 
-    try:
-      mv_target = path if not atomic else path + '.tmp'
-      cros_build_lib.sudo_run(['mv', write_path, mv_target],
-                              print_cmd=False, stderr=True)
-      Chown(mv_target, user='root', group='root')
-      if atomic:
-        cros_build_lib.sudo_run(['mv', mv_target, path],
+    else:
+      with tempfile.NamedTemporaryFile(mode=mode, delete=False) as temp:
+        write_path = temp.name
+        temp.writelines(write_wrapper(
+            cros_build_lib.iflatten_instance(content)))
+      os.chmod(write_path, 0o644)
+
+      try:
+        mv_target = path if not atomic else path + '.tmp'
+        cros_build_lib.sudo_run(['mv', write_path, mv_target],
                                 print_cmd=False, stderr=True)
+        Chown(mv_target, user='root', group='root')
+        if atomic:
+          cros_build_lib.sudo_run(['mv', mv_target, path],
+                                  print_cmd=False, stderr=True)
 
-    except cros_build_lib.RunCommandError:
-      SafeUnlink(write_path)
-      SafeUnlink(mv_target)
-      raise
+      except cros_build_lib.RunCommandError:
+        SafeUnlink(write_path)
+        SafeUnlink(mv_target)
+        raise
 
   else:
     # We have the right permissions, simply write the file in python.
diff --git a/lib/osutils_unittest.py b/lib/osutils_unittest.py
index 43f976a..619372d 100644
--- a/lib/osutils_unittest.py
+++ b/lib/osutils_unittest.py
@@ -90,10 +90,14 @@
         self.assertEqual('test', osutils.ReadFile(filename))
         self.assertEqual(0, os.stat(filename).st_uid)
 
-      # Appending to a file is not supported with sudo.
-      self.assertRaises(ValueError, osutils.WriteFile,
-                        os.path.join(root_owned_dir, 'nope'), 'data',
-                        sudo=True, mode='a')
+  def testSudoWriteAppend(self):
+    """Verify that we can write a file as sudo when appending."""
+    with osutils.TempDir(sudo_rm=True) as tempdir:
+      path = os.path.join(tempdir, 'foo')
+      osutils.WriteFile(path, 'one', sudo=True)
+      self.assertRaises(IOError, osutils.WriteFile, path, 'data')
+      osutils.WriteFile(path, 'two', mode='a', sudo=True)
+      self.assertEqual('onetwo', osutils.ReadFile(path))
 
   def testReadFileNonExistent(self):
     """Verify what happens if you ReadFile a file that isn't there."""