r2d2_cryptoki/
lib.rs

1#![warn(missing_docs)]
2#![doc = include_str!("../README.md")]
3
4use std::sync::{Arc, Mutex};
5
6pub use cryptoki;
7pub use r2d2;
8
9use cryptoki::{
10    context::{Function, Pkcs11},
11    error::RvError,
12    session::{Session, SessionState, UserType},
13    slot::{Limit, Slot},
14    types::AuthPin,
15};
16use r2d2::{CustomizeConnection, ManageConnection, NopConnectionCustomizer};
17
18/// Alias for this crate's instance of r2d2's Pool
19pub type Pool = r2d2::Pool<SessionManager>;
20/// Alias for this crate's instance of r2d2's PooledSession
21pub type PooledSession = r2d2::PooledConnection<SessionManager>;
22
23/// Manager holding all information necessary for opening new connections
24#[derive(Debug, Clone)]
25pub struct SessionManager {
26    pkcs11: Pkcs11,
27    slot: Slot,
28    session_state: SessionState,
29}
30
31/// Session types, holding the pin for the authenticated sessions
32#[derive(Debug, Clone)]
33pub enum SessionAuth {
34    /// [SessionState::RoPublic]
35    RoPublic,
36    /// [SessionState::RoUser]
37    RoUser(AuthPin),
38    /// [SessionState::RwPublic]
39    RwPublic,
40    /// [SessionState::RwUser]
41    RwUser(AuthPin),
42    /// [SessionState::RwSecurityOfficer]
43    RwSecurityOfficer(AuthPin),
44}
45
46/// Mandatory connection customizer for logins
47#[derive(Debug, Clone)]
48struct LoginCustomizer {
49    auth_pin: AuthPin,
50    user_type: UserType,
51    active_sessions: Arc<Mutex<u32>>,
52}
53
54impl SessionAuth {
55    fn as_state(&self) -> SessionState {
56        match self {
57            Self::RoPublic => SessionState::RoPublic,
58            Self::RoUser(_) => SessionState::RoUser,
59            Self::RwPublic => SessionState::RwPublic,
60            Self::RwUser(_) => SessionState::RwUser,
61            Self::RwSecurityOfficer(_) => SessionState::RwSecurityOfficer,
62        }
63    }
64
65    /// Returns the correct customizer to use for the specified session auth
66    pub fn into_customizer(self) -> Box<dyn CustomizeConnection<Session, cryptoki::error::Error>> {
67        match self {
68            Self::RoPublic | Self::RwPublic => Box::new(NopConnectionCustomizer),
69            Self::RoUser(auth_pin) | Self::RwUser(auth_pin) => Box::from(LoginCustomizer {
70                auth_pin,
71                user_type: UserType::User,
72                active_sessions: Default::default(),
73            }),
74            Self::RwSecurityOfficer(auth_pin) => Box::from(LoginCustomizer {
75                auth_pin,
76                user_type: UserType::So,
77                active_sessions: Default::default(),
78            }),
79        }
80    }
81}
82
83impl SessionManager {
84    /// # Example
85    /// ```no_run
86    ///  # use r2d2_cryptoki::{*, cryptoki::{context::*, types::AuthPin}};
87    ///  let pkcs11 = Pkcs11::new("libsofthsm2.so").unwrap();
88    ///  pkcs11 .initialize(CInitializeArgs::new(CInitializeFlags::OS_LOCKING_OK)).unwrap();
89    ///  let slots = pkcs11.get_slots_with_token().unwrap();
90    ///  let slot = slots.first().unwrap();
91    ///  let manager = SessionManager::new(pkcs11, *slot, &SessionAuth::RwUser(AuthPin::new("abcd".into())));
92    /// ```
93    pub fn new(pkcs11: Pkcs11, slot: Slot, session_auth: &SessionAuth) -> Self {
94        Self {
95            pkcs11,
96            slot,
97            session_state: session_auth.as_state(),
98        }
99    }
100
101    /// Returns the maximum number of sessions supported by the HSM.
102    ///
103    /// Arguments:
104    /// * `maximum`: A maximum number of sessions as `max_size` can return u32::max_value() which is probably more than what your application should use.
105    ///
106    /// # Example
107    /// ```no_run
108    ///  # use r2d2_cryptoki::{*, cryptoki::{context::*, types::AuthPin}};
109    ///  # let pkcs11 = Pkcs11::new("libsofthsm2.so").unwrap();
110    ///  # pkcs11.initialize(CInitializeArgs::new(CInitializeFlags::OS_LOCKING_OK));
111    ///  # let slots = pkcs11.get_slots_with_token().unwrap();
112    ///  # let slot = slots.first().unwrap();
113    ///  # let session_auth = SessionAuth::RwUser(AuthPin::new("fedcba".into()));
114    ///  # let manager = SessionManager::new(pkcs11, *slot, &session_auth);
115    ///  let pool_builder = Pool::builder().connection_customizer(session_auth.into_customizer());
116    ///  let pool_builder = if let Some(max_size) = manager.max_size(100).unwrap() {
117    ///     pool_builder.max_size(max_size)
118    ///  } else {
119    ///     pool_builder
120    ///  };
121    ///  let pool = pool_builder.build(manager).unwrap();
122    /// ```
123    pub fn max_size(&self, maximum: u32) -> Result<Option<u32>, cryptoki::error::Error> {
124        let token_info = self.pkcs11.get_token_info(self.slot)?;
125        let limit = token_info.max_session_count();
126        let res = match limit {
127            Limit::Max(m) => Some(m.try_into().unwrap_or(u32::MAX)),
128            Limit::Unavailable => None,
129            Limit::Infinite => Some(u32::MAX),
130        };
131        Ok(if let Some(true) = res.map(|r| r > maximum) {
132            Some(maximum)
133        } else {
134            res
135        })
136    }
137}
138
139impl ManageConnection for SessionManager {
140    type Connection = Session;
141
142    type Error = cryptoki::error::Error;
143
144    fn connect(&self) -> Result<Self::Connection, Self::Error> {
145        let session = match self.session_state {
146            SessionState::RoPublic | SessionState::RoUser => {
147                self.pkcs11.open_ro_session(self.slot)?
148            }
149            SessionState::RwPublic | SessionState::RwUser | SessionState::RwSecurityOfficer => {
150                self.pkcs11.open_rw_session(self.slot)?
151            }
152        };
153        Ok(session)
154    }
155
156    fn is_valid(&self, session: &mut Self::Connection) -> Result<(), Self::Error> {
157        let actual_state = session.get_session_info()?.session_state();
158        if actual_state != self.session_state {
159            Err(Self::Error::Pkcs11(
160                RvError::UserNotLoggedIn,
161                Function::GetSessionInfo,
162            ))
163        } else {
164            Ok(())
165        }
166    }
167
168    fn has_broken(&self, _session: &mut Self::Connection) -> bool {
169        // TODO find a way to check session state without reaching out to the HSM
170        false
171    }
172}
173
174impl CustomizeConnection<Session, cryptoki::error::Error> for LoginCustomizer {
175    fn on_acquire(&self, session: &mut Session) -> Result<(), cryptoki::error::Error> {
176        let mutex = self.active_sessions.clone();
177        let mut active = mutex.lock().unwrap_or_else(|e| e.into_inner());
178
179        // Login is global, once a session logs in, all sessions are logged in https://stackoverflow.com/a/40225885.
180        if *active == 0 {
181            match session.login(self.user_type, Some(&self.auth_pin)) {
182                // Can happen with poisoned mutex
183                Err(cryptoki::error::Error::Pkcs11(
184                    RvError::UserAlreadyLoggedIn,
185                    Function::Login,
186                )) => {}
187                res => res?,
188            };
189        };
190
191        // Increase after login to prefer login too many over too few
192        *active += 1;
193
194        Ok(())
195    }
196
197    fn on_release(&self, _: Session) {
198        let mutex = self.active_sessions.clone();
199        let mut active = mutex.lock().unwrap_or_else(|e| e.into_inner());
200        if *active > 0 {
201            *active -= 1;
202        }
203    }
204}
205
206#[cfg(test)]
207mod test {
208    use std::{
209        env, fs,
210        path::Path,
211        time::{Duration, Instant},
212    };
213
214    use cached::proc_macro::{cached, once};
215    use cryptoki::{
216        context::{CInitializeArgs, CInitializeFlags},
217        mechanism::Mechanism,
218        object::{Attribute, KeyType, ObjectClass},
219    };
220    use r2d2::PooledConnection;
221
222    use super::*;
223
224    #[derive(Clone, Hash, PartialEq, Eq)]
225    struct Config {
226        max_sessions: Option<u32>,
227        label: Vec<u8>,
228    }
229
230    const DEFAULT_PIN: &str = "abcde";
231
232    // Using cached to create only one pkcs11 ojbect, otherwise it segfaults.
233    #[once(sync_writes = true)]
234    fn default_pkcs11() -> Pkcs11 {
235        env::set_var("SOFTHSM2_CONF", "./test/softhsm2.conf");
236        let tokens_path = Path::new("./test/softhsm/tokens");
237        if tokens_path.exists() {
238            fs::remove_dir_all(tokens_path.to_str().unwrap()).unwrap();
239        }
240        fs::create_dir_all(tokens_path.to_str().unwrap()).unwrap();
241
242        let pkcs11 = Pkcs11::new("libsofthsm2.so").expect("Could not use pkcs11 library");
243        pkcs11
244            .initialize(CInitializeArgs::new(CInitializeFlags::OS_LOCKING_OK))
245            .expect("Could not initialize pkcs11");
246        pkcs11
247    }
248
249    #[cached(sync_writes = "default")]
250    fn default_token(pin: String) -> (Pkcs11, Slot) {
251        let pkcs11 = default_pkcs11();
252        let slot = {
253            let slots = pkcs11
254                .get_slots_with_token()
255                .expect("Could not get slots with token");
256            *slots.first().expect("Could not find a slot")
257        };
258        pkcs11
259            .init_token(slot, &pin.clone().into(), "token")
260            .expect("Could not initialize token");
261        let session = pkcs11.open_rw_session(slot).unwrap();
262        session
263            .login(cryptoki::session::UserType::So, Some(&pin.clone().into()))
264            .unwrap();
265        session.init_pin(&pin.into()).unwrap();
266
267        (pkcs11, slot)
268    }
269
270    /// A token on a slot of its own. Login is global to a token, so sharing
271    /// [default_token] would let the sessions other tests hold open decide whether
272    /// this one is logged in.
273    #[cached(sync_writes = "default")]
274    fn isolated_token(pin: String) -> (Pkcs11, Slot) {
275        // Claiming a slot is a read-modify-write over the slot list.
276        static CLAIM_SLOT: Mutex<()> = Mutex::new(());
277
278        // Initializing the shared token first makes SoftHSM expose a spare slot.
279        let (pkcs11, _) = default_token(DEFAULT_PIN.to_string());
280        let _guard = CLAIM_SLOT.lock().unwrap_or_else(|e| e.into_inner());
281        let initialized = pkcs11
282            .get_slots_with_initialized_token()
283            .expect("Could not get slots with initialized token");
284        let slot = pkcs11
285            .get_slots_with_token()
286            .expect("Could not get slots with token")
287            .into_iter()
288            .find(|slot| !initialized.contains(slot))
289            .expect("Could not find a spare slot to initialize an isolated token on");
290        pkcs11
291            .init_token(slot, &pin.clone().into(), "isolated")
292            .expect("Could not initialize token");
293        let session = pkcs11.open_rw_session(slot).unwrap();
294        session
295            .login(UserType::So, Some(&pin.clone().into()))
296            .unwrap();
297        session.init_pin(&pin.into()).unwrap();
298
299        (pkcs11, slot)
300    }
301
302    fn default_setup(config: Config) -> Pool {
303        let pin_string = DEFAULT_PIN.to_string();
304        let pin = AuthPin::new(pin_string.clone().into());
305        let (pkcs11, slot) = default_token(pin_string);
306
307        let login = SessionAuth::RwUser(pin);
308        let manager = SessionManager::new(pkcs11, slot, &login);
309        let pool_builder = Pool::builder().connection_customizer(login.into_customizer());
310        let pool_builder = if let Some(m) = config.max_sessions {
311            pool_builder.max_size(m)
312        } else {
313            pool_builder
314        };
315        let pool = pool_builder.build(manager).unwrap();
316
317        let mechanism = Mechanism::EccKeyPairGen;
318        let pub_key_template = vec![
319            Attribute::Token(true),
320            Attribute::Private(false),
321            Attribute::Derive(true),
322            Attribute::KeyType(KeyType::EC),
323            Attribute::Verify(true),
324            Attribute::EcParams(vec![
325                0x06, 0x08, 0x2A, 0x86, 0x48, 0xCE, 0x3D, 0x03, 0x01, 0x07,
326            ]),
327            Attribute::Label(config.label.clone()),
328        ];
329        let priv_key_template = vec![
330            Attribute::Token(true),
331            Attribute::Private(false),
332            Attribute::Sensitive(true),
333            Attribute::Extractable(false),
334            Attribute::Derive(true),
335            Attribute::Sign(true),
336            Attribute::Label(config.label),
337        ];
338
339        // sometimes raises an GeneralError
340        backoff::retry(
341            backoff::backoff::Constant::new(Duration::from_millis(25)),
342            || {
343                Ok(pool.get().unwrap().generate_key_pair(
344                    &mechanism,
345                    &pub_key_template,
346                    &priv_key_template,
347                )?)
348            },
349        )
350        .unwrap();
351        pool
352    }
353
354    fn sign(config: &Config, session: &PooledConnection<SessionManager>) -> Vec<u8> {
355        let template = vec![
356            Attribute::Class(ObjectClass::PRIVATE_KEY),
357            Attribute::Label(config.label.clone()),
358        ];
359        let objects = session.find_objects(&template).unwrap();
360        let private = objects.first().unwrap();
361        session
362            .sign(&Mechanism::Ecdsa, *private, "test_data".as_bytes())
363            .unwrap()
364    }
365    fn verify(config: &Config, session: &PooledConnection<SessionManager>, signature: &[u8]) {
366        let template = vec![
367            Attribute::Class(ObjectClass::PUBLIC_KEY),
368            Attribute::Label(config.label.clone()),
369        ];
370        let objects = session.find_objects(&template).unwrap();
371        let public = objects.first().unwrap();
372        session
373            .verify(
374                &Mechanism::Ecdsa,
375                *public,
376                "test_data".as_bytes(),
377                signature,
378            )
379            .unwrap();
380    }
381
382    #[test]
383    fn basic() {
384        let config = Config {
385            max_sessions: None,
386            label: "basic".into(),
387        };
388        let pool = default_setup(config.clone());
389        let sig = sign(&config, &pool.get().unwrap());
390        verify(&config, &pool.get().unwrap(), &sig);
391    }
392
393    fn basic_test(config: &Config, pool1: Pool) {
394        let pool2 = pool1.clone();
395        let config1 = config.clone();
396        let config2 = config.clone();
397        loom::thread::spawn(move || {
398            let sig = sign(&config1, &pool1.get().unwrap());
399            verify(&config1, &pool1.get().unwrap(), &sig);
400        });
401        let sig = sign(&config2, &pool2.get().unwrap());
402        verify(&config2, &pool2.get().unwrap(), &sig);
403    }
404
405    #[test]
406    fn basic_concurrency() {
407        loom::model(|| {
408            let config = Config {
409                max_sessions: None,
410                label: "basic_concurrency".into(),
411            };
412            let pool1 = default_setup(config.clone());
413            basic_test(&config, pool1);
414        });
415    }
416
417    #[test]
418    fn max_one_session() {
419        loom::model(|| {
420            let config = Config {
421                max_sessions: Some(1),
422                label: "max_one_session".into(),
423            };
424            let pool1 = default_setup(config.clone());
425            basic_test(&config, pool1);
426        });
427    }
428
429    fn session_state(session: &Session) -> SessionState {
430        session.get_session_info().unwrap().session_state()
431    }
432
433    /// Shrinking the pool must not disturb the login: `on_release` closes the
434    /// discarded session while the count is still non-zero, so nothing logs back in
435    /// and the surviving session has to keep the token logged in on its own.
436    #[test]
437    fn pool_stays_logged_in_while_shrinking() {
438        let pin_string = "baefc".to_string();
439        let pin = AuthPin::new(pin_string.clone().into());
440        let (pkcs11, slot) = isolated_token(pin_string);
441        let customizer = LoginCustomizer {
442            auth_pin: pin.clone(),
443            user_type: UserType::User,
444            active_sessions: Default::default(),
445        };
446        let active_sessions = customizer.active_sessions.clone();
447        let active = || *active_sessions.lock().unwrap();
448        let manager = SessionManager::new(pkcs11, slot, &SessionAuth::RwUser(pin));
449        let pool = Pool::builder()
450            .max_size(2)
451            // Without this the reaped session is replaced immediately and the count
452            // never actually drops.
453            .min_idle(Some(0))
454            .idle_timeout(Some(Duration::from_millis(1)))
455            .connection_customizer(Box::new(customizer))
456            .connection_timeout(Duration::from_secs(5))
457            .build(manager)
458            .unwrap();
459
460        // Hold one session so the reaper can only take the other.
461        let held = pool.get().unwrap();
462        drop(pool.get().unwrap());
463        assert_eq!(active(), 2);
464
465        // r2d2 fixes its reaper at 30s and does not expose the knob.
466        let deadline = Instant::now() + Duration::from_secs(120);
467        while active() != 1 {
468            assert!(
469                Instant::now() < deadline,
470                "reaper did not discard the idle session"
471            );
472            std::thread::sleep(Duration::from_millis(250));
473        }
474
475        assert_eq!(session_state(&held), SessionState::RwUser);
476
477        // The count is still 1, so this session is established without a login of its
478        // own and depends entirely on `held` having kept the token logged in.
479        let fresh = pool.get().unwrap();
480        assert_eq!(active(), 2);
481        assert_eq!(session_state(&fresh), SessionState::RwUser);
482    }
483
484    /// Login is global to a token, so logging it out invalidates every pooled session
485    /// at once: `is_valid` has to reject them and the pool has to log back in as it
486    /// replaces them.
487    #[test]
488    fn pool_recovers_from_token_logout() {
489        let pin_string = "cbafe".to_string();
490        let pin = AuthPin::new(pin_string.clone().into());
491        let (pkcs11, slot) = isolated_token(pin_string);
492        let login = SessionAuth::RwUser(pin);
493        let manager = SessionManager::new(pkcs11.clone(), slot, &login);
494        let pool = Pool::builder()
495            .max_size(5)
496            .connection_customizer(login.into_customizer())
497            .connection_timeout(Duration::from_secs(5))
498            .build(manager)
499            .unwrap();
500
501        assert_eq!(session_state(&pool.get().unwrap()), SessionState::RwUser);
502
503        let outsider = pkcs11.open_rw_session(slot).unwrap();
504        outsider.logout().unwrap();
505        drop(outsider);
506
507        let session = pool.get().unwrap();
508        assert_eq!(session_state(&session), SessionState::RwUser);
509    }
510
511    /// Returning a [PooledSession] to the pool does not call
512    /// [CustomizeConnection::on_release], and no pool configuration reaches it
513    /// deterministically, so the customizer is driven directly.
514    #[test]
515    fn login_state_across_session_drops() {
516        let pin_string = "fedcb".to_string();
517        let pin = AuthPin::new(pin_string.clone().into());
518        let (pkcs11, slot) = isolated_token(pin_string);
519
520        let customizer = LoginCustomizer {
521            auth_pin: pin,
522            user_type: UserType::User,
523            active_sessions: Default::default(),
524        };
525        let active_sessions = customizer.active_sessions.clone();
526        let active = || *active_sessions.lock().unwrap();
527
528        let mut first = pkcs11.open_rw_session(slot).unwrap();
529        let mut second = pkcs11.open_rw_session(slot).unwrap();
530
531        // The second session reuses the token-wide login.
532        customizer.on_acquire(&mut first).unwrap();
533        assert_eq!(active(), 1);
534        assert_eq!(session_state(&first), SessionState::RwUser);
535        customizer.on_acquire(&mut second).unwrap();
536        assert_eq!(active(), 2);
537        assert_eq!(session_state(&second), SessionState::RwUser);
538
539        customizer.on_release(first);
540        assert_eq!(active(), 1);
541        assert_eq!(session_state(&second), SessionState::RwUser);
542
543        customizer.on_release(second);
544        assert_eq!(active(), 0);
545
546        // With every session gone, the next acquire has to log in again.
547        let mut third = pkcs11.open_rw_session(slot).unwrap();
548        assert_eq!(
549            session_state(&third),
550            SessionState::RwPublic,
551            "token should be logged out once its last session is closed"
552        );
553        customizer.on_acquire(&mut third).unwrap();
554        assert_eq!(active(), 1);
555        assert_eq!(session_state(&third), SessionState::RwUser);
556
557        // If the count drifts below the number of live sessions, `on_acquire` logs in
558        // while the token already is. That must not be an error.
559        *active_sessions.lock().unwrap() = 0;
560        let mut fourth = pkcs11.open_rw_session(slot).unwrap();
561        assert_eq!(session_state(&fourth), SessionState::RwUser);
562        customizer.on_acquire(&mut fourth).unwrap();
563        assert_eq!(active(), 1);
564        assert_eq!(session_state(&fourth), SessionState::RwUser);
565    }
566
567    #[test]
568    fn multiple_operations_per_session() {
569        loom::model(|| {
570            let config = Config {
571                max_sessions: Some(1),
572                label: "multiple_operations_per_session".into(),
573            };
574            let config2 = config.clone();
575            let pool1 = default_setup(config.clone());
576            let pool2 = pool1.clone();
577            loom::thread::spawn(move || {
578                let session = pool1.get().unwrap();
579                let sig = sign(&config, &session);
580                verify(&config, &session, &sig);
581            });
582            let session = pool2.get().unwrap();
583            let sig = sign(&config2, &session);
584            verify(&config2, &session, &sig);
585        });
586    }
587}
r2d2_cryptoki/lib.rs

r2d2_cryptoki/
lib.rs
