From d70686bc098022bdcf9c98308de6fef31f03c510 Mon Sep 17 00:00:00 2001
From: Vunny Sodhi <vunny.sodhi@unikie.com>
Date: Tue, 31 Dec 2024 13:22:24 +0200
Subject: [PATCH] zfs: Address review comments

Signed-off-by: Vunny Sodhi <vunny.sodhi@unikie.com>
---
 modules/disko/disko-ab-partitions.nix | 13 +++----
 modules/disko/disko-zfs-postboot.nix  | 49 ++++++++++++++++++---------
 2 files changed, 38 insertions(+), 24 deletions(-)

diff --git a/modules/disko/disko-ab-partitions.nix b/modules/disko/disko-ab-partitions.nix
index 7311c9212..eb5fca0c4 100644
--- a/modules/disko/disko-ab-partitions.nix
+++ b/modules/disko/disko-ab-partitions.nix
@@ -45,6 +45,9 @@
       initrd.availableKernelModules = [ "zfs" ];
       supportedFilesystems = [ "zfs" ];
       zfs.extraPools = [ "zfs_data" ];
+      initrd.luks.devices.zfs_data = {
+        device = "/dev/disk/by-partlabel/disk-disk1-zfs_data";
+      };
     };
     disko = {
       # 8GB is the recommeneded minimum for ZFS, so we are using this for VMs to avoid `cp` oom errors.
@@ -102,14 +105,8 @@
               zfs_data = {
                 size = "100%";
                 content = {
-                  name = "zfs_data";
-                  type = "luks";
-                  # TODO: Have a better password mechanism later
-                  settings.keyFile = "${pkgs.writeText "password" "ghaf"}";
-                  content = {
-                    type = "zfs";
-                    pool = "zfs_data";
-                  };
+                  type = "zfs";
+                  pool = "zfs_data";
                 };
               };
             };
diff --git a/modules/disko/disko-zfs-postboot.nix b/modules/disko/disko-zfs-postboot.nix
index 0f4dbd8e5..013df9926 100644
--- a/modules/disko/disko-zfs-postboot.nix
+++ b/modules/disko/disko-zfs-postboot.nix
@@ -17,33 +17,50 @@ let
       set -xeuo pipefail
 
       # Check which physical disk is used by ZFS
-      ENCRYPTED_POOL=zfs_data
-      zpool import -f $ENCRYPTED_POOL
-      ZFS_POOLNAME=$(zpool list | grep -v NAME | grep $ENCRYPTED_POOL | awk '{print $1}')
-      ZFS_LOCATION=$(zpool status -P | grep dev | grep "$ZFS_POOLNAME" | awk '{print $1}')
+      ENCRYPTED_POOLNAME=zfs_data
+      zpool import -f "$ENCRYPTED_POOLNAME"
+      ZFS_POOLNAME=$(zpool list | grep -v NAME | grep $ENCRYPTED_POOLNAME | awk '{print $1}')
+      ZFS_LOCATION=$(zpool status "$ZFS_POOLNAME" -P | grep dev | awk '{print $1}')
 
       # Get the actual device path
-      P_DEVPATH=$(cryptsetup status "$ZFS_POOLNAME" | grep device | awk '{print $2}')
+      P_DEVPATH=$(readlink -f "$ZFS_LOCATION")
 
-      # Extract the partition number using regex
       if [[ "$P_DEVPATH" =~ [0-9]+$ ]]; then
         PARTNUM=$(echo "$P_DEVPATH" | grep -o '[0-9]*$')
-        PARENT_DISK=/dev/$(lsblk -no pkname "$P_DEVPATH" | head -n 1)
+        PARENT_DISK=/dev/$(lsblk -no pkname "$P_DEVPATH")
       else
-        echo "No partition number found in device path: $P_DEVPATH"
+       echo "No partition number found in device path: $P_DEVPATH"
       fi
 
-      # Fix GPT first
-      sgdisk "$PARENT_DISK" -e
+      set +o pipefail
+      # Check if zfs pool has luks headers
+      if (cryptsetup status "$ZFS_POOLNAME") | grep -q "is inactive"; then
+        # Fix GPT first
+        sgdisk "$PARENT_DISK" -e
 
-      # Call partprobe to update kernel's partitions
-      partprobe
+        # Call partprobe to update kernel's partitions
+        partprobe
 
-      # Extend the partition to use unallocated space
-      parted -s -a opt "$PARENT_DISK" "resizepart $PARTNUM 100%"
+        # Extend the partition to use unallocated space
+        parted -s -a opt "$PARENT_DISK" "resizepart $PARTNUM 100%"
 
-      # Extend ZFS pool to use newly allocated space
-      zpool online -e "$ZFS_POOLNAME" "$ZFS_LOCATION"
+        # Extend ZFS pool to use newly allocated space
+        zpool online -e "$ZFS_POOLNAME" "$ZFS_LOCATION"
+
+        # Format pool with LUKS
+        zpool export "$ZFS_POOLNAME"
+        # TODO: Remove hardcoded password and have better password mechanism for formating
+        echo -n ghaf | cryptsetup luksFormat --type luks2 -q "$ZFS_LOCATION"
+        echo -n ghaf | cryptsetup luksOpen "$ZFS_LOCATION" "$ZFS_POOLNAME" --persistent
+
+        # Create pool, datasets as luksFormat will erase pools, ZFS datasets stored on that partition
+        zpool create -f "$ZFS_POOLNAME" /dev/mapper/"$ZFS_POOLNAME"
+        zfs create -o quota=30G "$ZFS_POOLNAME"/vm_storage
+        zfs create -o quota=10G "$ZFS_POOLNAME"/reserved
+        zfs create -o quota=50G "$ZFS_POOLNAME"/gp_storage
+        zfs create "$ZFS_POOLNAME"/storagevm
+        zfs create "$ZFS_POOLNAME"/recovery
+      fi
     '';
   };