Updates for getting attribute info - #42

cryptoki/src/functions/object_management.rs

@@ -8,6 +8,7 @@ use crate::types::object::{Attribute, AttributeInfo, AttributeType, ObjectHandle
 use crate::types::session::Session;
 use crate::Result;
 use cryptoki_sys::*;
+use std::collections::HashMap;
 use std::convert::TryInto;
 
 // Search 10 elements at a time
@@ -97,39 +98,125 @@ impl<'a> Session<'a> {
     }
 
     /// Get the attribute info of an object: if the attribute is present and its size.
+    ///
+    /// # Arguments
+    ///
+    /// * `object` - The [ObjectHandle] used to reference the object
+    /// * `attributes` - The list of attributes to get the information of
+    ///
+    /// # Returns
+    ///
+    /// This function will return a Vector of [AttributeInfo] enums that will either contain
+    /// the size of the requested attribute, [AttributeInfo::TypeInvalid] if the attribute is not a
+    /// valid type for the object, or [AttributeInfo::Sensitive] if the requested attribute is
+    /// sensitive and will not be returned to the user.
+    ///
+    /// The list of returned attributes is 1-to-1 matched with the provided vector of attribute
+    /// types.  If you wish, you may create a hash table simply by:
+    ///
+    /// ```rust
+    /// use cryptoki::Pkcs11;
+    /// use cryptoki::types::locking::CInitializeArgs;
+    /// use cryptoki::types::object::AttributeType;
+    /// use cryptoki::types::session::UserType;
+    /// use cryptoki::types::SessionFlags;
+    /// use std::collections::HashMap;
+    /// use std::env;
+    ///
+    /// let pkcs11 = Pkcs11::new(
+    ///         env::var("PKCS11_SOFTHSM2_MODULE")
+    ///             .unwrap_or_else(|_| "/usr/local/lib/softhsm/libsofthsm2.so".to_string()),
+    ///     )
+    ///     .unwrap();
+    ///
+    /// pkcs11.initialize(CInitializeArgs::OsThreads).unwrap();
+    /// let slot = pkcs11.get_slots_with_token().unwrap().remove(0);
+    /// let mut flags = SessionFlags::new();
+    /// let _ = flags.set_rw_session(true).set_serial_session(true);
+    ///
+    /// pkcs11.set_pin(slot, "1234").unwrap();
+    /// let session = pkcs11.open_session_no_callback(slot, flags).unwrap();
+    /// session.login(UserType::User);
+    ///
+    /// let empty_attrib= vec![];
+    /// if let Some(object) = session.find_objects(&empty_attrib).unwrap().get(0) {
+    ///     let attribute_types = vec![
+    ///         AttributeType::Token,
+    ///         AttributeType::Private,
+    ///         AttributeType::Modulus,
+    ///         AttributeType::KeyType,
+    ///         AttributeType::Verify,];
+    ///
+    ///     let attribute_info = session.get_attribute_info(*object, &attribute_types).unwrap();
+    ///
+    ///     let hash = attribute_types
+    ///         .iter()
+    ///         .zip(attribute_info.iter())
+    ///         .collect::<HashMap<_, _>>();
+    /// }
+    /// ```
+    ///
+    /// Alternatively, you can call [Session::get_attribute_info_map], found below.
     pub fn get_attribute_info(
         &self,
         object: ObjectHandle,
         attributes: &[AttributeType],
     ) -> Result<Vec<AttributeInfo>> {
-        let mut template: Vec<CK_ATTRIBUTE> = attributes
-            .iter()
-            .map(|attr_type| CK_ATTRIBUTE {
-                type_: (*attr_type).into(),
+        let mut results = Vec::new();
+
+        for attrib in attributes.iter() {
+            let mut template: Vec<CK_ATTRIBUTE> = vec![CK_ATTRIBUTE {
+                type_: (*attrib).into(),
                 pValue: std::ptr::null_mut(),
                 ulValueLen: 0,
-            })
-            .collect();
+            }];
 
-        match unsafe {
-            Rv::from(get_pkcs11!(self.client(), C_GetAttributeValue)(
-                self.handle(),
-                object.handle(),
-                template.as_mut_ptr(),
-                template.len().try_into()?,
-            ))
-        } {
-            Rv::Ok
-            | Rv::Error(RvError::AttributeSensitive)
-            | Rv::Error(RvError::AttributeTypeInvalid) => Ok(template
-                .iter()
-                .map(|attr| match attr.ulValueLen {
-                    CK_UNAVAILABLE_INFORMATION => Ok(AttributeInfo::Unavailable),
-                    len => Ok(AttributeInfo::Available(len.try_into()?)),
-                })
-                .collect::<Result<Vec<AttributeInfo>>>()?),
-            Rv::Error(rv_error) => Err(rv_error.into()),
+            match unsafe {
+                Rv::from(get_pkcs11!(self.client(), C_GetAttributeValue)(
+                    self.handle(),
+                    object.handle(),
+                    template.as_mut_ptr(),
+                    template.len().try_into()?,
+                ))
+            } {
+                Rv::Ok => {
+                    results.push(AttributeInfo::Available(template[0].ulValueLen.try_into()?))
+                }
+                Rv::Error(RvError::AttributeSensitive) => results.push(AttributeInfo::Sensitive),
+                Rv::Error(RvError::AttributeTypeInvalid) => {
+                    results.push(AttributeInfo::TypeInvalid)
+                }
+                rv => rv.into_result()?,
+            }
         }
+        Ok(results)
+    }
+
+    /// Get the attribute info of an object: if the attribute is present and its size.
+    ///
+    /// # Arguments
+    ///
+    /// * `object` - The [ObjectHandle] used to reference the object
+    /// * `attributes` - The list of attributes to get the information of
+    ///
+    /// # Returns
+    ///
+    /// This function will return a HashMap of [AttributeType] and [AttributeInfo] enums that will
+    /// either contain the size of the requested attribute, [AttributeInfo::TypeInvalid] if the
+    /// attribute is not a valid type for the object, or [AttributeInfo::Sensitive] if the requested
+    /// attribute is sensitive and will not be returned to the user.
+    pub fn get_attribute_info_map(
+        &self,
+        object: ObjectHandle,
+        attributes: Vec<AttributeType>,
+    ) -> Result<HashMap<AttributeType, AttributeInfo>> {
+        let attrib_info = self.get_attribute_info(object, attributes.as_slice())?;
+
+        Ok(attributes
+            .iter()
+            .cloned()
+            .zip(attrib_info.iter().cloned())
+            .collect::<HashMap<_, _>>())
     }
 
     /// Get the attributes values of an object.

cryptoki/src/types/function.rs

@@ -226,7 +226,7 @@ pub enum RvError {
     SessionCount,
     /// The specified session handle was invalid at the time that the function was invoked.  Note that this can happen if the session’s token is removed before the function invocation, since removing a token closes all sessions with it.
     SessionHandleInvalid,
-    /// The specified token does not support parallel sessions.  This is a legacy error code—in Cryptoki Version 2.01 and up, no token supports parallel sessions.  CKR_SESSION_PARALLEL_NOT_SUPPORTED can only be returned by C_OpenSession, and it is only returned when C_OpenSession is called in a particular [deprecated] way.
+    /// The specified token does not support parallel sessions.  This is a legacy error code—in Cryptoki Version 2.01 and up, no token supports parallel sessions.  CKR_SESSION_PARALLEL_NOT_SUPPORTED can only be returned by C_OpenSession, and it is only returned when C_OpenSession is called in a particular `deprecated` way.
     SessionParallelNotSupported,
     /// The specified session was unable to accomplish the desired action because it is a read-only session.  This return value has lower priority than CKR_TOKEN_WRITE_PROTECTED.
     SessionReadOnly,

cryptoki/src/types/mechanism/mod.rs

@@ -52,12 +52,12 @@ impl MechanismType {
 
     // DES
     /// DES3
-    /// Note that DES3 is deprecated. See https://nvlpubs.nist.gov/nistpubs/SpecialPublications/NIST.SP.800-131Ar2.pdf section 2, p. 6.
+    /// Note that DES3 is deprecated. See <https://nvlpubs.nist.gov/nistpubs/SpecialPublications/NIST.SP.800-131Ar2.pdf> section 2, p. 6.
     pub const DES3_KEY_GEN: MechanismType = MechanismType {
         val: CKM_DES3_KEY_GEN,
     };
     /// DES3 ECB
-    /// Note that DES3 is deprecated. See https://nvlpubs.nist.gov/nistpubs/SpecialPublications/NIST.SP.800-131Ar2.pdf section 2, p. 6.
+    /// Note that DES3 is deprecated. See <https://nvlpubs.nist.gov/nistpubs/SpecialPublications/NIST.SP.800-131Ar2.pdf> section 2, p. 6.
     pub const DES3_ECB: MechanismType = MechanismType { val: CKM_DES3_ECB };
 
     // ECC

cryptoki/src/types/object.rs

@@ -13,7 +13,7 @@ use std::ffi::c_void;
 use std::fmt::Formatter;
 use std::ops::Deref;
 
-#[derive(Debug, Copy, Clone, Ord, PartialOrd, Eq, PartialEq)]
+#[derive(Debug, Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash)]
 #[non_exhaustive]
 /// Type of an attribute
 pub enum AttributeType {
@@ -1065,7 +1065,7 @@ impl KeyType {
     /// DES2 key
     pub const DES2: KeyType = KeyType { val: CKK_DES2 };
     /// DES3 secret
-    /// Note that DES3 is deprecated. See https://nvlpubs.nist.gov/nistpubs/SpecialPublications/NIST.SP.800-131Ar2.pdf section 2, p. 6.
+    /// Note that DES3 is deprecated. See <https://nvlpubs.nist.gov/nistpubs/SpecialPublications/NIST.SP.800-131Ar2.pdf> section 2, p. 6.
     pub const DES3: KeyType = KeyType { val: CKK_DES3 };
     /// CAST key
     pub const CAST: KeyType = KeyType { val: CKK_CAST };
@@ -1286,8 +1286,10 @@ impl TryFrom<CK_KEY_TYPE> for KeyType {
 #[derive(Debug, Copy, Clone)]
 /// Information about the attribute of an object
 pub enum AttributeInfo {
-    /// The attribute is not defined for the object
-    Unavailable,
+    /// The requested attribute is not a valid attribute for the object
+    TypeInvalid,
+    /// The value of the attribute is sensitive and will not be returned
+    Sensitive,
     /// The attribute is available to get from the object and has the specified size in bytes.
     Available(usize),
 }

cryptoki/tests/basic.rs

@@ -9,7 +9,7 @@ use cryptoki::types::object::{Attribute, AttributeInfo, AttributeType, KeyType,
 use cryptoki::types::session::{SessionState, UserType};
 use cryptoki::types::SessionFlags;
 use serial_test::serial;
-use std::error::Error;
+use std::collections::HashMap;
 use std::sync::Arc;
 use std::thread;
 
@@ -539,3 +539,105 @@ fn generate_random_test() -> Result<()> {
     assert!(!random_vec.iter().all(|&x| x == 0));
     Ok(())
 }
+
+#[test]
+#[serial]
+fn get_attribute_info_test() -> Result<()> {
+    let (pkcs11, slot) = init_pins();
+
+    let mut flags = SessionFlags::new();
+    let _ = flags.set_rw_session(true).set_serial_session(true);
+
+    // open a session
+    let session = pkcs11.open_session_no_callback(slot, flags)?;
+
+    // log in the session
+    session.login(UserType::User)?;
+
+    // get mechanism
+    let mechanism = Mechanism::RsaPkcsKeyPairGen;
+
+    let public_exponent: Vec<u8> = vec![0x01, 0x00, 0x01];
+    let modulus_bits = 2048;
+
+    // pub key template
+    let pub_key_template = vec![
+        Attribute::Token(false.into()),
+        Attribute::Private(false.into()),
+        Attribute::PublicExponent(public_exponent),
+        Attribute::ModulusBits(modulus_bits.into()),
+    ];
+
+    // priv key template
+    let priv_key_template = vec![
+        Attribute::Token(false.into()),
+        Attribute::Sensitive(true.into()),
+        Attribute::Extractable(false.into()),
+    ];
+
+    // generate a key pair
+    let (public, private) =
+        session.generate_key_pair(&mechanism, &pub_key_template, &priv_key_template)?;
+
+    let pub_attribs = vec![AttributeType::PublicExponent, AttributeType::Modulus];
+    let mut priv_attribs = pub_attribs.clone();
+    priv_attribs.push(AttributeType::PrivateExponent);
+
+    let attrib_info = session.get_attribute_info(public, &pub_attribs)?;
+    let hash = pub_attribs
+        .iter()
+        .zip(attrib_info.iter())
+        .collect::<HashMap<_, _>>();
+
+    if let AttributeInfo::Available(size) = hash[&AttributeType::Modulus] {
+        assert_eq!(*size, 2048 / 8);
+    } else {
+        panic!("Modulus should not return Unavailable for an RSA public key");
+    }
+
+    match hash[&AttributeType::PublicExponent] {
+        AttributeInfo::Available(_) => {}
+        _ => panic!("Public Exponent should not return Unavailable for an RSA public key"),
+    }
+
+    let attrib_info = session.get_attribute_info(private, &priv_attribs)?;
+    let hash = priv_attribs
+        .iter()
+        .zip(attrib_info.iter())
+        .collect::<HashMap<_, _>>();
+
+    if let AttributeInfo::Available(size) = hash[&AttributeType::Modulus] {
+        assert_eq!(*size, 2048 / 8);
+    } else {
+        panic!("Modulus should not return Unavailable on an RSA private key");
+    }
+
+    match hash[&AttributeType::PublicExponent] {
+        AttributeInfo::Available(_) => {}
+        _ => panic!("PublicExponent should not return Unavailable on an RSA private key"),
+    }
+
+    match hash[&AttributeType::PrivateExponent] {
+        AttributeInfo::Sensitive => {}
+        _ => panic!("Private Exponent of RSA private key should be sensitive"),
+    }
+
+    let hash = session.get_attribute_info_map(private, priv_attribs)?;
+    if let AttributeInfo::Available(size) = hash[&AttributeType::Modulus] {
+        assert_eq!(size, 2048 / 8);
+    } else {
+        panic!("Modulus should not return Unavailable on an RSA private key");
+    }
+
+    match hash[&AttributeType::PublicExponent] {
+        AttributeInfo::Available(_) => {}
+        _ => panic!("Public Exponent should not return Unavailable for an RSA private key"),
+    }
+
+    match hash[&AttributeType::PrivateExponent] {
+        AttributeInfo::Sensitive => {}
+        _ => panic!("Private Exponent of RSA private key should be sensitive"),
+    }
+
+    Ok(())
+}