mkosi-initrd: protect existing initrd image against errors

When `copy_tree()` ends up calling `cp`, if it fails because there is not enough
space, it leaves an incomplete initrd image in the output directory. If we are
writing the initrd image directly where a bootloader entry has an initrd
configured, the system will fail to boot. So, instead of copying the initrd
image directly to the output, copy it next to it in the same output directory,
and if the copy is successful, replace it.

Author: aafeijoo-suse

mkosi/initrd.py

@@ -7,6 +7,7 @@
 import os
 import platform
 import shutil
+import subprocess
 import sys
 import tempfile
 from pathlib import Path
@@ -18,7 +19,7 @@
 from mkosi.log import log_notice, log_setup
 from mkosi.run import find_binary, run, uncaught_exception_handler
 from mkosi.sandbox import __version__, umask
-from mkosi.tree import copy_tree
+from mkosi.tree import copy_tree, move_tree, rmtree
 from mkosi.util import PathString, mandatory_variable, resource_path
 
 
@@ -141,12 +142,22 @@ def initrd_finalize(staging_dir: str, output: str, output_dir: str) -> None:
     else:
         output_dir = os.fspath(Path.cwd())
 
-    log_notice(f"Copying {staging_dir}/{output} to {output_dir}/{output}")
+    staging = Path(f"{staging_dir}/{output}")
+    tmp = Path(f"{output_dir}/{output}.new")
+    final = Path(f"{output_dir}/{output}")
+
+    log_notice(f"Copying {staging} to {tmp}")
     # mkosi symlinks the expected output image, so dereference it
-    copy_tree(
-        Path(f"{staging_dir}/{output}").resolve(),
-        Path(f"{output_dir}/{output}"),
-    )
+    try:
+        copy_tree(staging.resolve(), tmp)
+    except subprocess.CalledProcessError:
+        rmtree(tmp)
+        raise
+
+    log_notice(f"Moving {tmp} to {final}")
+    if final.is_dir():
+        rmtree(final)
+    move_tree(tmp, final)
 
 
 def initrd_common_args(parser: argparse.ArgumentParser) -> None: