Manage C_CloseAllSessions function for multiple callers

Make C_CloseAllSessions work for different callers. Track the sessions
that each caller opens and close just those when C_CloseAllSessiosn is
called.

3 files changed, 268 insertions, 7 deletions

diff --git a/p11-kit/modules.c b/p11-kit/modules.c
index 20475f2..b124e0e 100644
--- a/p11-kit/modules.c
+++ b/p11-kit/modules.c
@@ -170,6 +170,8 @@ typedef struct _Module {
 typedef struct {
 	p11_virtual virt;
 	Module *mod;
+	bool initialized;
+	p11_dict *sessions;
 } Managed;
 
 /*
@@ -1408,13 +1410,31 @@ static CK_RV
 managed_C_Initialize (CK_X_FUNCTION_LIST *self,
                       CK_VOID_PTR init_args)
 {
-	Module *mod = ((Managed *)self)->mod;
+	Managed *managed = ((Managed *)self);
+	p11_dict *sessions;
 	CK_RV rv;
 
 	p11_debug ("in");
 	p11_lock ();
 
-	rv = initialize_module_inlock_reentrant (mod);
+	if (managed->initialized) {
+		rv = CKR_CRYPTOKI_ALREADY_INITIALIZED;
+
+	} else {
+		sessions = p11_dict_new (p11_dict_ulongptr_hash,
+		                         p11_dict_ulongptr_equal,
+		                         free, free);
+		if (!sessions)
+			rv = CKR_HOST_MEMORY;
+		else
+			rv = initialize_module_inlock_reentrant (managed->mod);
+		if (rv == CKR_OK) {
+			managed->sessions = sessions;
+			managed->initialized = true;
+		} else {
+			p11_dict_free (sessions);
+		}
+	}
 
 	p11_unlock ();
 	p11_debug ("out: %lu", rv);
@@ -1423,16 +1443,124 @@ managed_C_Initialize (CK_X_FUNCTION_LIST *self,
 }
 
 static CK_RV
+managed_track_session_inlock (p11_dict *sessions,
+                              CK_SLOT_ID slot_id,
+                              CK_SESSION_HANDLE session)
+{
+	void *key;
+	void *value;
+
+	key = memdup (&session, sizeof (CK_SESSION_HANDLE));
+	return_val_if_fail (key != NULL, CKR_HOST_MEMORY);
+
+	value = memdup (&slot_id, sizeof (CK_SESSION_HANDLE));
+	return_val_if_fail (value != NULL, CKR_HOST_MEMORY);
+
+	if (!p11_dict_set (sessions, key, value))
+		return_val_if_reached (CKR_HOST_MEMORY);
+
+	return CKR_OK;
+}
+
+static void
+managed_untrack_session_inlock (p11_dict *sessions,
+                                CK_SESSION_HANDLE session)
+{
+	p11_dict_remove (sessions, &session);
+}
+
+static CK_SESSION_HANDLE *
+managed_steal_sessions_inlock (p11_dict *sessions,
+                        bool matching_slot_id,
+                        CK_SLOT_ID slot_id,
+                        int *count)
+{
+	CK_SESSION_HANDLE *stolen;
+	CK_SESSION_HANDLE *key;
+	CK_SLOT_ID *value;
+	p11_dictiter iter;
+	int at, i;
+
+	assert (sessions != NULL);
+	assert (count != NULL);
+
+	stolen = calloc (p11_dict_size (sessions), sizeof (CK_SESSION_HANDLE));
+	return_val_if_fail (stolen != NULL, NULL);
+
+	at = 0;
+	p11_dict_iterate (sessions, &iter);
+	while (p11_dict_next (&iter, (void **)&key, (void **)&value)) {
+		if (!matching_slot_id || slot_id == *value)
+			stolen[at++] = *key;
+	}
+
+	/* Removed them all, clear the whole array */
+	if (at == p11_dict_size (sessions)) {
+		p11_dict_clear (sessions);
+
+	/* Only removed some, go through and remove those */
+	} else {
+		for (i = 0; i < at; i++) {
+			if (!p11_dict_remove (sessions, stolen + at))
+				assert_not_reached ();
+		}
+	}
+
+	*count = at;
+	return stolen;
+}
+
+static void
+managed_close_sessions (CK_X_FUNCTION_LIST *funcs,
+                        CK_SESSION_HANDLE *stolen,
+                        int count)
+{
+	CK_RV rv;
+	int i;
+
+	for (i = 0; i < count; i++) {
+		rv = funcs->C_CloseSession (funcs, stolen[i]);
+		if (rv != CKR_OK)
+			p11_message ("couldn't close session: %s", p11_kit_strerror (rv));
+	}
+}
+
+static CK_RV
 managed_C_Finalize (CK_X_FUNCTION_LIST *self,
                     CK_VOID_PTR reserved)
 {
-	Module *mod = ((Managed *)self)->mod;
+	Managed *managed = ((Managed *)self);
+	CK_SESSION_HANDLE *sessions;
+	int count;
 	CK_RV rv;
 
 	p11_debug ("in");
 	p11_lock ();
 
-	rv = finalize_module_inlock_reentrant (mod);
+	if (!managed->initialized) {
+		rv = CKR_CRYPTOKI_NOT_INITIALIZED;
+
+	} else {
+		sessions = managed_steal_sessions_inlock (managed->sessions, false, 0, &count);
+
+		if (sessions && count) {
+			/* WARNING: reentrancy can occur here */
+			p11_unlock ();
+			managed_close_sessions (&managed->mod->virt.funcs, sessions, count);
+			p11_lock ();
+		}
+
+		free (sessions);
+
+		/* WARNING: reentrancy can occur here */
+		rv = finalize_module_inlock_reentrant (managed->mod);
+
+		if (rv == CKR_OK) {
+			managed->initialized = false;
+			p11_dict_free (managed->sessions);
+			managed->sessions = NULL;
+		}
+	}
 
 	p11_unlock ();
 	p11_debug ("out: %lu", rv);
@@ -1440,6 +1568,69 @@ managed_C_Finalize (CK_X_FUNCTION_LIST *self,
 	return rv;
 }
 
+static CK_RV
+managed_C_OpenSession (CK_X_FUNCTION_LIST *self,
+                       CK_SLOT_ID slot_id,
+                       CK_FLAGS flags,
+                       CK_VOID_PTR application,
+                       CK_NOTIFY notify,
+                       CK_SESSION_HANDLE_PTR session)
+{
+	Managed *managed = ((Managed *)self);
+	CK_RV rv;
+
+	return_val_if_fail (session != NULL, CKR_ARGUMENTS_BAD);
+
+	self = &managed->mod->virt.funcs;
+	rv = self->C_OpenSession (self, slot_id, flags, application, notify, session);
+
+	if (rv == CKR_OK) {
+		p11_lock ();
+		rv = managed_track_session_inlock (managed->sessions, slot_id, *session);
+		p11_unlock ();
+	}
+
+	return rv;
+}
+
+static CK_RV
+managed_C_CloseSession (CK_X_FUNCTION_LIST *self,
+                        CK_SESSION_HANDLE session)
+{
+	Managed *managed = ((Managed *)self);
+	CK_RV rv;
+
+	self = &managed->mod->virt.funcs;
+	rv = self->C_CloseSession (self, session);
+
+	if (rv == CKR_OK) {
+		p11_lock ();
+		managed_untrack_session_inlock (managed->sessions, session);
+		p11_unlock ();
+	}
+
+	return rv;
+}
+
+static CK_RV
+managed_C_CloseAllSessions (CK_X_FUNCTION_LIST *self,
+                            CK_SLOT_ID slot_id)
+{
+	Managed *managed = ((Managed *)self);
+	CK_SESSION_HANDLE *stolen;
+	int count;
+
+	p11_lock ();
+	stolen = managed_steal_sessions_inlock (managed->sessions, true, slot_id, &count);
+	p11_unlock ();
+
+	self = &managed->mod->virt.funcs;
+	managed_close_sessions (self, stolen, count);
+	free (stolen);
+
+	return stolen ? CKR_OK : CKR_GENERAL_ERROR;
+}
+
 static void
 managed_free_inlock (void *data)
 {
@@ -1460,6 +1651,9 @@ managed_create_inlock (Module *mod)
 	                  &mod->virt, NULL);
 	managed->virt.funcs.C_Initialize = managed_C_Initialize;
 	managed->virt.funcs.C_Finalize = managed_C_Finalize;
+	managed->virt.funcs.C_CloseAllSessions = managed_C_CloseAllSessions;
+	managed->virt.funcs.C_CloseSession = managed_C_CloseSession;
+	managed->virt.funcs.C_OpenSession = managed_C_OpenSession;
 	managed->mod = mod;
 	mod->ref_count++;
 
diff --git a/p11-kit/tests/test-init.c b/p11-kit/tests/test-init.c
index f2347c8..dbeab46 100644
--- a/p11-kit/tests/test-init.c
+++ b/p11-kit/tests/test-init.c
@@ -152,7 +152,14 @@ test_recursive_initialization (CuTest *tc)
 	p11_unlock ();
 
 	rv = p11_kit_module_initialize (recursive_managed);
-	CuAssertTrue (tc, rv == CKR_FUNCTION_FAILED);
+	CuAssertIntEquals (tc, CKR_FUNCTION_FAILED, rv);
+
+	p11_lock ();
+
+	rv = p11_module_release_inlock_reentrant (recursive_managed);
+	CuAssertTrue (tc, rv == CKR_OK);
+
+	p11_unlock ();
 
 	p11_kit_be_loud ();
 }
diff --git a/p11-kit/tests/test-managed.c b/p11-kit/tests/test-managed.c
index 8a6a1e8..6184c44 100644
--- a/p11-kit/tests/test-managed.c
+++ b/p11-kit/tests/test-managed.c
@@ -105,15 +105,33 @@ test_initialize_finalize (CuTest *tc)
 	CK_FUNCTION_LIST_PTR module;
 	CK_RV rv;
 
-	module = setup_mock_module (tc, NULL);
+	p11_lock ();
+
+	rv = p11_module_load_inlock_reentrant (&mock_module, 0, &module);
+	CuAssertTrue (tc, rv == CKR_OK);
+	CuAssertPtrNotNull (tc, module);
+	CuAssertTrue (tc, p11_virtual_is_wrapper (module));
+
+	p11_unlock ();
 
 	rv = module->C_Initialize (NULL);
 	CuAssertTrue (tc, rv == CKR_OK);
 
+	rv = module->C_Initialize (NULL);
+	CuAssertTrue (tc, rv == CKR_CRYPTOKI_ALREADY_INITIALIZED);
+
 	rv = module->C_Finalize (NULL);
 	CuAssertTrue (tc, rv == CKR_OK);
 
-	teardown_mock_module (tc, module);
+	rv = module->C_Finalize (NULL);
+	CuAssertTrue (tc, rv == CKR_CRYPTOKI_NOT_INITIALIZED);
+
+	p11_lock ();
+
+	rv = p11_module_release_inlock_reentrant (module);
+	CuAssertTrue (tc, rv == CKR_OK);
+
+	p11_unlock ();
 }
 
 static void
@@ -137,6 +155,47 @@ test_initialize_fail (CuTest *tc)
 	CuAssertTrue (tc, rv == CKR_FUNCTION_FAILED);
 }
 
+static void
+test_separate_close_all_sessions (CuTest *tc)
+{
+	CK_FUNCTION_LIST *first;
+	CK_FUNCTION_LIST *second;
+	CK_SESSION_HANDLE s1;
+	CK_SESSION_HANDLE s2;
+	CK_SESSION_INFO info;
+	CK_RV rv;
+
+	first = setup_mock_module (tc, &s1);
+	second = setup_mock_module (tc, &s2);
+
+	rv = first->C_GetSessionInfo (s1, &info);
+	CuAssertTrue (tc, rv == CKR_OK);
+
+	rv = second->C_GetSessionInfo (s2, &info);
+	CuAssertTrue (tc, rv == CKR_OK);
+
+	first->C_CloseAllSessions (MOCK_SLOT_ONE_ID);
+	CuAssertTrue (tc, rv == CKR_OK);
+
+	rv = first->C_GetSessionInfo (s1, &info);
+	CuAssertTrue (tc, rv == CKR_SESSION_HANDLE_INVALID);
+
+	rv = second->C_GetSessionInfo (s2, &info);
+	CuAssertTrue (tc, rv == CKR_OK);
+
+	second->C_CloseAllSessions (MOCK_SLOT_ONE_ID);
+	CuAssertTrue (tc, rv == CKR_OK);
+
+	rv = first->C_GetSessionInfo (s1, &info);
+	CuAssertTrue (tc, rv == CKR_SESSION_HANDLE_INVALID);
+
+	rv = second->C_GetSessionInfo (s2, &info);
+	CuAssertTrue (tc, rv == CKR_SESSION_HANDLE_INVALID);
+
+	teardown_mock_module (tc, first);
+	teardown_mock_module (tc, second);
+}
+
 /* Bring in all the mock module tests */
 #include "test-mock.c"
 
@@ -153,6 +212,7 @@ main (void)
 
 	SUITE_ADD_TEST (suite, test_initialize_finalize);
 	SUITE_ADD_TEST (suite, test_initialize_fail);
+	SUITE_ADD_TEST (suite, test_separate_close_all_sessions);
 	test_mock_add_tests (suite);
 
 	p11_kit_be_quiet ();