From b6018cfb40c1dca4ca7b9c0e5700498e2d24946c Mon Sep 17 00:00:00 2001
From: Mykyta Poturai <mykyta_poturai@epam.com>
Date: Sat, 26 Oct 2024 15:49:10 +0300
Subject: [PATCH] rouge: Add support for resizing raw images

Add new "resize" option to the raw_image block, which allows rouge
to resize the raw_image to the size set in the block. It is enabled
by default and can be disabled manually.

This is useful when the image size needs to be changed frequently, or
multiple images need to be created from the same source files, as it
allows to skip the rootfs rebuild step.

Images are copied to a temporary directory before resizing to preserve
the original build artifacts. The temporary directory is created in the
current working directory to prevent filling the system's /tmp with
large images.

Signed-off-by: Mykyta Poturai <mykyta_poturai@epam.com>
---
 docs/rouge.rst              | 10 ++++++++--
 moulin/rouge/block_entry.py | 25 ++++++++++++++++++++++++-
 moulin/rouge/ext_utils.py   | 10 ++++++++++
 3 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/docs/rouge.rst b/docs/rouge.rst
index c634d04..5f248e1 100644
--- a/docs/rouge.rst
+++ b/docs/rouge.rst
@@ -218,6 +218,7 @@ partition (which is described below).
 
    type: raw_image # defines raw image block
    size: 400 MiB
+   resize: false
    image_path: "some/path/rootfs.ext4"
 
 :code:`image_path` is mandatory. This is a file to be included into
@@ -225,8 +226,13 @@ resulting image.
 
 :code:`size` is optional. If it is omitted, `rouge` will use size of
 file. If provided :code:`size` is smaller than file size, `rouge` will
-stop with an error. Thus, you can create block that is bigger than
-file, but not smaller.
+stop with an error. If provided :code:`size` is bigger than file size,
+`rouge` will try to resize the file to match :code:`size`.
+
+:code:`resize` is optional. If set to :code:`false`, it will prevent
+`rouge` from resizing the image to the size of the block. This is
+useful when you want to include a file that is smaller than the block
+and leave the rest of the block empty.
 
 Android Sparse Image Block
 ^^^^^^^^^^^^^^^^^^^^^^^^^^
diff --git a/moulin/rouge/block_entry.py b/moulin/rouge/block_entry.py
index 61e8135..e8dd612 100644
--- a/moulin/rouge/block_entry.py
+++ b/moulin/rouge/block_entry.py
@@ -9,6 +9,7 @@
 import shutil
 import logging
 import itertools
+import subprocess
 from typing import List, Tuple, NamedTuple, cast
 from tempfile import NamedTemporaryFile, TemporaryDirectory
 
@@ -129,12 +130,14 @@ def __init__(self, node: YamlValue, **kwargs):
         self._node = node
         self._fname = self._node["image_path"].as_str
         self._size = 0
+        self._resize = True
 
     def _complete_init(self):
         mark = self._node["image_path"].mark
         if not os.path.exists(self._fname):
             raise YAMLProcessingError(f"Can't find file '{self._fname}'", mark)
         fsize = os.path.getsize(self._fname)
+        self._resize = self._node.get("resize", True).as_bool
         size_node = self._node.get("size", None)
         if size_node:
             self._size = _parse_size(size_node)
@@ -154,7 +157,27 @@ def size(self) -> int:
     def write(self, fp, offset):
         if not self._size:
             self._complete_init()
-        ext_utils.dd(self._fname, fp, offset)
+
+        fsize = os.path.getsize(self._fname)
+
+        if (self._resize and fsize < self._size):
+            # Not using default /tmp to prevent filling ram with huge images
+            with TemporaryDirectory(dir=".") as tmpd:
+                shutil.copy(self._fname, tmpd)
+                with open(os.path.join(tmpd, os.path.basename(self._fname)), "rb+") as data:
+                    data.truncate(self._size)
+                try:
+                    ext_utils.resize2fs(os.path.join(tmpd, os.path.basename(self._fname)))
+                except subprocess.CalledProcessError as e:
+                    log.error(
+                        "%s is not resizable. If you don't need to resize it, set 'resize: false' in the yaml.",
+                        self._fname
+                    )
+                    raise e
+
+                ext_utils.dd(os.path.join(tmpd, os.path.basename(self._fname)), fp, offset)
+        else:
+            ext_utils.dd(self._fname, fp, offset)
 
     def get_deps(self) -> List[str]:
         "Return list of dependencies needed to build this block"
diff --git a/moulin/rouge/ext_utils.py b/moulin/rouge/ext_utils.py
index 1ac7ab8..ef5a568 100644
--- a/moulin/rouge/ext_utils.py
+++ b/moulin/rouge/ext_utils.py
@@ -93,3 +93,13 @@ def mmd(img: BinaryIO, folders: list):
     args.extend(folders)
 
     _run_cmd(args)
+
+
+def resize2fs(img: str, size: Optional[int] = None):
+    "Resize fs image to the given size"
+    args = ["resize2fs", img]
+
+    if size:
+        args.append(str(size))
+
+    _run_cmd(args)