Incremental object search

cryptoki/src/session/mod.rs

@@ -23,6 +23,7 @@ mod session_management;
 mod signing_macing;
 mod slot_token_management;
 
+pub use object_management::FindObjects;
 pub use session_info::{SessionInfo, SessionState};
 
 /// Type that identifies a session
@@ -126,11 +127,57 @@ impl Session {
         session_management::get_session_info(self)
     }
 
-    /// Search for session objects matching a template
-    pub fn find_objects(&self, template: &[Attribute]) -> Result<Vec<ObjectHandle>> {
+    /// Search for token and session objects matching a template
+    pub fn find_objects(&mut self, template: &[Attribute]) -> Result<Vec<ObjectHandle>> {
         object_management::find_objects(self, template)
     }
 
+    /// Initiate a search for token and session objects matching a template
+    ///
+    /// # Arguments
+    ///
+    /// * `template` - The list of attributes to match
+    ///
+    /// # Returns
+    ///
+    /// This function returns a [FindObjects], which represents an ongoing search.  The
+    /// lifetime of this search is tied to a mutable borrow of the session, so that there
+    /// may only be one search per session at once.  When the [FindObjects] is dropped,
+    /// the search is ended.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// use cryptoki::error::Result;
+    /// use cryptoki::object::{Attribute, AttributeType};
+    /// use cryptoki::session::Session;
+    ///
+    /// const BATCH_SIZE: usize = 10;
+    ///
+    /// fn print_object_labels(session: &mut Session, template: &[Attribute]) -> Result<()> {
+    ///     // Initiate the search.
+    ///     let mut search = session.find_objects_init(template)?;
+    ///
+    ///     // Iterate over batches of results, while find_next returns a non-empty batch
+    ///     while let ref objects @ [_, ..] = search.find_next(BATCH_SIZE)?[..] {
+    ///         // Iterate over objects in the batch.
+    ///         for &object in objects {
+    ///             // Look up the label for the object.  We can't use `session` directly here,
+    ///             // since it's mutably borrowed by search.  Instead, use `search.session()`.
+    ///             let attrs = search.session().get_attributes(object, &[AttributeType::Label])?;
+    ///             if let Some(Attribute::Label(label)) = attrs.get(0) {
+    ///                 println!("Found object: {}", String::from_utf8_lossy(&label));
+    ///             }
+    ///         }
+    ///     }
+    ///
+    ///     Ok(())
+    /// }
+    /// ```
+    pub fn find_objects_init<'a>(&'a mut self, template: &[Attribute]) -> Result<FindObjects<'a>> {
+        object_management::find_objects_init(self, template)
+    }
+
     /// Create a new object
     pub fn create_object(&self, template: &[Attribute]) -> Result<ObjectHandle> {
         object_management::create_object(self, template)
@@ -175,7 +222,7 @@ impl Session {
     /// pkcs11.initialize(CInitializeArgs::OsThreads).unwrap();
     /// let slot = pkcs11.get_slots_with_token().unwrap().remove(0);
     ///
-    /// let session = pkcs11.open_ro_session(slot).unwrap();
+    /// let mut session = pkcs11.open_ro_session(slot).unwrap();
     /// session.login(UserType::User, Some("fedcba"));
     ///
     /// let empty_attrib= vec![];

cryptoki/src/session/object_management.rs

@@ -6,62 +6,107 @@ use crate::error::{Result, Rv, RvError};
 use crate::object::{Attribute, AttributeInfo, AttributeType, ObjectHandle};
 use crate::session::Session;
 use cryptoki_sys::*;
+use log::error;
 use std::collections::HashMap;
 use std::convert::TryInto;
 
-// Search 10 elements at a time
-const MAX_OBJECT_COUNT: usize = 10;
-
-// See public docs on stub in parent mod.rs
-#[inline(always)]
-pub(super) fn find_objects(session: &Session, template: &[Attribute]) -> Result<Vec<ObjectHandle>> {
-    let mut template: Vec<CK_ATTRIBUTE> = template.iter().map(|attr| attr.into()).collect();
-
-    unsafe {
-        Rv::from(get_pkcs11!(session.client(), C_FindObjectsInit)(
-            session.handle(),
-            template.as_mut_ptr(),
-            template.len().try_into()?,
-        ))
-        .into_result()?;
-    }
-
-    let mut object_handles = [0; MAX_OBJECT_COUNT];
-    let mut object_count = 0;
-    let mut objects = Vec::new();
+/// Represents an ongoing object search
+///
+/// See the documentation for [Session::find_objects_init].
+#[derive(Debug)]
+pub struct FindObjects<'a> {
+    session: &'a mut Session,
+}
 
-    unsafe {
-        Rv::from(get_pkcs11!(session.client(), C_FindObjects)(
-            session.handle(),
-            object_handles.as_mut_ptr() as CK_OBJECT_HANDLE_PTR,
-            MAX_OBJECT_COUNT.try_into()?,
-            &mut object_count,
-        ))
-        .into_result()?;
-    }
+impl<'a> FindObjects<'a> {
+    /// Continue an ongoing object search
+    ///
+    /// # Arguments
+    ///
+    /// * `max_objects` - The maximum number of objects to return
+    ///
+    /// # Returns
+    ///
+    /// This function returns up to `max_objects` objects.  If there are no remaining
+    /// objects, or `max_objects` is 0, then it returns an empty vector.
+    pub fn find_next(&mut self, max_objects: usize) -> Result<Vec<ObjectHandle>> {
+        if max_objects == 0 {
+            return Ok(vec![]);
+        }
 
-    while object_count > 0 {
-        objects.extend_from_slice(&object_handles[..object_count.try_into()?]);
+        let mut object_handles = Vec::with_capacity(max_objects);
+        let mut object_count = 0;
 
         unsafe {
-            Rv::from(get_pkcs11!(session.client(), C_FindObjects)(
-                session.handle(),
-                object_handles.as_mut_ptr() as CK_OBJECT_HANDLE_PTR,
-                MAX_OBJECT_COUNT.try_into()?,
+            Rv::from(get_pkcs11!(self.session.client(), C_FindObjects)(
+                self.session.handle(),
+                object_handles.as_mut_ptr(),
+                max_objects.try_into()?,
                 &mut object_count,
             ))
             .into_result()?;
+            object_handles.set_len(object_count.try_into()?)
         }
+
+        Ok(object_handles.into_iter().map(ObjectHandle::new).collect())
     }
 
+    /// Get the session associated to the search
+    pub fn session(&self) -> &Session {
+        self.session
+    }
+}
+
+impl<'a> Drop for FindObjects<'a> {
+    fn drop(&mut self) {
+        if let Err(e) = find_objects_final_private(self.session) {
+            error!("Failed to terminate object search: {}", e);
+        }
+    }
+}
+
+fn find_objects_final_private(session: &Session) -> Result<()> {
     unsafe {
         Rv::from(get_pkcs11!(session.client(), C_FindObjectsFinal)(
             session.handle(),
         ))
+        .into_result()
+    }
+}
+
+// See public docs on stub in parent mod.rs
+#[inline(always)]
+pub(super) fn find_objects_init<'a>(
+    session: &'a mut Session,
+    template: &[Attribute],
+) -> Result<FindObjects<'a>> {
+    let mut template: Vec<CK_ATTRIBUTE> = template.iter().map(|attr| attr.into()).collect();
+    unsafe {
+        Rv::from(get_pkcs11!(session.client(), C_FindObjectsInit)(
+            session.handle(),
+            template.as_mut_ptr(),
+            template.len().try_into()?,
+        ))
         .into_result()?;
     }
+    Ok(FindObjects { session })
+}
 
-    let objects = objects.into_iter().map(ObjectHandle::new).collect();
+// Search 10 elements at a time
+const MAX_OBJECT_COUNT: usize = 10;
+
+// See public docs on stub in parent mod.rs
+#[inline(always)]
+pub(super) fn find_objects(
+    session: &mut Session,
+    template: &[Attribute],
+) -> Result<Vec<ObjectHandle>> {
+    let mut search = session.find_objects_init(template)?;
+    let mut objects = Vec::new();
+
+    while let ref new_objects @ [_, ..] = search.find_next(MAX_OBJECT_COUNT)?[..] {
+        objects.extend_from_slice(new_objects)
+    }
 
     Ok(objects)
 }