chore: One more pass at guest_disk simplification

rs/ic_os/os_tools/guest_disk/src/crypt.rs

@@ -26,14 +26,22 @@ const PBKDF_ITERATIONS: u32 = 1000;
 pub const LUKS2_N_KEYSLOTS: u32 = 32;
 /// Number of tokens supported by LUKS2
 pub const LUKS2_N_TOKENS: u32 = 32;
+/// LUKS2 token type identifier for our key slot metadata.
 pub const IC_KEY_TOKEN_TYPE: &str = "ic-key-metadata";
 
+/// All encrypted partitions contain a single active keyslot with maximum 1 active token.
+/// (SEV-based derivation has 1 token, generated keys have no associated tokens)
+pub const SINGLE_KEYSLOT_INDEX: u32 = 0;
+pub const SINGLE_TOKEN_INDEX: u32 = 0;
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct KeyslotMetadata {
-    // token_type and keyslots are standard fields in LUKS tokens
+pub struct KeyslotToken {
+    /// LUKS2 token type — must be set to [`IC_KEY_TOKEN_TYPE`].
     #[serde(rename = "type")]
     token_type: String,
+    /// Key slots this token is associated with (LUKS2 requires this field).
     pub keyslots: Vec<String>,
+    /// Metadata used to derive the key this token refers to.
     pub sev_metadata: SevMetadata,
     // Note: this type is serialized and stored on disk. When adding a new field, make sure to
     // set the type to Optional or mark it with #[serde(default)].
@@ -47,26 +55,14 @@ pub struct SevMetadata {
     // set the type to Optional or mark it with #[serde(default)].
 }
 
-impl KeyslotMetadata {
-    pub fn new_sev(keyslot: u32, sev_metadata: SevMetadata) -> Self {
+impl KeyslotToken {
+    pub fn new_sev(sev_metadata: SevMetadata) -> Self {
         Self {
             token_type: IC_KEY_TOKEN_TYPE.to_string(),
-            keyslots: vec![keyslot.to_string()],
+            keyslots: vec![SINGLE_KEYSLOT_INDEX.to_string()],
             sev_metadata,
         }
     }
-
-    pub fn keyslot(&self) -> Result<u32> {
-        if self.keyslots.len() != 1 {
-            bail!(
-                "Token must have exactly one keyslot, but found {:?}",
-                self.keyslots
-            );
-        }
-        self.keyslots[0]
-            .parse()
-            .context("Keyslot is not a valid number")
-    }
 }
 
 #[derive(Debug)]
@@ -135,9 +131,8 @@ fn obtain_crypt_device_handle_with_detached_header(
         .context("Failed to initialize cryptographic device with detached header")
 }
 
-/// Initializes a cryptographic device at the specified path with LUKS2 format and activates it
-/// using the provided name and encryption key.
-/// Returns the crypt device and the activated keyslot.
+/// Activates the cryptographic device at the specified path under the given name,
+/// using the provided encryption key.
 pub fn activate_crypt_device(
     device_path: &Path,
     header_location: LuksHeaderLocation,
@@ -146,7 +141,7 @@ pub fn activate_crypt_device(
     flags: CryptActivate,
     verify_luks_params: bool,
     metrics_registry: Option<&Registry>,
-) -> Result<(CryptDevice, u32)> {
+) -> Result<()> {
     let mut crypt_device = open_luks2_device(device_path, header_location)?;
 
     let luks_parameters = extract_luks_parameters(&mut crypt_device);
@@ -166,7 +161,7 @@ pub fn activate_crypt_device(
         }
     }
 
-    Ok((crypt_device, active_keyslot))
+    Ok(())
 }
 
 /// Deactivates the cryptographic device with the given name.
@@ -208,7 +203,7 @@ pub fn format_crypt_device(
     device_path: &Path,
     header_location: LuksHeaderLocation,
     encryption_key: &[u8],
-) -> Result<(CryptDevice, u32)> {
+) -> Result<CryptDevice> {
     if let LuksHeaderLocation::Detached(header_path) = header_location {
         File::create(header_path)
             .context("Failed to create detached LUKS header file")?
@@ -236,12 +231,17 @@ pub fn format_crypt_device(
             None,
         )
         .context("Failed to call format")?;
-    let keyslot = crypt_device
+    crypt_device
         .keyslot_handle()
-        .add_by_key(None, None, encryption_key, CryptVolumeKey::empty())
+        .add_by_key(
+            Some(SINGLE_KEYSLOT_INDEX),
+            None,
+            encryption_key,
+            CryptVolumeKey::empty(),
+        )
         .context("Could not add key to cryptographic device")?;
 
-    Ok((crypt_device, keyslot))
+    Ok(crypt_device)
 }
 
 /// Opens a LUKS2 device at the specified path, loads its context, and prepares handle-local
@@ -470,62 +470,41 @@ pub(crate) fn destroy_keyslots_except_first(crypt_device: &mut CryptDevice) -> R
 
 fn remove_all_tokens(crypt_device: &mut CryptDevice) -> Result<()> {
     for token_id in 0..LUKS2_N_TOKENS {
-        crypt_device
-            .token_handle()
-            .json_set(TokenInput::RemoveToken(token_id))
-            .with_context(|| format!("Failed to remove IC key metadata token {token_id}"))?;
+        if !matches!(
+            crypt_device.token_handle().status(token_id),
+            Ok(CryptTokenInfo::Inactive)
+        ) {
+            crypt_device
+                .token_handle()
+                .json_set(TokenInput::RemoveToken(token_id))
+                .with_context(|| format!("Failed to remove IC key metadata token {token_id}"))?;
+        }
     }
 
     Ok(())
 }
 
-fn ic_key_token_ids(crypt_device: &mut CryptDevice) -> Vec<u32> {
-    (0..LUKS2_N_TOKENS)
-        .filter(|&token_id| {
-            matches!(
-                crypt_device.token_handle().status(token_id),
-                Ok(CryptTokenInfo::External(ref token_type)
-                    | CryptTokenInfo::ExternalUnknown(ref token_type))
-                    if token_type == IC_KEY_TOKEN_TYPE
-            )
-        })
-        .collect()
-}
-
-pub fn read_keyslot_metadata(crypt_device: &mut CryptDevice) -> Result<Vec<KeyslotMetadata>> {
-    ic_key_token_ids(crypt_device)
-        .into_iter()
-        .map(|token_id| {
-            let json = crypt_device
-                .token_handle()
-                .json_get(token_id)
-                .with_context(|| format!("Failed to read IC key metadata token {token_id}"))?;
-            let metadata = serde_json::from_value::<KeyslotMetadata>(json)
-                .with_context(|| format!("Failed to parse IC key metadata token {token_id}"))?;
-            metadata
-                .keyslot()
-                .with_context(|| format!("Invalid keyslot in IC key metadata token {token_id}"))?;
-            Ok(metadata)
-        })
-        .collect()
+pub fn read_single_keyslot_token(crypt_device: &mut CryptDevice) -> Result<KeyslotToken> {
+    let json = crypt_device
+        .token_handle()
+        .json_get(SINGLE_TOKEN_INDEX)
+        .with_context(|| format!("Failed to read IC key metadata token {SINGLE_TOKEN_INDEX}"))?;
+    serde_json::from_value::<KeyslotToken>(json)
+        .with_context(|| format!("Failed to parse IC key metadata token {SINGLE_TOKEN_INDEX}"))
 }
 
-pub fn add_sev_metadata(
-    crypt_device: &mut CryptDevice,
-    keyslot: u32,
-    sev_metadata: SevMetadata,
-) -> Result<()> {
+pub fn add_sev_metadata(crypt_device: &mut CryptDevice, sev_metadata: SevMetadata) -> Result<()> {
     // TODO: Legacy headers may carry more than one IC key metadata token. Once all nodes
     // have been updated (i.e., the num_tokens metric is 1 everywhere), this removal can
     // be deleted so that only one token is written.
     remove_all_tokens(crypt_device)?;
 
-    let json = serde_json::to_value(KeyslotMetadata::new_sev(keyslot, sev_metadata))
+    let json = serde_json::to_value(KeyslotToken::new_sev(sev_metadata))
         .context("Failed to serialize key slot metadata")?;
 
     crypt_device
         .token_handle()
-        .json_set(TokenInput::ReplaceToken(0, &json))
+        .json_set(TokenInput::ReplaceToken(SINGLE_TOKEN_INDEX, &json))
         .context("Failed to write LUKS2 token")?;
 
     Ok(())

rs/ic_os/os_tools/guest_disk/src/generated_key.rs

@@ -14,7 +14,7 @@ const GENERATED_KEY_SIZE_BYTES: usize = 16;
 
 pub struct GeneratedKeyDiskEncryption<'a> {
     pub key_path: &'a Path,
-    pub metrics_registry: &'a Registry,
+    pub metrics_registry: Registry,
 }
 
 impl DiskEncryption for GeneratedKeyDiskEncryption<'_> {
@@ -32,7 +32,7 @@ impl DiskEncryption for GeneratedKeyDiskEncryption<'_> {
             // execution environment)
             /*verify_luks_params=*/
             false,
-            Some(self.metrics_registry),
+            Some(&self.metrics_registry),
         )
         .context("Failed to initialize crypt device")?;
 

rs/ic_os/os_tools/guest_disk/src/main.rs

@@ -88,7 +88,7 @@ fn run(
     } else {
         Box::new(GeneratedKeyDiskEncryption {
             key_path: generated_key_path,
-            metrics_registry: &metrics_registry,
+            metrics_registry: metrics_registry.clone(),
         })
     };
     let partition = args.partition();

rs/ic_os/os_tools/guest_disk/src/sev.rs

@@ -1,6 +1,6 @@
 use crate::crypt::{
-    LuksHeaderLocation, SevMetadata, activate_crypt_device, add_sev_metadata, check_encryption_key,
-    destroy_keyslots_except_first, format_crypt_device, open_luks2_device,
+    LuksHeaderLocation, SINGLE_KEYSLOT_INDEX, SevMetadata, activate_crypt_device, add_sev_metadata,
+    check_encryption_key, destroy_keyslots_except_first, format_crypt_device, open_luks2_device,
 };
 use crate::{DiskEncryption, Partition, activate_flags};
 use anyhow::{Context, Result, bail};
@@ -22,6 +22,15 @@ pub struct SevDiskEncryption {
     pub metrics_registry: Registry,
 }
 
+impl SevDiskEncryption {
+    fn header_location(&self, partition: Partition) -> LuksHeaderLocation<'_> {
+        match partition {
+            Partition::Store => LuksHeaderLocation::Detached(&self.store_luks_header_path),
+            Partition::Var => LuksHeaderLocation::Attached,
+        }
+    }
+}
+
 impl DiskEncryption for SevDiskEncryption {
     fn open(&mut self, device_path: &Path, partition: Partition, crypt_name: &str) -> Result<()> {
         let key = derive_key_from_sev_measurement(
@@ -30,33 +39,16 @@ impl DiskEncryption for SevDiskEncryption {
         )
         .context("Failed to derive SEV key for disk encryption")?;
 
-        match partition {
-            Partition::Var => {
-                activate_crypt_device(
-                    device_path,
-                    LuksHeaderLocation::Attached,
-                    crypt_name,
-                    key.as_bytes(),
-                    activate_flags(partition),
-                    /*verify_luks_params=*/ true,
-                    Some(&self.metrics_registry),
-                )
-                .context("Failed to open crypt device for var partition")?;
-            }
-
-            Partition::Store => {
-                activate_crypt_device(
-                    device_path,
-                    LuksHeaderLocation::Detached(&self.store_luks_header_path),
-                    crypt_name,
-                    key.as_bytes(),
-                    activate_flags(partition),
-                    /*verify_luks_params=*/ true,
-                    Some(&self.metrics_registry),
-                )
-                .context("Failed to initialize crypt device for store partition")?;
-            }
-        }
+        activate_crypt_device(
+            device_path,
+            self.header_location(partition),
+            crypt_name,
+            key.as_bytes(),
+            activate_flags(partition),
+            /*verify_luks_params=*/ true,
+            Some(&self.metrics_registry),
+        )
+        .with_context(|| format!("Failed to open the {partition:?} partition"))?;
 
         Ok(())
     }
@@ -70,24 +62,18 @@ impl DiskEncryption for SevDiskEncryption {
 
         let sev_metadata = get_sev_metadata_for_luks(self.sev_firmware.as_mut())?;
 
-        let header_location = match partition {
-            Partition::Store => {
-                if self.store_luks_header_path.exists() {
-                    bail!(
-                        "Refusing to format Store because detached LUKS header {} already exists. \
-                        Remove the stale header first if you really want to reformat the device.",
-                        self.store_luks_header_path.display()
-                    );
-                }
-                LuksHeaderLocation::Detached(&self.store_luks_header_path)
-            }
-            Partition::Var => LuksHeaderLocation::Attached,
-        };
+        if partition == Partition::Store && self.store_luks_header_path.exists() {
+            bail!(
+                "Refusing to format Store because detached LUKS header {} already exists. \
+                Remove the stale header first if you really want to reformat the device.",
+                self.store_luks_header_path.display()
+            );
+        }
 
-        let (mut crypt_device, keyslot) =
-            format_crypt_device(device_path, header_location, key.as_bytes())
+        let mut crypt_device =
+            format_crypt_device(device_path, self.header_location(partition), key.as_bytes())
                 .context("Failed to format partition")?;
-        add_sev_metadata(&mut crypt_device, keyslot, sev_metadata)
+        add_sev_metadata(&mut crypt_device, sev_metadata)
             .context("Failed to write SEV keyslot metadata")?;
 
         Ok(())
@@ -146,12 +132,19 @@ pub fn rekey(
     // new key. Fails if the old key does not unlock any keyslot.
     crypt_device
         .keyslot_handle()
-        .change_by_passphrase(None, Some(0), old_key, new_key.as_bytes())
+        .change_by_passphrase(
+            // TODO: after all nodes have a single keyslot at SINGLE_KEYSLOT_INDEX, change this to
+            //  Some(SINGLE_KEYSLOT_INDEX)
+            None,
+            Some(SINGLE_KEYSLOT_INDEX),
+            old_key,
+            new_key.as_bytes(),
+        )
         .context("Failed to replace the received key with the new SEV-derived key")?;
     // Removes the keyslots that legacy headers may still carry.
     // TODO: remove it (see comment on destroy_keyslots_except_first).
     destroy_keyslots_except_first(&mut crypt_device)?;
-    add_sev_metadata(&mut crypt_device, 0, sev_metadata)
+    add_sev_metadata(&mut crypt_device, sev_metadata)
         .context("Failed to write SEV keyslot metadata")?;
 
     Ok(())