3 files changed, 295 insertions, 27 deletions

diff --git a/p11-kit/iter.c b/p11-kit/iter.c
index 4caf5d7..0e4ca6e 100644
--- a/p11-kit/iter.c
+++ b/p11-kit/iter.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2013 Red Hat Inc.
+ * Copyright (C) 2013,2016 Red Hat Inc.
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions
@@ -55,7 +55,7 @@ typedef struct _Callback {
 /**
  * P11KitIter:
  *
- * Used to iterate over PKCS\#11 objects.
+ * Used to iterate over PKCS\#11 objects, tokens, slots, and modules.
  */
 struct p11_kit_iter {
 
@@ -82,12 +82,15 @@ struct p11_kit_iter {
 	CK_ULONG saw_objects;
 
 	/* The current iteration */
+	P11KitIterKind kind;
 	CK_FUNCTION_LIST_PTR module;
 	CK_SLOT_ID slot;
 	CK_SESSION_HANDLE session;
 	CK_OBJECT_HANDLE object;
 	CK_SLOT_INFO slot_info;
 	CK_TOKEN_INFO token_info;
+	int move_next_session_state;
+	int iter_next_state;
 
 	/* And various flags */
 	unsigned int searching : 1;
@@ -97,14 +100,33 @@ struct p11_kit_iter {
 	unsigned int keep_session : 1;
 	unsigned int preload_results : 1;
 	unsigned int want_writable : 1;
+	unsigned int with_modules : 1;
+	unsigned int with_slots : 1;
+	unsigned int with_tokens : 1;
+	unsigned int with_objects : 1;
 };
 
 /**
+ * P11KitIterKind:
+ * @P11_KIT_ITER_KIND_MODULE: The iterator is pointing to a module.
+ * @P11_KIT_ITER_KIND_SLOT: The iterator is pointing to a slot.
+ * @P11_KIT_ITER_KIND_TOKEN: The iterator is pointing to a token.
+ * @P11_KIT_ITER_KIND_OBJECT: The iterator is pointing to an object.
+ * @P11_KIT_ITER_KIND_UNKNOWN: The iterator doesn't point to anything.
+ *
+ * The kind of the current match.
+ */
+
+/**
  * P11KitIterBehavior:
  * @P11_KIT_ITER_BUSY_SESSIONS: Allow the iterator's sessions to be
  *   in a busy state when the iterator returns an object.
  * @P11_KIT_ITER_WANT_WRITABLE: Try to open read-write sessions when
- *   iterating over obojects.
+ *   iterating over objects.
+ * @P11_KIT_ITER_WITH_MODULES: Stop at each module while iterating.
+ * @P11_KIT_ITER_WITH_SLOTS: Stop at each slot while iterating.
+ * @P11_KIT_ITER_WITH_TOKENS: Stop at each token while iterating.
+ * @P11_KIT_ITER_WITHOUT_OBJECTS: Ignore objects while iterating.
  *
  * Various flags controlling the behavior of the iterator.
  */
@@ -139,6 +161,10 @@ p11_kit_iter_new (P11KitUri *uri,
 
 	iter->want_writable = !!(behavior & P11_KIT_ITER_WANT_WRITABLE);
 	iter->preload_results = !(behavior & P11_KIT_ITER_BUSY_SESSIONS);
+	iter->with_modules = !!(behavior & P11_KIT_ITER_WITH_MODULES);
+	iter->with_slots = !!(behavior & P11_KIT_ITER_WITH_SLOTS);
+	iter->with_tokens = !!(behavior & P11_KIT_ITER_WITH_TOKENS);
+	iter->with_objects = !(behavior & P11_KIT_ITER_WITHOUT_OBJECTS);
 
 	p11_kit_iter_set_uri (iter, uri);
 	return iter;
@@ -334,6 +360,9 @@ finish_iterating (P11KitIter *iter,
 	p11_array_clear (iter->modules);
 
 	iter->iterating = 0;
+	iter->move_next_session_state = 0;
+	iter->iter_next_state = 0;
+	iter->kind = P11_KIT_ITER_KIND_UNKNOWN;
 	return rv;
 }
 
@@ -470,6 +499,10 @@ call_all_filters (P11KitIter *iter,
 	return CKR_OK;
 }
 
+#define COROUTINE_BEGIN(name) switch (iter->name ## _state) { case 0:
+#define COROUTINE_RETURN(name,i,x) do { iter->name ## _state = i; return x; case i:; } while (0)
+#define COROUTINE_END(name) }
+
 static CK_RV
 move_next_session (P11KitIter *iter)
 {
@@ -478,6 +511,8 @@ move_next_session (P11KitIter *iter)
 	CK_INFO minfo;
 	CK_RV rv;
 
+	COROUTINE_BEGIN (move_next_session);
+
 	finish_slot (iter);
 
 	/* If we have no more slots, then move to next module */
@@ -497,23 +532,31 @@ move_next_session (P11KitIter *iter)
 		if (rv != CKR_OK || !p11_match_uri_module_info (&iter->match_module, &minfo))
 			continue;
 
-		rv = (iter->module->C_GetSlotList) (CK_TRUE, NULL, &num_slots);
-		if (rv != CKR_OK)
-			return finish_iterating (iter, rv);
+		if (iter->with_modules) {
+			iter->kind = P11_KIT_ITER_KIND_MODULE;
+			COROUTINE_RETURN (move_next_session, 1, CKR_OK);
+		}
+
+		if (iter->with_slots || iter->with_tokens || iter->with_objects) {
+			rv = (iter->module->C_GetSlotList) (CK_TRUE, NULL, &num_slots);
+			if (rv != CKR_OK)
+				return finish_iterating (iter, rv);
 
-		iter->slots = realloc (iter->slots, sizeof (CK_SLOT_ID) * (num_slots + 1));
-		return_val_if_fail (iter->slots != NULL, CKR_HOST_MEMORY);
+			iter->slots = realloc (iter->slots, sizeof (CK_SLOT_ID) * (num_slots + 1));
+			return_val_if_fail (iter->slots != NULL, CKR_HOST_MEMORY);
 
-		rv = (iter->module->C_GetSlotList) (CK_TRUE, iter->slots, &num_slots);
-		if (rv != CKR_OK)
-			return finish_iterating (iter, rv);
+			rv = (iter->module->C_GetSlotList) (CK_TRUE, iter->slots, &num_slots);
+			if (rv != CKR_OK)
+				return finish_iterating (iter, rv);
 
-		iter->num_slots = num_slots;
-		assert (iter->saw_slots == 0);
+			iter->num_slots = num_slots;
+			assert (iter->saw_slots == 0);
+		}
 	}
 
 	/* Move to the next slot, and open a session on it */
-	while (iter->saw_slots < iter->num_slots) {
+	while ((iter->with_slots || iter->with_tokens || iter->with_objects) &&
+	       iter->saw_slots < iter->num_slots) {
 		iter->slot = iter->slots[iter->saw_slots++];
 
 		assert (iter->module != NULL);
@@ -522,9 +565,17 @@ move_next_session (P11KitIter *iter)
 		rv = (iter->module->C_GetSlotInfo) (iter->slot, &iter->slot_info);
 		if (rv != CKR_OK || !p11_match_uri_slot_info (&iter->match_slot, &iter->slot_info))
 			continue;
+		if (iter->with_slots) {
+			iter->kind = P11_KIT_ITER_KIND_SLOT;
+			COROUTINE_RETURN (move_next_session, 2, CKR_OK);
+		}
 		rv = (iter->module->C_GetTokenInfo) (iter->slot, &iter->token_info);
 		if (rv != CKR_OK || !p11_match_uri_token_info (&iter->match_token, &iter->token_info))
 			continue;
+		if (iter->with_tokens) {
+			iter->kind = P11_KIT_ITER_KIND_TOKEN;
+			COROUTINE_RETURN (move_next_session, 3, CKR_OK);
+		}
 
 		session_flags = CKF_SERIAL_SESSION;
 
@@ -537,11 +588,17 @@ move_next_session (P11KitIter *iter)
 		if (rv != CKR_OK)
 			return finish_iterating (iter, rv);
 
-		if (iter->session != 0)
+		if (iter->session != 0) {
+			iter->move_next_session_state = 0;
+			iter->kind = P11_KIT_ITER_KIND_UNKNOWN;
 			return CKR_OK;
+		}
 	}
 
+	COROUTINE_END (move_next_session);
+
 	/* Otherwise try again */
+	iter->move_next_session_state = 0;
 	return move_next_session (iter);
 }
 
@@ -572,36 +629,54 @@ p11_kit_iter_next (P11KitIter *iter)
 
 	return_val_if_fail (iter->iterating, CKR_OPERATION_NOT_INITIALIZED);
 
+	COROUTINE_BEGIN (iter_next);
+
 	iter->object = 0;
 
 	if (iter->match_nothing)
 		return finish_iterating (iter, CKR_CANCEL);
 
+	if (!(iter->with_modules || iter->with_slots || iter->with_tokens || iter->with_objects))
+		return finish_iterating (iter, CKR_CANCEL);
+
 	/*
 	 * If we have outstanding objects, then iterate one through those
 	 * Note that we pass each object through the filters, and only
 	 * assume it's iterated if it matches
 	 */
-	while (iter->saw_objects < iter->num_objects) {
+	while (iter->with_objects && iter->saw_objects < iter->num_objects) {
 		iter->object = iter->objects[iter->saw_objects++];
 
 		rv = call_all_filters (iter, &matches);
 		if (rv != CKR_OK)
 			return finish_iterating (iter, rv);
 
-		if (matches)
-			return CKR_OK;
+		if (matches && iter->with_objects) {
+			iter->kind = P11_KIT_ITER_KIND_OBJECT;
+			COROUTINE_RETURN (iter_next, 1, CKR_OK);
+		}
 	}
 
-	/* If we have finished searching then move to next session */
-	if (iter->searched) {
-		rv = move_next_session (iter);
-		if (rv != CKR_OK)
-			return finish_iterating (iter, rv);
+	/* Move to next session, if we have finished searching
+	 * objects, or we are looking for modules/slots/tokens */
+	if ((iter->with_objects && iter->searched) ||
+	    (!iter->with_objects &&
+	     (iter->with_modules || iter->with_slots || iter->with_tokens))) {
+		/* Use iter->kind as the sentinel to detect the case where
+		 * any match (except object) is successful in
+		 * move_next_session() */
+		do {
+			iter->kind = P11_KIT_ITER_KIND_UNKNOWN;
+			rv = move_next_session (iter);
+			if (rv != CKR_OK)
+				return finish_iterating (iter, rv);
+			if (iter->kind != P11_KIT_ITER_KIND_UNKNOWN)
+				COROUTINE_RETURN (iter_next, 2, CKR_OK);
+		} while (iter->move_next_session_state > 0);
 	}
 
 	/* Ready to start searching */
-	if (!iter->searching && !iter->searched) {
+	if (iter->with_objects && !iter->searching && !iter->searched) {
 		count = p11_attrs_count (iter->match_attrs);
 		rv = (iter->module->C_FindObjectsInit) (iter->session, iter->match_attrs, count);
 		if (rv != CKR_OK)
@@ -611,7 +686,7 @@ p11_kit_iter_next (P11KitIter *iter)
 	}
 
 	/* If we have searched on this session then try to continue */
-	if (iter->searching) {
+	if (iter->with_objects && iter->searching) {
 		assert (iter->module != NULL);
 		assert (iter->session != 0);
 		iter->num_objects = 0;
@@ -650,11 +725,35 @@ p11_kit_iter_next (P11KitIter *iter)
 		}
 	}
 
+	COROUTINE_END (iter_next);
+
 	/* Try again */
+	iter->iter_next_state = 0;
+	iter->move_next_session_state = 0;
+	iter->kind = P11_KIT_ITER_KIND_UNKNOWN;
 	return p11_kit_iter_next (iter);
 }
 
 /**
+ * p11_kit_iter_get_kind:
+ * @iter: the iterator
+ *
+ * Get the kind of the current match (a module, slot, token, or an
+ * object).
+ *
+ * This can only be called after p11_kit_iter_next() succeeds.
+ *
+ * Returns: a #P11KitIterKind value
+ */
+P11KitIterKind
+p11_kit_iter_get_kind (P11KitIter *iter)
+{
+	return_val_if_fail (iter != NULL, P11_KIT_ITER_KIND_UNKNOWN);
+	return_val_if_fail (iter->iterating, P11_KIT_ITER_KIND_UNKNOWN);
+	return iter->kind;
+}
+
+/**
  * p11_kit_iter_get_module:
  * @iter: the iterator
  *
diff --git a/p11-kit/iter.h b/p11-kit/iter.h
index 3f51041..fabcd2f 100644
--- a/p11-kit/iter.h
+++ b/p11-kit/iter.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013 Red Hat, Inc
+ * Copyright (c) 2013,2016 Red Hat, Inc
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions
@@ -49,8 +49,20 @@ typedef struct p11_kit_iter P11KitIter;
 typedef P11KitIter p11_kit_iter;
 
 typedef enum {
+	P11_KIT_ITER_KIND_MODULE,
+	P11_KIT_ITER_KIND_SLOT,
+	P11_KIT_ITER_KIND_TOKEN,
+	P11_KIT_ITER_KIND_OBJECT,
+	P11_KIT_ITER_KIND_UNKNOWN = -1,
+} P11KitIterKind;
+
+typedef enum {
 	P11_KIT_ITER_BUSY_SESSIONS = 1 << 1,
 	P11_KIT_ITER_WANT_WRITABLE = 1 << 2,
+	P11_KIT_ITER_WITH_MODULES = 1 << 3,
+	P11_KIT_ITER_WITH_SLOTS = 1 << 4,
+	P11_KIT_ITER_WITH_TOKENS = 1 << 5,
+	P11_KIT_ITER_WITHOUT_OBJECTS = 1 << 6,
 } P11KitIterBehavior;
 
 typedef CK_RV      (* p11_kit_iter_callback)                (P11KitIter *iter,
@@ -84,6 +96,8 @@ void                  p11_kit_iter_begin_with               (P11KitIter *iter,
 
 CK_RV                 p11_kit_iter_next                     (P11KitIter *iter);
 
+P11KitIterKind        p11_kit_iter_get_kind                 (P11KitIter *iter);
+
 CK_FUNCTION_LIST_PTR  p11_kit_iter_get_module               (P11KitIter *iter);
 
 CK_SLOT_ID            p11_kit_iter_get_slot                 (P11KitIter *iter);
diff --git a/p11-kit/test-iter.c b/p11-kit/test-iter.c
index 4ce480a..ac31fbe 100644
--- a/p11-kit/test-iter.c
+++ b/p11-kit/test-iter.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013, Red Hat Inc.
+ * Copyright (c) 2013,2016 Red Hat Inc.
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions
@@ -49,6 +49,8 @@
 #include <stdio.h>
 #include <stdlib.h>
 
+#define ELEMS(x) (sizeof (x) / sizeof (x[0]))
+
 static CK_FUNCTION_LIST_PTR_PTR
 initialize_and_get_modules (void)
 {
@@ -697,6 +699,44 @@ test_module_mismatch (void)
 }
 
 static void
+test_module_only (void)
+{
+	CK_FUNCTION_LIST_PTR *modules;
+	P11KitIter *iter;
+	P11KitUri *uri;
+	CK_RV rv;
+	int count;
+	int ret;
+
+	modules = initialize_and_get_modules ();
+
+	uri = p11_kit_uri_new ();
+	ret = p11_kit_uri_parse ("pkcs11:library-description=MOCK%20LIBRARY", P11_KIT_URI_FOR_MODULE, uri);
+	assert_num_eq (P11_KIT_URI_OK, ret);
+
+	iter = p11_kit_iter_new (uri, P11_KIT_ITER_WITH_MODULES | P11_KIT_ITER_WITHOUT_OBJECTS);
+	p11_kit_uri_free (uri);
+
+	p11_kit_iter_begin (iter, modules);
+
+	count = 0;
+	while ((rv = p11_kit_iter_next (iter)) == CKR_OK) {
+		P11KitIterKind kind = p11_kit_iter_get_kind (iter);
+		assert_num_eq (P11_KIT_ITER_KIND_MODULE, kind);
+		count++;
+	}
+
+	assert (rv == CKR_CANCEL);
+
+	/* Three modules, each with 1 slot, and 3 public objects */
+	assert_num_eq (3, count);
+
+	p11_kit_iter_free (iter);
+
+	finalize_and_free_modules (modules);
+}
+
+static void
 test_slot_match (void)
 {
 	CK_FUNCTION_LIST_PTR *modules;
@@ -767,6 +807,44 @@ test_slot_mismatch (void)
 }
 
 static void
+test_slot_only (void)
+{
+	CK_FUNCTION_LIST_PTR *modules;
+	P11KitIter *iter;
+	P11KitUri *uri;
+	CK_RV rv;
+	int count;
+	int ret;
+
+	modules = initialize_and_get_modules ();
+
+	uri = p11_kit_uri_new ();
+	ret = p11_kit_uri_parse ("pkcs11:slot-manufacturer=TEST%20MANUFACTURER", P11_KIT_URI_FOR_SLOT, uri);
+	assert_num_eq (P11_KIT_URI_OK, ret);
+
+	iter = p11_kit_iter_new (uri, P11_KIT_ITER_WITH_SLOTS | P11_KIT_ITER_WITHOUT_OBJECTS);
+	p11_kit_uri_free (uri);
+
+	p11_kit_iter_begin (iter, modules);
+
+	count = 0;
+	while ((rv = p11_kit_iter_next (iter)) == CKR_OK) {
+		P11KitIterKind kind = p11_kit_iter_get_kind (iter);
+		assert_num_eq (P11_KIT_ITER_KIND_SLOT, kind);
+		count++;
+	}
+
+	assert (rv == CKR_CANCEL);
+
+	/* Three modules, each with 1 slot, and 3 public objects */
+	assert_num_eq (3, count);
+
+	p11_kit_iter_free (iter);
+
+	finalize_and_free_modules (modules);
+}
+
+static void
 test_slot_match_by_id (void)
 {
 	CK_FUNCTION_LIST_PTR *modules;
@@ -943,6 +1021,44 @@ test_token_mismatch (void)
 }
 
 static void
+test_token_only (void)
+{
+	CK_FUNCTION_LIST_PTR *modules;
+	P11KitIter *iter;
+	P11KitUri *uri;
+	CK_RV rv;
+	int count;
+	int ret;
+
+	modules = initialize_and_get_modules ();
+
+	uri = p11_kit_uri_new ();
+	ret = p11_kit_uri_parse ("pkcs11:manufacturer=TEST%20MANUFACTURER", P11_KIT_URI_FOR_TOKEN, uri);
+	assert_num_eq (P11_KIT_URI_OK, ret);
+
+	iter = p11_kit_iter_new (uri, P11_KIT_ITER_WITH_TOKENS | P11_KIT_ITER_WITHOUT_OBJECTS);
+	p11_kit_uri_free (uri);
+
+	p11_kit_iter_begin (iter, modules);
+
+	count = 0;
+	while ((rv = p11_kit_iter_next (iter)) == CKR_OK) {
+		P11KitIterKind kind = p11_kit_iter_get_kind (iter);
+		assert_num_eq (P11_KIT_ITER_KIND_TOKEN, kind);
+		count++;
+	}
+
+	assert (rv == CKR_CANCEL);
+
+	/* Three modules, each with 1 slot, and 3 public objects */
+	assert_num_eq (3, count);
+
+	p11_kit_iter_free (iter);
+
+	finalize_and_free_modules (modules);
+}
+
+static void
 test_token_info (void)
 {
 	CK_FUNCTION_LIST_PTR *modules;
@@ -1464,6 +1580,41 @@ test_destroy_object (void)
 	finalize_and_free_modules (modules);
 }
 
+/* Test all combinations of P11_KIT_ITER_WITH_{TOKENS,SLOTS,MODULES}
+ * and P11_KIT_ITER_WITHOUT_OBJECTS, against three modules, each
+ * with 1 slot, and 3 public objects */
+static void
+test_exhaustive_match (void)
+{
+	CK_FUNCTION_LIST_PTR *modules;
+	P11KitIter *iter;
+	CK_RV rv;
+	int counts[] = {
+		9, 12, 12, 15, 12, 15, 15, 18, 0, 3, 3, 6, 3, 6, 6, 9
+	};
+	int count;
+	int i;
+
+	for (i = 0; i < ELEMS (counts); i++) {
+		modules = initialize_and_get_modules ();
+
+		iter = p11_kit_iter_new (NULL, (P11KitIterBehavior) i << 3);
+		p11_kit_iter_begin (iter, modules);
+
+		count = 0;
+		while ((rv = p11_kit_iter_next (iter)) == CKR_OK)
+			count++;
+
+		assert (rv == CKR_CANCEL);
+
+		assert_num_eq (counts[i], count);
+
+		p11_kit_iter_free (iter);
+
+		finalize_and_free_modules (modules);
+	}
+}
+
 int
 main (int argc,
       char *argv[])
@@ -1487,13 +1638,16 @@ main (int argc,
 	p11_test (test_token_match, "/iter/test_token_match");
 	p11_test (test_token_mismatch, "/iter/test_token_mismatch");
 	p11_test (test_token_info, "/iter/token-info");
+	p11_test (test_token_only, "/iter/test_token_only");
 	p11_test (test_slot_match, "/iter/test_slot_match");
 	p11_test (test_slot_mismatch, "/iter/test_slot_mismatch");
 	p11_test (test_slot_match_by_id, "/iter/test_slot_match_by_id");
 	p11_test (test_slot_mismatch_by_id, "/iter/test_slot_mismatch_by_id");
 	p11_test (test_slot_info, "/iter/slot-info");
+	p11_test (test_slot_only, "/iter/test_slot_only");
 	p11_test (test_module_match, "/iter/test_module_match");
 	p11_test (test_module_mismatch, "/iter/test_module_mismatch");
+	p11_test (test_module_only, "/iter/test_module_only");
 	p11_test (test_getslotlist_fail_first, "/iter/test_getslotlist_fail_first");
 	p11_test (test_getslotlist_fail_late, "/iter/test_getslotlist_fail_late");
 	p11_test (test_open_session_fail, "/iter/test_open_session_fail");
@@ -1507,6 +1661,7 @@ main (int argc,
 	p11_testx (test_many, "", "/iter/test-many");
 	p11_testx (test_many, "busy-sessions", "/iter/test-many-busy");
 	p11_test (test_destroy_object, "/iter/destroy-object");
+	p11_test (test_exhaustive_match, "/iter/test_exhaustive_match");
 
 	return p11_test_run (argc, argv);
 }