#include "p11-tests.h"
#include "p11-tests-lib.h"

#include <string.h>

static CK_OBJECT_HANDLE
find_non_key_object (CK_SESSION_HANDLE session)
{
	CK_OBJECT_HANDLE_PTR objects;
	CK_OBJECT_HANDLE result = CK_INVALID;
	CK_ATTRIBUTE attr;
	CK_OBJECT_CLASS klass;
	CK_ULONG n_objects, i;

	objects = p11t_object_find (session, NULL, 0, &n_objects);
	for (i = 0; i < n_objects; ++i) {
		attr.type = CKA_CLASS;
		attr.ulValueLen = sizeof (klass);
		attr.pValue = &klass;

		if (p11t_object_get (session, objects[i], &attr, 1) &&
		    attr.ulValueLen == sizeof (klass) &&
		    klass != CKO_PUBLIC_KEY && klass != CKO_PRIVATE_KEY) {
			result = objects[i];
			break;
		}
	}

	free (objects);
	return result;
}

static CK_OBJECT_HANDLE
find_non_method_key (CK_SESSION_HANDLE session, CK_MECHANISM_TYPE mech,
                     CK_ATTRIBUTE_TYPE method)
{
	CK_BYTE buffer[P11T_BLOCK];
	CK_OBJECT_HANDLE_PTR objects;
	CK_OBJECT_HANDLE result = CK_INVALID;
	CK_MECHANISM_TYPE_PTR mechanisms;
	CK_ATTRIBUTE attrs[2];
	CK_BBOOL can_method;
	CK_ULONG n_objects, i;
	CK_ULONG n_mechanisms, j;

	/* All objects */
	objects = p11t_object_find (session, NULL, 0, &n_objects);
	for (i = 0; i < n_objects && result == CK_INVALID; ++i) {

		attrs[0].type = method;
		attrs[0].pValue = &can_method;
		attrs[0].ulValueLen = sizeof (can_method);
		attrs[1].type = CKA_ALLOWED_MECHANISMS;
		attrs[1].pValue = buffer;
		attrs[1].ulValueLen = sizeof (buffer);

		if (p11t_object_get (session, objects[i], attrs, 2) &&
		    attrs[1].ulValueLen != CK_INVALID) {

			/* Skip over ones that can do this method */
			if (attrs[0].ulValueLen == sizeof (can_method) && can_method)
				continue;

			/* Find one that can do this mechanism */
			mechanisms = (CK_MECHANISM_TYPE_PTR)buffer;
			n_mechanisms = attrs[1].ulValueLen / sizeof (CK_MECHANISM_TYPE);
			for (j = 0; j < n_mechanisms; ++j) {
				if (mech == mechanisms[j]) {
					result = objects[i];
					break;
				}
			}
		}
	}

	free (objects);
	return result;
}

static CK_OBJECT_HANDLE
find_non_mechanism_key (CK_SESSION_HANDLE session, CK_MECHANISM_TYPE mech,
                        CK_ATTRIBUTE_TYPE method)
{
	CK_BYTE buffer[P11T_BLOCK];
	CK_OBJECT_HANDLE_PTR objects;
	CK_OBJECT_HANDLE result = CK_INVALID;
	CK_MECHANISM_TYPE_PTR mechanisms;
	CK_ATTRIBUTE attr;
	CK_BBOOL val_true = CK_TRUE;
	CK_ULONG n_objects, i;
	CK_ULONG n_mechanisms, j;
	CK_BBOOL has;

	attr.type = method;
	attr.ulValueLen = sizeof (val_true);
	attr.pValue = &val_true;

	objects = p11t_object_find (session, &attr, 1, &n_objects);
	for (i = 0; i < n_objects; ++i) {

		/* Only allow keys that don't match our mechanism */
		attr.type = CKA_ALLOWED_MECHANISMS;
		attr.pValue = buffer;
		attr.ulValueLen = sizeof (buffer);

		if (p11t_object_get (session, objects[i], &attr, 1) &&
		    attr.ulValueLen != CK_INVALID) {

			has = CK_FALSE;

			mechanisms = (CK_MECHANISM_TYPE_PTR)buffer;
			n_mechanisms = attr.ulValueLen / sizeof (CK_MECHANISM_TYPE);
			for (j = 0; j < n_mechanisms; ++j) {
				if (mech == mechanisms[j]) {
					has = CK_TRUE;
					break;
				}
			}

			if (has == CK_FALSE) {
				result = objects[i];
				break;
			}
		}
	}

	free (objects);
	return result;
}

static CK_OBJECT_HANDLE
find_applicable_key (CK_SESSION_HANDLE session, CK_MECHANISM_TYPE mech,
                     CK_ATTRIBUTE_TYPE method)
{
	CK_BYTE buffer[P11T_BLOCK];
	CK_OBJECT_HANDLE_PTR objects;
	CK_OBJECT_HANDLE result = CK_INVALID;
	CK_MECHANISM_TYPE_PTR mechanisms;
	CK_ATTRIBUTE attr;
	CK_BBOOL val_true = CK_TRUE;
	CK_ULONG n_objects, i;
	CK_ULONG n_mechanisms, j;

	attr.type = method;
	attr.ulValueLen = sizeof (val_true);
	attr.pValue = &val_true;

	objects = p11t_object_find (session, &attr, 1, &n_objects);
	for (i = 0; i < n_objects && result == CK_INVALID; ++i) {

		/* Only allow keys that don't match our mechanism */
		attr.type = CKA_ALLOWED_MECHANISMS;
		attr.pValue = buffer;
		attr.ulValueLen = sizeof (buffer);

		if (p11t_object_get (session, objects[i], &attr, 1) &&
		    attr.ulValueLen != CK_INVALID) {
			mechanisms = (CK_MECHANISM_TYPE_PTR)buffer;
			n_mechanisms = attr.ulValueLen / sizeof (CK_MECHANISM_TYPE);
			for (j = 0; j < n_mechanisms; ++j) {
				if (mech == mechanisms[j]) {
					result = objects[i];
					break;
				}
			}
		}
	}

	free (objects);
	return result;
}


static int
test_without_init (CK_SESSION_HANDLE session, CK_ATTRIBUTE_TYPE method)
{
	CK_BYTE_PTR input;
	CK_ULONG n_input;
	CK_ULONG n_output;
	CK_RV rv;

	/*
	 * No operation should be valid, therefore input is neither
	 * valid or invalid, and doesn't have to conform to a
	 * mechanism, since none is initialized.
	 *
	 * We just send some in to prevent CKR_ARGUMENTS_BAD
	 */
	input = (CK_BYTE_PTR)p11t_test_data;
	n_input = 10;

	P11T_SECTION ("C_Encrypt");

	if (method != CKA_ENCRYPT) {
		rv = (p11t_module_funcs->C_Encrypt) (session, input, n_input, NULL, &n_output);
		P11T_CHECK_RV ("Without calling C_EncyrptInit", rv, CKR_OPERATION_NOT_INITIALIZED);
	}

	P11T_SECTION ("C_Decrypt");

	if (method != CKA_DECRYPT) {
		rv = (p11t_module_funcs->C_Decrypt) (session, input, n_input, NULL, &n_output);
		P11T_CHECK_RV ("Without calling C_DecryptInit", rv, CKR_OPERATION_NOT_INITIALIZED);
	}

	P11T_SECTION ("C_Sign");

	if (method != CKA_SIGN) {
		rv = (p11t_module_funcs->C_Sign) (session, input, n_input, NULL, &n_output);
		P11T_CHECK_RV ("Without calling C_SignInit", rv, CKR_OPERATION_NOT_INITIALIZED);
	}

	P11T_SECTION ("C_SignRecover");

	if (method != CKA_SIGN) {
		rv = (p11t_module_funcs->C_Sign) (session, input, n_input, NULL, &n_output);
		P11T_CHECK_RV ("Without calling C_SignRecoverInit", rv, CKR_OPERATION_NOT_INITIALIZED);
	}

	P11T_SECTION ("C_Verify");

	if (method != CKA_VERIFY) {
		rv = (p11t_module_funcs->C_Verify) (session, input, n_input, input, n_input);
		P11T_CHECK_RV ("Without calling C_VerifyInit", rv, CKR_OPERATION_NOT_INITIALIZED);
	}

	P11T_SECTION ("C_VerifyRecover");

	if (method != CKA_SIGN) {
		rv = (p11t_module_funcs->C_Sign) (session, input, n_input, NULL, &n_output);
		P11T_CHECK_RV ("Without calling C_VerifyRecoverInit", rv, CKR_OPERATION_NOT_INITIALIZED);
	}

	return CONTINUE;
}

static int
prepare_raw_data (CK_SESSION_HANDLE session, CK_MECHANISM_TYPE mech, CK_ATTRIBUTE_TYPE method,
                  CK_OBJECT_HANDLE handle, CK_BYTE_PTR data, CK_ULONG_PTR n_data)
{
	/*
	 * Obviously this data isn't valid, but should be
	 * indistinguishable from possibly valid data.
	 */

	if (mech == CKM_RSA_PKCS) {

		/* 11 bytes shorter than key is valid */
		if (method == CKA_SIGN || method == CKA_ENCRYPT) {
			memcpy (data, p11t_test_data, 20);
			*n_data = 20;
			return CONTINUE;
		}

	} else if (mech == CKM_RSA_X_509) {

		/* Any data shorter than the key length is valid */
		if (method == CKA_SIGN || method == CKA_ENCRYPT) {
			memcpy (data, p11t_test_data, 36);
			*n_data = 36;
			return CONTINUE;
		}
	}  else if (mech == CKM_DSA) {

		/* Input data must be 20 bytes */
		if (method == CKA_SIGN) {
			memcpy (data, p11t_test_data, 20);
			*n_data = 20;
			return CONTINUE;
		}
	}

	/* Add more mechanisms here */
	return STOP;
}

static int
prepare_raw_invalid_len (CK_SESSION_HANDLE session, CK_MECHANISM_TYPE mech, CK_ATTRIBUTE_TYPE method,
                         CK_OBJECT_HANDLE handle, CK_BYTE_PTR data, CK_ULONG_PTR n_data)
{
	/* This should be too long to sign for any current mechanism */
	memcpy (data, p11t_test_data, P11T_BLOCK);
	*n_data = P11T_BLOCK;
	return CONTINUE;
}

static int
prepare_raw_invalid (CK_SESSION_HANDLE session, CK_MECHANISM_TYPE mech, CK_ATTRIBUTE_TYPE method,
                     CK_OBJECT_HANDLE handle, CK_BYTE_PTR data, CK_ULONG_PTR n_data)
{
	/* Add more mechanisms here */
	return STOP;
}

/* --------------------------------------------------------------------------------
 * ENCRYPT
 */

static CK_RV
run_encrypt_login_if_necessary (CK_SESSION_HANDLE session, CK_BYTE_PTR input, CK_ULONG n_input,
                                CK_BYTE_PTR output, CK_ULONG_PTR n_output)
{
	CK_RV rv;

	rv = (p11t_module_funcs->C_Encrypt) (session, input, n_input, output, n_output);
	if (rv == CKR_USER_NOT_LOGGED_IN) {
		rv = p11t_key_login_context_specific (session, CK_INVALID);
		if (rv == CKR_OK)
			rv = (p11t_module_funcs->C_Encrypt) (session, input, n_input, output, n_output);
	}

	return rv;
}

int
p11t_crypto_test_encrypt (CK_SESSION_HANDLE session, CK_MECHANISM_TYPE mech)
{
	CK_OBJECT_HANDLE handle;
	CK_MECHANISM mechanism;
	CK_BYTE input[P11T_BLOCK];
	CK_BYTE invalid[P11T_BLOCK];
	CK_BYTE output[P11T_BLOCK];
	CK_ULONG n_input, n_invalid, n_output, count;
	CK_RV rv;

	if (!p11t_test_unexpected)
		return CONTINUE;

	/* An valid mechanism */
	mechanism.mechanism = mech;
	mechanism.pParameter = NULL;
	mechanism.ulParameterLen = 0;

	test_without_init (session, 0);

	P11T_SECTION ("C_EncryptInit");

	rv = (p11t_module_funcs->C_EncryptInit)(session, &mechanism, (CK_ULONG)-9);
	P11T_CHECK_RV ("Invalid handle", rv, CKR_OBJECT_HANDLE_INVALID);

	handle = find_non_key_object (session);
	if (handle != CK_INVALID) {
		rv = (p11t_module_funcs->C_EncryptInit) (session, &mechanism, handle);
		P11T_CHECK_RV ("Invalid object type", rv, CKR_KEY_HANDLE_INVALID);
	}

	handle = find_non_method_key (session, mech, CKA_ENCRYPT);
	if (handle != CK_INVALID) {
		rv = (p11t_module_funcs->C_EncryptInit) (session, &mechanism, handle);
		P11T_CHECK_RV ("Key that cannot encrypt", rv, CKR_KEY_FUNCTION_NOT_PERMITTED);
	}

	handle = find_non_mechanism_key (session, mech, CKA_ENCRYPT);
	if (handle != CK_INVALID) {
		rv = (p11t_module_funcs->C_EncryptInit) (session, &mechanism, handle);
		P11T_CHECK_RV ("Key from wrong mechanism", rv, CKR_KEY_TYPE_INCONSISTENT);
	}

	/* Finally find a valid object */
	handle = find_applicable_key (session, mech, CKA_ENCRYPT);
	if (handle == CK_INVALID)
		return CONTINUE;

	rv = (p11t_module_funcs->C_EncryptInit)((CK_ULONG)-5, &mechanism, handle);
	P11T_CHECK_RV ("Invalid session", rv, CKR_SESSION_HANDLE_INVALID);

	rv = (p11t_module_funcs->C_EncryptInit)(session, NULL, handle);
	P11T_CHECK_RV ("Null mechanism", rv, CKR_ARGUMENTS_BAD);

	/* The real initialize here */

	rv = (p11t_module_funcs->C_EncryptInit)(session, &mechanism, handle);
	P11T_CHECK_RV ("A normal call", rv, CKR_OK);


	/* All other crypto operations shouldn't be valid */
	test_without_init (session, CKA_ENCRYPT);


	/* Prepare some input to encrypt */
	n_input = P11T_BLOCK;
	if (!prepare_raw_data (session, mech, CKA_ENCRYPT, handle, input, &n_input))
		return CONTINUE;

	P11T_SECTION ("C_Encrypt");

	/* All the invalid operations first */

	rv = run_encrypt_login_if_necessary ((CK_ULONG)-5, input, n_input, NULL, &n_output);
	P11T_CHECK_RV ("Invalid session", rv, CKR_SESSION_HANDLE_INVALID);

	/*
	 * This does not terminate the current operation, since we
	 * didn't give it a valid session. The below invalid tests should
	 * terminate the current operation.
	 */

	rv = run_encrypt_login_if_necessary (session, NULL, n_input, NULL, &n_output);
	P11T_CHECK_RV ("Null input", rv, CKR_ARGUMENTS_BAD);

	rv = (p11t_module_funcs->C_EncryptInit) (session, &mechanism, handle);
	P11T_CHECK_RV ("Error result terminates operation", rv, CKR_OK);

	rv = run_encrypt_login_if_necessary (session, input, n_input, NULL, NULL);
	P11T_CHECK_RV ("Null output length", rv, CKR_ARGUMENTS_BAD);

	rv = (p11t_module_funcs->C_EncryptInit) (session, &mechanism, handle);
	P11T_CHECK_RV ("Error result terminates operation", rv, CKR_OK);

	/* Now finally a valid call, albeit without output input */
	rv = run_encrypt_login_if_necessary (session, input, n_input, NULL, &count);
	P11T_CHECK_RV ("Call without output buffer", rv, CKR_OK);

	/* This should not have cancelled the operation */

	n_output = 2;
	rv = run_encrypt_login_if_necessary (session, input, n_input, output, &n_output);
	P11T_CHECK_RV ("Call with too small output buffer", rv, CKR_BUFFER_TOO_SMALL);

	/* Again this should not have cancelled the operation */

	/* Finally do the actual encrypting */
	n_output = count;
	rv = run_encrypt_login_if_necessary (session, input, n_input, output, &n_output);
	P11T_CHECK_RV ("Normal call with output buffer", rv, CKR_OK);

	/* Some more invalid checks */
	n_invalid = P11T_BLOCK;
	if (prepare_raw_invalid_len (session, mech, CKA_ENCRYPT, handle, invalid, &n_invalid)) {
		rv = (p11t_module_funcs->C_EncryptInit) (session, &mechanism, handle);
		P11T_CHECK_RV ("Error result terminates operation", rv, CKR_OK);

		n_output = count;
		rv = run_encrypt_login_if_necessary (session, invalid, n_invalid, output, &n_output);
		P11T_CHECK_RV ("Bad input length", rv, CKR_DATA_LEN_RANGE);
	}

	n_invalid = P11T_BLOCK;
	if (prepare_raw_invalid (session, mech, CKA_ENCRYPT, handle, invalid, &n_invalid)) {
		rv = (p11t_module_funcs->C_EncryptInit) (session, &mechanism, handle);
		P11T_CHECK_RV ("Error result terminates operation", rv, CKR_OK);

		n_output = count;
		rv = run_encrypt_login_if_necessary (session, invalid, n_invalid, output, &n_output);
		P11T_CHECK_RV ("Invalid input", rv, CKR_DATA_INVALID);
	}

	/* At this point no crypto operations should be valid */

	test_without_init (session, 0);

	return CONTINUE;
}

/* -----------------------------------------------------------------------------------------------
 * DECRYPT
 */

static int
prepare_decrypt_data (CK_SESSION_HANDLE session, CK_MECHANISM_TYPE mech,
                      CK_OBJECT_HANDLE handle, CK_BYTE_PTR data, CK_ULONG_PTR n_data)
{
	CK_BYTE raw[P11T_BLOCK];
	CK_MECHANISM mechanism;
	CK_ULONG n_raw;
	CK_OBJECT_HANDLE pubkey;
	CK_RV rv;

	mechanism.mechanism = mech;
	mechanism.pParameter = NULL;
	mechanism.ulParameterLen = 0;

	pubkey = p11t_key_get_public (session, handle);
	if (pubkey == CK_INVALID)
		return STOP;

	/* Send for some encryptable data */
	if (!prepare_raw_data (session, mech, CKA_ENCRYPT, handle, raw, &n_raw))
		return STOP;

	/* And encrypt it */
	rv = (p11t_module_funcs->C_EncryptInit) (session, &mechanism, pubkey);
	if (rv != CKR_OK)
		return STOP;

	rv = run_encrypt_login_if_necessary (session, raw, n_raw, data, n_data);
	if (rv != CKR_OK)
		return STOP;

	return CONTINUE;
}

static CK_RV
run_decrypt_login_if_necessary (CK_SESSION_HANDLE session, CK_BYTE_PTR input, CK_ULONG n_input,
                                CK_BYTE_PTR output, CK_ULONG_PTR n_output)
{
	CK_RV rv;

	rv = (p11t_module_funcs->C_Decrypt) (session, input, n_input, output, n_output);
	if (rv == CKR_USER_NOT_LOGGED_IN) {
		rv = p11t_key_login_context_specific (session, CK_INVALID);
		if (rv == CKR_OK)
			rv = (p11t_module_funcs->C_Decrypt) (session, input, n_input, output, n_output);
	}

	return rv;
}

int
p11t_crypto_test_decrypt (CK_SESSION_HANDLE session, CK_MECHANISM_TYPE mech)
{
	CK_OBJECT_HANDLE handle;
	CK_MECHANISM mechanism;
	CK_BYTE input[P11T_BLOCK];
	CK_BYTE invalid[P11T_BLOCK];
	CK_BYTE output[P11T_BLOCK];
	CK_ULONG n_input, n_invalid, n_output, count;
	CK_RV rv;

	if (!p11t_test_unexpected)
		return CONTINUE;

	/* An valid mechanism */
	mechanism.mechanism = mech;
	mechanism.pParameter = NULL;
	mechanism.ulParameterLen = 0;

	test_without_init (session, 0);

	P11T_SECTION ("C_DecryptInit");

	rv = (p11t_module_funcs->C_DecryptInit)(session, &mechanism, (CK_ULONG)-9);
	P11T_CHECK_RV ("Invalid handle", rv, CKR_OBJECT_HANDLE_INVALID);

	handle = find_non_key_object (session);
	if (handle != CK_INVALID) {
		rv = (p11t_module_funcs->C_DecryptInit) (session, &mechanism, handle);
		P11T_CHECK_RV ("Invalid object type", rv, CKR_KEY_HANDLE_INVALID);
	}

	handle = find_non_method_key (session, mech, CKA_DECRYPT);
	if (handle != CK_INVALID) {
		rv = (p11t_module_funcs->C_DecryptInit) (session, &mechanism, handle);
		P11T_CHECK_RV ("Key that cannot decrypt", rv, CKR_KEY_FUNCTION_NOT_PERMITTED);
	}

	handle = find_non_mechanism_key (session, mech, CKA_DECRYPT);
	if (handle != CK_INVALID) {
		rv = (p11t_module_funcs->C_DecryptInit) (session, &mechanism, handle);
		P11T_CHECK_RV ("Key from wrong mechanism", rv, CKR_KEY_TYPE_INCONSISTENT);
	}

	/* Finally find a valid object */
	handle = find_applicable_key (session, mech, CKA_DECRYPT);
	if (handle == CK_INVALID)
		return CONTINUE;

	/* Prepare some input to decrypt */
	n_input = P11T_BLOCK;
	if (!prepare_decrypt_data (session, mech, handle, input, &n_input))
		return CONTINUE;

	rv = (p11t_module_funcs->C_DecryptInit)((CK_ULONG)-5, &mechanism, handle);
	P11T_CHECK_RV ("Invalid session", rv, CKR_SESSION_HANDLE_INVALID);

	rv = (p11t_module_funcs->C_DecryptInit)(session, NULL, handle);
	P11T_CHECK_RV ("Null mechanism", rv, CKR_ARGUMENTS_BAD);

	/* The real initialize here */

	rv = (p11t_module_funcs->C_DecryptInit)(session, &mechanism, handle);
	P11T_CHECK_RV ("A normal call", rv, CKR_OK);


	/* All other crypto operations shouldn't be valid */
	test_without_init (session, CKA_DECRYPT);


	P11T_SECTION ("C_Decrypt");

	/* All the invalid operations first */

	rv = run_decrypt_login_if_necessary ((CK_ULONG)-5, input, n_input, NULL, &n_output);
	P11T_CHECK_RV ("Invalid session", rv, CKR_SESSION_HANDLE_INVALID);

	/*
	 * This does not terminate the current operation, since we
	 * didn't give it a valid session. The below invalid tests should
	 * terminate the current operation.
	 */

	rv = run_decrypt_login_if_necessary (session, NULL, n_input, NULL, &n_output);
	P11T_CHECK_RV ("Null input", rv, CKR_ARGUMENTS_BAD);

	rv = (p11t_module_funcs->C_DecryptInit) (session, &mechanism, handle);
	P11T_CHECK_RV ("Error result terminates operation", rv, CKR_OK);

	rv = run_decrypt_login_if_necessary (session, input, n_input, NULL, NULL);
	P11T_CHECK_RV ("Null output length", rv, CKR_ARGUMENTS_BAD);

	rv = (p11t_module_funcs->C_DecryptInit) (session, &mechanism, handle);
	P11T_CHECK_RV ("Error result terminates operation", rv, CKR_OK);

	/* Now finally a valid call, albeit without output input */
	rv = run_decrypt_login_if_necessary (session, input, n_input, NULL, &count);
	P11T_CHECK_RV ("Call without output buffer", rv, CKR_OK);

	/* This should not have cancelled the operation */

	n_output = 2;
	rv = run_decrypt_login_if_necessary (session, input, n_input, output, &n_output);
	P11T_CHECK_RV ("Call with too small output buffer", rv, CKR_BUFFER_TOO_SMALL);

	/* Again this should not have cancelled the operation */

	/* Finally do the actual decrypting */
	n_output = count;
	rv = run_decrypt_login_if_necessary (session, input, n_input, output, &n_output);
	P11T_CHECK_RV ("Normal call with output buffer", rv, CKR_OK);

	/* Some more invalid checks */
	n_invalid = P11T_BLOCK;
	if (prepare_raw_invalid_len (session, mech, CKA_DECRYPT, handle, invalid, &n_invalid)) {
		rv = (p11t_module_funcs->C_DecryptInit) (session, &mechanism, handle);
		P11T_CHECK_RV ("Error result terminates operation", rv, CKR_OK);

		n_output = count;
		rv = run_decrypt_login_if_necessary (session, invalid, n_invalid, output, &n_output);
		P11T_CHECK_RV ("Bad input length", rv, CKR_DATA_LEN_RANGE);
	}

	n_invalid = P11T_BLOCK;
	if (prepare_raw_invalid (session, mech, CKA_DECRYPT, handle, invalid, &n_invalid)) {
		rv = (p11t_module_funcs->C_DecryptInit) (session, &mechanism, handle);
		P11T_CHECK_RV ("Error result terminates operation", rv, CKR_OK);

		n_output = count;
		rv = run_decrypt_login_if_necessary (session, invalid, n_invalid, output, &n_output);
		P11T_CHECK_RV ("Invalid input", rv, CKR_DATA_INVALID);
	}

	/* At this point no crypto operations should be valid */

	test_without_init (session, 0);

	return CONTINUE;
}

/* -----------------------------------------------------------------------------------------------
 * SIGNING
 */

static CK_RV
run_sign_login_if_necessary (CK_SESSION_HANDLE session, CK_BYTE_PTR input, CK_ULONG n_input,
                             CK_BYTE_PTR output, CK_ULONG_PTR n_output)
{
	CK_RV rv;

	rv = (p11t_module_funcs->C_Sign) (session, input, n_input, output, n_output);
	if (rv == CKR_USER_NOT_LOGGED_IN) {
		rv = p11t_key_login_context_specific (session, CK_INVALID);
		if (rv == CKR_OK)
			rv = (p11t_module_funcs->C_Sign) (session, input, n_input, output, n_output);
	}

	return rv;
}

int
p11t_crypto_test_sign (CK_SESSION_HANDLE session, CK_MECHANISM_TYPE mech)
{
	CK_OBJECT_HANDLE handle;
	CK_MECHANISM mechanism;
	CK_BYTE input[P11T_BLOCK];
	CK_BYTE invalid[P11T_BLOCK];
	CK_BYTE output[P11T_BLOCK];
	CK_ULONG n_input, n_invalid, n_output, count;
	CK_RV rv;

	if (!p11t_test_unexpected)
		return CONTINUE;

	/* An valid mechanism */
	mechanism.mechanism = mech;
	mechanism.pParameter = NULL;
	mechanism.ulParameterLen = 0;

	test_without_init (session, 0);

	P11T_SECTION ("C_SignInit");

	rv = (p11t_module_funcs->C_SignInit)(session, &mechanism, (CK_ULONG)-9);
	P11T_CHECK_RV ("Invalid handle", rv, CKR_OBJECT_HANDLE_INVALID);

	handle = find_non_key_object (session);
	if (handle != CK_INVALID) {
		rv = (p11t_module_funcs->C_SignInit) (session, &mechanism, handle);
		P11T_CHECK_RV ("Invalid object type", rv, CKR_KEY_HANDLE_INVALID);
	}

	handle = find_non_method_key (session, mech, CKA_SIGN);
	if (handle != CK_INVALID) {
		rv = (p11t_module_funcs->C_SignInit) (session, &mechanism, handle);
		P11T_CHECK_RV ("Key that's not signable", rv, CKR_KEY_FUNCTION_NOT_PERMITTED);
	}

	handle = find_non_mechanism_key (session, mech, CKA_SIGN);
	if (handle != CK_INVALID) {
		rv = (p11t_module_funcs->C_SignInit) (session, &mechanism, handle);
		P11T_CHECK_RV ("Key from wrong mechanism", rv, CKR_KEY_TYPE_INCONSISTENT);
	}

	/* Finally find a valid object */
	handle = find_applicable_key (session, mech, CKA_SIGN);
	if (handle == CK_INVALID)
		return CONTINUE;

	/* Prepare some input to sign */
	n_input = P11T_BLOCK;
	if (!prepare_raw_data (session, mech, CKA_SIGN, handle, input, &n_input))
		return CONTINUE;


	rv = (p11t_module_funcs->C_SignInit)((CK_ULONG)-5, &mechanism, handle);
	P11T_CHECK_RV ("Invalid session", rv, CKR_SESSION_HANDLE_INVALID);

	rv = (p11t_module_funcs->C_SignInit)(session, NULL, handle);
	P11T_CHECK_RV ("Null mechanism", rv, CKR_ARGUMENTS_BAD);

	/* The real initialize here */

	rv = (p11t_module_funcs->C_SignInit)(session, &mechanism, handle);
	P11T_CHECK_RV ("A normal call", rv, CKR_OK);


	/* All other crypto operations shouldn't be valid */
	test_without_init (session, CKA_SIGN);


	P11T_SECTION ("C_Sign");

	/* All the invalid signs first */

	rv = run_sign_login_if_necessary ((CK_ULONG)-5, input, n_input, NULL, &n_output);
	P11T_CHECK_RV ("Invalid session", rv, CKR_SESSION_HANDLE_INVALID);

	/*
	 * This does not terminate the current operation, since we
	 * didn't give it a valid session. The below invalid tests should
	 * terminate the current operation.
	 */

	rv = run_sign_login_if_necessary (session, NULL, n_input, NULL, &n_output);
	P11T_CHECK_RV ("Null input", rv, CKR_ARGUMENTS_BAD);

	rv = (p11t_module_funcs->C_SignInit) (session, &mechanism, handle);
	P11T_CHECK_RV ("Error result terminates operation", rv, CKR_OK);

	rv = run_sign_login_if_necessary (session, input, n_input, NULL, NULL);
	P11T_CHECK_RV ("Null output output length", rv, CKR_ARGUMENTS_BAD);

	rv = (p11t_module_funcs->C_SignInit) (session, &mechanism, handle);
	P11T_CHECK_RV ("Error result terminates operation", rv, CKR_OK);

	/* Now finally a valid call, albeit without output input */
	rv = run_sign_login_if_necessary (session, input, n_input, NULL, &count);
	P11T_CHECK_RV ("Call without output buffer", rv, CKR_OK);

	/* This should not have cancelled the operation */

	n_output = 2;
	rv = run_sign_login_if_necessary (session, input, n_input, output, &n_output);
	P11T_CHECK_RV ("Call with too small output buffer", rv, CKR_BUFFER_TOO_SMALL);

	/* Again this should not have cancelled the operation */

	/* Finally do the actual signing */
	n_output = count;
	rv = run_sign_login_if_necessary (session, input, n_input, output, &n_output);
	P11T_CHECK_RV ("Normal call with output buffer", rv, CKR_OK);

	/* Some more invalid checks */
	n_invalid = P11T_BLOCK;
	if (prepare_raw_invalid_len (session, mech, CKA_SIGN, handle, invalid, &n_invalid)) {
		rv = (p11t_module_funcs->C_SignInit) (session, &mechanism, handle);
		P11T_CHECK_RV ("Error result terminates operation", rv, CKR_OK);

		n_output = count;
		rv = run_sign_login_if_necessary (session, invalid, n_invalid, output, &n_output);
		P11T_CHECK_RV ("Bad input length", rv, CKR_DATA_LEN_RANGE);
	}

	n_invalid = P11T_BLOCK;
	if (prepare_raw_invalid (session, mech, CKA_SIGN, handle, invalid, &n_invalid)) {
		rv = (p11t_module_funcs->C_SignInit) (session, &mechanism, handle);
		P11T_CHECK_RV ("Error result terminates operation", rv, CKR_OK);

		n_output = count;
		rv = run_sign_login_if_necessary (session, invalid, n_invalid, output, &n_output);
		P11T_CHECK_RV ("Invalid input", rv, CKR_DATA_INVALID);
	}

	/* At this point no crypto operations should be valid */

	test_without_init (session, 0);

	return CONTINUE;
}


/* -----------------------------------------------------------------------------------------------
 * VERIFYING
 */

static int
prepare_verify_data (CK_SESSION_HANDLE session, CK_MECHANISM_TYPE mech,
                     CK_OBJECT_HANDLE handle, CK_BYTE_PTR input, CK_ULONG n_input,
                     CK_BYTE_PTR signature, CK_ULONG_PTR n_signature)
{
	CK_MECHANISM mechanism;
	CK_OBJECT_HANDLE privkey;
	CK_RV rv;

	mechanism.mechanism = mech;
	mechanism.pParameter = NULL;
	mechanism.ulParameterLen = 0;

	privkey = p11t_key_get_private (session, handle);
	if (privkey == CK_INVALID)
		return STOP;

	/* And sign it */
	rv = (p11t_module_funcs->C_SignInit) (session, &mechanism, privkey);
	if (rv != CKR_OK)
		return STOP;

	rv = run_sign_login_if_necessary (session, input, n_input, signature, n_signature);
	if (rv != CKR_OK)
		return STOP;

	return CONTINUE;
}

static int
prepare_verify_invalid_len (CK_SESSION_HANDLE session, CK_MECHANISM_TYPE mech,
                            CK_OBJECT_HANDLE handle, CK_ULONG_PTR n_signature)
{
	if (mech == CKM_DSA) {

		/* Signature must be 40 bytes */
		*n_signature = 5;
		return CONTINUE;

	} else if (mech == CKM_RSA_PKCS) {

		/* Signature too long */
		*n_signature = P11T_BLOCK;
		return CONTINUE;
	} else if (mech == CKM_RSA_X_509) {

		/* Signature too long */
		*n_signature = P11T_BLOCK;
		return CONTINUE;
	}

	return STOP;
}


static CK_RV
run_verify_login_if_necessary (CK_SESSION_HANDLE session, CK_BYTE_PTR input, CK_ULONG n_input,
                              CK_BYTE_PTR output, CK_ULONG n_output)
{
	CK_RV rv;

	rv = (p11t_module_funcs->C_Verify) (session, input, n_input, output, n_output);
	if (rv == CKR_USER_NOT_LOGGED_IN) {
		rv = p11t_key_login_context_specific (session, CK_INVALID);
		if (rv == CKR_OK)
			rv = (p11t_module_funcs->C_Verify) (session, input, n_input, output, n_output);
	}

	return rv;
}

int
p11t_crypto_test_verify (CK_SESSION_HANDLE session, CK_MECHANISM_TYPE mech)
{
	CK_OBJECT_HANDLE handle;
	CK_MECHANISM mechanism;
	CK_BYTE input[P11T_BLOCK];
	CK_BYTE invalid[P11T_BLOCK];
	CK_BYTE signature[P11T_BLOCK];
	CK_ULONG n_input, n_invalid, n_signature;
	CK_RV rv;

	if (!p11t_test_unexpected)
		return CONTINUE;

	/* An valid mechanism */
	mechanism.mechanism = mech;
	mechanism.pParameter = NULL;
	mechanism.ulParameterLen = 0;

	test_without_init (session, 0);

	P11T_SECTION ("C_VerifyInit");

	rv = (p11t_module_funcs->C_VerifyInit)(session, &mechanism, (CK_ULONG)-9);
	P11T_CHECK_RV ("Invalid handle", rv, CKR_OBJECT_HANDLE_INVALID);

	handle = find_non_key_object (session);
	if (handle != CK_INVALID) {
		rv = (p11t_module_funcs->C_VerifyInit) (session, &mechanism, handle);
		P11T_CHECK_RV ("Invalid object type", rv, CKR_KEY_HANDLE_INVALID);
	}

	handle = find_non_method_key (session, mech, CKA_VERIFY);
	if (handle != CK_INVALID) {
		rv = (p11t_module_funcs->C_VerifyInit) (session, &mechanism, handle);
		P11T_CHECK_RV ("Key that's not verifiable", rv, CKR_KEY_FUNCTION_NOT_PERMITTED);
	}

	handle = find_non_mechanism_key (session, mech, CKA_VERIFY);
	if (handle != CK_INVALID) {
		rv = (p11t_module_funcs->C_VerifyInit) (session, &mechanism, handle);
		P11T_CHECK_RV ("Key from wrong mechanism", rv, CKR_KEY_TYPE_INCONSISTENT);
	}

	/* Finally find a valid object */
	handle = find_applicable_key (session, mech, CKA_VERIFY);
	if (handle == CK_INVALID)
		return CONTINUE;

	/* Prepare some input to verify */
	n_input = P11T_BLOCK;
	if (!prepare_raw_data (session, mech, CKA_SIGN, handle, input, &n_input))
		return CONTINUE;

	/* Prepare a signature to verify */
	n_signature = P11T_BLOCK;
	if (!prepare_verify_data (session, mech, handle, input, n_input, signature, &n_signature))
		return CONTINUE;

	rv = (p11t_module_funcs->C_VerifyInit)((CK_ULONG)-5, &mechanism, handle);
	P11T_CHECK_RV ("Invalid session", rv, CKR_SESSION_HANDLE_INVALID);

	rv = (p11t_module_funcs->C_VerifyInit)(session, NULL, handle);
	P11T_CHECK_RV ("Null mechanism", rv, CKR_ARGUMENTS_BAD);

	/* The real initialize here */

	rv = (p11t_module_funcs->C_VerifyInit)(session, &mechanism, handle);
	P11T_CHECK_RV ("A normal call", rv, CKR_OK);


	/* All other crypto operations shouldn't be valid */
	test_without_init (session, CKA_VERIFY);

	P11T_SECTION ("C_Verify");

	/* All the invalid verifys first */

	rv = run_verify_login_if_necessary ((CK_ULONG)-5, input, n_input, signature, n_signature);
	P11T_CHECK_RV ("Invalid session", rv, CKR_SESSION_HANDLE_INVALID);

	/*
	 * This does not terminate the current operation, since we
	 * didn't give it a valid session. The below invalid tests should
	 * terminate the current operation.
	 */

	rv = run_verify_login_if_necessary (session, NULL, n_input, signature, n_signature);
	P11T_CHECK_RV ("Null input", rv, CKR_ARGUMENTS_BAD);

	rv = (p11t_module_funcs->C_VerifyInit) (session, &mechanism, handle);
	P11T_CHECK_RV ("Error result terminates operation", rv, CKR_OK);

	rv = run_verify_login_if_necessary (session, NULL, n_input, NULL, n_signature);
	P11T_CHECK_RV ("Null signature", rv, CKR_ARGUMENTS_BAD);

	rv = (p11t_module_funcs->C_VerifyInit) (session, &mechanism, handle);
	P11T_CHECK_RV ("Error result terminates operation", rv, CKR_OK);

	/* Finally do the actual verifying */
	rv = run_verify_login_if_necessary (session, input, n_input, signature, n_signature);
	P11T_CHECK_RV ("Normal call", rv, CKR_OK);

	/* Some more invalid checks */
	n_invalid = P11T_BLOCK;
	if (prepare_raw_invalid_len (session, mech, CKA_SIGN, handle, invalid, &n_invalid)) {
		rv = (p11t_module_funcs->C_VerifyInit) (session, &mechanism, handle);
		P11T_CHECK_RV ("Successful result terminates operation", rv, CKR_OK);

		rv = run_verify_login_if_necessary (session, invalid, n_invalid, signature, n_signature);
		P11T_CHECK_RV ("Bad input length", rv, CKR_DATA_LEN_RANGE);
	}

	n_invalid = P11T_BLOCK;
	if (prepare_verify_invalid_len (session, mech, handle, &n_invalid)) {
		rv = (p11t_module_funcs->C_VerifyInit) (session, &mechanism, handle);
		P11T_CHECK_RV ("Successful result terminates operation", rv, CKR_OK);

		rv = run_verify_login_if_necessary (session, input, n_input, signature, n_invalid);
		P11T_CHECK_RV ("Bad input length", rv, CKR_SIGNATURE_LEN_RANGE);
	}

	n_invalid = P11T_BLOCK;
	if (prepare_raw_invalid (session, mech, CKA_SIGN, handle, invalid, &n_invalid)) {
		rv = (p11t_module_funcs->C_VerifyInit) (session, &mechanism, handle);
		P11T_CHECK_RV ("Successful result terminates operation", rv, CKR_OK);

		rv = run_verify_login_if_necessary (session, invalid, n_invalid, signature, n_signature);
		P11T_CHECK_RV ("Invalid input", rv, CKR_DATA_INVALID);
	}

	/* Now corrupt the signature and check */
	signature[0] = ~(signature[0]);

	rv = (p11t_module_funcs->C_VerifyInit) (session, &mechanism, handle);
	P11T_CHECK_RV ("Successful result terminates operation", rv, CKR_OK);

	rv = run_verify_login_if_necessary (session, input, n_input, signature, n_signature);
	P11T_CHECK_RV ("Corrupted signature", rv, CKR_SIGNATURE_INVALID);


	/* Now corrupt the data and check */
	signature[0] = ~(signature[0]);
	input[0] = ~(input[0]);

	rv = (p11t_module_funcs->C_VerifyInit) (session, &mechanism, handle);
	P11T_CHECK_RV ("Error result terminates operation", rv, CKR_OK);

	rv = run_verify_login_if_necessary (session, input, n_input, signature, n_signature);
	P11T_CHECK_RV ("Corrupted signature", rv, CKR_SIGNATURE_INVALID);


	/* At this point no crypto operations should be valid */

	test_without_init (session, 0);

	return CONTINUE;
}