Begin initial attempt at getting zfs setup working

rust/program/disko-zpool-expand/src/main.rs

@@ -1,117 +0,0 @@
-use serde::Deserialize;
-use std::{collections::BTreeMap, path::PathBuf};
-
-use anyhow::Context as _;
-use clap::{Parser, Subcommand};
-
-fn main() {
-    common::entrypoint(program);
-}
-
-#[derive(Parser, Debug)]
-#[command(version, about, long_about = None)]
-pub struct Args {
-    #[command(subcommand)]
-    pub command: Commands,
-}
-
-#[derive(Debug, Subcommand)]
-pub enum Commands {
-    /// Expands the partitions based on a zpool and brings the pool up to the new size.
-    #[command(name = "expand-zpool")]
-    ExpandZpool(ExpandZpool),
-}
-
-#[derive(Debug, Clone, clap::Args)]
-pub struct ExpandZpool {
-    /// Name of the pool to expand
-    pool_name: String,
-}
-
-fn program() -> anyhow::Result<()> {
-    let args = Args::parse();
-    match args.command {
-        Commands::ExpandZpool(pool) => expand_zpool(pool),
-    }
-}
-
-#[derive(Deserialize)]
-struct ZpoolStatus {
-    pools: BTreeMap<String, ZpoolStatusPool>,
-}
-
-#[derive(Deserialize)]
-struct ZpoolStatusPool {
-    state: Option<ZpoolState>,
-    vdevs: BTreeMap<String, ZpoolStatusVdev>,
-}
-
-#[derive(Clone, Copy, Deserialize, PartialEq)]
-enum ZpoolState {
-    #[serde(rename = "ONLINE")]
-    Online,
-}
-
-#[derive(Deserialize)]
-struct ZpoolStatusVdev {
-    vdevs: BTreeMap<String, ZpoolStatusVdevVdev>,
-}
-#[derive(Deserialize)]
-struct ZpoolStatusVdevVdev {
-    path: PathBuf,
-}
-
-fn expand_zpool(p: ExpandZpool) -> anyhow::Result<()> {
-    let mut proc = common::proc::Command::new("zpool");
-    proc.args(["status", "--json", &p.pool_name]);
-    let result: ZpoolStatus = proc
-        .try_spawn_to_json()
-        .context("Could not get zpool status")?;
-
-    let pool = result
-        .pools
-        .get(&p.pool_name)
-        .context("Could not find requested pool in status output")?;
-
-    if !pool
-        .state
-        .as_ref()
-        .is_some_and(|st| *st == ZpoolState::Online)
-    {
-        return Err(anyhow::format_err!("Zpool {} is not online", p.pool_name));
-    }
-
-    for vdev in pool.vdevs.values() {
-        for vdev in vdev.vdevs.values() {
-            let partition_dev = vdev.path.display().to_string();
-            let Some(dev) = partition_dev.strip_suffix("-part1") else {
-                return Err(anyhow::format_err!(
-                    "Expected vdev path {} to end with -part1",
-                    vdev.path.display()
-                ));
-            };
-            let mut proc = common::proc::Command::new("growpart");
-            proc.args([dev, "1"]);
-            let (stdout, _stderr, status) = proc.spawn_into_parts()?;
-            if !status.success() && !stdout.starts_with("NOCHANGE: ") {
-                return Err(anyhow::format_err!(
-                    "Could not resize partitin for {}, err: {stdout}",
-                    vdev.path.display()
-                ));
-            }
-            // let name = partition_dev
-            //     .split("/")
-            //     .last()
-            //     .expect("Should always have at least one element");
-            let mut proc = common::proc::Command::new("zpool");
-            proc.args(["online", "-e", &p.pool_name, &partition_dev]);
-            proc.try_spawn_to_string().with_context(|| {
-                format!(
-                    "Could not bring zpool {} online with expand flag",
-                    p.pool_name
-                )
-            })?;
-        }
-    }
-    Ok(())
-}

rust/program/zpool-setup/Cargo.toml

@@ -1,8 +1,8 @@
 [package]
-name = "disko-zpool-expand"
+name = "zpool-setup"
 edition = "2024"
 version = "1.0.0"
-metadata.crane.name = "disko-zpool-expand"
+metadata.crane.name = "zpool-setup"
 
 [dependencies]
 anyhow = { workspace = true }

rust/program/zpool-setup/src/cli.rs

@@ -0,0 +1,93 @@
+use std::{borrow::Cow, collections::BTreeMap, path::PathBuf, str::FromStr};
+
+use serde::Deserialize;
+
+use crate::disk_mapping::DiskMapping;
+
+#[derive(Clone, Copy, Debug, Deserialize, PartialEq)]
+pub enum VdevMode {
+    #[serde(rename = "mirror")]
+    Mirror,
+    #[serde(rename = "raidz")]
+    Raidz,
+    #[serde(rename = "raidz1")]
+    Raidz1,
+    #[serde(rename = "raidz2")]
+    Raidz2,
+    #[serde(rename = "raidz3")]
+    Raidz3,
+}
+
+impl VdevMode {
+    fn str(&self) -> &'static str {
+        match self {
+            Self::Mirror => "mirror",
+            Self::Raidz => "raidz",
+            Self::Raidz1 => "raidz1",
+            Self::Raidz2 => "raidz2",
+            Self::Raidz3 => "raidz3",
+        }
+    }
+}
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct Vdev {
+    pub mode: VdevMode,
+    pub members: Vec<String>,
+}
+
+impl Vdev {
+    pub fn cli_args(&self, disk_mapper: &DiskMapping) -> anyhow::Result<Vec<Cow<'static, str>>> {
+        let mut args = Vec::with_capacity(self.members.len() + 1);
+        if self.members.len() > 1 || self.mode != VdevMode::Mirror {
+            args.push(Cow::Borrowed(self.mode.str()));
+        }
+        for member in self.members.iter() {
+            let resolved = disk_mapper.resolve(member)?;
+            args.push(resolved.into());
+        }
+        Ok(args)
+    }
+}
+
+#[derive(Clone, Debug, Deserialize)]
+#[serde(transparent)]
+pub struct Vdevs(pub Vec<Vdev>);
+
+impl FromStr for Vdevs {
+    type Err = anyhow::Error;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        common::json::from_str(s)
+    }
+}
+
+#[derive(Clone, Debug, Deserialize)]
+#[serde(transparent)]
+pub struct Options(pub BTreeMap<String, String>);
+
+impl FromStr for Options {
+    type Err = anyhow::Error;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        common::json::from_str(s)
+    }
+}
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct Dataset {
+    pub options: Options,
+    pub mountpoint: Option<PathBuf>,
+}
+
+#[derive(Clone, Debug, Deserialize)]
+#[serde(transparent)]
+pub struct Datasets(pub BTreeMap<String, Dataset>);
+
+impl FromStr for Datasets {
+    type Err = anyhow::Error;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        common::json::from_str(s)
+    }
+}

rust/program/zpool-setup/src/disk_mapping.rs

@@ -0,0 +1,41 @@
+use std::collections::BTreeMap;
+
+use serde::Deserialize;
+
+#[derive(Debug, Deserialize)]
+#[serde(transparent)]
+pub struct DiskMapping(DiskmappingFile);
+
+impl DiskMapping {
+    pub fn resolve(&self, name: &str) -> anyhow::Result<String> {
+        let resolved = self
+            .0
+            .disks
+            .get(name)
+            .ok_or_else(|| anyhow::format_err!("No mapping for disk named {}", name))?;
+
+        Ok(self.0.template.execute(resolved.linux_device.as_str()))
+    }
+}
+
+#[derive(Debug, Deserialize)]
+struct DiskmappingFile {
+    disks: BTreeMap<String, Disk>,
+    template: DeviceTemplate,
+}
+
+#[derive(Debug, Deserialize)]
+struct Disk {
+    #[serde(rename = "linuxDevice")]
+    linux_device: String,
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(transparent)]
+struct DeviceTemplate(String);
+
+impl DeviceTemplate {
+    pub fn execute(&self, name: &str) -> String {
+        self.0.replace("{id}", name)
+    }
+}

rust/program/zpool-setup/src/main.rs

@@ -0,0 +1,200 @@
+use serde::Deserialize;
+use std::{collections::BTreeMap, path::PathBuf};
+
+use anyhow::Context as _;
+use clap::{Parser, Subcommand};
+
+mod cli;
+mod disk_mapping;
+mod zfs;
+
+fn main() {
+    common::entrypoint(program);
+}
+
+#[derive(Parser, Debug)]
+#[command(version, about, long_about = None)]
+pub struct Args {
+    #[command(subcommand)]
+    pub command: Commands,
+}
+
+#[derive(Debug, Subcommand)]
+pub enum Commands {
+    /// Expands the partitions based on a zpool and brings the pool up to the new size.
+    #[command(name = "setup")]
+    Setup(SetupZpool),
+}
+
+#[derive(Debug, Clone, clap::Args)]
+pub struct SetupZpool {
+    /// Openbao mount of the encryption key for the pool. Can only omit during test.
+    #[arg(long = "encryption-key-mount")]
+    encryption_key_mount: Option<String>,
+
+    /// Openbao name of the encryption key for the pool. Can only omit during test.
+    #[arg(long = "encryption-key-name")]
+    encryption_key_name: Option<String>,
+
+    /// Openbao name of the encryption field for the pool. Can only omit during test.
+    #[arg(long = "encryption-key-name")]
+    encryption_key_field: Option<String>,
+
+    /// Vdevs of the pool
+    #[arg(long = "vdevs")]
+    vdevs: cli::Vdevs,
+
+    /// Options of the pool
+    #[arg(long = "zpool-options")]
+    zpool_options: cli::Options,
+
+    /// Options of the root file system
+    #[arg(long = "root-fs-options")]
+    root_fs_options: cli::Options,
+
+    /// Datasets the pool should have
+    #[arg(long = "datasets")]
+    datasets: cli::Datasets,
+
+    /// Name of the pool to expand
+    pool_name: String,
+}
+
+struct TempDir {
+    path: PathBuf,
+}
+
+impl TempDir {
+    pub fn try_new(template: &str) -> anyhow::Result<Self> {
+        let mut proc = common::proc::Command::new("mktemp");
+        proc.args(["-dt", template]);
+        let path: PathBuf = proc.try_spawn_to_string()?.into();
+        common::fs::create_dir_recursive(&path)?;
+        Ok(Self { path })
+    }
+}
+
+impl Drop for TempDir {
+    fn drop(&mut self) {
+        common::fs::remove_dir_recursive(&self.path).expect("Could not clean up after temp dir");
+    }
+}
+
+impl SetupZpool {
+    fn encryption_key(&self) -> anyhow::Result<String> {
+        let is_test = common::env::read_env("ZFS_TEST").is_ok_and(|t| t == "true");
+        if is_test {
+            return Ok(String::from("testtest"));
+        }
+        let role_id_file = common::env::read_path_env("VAULT_ROLE_ID_FILE")?;
+        let role_id = common::fs::read_to_string(&role_id_file)?;
+        let secret_id_file = common::env::read_path_env("VAULT_SECRET_ID_FILE")?;
+        let secret_id = common::fs::read_to_string(&secret_id_file)?;
+        let tmpdir = TempDir::try_new("zpool-setup.XXXXXX")?;
+        common::fs::write_file_string(
+            &tmpdir.path.join(".vault"),
+            "token_helper = \"/bin/true\"",
+            common::fs::user_only_file_permissions(),
+        )?;
+        let mut login_proc = common::proc::Command::new("bao");
+        login_proc.env("HOME", tmpdir.path.display().to_string());
+        login_proc.args(["write", "-field=token", "auth/approle/login"]);
+        login_proc.args([
+            format!("role_id={role_id}"),
+            format!("secret_id={secret_id}"),
+        ]);
+        let vault_token = login_proc.try_spawn_to_string()?;
+        let (field, name, mount) = match (
+            self.encryption_key_field.as_deref(),
+            self.encryption_key_name.as_deref(),
+            self.encryption_key_mount.as_deref(),
+        ) {
+            (Some(field), Some(name), Some(mount)) => (field, name, mount),
+            _ => {
+                return Err(anyhow::format_err!(
+                    "Missing one of --encryption-key-mount, --encryption-key-name, --encryption-key-field"
+                ));
+            }
+        };
+        let mut proc = common::proc::Command::new("bao");
+        proc.env("HOME", tmpdir.path.display().to_string());
+        proc.env_sensitive("VAULT_TOKEN", vault_token);
+        proc.args(["kv", "get"]);
+        proc.arg(format!("-field={field}"));
+        proc.arg(format!("-mount={mount}"));
+        proc.arg(name);
+
+        proc.try_spawn_to_string()
+    }
+}
+
+fn program() -> anyhow::Result<()> {
+    let args = Args::parse();
+    match args.command {
+        Commands::Setup(setup) => setup_zpool(setup),
+    }
+}
+
+#[derive(Deserialize)]
+struct ZpoolStatus {
+    pools: BTreeMap<String, ZpoolStatusPool>,
+}
+
+#[derive(Deserialize)]
+struct ZpoolStatusPool {
+    state: Option<ZpoolState>,
+}
+
+#[derive(Clone, Copy, Deserialize, PartialEq)]
+enum ZpoolState {
+    #[serde(rename = "ONLINE")]
+    Online,
+}
+
+fn setup_zpool(p: SetupZpool) -> anyhow::Result<()> {
+    let disk_mapping_file = common::env::read_path_env("DISK_MAPPING_FILE")?;
+    let disk_mapping = common::fs::read_to_string(&disk_mapping_file)?;
+    let disk_mapping = common::json::from_str(&disk_mapping)?;
+    if !zfs::import_pool(&p.pool_name)? {
+        let encryption_key = p.encryption_key()?;
+        zfs::create_pool(&p, &disk_mapping, &encryption_key)?;
+        for (name, dataset) in p.datasets.0.iter() {
+            zfs::create_dataset_recursive(&p.pool_name, name, dataset)?;
+        }
+        zfs::mount_all(&p.pool_name)?;
+        return Ok(());
+    }
+    let mut proc: common::proc::Command = common::proc::Command::new("zpool");
+    proc.args(["status", "--json", &p.pool_name]);
+    let result: ZpoolStatus = proc
+        .try_spawn_to_json()
+        .context("Could not get zpool status")?;
+
+    let pool = result
+        .pools
+        .get(&p.pool_name)
+        .context("Could not find requested pool in status output")?;
+
+    if !pool
+        .state
+        .as_ref()
+        .is_some_and(|st| *st == ZpoolState::Online)
+    {
+        return Err(anyhow::format_err!("Zpool {} is not online", p.pool_name));
+    }
+
+    for vdev in p.vdevs.0.iter() {
+        for member in vdev.members.iter() {
+            let resolved = disk_mapping.resolve(member)?;
+            zfs::resize_disk(&p.pool_name, &resolved)?;
+        }
+    }
+    if zfs::encryption_key_needs_load(&p.pool_name)? {
+        let encryption_key = p.encryption_key()?;
+        zfs::load_key(&p.pool_name, &encryption_key)?;
+    }
+    // TODO: Update pool options, and all fs options, and create missing datasets.
+    //       Maybe for extranous datasets, set mountpoint=none ?
+    zfs::mount_all(&p.pool_name)?;
+    Ok(())
+}

rust/program/zpool-setup/src/zfs.rs

@@ -0,0 +1,174 @@
+use std::collections::BTreeMap;
+
+use anyhow::Context as _;
+use common::proc::Command;
+use serde::Deserialize;
+
+use crate::{SetupZpool, cli::Dataset, disk_mapping::DiskMapping};
+
+#[derive(Debug, Deserialize, PartialEq)]
+enum ZpoolState {
+    #[serde(rename = "ONLINE")]
+    Online,
+}
+
+pub fn import_pool(name: &str) -> anyhow::Result<bool> {
+    // Test if the pool exists and is already imported
+    let mut exists_proc = Command::new("zpool");
+    exists_proc.args(["status", name]);
+    if exists_proc.try_spawn_to_bytes().is_ok() {
+        return Ok(true);
+    }
+    // Try to import the pool if it exists
+    let mut proc = Command::new("zpool");
+    proc.args(["import", name]);
+
+    let (_stdout, stderr, exit_code) = proc.spawn_into_parts()?;
+    if exit_code.success() {
+        return Ok(true);
+    }
+    if stderr.contains("no such pool available") {
+        // The pool doesn't exist
+        return Ok(false);
+    }
+    Err(anyhow::format_err!(
+        "Could not import pool {name}, stderr: {stderr}"
+    ))
+}
+
+pub fn create_pool(
+    zpool: &SetupZpool,
+    disk_mapping: &DiskMapping,
+    encryption_key: &str,
+) -> anyhow::Result<()> {
+    let mut proc = Command::new("zpool");
+    proc.args([
+        "create",
+        zpool.pool_name.as_str(),
+        "-m",
+        "none",
+        "-o",
+        "feature@device_removal=enabled",
+        "-o",
+        "feature@draid=enabled",
+        "-o",
+        "feature@raidz_expansion=enabled",
+        "-o",
+        "feature@zilsaxattr=enabled",
+        "-o",
+        "feature@zstd_compress=enabled",
+        "-o",
+        "cachefile=none",
+    ]);
+
+    for (key, value) in zpool.zpool_options.0.iter() {
+        proc.args(["-o", &format!("{key}={value}")]);
+    }
+
+    for (key, value) in zpool.root_fs_options.0.iter() {
+        proc.args(["-O", &format!("{key}={value}")]);
+    }
+
+    proc.args([
+        "-O",
+        "encryption=aes-256-gcm",
+        "-O",
+        "keyformat=passphrase",
+        "-O",
+        "keylocation=prompt",
+    ]);
+
+    for vdev in zpool.vdevs.0.iter() {
+        proc.args(vdev.cli_args(disk_mapping)?.into_iter());
+    }
+
+    proc.stdin_string(encryption_key);
+    proc.try_spawn_to_bytes()?;
+
+    Ok(())
+}
+
+pub fn create_dataset_recursive(
+    pool_name: &str,
+    dataset_name: &str,
+    dataset: &Dataset,
+) -> anyhow::Result<()> {
+    let mut proc = Command::new("zfs");
+    let name = format!("{pool_name}/{dataset_name}");
+    proc.args(["create", "-p", "-u"]);
+    if let Some(mountpoint) = dataset.mountpoint.as_deref() {
+        proc.arg("-o");
+        proc.arg(format!("mountpoint={}", mountpoint.display()));
+    }
+    for (key, value) in dataset.options.0.iter() {
+        proc.arg("-o");
+        proc.arg(format!("{key}={value}"));
+    }
+
+    proc.arg(name);
+
+    let _ = proc.try_spawn_to_bytes()?;
+    Ok(())
+}
+
+pub fn mount_all(pool: &str) -> anyhow::Result<()> {
+    let mut proc = Command::new("zfs");
+    proc.args(["mount", "-R", pool]);
+    proc.try_spawn_to_bytes()?;
+    Ok(())
+}
+
+pub fn resize_disk(pool_name: &str, device: &str) -> anyhow::Result<()> {
+    let mut proc = Command::new("zpool");
+    proc.args(["online", "-e", pool_name, device]);
+    let _ = proc.try_spawn_to_bytes().with_context(|| {
+        format!("Could not bring zpool {pool_name} online with expand flag for device {device}",)
+    })?;
+    Ok(())
+}
+
+pub fn load_key(pool_name: &str, encryption_key: &str) -> anyhow::Result<()> {
+    let mut proc = Command::new("zfs");
+    proc.args(["load-key", "-r", "-L", "prompt", pool_name]);
+    proc.stdin_bytes(encryption_key);
+    proc.try_spawn_to_string()?;
+    Ok(())
+}
+
+pub fn encryption_key_needs_load(pool_name: &str) -> anyhow::Result<bool> {
+    #[derive(Deserialize)]
+    struct PoolEncStatus {
+        datasets: BTreeMap<String, PoolEncStatusDataset>,
+    }
+
+    #[derive(Deserialize)]
+    struct PoolEncStatusDataset {
+        properties: PoolEncStatusDatasetProperties,
+    }
+
+    #[derive(Deserialize)]
+    struct PoolEncStatusDatasetProperties {
+        keystatus: PoolEncStatusDatasetProperty,
+    }
+
+    #[derive(Deserialize)]
+    struct PoolEncStatusDatasetProperty {
+        value: Option<String>,
+    }
+
+    // "$(zfs list -j -o keystatus zroot/mailserver | jq --raw-output '.datasets."zroot/mailserver".properties.keystatus.value')" == "unavailable"
+    let mut proc = Command::new("zfs");
+    proc.args(["list", "-j", "-o", "keystatus", pool_name]);
+    let json: PoolEncStatus = proc.try_spawn_to_json()?;
+    let pool = json
+        .datasets
+        .get(pool_name)
+        .ok_or_else(|| anyhow::format_err!("Pool {pool_name} not found in status output"))?;
+
+    Ok(pool
+        .properties
+        .keystatus
+        .value
+        .as_deref()
+        .is_some_and(|v| v == "unavailable"))
+}