feat: disable key fetching if the application asks for it

Nitrokey · Sep 18, 2023 · 61df9dc · 61df9dc
1 parent 6f570d3
commit 61df9dc
Show file tree

Hide file tree

Showing 4 changed files with 75 additions and 49 deletions.
diff --git a/pkcs11/src/api/mod.rs b/pkcs11/src/api/mod.rs
@@ -15,7 +15,7 @@ pub mod verify;
 
 use crate::{
     backend::events::{fetch_slots_state, EventsManager},
-    data::{self, DEVICE, EVENTS_MANAGER, TOKENS_STATE},
+    data::{self, DEVICE, EVENTS_MANAGER, THREADS_ALLOWED, TOKENS_STATE},
     defs, padded_str,
 };
 use cryptoki_sys::{CK_INFO, CK_INFO_PTR, CK_RV, CK_VOID_PTR};
@@ -69,7 +69,7 @@ pub extern "C" fn C_Initialize(pInitArgs: CK_VOID_PTR) -> CK_RV {
 
         // currently we are using tokio that needs to create threads, so if the programs forbids us to create threads we return an error
         if flags & cryptoki_sys::CKF_LIBRARY_CANT_CREATE_OS_THREADS != 0 {
-            return cryptoki_sys::CKR_NEED_TO_CREATE_THREADS;
+            *THREADS_ALLOWED.lock().unwrap() = false;
         }
     }
 

diff --git a/pkcs11/src/backend/key.rs b/pkcs11/src/backend/key.rs
@@ -4,7 +4,7 @@ use std::{
 };
 
 use super::{
-    db::{self, attr::CkRawAttrTemplate, Object},
+    db::{self, attr::CkRawAttrTemplate, Db, Object},
     login::{self, LoginCtx},
     Error,
 };
@@ -555,3 +555,44 @@ fn extract_key_id_location_header(headers: HashMap<String, String>) -> Result<St
         .to_string();
     Ok(key_id)
 }
+
+pub type WorkResult = Result<Vec<(CK_ULONG, Object)>, Error>;
+
+pub fn fetch_loop(
+    keys: Arc<Mutex<Vec<nethsm_sdk_rs::models::KeyItem>>>,
+    db: Arc<Mutex<Db>>,
+    login_ctx: LoginCtx,
+    results: Arc<Mutex<Vec<WorkResult>>>,
+    kind: Option<ObjectKind>,
+) {
+    while let Some(key) = keys.lock().unwrap().pop() {
+        let key_id = key.key.clone();
+
+        if matches!(
+            kind,
+            None | Some(ObjectKind::Other)
+                | Some(ObjectKind::PrivateKey)
+                | Some(ObjectKind::PublicKey)
+                | Some(ObjectKind::SecretKey)
+        ) {
+            let login_ctx = login_ctx.clone();
+            let db = db.clone();
+            let res = fetch_key(&key_id, None, login_ctx, db);
+            results.lock().unwrap().push(res);
+        }
+
+        if matches!(kind, None | Some(ObjectKind::Certificate)) {
+            let login_ctx = login_ctx.clone();
+            let db = db.clone();
+            let res = match fetch_certificate(&key_id, None, login_ctx, db) {
+                Ok(vec) => Ok(vec),
+                Err(err) => {
+                    debug!("Failed to fetch certificate: {:?}", err);
+                    Ok(Vec::new())
+                }
+            };
+
+            results.lock().unwrap().push(res);
+        }
+    }
+}
diff --git a/pkcs11/src/backend/session.rs b/pkcs11/src/backend/session.rs
@@ -5,29 +5,31 @@ use std::{
 
 use cryptoki_sys::{
     CKR_OK, CK_FLAGS, CK_OBJECT_HANDLE, CK_RV, CK_SESSION_HANDLE, CK_SESSION_INFO, CK_SLOT_ID,
-    CK_ULONG, CK_USER_TYPE,
+    CK_USER_TYPE,
 };
 use log::{debug, error, trace};
 use nethsm_sdk_rs::apis::default_api;
 
 use crate::{
     backend::{login::UserMode, Error},
     config::device::Slot,
+    data::THREADS_ALLOWED,
 };
 
 use super::{
     db::{attr::CkRawAttrTemplate, object::ObjectKind, Db, Object},
     decrypt::DecryptCtx,
     encrypt::EncryptCtx,
-    key::{create_key_from_template, fetch_certificate, fetch_key, generate_key_from_template},
+    key::{
+        create_key_from_template, fetch_certificate, fetch_key, generate_key_from_template,
+        WorkResult,
+    },
     login::LoginCtx,
     mechanism::Mechanism,
     object::{EnumCtx, KeyRequirements},
     sign::SignCtx,
 };
 
-type WorkResult = Result<Vec<(CK_ULONG, Object)>, Error>;
-
 #[derive(Debug)]
 pub struct SessionManager {
     pub sessions: HashMap<CK_SESSION_HANDLE, Arc<Mutex<Session>>>,
@@ -520,54 +522,35 @@ impl Session {
 
         let keys = Arc::new(std::sync::Mutex::new(keys));
 
-        let mut thread_handles = Vec::new();
-
-        // 4 threads
-
-        for _ in 0..4 {
-            let keys = keys.clone();
+        if *THREADS_ALLOWED.lock()? {
+            let mut thread_handles = Vec::new();
 
-            let results = results.clone();
+            // 4 threads
 
-            let login_ctx = self.login_ctx.clone();
-            let db = self.db.clone();
+            for _ in 0..4 {
+                let keys = keys.clone();
 
-            thread_handles.push(std::thread::spawn(move || {
-                while let Some(key) = keys.lock().unwrap().pop() {
-                    let key_id = key.key.clone();
+                let results = results.clone();
 
-                    if matches!(
-                        kind,
-                        None | Some(ObjectKind::Other)
-                            | Some(ObjectKind::PrivateKey)
-                            | Some(ObjectKind::PublicKey)
-                            | Some(ObjectKind::SecretKey)
-                    ) {
-                        let login_ctx = login_ctx.clone();
-                        let db = db.clone();
-                        let res = fetch_key(&key_id, None, login_ctx, db);
-                        results.lock().unwrap().push(res);
-                    }
+                let login_ctx = self.login_ctx.clone();
+                let db = self.db.clone();
 
-                    if matches!(kind, None | Some(ObjectKind::Certificate)) {
-                        let login_ctx = login_ctx.clone();
-                        let db = db.clone();
-                        let res = match fetch_certificate(&key_id, None, login_ctx, db) {
-                            Ok(vec) => Ok(vec),
-                            Err(err) => {
-                                debug!("Failed to fetch certificate: {:?}", err);
-                                Ok(Vec::new())
-                            }
-                        };
-
-                        results.lock().unwrap().push(res);
-                    }
-                }
-            }));
-        }
+                thread_handles.push(std::thread::spawn(move || {
+                    super::key::fetch_loop(keys, db, login_ctx, results, kind)
+                }));
+            }
 
-        for handle in thread_handles {
-            handle.join().unwrap();
+            for handle in thread_handles {
+                handle.join().unwrap();
+            }
+        } else {
+            super::key::fetch_loop(
+                keys,
+                self.db.clone(),
+                self.login_ctx.clone(),
+                results.clone(),
+                kind,
+            );
         }
 
         let mut handles = Vec::new();

diff --git a/pkcs11/src/data.rs b/pkcs11/src/data.rs
@@ -32,6 +32,8 @@ lazy_static! {
     pub static ref EVENTS_MANAGER : Arc<RwLock<EventsManager>> = Arc::new(RwLock::new(EventsManager::new()));
     // Token present or not (true = present)
     pub static ref TOKENS_STATE : Arc<Mutex<std::collections::HashMap<CK_SLOT_ID, bool>>> = Arc::new(Mutex::new(std::collections::HashMap::new()));
+    // If the calling application allows threads to be used
+    pub static ref THREADS_ALLOWED : Arc<Mutex<bool>> = Arc::new(Mutex::new(true));
 }
 pub static mut FN_LIST: CK_FUNCTION_LIST = CK_FUNCTION_LIST {
     version: DEVICE_VERSION,