mkosi-initrd: protect old initrd image against errors

When `copy_tree()` ends up calling `cp`, if it fails because there is not enough
space, it leaves and 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, backup the old initrd image before
calling `copy_tree()`, and restore it if that function throws an exception.

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, rmtree
 from mkosi.util import PathString, mandatory_variable, resource_path
 
 
@@ -141,12 +142,32 @@ def initrd_finalize(staging_dir: str, output: str, output_dir: str) -> None:
     else:
         output_dir = os.fspath(Path.cwd())
 
+    # backup current image
+    if Path(f"{output_dir}/{output}").exists():
+        log_notice(f"Backup {output_dir}/{output} to {staging_dir}/{output}.bak")
+        copy_tree(
+            Path(f"{output_dir}/{output}"),
+            Path(f"{staging_dir}/{output}.bak"),
+        )
+
     log_notice(f"Copying {staging_dir}/{output} to {output_dir}/{output}")
-    # mkosi symlinks the expected output image, so dereference it
-    copy_tree(
-        Path(f"{staging_dir}/{output}").resolve(),
-        Path(f"{output_dir}/{output}"),
-    )
+    try:
+        # mkosi symlinks the expected output image, so dereference it
+        copy_tree(
+            Path(f"{staging_dir}/{output}").resolve(),
+            Path(f"{output_dir}/{output}"),
+        )
+    except subprocess.CalledProcessError:
+        # restore image backup on error
+        if Path(f"{staging_dir}/{output}.bak").exists():
+            log_notice(f"Removing incomplete {output_dir}/{output}")
+            rmtree(Path(f"{output_dir}/{output}"))
+            log_notice(f"Restoring backup {staging_dir}/{output}.bak to {output_dir}/{output}")
+            copy_tree(
+                Path(f"{staging_dir}/{output}.bak"),
+                Path(f"{output_dir}/{output}"),
+            )
+        raise
 
 
 def initrd_common_args(parser: argparse.ArgumentParser) -> None: