Session.find_objects(): implement fewer calls to C_FindObjects()

Signed-off-by: Eric Devolder <[email protected]>

cryptoki/src/session/object_management.rs

@@ -44,6 +44,13 @@ impl Session {
         while object_count > 0 {
             objects.extend_from_slice(&object_handles[..object_count.try_into()?]);
 
+            // If the returned slice is not full, there are no pending object handles to be returned.
+            // In which case, exit loop.
+            // This avoids, in many situations, an unecessary API call with 0 object handles returned.
+            if (object_count as usize) < MAX_OBJECT_COUNT {
+                break;
+            }
+
             unsafe {
                 Rv::from(get_pkcs11!(self.client(), C_FindObjects)(
                     self.handle(),

cryptoki/tests/basic.rs

@@ -309,6 +309,63 @@ fn get_token_info() -> TestResult {
     Ok(())
 }
 
+#[test]
+#[serial]
+fn session_find_objects() {
+    let (pkcs11, slot) = init_pins();
+    // open a session
+    let session = pkcs11.open_rw_session(slot).unwrap();
+
+    // log in the session
+    session
+        .login(UserType::User, Some(&AuthPin::new(USER_PIN.into())))
+        .unwrap();
+
+    // we generate 11 keys with the same CKA_ID
+    // we will check 3 different use cases, this will cover all cases for Session.find_objects
+    // find 11 keys with the same CKA_ID => should result in two internal iterations over MAX_OBJECT_COUNT
+    // find 10 keys with the same CKA_ID => should result in two internal iterations over MAX_OBJECT_COUNT
+    // find 9 keys with the same CKA_ID  => should result in one internal iteration over MAX_OBJECT_COUNT
+
+    (1..=11).for_each(|i| {
+        let key_template = vec![
+            Attribute::Token(true),
+            Attribute::Encrypt(true),
+            Attribute::Label(format!("key_{}", i).as_bytes().to_vec()),
+            Attribute::Id("12345678".as_bytes().to_vec()),          // reusing the same CKA_ID
+        ];
+
+        // generate a secret key
+        let _key = session
+            .generate_key(&Mechanism::Des3KeyGen, &key_template)
+            .unwrap();
+    });
+
+
+    // retrieve the keys by searching for them
+    let key_search_template = vec![
+        Attribute::Token(true),
+        Attribute::Id("12345678".as_bytes().to_vec()),
+        Attribute::Class(ObjectClass::SECRET_KEY),
+        Attribute::KeyType(KeyType::DES3),
+    ];
+
+    let mut found_keys = session.find_objects(&key_search_template).unwrap();
+    assert_eq!(found_keys.len(), 11);
+
+    // destroy one key
+    session.destroy_object(found_keys.pop().unwrap()).unwrap();
+
+    let mut found_keys = session.find_objects(&key_search_template).unwrap();
+    assert_eq!(found_keys.len(), 10);
+
+    // destroy another key
+    session.destroy_object(found_keys.pop().unwrap()).unwrap();
+    let found_keys = session.find_objects(&key_search_template).unwrap();
+    assert_eq!(found_keys.len(), 9);
+
+}
+
 #[test]
 #[serial]
 fn wrap_and_unwrap_key() {