Target the actual NVMe controller in Apple suspend fix

install/hardware/apple/fix-suspend-nvme.sh

@@ -2,29 +2,67 @@
 # This prevents NVMe drives from failing to wake from sleep properly
 MACBOOK_MODEL=$(cat /sys/class/dmi/id/product_name 2>/dev/null || true)
 
+find_nvme_pci_device() {
+  local sys_root="${1:-/sys}"
+  local controller=""
+  local device_path=""
+  local pci_device=""
+
+  for controller in "$sys_root"/class/nvme/nvme*; do
+    [[ -e $controller/device ]] || continue
+
+    device_path=$(readlink -f "$controller/device" 2>/dev/null) || continue
+    pci_device=${device_path##*/}
+
+    [[ $pci_device =~ ^[0-9A-Fa-f]{4}:[0-9A-Fa-f]{2}:[0-9A-Fa-f]{2}\.[0-7]$ ]] || continue
+    [[ -f $sys_root/bus/pci/devices/$pci_device/d3cold_allowed ]] || continue
+
+    printf '%s\n' "$pci_device"
+    return 0
+  done
+
+  return 1
+}
+
 if [[ $MACBOOK_MODEL =~ MacBook(8,1|9,1|10,1)|MacBookPro13,[123]|MacBookPro14,[123] ]]; then
   echo "Detected MacBook model: $MACBOOK_MODEL"
 
-  NVME_DEVICE="/sys/bus/pci/devices/0000:01:00.0/d3cold_allowed"
+  NVME_PCI=$(find_nvme_pci_device || true)
+
+  if [[ -n $NVME_PCI ]]; then
+    NVME_DEVICE="/sys/bus/pci/devices/$NVME_PCI/d3cold_allowed"
+    SERVICE_FILE="/etc/systemd/system/omarchy-nvme-suspend-fix.service"
+    LEGACY_PCI="0000:01:00.0"
+    LEGACY_DEVICE="/sys/bus/pci/devices/$LEGACY_PCI/d3cold_allowed"
+
+    echo "Applying NVMe suspend fix to $NVME_PCI..."
 
-  if [[ -f $NVME_DEVICE ]]; then
-    echo "Applying NVMe suspend fix..."
+    # Older Omarchy installs always targeted 01:00.0. On dGPU MacBooks that can
+    # be the graphics controller instead of the NVMe. Only undo that setting
+    # when the existing Omarchy unit proves it contains the legacy target.
+    if [[ $NVME_PCI != "$LEGACY_PCI" && -f $LEGACY_DEVICE && -f $SERVICE_FILE ]] &&
+      { grep -Fq '0000\:01\:00.0/d3cold_allowed' "$SERVICE_FILE" ||
+        grep -Fq '0000:01:00.0/d3cold_allowed' "$SERVICE_FILE"; }; then
+      echo "Restoring D3cold on legacy non-NVMe PCI target $LEGACY_PCI..."
+      echo 1 | sudo tee "$LEGACY_DEVICE" >/dev/null
+    fi
 
     sudo mkdir -p /etc/systemd/system
-    sudo tee /etc/systemd/system/omarchy-nvme-suspend-fix.service >/dev/null <<'EOF'
+    sudo tee "$SERVICE_FILE" >/dev/null <<EOF
 [Unit]
 Description=Omarchy NVMe Suspend Fix for MacBook
 
 [Service]
-ExecStart=/bin/bash -c 'echo 0 > /sys/bus/pci/devices/0000\:01\:00.0/d3cold_allowed'
+ExecStart=/bin/bash -c 'echo 0 > $NVME_DEVICE'
 
 [Install]
 WantedBy=multi-user.target
 EOF
 
+    sudo systemctl daemon-reload
     sudo systemctl enable omarchy-nvme-suspend-fix.service
   else
-    echo "Warning: NVMe device not found at expected PCI address (0000:01:00.0)"
+    echo "Warning: No PCI NVMe controller with d3cold_allowed was found"
     echo "This fix may not be needed for this MacBook model"
   fi
 fi

test/shell.d/apple-nvme-suspend-test.sh

@@ -0,0 +1,57 @@
+#!/bin/bash
+
+set -euo pipefail
+
+source "$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)/base-test.sh"
+
+script="$ROOT/install/hardware/apple/fix-suspend-nvme.sh"
+
+finder=$(awk '
+  $0 == "find_nvme_pci_device() {" { inside = 1 }
+  inside { print }
+  inside && $0 == "}" { exit }
+' "$script")
+[[ -n $finder ]] || fail "Apple NVMe suspend fix exposes a controller discovery helper"
+eval "$finder"
+
+test_tmp=$(mktemp -d)
+trap 'rm -rf "$test_tmp"' EXIT
+fake_sys="$test_tmp/sys"
+mkdir -p "$fake_sys/class/nvme/nvme0" \
+  "$fake_sys/bus/pci/devices/0000:02:00.0"
+touch "$fake_sys/bus/pci/devices/0000:02:00.0/d3cold_allowed"
+ln -s ../../../bus/pci/devices/0000:02:00.0 "$fake_sys/class/nvme/nvme0/device"
+
+[[ $(find_nvme_pci_device "$fake_sys") == 0000:02:00.0 ]] ||
+  fail "Apple NVMe suspend fix follows the actual NVMe controller PCI address"
+pass "Apple NVMe suspend fix follows the actual NVMe controller PCI address"
+
+rm "$fake_sys/class/nvme/nvme0/device"
+mkdir -p "$fake_sys/devices/platform/not-a-pci-device" \
+  "$fake_sys/class/nvme/nvme1" \
+  "$fake_sys/bus/pci/devices/0000:03:00.0"
+ln -s ../../../devices/platform/not-a-pci-device "$fake_sys/class/nvme/nvme0/device"
+touch "$fake_sys/bus/pci/devices/0000:03:00.0/d3cold_allowed"
+ln -s ../../../bus/pci/devices/0000:03:00.0 "$fake_sys/class/nvme/nvme1/device"
+
+[[ $(find_nvme_pci_device "$fake_sys") == 0000:03:00.0 ]] ||
+  fail "Apple NVMe suspend fix rejects non-PCI controller paths"
+pass "Apple NVMe suspend fix rejects non-PCI controller paths"
+
+rm "$fake_sys/bus/pci/devices/0000:03:00.0/d3cold_allowed"
+if find_nvme_pci_device "$fake_sys" >/dev/null; then
+  fail "Apple NVMe suspend fix requires a d3cold_allowed target"
+fi
+pass "Apple NVMe suspend fix requires a d3cold_allowed target"
+
+grep -Fq 'NVME_PCI=$(find_nvme_pci_device || true)' "$script" ||
+  fail "Apple NVMe suspend fix uses discovered controller at install time"
+grep -Fq "ExecStart=/bin/bash -c 'echo 0 > \$NVME_DEVICE'" "$script" ||
+  fail "Apple NVMe suspend unit targets the discovered controller"
+pass "Apple NVMe suspend unit targets the discovered controller"
+
+grep -Fq 'NVME_PCI != "$LEGACY_PCI"' "$script" ||
+  fail "Apple NVMe suspend fix only restores a distinct legacy target"
+grep -Fq "grep -Fq '0000\\:01\\:00.0/d3cold_allowed' \"\$SERVICE_FILE\"" "$script" ||
+  fail "Apple NVMe suspend fix verifies the old Omarchy unit before restoring D3cold"
+pass "Apple NVMe suspend fix gates legacy D3cold restoration on the old unit"