From 571de4064b73008b47cbf7d14b805c9a4fce1484 Mon Sep 17 00:00:00 2001
From: Gabriel Adrian Samfira <gsamfira@cloudbasesolutions.com>
Date: Tue, 22 Feb 2022 15:31:26 +0000
Subject: [PATCH] Make image compression format configurable

This change adds a new flag called --image_compression_format which
allows us to output the final VM image as one of the supported formats:

bz2 (default), gz, zip or none

if the compression format is "none" or "", the image will not be compressed.

Signed-off-by: Gabriel Adrian Samfira <gsamfira@cloudbasesolutions.com>
---
 build_library/release_util.sh  | 40 +++++++++++++++++++++++++++++-----
 build_library/vm_image_util.sh |  4 +++-
 2 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/build_library/release_util.sh b/build_library/release_util.sh
index f1ecc8acc0..9905645c33 100644
--- a/build_library/release_util.sh
+++ b/build_library/release_util.sh
@@ -7,6 +7,7 @@ UPLOAD_ROOT=
 UPLOAD_PATH=
 TORCX_UPLOAD_ROOT=
 UPLOAD_DEFAULT=${FLAGS_FALSE}
+DEFAULT_IMAGE_COMPRESSION_FORMAT="bz2"
 
 # Default upload root can be overridden from the environment.
 _user="${USER}"
@@ -15,9 +16,6 @@ _user="${USER}"
 : ${FLATCAR_TORCX_UPLOAD_ROOT:=${FLATCAR_UPLOAD_ROOT}/torcx}
 unset _user
 
-IMAGE_ZIPPER="lbzip2 --compress --keep"
-IMAGE_ZIPEXT=".bz2"
-
 DEFINE_boolean parallel ${FLAGS_TRUE} \
   "Enable parallelism in gsutil."
 DEFINE_boolean upload ${UPLOAD_DEFAULT} \
@@ -42,6 +40,38 @@ DEFINE_string sign "" \
   "Sign all files to be uploaded with the given GPG key."
 DEFINE_string sign_digests "" \
   "Sign image DIGESTS files with the given GPG key."
+DEFINE_string image_compression_format "${DEFAULT_IMAGE_COMPRESSION_FORMAT}" \
+  "Compress the resulting images using this format. Options are: none, bz2, gz, zip"
+
+
+compress_file() {
+    local filepath="$1"
+
+    [ ! -f "${filepath}" ] && die "Image file ${filepath} does not exist"
+
+    case "${FLAGS_image_compression_format}" in
+    "none"|"")
+        echo -n "${filepath}"
+        return 0
+        ;;
+    "bz2")
+        IMAGE_ZIPPER="lbzip2 --compress --keep"
+        ;;
+    "gz")
+        IMAGE_ZIPPER="pigz --keep"
+        ;;
+    "zip")
+        IMAGE_ZIPPER="pigz --keep --zip"
+        ;;
+    *)
+        die "Unsupported compression format ${FLAGS_image_compression_format}"
+        ;;
+    esac
+
+    ${IMAGE_ZIPPER} -f "${filepath}" 2>&1 >/dev/null || die "failed to compress ${filepath}"
+
+    echo -n "${filepath}.${FLAGS_image_compression_format}"
+}
 
 check_gsutil_opts() {
     [[ ${FLAGS_upload} -eq ${FLAGS_TRUE} ]] || return 0
@@ -217,8 +247,8 @@ upload_image() {
         # Compress disk images
         if [[ "${filename}" =~ \.(img|bin|vdi|vhd|vmdk)$ ]]; then
             info "Compressing ${filename##*/}"
-            $IMAGE_ZIPPER -f "${filename}"
-            uploads+=( "${filename}${IMAGE_ZIPEXT}" )
+            COMPRESSED_FILENAME=$(compress_file "${filename}")
+            uploads+=( "$COMPRESSED_FILENAME" )
         else
             uploads+=( "${filename}" )
         fi
diff --git a/build_library/vm_image_util.sh b/build_library/vm_image_util.sh
index 6b5ad1744d..a2798c8d76 100644
--- a/build_library/vm_image_util.sh
+++ b/build_library/vm_image_util.sh
@@ -1227,7 +1227,9 @@ vm_upload() {
 
     # If upload_images compressed $legacy_uploaded be sure to add .bz2
     if [[ "${legacy_uploaded}" =~ \.(img|bin|vdi|vhd|vmdk)$ ]]; then
-        legacy_uploaded+="${IMAGE_ZIPEXT}"
+        if ! [[ "${FLAGS_image_compression_format}" =~ ^(""|none)$ ]]; then
+            legacy_uploaded+=".${FLAGS_image_compression_format}"
+        fi
     fi
 
     local legacy_digests="${legacy_uploaded}.DIGESTS"