feat: per-slot database

pkcs11/src/backend/db/mod.rs

@@ -5,48 +5,86 @@
 
 pub mod attr;
 pub mod object;
-use std::collections::HashMap;
+use std::{collections::HashMap, time::SystemTime};
 
 use cryptoki_sys::CK_OBJECT_HANDLE;
 pub use object::{Object, ObjectHandle};
 
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub struct Db {
     objects: HashMap<CK_OBJECT_HANDLE, Object>,
     next_handle: CK_OBJECT_HANDLE,
+    last_fetchall_timestamp: Option<SystemTime>,
 }
 
 impl Db {
     pub fn new() -> Self {
         Self {
             objects: HashMap::new(),
-            next_handle: 1,
+            next_handle: 0,
+            last_fetchall_timestamp: None,
+        }
+    }
+
+    pub fn fetched_all_keys(&self) -> bool {
+        self.last_fetchall_timestamp
+            .map(|last| {
+                last.elapsed()
+                    // cache for 5 minutes
+                    .map(|elapsed| elapsed.as_secs() < 300)
+                    .unwrap_or(false)
+            })
+            .unwrap_or(false)
+    }
+
+    pub fn set_fetched_all_keys(&mut self, fetched_all_keys: bool) {
+        if fetched_all_keys {
+            self.last_fetchall_timestamp = Some(SystemTime::now());
+        } else {
+            self.last_fetchall_timestamp = None;
         }
     }
 
     pub fn clear(&mut self) {
+        self.set_fetched_all_keys(false);
         self.objects.clear();
     }
 
-    pub fn enumerate(&self) -> impl Iterator<Item = (ObjectHandle, &Object)> {
+    pub fn enumerate(&self) -> impl Iterator<Item = (CK_OBJECT_HANDLE, &Object)> {
         self.objects
             .iter()
             .enumerate()
-            .map(|(_, (handle, object))| (ObjectHandle::from(*handle), object))
+            .map(|(_, (handle, object))| (*handle, object))
     }
 
-    pub fn add_object(&mut self, object: Object) -> (CK_OBJECT_HANDLE, &Object) {
+    pub fn add_object(&mut self, object: Object) -> (CK_OBJECT_HANDLE, Object) {
+        // check if the object already exists
+
+        let found = self
+            .objects
+            .iter_mut()
+            .find(|(_, obj)| obj.id == object.id && obj.kind == object.kind);
+
+        if let Some((handle, obj)) = found {
+            *obj = object;
+            return (*handle, obj.clone());
+        }
+
+        // increment the handle
+
         let handle = self.next_handle;
-        self.objects.insert(handle, object);
         self.next_handle += 1;
-        (handle, self.objects.get(&handle).unwrap())
+
+        self.objects.insert(handle, object);
+
+        (handle, self.objects.get(&handle).unwrap().clone())
     }
 
-    pub fn object(&self, handle: ObjectHandle) -> Option<&Object> {
-        self.objects.get(&CK_OBJECT_HANDLE::from(handle))
+    pub fn object(&self, handle: CK_OBJECT_HANDLE) -> Option<&Object> {
+        self.objects.get(&handle)
     }
 
-    pub fn remove(&mut self, handle: ObjectHandle) -> Option<Object> {
-        self.objects.remove(&CK_OBJECT_HANDLE::from(handle))
+    pub fn remove(&mut self, handle: CK_OBJECT_HANDLE) -> Option<Object> {
+        self.objects.remove(&handle)
     }
 }

pkcs11/src/backend/session.rs

@@ -1,8 +1,11 @@
-use std::{collections::HashMap, sync::Arc};
+use std::{
+    collections::HashMap,
+    sync::{Arc, Mutex},
+};
 
 use cryptoki_sys::{
-    CKR_DEVICE_ERROR, CKR_DEVICE_MEMORY, CKR_OK, CKR_USER_NOT_LOGGED_IN, CK_FLAGS,
-    CK_OBJECT_HANDLE, CK_RV, CK_SESSION_HANDLE, CK_SLOT_ID, CK_USER_TYPE,
+    CKR_DEVICE_ERROR, CKR_DEVICE_MEMORY, CKR_KEY_HANDLE_INVALID, CKR_OK, CKR_USER_NOT_LOGGED_IN,
+    CK_FLAGS, CK_OBJECT_HANDLE, CK_RV, CK_SESSION_HANDLE, CK_SLOT_ID, CK_USER_TYPE,
 };
 use log::{debug, error, trace};
 use openapi::apis::{
@@ -16,7 +19,7 @@ use crate::{
 };
 
 use super::{
-    db::{self, attr::CkRawAttrTemplate, object::ObjectKind, Db, Object, ObjectHandle},
+    db::{self, attr::CkRawAttrTemplate, object::ObjectKind, Db, Object},
     decrypt::DecryptCtx,
     encrypt::EncryptCtx,
     key::{create_key_from_template, generate_key_from_template},
@@ -88,8 +91,7 @@ pub struct Session {
     pub login_ctx: LoginCtx,
     pub flags: CK_FLAGS,
     pub device_error: CK_RV,
-    pub fetched_all_keys: bool,
-    pub db: Db,
+    pub db: Arc<Mutex<Db>>,
     pub sign_ctx: Option<SignCtx>,
     pub encrypt_ctx: Option<EncryptCtx>,
     pub decrypt_ctx: Option<DecryptCtx>,
@@ -110,8 +112,7 @@ impl Session {
             login_ctx,
             slot_id,
             flags,
-            fetched_all_keys: false,
-            db: Db::new(),
+            db: slot.db.clone(),
             device_error: CKR_OK,
             sign_ctx: None,
             encrypt_ctx: None,
@@ -179,19 +180,23 @@ impl Session {
 
         trace!("sign_init() called with key handle {}", key_handle);
         trace!("sign_init() called with mechanism {:?}", mechanism);
-        trace!("sign_init() db size : {}", self.db.enumerate().count());
+
+        let db = match self.db.lock() {
+            Ok(db) => db,
+            Err(err) => {
+                error!("Failed to lock db: {:?}", err);
+                return CKR_DEVICE_ERROR;
+            }
+        };
 
         // get key id from the handle
 
-        let key = match self
-            .db
-            .object(ObjectHandle::from(key_handle))
-            .ok_or(cryptoki_sys::CKR_KEY_HANDLE_INVALID)
-        {
-            Ok(object) => object,
-            Err(err) => {
-                error!("Failed to get key: {:?}", err);
-                return err;
+        let key = match db.object(key_handle) {
+            Some(object) => object,
+
+            None => {
+                error!("Failed to get key: invalid handle");
+                return CKR_KEY_HANDLE_INVALID;
             }
         };
 
@@ -245,17 +250,22 @@ impl Session {
             return cryptoki_sys::CKR_OPERATION_ACTIVE;
         }
 
+        let db = match self.db.lock() {
+            Ok(db) => db,
+            Err(err) => {
+                error!("Failed to lock db: {:?}", err);
+                return CKR_DEVICE_ERROR;
+            }
+        };
+
         // get key id from the handle
 
-        let key = match self
-            .db
-            .object(ObjectHandle::from(key_handle))
-            .ok_or(cryptoki_sys::CKR_KEY_HANDLE_INVALID)
-        {
-            Ok(object) => object,
-            Err(err) => {
-                error!("Failed to get key: {:?}", err);
-                return err;
+        let key = match db.object(key_handle) {
+            Some(object) => object,
+
+            None => {
+                error!("Failed to get key: invalid handle");
+                return CKR_KEY_HANDLE_INVALID;
             }
         };
 
@@ -314,17 +324,22 @@ impl Session {
             return cryptoki_sys::CKR_OPERATION_ACTIVE;
         }
 
+        let db = match self.db.lock() {
+            Ok(db) => db,
+            Err(err) => {
+                error!("Failed to lock db: {:?}", err);
+                return CKR_DEVICE_ERROR;
+            }
+        };
+
         // get key id from the handle
 
-        let key = match self
-            .db
-            .object(ObjectHandle::from(key_handle))
-            .ok_or(cryptoki_sys::CKR_KEY_HANDLE_INVALID)
-        {
-            Ok(object) => object,
-            Err(err) => {
-                error!("Failed to get key: {:?}", err);
-                return err;
+        let key = match db.object(key_handle) {
+            Some(object) => object,
+
+            None => {
+                error!("Failed to get key: invalid handle");
+                return CKR_KEY_HANDLE_INVALID;
             }
         };
 
@@ -379,8 +394,16 @@ impl Session {
         self.decrypt_ctx = None;
     }
 
-    pub fn get_object(&self, handle: CK_OBJECT_HANDLE) -> Option<&Object> {
-        self.db.object(ObjectHandle::from(handle))
+    pub fn get_object(&self, handle: CK_OBJECT_HANDLE) -> Option<Object> {
+        let db = match self.db.lock() {
+            Ok(db) => db,
+            Err(err) => {
+                error!("Failed to lock db: {:?}", err);
+                return None;
+            }
+        };
+
+        db.object(handle).cloned()
     }
 
     pub(super) fn find_key(
@@ -416,16 +439,25 @@ impl Session {
     }
 
     fn fetch_all_keys(&mut self) -> Result<Vec<(CK_OBJECT_HANDLE, Object)>, CK_RV> {
-        if self.fetched_all_keys {
-            return Ok(self
-                .db
-                .enumerate()
-                .map(|(handle, obj)| (handle.into(), obj.clone()))
-                .collect());
-        }
+        {
+            let mut db = match self.db.lock() {
+                Ok(db) => db,
+                Err(err) => {
+                    error!("Failed to lock db: {:?}", err);
+                    return Err(CKR_DEVICE_ERROR);
+                }
+            };
 
-        // clear the db to not have any double entries
-        self.db.clear();
+            if db.fetched_all_keys() {
+                return Ok(db
+                    .enumerate()
+                    .map(|(handle, obj)| (handle, obj.clone()))
+                    .collect());
+            }
+
+            // clear the db to not have any double entries
+            db.clear();
+        }
 
         if !self
             .login_ctx
@@ -462,6 +494,15 @@ impl Session {
                 handles.push((handle, object));
             }
         }
+        let mut db = match self.db.lock() {
+            Ok(db) => db,
+            Err(err) => {
+                error!("Failed to lock db: {:?}", err);
+                return Err(CKR_DEVICE_ERROR);
+            }
+        };
+        db.set_fetched_all_keys(true);
+
         Ok(handles)
     }
 
@@ -489,9 +530,16 @@ impl Session {
             CKR_DEVICE_ERROR
         })?;
 
-        let r = self.db.add_object(object);
+        let r = self
+            .db
+            .lock()
+            .map_err(|err| {
+                error!("Failed to lock db: {:?}", err);
+                CKR_DEVICE_ERROR
+            })?
+            .add_object(object);
 
-        Ok((r.0, r.1.clone()))
+        Ok(r)
     }
 
     // we need the raw id when the CKA_KEY_ID doesn't parse to an alphanumeric string
@@ -534,8 +582,16 @@ impl Session {
 
         let mut result = Vec::new();
 
+        let mut db = match self.db.lock() {
+            Ok(db) => db,
+            Err(err) => {
+                error!("Failed to lock db: {:?}", err);
+                return Err(CKR_DEVICE_ERROR);
+            }
+        };
+
         for object in objects {
-            let r = self.db.add_object(object.clone());
+            let r = db.add_object(object.clone());
             result.push((r.0, r.1.clone()));
         }
 
@@ -577,7 +633,15 @@ impl Session {
             return Err(CKR_USER_NOT_LOGGED_IN);
         }
 
-        let key = self.db.object(ObjectHandle::from(handle)).ok_or_else(|| {
+        let mut db = match self.db.lock() {
+            Ok(db) => db,
+            Err(err) => {
+                error!("Failed to lock db: {:?}", err);
+                return Err(CKR_DEVICE_ERROR);
+            }
+        };
+
+        let key = db.object(handle).ok_or_else(|| {
             error!("Failed to delete object: invalid handle");
             CKR_DEVICE_ERROR
         })?;
@@ -612,7 +676,7 @@ impl Session {
             }
         }
 
-        self.db.remove(ObjectHandle::from(handle)).ok_or_else(|| {
+        db.remove(handle).ok_or_else(|| {
             error!("Failed to delete object: invalid handle");
             CKR_DEVICE_ERROR
         })?;

pkcs11/src/config/device.rs

@@ -1,7 +1,9 @@
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 
 use openapi::apis::configuration::Configuration;
 
+use crate::backend::db::Db;
+
 use super::config_file::UserConfig;
 
 #[derive(Debug, Clone)]
@@ -22,6 +24,7 @@ pub struct Slot {
     pub instances: Vec<Configuration>,
     pub operator: Option<UserConfig>,
     pub administrator: Option<UserConfig>,
+    pub db: Arc<Mutex<Db>>,
 }
 
 impl Slot {