From 8e3e89e52428da908f6bc8e1bb1a0d04cc125ae6 Mon Sep 17 00:00:00 2001
From: Kroese <kroese@users.noreply.github.com>
Date: Sat, 9 Dec 2023 02:14:28 +0100
Subject: [PATCH] feat: Support for qcow2 disks

- Make getSize use the stat function for raw disks, so that it still works even when qemu-utils is not installed
- Make convertDisk only look for known extensions, so that potential backup files like data.bak are not converted
---
 src/disk.sh | 77 ++++++++++++++++++++++++++++++++++++-----------------
 1 file changed, 52 insertions(+), 25 deletions(-)

diff --git a/src/disk.sh b/src/disk.sh
index 954b90de..1fabe2db 100644
--- a/src/disk.sh
+++ b/src/disk.sh
@@ -53,11 +53,27 @@ ext2fmt() {
 
 getSize() {
   local DISK_FILE=$1
+  local DISK_EXT
+  local DISK_FMT
 
-  qemu-img info "${DISK_FILE}" -f "${DISK_FMT}" | grep '^virtual size: ' | sed 's/.*(\(.*\) bytes)/\1/'
+  DISK_EXT="$(echo "${DISK_FILE//*./}" | sed 's/^.*\.//')"
+  DISK_FMT="$(ext2fmt "${DISK_EXT}")"
+
+  case "${DISK_FMT,,}" in
+    raw)
+      stat -c%s "${DISK_FILE}"
+      ;;
+    qcow2)
+      qemu-img info "${DISK_FILE}" -f "${DISK_FMT}" | grep '^virtual size: ' | sed 's/.*(\(.*\) bytes)/\1/'
+      ;;
+    *)
+      error "Unrecognized disk format: ${DISK_FMT}" && exit 88
+      ;;
+  esac
 }
 
-doResize() {
+resizeDisk() {
+
   local GB
   local REQ
   local SPACE
@@ -112,6 +128,7 @@ doResize() {
 }
 
 convertDisk() {
+
   local CONV_FLAGS=""
   local SOURCE_FILE=$1
   local SOURCE_FMT=$2
@@ -119,8 +136,9 @@ convertDisk() {
   local DST_FMT=$4
 
   case "${DST_FMT}" in
-  qcow2)
-    CONV_FLAGS="${CONV_FLAGS} -c"
+    qcow2)
+      CONV_FLAGS="${CONV_FLAGS} -c"
+      ;;
   esac
 
   # shellcheck disable=SC2086
@@ -128,6 +146,7 @@ convertDisk() {
 }
 
 createDisk() {
+
   local GB
   local SPACE
   local SPACE_GB
@@ -143,7 +162,7 @@ createDisk() {
         # Create an empty file
         if ! truncate -s "${DISK_SPACE}" "${DISK_FILE}"; then
           rm -f "${DISK_FILE}"
-          error "Could not create a ${DISK_SPACE} file for ${DISK_DESC} (${DISK_FILE})" && exit 87
+          error "Could not create a ${DISK_SPACE} ${DISK_FMT} file for ${DISK_DESC} (${DISK_FILE})" && exit 87
         fi
 
       else
@@ -161,7 +180,7 @@ createDisk() {
         if ! fallocate -l "${DISK_SPACE}" "${DISK_FILE}"; then
           if ! truncate -s "${DISK_SPACE}" "${DISK_FILE}"; then
             rm -f "${DISK_FILE}"
-            error "Could not create a ${DISK_SPACE} file for ${DISK_DESC} (${DISK_FILE})" && exit 87
+            error "Could not create a ${DISK_SPACE} ${DISK_FMT} file for ${DISK_DESC} (${DISK_FILE})" && exit 87
           fi
         fi
 
@@ -169,7 +188,7 @@ createDisk() {
       ;;
     qcow2)
       if ! qemu-img create -f "$DISK_FMT" -- "${DISK_FILE}" "${DISK_SPACE}" ; then
-        error "Could not create a ${DISK_SPACE} byte ${DISK_FMT} file for ${DISK_DESC} (${DISK_FILE})" && exit 89
+        error "Could not create a ${DISK_SPACE} ${DISK_FMT} file for ${DISK_DESC} (${DISK_FILE})" && exit 89
       fi
       ;;
   esac
@@ -181,7 +200,6 @@ addDisk () {
   local CUR_SIZE
   local DATA_SIZE
   local DISK_FILE
-  local DISK_ROOT
   local DISK_ID=$1
   local DISK_BASE=$2
   local DISK_EXT=$3
@@ -193,26 +211,31 @@ addDisk () {
 
   DISK_FILE="${DISK_BASE}.${DISK_EXT}"
 
-  DISK_ROOT="$(basename -- "${DISK_BASE}")"
-
   DIR=$(dirname "${DISK_FILE}")
   [ ! -d "${DIR}" ] && return 0
 
   if ! [ -f "${DISK_FILE}" ] ; then
-    local OTHER_FORMS
-    OTHER_FORMS="$(find "${DIR}" -maxdepth 1 | sed -n -- "/\/${DISK_ROOT}\./p" | sed -- "/\.${DISK_EXT}$/d")"
-
-    if [[ -n "${OTHER_FORMS}" ]] ; then
-      local SOURCE_FILE
-      local SOURCE_EXT
-      local SOURCE_FMT
-      SOURCE_FILE="$(echo "${OTHER_FORMS}" | head -n1)"
-      SOURCE_EXT="$(echo "${SOURCE_FILE//*./}" | sed 's/^.*\.//')"
-      SOURCE_FMT="$(ext2fmt "${SOURCE_EXT}")"
-      info "Other disk formats detected for ${DISK_DESC} (${OTHER_FORMS//$'\n'/, }), converting ${SOURCE_FILE}"
-      if ! convertDisk "${SOURCE_FILE}" "${SOURCE_FMT}" "${DISK_FILE}" "${DISK_FMT}" ; then
-        info "Disk conversion failed, creating new disk image as fallback"
-        rm "${DISK_FILE}"
+    local PREV_EXT
+    local PREV_FMT
+    local PREV_FILE
+
+    if [[ "${DISK_FMT,,}" != "raw" ]]; then
+      PREV_FMT="raw"
+    else
+      PREV_FMT="qcow2"
+    fi
+    PREV_EXT="$(fmt2ext "${PREV_FMT}")"
+    PREV_FILE="${DISK_BASE}.${PREV_EXT}"
+
+    if [ -f "${PREV_FILE}" ] ; then
+      info "Disk format change detected for ${DISK_DESC} (${PREV_FMT} to ${DISK_FMT}), converting ${PREV_FILE} ..."
+      
+      if ! convertDisk "${PREV_FILE}" "${PREV_FMT}" "${DISK_FILE}" "${DISK_FMT}" ; then
+        info "Disk conversion failed, creating new disk image as fallback."
+        rm -f "${DISK_FILE}"
+      else
+        info "Disk conversion completed succesfully, removing ${PREV_FILE} ..."
+        rm -f "${PREV_FILE}"
       fi
     fi
   fi
@@ -222,13 +245,17 @@ addDisk () {
   DATA_SIZE=$(numfmt --from=iec "${DISK_SPACE}")
 
   if [ -f "${DISK_FILE}" ]; then
+
     CUR_SIZE=$(getSize "${DISK_FILE}")
 
     if [ "$DATA_SIZE" -gt "$CUR_SIZE" ]; then
-      doResize "${DISK_FILE}" "${CUR_SIZE}" "${DATA_SIZE}" "${DISK_SPACE}" "${DISK_DESC}" "${DISK_FMT}" || exit $?
+      resizeDisk "${DISK_FILE}" "${CUR_SIZE}" "${DATA_SIZE}" "${DISK_SPACE}" "${DISK_DESC}" "${DISK_FMT}" || exit $?
     fi
+
   else
+
     createDisk "${DISK_FILE}" "${DISK_SPACE}" "${DISK_DESC}" "${DISK_FMT}" || exit $?
+
   fi
 
   DISK_OPTS="${DISK_OPTS} \