pv: Refactor 'get_verified_hkds'

Get rid of 'read_hkd' by refactoring the 'get_verified_hkds' function.
For this a new HkdLoader::load_and_verify is introduced that is a
reworked version of the original code.

In addition, add test cases for testing all the edge cases.

Assisted-by: IBM Bob:1.0.6
Signed-off-by: Marc Hartmayer <marc@linux.ibm.com>
Reviewed-by: Steffen Eiden <seiden@linux.ibm.com>
Signed-off-by: Steffen Eiden <seiden@linux.ibm.com>

rust/Cargo.lock

@@ -1222,6 +1222,7 @@ dependencies = [
  "clap",
  "libc",
  "log",
+ "openssl",
  "s390_pv",
  "serde",
  "serde_json",

rust/pv/src/lib.rs

@@ -79,7 +79,7 @@ pub mod pem {
 pub mod misc {
     pub use pv_core::misc::*;
 
-    pub use crate::utils::{read_certs, read_hkd};
+    pub use crate::utils::read_certs;
 }
 
 pub use error::{Error, Result};

rust/pv/src/utils.rs

@@ -1,14 +1,9 @@
-use std::path::Path;
-
 // SPDX-License-Identifier: MIT
 //
 // Copyright IBM Corp. 2023
-use log::{error, info};
 use openssl::error::ErrorStack;
 use openssl::x509::{X509Crl, X509};
-use pv_core::misc::read_file;
 
-use crate::req::{HostKey, HybridPKey};
 use crate::{Error, Result};
 
 /// Read all CRLs from the buffer and parse them into a vector.
@@ -36,49 +31,6 @@ pub fn read_certs<T: AsRef<[u8]>>(buf: T) -> Result<Vec<X509>, ErrorStack> {
         .or_else(|_| X509::stack_from_pem(buf.as_ref()))
 }
 
-/// Read a host-key document from a file.
-///
-/// # Errors
-///
-/// This function will return an error if:
-/// - The file cannot be read
-/// - The content is not valid PEM or DER format
-/// - The file contains no certificates or more than 2 certificates
-/// - The public key cannot be extracted from the certificate(s)
-pub fn read_hkd<P: AsRef<Path>>(path: P) -> Result<HostKey> {
-    let path = path.as_ref();
-    let hk = read_file(path, "host-key document")?;
-    let certs = read_certs(&hk).map_err(|source| Error::HkdNotPemOrDer {
-        hkd: path.display().to_string(),
-        source,
-    })?;
-    if certs.is_empty() {
-        return Err(Error::NoHkdInFile(path.display().to_string()));
-    }
-    let c1 = certs.first().unwrap();
-    match certs.len() {
-        1 => {
-            info!("Using version 1 of the host-key document format");
-            Ok(HostKey::V1(c1.public_key()?))
-        }
-        2 => {
-            info!("Using version 2 of the host-key document format");
-            let c2 = &certs[1];
-            Ok(HostKey::V2(HybridPKey::new(
-                c1.public_key()?,
-                c2.public_key()?,
-            )?))
-        }
-        _ => {
-            error!(
-                "Invalid host-key document '{}': it contains more than two certificates, which is not supported by any host-key document format.",
-                    path.display()
-            );
-            Err(Error::WrongNumberOfKeys(path.display().to_string()))
-        }
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use crate::test_utils::*;

rust/pvimg/examples/create-sehdr/main.rs

@@ -12,13 +12,16 @@ use std::str::FromStr;
 use anyhow::{anyhow, Context, Error};
 use clap::{Parser, ValueEnum, ValueHint};
 use log::info;
-use pv::misc::{decode_hex, open_file, read_file, read_hkd, try_parse_u64};
-use pv::request::{HostKey, SymKeyType};
+use pv::misc::{decode_hex, open_file, read_file, try_parse_u64};
+use pv::request::{HostKey, NoVerifyHkd, SymKeyType};
 use pv::Result;
 use pvimg::misc::PSW;
 use pvimg::secured_comp::{ComponentTrait, Layout, SecuredComponentBuilder};
 use pvimg::uvdata::{BuilderTrait, SeHdrBuilder, SeHdrControlFlags, SeHdrVersion, SeTarget};
-use utils::{AtomicFile, AtomicFileOperation, HexSlice, PvLogger, VerbosityOptions};
+use utils::{
+    AtomicFile, AtomicFileOperation, HexSlice, HkdLoader, HkdVersionSelection, PvLogger,
+    VerbosityOptions,
+};
 
 /// Converts the hexstring into a byte vector.
 ///
@@ -245,7 +248,7 @@ fn main() -> anyhow::Result<()> {
             "Use the file '{}' as a host key document",
             hkd_path.display()
         );
-        let cert = read_hkd(&hkd_path)?;
+        let cert = HkdLoader::load_and_verify(&hkd_path, &NoVerifyHkd, HkdVersionSelection::Auto)?;
         target_pub_keys.push(cert);
     }
     let version: SeHdrVersion = args

rust/pvimg/src/cmd/test.rs

@@ -6,10 +6,12 @@ use std::path::Path;
 
 use anyhow::Result;
 use log::{info, warn};
-use pv::misc::{open_file, read_hkd};
+use pv::misc::open_file;
+use pv::request::NoVerifyHkd;
 use pv::{FileAccessErrorType, PvCoreError};
 use pvimg::error::{Error, OwnExitCode};
 use pvimg::uvdata::{KeyExchangeTrait, SeHdr, UvKeyHashesV1};
+use utils::hkd::{HkdLoader, HkdVersionSelection};
 use utils::HexSlice;
 
 use crate::cli::TestArgs;
@@ -72,7 +74,7 @@ where
 
     let mut result = false;
     for path in host_key_documents {
-        let hkd = read_hkd(path)?;
+        let hkd = HkdLoader::load_and_verify(path, &NoVerifyHkd, HkdVersionSelection::Auto)?;
         if hdr.contains(hkd)? {
             result = true;
             log_println!(

rust/utils/Cargo.toml

@@ -9,6 +9,7 @@ chrono = { version = "0.4.44", default-features = false, features = ["std"] }
 clap = { version ="4.6", features = ["derive", "wrap_help"] }
 libc = "0.2.186"
 log = { version = "0.4.29", features = ["std", "release_max_level_debug"] }
+openssl = "0.10"
 pv = { path = "../pv", package = "s390_pv" }
 serde = { version = "1.0.228"}
 

rust/utils/src/cli.rs

@@ -10,13 +10,12 @@ use std::str::FromStr;
 
 use clap::builder::{EnumValueParser, PossibleValue, TypedValueParser};
 use clap::{Arg, ArgAction, ArgGroup, Args, Command, ValueEnum, ValueHint};
-use log::{info, warn, LevelFilter};
-use openssl::Nid;
-use pv::misc::{create_file, open_file, read_certs, read_file};
-use pv::request::openssl::pkey::{KeyType, PKeyRef, Public};
-use pv::request::{openssl, HkdVerifier, HostKey, HybridPKey};
-use pv::{Error, Result};
-use utils_macros::{ValueEnumDisplay, ValueEnumFromStr};
+use log::LevelFilter;
+use pv::misc::{create_file, open_file};
+use pv::request::{HkdVerifier, HostKey};
+use pv::Result;
+
+use crate::hkd::{HkdLoader, HkdVersionSelection};
 
 /// Generic version selection for CLI
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -133,17 +132,6 @@ where
     }
 }
 
-/// Host key document version for CLI
-#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum, ValueEnumDisplay, ValueEnumFromStr)]
-pub enum HkdVersion {
-    /// Version 1 - uses traditional cryptographic keys
-    Classical,
-    /// Version 2 - uses hybrid (post-quantum) cryptographic keys
-    Hybrid,
-}
-
-pub type HkdVersionSelection = AutoOrExplicit<HkdVersion>;
-
 /// CLI Argument collection for handling host-keys, IBM signing keys, and certificates.
 #[derive(Args, Debug, Clone, PartialEq, Eq, Default)]
 #[command(
@@ -239,21 +227,6 @@ impl CertificateOptions {
         }
     }
 
-    fn is_ec_p521_key(key: &PKeyRef<Public>) -> bool {
-        if key.id() == openssl::pkey::Id::EC {
-            let ec_key = key.ec_key().unwrap();
-            let group = ec_key.group();
-            let curve_nid = group.curve_name().unwrap();
-            curve_nid == Nid::SECP521R1
-        } else {
-            false
-        }
-    }
-
-    fn is_mlkem1024_key(key: &PKeyRef<Public>) -> bool {
-        key.is_a(KeyType::ML_KEM_1024)
-    }
-
     /// Read the hybrid host-keys specified and verifies them if required
     ///
     /// - `protectee`: what you want to create. e.g. add-secret request or SE-image
@@ -273,67 +246,8 @@ impl CertificateOptions {
 
         let mut res = Vec::with_capacity(hkds.len());
         for hkd in hkds {
-            let hk = read_file(hkd, "host-key document")?;
-            let certs = read_certs(&hk).map_err(|source| Error::HkdNotPemOrDer {
-                hkd: hkd.display().to_string(),
-                source,
-            })?;
-            if certs.is_empty() {
-                return Err(Error::NoHkdInFile(hkd.display().to_string()));
-            }
-            let required_cert_count = match requested_version {
-                HkdVersionSelection::Auto => {
-                    info!("Auto-detecting version of the host-key document format");
-                    match certs.len() {
-                        1 => 1,
-                        2 => 2,
-                        _ => {
-                            warn!(
-                                "The host-key document in '{}' contains more than two certificates!",
-                                hkd.display()
-                            );
-                            return Err(Error::WrongNumberOfKeys(hkd.display().to_string()));
-                        }
-                    }
-                }
-                HkdVersionSelection::Explicit(HkdVersion::Classical) => {
-                    info!("Using version 1 of the host-key document format");
-                    1
-                }
-                HkdVersionSelection::Explicit(HkdVersion::Hybrid) => {
-                    info!("Using version 2 of the host-key document format");
-                    2
-                }
-            };
-            if required_cert_count != certs.len() {
-                return Err(Error::WrongNumberOfKeys(hkd.display().to_string()));
-            }
-            let c1 = certs.first().unwrap();
-            if !Self::is_ec_p521_key(c1.public_key()?.as_ref()) {
-                return Err(Error::InvalidHkd(
-                    "First key must be a EC-p521 key".to_string(),
-                ));
-            }
-            verifier.verify(c1)?;
-            match certs.len() {
-                1 => {
-                    res.push(HostKey::V1(c1.public_key()?));
-                }
-                2 => {
-                    let c2 = &certs[1];
-                    if !Self::is_mlkem1024_key(c2.public_key()?.as_ref()) {
-                        return Err(Error::InvalidHkd(
-                            "Second key must be a ML-KEM 1024 key".to_string(),
-                        ));
-                    }
-                    verifier.verify(c2)?;
-                    res.push(HostKey::V2(HybridPKey::new(
-                        c1.public_key()?,
-                        c2.public_key()?,
-                    )?))
-                }
-                _ => unreachable!("Already checked"),
-            }
+            let host_key = HkdLoader::load_and_verify(hkd, verifier.as_ref(), requested_version)?;
+            res.push(host_key);
         }
         Ok(res)
     }

rust/utils/src/hkd.rs

@@ -0,0 +1,162 @@
+// SPDX-License-Identifier: MIT
+//
+// Copyright IBM Corp.
+
+use std::fmt::{Display, Formatter};
+use std::path::Path;
+
+use log::{error, info};
+use openssl::nid::Nid;
+use openssl::pkey::{Id, KeyType, PKeyRef, Public};
+use openssl::x509::X509;
+use pv::misc::{read_certs, read_file};
+use pv::request::{HkdVerifier, HostKey, HybridPKey};
+use pv::{Error, Result};
+
+use crate::AutoOrExplicit;
+
+/// Host key document version
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum HkdVersion {
+    /// Version 1 - uses traditional cryptographic keys (1 certificate)
+    Classical,
+    /// Version 2 - uses hybrid (post-quantum) cryptographic keys (2 certificates)
+    Hybrid,
+}
+
+impl HkdVersion {
+    /// Get the required certificate count for this version
+    pub fn cert_count(self) -> usize {
+        match self {
+            HkdVersion::Classical => 1,
+            HkdVersion::Hybrid => 2,
+        }
+    }
+}
+
+impl Display for HkdVersion {
+    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
+        match self {
+            HkdVersion::Classical => write!(f, "classical"),
+            HkdVersion::Hybrid => write!(f, "hybrid"),
+        }
+    }
+}
+
+pub type HkdVersionSelection = AutoOrExplicit<HkdVersion>;
+
+/// Helper struct for loading and verifying host-key documents
+pub struct HkdLoader;
+
+impl HkdLoader {
+    fn detect_version(path: &Path, certs: &Vec<X509>) -> Result<HkdVersion> {
+        info!("Auto-detecting version of the host-key document format");
+        Ok(match certs.len() {
+            1 => HkdVersion::Classical,
+            2 => HkdVersion::Hybrid,
+            _ => {
+                error!(
+                "Invalid host-key document '{}': it contains more than two certificates, which is not supported by any host-key document format.",
+                path.display()
+            );
+                return Err(Error::WrongNumberOfKeys(path.display().to_string()));
+            }
+        })
+    }
+
+    fn validate_version(path: &Path, certs: &Vec<X509>, version: HkdVersion) -> Result<HkdVersion> {
+        if certs.len() != version.cert_count() {
+            error!(
+                "Host-key document '{}' is not a {} host-key document.",
+                path.display(),
+                version,
+            );
+            return Err(Error::WrongNumberOfKeys(path.display().to_string()));
+        }
+        Ok(version)
+    }
+
+    fn is_ec_p521_key(key: &PKeyRef<Public>) -> bool {
+        if key.id() == Id::EC {
+            if let Ok(ec_key) = key.ec_key() {
+                let group = ec_key.group();
+                if let Some(curve_nid) = group.curve_name() {
+                    return curve_nid == Nid::SECP521R1;
+                }
+            }
+        }
+        false
+    }
+
+    fn is_mlkem1024_key(key: &PKeyRef<Public>) -> bool {
+        key.is_a(KeyType::ML_KEM_1024)
+    }
+
+    /// Load and verify a host-key document from a file
+    ///
+    /// # Errors
+    ///
+    /// This function will return an error if:
+    /// - The file cannot be read
+    /// - The content is not valid PEM or DER format
+    /// - The file contains no certificates or wrong number of certificates
+    /// - The public key cannot be extracted from the certificate(s)
+    /// - The verification fails
+    /// - Key types are invalid
+    pub fn load_and_verify<P: AsRef<Path>>(
+        path: P,
+        verifier: &dyn HkdVerifier,
+        requested_version: HkdVersionSelection,
+    ) -> Result<HostKey> {
+        let path = path.as_ref();
+        let hk = read_file(path, "host-key document")?;
+        let certs = read_certs(&hk).map_err(|source| Error::HkdNotPemOrDer {
+            hkd: path.display().to_string(),
+            source,
+        })?;
+
+        if certs.is_empty() {
+            return Err(Error::NoHkdInFile(path.display().to_string()));
+        }
+
+        let version = match requested_version {
+            HkdVersionSelection::Auto => Self::detect_version(path, &certs)?,
+            HkdVersionSelection::Explicit(version) => {
+                Self::validate_version(path, &certs, version)?
+            }
+        };
+
+        info!("Using {version} host-key document format");
+
+        // SAFETY: certs is guaranteed to be non-empty due to the check
+        let c1 = certs
+            .first()
+            .expect("Certificate list validated as non-empty");
+
+        if !Self::is_ec_p521_key(c1.public_key()?.as_ref()) {
+            return Err(Error::InvalidHkd(
+                "First key must be a EC-p521 key".to_string(),
+            ));
+        }
+        verifier.verify(c1)?;
+
+        match version {
+            HkdVersion::Classical => Ok(HostKey::V1(c1.public_key()?)),
+            HkdVersion::Hybrid => {
+                let c2 = &certs
+                    .get(1)
+                    .expect("Certificate list length was already checked");
+                if !Self::is_mlkem1024_key(c2.public_key()?.as_ref()) {
+                    return Err(Error::InvalidHkd(
+                        "Second key must be a ML-KEM 1024 key".to_string(),
+                    ));
+                }
+                verifier.verify(c2)?;
+                Ok(HostKey::V2(HybridPKey::new(
+                    c1.public_key()?,
+                    c2.public_key()?,
+                )?))
+            }
+        }
+    }
+}

rust/utils/src/lib.rs

@@ -7,6 +7,7 @@ mod cli;
 mod exit_code;
 mod file;
 mod hexslice;
+pub mod hkd;
 mod hostname;
 mod json;
 mod log;
@@ -19,12 +20,13 @@ pub use utils_macros::{ControlFlag, ValueEnumDisplay, ValueEnumFromStr};
 pub use crate::cli::{
     combined_path_opt, combined_path_req, get_reader_from_cli_file_arg,
     get_writer_from_cli_file_arg, print_cli_error, print_error, AutoOrExplicit,
-    AutoOrExplicitParser, CertificateOptions, DeprecatedVerbosityOptions, HkdVersion,
-    HkdVersionSelection, VerbosityOptions, STDIN, STDOUT,
+    AutoOrExplicitParser, CertificateOptions, DeprecatedVerbosityOptions, VerbosityOptions, STDIN,
+    STDOUT,
 };
 pub use crate::exit_code::{docstring, ExitCodeDoc, ExitCodeTrait, ExitCodeVariantDoc};
 pub use crate::file::{AtomicFile, AtomicFileOperation};
 pub use crate::hexslice::HexSlice;
+pub use crate::hkd::{HkdLoader, HkdVersion, HkdVersionSelection};
 pub use crate::hostname::gethostname;
 pub use crate::json::S390ToolsMetaData;
 pub use crate::log::PvLogger;

rust/utils/tests/malformed_hkd.rs

@@ -0,0 +1,259 @@
+// SPDX-License-Identifier: MIT
+//
+// Copyright IBM Corp.
+
+//! Integration tests for malformed host key document handling
+
+use std::fs::{self, File};
+use std::io::Write;
+use std::path::PathBuf;
+
+use pv::request::NoVerifyHkd;
+use utils::{AutoOrExplicit, HkdLoader, HkdVersion, TemporaryDirectory};
+
+/// Path to test certificate assets
+fn cert_asset_path(name: &str) -> PathBuf {
+    PathBuf::from(env!("CARGO_MANIFEST_DIR"))
+        .parent()
+        .unwrap()
+        .join("pv/tests/assets/cert")
+        .join(name)
+}
+
+#[test]
+fn test_empty_hkd_file() {
+    let temp_dir = TemporaryDirectory::new().unwrap();
+    let file_path = temp_dir.path().join("empty.hkd");
+    File::create(&file_path).unwrap();
+
+    let result = HkdLoader::load_and_verify(&file_path, &NoVerifyHkd, AutoOrExplicit::Auto);
+    assert!(result.is_err(), "Empty HKD file should be rejected");
+
+    let err = result.unwrap_err();
+    assert!(
+        matches!(err, pv::Error::NoHkdInFile(_)),
+        "Expected NoHkdInFile error, got: {:?}",
+        err
+    );
+}
+
+#[test]
+fn test_invalid_pem_format() {
+    // Create a file with invalid PEM content
+    let temp_dir = TemporaryDirectory::new().unwrap();
+    let file_path = temp_dir.path().join("invalid.hkd");
+    let mut temp_file = File::create(&file_path).unwrap();
+    writeln!(temp_file, "-----BEGIN CERTIFICATE-----").unwrap();
+    writeln!(temp_file, "INVALID_BASE64_CONTENT!!!").unwrap();
+    writeln!(temp_file, "-----END CERTIFICATE-----").unwrap();
+    temp_file.flush().unwrap();
+
+    // Test that invalid PEM is properly rejected with HkdNotPemOrDer error
+    let result = HkdLoader::load_and_verify(&file_path, &NoVerifyHkd, AutoOrExplicit::Auto);
+    assert!(result.is_err(), "Invalid PEM format should be rejected");
+
+    let err = result.unwrap_err();
+    assert!(
+        matches!(err, pv::Error::HkdNotPemOrDer { .. }),
+        "Expected HkdNotPemOrDer error, got: {:?}",
+        err
+    );
+}
+
+#[test]
+fn test_wrong_number_of_certificates_v1() {
+    // Create a file with three valid certificates (invalid count for v1, expects 1)
+    let temp_dir = TemporaryDirectory::new().unwrap();
+    let file_path = temp_dir.path().join("wrong_count_v1.hkd");
+    let mut temp_file = File::create(&file_path).unwrap();
+
+    let cert1 = fs::read_to_string(cert_asset_path("host.crt")).unwrap();
+    let cert2 = fs::read_to_string(cert_asset_path("ibm.crt")).unwrap();
+    let cert3 = fs::read_to_string(cert_asset_path("root_ca.crt")).unwrap();
+
+    write!(temp_file, "{}{}{}", cert1, cert2, cert3).unwrap();
+    temp_file.flush().unwrap();
+
+    let result = HkdLoader::load_and_verify(
+        &file_path,
+        &NoVerifyHkd,
+        AutoOrExplicit::Explicit(HkdVersion::Classical),
+    );
+    assert!(
+        result.is_err(),
+        "Wrong number of certificates for v1 should be rejected"
+    );
+
+    let err = result.unwrap_err();
+    assert!(
+        matches!(err, pv::Error::WrongNumberOfKeys(_)),
+        "Expected WrongNumberOfKeys error, got: {:?}",
+        err
+    );
+}
+
+#[test]
+fn test_wrong_number_of_certificates_v2() {
+    // Use existing single certificate file directly (invalid count for v2, expects 2)
+    let file_path = cert_asset_path("host.crt");
+
+    // Test that wrong number of certificates is properly rejected with WrongNumberOfKeys error
+    let result = HkdLoader::load_and_verify(
+        &file_path,
+        &NoVerifyHkd,
+        AutoOrExplicit::Explicit(HkdVersion::Hybrid),
+    );
+    assert!(
+        result.is_err(),
+        "Wrong number of certificates for v2 should be rejected"
+    );
+
+    let err = result.unwrap_err();
+    assert!(
+        matches!(err, pv::Error::WrongNumberOfKeys(_)),
+        "Expected WrongNumberOfKeys error, got: {:?}",
+        err
+    );
+}
+
+#[test]
+fn test_corrupted_certificate() {
+    // Create a file with corrupted certificate data
+    let temp_dir = TemporaryDirectory::new().unwrap();
+    let file_path = temp_dir.path().join("corrupted.hkd");
+    let mut temp_file = File::create(&file_path).unwrap();
+    writeln!(temp_file, "-----BEGIN CERTIFICATE-----").unwrap();
+    writeln!(
+        temp_file,
+        "MIICdTCCAd6gAwIBAgIBADANBgkqhkiG9w0BAQsFADBQMQswCQYDVQQGEwJVUzEL"
+    )
+    .unwrap();
+    writeln!(temp_file, "CORRUPTED_DATA_HERE").unwrap();
+    writeln!(temp_file, "-----END CERTIFICATE-----").unwrap();
+    temp_file.flush().unwrap();
+
+    let result = HkdLoader::load_and_verify(&file_path, &NoVerifyHkd, AutoOrExplicit::Auto);
+    assert!(result.is_err(), "Corrupted certificate should be rejected");
+
+    let err = result.unwrap_err();
+    assert!(
+        matches!(err, pv::Error::HkdNotPemOrDer { .. }),
+        "Expected HkdNotPemOrDer error, got: {:?}",
+        err
+    );
+}
+
+#[test]
+fn test_wrong_key_type_v1() {
+    // Use RSA certificate instead of EC-p521 for v1 (should fail)
+    let file_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
+        .parent()
+        .unwrap()
+        .join("pv/tests/assets/keys/rsa2048.crt");
+
+    let result = HkdLoader::load_and_verify(
+        &file_path,
+        &NoVerifyHkd,
+        AutoOrExplicit::Explicit(HkdVersion::Classical),
+    );
+    assert!(
+        result.is_err(),
+        "RSA key should be rejected for v1 (expects EC-p521)"
+    );
+
+    let err = result.unwrap_err();
+    if let pv::Error::InvalidHkd(msg) = err {
+        assert_eq!(
+            msg, "First key must be a EC-p521 key",
+            "Error message should indicate EC-p521 requirement"
+        );
+    } else {
+        panic!(
+            "Expected InvalidHkd error for wrong key type, got: {:?}",
+            err
+        );
+    }
+}
+
+#[test]
+fn test_wrong_key_type_v2() {
+    // Create a file with two RSA certificates instead of EC-p521 + ML-KEM for v2
+    let temp_dir = TemporaryDirectory::new().unwrap();
+    let file_path = temp_dir.path().join("wrong_key_v2.hkd");
+    let mut temp_file = File::create(&file_path).unwrap();
+
+    // Read two RSA certificates and concatenate them
+    let rsa_cert_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
+        .parent()
+        .unwrap()
+        .join("pv/tests/assets/keys/rsa2048.crt");
+    let rsa_cert = fs::read_to_string(&rsa_cert_path).unwrap();
+
+    write!(temp_file, "{}{}", rsa_cert, rsa_cert).unwrap();
+    temp_file.flush().unwrap();
+
+    // Test that wrong key types are properly rejected with InvalidHkd error
+    let result = HkdLoader::load_and_verify(
+        &file_path,
+        &NoVerifyHkd,
+        AutoOrExplicit::Explicit(HkdVersion::Hybrid),
+    );
+    assert!(
+        result.is_err(),
+        "RSA keys should be rejected for v2 (expects EC-p521 + ML-KEM)"
+    );
+
+    let err = result.unwrap_err();
+    if let pv::Error::InvalidHkd(msg) = err {
+        assert_eq!(
+            msg, "First key must be a EC-p521 key",
+            "Error message should indicate EC-p521 requirement"
+        );
+    } else {
+        panic!(
+            "Expected InvalidHkd error for wrong key types, got: {:?}",
+            err
+        );
+    }
+}
+
+#[test]
+fn test_wrong_second_key_type_v2() {
+    // Create a file with EC-p521 + EC (wrong) instead of EC-p521 + ML-KEM for v2
+    let temp_dir = TemporaryDirectory::new().unwrap();
+    let file_path = temp_dir.path().join("wrong_second_key_v2.hkd");
+    let mut temp_file = File::create(&file_path).unwrap();
+
+    let ec_p521_cert = fs::read_to_string(cert_asset_path("host.crt")).unwrap();
+    let ec_cert_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
+        .parent()
+        .unwrap()
+        .join("pv/tests/assets/keys/host.ec.crt");
+    let ec_cert = fs::read_to_string(&ec_cert_path).unwrap();
+
+    write!(temp_file, "{}{}", ec_p521_cert, ec_cert).unwrap();
+    temp_file.flush().unwrap();
+
+    let result = HkdLoader::load_and_verify(
+        &file_path,
+        &NoVerifyHkd,
+        AutoOrExplicit::Explicit(HkdVersion::Hybrid),
+    );
+    assert!(
+        result.is_err(),
+        "EC key should be rejected as second key for v2 (expects ML-KEM-1024)"
+    );
+
+    let err = result.unwrap_err();
+    if let pv::Error::InvalidHkd(msg) = err {
+        assert_eq!(
+            msg, "Second key must be a ML-KEM 1024 key",
+            "Error message should indicate ML-KEM-1024 requirement"
+        );
+    } else {
+        panic!(
+            "Expected InvalidHkd error for wrong second key type, got: {:?}",
+            err
+        );
+    }
+}