Extend public_key imports for KEMs

Author: beurdouche

bindings/bindings.toml

@@ -256,7 +256,6 @@ enums = [
 opaque = [
     "PK11ContextStr",
     "PK11SlotInfoStr",
-    "SECKEYPublicKeyStr",
 ]
 variables = [
     "AES_BLOCK_SIZE",

src/kem.rs

@@ -306,6 +306,60 @@ pub(crate) fn mlkem_decapsulate(
     Ok(shared_secret)
 }
 
+// ============================================================================
+// ML-KEM Public Key Import
+// ============================================================================
+
+/// Import an ML-KEM public key from raw bytes.
+///
+/// Constructs a `SECKEYPublicKey` with `keyType = kyberKey` and the
+/// appropriate `KyberParams`, then uses `SECKEY_CopyPublicKey` to create
+/// a properly arena-backed copy, and `PK11_ImportPublicKey` to register
+/// it in PKCS#11.
+pub fn import_mlkem_public_key(
+    raw_key: &[u8],
+    params: MlKemParameterSet,
+) -> Res<PublicKey> {
+    init()?;
+
+    let expected_size = params.public_key_bytes();
+    if raw_key.len() != expected_size {
+        return Err(Error::InvalidInput);
+    }
+
+    let kyber_params = match params {
+        MlKemParameterSet::MlKem768 => p11::KyberParams_params_ml_kem768,
+        MlKemParameterSet::MlKem1024 => p11::KyberParams_params_ml_kem1024,
+    };
+
+    unsafe {
+        let mut temp_key: p11::SECKEYPublicKey = std::mem::zeroed();
+        temp_key.keyType = p11::KeyType_kyberKey;
+        temp_key.u.kyber.as_mut().params = kyber_params;
+        temp_key.u.kyber.as_mut().publicValue = SECItem {
+            type_: crate::SECItemType::siBuffer,
+            data: raw_key.as_ptr() as *mut u8,
+            len: raw_key.len() as u32,
+        };
+
+        // SECKEY_CopyPublicKey allocates a new arena and deep-copies
+        let copied: PublicKey = p11::SECKEY_CopyPublicKey(&temp_key).into_result()?;
+
+        // Register in PKCS#11
+        let slot = Slot::internal()?;
+        let handle = p11::PK11_ImportPublicKey(
+            *slot,
+            *copied,
+            crate::PR_FALSE,
+        );
+        if handle == pkcs11_bindings::CK_INVALID_HANDLE {
+            return Err(Error::InvalidInput);
+        }
+
+        Ok(copied)
+    }
+}
+
 // ============================================================================
 // Public API
 // ============================================================================
@@ -616,6 +670,69 @@ mod tests {
     // Legacy API tests (for backwards compatibility)
     // ========================================================================
 
+    // ========================================================================
+    // Import tests
+    // ========================================================================
+
+    #[test]
+    fn test_mlkem768_import_public_key_roundtrip() {
+        use crate::err::secstatus_to_res;
+        use crate::util::SECItemMut;
+
+        fixture_init();
+
+        // Generate a keypair
+        let keypair = generate_keypair(KemParameterSet::MlKem768).unwrap();
+        let KemKeypair::MlKem768 {
+            ref public,
+            ref private,
+        } = keypair
+        else {
+            panic!("Expected MlKem768 keypair");
+        };
+
+        // Export public key bytes via PK11_ReadRawAttribute(CKA_VALUE)
+        let mut key_item = SECItemMut::make_empty();
+        secstatus_to_res(unsafe {
+            p11::PK11_ReadRawAttribute(
+                p11::PK11ObjectType::PK11_TypePubKey,
+                (**public).cast(),
+                pkcs11_bindings::CKA_VALUE,
+                key_item.as_mut(),
+            )
+        })
+        .unwrap();
+        let pk_bytes = key_item.as_slice().to_owned();
+        assert_eq!(pk_bytes.len(), MlKemParameterSet::MlKem768.public_key_bytes());
+
+        // Import from raw bytes
+        let imported_pk =
+            import_mlkem_public_key(&pk_bytes, MlKemParameterSet::MlKem768).unwrap();
+
+        // Encapsulate with imported public key
+        let target = p11::CKM_HKDF_DERIVE.into();
+        let (ss_key, ciphertext) = mlkem_encapsulate(&imported_pk, target).unwrap();
+        let shared_secret = ss_key.key_data().unwrap().to_vec();
+
+        // Decapsulate with original private key
+        let dec_key = mlkem_decapsulate(private, &ciphertext, target).unwrap();
+        let decap_ss = dec_key.key_data().unwrap().to_vec();
+
+        assert_eq!(shared_secret, decap_ss);
+    }
+
+    #[test]
+    fn test_mlkem_import_invalid_size() {
+        fixture_init();
+
+        let result = import_mlkem_public_key(&[0u8; 100], MlKemParameterSet::MlKem768);
+        assert!(result.is_err());
+    }
+
+    // ========================================================================
+    // Legacy API tests (for backwards compatibility)
+    // ========================================================================
+
     #[test]
     fn test_mlkem_parameter_set_sizes() {
         // ML-KEM-768 sizes

src/kem_combiners.rs

@@ -154,8 +154,8 @@ fn xwing_combiner(
 /// RFC 8410 OID for X25519: 1.3.101.110
 const RFC8410_OID_X25519: &[u8] = &[0x2b, 0x65, 0x6e];
 
-/// Import a raw X25519 public key.
-fn import_x25519_public_key(raw_key: &[u8; 32]) -> Res<PublicKey> {
+/// Import a raw X25519 public key from 32 bytes.
+pub fn import_x25519_public_key(raw_key: &[u8; 32]) -> Res<PublicKey> {
     // Build SPKI structure for X25519 per RFC 8410
     // SEQUENCE {
     //   SEQUENCE {
@@ -194,6 +194,31 @@ fn import_x25519_public_key(raw_key: &[u8; 32]) -> Res<PublicKey> {
     import_ec_public_key_from_spki(&spki)
 }
 
+/// Import an X-Wing (ML-KEM-768 + X25519) public key from raw bytes.
+///
+/// Expects 1216 bytes: 1184 bytes ML-KEM-768 || 32 bytes X25519.
+/// Returns a tuple of (mlkem_public, x25519_public) `PublicKey` objects
+/// suitable for passing to `xwing_encapsulate()`.
+pub fn import_xwing_public_key(
+    raw_key: &[u8],
+) -> Res<(PublicKey, PublicKey)> {
+    use crate::kem::{import_mlkem_public_key, MlKemParameterSet};
+
+    if raw_key.len() != XWING_MLKEM768_X25519_PUBLIC_KEY_SIZE {
+        return Err(Error::InvalidInput);
+    }
+
+    let mlkem_pk_bytes = &raw_key[..MLKEM768_PUBLIC_KEY_SIZE];
+    let x25519_pk_bytes: &[u8; 32] = raw_key[MLKEM768_PUBLIC_KEY_SIZE..]
+        .try_into()
+        .map_err(|_| Error::InvalidInput)?;
+
+    let mlkem_pk = import_mlkem_public_key(mlkem_pk_bytes, MlKemParameterSet::MlKem768)?;
+    let x25519_pk = import_x25519_public_key(x25519_pk_bytes)?;
+
+    Ok((mlkem_pk, x25519_pk))
+}
+
 /// Encapsulate using an X-Wing public key.
 ///
 /// This function generates a random shared secret and encapsulates it using
@@ -449,4 +474,62 @@ mod tests {
         );
         assert!(result.is_err());
     }
+
+    #[test]
+    fn test_xwing_import_public_key_roundtrip() {
+        use crate::err::secstatus_to_res;
+        use crate::p11;
+        use crate::util::SECItemMut;
+
+        fixture_init();
+
+        let keypair = XWingKeyPair::generate().unwrap();
+
+        // Export ML-KEM public key bytes via PK11_ReadRawAttribute(CKA_VALUE)
+        let mut key_item = SECItemMut::make_empty();
+        secstatus_to_res(unsafe {
+            p11::PK11_ReadRawAttribute(
+                p11::PK11ObjectType::PK11_TypePubKey,
+                (*keypair.mlkem_public).cast(),
+                pkcs11_bindings::CKA_VALUE,
+                key_item.as_mut(),
+            )
+        })
+        .unwrap();
+        let mlkem_pk_bytes = key_item.as_slice().to_owned();
+
+        // Export X25519 public key bytes
+        let x25519_raw = extract_x25519_public_key_bytes(&keypair.x25519_public).unwrap();
+
+        // Concatenate: ML-KEM-768 pk || X25519 pk
+        let mut combined = Vec::with_capacity(XWING_MLKEM768_X25519_PUBLIC_KEY_SIZE);
+        combined.extend_from_slice(&mlkem_pk_bytes);
+        combined.extend_from_slice(&x25519_raw);
+        assert_eq!(combined.len(), XWING_MLKEM768_X25519_PUBLIC_KEY_SIZE);
+
+        // Import
+        let (imported_mlkem, imported_x25519) = import_xwing_public_key(&combined).unwrap();
+
+        // Encapsulate with imported keys
+        let encap = xwing_encapsulate(&imported_mlkem, &imported_x25519).unwrap();
+
+        // Decapsulate with original private keys
+        let ss = xwing_decapsulate(
+            &encap.ciphertext,
+            &keypair.mlkem_private,
+            &keypair.x25519_private,
+            &keypair.x25519_public,
+        )
+        .unwrap();
+
+        assert_eq!(encap.shared_secret, ss);
+    }
+
+    #[test]
+    fn test_xwing_import_invalid_size() {
+        fixture_init();
+
+        let result = import_xwing_public_key(&[0u8; 100]);
+        assert!(result.is_err());
+    }
 }

src/lib.rs

@@ -87,12 +87,13 @@ pub use self::{
     ext::{ExtensionHandler, ExtensionHandlerResult, ExtensionWriterResult},
     kem::{
         decapsulate as kem_decapsulate, encapsulate as kem_encapsulate,
-        generate_keypair as kem_generate_keypair, KemEncapResult, KemKeypair, KemParameterSet,
-        MlKemKeypair, MlKemParameterSet,
+        generate_keypair as kem_generate_keypair, import_mlkem_public_key, KemEncapResult,
+        KemKeypair, KemParameterSet, MlKemKeypair, MlKemParameterSet,
     },
     kem_combiners::{
-        xwing_decapsulate, xwing_encapsulate, XWingEncapResult, XWingKeyPair,
-        XWING_MLKEM768_X25519_CIPHERTEXT_SIZE, XWING_MLKEM768_X25519_PUBLIC_KEY_SIZE, XWING_MLKEM768_X25519_SECRET_KEY_SIZE,
+        import_x25519_public_key, import_xwing_public_key, xwing_decapsulate, xwing_encapsulate,
+        XWingEncapResult, XWingKeyPair, XWING_MLKEM768_X25519_CIPHERTEXT_SIZE,
+        XWING_MLKEM768_X25519_PUBLIC_KEY_SIZE, XWING_MLKEM768_X25519_SECRET_KEY_SIZE,
         XWING_MLKEM768_X25519_SHARED_SECRET_SIZE,
     },
     p11::{random, randomize, PrivateKey, PublicKey, SymKey},