From 6672046b6f86cb4ed9fb1862632976dbf135d1f8 Mon Sep 17 00:00:00 2001
From: Stefan Agner <stefan@agner.ch>
Date: Fri, 4 Dec 2020 20:55:35 +0100
Subject: [PATCH] Make the datactl command more robust (#1059)

* Make the datactl command more robust

Validate target disk (partition) size to avoid a copy attempt which will
fail. If e2image operation fails, make sure the leftover copy is not
regonized as data partition.

* Fix hassos-data service device unit dependencies
---
 .../rootfs-overlay/usr/bin/datactl              | 17 +++++++++++++++++
 .../usr/lib/systemd/system/hassos-data.service  |  4 ++--
 .../rootfs-overlay/usr/libexec/hassos-data      |  9 +++++++++
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/buildroot-external/rootfs-overlay/usr/bin/datactl b/buildroot-external/rootfs-overlay/usr/bin/datactl
index ba4e0b375..4a31fb9f5 100755
--- a/buildroot-external/rootfs-overlay/usr/bin/datactl
+++ b/buildroot-external/rootfs-overlay/usr/bin/datactl
@@ -9,6 +9,11 @@ set -e
 DATA_DEVICE_CHILD="$(findmnt --noheadings --output=source /mnt/data)"
 DATA_DEVICE_ROOT="/dev/$(lsblk -no pkname "${DATA_DEVICE_CHILD}")"
 
+if [ "${DATA_DEVICE_ROOT}" = "" ]; then
+    echo "[ERROR] Data disk device not found!"
+    exit 1
+fi
+
 # Move command
 if [ "${1}" = "move" ] && [ -e "${2}" ]; then
     NEW_DEVICE_ROOT="${2}"
@@ -39,8 +44,20 @@ if [ "${1}" = "move" ] && [ -e "${2}" ]; then
         -u "0:a52a4597-fa3a-4851-aefd-2fbe9f849079" \
         "${NEW_DEVICE_ROOT}"
 
+    # Since we create a new partition table etc. we are guaranteed the target
+    # partition is partition 1
+    NEW_DEVICE_PART_SIZE=$(cat "/sys/class/block/$(basename "${NEW_DEVICE_ROOT}")1/size")
+    OLD_DEVICE_PART_SIZE=$(cat "/sys/class/block/$(basename "${DATA_DEVICE_CHILD}")/size")
+
+    if [ "${NEW_DEVICE_PART_SIZE}" -lt "${OLD_DEVICE_PART_SIZE}" ]; then
+        echo "[INFO] Target device too small!"
+        sgdisk --zap-all "${NEW_DEVICE_ROOT}"
+        exit 1
+    fi
+
     touch "/mnt/overlay/move-data"
     cat << EOF
+
 Disk ${NEW_DEVICE_ROOT} has been prepared to be used as data drive and the data
 move has been scheduled for the next reboot. Please reboot the device now and
 make sure to leave the disk connected to the system from now on.
diff --git a/buildroot-external/rootfs-overlay/usr/lib/systemd/system/hassos-data.service b/buildroot-external/rootfs-overlay/usr/lib/systemd/system/hassos-data.service
index 5d7598bbf..f2b4660c0 100644
--- a/buildroot-external/rootfs-overlay/usr/lib/systemd/system/hassos-data.service
+++ b/buildroot-external/rootfs-overlay/usr/lib/systemd/system/hassos-data.service
@@ -3,9 +3,9 @@ Description=Home Assistant OS data partition migration
 DefaultDependencies=no
 RefuseManualStart=true
 RefuseManualStop=true
-Requires=mnt-overlay.mount dev-disk-by\x2dlabel-hassos\x2ddata.device dev-disk-by\x2dpartlabel-hassos\x2ddata\x2dexternal.service
+Requires=mnt-overlay.mount dev-disk-by\x2dlabel-hassos\x2ddata.device dev-disk-by\x2dpartlabel-hassos\x2ddata\x2dexternal.device
 Wants=hassos-expand.service
-After=mnt-overlay.mount dev-disk-by\x2dlabel-hassos\x2ddata.device systemd-fsck@dev-disk-by\x2dlabel-hassos\x2ddata.service dev-disk-by\x2dpartlabel-hassos\x2ddata\x2dexternal.service
+After=mnt-overlay.mount dev-disk-by\x2dlabel-hassos\x2ddata.device systemd-fsck@dev-disk-by\x2dlabel-hassos\x2ddata.service dev-disk-by\x2dpartlabel-hassos\x2ddata\x2dexternal.device
 Before=hassos-expand.service
 ConditionPathExists=/mnt/overlay/move-data
 
diff --git a/buildroot-external/rootfs-overlay/usr/libexec/hassos-data b/buildroot-external/rootfs-overlay/usr/libexec/hassos-data
index 0e1b74023..83a38c556 100755
--- a/buildroot-external/rootfs-overlay/usr/libexec/hassos-data
+++ b/buildroot-external/rootfs-overlay/usr/libexec/hassos-data
@@ -10,9 +10,18 @@ OLD_DEVICE_CHILD="$(readlink -f "/dev/disk/by-label/hassos-data")"
 # Rely on systemd-udev symlinks to find external data partition by partlabel
 NEW_DEVICE_CHILD="$(readlink -f "/dev/disk/by-partlabel/hassos-data-external")"
 
+NEW_DEVICE_PART_SIZE=$(cat "/sys/class/block/$(basename "${NEW_DEVICE_CHILD}")/size")
+OLD_DEVICE_PART_SIZE=$(cat "/sys/class/block/$(basename "${OLD_DEVICE_CHILD}")/size")
+if [ "${NEW_DEVICE_PART_SIZE}" -lt "${OLD_DEVICE_PART_SIZE}" ]; then
+    echo "[INFO] Target device too small!"
+    exit 1
+fi
+
 echo "[INFO] Moving data from ${OLD_DEVICE_CHILD} to ${NEW_DEVICE_CHILD}"
 if ! e2image -ra -p "${OLD_DEVICE_CHILD}" "${NEW_DEVICE_CHILD}"; then
     echo "[ERROR] Copying data partition to external device failed!"
+    # Rename partition to make sure HAOS does not try to mount a failed copy attempt.
+    e2label "${NEW_DEVICE_CHILD}" hassos-data-failed || true
     exit 1
 fi