From 66c6af46dd6d79357e86b1d41b7ac087df62dde7 Mon Sep 17 00:00:00 2001
From: Lleyton Gray <lleyton@fyralabs.com>
Date: Wed, 22 May 2024 00:07:07 -0700
Subject: [PATCH] feat: allow setting GPT partition flags

---
 src/config.rs | 54 ++++++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 51 insertions(+), 3 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index 37686d8..ef16175 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -458,6 +458,16 @@ impl PartitionLayout {
 			trace!("parted -s {disk:?} type {i} {part_type_uuid}");
 			cmd_lib::run_cmd!(parted -s $disk type $i $part_type_uuid 2>&1)?;
 
+			if let Some(flags) = &part.flags {
+				debug!("Setting partition attribute flags");
+
+				for flag in flags {
+					let position = flag.flag_position();
+					trace!("parted -s {disk:?} toggle {i} {position}");
+					cmd_lib::run_cmd!(parted -s $disk toggle $i $position 2>&1)?;
+				}
+			}
+
 			if part.filesystem == "efi" {
 				debug!("Setting esp on for efi partition");
 				trace!("parted -s {disk:?} set {i} esp on");
@@ -514,6 +524,7 @@ fn test_partlay() {
 	partlay.add_partition(Partition {
 		label: Some("EFI".to_string()),
 		partition_type: PartitionType::Esp,
+		flags: None,
 		size: Some(ByteSize::mib(100)),
 		filesystem: "efi".to_string(),
 		mountpoint: "/boot/efi".to_string(),
@@ -523,6 +534,7 @@ fn test_partlay() {
 	partlay.add_partition(Partition {
 		label: Some("boot".to_string()),
 		partition_type: PartitionType::Xbootldr,
+		flags: None,
 		size: Some(ByteSize::gib(100)),
 		filesystem: "ext4".to_string(),
 		mountpoint: "/boot".to_string(),
@@ -532,6 +544,7 @@ fn test_partlay() {
 	partlay.add_partition(Partition {
 		label: Some("ROOT".to_string()),
 		partition_type: PartitionType::Root,
+		flags: None,
 		size: Some(ByteSize::gib(100)),
 		filesystem: "ext4".to_string(),
 		mountpoint: "/".to_string(),
@@ -569,6 +582,7 @@ fn test_partlay() {
 			Partition {
 				label: Some("ROOT".to_string()),
 				partition_type: PartitionType::Root,
+				flags: None,
 				size: Some(ByteSize::gib(100)),
 				filesystem: "ext4".to_string(),
 				mountpoint: "/".to_string(),
@@ -580,6 +594,7 @@ fn test_partlay() {
 			Partition {
 				label: Some("boot".to_string()),
 				partition_type: PartitionType::Xbootldr,
+				flags: None,
 				size: Some(ByteSize::gib(100)),
 				filesystem: "ext4".to_string(),
 				mountpoint: "/boot".to_string(),
@@ -591,6 +606,7 @@ fn test_partlay() {
 			Partition {
 				label: Some("EFI".to_string()),
 				partition_type: PartitionType::Esp,
+				flags: None,
 				size: Some(ByteSize::mib(100)),
 				filesystem: "efi".to_string(),
 				mountpoint: "/boot/efi".to_string(),
@@ -634,7 +650,7 @@ pub enum PartitionType {
 }
 
 impl PartitionType {
-	/// Get the GPT parition type GUID
+	/// Get the GPT partition type GUID
 	fn uuid(&self, target_arch: &str) -> String {
 		// https://uapi-group.org/specifications/specs/discoverable_partitions_specification/#partition-names
 		match self {
@@ -657,13 +673,45 @@ impl PartitionType {
 	}
 }
 
+/// Represents GPT partition attrbite flags which can be used, from https://uapi-group.org/specifications/specs/discoverable_partitions_specification/#partition-attribute-flags.
+#[derive(Deserialize, Debug, Clone, Serialize, PartialEq, Eq)]
+#[serde(rename_all = "kebab-case")]
+pub enum PartitionFlag {
+	/// Disable auto discovery for the partition, preventing automatic mounting
+	NoAuto,
+	/// Mark partition for mounting as read-only
+	ReadOnly,
+	/// Enable automatically growing the underlying file system when mounted
+	GrowFs,
+	/// An arbitrary GPT attribute flag position, 0 - 63
+	#[serde(untagged)]
+	FlagPosition(u8),
+}
+
+impl PartitionFlag {
+	/// Get the position offset for this flag
+	fn flag_position(&self) -> u8 {
+		// https://uapi-group.org/specifications/specs/discoverable_partitions_specification/#partition-attribute-flags
+		match &self {
+			PartitionFlag::NoAuto => 63,
+			PartitionFlag::ReadOnly => 60,
+			PartitionFlag::GrowFs => 59,
+			PartitionFlag::FlagPosition(position @ 0..=63) => *position,
+			_ => unimplemented!(),
+		}
+	}
+}
+
 #[derive(Deserialize, Debug, Clone, Serialize, PartialEq, Eq)]
 pub struct Partition {
 	pub label: Option<String>,
-	// Partition type
+	/// Partition type
 	#[serde(rename = "type")]
 	pub partition_type: PartitionType,
-	// If not specified, the partition will be created at the end of the disk (100%)
+	/// GPT partition attribute flags to add
+	// todo: maybe represent this as a bitflag number, parted consumes the positions so I'm doing this for now
+	pub flags: Option<Vec<PartitionFlag>>,
+	/// If not specified, the partition will be created at the end of the disk (100%)
 	pub size: Option<ByteSize>,
 	/// Filesystem of the partition
 	pub filesystem: String,