From 6a2834b3e64116d90de3d3b5a8be7c734d4122a5 Mon Sep 17 00:00:00 2001
From: truffle <truffle074@gmail.com>
Date: Fri, 15 Nov 2024 11:24:39 +1100
Subject: [PATCH] adds support for late resume controlled by the 'late_resume'
 parameter, will move the trigger point for resume to after devices have been
 decrypted or discovered allows resume from encrypted swap or devices
 requiring extra handling be visible to the kernel

---
 docs/configuration.md   |  4 ++++
 src/ugrd/fs/resume.py   | 14 ++++++++++++--
 src/ugrd/fs/resume.toml |  6 ++++++
 3 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/docs/configuration.md b/docs/configuration.md
index 2bb7eab..c0fae93 100644
--- a/docs/configuration.md
+++ b/docs/configuration.md
@@ -276,6 +276,10 @@ Importing this module will run `btrfs device scan` and pull btrfs modules.
 
 When enabled, attempts to resume after hibernation if resume= is passed on the kernel command line.
 
+> Please use the following option with **CAUTION** as it can be unstable in certain configurations! Any writes to disks that occur pre-resume run the risk of causing system instability! For more information have a read of the warnings in the [kernel docs](https://www.kernel.org/doc/html/latest/power/swsusp.html).
+
+* `late_resume` (true) When enabled will attempt to resume from hibernation after decryption and device mapping, allowing resume from encrypted or otherwise hidden swap devices.
+
 ### Cryptographic modules
 
 Several cryptographic modules are provided, mostly to assist in mounting encrypted volumes and handling keyfiles.
diff --git a/src/ugrd/fs/resume.py b/src/ugrd/fs/resume.py
index 38f5ddd..fa9ba08 100644
--- a/src/ugrd/fs/resume.py
+++ b/src/ugrd/fs/resume.py
@@ -1,7 +1,8 @@
 __version__ = "0.4.0"
 
+from zenlib.util import contains, unset
 
-def handle_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.
@@ -17,7 +18,7 @@ def handle_resume(self) -> None:
     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 "PARTUUID="; then',  # resolve partuuid to device
+        '    if echo "$resumeval" | grep -q -E "^PARTUUID=|^UUID="; then',  # resolve partuuid to device
         '        resume=$(blkid -t "$resumeval" -o device)',
         "    else",
         "        resume=$resumeval",
@@ -35,3 +36,12 @@ def handle_resume(self) -> None:
         "    fi",
         "fi",
     ]
+
+@unset('late_resume')
+def handle_resume(self) -> None:
+    return _resume(self)
+
+@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
diff --git a/src/ugrd/fs/resume.toml b/src/ugrd/fs/resume.toml
index c4732dd..8c5d266 100644
--- a/src/ugrd/fs/resume.toml
+++ b/src/ugrd/fs/resume.toml
@@ -2,3 +2,9 @@ cmdline_strings = [ "resume" ]
 
 [imports.init_early] 
 "ugrd.fs.resume" = [ "handle_resume" ]
+
+[imports.init_premount]
+"ugrd.fs.resume" = [ "handle_late_resume"]
+
+[custom_parameters]
+late_resume = "bool"
\ No newline at end of file