From a8c1143eee01e418d488fd704258f2f8d7a88906 Mon Sep 17 00:00:00 2001
From: Dan Bungert <daniel.bungert@canonical.com>
Date: Thu, 22 Feb 2024 18:40:25 -0700
Subject: [PATCH] filesystem: create zpool with encryption info

---
 subiquity/common/filesystem/manipulator.py     | 18 +++++++++++++++++-
 subiquity/models/filesystem.py                 |  6 ++++++
 subiquity/server/controllers/filesystem.py     | 13 +++++++++++--
 .../controllers/tests/test_filesystem.py       |  5 +++++
 4 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/subiquity/common/filesystem/manipulator.py b/subiquity/common/filesystem/manipulator.py
index 8c6d5456..ba6d106f 100644
--- a/subiquity/common/filesystem/manipulator.py
+++ b/subiquity/common/filesystem/manipulator.py
@@ -20,6 +20,7 @@ from curtin.block import get_resize_fstypes
 from subiquity.common.filesystem import boot, gaps
 from subiquity.common.types import Bootloader
 from subiquity.models.filesystem import Partition, align_up
+from subiquitycore.utils import write_named_tempfile
 
 log = logging.getLogger("subiquity.common.filesystem.manipulator")
 
@@ -188,7 +189,16 @@ class FilesystemManipulator:
         self.create_filesystem(dmc, dict(fstype="swap"))
         return dmc
 
-    def create_zpool(self, device, pool, mountpoint, boot=False, canmount="on"):
+    def create_zpool(
+        self,
+        device,
+        pool,
+        mountpoint,
+        boot=False,
+        canmount="on",
+        encryption_style=None,
+        key=None,
+    ):
         fs_properties = dict(
             atime=None,
             acltype="posixacl",
@@ -201,6 +211,10 @@ class FilesystemManipulator:
             xattr="sa",
         )
 
+        keyfile = None
+        if key is not None:
+            keyfile = write_named_tempfile("zpool-key-", key)
+
         pool_properties = dict(ashift=12, autotrim="on", version=None)
         default_features = True
         if boot:
@@ -217,6 +231,8 @@ class FilesystemManipulator:
             default_features=default_features,
             fs_properties=fs_properties,
             pool_properties=pool_properties,
+            encryption_style=encryption_style,
+            keyfile=keyfile,
         )
 
     def delete(self, obj):
diff --git a/subiquity/models/filesystem.py b/subiquity/models/filesystem.py
index 1e1fa7e6..4e299195 100644
--- a/subiquity/models/filesystem.py
+++ b/subiquity/models/filesystem.py
@@ -1310,6 +1310,8 @@ class ZPool:
     fs_properties: Optional[dict] = None
 
     default_features: Optional[bool] = True
+    encryption_style: Optional[str] = None
+    keyfile: Optional[str] = None
 
     component_name = "vdev"
 
@@ -2255,6 +2257,8 @@ class FilesystemModel:
         default_features=True,
         fs_properties=None,
         pool_properties=None,
+        encryption_style=None,
+        keyfile=None,
     ):
         zpool = ZPool(
             m=self,
@@ -2264,6 +2268,8 @@ class FilesystemModel:
             default_features=default_features,
             pool_properties=pool_properties,
             fs_properties=fs_properties,
+            encryption_style=encryption_style,
+            keyfile=keyfile,
         )
         self._actions.append(zpool)
         return zpool
diff --git a/subiquity/server/controllers/filesystem.py b/subiquity/server/controllers/filesystem.py
index 6ec06ceb..2d18abd3 100644
--- a/subiquity/server/controllers/filesystem.py
+++ b/subiquity/server/controllers/filesystem.py
@@ -565,7 +565,9 @@ class FilesystemController(SubiquityController, FilesystemManipulator):
         bootfs_size = align_up(sizes.get_bootfs_size(gap.size), part_align)
         gap_boot, gap_rest = gap.split(bootfs_size)
         bpart = self.create_partition(device, gap_boot, dict(fstype=None))
-        encrypted = choice.password is not None
+        encryption_style = None
+        if encrypted := choice.password is not None:
+            encryption_style = "luks_keystore"
 
         avail = gap_rest.size - self._info.min_size
         swap_size = align_down(swap.suggested_swapsize(avail=avail), part_align)
@@ -586,7 +588,14 @@ class FilesystemController(SubiquityController, FilesystemManipulator):
         bpool.create_zfs("BOOT", canmount="off", mountpoint="none")
         bpool.create_zfs(f"BOOT/ubuntu_{uuid}", mountpoint="/boot")
 
-        rpool = self.create_zpool(rpart, "rpool", "/", canmount="off")
+        rpool = self.create_zpool(
+            rpart,
+            "rpool",
+            "/",
+            canmount="off",
+            encryption_style=encryption_style,
+            key=choice.password,
+        )
         rpool.create_zfs("ROOT", canmount="off", mountpoint="none")
         rpool.create_zfs(f"ROOT/ubuntu_{uuid}", mountpoint="/")
         rpool.create_zfs(f"ROOT/ubuntu_{uuid}/var", canmount="off")
diff --git a/subiquity/server/controllers/tests/test_filesystem.py b/subiquity/server/controllers/tests/test_filesystem.py
index ab767374..4f9768eb 100644
--- a/subiquity/server/controllers/tests/test_filesystem.py
+++ b/subiquity/server/controllers/tests/test_filesystem.py
@@ -581,6 +581,8 @@ class TestGuided(IsolatedAsyncioTestCase):
         [rpool] = self.model._all(type="zpool", pool="rpool")
         self.assertIsNone(rpool.path)
         self.assertEqual([root], rpool.vdevs)
+        self.assertIsNone(rpool.encryption_style)
+        self.assertIsNone(rpool.keyfile)
         [bpool] = self.model._all(type="zpool", pool="bpool")
         self.assertIsNone(bpool.path)
         self.assertEqual([boot], bpool.vdevs)
@@ -618,6 +620,9 @@ class TestGuided(IsolatedAsyncioTestCase):
         [rpool] = self.model._all(type="zpool", pool="rpool")
         self.assertIsNone(rpool.path)
         self.assertEqual([root], rpool.vdevs)
+        self.assertEqual("luks_keystore", rpool.encryption_style)
+        with open(rpool.keyfile) as fp:
+            self.assertEqual("passw0rd", fp.read())
         [bpool] = self.model._all(type="zpool", pool="bpool")
         self.assertIsNone(bpool.path)
         self.assertEqual([boot], bpool.vdevs)