From f998687833b4ca63bd63d11d451aea311f23fa5f Mon Sep 17 00:00:00 2001
From: xuliping <xuliping@microsoft.com>
Date: Fri, 21 Feb 2025 06:07:26 +0000
Subject: [PATCH 1/2] use pzstd for FILESYSTEM_DOCKERFS when defined
 BUILD_REDUCE_IMAGE_SIZE

---
 build_debian.sh                      | 12 +++++++++++-
 files/dsc/install_debian.j2          |  4 ++++
 files/initramfs-tools/file           | 18 ++++++++++++++++++
 files/initramfs-tools/pzstd          | 16 ++++++++++++++++
 files/initramfs-tools/union-mount.j2 | 14 ++++++++++----
 5 files changed, 59 insertions(+), 5 deletions(-)
 create mode 100644 files/initramfs-tools/file
 create mode 100644 files/initramfs-tools/pzstd

diff --git a/build_debian.sh b/build_debian.sh
index 4b915d2d7e35..3102c45edcbd 100755
--- a/build_debian.sh
+++ b/build_debian.sh
@@ -392,6 +392,12 @@ sudo LANG=C DEBIAN_FRONTEND=noninteractive chroot $FILESYSTEM_ROOT apt-get -y in
     zstd                    \
     nvme-cli
 
+sudo cp files/initramfs-tools/pzstd $FILESYSTEM_ROOT/etc/initramfs-tools/hooks/pzstd
+sudo chmod +x $FILESYSTEM_ROOT/etc/initramfs-tools/hooks/pzstd
+
+sudo cp files/initramfs-tools/file $FILESYSTEM_ROOT/etc/initramfs-tools/hooks/file
+sudo chmod +x $FILESYSTEM_ROOT/etc/initramfs-tools/hooks/file
+
 # Have systemd create the auditd log directory
 sudo mkdir -p ${FILESYSTEM_ROOT}/etc/systemd/system/auditd.service.d
 sudo tee ${FILESYSTEM_ROOT}/etc/systemd/system/auditd.service.d/log-directory.conf >/dev/null <<EOF
@@ -855,7 +861,11 @@ if [[ $MULTIARCH_QEMU_ENVIRON == y || $CROSS_BUILD_ENVIRON == y ]]; then
 fi
 
 ## Compress docker files
-pushd $FILESYSTEM_ROOT && sudo tar -I pigz -cf $OLDPWD/$FILESYSTEM_DOCKERFS -C ${DOCKERFS_PATH}var/lib/docker .; popd
+if [ "$BUILD_REDUCE_IMAGE_SIZE" = "y" ]; then
+    pushd $FILESYSTEM_ROOT && sudo tar -I pzstd -cf $OLDPWD/$FILESYSTEM_DOCKERFS -C ${DOCKERFS_PATH}var/lib/docker .; popd
+else
+    pushd $FILESYSTEM_ROOT && sudo tar -I pigz -cf $OLDPWD/$FILESYSTEM_DOCKERFS -C ${DOCKERFS_PATH}var/lib/docker .; popd
+fi
 
 ## Compress together with /boot, /var/lib/docker and $PLATFORM_DIR as an installer payload zip file
 pushd $FILESYSTEM_ROOT && sudo tar -I pigz -cf platform.tar.gz -C $PLATFORM_DIR . && sudo zip -n .gz $OLDPWD/$INSTALLER_PAYLOAD -r boot/ platform.tar.gz; popd
diff --git a/files/dsc/install_debian.j2 b/files/dsc/install_debian.j2
index def175249764..a8397c974c98 100755
--- a/files/dsc/install_debian.j2
+++ b/files/dsc/install_debian.j2
@@ -15,7 +15,11 @@ image_dir=image-$image_version
 
 INSTALLER_PAYLOAD=fs.zip
 DOCKERFS_DIR=docker
+{% if BUILD_REDUCE_IMAGE_SIZE == "y" -%}
+FILESYSTEM_DOCKERFS=dockerfs.tar.zstd
+{%- else -%}
 FILESYSTEM_DOCKERFS=dockerfs.tar.gz
+{%- endif %}
 BL_CONF=boot.conf
 
 DATA_PARTUUID=6ED62003-DD8D-44B8-9538-0A2B7C7E628F
diff --git a/files/initramfs-tools/file b/files/initramfs-tools/file
new file mode 100644
index 000000000000..cff6acd18fd7
--- /dev/null
+++ b/files/initramfs-tools/file
@@ -0,0 +1,18 @@
+#!/bin/sh
+set -e
+PREREQ=""
+prereqs() {
+    echo "$PREREQ"
+}
+case "$1" in
+    prereqs)
+        prereqs
+        exit 0
+        ;;
+esac
+. /usr/share/initramfs-tools/hook-functions
+# Include file binary
+copy_exec /usr/bin/file /usr/bin
+# Include magic database
+copy_exec /usr/lib/file/magic.mgc /etc
+exit 0
diff --git a/files/initramfs-tools/pzstd b/files/initramfs-tools/pzstd
new file mode 100644
index 000000000000..561319dd941b
--- /dev/null
+++ b/files/initramfs-tools/pzstd
@@ -0,0 +1,16 @@
+#!/bin/sh
+set -e
+PREREQ=""
+prereqs() {
+    echo "$PREREQ"
+}
+case "$1" in
+    prereqs)
+        prereqs
+        exit 0
+        ;;
+esac
+. /usr/share/initramfs-tools/hook-functions
+# Include pzstd binary
+copy_exec /usr/bin/pzstd /usr/bin
+exit 0
\ No newline at end of file
diff --git a/files/initramfs-tools/union-mount.j2 b/files/initramfs-tools/union-mount.j2
index 8f8abb8f6af5..cd27a9728c5e 100644
--- a/files/initramfs-tools/union-mount.j2
+++ b/files/initramfs-tools/union-mount.j2
@@ -1,4 +1,4 @@
-#!/bin/sh -e
+#!/bin/sh
 
 PREREQS="varlog"
 
@@ -130,9 +130,15 @@ extract_dockerfs()
 {
     echo "Extracting {{ FILESYSTEM_DOCKERFS }}"
     if [ -f "${rootmnt}/host/$image_dir/{{ FILESYSTEM_DOCKERFS }}" ] && [ "$secureboot" = false ]; then
-        # Extract dockerfs.tar.gz into /var/lib/docker unless the system booted with secureboot
-        # In secureboot dockerfs.tar.gz cannot be trusted as it does not have a signature
-        tar xz --numeric-owner -f ${rootmnt}/host/$image_dir/{{ FILESYSTEM_DOCKERFS }} -C ${rootmnt}/var/lib/docker
+        # Check if the file is zstd compressed
+        file_type=$(file -b --mime-type "${rootmnt}/host/$image_dir/{{ FILESYSTEM_DOCKERFS }}")
+        if [ "$file_type" = "application/zstd" ]; then
+            echo "Detected zstd compression, extracting with pzstd..."
+            pzstd -d -q ${rootmnt}/host/$image_dir/{{ FILESYSTEM_DOCKERFS }} -c | tar x --numeric-owner -C ${rootmnt}/var/lib/docker
+        else
+            echo "Using default extraction method (gzip assumed)..."
+            tar xz --numeric-owner -f "${rootmnt}/host/$image_dir/{{ FILESYSTEM_DOCKERFS }}" -C "${rootmnt}/var/lib/docker"
+        fi
     elif [ "$bootloader" = "aboot" ] && unzip -l "$swi_path" | grep -q {{ FILESYSTEM_DOCKERFS }}; then
         # Aboot swi images also support extracting dockerfs.tar.gz directly from them
         unzip -qp "$swi_path" {{ FILESYSTEM_DOCKERFS }} | tar xz --numeric-owner -C ${rootmnt}/var/lib/docker

From ed35d0f80e4583dd0cf4b9d562413b2f5f06bfc0 Mon Sep 17 00:00:00 2001
From: xuliping <xuliping@microsoft.com>
Date: Wed, 26 Feb 2025 02:06:55 +0000
Subject: [PATCH 2/2] update, add back -e for union-mount.j2

---
 files/initramfs-tools/pzstd          | 2 +-
 files/initramfs-tools/union-mount.j2 | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/files/initramfs-tools/pzstd b/files/initramfs-tools/pzstd
index 561319dd941b..8775d480b62c 100644
--- a/files/initramfs-tools/pzstd
+++ b/files/initramfs-tools/pzstd
@@ -13,4 +13,4 @@ esac
 . /usr/share/initramfs-tools/hook-functions
 # Include pzstd binary
 copy_exec /usr/bin/pzstd /usr/bin
-exit 0
\ No newline at end of file
+exit 0
diff --git a/files/initramfs-tools/union-mount.j2 b/files/initramfs-tools/union-mount.j2
index cd27a9728c5e..de2cb5c75961 100644
--- a/files/initramfs-tools/union-mount.j2
+++ b/files/initramfs-tools/union-mount.j2
@@ -1,4 +1,4 @@
-#!/bin/sh
+#!/bin/sh -e
 
 PREREQS="varlog"