Added derive_keys function to get additional derived keys from mechanism params

Signed-off-by: Jacob Prud'homme <[email protected]>

Author: jacobprudhomme

cryptoki/src/mechanism/kbkdf.rs

@@ -3,7 +3,7 @@
 //! Mechanisms of NIST key-based key derive functions (SP 800-108, informally KBKDF)
 //! See: <https://docs.oasis-open.org/pkcs11/pkcs11-curr/v3.0/os/pkcs11-curr-v3.0-os.html#_Toc30061446>
 
-use std::{convert::TryInto, marker::PhantomData, ptr};
+use core::{convert::TryInto, marker::PhantomData, ptr, slice};
 
 use cryptoki_sys::{
     CK_ATTRIBUTE, CK_ATTRIBUTE_PTR, CK_DERIVED_KEY, CK_DERIVED_KEY_PTR, CK_OBJECT_HANDLE,
@@ -158,14 +158,12 @@ impl<'a> KbkdfCounterParams<'a> {
     pub fn new(
         prf_mechanism: MechanismType,
         prf_data_params: Vec<PrfDataParam<'a>>,
-        additional_derived_keys: Vec<DerivedKey<'a>>,
+        mut additional_derived_keys: Vec<DerivedKey<'a>>,
     ) -> Self {
         let prf_data_params: Vec<CK_PRF_DATA_PARAM> =
-            prf_data_params.into_iter().map(Into::into).collect();
-        let additional_derived_keys: Vec<CK_DERIVED_KEY> = additional_derived_keys
-            .into_iter()
-            .map(Into::into)
-            .collect();
+            prf_data_params.iter().map(Into::into).collect();
+        let additional_derived_keys: Vec<CK_DERIVED_KEY> =
+            additional_derived_keys.iter_mut().map(Into::into).collect();
 
         Self {
             inner: cryptoki_sys::CK_SP800_108_KDF_PARAMS {
@@ -184,6 +182,23 @@ impl<'a> KbkdfCounterParams<'a> {
             _marker: PhantomData,
         }
     }
+
+    /// The additional keys derived by the KDF, as per the params
+    pub fn additional_derived_keys(&self) -> Vec<CK_OBJECT_HANDLE> {
+        let derived_keys = unsafe {
+            slice::from_raw_parts(
+                self.inner.pAdditionalDerivedKeys,
+                self.inner.ulAdditionalDerivedKeys as _,
+            )
+        };
+
+        unsafe {
+            derived_keys
+                .iter()
+                .map(|derived_key| *derived_key.phKey)
+                .collect()
+        }
+    }
 }
 
 /// NIST SP 800-108 (aka KBKDF) feedback-mode parameters.
@@ -214,14 +229,12 @@ impl<'a> KbkdfFeedbackParams<'a> {
         prf_mechanism: MechanismType,
         prf_data_params: Vec<PrfDataParam<'a>>,
         iv: Option<&'a [u8]>,
-        additional_derived_keys: Vec<DerivedKey<'a>>,
+        mut additional_derived_keys: Vec<DerivedKey<'a>>,
     ) -> Self {
         let prf_data_params: Vec<CK_PRF_DATA_PARAM> =
-            prf_data_params.into_iter().map(Into::into).collect();
-        let additional_derived_keys: Vec<CK_DERIVED_KEY> = additional_derived_keys
-            .into_iter()
-            .map(Into::into)
-            .collect();
+            prf_data_params.iter().map(Into::into).collect();
+        let additional_derived_keys: Vec<CK_DERIVED_KEY> =
+            additional_derived_keys.iter_mut().map(Into::into).collect();
 
         Self {
             inner: cryptoki_sys::CK_SP800_108_FEEDBACK_KDF_PARAMS {
@@ -246,6 +259,23 @@ impl<'a> KbkdfFeedbackParams<'a> {
             _marker: PhantomData,
         }
     }
+
+    /// The additional keys derived by the KDF, as per the params
+    pub fn additional_derived_keys(&self) -> Vec<CK_OBJECT_HANDLE> {
+        let derived_keys = unsafe {
+            slice::from_raw_parts(
+                self.inner.pAdditionalDerivedKeys,
+                self.inner.ulAdditionalDerivedKeys as _,
+            )
+        };
+
+        unsafe {
+            derived_keys
+                .iter()
+                .map(|derived_key| *derived_key.phKey)
+                .collect()
+        }
+    }
 }
 
 /// NIST SP 800-108 (aka KBKDF) double pipeline-mode parameters.
@@ -273,14 +303,12 @@ impl<'a> KbkdfDoublePipelineParams<'a> {
     pub fn new(
         prf_mechanism: MechanismType,
         prf_data_params: Vec<PrfDataParam<'a>>,
-        additional_derived_keys: Vec<DerivedKey<'a>>,
+        mut additional_derived_keys: Vec<DerivedKey<'a>>,
     ) -> Self {
         let prf_data_params: Vec<CK_PRF_DATA_PARAM> =
-            prf_data_params.into_iter().map(Into::into).collect();
-        let additional_derived_keys: Vec<CK_DERIVED_KEY> = additional_derived_keys
-            .into_iter()
-            .map(Into::into)
-            .collect();
+            prf_data_params.iter().map(Into::into).collect();
+        let additional_derived_keys: Vec<CK_DERIVED_KEY> =
+            additional_derived_keys.iter_mut().map(Into::into).collect();
 
         Self {
             inner: cryptoki_sys::CK_SP800_108_KDF_PARAMS {
@@ -299,10 +327,27 @@ impl<'a> KbkdfDoublePipelineParams<'a> {
             _marker: PhantomData,
         }
     }
+
+    /// The additional keys derived by the KDF, as per the params
+    pub fn additional_derived_keys(&self) -> Vec<CK_OBJECT_HANDLE> {
+        let derived_keys = unsafe {
+            slice::from_raw_parts(
+                self.inner.pAdditionalDerivedKeys,
+                self.inner.ulAdditionalDerivedKeys as _,
+            )
+        };
+
+        unsafe {
+            derived_keys
+                .iter()
+                .map(|derived_key| *derived_key.phKey)
+                .collect()
+        }
+    }
 }
 
-impl<'a> From<PrfDataParam<'a>> for CK_PRF_DATA_PARAM {
-    fn from(value: PrfDataParam<'a>) -> Self {
+impl<'a> From<&PrfDataParam<'a>> for CK_PRF_DATA_PARAM {
+    fn from(value: &PrfDataParam<'a>) -> Self {
         Self {
             type_: match value {
                 PrfDataParam::IterationVariable => CK_SP800_108_ITERATION_VARIABLE,
@@ -312,8 +357,8 @@ impl<'a> From<PrfDataParam<'a>> for CK_PRF_DATA_PARAM {
             },
             pValue: match value {
                 PrfDataParam::IterationVariable => ptr::null_mut(),
-                PrfDataParam::Counter(inner) => &inner as *const _ as *mut _,
-                PrfDataParam::DkmLength(inner) => &inner as *const _ as *mut _,
+                PrfDataParam::Counter(inner) => inner as *const _ as *mut _,
+                PrfDataParam::DkmLength(inner) => inner as *const _ as *mut _,
                 PrfDataParam::ByteArray(data) => data.as_ptr() as *mut _,
             },
             ulValueLen: match value {
@@ -331,17 +376,17 @@ impl<'a> From<PrfDataParam<'a>> for CK_PRF_DATA_PARAM {
     }
 }
 
-impl<'a> From<PrfCounterDataParam<'a>> for CK_PRF_DATA_PARAM {
-    fn from(value: PrfCounterDataParam<'a>) -> Self {
+impl<'a> From<&PrfCounterDataParam<'a>> for CK_PRF_DATA_PARAM {
+    fn from(value: &PrfCounterDataParam<'a>) -> Self {
         Self {
             type_: match value {
                 PrfCounterDataParam::IterationVariable(_) => CK_SP800_108_ITERATION_VARIABLE,
                 PrfCounterDataParam::DkmLength(_) => CK_SP800_108_DKM_LENGTH,
                 PrfCounterDataParam::ByteArray(_) => CK_SP800_108_BYTE_ARRAY,
             },
             pValue: match value {
-                PrfCounterDataParam::IterationVariable(inner) => &inner as *const _ as *mut _,
-                PrfCounterDataParam::DkmLength(inner) => &inner as *const _ as *mut _,
+                PrfCounterDataParam::IterationVariable(inner) => inner as *const _ as *mut _,
+                PrfCounterDataParam::DkmLength(inner) => inner as *const _ as *mut _,
                 PrfCounterDataParam::ByteArray(data) => data.as_ptr() as *mut _,
             },
             ulValueLen: match value {
@@ -360,9 +405,9 @@ impl<'a> From<PrfCounterDataParam<'a>> for CK_PRF_DATA_PARAM {
     }
 }
 
-impl<'a> From<DerivedKey<'a>> for CK_DERIVED_KEY {
-    fn from(mut value: DerivedKey<'a>) -> Self {
-        let template: Vec<CK_ATTRIBUTE> = value.template.iter().map(|attr| attr.into()).collect();
+impl<'a> From<&mut DerivedKey<'a>> for CK_DERIVED_KEY {
+    fn from(value: &mut DerivedKey<'a>) -> Self {
+        let template: Vec<CK_ATTRIBUTE> = value.template.iter().map(Into::into).collect();
 
         Self {
             pTemplate: template.as_ptr() as CK_ATTRIBUTE_PTR,

cryptoki/src/mechanism/mod.rs

@@ -12,7 +12,6 @@ mod mechanism_info;
 pub mod rsa;
 pub mod vendor_defined;
 
-use crate::error::Error;
 use cryptoki_sys::*;
 use log::error;
 use std::convert::{TryFrom, TryInto};
@@ -23,7 +22,9 @@ use std::ops::Deref;
 use std::ptr::null_mut;
 use vendor_defined::VendorDefinedMechanism;
 
+use crate::error::Error;
 use crate::mechanism::rsa::PkcsOaepParams;
+use crate::object::ObjectHandle;
 pub use mechanism_info::MechanismInfo;
 
 #[derive(Copy, Debug, Clone, PartialEq, Eq)]
@@ -1283,3 +1284,25 @@ impl MessageParam<'_> {
         }
     }
 }
+
+/// Trait for mechanism types that define additional keys to be derived in their parameters
+pub trait HasAdditionalDerivedKeys {
+    /// Get the object handles for the additional keys that were derived
+    fn additional_derived_keys(&self) -> Vec<ObjectHandle>;
+}
+
+impl HasAdditionalDerivedKeys for &Mechanism<'_> {
+    fn additional_derived_keys(&self) -> Vec<ObjectHandle> {
+        let additional_derived_keys = match self {
+            Mechanism::KbkdfCounter(params) => params.additional_derived_keys(),
+            Mechanism::KbkdfFeedback(params) => params.additional_derived_keys(),
+            Mechanism::KbkdfDoublePipeline(params) => params.additional_derived_keys(),
+            _ => unimplemented!("The given mechanism doesn't define additional keys to derive"), // TODO: this or return an option?
+        };
+
+        additional_derived_keys
+            .into_iter()
+            .map(ObjectHandle::new)
+            .collect()
+    }
+}

cryptoki/src/session/key_management.rs

@@ -4,7 +4,7 @@
 
 use crate::context::Function;
 use crate::error::{Result, Rv};
-use crate::mechanism::Mechanism;
+use crate::mechanism::{HasAdditionalDerivedKeys as _, Mechanism};
 use crate::object::{Attribute, ObjectHandle};
 use crate::session::Session;
 use cryptoki_sys::{CK_ATTRIBUTE, CK_MECHANISM, CK_MECHANISM_PTR};
@@ -93,6 +93,33 @@ impl Session {
         Ok(ObjectHandle::new(handle))
     }
 
+    /// Derives multiple keys from a base key
+    pub fn derive_keys(
+        &self,
+        mechanism: &Mechanism,
+        base_key: ObjectHandle,
+        template: &[Attribute],
+    ) -> Result<(ObjectHandle, Vec<ObjectHandle>)> {
+        let mut c_mechanism: CK_MECHANISM = mechanism.into();
+        let mut template: Vec<CK_ATTRIBUTE> = template.iter().map(|attr| attr.into()).collect();
+        let mut handle = 0;
+        unsafe {
+            Rv::from(get_pkcs11!(self.client(), C_DeriveKey)(
+                self.handle(),
+                &mut c_mechanism as CK_MECHANISM_PTR,
+                base_key.handle(),
+                template.as_mut_ptr(),
+                template.len().try_into()?,
+                &mut handle,
+            ))
+            .into_result(Function::DeriveKey)?;
+        }
+
+        let additional_key_handles = mechanism.additional_derived_keys();
+
+        Ok((ObjectHandle::new(handle), additional_key_handles))
+    }
+
     /// Wrap key
     pub fn wrap_key(
         &self,