Implement testing of DH key pair generation

6 files changed, 368 insertions, 197 deletions

diff --git a/src/Makefile.am b/src/Makefile.am
index e86bd04..c4d34ed 100644
--- a/src/Makefile.am
+++ b/src/Makefile.am
@@ -11,6 +11,7 @@ libp11_tests_la_SOURCES = \
 	config.c \
 	crypto.c \
 	dsa.c \
+	dh.c \
 	key.c \
 	module.c \
 	msg.c \
diff --git a/src/dh.c b/src/dh.c
new file mode 100644
index 0000000..3325caa
--- /dev/null
+++ b/src/dh.c
@@ -0,0 +1,91 @@
+
+#include "p11-tests-lib.h"
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#include <openssl/dh.h>
+
+static void
+test_dh_key_pair_gen (CK_SLOT_ID slot, CK_MECHANISM_TYPE mech, CK_MECHANISM_INFO_PTR info)
+{
+	CK_SESSION_HANDLE session;
+
+	session = p11t_session_open(slot, 0);
+	if(session == CK_INVALID || !p11t_session_login (session))
+		return;
+
+	p11t_dh_test_generate_pair (session);
+}
+
+void
+p11t_dh_tests (void)
+{
+	p11t_slot_for_each_mech (CKM_DH_PKCS_KEY_PAIR_GEN, test_dh_key_pair_gen);
+}
+
+int
+p11t_dh_test_generate_pair (CK_SESSION_HANDLE session)
+{
+	CK_ATTRIBUTE pub_attrs[3];
+	CK_ATTRIBUTE priv_attrs[3];
+	CK_MECHANISM mech;
+	CK_OBJECT_HANDLE pub_key;
+	CK_OBJECT_HANDLE priv_key;
+	CK_ULONG bits;
+	CK_RV rv;
+	DH *dh;
+
+	P11T_SECTION ("C_GenerateKeyPair");
+
+	dh = DH_generate_parameters (256, 2, NULL, NULL);
+	assert (dh);
+
+	pub_attrs[0].type = CKA_PRIME;
+	pub_attrs[0].ulValueLen = BN_num_bytes (dh->p);
+	pub_attrs[0].pValue = alloca (pub_attrs[0].ulValueLen);
+	BN_bn2bin (dh->p, (unsigned char*)pub_attrs[0].pValue);
+
+	pub_attrs[1].type = CKA_BASE;
+	pub_attrs[1].ulValueLen = BN_num_bytes (dh->g);
+	pub_attrs[1].pValue = alloca (pub_attrs[1].ulValueLen);
+	BN_bn2bin (dh->g, (unsigned char*)pub_attrs[1].pValue);
+
+	priv_attrs[0].type = CKA_PRIME;
+	priv_attrs[0].ulValueLen = BN_num_bytes (dh->p);
+	priv_attrs[0].pValue = alloca (priv_attrs[0].ulValueLen);
+	BN_bn2bin (dh->p, (unsigned char*)priv_attrs[0].pValue);
+
+	priv_attrs[1].type = CKA_BASE;
+	priv_attrs[1].ulValueLen = BN_num_bytes (dh->g);
+	priv_attrs[1].pValue = alloca (priv_attrs[1].ulValueLen);
+	BN_bn2bin (dh->g, (unsigned char*)priv_attrs[1].pValue);
+
+	priv_attrs[2].type = CKA_VALUE_BITS;
+	priv_attrs[2].ulValueLen = sizeof (bits);
+	priv_attrs[2].pValue = &bits;
+
+	mech.mechanism = CKM_DH_PKCS_KEY_PAIR_GEN;
+	mech.pParameter = NULL;
+	mech.ulParameterLen = 0;
+
+	if (p11t_test_unexpected) {
+		rv = (p11t_module_funcs->C_GenerateKeyPair) (session, &mech, pub_attrs, 0, priv_attrs, 0, &pub_key, &priv_key);
+		P11T_CHECK_RV ("DH Key Pair without CKA_PRIME", rv, CKR_TEMPLATE_INCOMPLETE);
+
+		bits = 1024;
+		rv = (p11t_module_funcs->C_GenerateKeyPair) (session, &mech, pub_attrs, 2, priv_attrs, 3, &pub_key, &priv_key);
+		P11T_CHECK_RV ("DH Key Pair with CKA_VALUE_BITS larger than CKA_PRIME", rv, CKR_TEMPLATE_INCONSISTENT);
+	}
+
+	bits = 256;
+	rv = (p11t_module_funcs->C_GenerateKeyPair) (session, &mech, pub_attrs, 2, priv_attrs, 3, &pub_key, &priv_key);
+	P11T_CHECK_RV ("DH Key Pair", rv, CKR_OK);
+
+	p11t_key_test (session, pub_key, CKO_PUBLIC_KEY);
+	p11t_key_test (session, priv_key, CKO_PRIVATE_KEY);
+
+	/* Test corner cases */
+	return p11t_object_generate_pair_bad (session, CKM_DH_PKCS_KEY_PAIR_GEN);
+}
diff --git a/src/key.c b/src/key.c
index 0f6aa38..a305215 100644
--- a/src/key.c
+++ b/src/key.c
@@ -14,17 +14,21 @@ test_key_attributes(CK_SLOT_ID slot, CK_SESSION_HANDLE session, CK_OBJECT_HANDLE
 	CK_MECHANISM_TYPE_PTR mechanisms;
 	CK_ULONG n_mechanisms, i;
 	CK_BBOOL bval, local;
+	CK_KEY_TYPE ktype;
 	CK_DATE date;
 	CK_RV rv;
 
 	P11T_SECTION("Key Object");
 
 	attr.type = CKA_KEY_TYPE;
-	attr.ulValueLen = sizeof(*key_type);
-	attr.pValue = key_type;
+	attr.ulValueLen = sizeof(ktype);
+	attr.pValue = &ktype;
 	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
 	P11T_CHECK_RV("CKA_KEY_TYPE", rv, CKR_OK);
 
+	if (key_type)
+		*key_type = ktype;
+
 	attr.type = CKA_ID;
 	attr.ulValueLen = sizeof(buffer);
 	attr.pValue = buffer;
@@ -74,7 +78,7 @@ test_key_attributes(CK_SLOT_ID slot, CK_SESSION_HANDLE session, CK_OBJECT_HANDLE
 		{
 			if(!local)
 				P11T_CHECK_FAIL_MSG("CKA_KEY_GEN_MECHANISM", "Non locally generated key has key gen mechanism");
-			if(!p11t_slot_get_mech_info(slot, mech_type))
+			if(slot && !p11t_slot_get_mech_info(slot, mech_type))
 				P11T_CHECK_FAIL_MSG("CKA_KEY_GEN_MECHANISM", "Mechanism not present on slot");
 		}
 
@@ -91,7 +95,7 @@ test_key_attributes(CK_SLOT_ID slot, CK_SESSION_HANDLE session, CK_OBJECT_HANDLE
 		n_mechanisms = attr.ulValueLen / sizeof(CK_MECHANISM_TYPE);
 		for(i = 0; i < n_mechanisms; ++i)
 		{
-			if(!p11t_slot_get_mech_info(slot, mechanisms[i]))
+			if(slot && !p11t_slot_get_mech_info(slot, mechanisms[i]))
 				P11T_CHECK_FAIL_MSG("CKA_ALLOWED_MECHANISMS", "Mechanism not present in slot");
 		}
 	}
@@ -100,6 +104,168 @@ test_key_attributes(CK_SLOT_ID slot, CK_SESSION_HANDLE session, CK_OBJECT_HANDLE
 }
 
 static int
+test_public_attributes(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE object)
+{
+	CK_BYTE buffer[4096];
+	CK_ATTRIBUTE attr;
+	CK_BBOOL bval, wrap;
+	const char *msg;
+	CK_RV rv;
+
+	P11T_SECTION("CKO_PUBLIC_KEY");
+
+	attr.type = CKA_SUBJECT;
+	attr.pValue = buffer;
+	attr.ulValueLen = sizeof(buffer);
+	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
+	P11T_CHECK_RV("CKA_SUBJECT", rv, CKR_OK);
+
+	if(attr.ulValueLen)
+	{
+		msg = p11t_certificate_validate_dn(attr.pValue, attr.ulValueLen);
+		if(msg != NULL)
+			P11T_CHECK_FAIL_MSG("CKA_SUBJECT", msg);
+	}
+
+	attr.type = CKA_ENCRYPT;
+	attr.pValue = &bval;
+	attr.ulValueLen = sizeof(bval);
+	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
+	P11T_CHECK_RV("CKA_ENCRYPT", rv, CKR_OK);
+	P11T_CHECK_BOOL("CKA_ENCRYPT", bval);
+
+	attr.type = CKA_VERIFY;
+	attr.pValue = &bval;
+	attr.ulValueLen = sizeof(bval);
+	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
+	P11T_CHECK_RV("CKA_VERIFY", rv, CKR_OK);
+	P11T_CHECK_BOOL("CKA_VERIFY", bval);
+
+	attr.type = CKA_VERIFY_RECOVER;
+	attr.pValue = &bval;
+	attr.ulValueLen = sizeof(bval);
+	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
+	P11T_CHECK_RV("CKA_VERIFY_RECOVER", rv, CKR_OK);
+	P11T_CHECK_BOOL("CKA_VERIFY_RECOVER", bval);
+
+	attr.type = CKA_TRUSTED;
+	attr.pValue = &bval;
+	attr.ulValueLen = sizeof(bval);
+	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
+	P11T_CHECK_RV("CKA_TRUSTED", rv, CKR_OK);
+	P11T_CHECK_BOOL("CKA_TRUSTED", bval);
+
+	attr.type = CKA_WRAP;
+	attr.pValue = &wrap;
+	attr.ulValueLen = sizeof(wrap);
+	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
+	P11T_CHECK_RV("CKA_WRAP", rv, CKR_OK);
+	P11T_CHECK_BOOL("CKA_WRAP", wrap);
+
+	return CONTINUE;
+}
+
+static int
+test_private_attributes(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE object)
+{
+	CK_BYTE buffer[4096];
+	CK_ATTRIBUTE attr;
+	CK_BBOOL bval;
+	const char *msg;
+	CK_RV rv;
+
+	P11T_SECTION("CKO_PRIVATE_KEY");
+
+	attr.type = CKA_SUBJECT;
+	attr.pValue = buffer;
+	attr.ulValueLen = sizeof(buffer);
+	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
+	P11T_CHECK_RV("CKA_SUBJECT", rv, CKR_OK);
+
+	if(attr.ulValueLen)
+	{
+		msg = p11t_certificate_validate_dn(attr.pValue, attr.ulValueLen);
+		if(msg != NULL)
+			P11T_CHECK_FAIL_MSG("CKA_SUBJECT", msg);
+	}
+
+	attr.type = CKA_SENSITIVE;
+	attr.pValue = &bval;
+	attr.ulValueLen = sizeof(bval);
+	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
+	P11T_CHECK_RV("CKA_SENSITIVE", rv, CKR_OK);
+	P11T_CHECK_BOOL("CKA_SENSITIVE", bval);
+
+	attr.type = CKA_DECRYPT;
+	attr.pValue = &bval;
+	attr.ulValueLen = sizeof(bval);
+	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
+	P11T_CHECK_RV("CKA_DECRYPT", rv, CKR_OK);
+	P11T_CHECK_BOOL("CKA_DECRYPT", bval);
+
+	attr.type = CKA_SIGN;
+	attr.pValue = &bval;
+	attr.ulValueLen = sizeof(bval);
+	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
+	P11T_CHECK_RV("CKA_SIGN", rv, CKR_OK);
+	P11T_CHECK_BOOL("CKA_SIGN", bval);
+
+	attr.type = CKA_SIGN_RECOVER;
+	attr.pValue = &bval;
+	attr.ulValueLen = sizeof(bval);
+	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
+	P11T_CHECK_RV("CKA_SIGN_RECOVER", rv, CKR_OK);
+	P11T_CHECK_BOOL("CKA_SIGN_RECOVER", bval);
+
+	attr.type = CKA_UNWRAP;
+	attr.pValue = &bval;
+	attr.ulValueLen = sizeof(bval);
+	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
+	P11T_CHECK_RV("CKA_UNWRAP", rv, CKR_OK);
+	P11T_CHECK_BOOL("CKA_UNWRAP", bval);
+
+	attr.type = CKA_EXTRACTABLE;
+	attr.pValue = &bval;
+	attr.ulValueLen = sizeof(bval);
+	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
+	P11T_CHECK_RV("CKA_EXTRACTABLE", rv, CKR_OK);
+	P11T_CHECK_BOOL("CKA_EXTRACTABLE", bval);
+
+	attr.type = CKA_ALWAYS_SENSITIVE;
+	attr.pValue = &bval;
+	attr.ulValueLen = sizeof(bval);
+	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
+	P11T_CHECK_RV("CKA_ALWAYS_SENSITIVE", rv, CKR_OK);
+	P11T_CHECK_BOOL("CKA_ALWAYS_SENSITIVE", bval);
+
+	attr.type = CKA_NEVER_EXTRACTABLE;
+	attr.pValue = &bval;
+	attr.ulValueLen = sizeof(bval);
+	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
+	P11T_CHECK_RV("CKA_NEVER_EXTRACTABLE", rv, CKR_OK);
+	P11T_CHECK_BOOL("CKA_NEVER_EXTRACTABLE", bval);
+
+	if(p11t_test_unexpected)
+	{
+		attr.type = CKA_WRAP_WITH_TRUSTED;
+		attr.pValue = &bval;
+		attr.ulValueLen = sizeof(bval);
+		rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
+		P11T_CHECK_RV("CKA_WRAP_WITH_TRUSTED", rv, CKR_OK);
+		P11T_CHECK_BOOL("CKA_WRAP_WITH_TRUSTED", bval);
+
+		attr.type = CKA_ALWAYS_AUTHENTICATE;
+		attr.pValue = &bval;
+		attr.ulValueLen = sizeof(bval);
+		rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
+		P11T_CHECK_RV("CKA_ALWAYS_AUTHENTICATE", rv, CKR_OK);
+		P11T_CHECK_BOOL("CKA_ALWAYS_AUTHENTICATE", bval);
+	}
+
+	return CONTINUE;
+}
+
+static int
 test_rsa_public(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE object)
 {
 	CK_ATTRIBUTE attr;
@@ -591,8 +757,7 @@ test_rsa_public_create(CK_SESSION_HANDLE session, CK_BBOOL token, RSA *rsa, CK_O
 		return CONTINUE;
 	P11T_CHECK_RV("RSA Public Key", rv, CKR_OK);
 
-	test_rsa_public (session, object);
-
+	p11t_key_test (session, object, CKO_PUBLIC_KEY);
 	*result = object;
 	return CONTINUE;
 }
@@ -655,8 +820,7 @@ test_rsa_private_create(CK_SESSION_HANDLE session, CK_BBOOL token, RSA *rsa, CK_
 		return STOP;
 	P11T_CHECK_RV("RSA Private Key", rv, CKR_OK);
 
-	test_rsa_private (session, object);
-
+	p11t_key_test (session, object, CKO_PRIVATE_KEY);
 	*result = object;
 	return CONTINUE;
 }
@@ -746,8 +910,7 @@ test_dsa_public_create(CK_SESSION_HANDLE session, CK_BBOOL token, DSA *dsa, CK_O
 		return CONTINUE;
 	P11T_CHECK_RV("DSA Public Key", rv, CKR_OK);
 
-	test_dsa_public (session, object);
-
+	p11t_key_test (session, object, CKO_PUBLIC_KEY);
 	*result = object;
 	return CONTINUE;
 }
@@ -805,8 +968,7 @@ test_dsa_private_create(CK_SESSION_HANDLE session, CK_BBOOL token, DSA *dsa, CK_
 		return CONTINUE;
 	P11T_CHECK_RV("DSA Private Key", rv, CKR_OK);
 
-	test_dsa_private (session, object);
-
+	p11t_key_test (session, object, CKO_PRIVATE_KEY);
 	*result = object;
 	return CONTINUE;
 }
@@ -859,8 +1021,7 @@ test_dh_public_create(CK_SESSION_HANDLE session, CK_BBOOL token, DH *dh, CK_OBJE
 		return CONTINUE;
 	P11T_CHECK_RV("DH Public Key", rv, CKR_OK);
 
-	test_dh_public (session, object);
-
+	p11t_key_test (session, object, CKO_PUBLIC_KEY);
 	*result = object;
 	return CONTINUE;
 }
@@ -913,8 +1074,7 @@ test_dh_private_create(CK_SESSION_HANDLE session, CK_BBOOL token, DH *dh, CK_OBJ
 		return CONTINUE;
 	P11T_CHECK_RV("DH Private Key", rv, CKR_OK);
 
-	test_dh_private (session, object);
-
+	p11t_key_test (session, object, CKO_PRIVATE_KEY);
 	*result = object;
 	return CONTINUE;
 }
@@ -1017,163 +1177,46 @@ test_create_unexpected(CK_SESSION_HANDLE session)
 	return CONTINUE;
 }
 
-static int
-test_public_attributes(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE object)
+int
+p11t_key_test (CK_SESSION_HANDLE session, CK_OBJECT_HANDLE object, CK_OBJECT_CLASS klass)
 {
-	CK_BYTE buffer[4096];
+	CK_OBJECT_CLASS key_class;
 	CK_ATTRIBUTE attr;
-	CK_BBOOL bval, wrap;
-	const char *msg;
-	CK_RV rv;
-
-	P11T_SECTION("CKO_PUBLIC_KEY");
-
-	attr.type = CKA_SUBJECT;
-	attr.pValue = buffer;
-	attr.ulValueLen = sizeof(buffer);
-	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
-	P11T_CHECK_RV("CKA_SUBJECT", rv, CKR_OK);
-
-	if(attr.ulValueLen)
-	{
-		msg = p11t_certificate_validate_dn(attr.pValue, attr.ulValueLen);
-		if(msg != NULL)
-			P11T_CHECK_FAIL_MSG("CKA_SUBJECT", msg);
-	}
-
-	attr.type = CKA_ENCRYPT;
-	attr.pValue = &bval;
-	attr.ulValueLen = sizeof(bval);
-	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
-	P11T_CHECK_RV("CKA_ENCRYPT", rv, CKR_OK);
-	P11T_CHECK_BOOL("CKA_ENCRYPT", bval);
-
-	attr.type = CKA_VERIFY;
-	attr.pValue = &bval;
-	attr.ulValueLen = sizeof(bval);
-	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
-	P11T_CHECK_RV("CKA_VERIFY", rv, CKR_OK);
-	P11T_CHECK_BOOL("CKA_VERIFY", bval);
-
-	attr.type = CKA_VERIFY_RECOVER;
-	attr.pValue = &bval;
-	attr.ulValueLen = sizeof(bval);
-	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
-	P11T_CHECK_RV("CKA_VERIFY_RECOVER", rv, CKR_OK);
-	P11T_CHECK_BOOL("CKA_VERIFY_RECOVER", bval);
-
-	attr.type = CKA_TRUSTED;
-	attr.pValue = &bval;
-	attr.ulValueLen = sizeof(bval);
-	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
-	P11T_CHECK_RV("CKA_TRUSTED", rv, CKR_OK);
-	P11T_CHECK_BOOL("CKA_TRUSTED", bval);
-
-	attr.type = CKA_WRAP;
-	attr.pValue = &wrap;
-	attr.ulValueLen = sizeof(wrap);
-	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
-	P11T_CHECK_RV("CKA_WRAP", rv, CKR_OK);
-	P11T_CHECK_BOOL("CKA_WRAP", wrap);
-
-	return CONTINUE;
-}
-
-static int
-test_private_attributes(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE object)
-{
-	CK_BYTE buffer[4096];
-	CK_ATTRIBUTE attr;
-	CK_BBOOL bval;
-	const char *msg;
+	CK_SESSION_INFO info;
+	CK_KEY_TYPE key_type;
 	CK_RV rv;
 
-	P11T_SECTION("CKO_PRIVATE_KEY");
-
-	attr.type = CKA_SUBJECT;
-	attr.pValue = buffer;
-	attr.ulValueLen = sizeof(buffer);
-	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
-	P11T_CHECK_RV("CKA_SUBJECT", rv, CKR_OK);
-
-	if(attr.ulValueLen)
-	{
-		msg = p11t_certificate_validate_dn(attr.pValue, attr.ulValueLen);
-		if(msg != NULL)
-			P11T_CHECK_FAIL_MSG("CKA_SUBJECT", msg);
-	}
-
-	attr.type = CKA_SENSITIVE;
-	attr.pValue = &bval;
-	attr.ulValueLen = sizeof(bval);
-	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
-	P11T_CHECK_RV("CKA_SENSITIVE", rv, CKR_OK);
-	P11T_CHECK_BOOL("CKA_SENSITIVE", bval);
-
-	attr.type = CKA_DECRYPT;
-	attr.pValue = &bval;
-	attr.ulValueLen = sizeof(bval);
-	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
-	P11T_CHECK_RV("CKA_DECRYPT", rv, CKR_OK);
-	P11T_CHECK_BOOL("CKA_DECRYPT", bval);
-
-	attr.type = CKA_SIGN;
-	attr.pValue = &bval;
-	attr.ulValueLen = sizeof(bval);
-	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
-	P11T_CHECK_RV("CKA_SIGN", rv, CKR_OK);
-	P11T_CHECK_BOOL("CKA_SIGN", bval);
-
-	attr.type = CKA_SIGN_RECOVER;
-	attr.pValue = &bval;
-	attr.ulValueLen = sizeof(bval);
-	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
-	P11T_CHECK_RV("CKA_SIGN_RECOVER", rv, CKR_OK);
-	P11T_CHECK_BOOL("CKA_SIGN_RECOVER", bval);
-
-	attr.type = CKA_UNWRAP;
-	attr.pValue = &bval;
-	attr.ulValueLen = sizeof(bval);
-	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
-	P11T_CHECK_RV("CKA_UNWRAP", rv, CKR_OK);
-	P11T_CHECK_BOOL("CKA_UNWRAP", bval);
-
-	attr.type = CKA_EXTRACTABLE;
-	attr.pValue = &bval;
-	attr.ulValueLen = sizeof(bval);
-	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
-	P11T_CHECK_RV("CKA_EXTRACTABLE", rv, CKR_OK);
-	P11T_CHECK_BOOL("CKA_EXTRACTABLE", bval);
-
-	attr.type = CKA_ALWAYS_SENSITIVE;
-	attr.pValue = &bval;
-	attr.ulValueLen = sizeof(bval);
-	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
-	P11T_CHECK_RV("CKA_ALWAYS_SENSITIVE", rv, CKR_OK);
-	P11T_CHECK_BOOL("CKA_ALWAYS_SENSITIVE", bval);
-
-	attr.type = CKA_NEVER_EXTRACTABLE;
-	attr.pValue = &bval;
-	attr.ulValueLen = sizeof(bval);
-	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
-	P11T_CHECK_RV("CKA_NEVER_EXTRACTABLE", rv, CKR_OK);
-	P11T_CHECK_BOOL("CKA_NEVER_EXTRACTABLE", bval);
-
-	if(p11t_test_unexpected)
-	{
-		attr.type = CKA_WRAP_WITH_TRUSTED;
-		attr.pValue = &bval;
-		attr.ulValueLen = sizeof(bval);
-		rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
-		P11T_CHECK_RV("CKA_WRAP_WITH_TRUSTED", rv, CKR_OK);
-		P11T_CHECK_BOOL("CKA_WRAP_WITH_TRUSTED", bval);
+	P11T_SECTION ("Key Object");
+	attr.type = CKA_CLASS;
+	attr.ulValueLen = sizeof (key_class);
+	attr.pValue = &key_class;
+	rv = (p11t_module_funcs->C_GetAttributeValue) (session, object, &attr, 1);
+	P11T_CHECK_RV ("CKA_CLASS", rv, CKR_OK);
+	if (klass && klass != key_class)
+		P11T_CHECK_FAIL_MSG ("CKA_CLASS", "Key class is not as expected");
 
-		attr.type = CKA_ALWAYS_AUTHENTICATE;
-		attr.pValue = &bval;
-		attr.ulValueLen = sizeof(bval);
-		rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, &attr, 1);
-		P11T_CHECK_RV("CKA_ALWAYS_AUTHENTICATE", rv, CKR_OK);
-		P11T_CHECK_BOOL("CKA_ALWAYS_AUTHENTICATE", bval);
+	rv = (p11t_module_funcs->C_GetSessionInfo) (session, &info);
+	if (rv != CKR_OK)
+		info.slotID = 0;
+
+	if (test_key_attributes (info.slotID, session, object, &key_type)) {
+		if (key_class == CKO_PRIVATE_KEY) {
+			test_private_attributes (session, object);
+			if (key_type == CKK_RSA)
+				test_rsa_private (session, object);
+			else if (key_type == CKK_DSA)
+				test_dsa_private (session, object);
+			else if (key_type == CKK_DH)
+				test_dh_private (session, object);
+		} else if (key_class == CKO_PUBLIC_KEY) {
+			test_public_attributes (session, object);
+			if (key_type == CKK_RSA)
+				test_rsa_public (session, object);
+			else if (key_type == CKK_DSA)
+				test_dsa_public (session, object);
+			else if (key_type == CKK_DH)
+				test_dh_public (session, object);
+		}
 	}
 
 	return CONTINUE;
@@ -1185,11 +1228,9 @@ p11t_key_tests(void)
 	CK_OBJECT_CLASS klass;
 	CK_OBJECT_HANDLE_PTR objects;
 	CK_SESSION_HANDLE session;
-	CK_OBJECT_HANDLE object;
 	CK_ATTRIBUTE attrs[1];
 	CK_ULONG j, i, n_objects;
 	CK_SLOT_ID slot;
-	CK_KEY_TYPE key_type;
 
 	attrs[0].type = CKA_CLASS;
 	attrs[0].ulValueLen = sizeof(klass);
@@ -1206,19 +1247,7 @@ p11t_key_tests(void)
 
 		objects = p11t_object_find(session, attrs, 1, &n_objects);
 		for(i = 0; objects && i < n_objects; ++i)
-		{
-			object = objects[i];
-			if(test_key_attributes(slot, session, object, &key_type))
-			{
-				test_public_attributes(session, object);
-				if(key_type == CKK_RSA)
-					test_rsa_public(session, object);
-				else if(key_type == CKK_DSA)
-					test_dsa_public(session, object);
-				else if(key_type == CKK_DH)
-					test_dh_public(session, object);
-			}
-		}
+			p11t_key_test (session, objects[i], klass);
 
 		if(p11t_test_write_session)
 			test_create_unexpected(session);
@@ -1231,19 +1260,7 @@ p11t_key_tests(void)
 
 			objects = p11t_object_find(session, attrs, 1, &n_objects);
 			for(i = 0; objects && i < n_objects; ++i)
-			{
-				object = objects[i];
-				if(test_key_attributes(slot, session, object, &key_type))
-				{
-					test_private_attributes(session, object);
-					if(key_type == CKK_RSA)
-						test_rsa_private(session, object);
-					else if(key_type == CKK_DSA)
-						test_dsa_private(session, object);
-					else if (key_type == CKK_DH)
-						test_dh_private (session, object);
-				}
-			}
+				p11t_key_test (session, objects[i], klass);
 			free(objects);
 
 			if (p11t_test_write_session) {
diff --git a/src/object.c b/src/object.c
index aee8363..4252bbc 100644
--- a/src/object.c
+++ b/src/object.c
@@ -396,6 +396,56 @@ p11t_object_tests(void)
 	}
 }
 
+int
+p11t_object_generate_pair_bad (CK_SESSION_HANDLE session, CK_MECHANISM_TYPE mech)
+{
+	CK_OBJECT_CLASS pub_class = CKO_PUBLIC_KEY;
+	CK_OBJECT_CLASS priv_class = CKO_PRIVATE_KEY;
+	CK_ATTRIBUTE pub_attrs[1];
+	CK_ATTRIBUTE priv_attrs[1];
+	CK_MECHANISM mechanism;
+	CK_OBJECT_HANDLE pub_key, priv_key;
+	CK_RV rv;
+
+	if (!p11t_test_unexpected)
+		return CONTINUE;
+
+	/* An valid mechanism */
+	mechanism.mechanism = mech;
+	mechanism.pParameter = NULL;
+	mechanism.ulParameterLen = 0;
+
+	pub_attrs[0].type = CKA_CLASS;
+	pub_attrs[0].pValue = &pub_class;
+	pub_attrs[0].ulValueLen = sizeof (pub_class);
+
+	priv_attrs[0].type = CKA_CLASS;
+	priv_attrs[0].pValue = &priv_class;
+	priv_attrs[0].ulValueLen = sizeof (priv_class);
+
+	P11T_SECTION ("C_GenerateKeyPair");
+
+	rv = (p11t_module_funcs->C_GenerateKeyPair) ((CK_ULONG)-9, &mechanism, pub_attrs, 1, priv_attrs, 1, &pub_key, &priv_key);
+	P11T_CHECK_RV ("Invalid Session", rv, CKR_SESSION_HANDLE_INVALID);
+
+	rv = (p11t_module_funcs->C_GenerateKeyPair) (session, NULL, pub_attrs, 1, priv_attrs, 1, &pub_key, &priv_key);
+	P11T_CHECK_RV ("Missing mechanism", rv, CKR_ARGUMENTS_BAD);
+
+	rv = (p11t_module_funcs->C_GenerateKeyPair) (session, &mechanism, NULL, 1, priv_attrs, 1, &pub_key, &priv_key);
+	P11T_CHECK_RV ("Missing public attributes", rv, CKR_ARGUMENTS_BAD);
+
+	rv = (p11t_module_funcs->C_GenerateKeyPair) (session, &mechanism, pub_attrs, 1, NULL, 1, &pub_key, &priv_key);
+	P11T_CHECK_RV ("Missing private attributes", rv, CKR_ARGUMENTS_BAD);
+
+	rv = (p11t_module_funcs->C_GenerateKeyPair) (session, &mechanism, pub_attrs, 1, priv_attrs, 1, NULL, &priv_key);
+	P11T_CHECK_RV ("Missing public key handle", rv, CKR_ARGUMENTS_BAD);
+
+	rv = (p11t_module_funcs->C_GenerateKeyPair) (session, &mechanism, pub_attrs, 1, priv_attrs, 1, &pub_key, NULL);
+	P11T_CHECK_RV ("Missing private key handle", rv, CKR_ARGUMENTS_BAD);
+
+	return CONTINUE;
+}
+
 /* ----------------------------------------------------------------------------------
  * METHODS
  */
diff --git a/src/p11-tests-lib.c b/src/p11-tests-lib.c
index 4e51a12..269be11 100644
--- a/src/p11-tests-lib.c
+++ b/src/p11-tests-lib.c
@@ -85,6 +85,7 @@ p11_tests_perform (CK_FUNCTION_LIST_PTR module)
 	p11t_certificate_tests();
 	p11t_rsa_tests();
 	p11t_dsa_tests();
+	p11t_dh_tests ();
 
 	/* Remaining module tests */
 	p11t_module_finalize();
diff --git a/src/p11-tests-lib.h b/src/p11-tests-lib.h
index ac35204..500f75c 100644
--- a/src/p11-tests-lib.h
+++ b/src/p11-tests-lib.h
@@ -132,6 +132,15 @@ void p11t_dsa_tests(void);
 void p11t_dsa_test_public_key(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE handle);
 void p11t_dsa_test_private_key(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE handle);
 
+int p11t_dh_test_generate_pair (CK_SESSION_HANDLE session);
+
+/* -------------------------------------------------------------------
+ * dh.c
+ */
+
+void p11t_dh_tests (void);
+int  p11t_dh_test_generate_pair (CK_SESSION_HANDLE session);
+
 /* -------------------------------------------------------------------
  * key.c
  */
@@ -148,6 +157,7 @@ DSA* p11t_key_export_public_dsa(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE key)
 CK_RV p11t_key_login_context_specific (CK_SESSION_HANDLE session, CK_OBJECT_HANDLE key);
 
 void p11t_key_tests(void);
+int p11t_key_test (CK_SESSION_HANDLE session, CK_OBJECT_HANDLE object, CK_OBJECT_CLASS klass);
 
 /* -------------------------------------------------------------------
  * msg.c
@@ -193,7 +203,8 @@ int p11t_object_get(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE object,
 CK_OBJECT_HANDLE p11t_object_find_one(CK_SESSION_HANDLE session, CK_ATTRIBUTE_PTR attrs,
                                       CK_ULONG n_attrs);
 
-void p11t_object_tests(void);
+void p11t_object_tests (void);
+int p11t_object_generate_pair_bad (CK_SESSION_HANDLE session, CK_MECHANISM_TYPE mech);
 
 /* -------------------------------------------------------------------
  * rsa.c