Skip to content

Use ConnectionCustomizer to ensure login is only done when needed #4

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sbihel merged 3 commits into from
Sep 8, 2025
Merged

Use ConnectionCustomizer to ensure login is only done when needed #4

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3b70ea5
Simplify token_info.max_session_count()
middagj Jul 11, 2025
617a966
Use ConnectionCustomizer to ensure login is only done when needed
middagj Jul 11, 2025
426be94
Fix README example
sbihel Sep 8, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 9 additions & 2 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,11 @@

Session pool manager for [cryptoki](https://github.com/parallaxsecond/rust-cryptoki/).

Cryptoki has a single login state for all sessions.
Only when all sessions are closed, a login is needed again.
This library requires a `ConnectionCustomizer` on the pool to ensure login is only done when needed.
The `SessionAuth` can be converted into the appropriate `ConnectionCustomizer`.

## Example

```rust no_run
Expand All @@ -14,9 +19,11 @@ let pkcs11 = Pkcs11::new("libsofthsm2.so").unwrap();
pkcs11.initialize(CInitializeArgs::OsThreads).unwrap();
let slots = pkcs11.get_slots_with_token().unwrap();
let slot = slots.first().unwrap();
let manager = SessionManager::new(pkcs11, *slot, SessionType::RwUser(AuthPin::new("fedcba".to_string())));
let session_auth = SessionAuth::RwUser(AuthPin::new("fedcba".to_string()));
let manager = SessionManager::new(pkcs11, *slot, &session_auth);

let pool = r2d2::Pool::builder().build(manager).unwrap();
let pool_builder = Pool::builder().connection_customizer(session_auth.into_customizer());
let pool = pool_builder.build(manager).unwrap();

let session = pool.get().unwrap();
println!("{:?}", session.get_session_info().unwrap());
Expand Down
112 changes: 77 additions & 35 deletions src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
#![warn(missing_docs)]
#![doc = include_str!("../README.md")]

use std::sync::{Arc, Mutex};

pub use cryptoki;
pub use r2d2;

Expand All @@ -11,7 +13,7 @@ use cryptoki::{
slot::{Limit, Slot},
types::AuthPin,
};
use r2d2::ManageConnection;
use r2d2::{CustomizeConnection, ManageConnection, NopConnectionCustomizer};

/// Alias for this crate's instance of r2d2's Pool
pub type Pool = r2d2::Pool<SessionManager>;
Expand All @@ -23,12 +25,12 @@ pub type PooledSession = r2d2::PooledConnection<SessionManager>;
pub struct SessionManager {
pkcs11: Pkcs11,
slot: Slot,
session_type: SessionType,
session_state: SessionState,
}

/// Session types, holding the pin for the authenticated sessions
#[derive(Debug, Clone)]
pub enum SessionType {
pub enum SessionAuth {
/// [SessionState::RoPublic]
RoPublic,
/// [SessionState::RoUser]
Expand All @@ -41,7 +43,15 @@ pub enum SessionType {
RwSecurityOfficer(AuthPin),
}

impl SessionType {
/// Mandatory connection customizer for logins
#[derive(Debug, Clone)]
struct LoginCustomizer {
auth_pin: AuthPin,
user_type: UserType,
active_sessions: Arc<Mutex<u32>>,
}

impl SessionAuth {
fn as_state(&self) -> SessionState {
match self {
Self::RoPublic => SessionState::RoPublic,
Expand All @@ -51,6 +61,23 @@ impl SessionType {
Self::RwSecurityOfficer(_) => SessionState::RwSecurityOfficer,
}
}

/// Returns the correct customizer to use for the specified session auth
pub fn into_customizer(self) -> Box<dyn CustomizeConnection<Session, cryptoki::error::Error>> {
match self {
Self::RoPublic | Self::RwPublic => Box::new(NopConnectionCustomizer),
Self::RoUser(auth_pin) | Self::RwUser(auth_pin) => Box::from(LoginCustomizer {
auth_pin,
user_type: UserType::User,
active_sessions: Default::default(),
}),
Self::RwSecurityOfficer(auth_pin) => Box::from(LoginCustomizer {
auth_pin,
user_type: UserType::So,
active_sessions: Default::default(),
}),
}
}
}

impl SessionManager {
Expand All @@ -61,13 +88,13 @@ impl SessionManager {
/// pkcs11 .initialize(CInitializeArgs::OsThreads).unwrap();
/// let slots = pkcs11.get_slots_with_token().unwrap();
/// let slot = slots.first().unwrap();
/// let manager = SessionManager::new(pkcs11, *slot, SessionType::RwUser(AuthPin::new("abcd".to_string())));
/// let manager = SessionManager::new(pkcs11, *slot, &SessionAuth::RwUser(AuthPin::new("abcd".to_string())));
/// ```
pub fn new(pkcs11: Pkcs11, slot: Slot, session_type: SessionType) -> Self {
pub fn new(pkcs11: Pkcs11, slot: Slot, session_auth: &SessionAuth) -> Self {
Self {
pkcs11,
slot,
session_type,
session_state: session_auth.as_state(),
}
}

Expand All @@ -83,8 +110,9 @@ impl SessionManager {
/// # pkcs11.initialize(CInitializeArgs::OsThreads);
/// # let slots = pkcs11.get_slots_with_token().unwrap();
/// # let slot = slots.first().unwrap();
/// # let manager = SessionManager::new(pkcs11, *slot, SessionType::RwUser(AuthPin::new("fedcba".to_string())));
/// let pool_builder = r2d2::Pool::builder();
/// # let session_auth = SessionAuth::RwUser(AuthPin::new("fedcba".to_string()));
/// # let manager = SessionManager::new(pkcs11, *slot, &session_auth);
/// let pool_builder = Pool::builder().connection_customizer(session_auth.into_customizer());
/// let pool_builder = if let Some(max_size) = manager.max_size(100).unwrap() {
/// pool_builder.max_size(max_size)
/// } else {
Expand All @@ -94,12 +122,7 @@ impl SessionManager {
/// ```
pub fn max_size(&self, maximum: u32) -> Result<Option<u32>, cryptoki::error::Error> {
let token_info = self.pkcs11.get_token_info(self.slot)?;
let limit = match self.session_type {
SessionType::RoPublic | SessionType::RoUser(_) => token_info.max_session_count(),
SessionType::RwPublic | SessionType::RwUser(_) | SessionType::RwSecurityOfficer(_) => {
token_info.max_session_count()
}
};
let limit = token_info.max_session_count();
let res = match limit {
Limit::Max(m) => Some(m.try_into().unwrap_or(u32::MAX)),
Limit::Unavailable => None,
Expand All @@ -118,35 +141,21 @@ impl ManageConnection for SessionManager {

type Error = cryptoki::error::Error;

// Login is global, once a session logs in, all sessions are logged in https://stackoverflow.com/a/40225885.
// TODO cryptoki automatically logs out on Drop, so this is ineficient and we will need to find a better way to check the login state when we start having a pool of sessions
fn connect(&self) -> Result<Self::Connection, Self::Error> {
let session = match self.session_type {
SessionType::RoPublic | SessionType::RoUser(_) => {
let session = match self.session_state {
SessionState::RoPublic | SessionState::RoUser => {
self.pkcs11.open_ro_session(self.slot)?
}
SessionType::RwPublic | SessionType::RwUser(_) | SessionType::RwSecurityOfficer(_) => {
SessionState::RwPublic | SessionState::RwUser | SessionState::RwSecurityOfficer => {
self.pkcs11.open_rw_session(self.slot)?
}
};
let maybe_user_info = match &self.session_type {
SessionType::RoPublic | SessionType::RwPublic => None,
SessionType::RoUser(pin) | SessionType::RwUser(pin) => Some((UserType::User, pin)),
SessionType::RwSecurityOfficer(pin) => Some((UserType::So, pin)),
};
if let Some(user_type) = maybe_user_info {
match session.login(user_type.0, Some(user_type.1)) {
Err(Self::Error::Pkcs11(RvError::UserAlreadyLoggedIn, Function::Login)) => {}
res => res?,
};
}
Ok(session)
}

fn is_valid(&self, session: &mut Self::Connection) -> Result<(), Self::Error> {
let actual_state = session.get_session_info()?.session_state();
let expected_state = &self.session_type;
if actual_state != expected_state.as_state() {
if actual_state != self.session_state {
Err(Self::Error::Pkcs11(
RvError::UserNotLoggedIn,
Function::GetSessionInfo,
Expand All @@ -162,6 +171,38 @@ impl ManageConnection for SessionManager {
}
}

impl CustomizeConnection<Session, cryptoki::error::Error> for LoginCustomizer {
fn on_acquire(&self, session: &mut Session) -> Result<(), cryptoki::error::Error> {
let mutex = self.active_sessions.clone();
let mut active = mutex.lock().unwrap_or_else(|e| e.into_inner());

// Login is global, once a session logs in, all sessions are logged in https://stackoverflow.com/a/40225885.
if *active == 0 {
match session.login(self.user_type, Some(&self.auth_pin)) {
// Can happen with poisoned mutex
Err(cryptoki::error::Error::Pkcs11(
RvError::UserAlreadyLoggedIn,
Function::Login,
)) => {}
res => res?,
};
};

// Increase after login to prefer login too many over too few
*active += 1;

Ok(())
}

fn on_release(&self, _: Session) {
let mutex = self.active_sessions.clone();
let mut active = mutex.lock().unwrap_or_else(|e| e.into_inner());
if *active > 0 {
*active -= 1;
}
}
}

#[cfg(test)]
mod test {
use std::{env, fs, path::Path, time::Duration};
Expand Down Expand Up @@ -225,8 +266,9 @@ mod test {
let pin = AuthPin::new(pin_string.clone());
let (pkcs11, slot) = default_token(pin_string);

let manager = SessionManager::new(pkcs11, slot, SessionType::RwUser(pin));
let pool_builder = r2d2::Pool::builder();
let login = SessionAuth::RwUser(pin);
let manager = SessionManager::new(pkcs11, slot, &login);
let pool_builder = Pool::builder().connection_customizer(login.into_customizer());
let pool_builder = if let Some(m) = config.max_sessions {
pool_builder.max_size(m)
} else {
Expand Down