Make Session hold a reference to a Pkcs11 or an owned Pkcs11

cryptoki/src/context/session_management.rs

@@ -2,7 +2,8 @@
 // SPDX-License-Identifier: Apache-2.0
 //! Session management functions
 
-use cryptoki_sys::{CKF_RW_SESSION, CKF_SERIAL_SESSION};
+use cryptoki_sys::{CKF_RW_SESSION, CKF_SERIAL_SESSION, CK_SESSION_HANDLE};
+use std::sync::Arc;
 
 use crate::context::Pkcs11;
 use crate::error::{Result, Rv};
@@ -13,7 +14,7 @@ use super::Function;
 
 impl Pkcs11 {
     #[inline(always)]
-    fn open_session(&self, slot_id: Slot, read_write: bool) -> Result<Session<'_>> {
+    fn open_session(&self, slot_id: Slot, read_write: bool) -> Result<CK_SESSION_HANDLE> {
         let mut session_handle = 0;
 
         let flags = if read_write {
@@ -33,7 +34,7 @@ impl Pkcs11 {
             .into_result(Function::OpenSession)?;
         }
 
-        Ok(Session::new(session_handle, self))
+        Ok(session_handle)
     }
 
     /// Open a new Read-Only session
@@ -63,13 +64,33 @@ impl Pkcs11 {
     /// # let _ = session; Ok(()) }
     /// ```
     pub fn open_ro_session(&self, slot_id: Slot) -> Result<Session<'_>> {
-        self.open_session(slot_id, false)
+        let session_handle = self.open_session(slot_id, false)?;
+        Ok(Session::new(session_handle, self))
+    }
+
+    /// Open a new Read-Only session
+    ///
+    /// This is like [`Pkcs11::open_ro_session`], but the session holds
+    /// an `Arc<Pkcs11>` instead of a `&Pkcs11`.
+    pub fn open_owned_ro_session(self: Arc<Self>, slot_id: Slot) -> Result<Session<'static>> {
+        let session_handle = self.open_session(slot_id, false)?;
+        Ok(Session::new(session_handle, self))
     }
 
     /// Open a new Read/Write session
     ///
     /// Note: No callback is set when opening the session.
     pub fn open_rw_session(&self, slot_id: Slot) -> Result<Session<'_>> {
-        self.open_session(slot_id, true)
+        let session_handle = self.open_session(slot_id, true)?;
+        Ok(Session::new(session_handle, self))
+    }
+
+    /// Open a new Read/Write session
+    ///
+    /// This is like [`Pkcs11::open_rw_session`], but the session
+    /// holds an `Arc<Pkcs11>` instead of a `&Pkcs11`.
+    pub fn open_owned_rw_session(self: Arc<Self>, slot_id: Slot) -> Result<Session<'static>> {
+        let session_handle = self.open_session(slot_id, true)?;
+        Ok(Session::new(session_handle, self))
     }
 }

cryptoki/src/session/mod.rs

@@ -8,6 +8,7 @@ use crate::error::Result;
 use cryptoki_sys::*;
 use std::fmt::Formatter;
 use std::marker::PhantomData;
+use std::sync::Arc;
 
 mod decryption;
 mod digesting;
@@ -28,6 +29,60 @@ pub use object_management::ObjectHandleIterator;
 pub use session_info::{SessionInfo, SessionState};
 pub use validation::ValidationFlagsType;
 
+/// A wrapper type that contains a reference or an owned value.
+///
+/// This wrapper type implements clone by wrapping the owned value in
+/// an `Arc`.
+#[derive(Clone)]
+pub enum MaybeOwned<'a, T>
+where
+    T: ?Sized,
+{
+    /// A reference to the thing.
+    Ref(&'a T),
+    /// An owned value.
+    Arc(Arc<T>),
+}
+
+impl<'a, T> AsRef<T> for MaybeOwned<'a, T>
+where
+    T: ?Sized,
+{
+    fn as_ref(&self) -> &T {
+        match self {
+            MaybeOwned::Ref(thing) => thing,
+            MaybeOwned::Arc(ref thing) => thing,
+        }
+    }
+}
+
+impl<'a, T> std::fmt::Debug for MaybeOwned<'a, T>
+where
+    T: std::fmt::Debug,
+{
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        self.as_ref().fmt(f)
+    }
+}
+
+impl<T> From<T> for MaybeOwned<'static, T> {
+    fn from(thing: T) -> MaybeOwned<'static, T> {
+        MaybeOwned::Arc(Arc::new(thing))
+    }
+}
+
+impl<T> From<Arc<T>> for MaybeOwned<'static, T> {
+    fn from(thing: Arc<T>) -> MaybeOwned<'static, T> {
+        MaybeOwned::Arc(thing)
+    }
+}
+
+impl<'a, T> From<&'a T> for MaybeOwned<'a, T> {
+    fn from(thing: &'a T) -> MaybeOwned<'a, T> {
+        MaybeOwned::Ref(thing)
+    }
+}
+
 /// Type that identifies a session
 ///
 /// It will automatically get closed (and logout) on drop.
@@ -37,7 +92,7 @@ pub use validation::ValidationFlagsType;
 #[derive(Debug)]
 pub struct Session<'a> {
     handle: CK_SESSION_HANDLE,
-    client: &'a Pkcs11,
+    client: MaybeOwned<'a, Pkcs11>,
     // This is not used but to prevent Session to automatically implement Send and Sync
     _guard: PhantomData<*mut u32>,
 }
@@ -61,10 +116,13 @@ impl<'a> std::fmt::UpperHex for Session<'a> {
 }
 
 impl<'a> Session<'a> {
-    pub(crate) fn new(handle: CK_SESSION_HANDLE, client: &'a Pkcs11) -> Self {
+    pub(crate) fn new(
+        handle: CK_SESSION_HANDLE,
+        client: impl Into<MaybeOwned<'a, Pkcs11>>,
+    ) -> Self {
         Session {
             handle,
-            client,
+            client: client.into(),
             _guard: PhantomData,
         }
     }
@@ -83,7 +141,10 @@ impl<'a> Session<'a> {
     }
 
     pub(crate) fn client(&self) -> &Pkcs11 {
-        self.client
+        match self.client {
+            MaybeOwned::Ref(pkcs11) => pkcs11,
+            MaybeOwned::Arc(ref pkcs11) => pkcs11,
+        }
     }
 }