From 872fe17d7e39d27e7ae74aa47c4793eecfc90101 Mon Sep 17 00:00:00 2001
From: truffle <truffle074@gmail.com>
Date: Wed, 11 Dec 2024 16:42:59 +1100
Subject: [PATCH] more safety checks and features currently it's more of a
 rearrangment of the original code, with some added safety/integrity checks
 `resume` is now a separate function, the plan is to have the code for
 actually resuming within this function and code for locating the device
 outside. in the more general case, if `late_resume` is enabled, the initram
 will check at both the `init_early` and again at the `init_premount` stage
 for the presence of the resume device. Once host-only is implemented then
 it'll be able to determine at build time which state it'll need to check at.
 For safety checks, currently it checks whether the device is a valid block
 device, whether resume has already been attempted previous (based on the
 content of `/sys/power/resume`) & whether there are are any other mounted
 block devices.

---
 src/ugrd/fs/resume.py   | 67 ++++++++++++++++++++++++++---------------
 src/ugrd/fs/resume.toml |  7 ++++-
 2 files changed, 48 insertions(+), 26 deletions(-)

diff --git a/src/ugrd/fs/resume.py b/src/ugrd/fs/resume.py
index fa9ba08..05e4d28 100644
--- a/src/ugrd/fs/resume.py
+++ b/src/ugrd/fs/resume.py
@@ -2,7 +2,7 @@
 
 from zenlib.util import contains, unset
 
-def _resume(self) -> None:
+def resume(self) -> None:
     """Returns a bash script handling resume from hibernation.
     Checks that /sys/power/resume is writable, resume= is set, and noresume is not set, if so,
     checks if PARTUUID= is in the resume var, and tries to use blkid to find the resume device.
@@ -16,32 +16,49 @@ def _resume(self) -> None:
     Distinguising between a fresh boot and missing/borked hibernation image is not possible at run time.
     """
     return [
-        "resumeval=$(readvar resume)",  # read the cmdline resume var
-        'if ! check_var noresume && [ -n "$resumeval" ] && [ -w /sys/power/resume ]; then',
-        '    if echo "$resumeval" | grep -q -E "^PARTUUID=|^UUID="; then',  # resolve partuuid to device
-        '        resume=$(blkid -t "$resumeval" -o device)',
-        "    else",
-        "        resume=$resumeval",
-        "    fi",
-        '    if [ -e "$resume" ]; then',  # Check if the resume device exists
-        '        einfo "Resuming from: $resume"',
-        '        echo -n "$resume" > /sys/power/resume',  # Attempt to resume
-        '        ewarn "Failed to resume from: $resume"',
-        "    else",
-        '        ewarn "Resume device not found: $resume)"',  # Warn if the resume device does not exist
-        r'        eerror "Block devices:\n$(blkid)"',
-        '        eerror "If you wish to continue booting, remove the resume= kernel parameter."',
-        '''        eerror " or run 'setvar noresume 1' from the recovery shell to skip resuming."''',
-        '        rd_fail "Failed to resume from $(readvar resume)."',
-        "    fi",
-        "fi",
+        # Check resume support
+        '[ -n "$1" ] || (ewarn "No device?" ; return 1)',
+        '[ -w /sys/power/resume ] || (ewarn "Kernel does not support resume!" ; return 1)',
+        '[[ ! "$(cat /sys/power/resume)" == "0:0" ]] || ewarn "/sys/power/resume not empty, resume has already been attempted!"',
+        
+        # Safety checks
+        'if [ ! -z $(lsblk -Q MOUNTPOINT)] ; then',
+        r'    eerror "Cannot safely resume with mounted block devices:\n$(lsblk -Q MOUNTPOINT -no PATH)"',
+        '    return 1',
+        'fi',
+
+        '[ -b "$1" ] || (ewarn "\'$1\' is not a valid block device!" ; return 1)',
+
+        'einfo "Resuming from: $1"',
+        'echo -n "$1" > /sys/power/resume',
+        'einfo "No image on $resume"',
+        'return 0',
     ]
 
-@unset('late_resume')
-def handle_resume(self) -> None:
-    return _resume(self)
+def handle_early_resume(self) -> None:
+    return [
+        'resumeval=$(readvar resume)',
+        'if ! check_var noresume && [ -n "$resumeval" ] ; then',
+        '    local resume',
+        '    if echo "$1" | grep -qE "^PARTUUID=|^UUID="; then',
+        '        resume=$(blkid -t "$resumeval" -o device)',
+        '    else',
+        '        resume="$1"',
+        '    fi',
+
+        '    if [ -z $resume ] || ! [ -b $resume ] ; then',
+        '        einfo "Resume device \'$1\' not found!"',
+        '    else',
+        '        if ! resume "$resume" ; then',
+        '            eerror "If you wish to continue booting, remove the resume= kernel parameter."',
+        '''             eerror " or run 'setvar noresume 1' from the recovery shell to skip resuming."''',
+        '            rd_fail "Failed to resume from $(readvar resume)."',
+        '        fi',
+        '    fi',
+        'fi',
+    ]
 
 @contains('late_resume')
 def handle_late_resume(self) -> None:
-    self.logger.warning("WARNING: Late resume enabled, can cause instability if not used with caution. Highly recommended to read configuration documentation before use.")
-    return _resume(self)
\ No newline at end of file
+    self.logger.warning("[late_resume] enabled, this can result in data loss if filesystems are modified before resuming. Read the docs for more info.")
+    return handle_early_resume(self)  # At the moment it's the same code but delayed, different features are planned
\ No newline at end of file
diff --git a/src/ugrd/fs/resume.toml b/src/ugrd/fs/resume.toml
index 8c5d266..a3ea830 100644
--- a/src/ugrd/fs/resume.toml
+++ b/src/ugrd/fs/resume.toml
@@ -1,10 +1,15 @@
 cmdline_strings = [ "resume" ]
 
+binaries = [ 'lsblk' ]
+
 [imports.init_early] 
-"ugrd.fs.resume" = [ "handle_resume" ]
+"ugrd.fs.resume" = [ "handle_early_resume" ]
 
 [imports.init_premount]
 "ugrd.fs.resume" = [ "handle_late_resume"]
 
+[imports.functions]
+"ugrd.fs.resume" = ["resume"]
+
 [custom_parameters]
 late_resume = "bool"
\ No newline at end of file