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.

diff --git a/mkosi/initrd.py b/mkosi/initrd.py
index 782ffeda0ef0f3933ef7802174acb4357da8b88f..7b81c8243f5b0235a72a6262d2102121c7752141 100644 (file)
--- a/mkosi/initrd.py
+++ b/mkosi/initrd.py
@@ -7,6 +7,7 @@ import logging
 import os
 import platform
 import shutil
+import subprocess
 import sys
 import tempfile
 from pathlib import Path
@@ -18,7 +19,7 @@ from mkosi.documentation import show_docs
 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
 
 
@@ -187,12 +188,21 @@ def initrd_finalize(staging_dir: Path, output: str, output_dir: Optional[Path])
     else:
         output_dir = Path.cwd()
 
-    log_notice(f"Copying {staging_dir / output} to {output_dir / output}")
+    staging = staging_dir / output
+    tmp = output_dir / f"{output}.new"
+    final = output_dir / output
+
+    log_notice(f"Copying {staging} to {tmp}")
     # mkosi symlinks the expected output image, so dereference it
-    copy_tree(
-        (staging_dir / output).resolve(),
-        output_dir / output,
-    )
+    try:
+        copy_tree(staging.resolve(), tmp)
+    except subprocess.CalledProcessError:
+        rmtree(tmp)
+        raise
+
+    log_notice(f"Moving {tmp} to {final}")
+    rmtree(final)
+    move_tree(tmp, final)
 
 
 def initrd_common_args(parser: argparse.ArgumentParser) -> None: