#include "p11-tests-lib.h"

#include <stdlib.h>
#include <string.h>

/* An very likey to be invalid attribute */
#define CKA_INVALID 0x80FFFFFFUL

/* ----------------------------------------------------------------------------------
 * TESTS
 */

static int
test_create_object(CK_SESSION_HANDLE session_rw, CK_SESSION_HANDLE session_ro)
{
	P11T_SECTION("C_CreateObject");
	P11T_CHECK_NOTE("Not Tested");
	return CONTINUE;
}

static int
test_copy_object(CK_SESSION_HANDLE session_rw, CK_SESSION_HANDLE session_ro)
{
	P11T_SECTION("C_CopyObject");
	P11T_CHECK_NOTE("Not Tested");
	return CONTINUE;
}

static int
test_destroy_object(CK_SESSION_HANDLE session_rw, CK_SESSION_HANDLE session_ro)
{
	P11T_SECTION("C_DestroyObject");
	P11T_CHECK_NOTE("Not Tested");
	return CONTINUE;
}

static int
compar_object_handles(const void *one, const void *two)
{
	CK_OBJECT_HANDLE a = *((CK_SLOT_ID*)one);
	CK_OBJECT_HANDLE b = *((CK_SLOT_ID*)two);
	if(a == b)
		return 0;
	if(a > b)
		return 1;
	return -1;
}

static int
test_find_consistent(CK_SESSION_HANDLE session)
{
	CK_OBJECT_HANDLE_PTR objects;
	CK_OBJECT_HANDLE_PTR check;
	CK_ULONG n_objects;
	CK_ULONG i, n_check;

	objects = p11t_object_find(session, NULL, 0, &n_objects);
	if(!objects)
		return CONTINUE;

	qsort(objects, n_objects, sizeof(CK_OBJECT_HANDLE), compar_object_handles);

	for(i = 0; i < 10; ++i)
	{
		check = p11t_object_find(session, NULL, 0, &n_check);
		if(check == NULL)
			P11T_CHECK_FAIL_MSG("Consistently finds same objects", "no objects found");
		if(n_objects != n_check)
			P11T_CHECK_FAIL_MSG("Consistently finds same objects", "wrong number of objects found");
		qsort(check, n_check, sizeof(CK_OBJECT_HANDLE), compar_object_handles);
		if(memcmp(objects, check, n_check * sizeof(CK_OBJECT_HANDLE)) != 0)
			P11T_CHECK_FAIL_MSG("Consistently finds same objects", "different objects found");

		free(check);
	}

	free(objects);

	return CONTINUE;
}

static int
test_find_objects(CK_SESSION_HANDLE session)
{
	CK_OBJECT_HANDLE objects[1024];
	CK_OBJECT_HANDLE extra;
	CK_ATTRIBUTE attr;
	CK_SESSION_INFO sinfo;
	CK_ULONG object_count, count, i;
	CK_BBOOL is_private;
	CK_RV rv;

	assert(p11t_module_funcs);

	P11T_SECTION("C_FindObjectsInit");

	if(p11t_test_unexpected)
	{
		rv = (p11t_module_funcs->C_FindObjectsInit)((CK_SESSION_HANDLE)-99, &attr, 0);
		P11T_CHECK_RV("Invalid session", rv, CKR_SESSION_HANDLE_INVALID);

		rv = (p11t_module_funcs->C_FindObjectsInit)(session, NULL, 1);
		P11T_CHECK_RV("Attribute count without buffer", rv, CKR_ARGUMENTS_BAD);

	}

	rv = (p11t_module_funcs->C_FindObjectsInit)(session, &attr, 0);
	P11T_CHECK_RV("Find all objects", rv, CKR_OK);


	P11T_SECTION("C_FindObjects");

	if(p11t_test_unexpected)
	{
		rv = (p11t_module_funcs->C_FindObjects)((CK_SESSION_HANDLE)-99, NULL, 0, &object_count);
		P11T_CHECK_RV("Invalid session", rv, CKR_SESSION_HANDLE_INVALID);

		rv = (p11t_module_funcs->C_FindObjects)(session, NULL, 0, NULL);
		P11T_CHECK_RV("Null object count", rv, CKR_ARGUMENTS_BAD);

		rv = (p11t_module_funcs->C_FindObjects)(session, objects, 1, &count);
		P11T_CHECK_RV("Retrieve a single object before remainder", rv, CKR_OK);
		if(count != 0)
			P11T_CHECK_ULONG("If asked for a single object, return one.", count, 1);
	}

	do
	{
		rv = (p11t_module_funcs->C_FindObjects)(session, objects, 1024, &count);
		P11T_CHECK_RV("Retrieve remaining objects", rv, CKR_OK);
	}
	while(count == 1024);

	/* Get the session state */
	rv = (p11t_module_funcs->C_GetSessionInfo) (session, &sinfo);
	P11T_CHECK_RV("Check private objects not found in public session", rv, CKR_OK);

	/* Check that they're not private, if the session is logged in */
	if (sinfo.state == CKS_RO_PUBLIC_SESSION || sinfo.state == CKS_RW_PUBLIC_SESSION) {
		for (i = 0; i < count; ++i) {
			attr.type = CKA_PRIVATE;
			attr.ulValueLen = sizeof (is_private);
			attr.pValue = &is_private;
			rv = (p11t_module_funcs->C_GetAttributeValue) (session, objects[i], &attr, 1);
			P11T_CHECK_RV("Check private objects not found in public session", rv, CKR_OK);
			if (is_private)
				P11T_CHECK_FAIL("Check private objects not found in public session");
		}
	}

	if(p11t_test_unexpected)
	{
		rv = (p11t_module_funcs->C_FindObjects(session, &extra, 1, &count));
		P11T_CHECK_RV("Extra call after retrieving all objects", rv, CKR_OK);
		P11T_CHECK_ULONG("Should return no objects in extra call", count, 0);
	}

	P11T_SECTION("C_FindObjectsFinal");

	if(p11t_test_unexpected)
	{
		rv = (p11t_module_funcs->C_FindObjectsFinal)((CK_SESSION_HANDLE)-88);
		P11T_CHECK_RV("Invalid session", rv, CKR_SESSION_HANDLE_INVALID);
	}

	rv = (p11t_module_funcs->C_FindObjectsFinal)(session);
	P11T_CHECK_RV("Normal call", rv, CKR_OK);

	if(p11t_test_unexpected)
	{
		rv = (p11t_module_funcs->C_FindObjectsFinal)(session);
		P11T_CHECK_RV("Extra call", rv, CKR_OPERATION_NOT_INITIALIZED);

		P11T_SECTION("C_FindObjects");

		rv = (p11t_module_funcs->C_FindObjects)(session, &extra, 1, &count);
		P11T_CHECK_RV("Out of order call", rv, CKR_OPERATION_NOT_INITIALIZED);
	}

	return CONTINUE;
}

static int
test_get_attribute_value(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE object)
{
	CK_OBJECT_CLASS klass;
	CK_ATTRIBUTE attrs[5];
	CK_BYTE buffer[8];
	CK_BBOOL bools[3];
	CK_RV rv;

	assert(p11t_module_funcs);

	P11T_SECTION("C_GetAttributeValue");

	memset(&attrs, 0, sizeof(attrs));
	attrs[0].type = CKA_CLASS;

	if(p11t_test_unexpected)
	{
		rv = (p11t_module_funcs->C_GetAttributeValue)((CK_SESSION_HANDLE)-57, object, attrs, 1);
		P11T_CHECK_RV("Invalid session", rv, CKR_SESSION_HANDLE_INVALID);

		rv = (p11t_module_funcs->C_GetAttributeValue)(session, (CK_OBJECT_HANDLE)-58, attrs, 1);
		P11T_CHECK_RV("Invalid object", rv, CKR_OBJECT_HANDLE_INVALID);

		rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, NULL, 1);
		P11T_CHECK_RV("No template", rv, CKR_ARGUMENTS_BAD);

		attrs[0].pValue = buffer;
		attrs[0].ulValueLen = 2;
		rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, attrs, 1);
		P11T_CHECK_RV("Buffer too small", rv, CKR_BUFFER_TOO_SMALL);
		P11T_CHECK_ULONG("Buffer too small should return -1", attrs[0].ulValueLen, (CK_ULONG)-1);
	}

	attrs[0].pValue = NULL;
	attrs[0].ulValueLen = 0;
	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, attrs, 1);
	P11T_CHECK_RV("Retrieve attribute length", rv, CKR_OK);
	P11T_CHECK_ULONG("Size of object class attribute.", attrs[0].ulValueLen, sizeof(CK_OBJECT_CLASS));

	attrs[0].pValue = buffer;
	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, attrs, 1);
	P11T_CHECK_RV("Retrieve single attribute", rv, CKR_OK);

	/* The actual class of this object */
	klass = *((CK_OBJECT_CLASS*)attrs[0].pValue);

	if(p11t_test_unexpected)
	{
		/*
		 * Now we do a bit of testing to retrieve a valid and invalid attribute
		 * at the same time.
		 */

		attrs[0].type = CKA_INVALID;
		attrs[1].type = CKA_CLASS;
		attrs[0].ulValueLen = attrs[1].ulValueLen = 0;
		attrs[0].pValue = attrs[1].pValue = NULL;

		rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, attrs, 2);
		P11T_CHECK_RV("With one invalid attribute, no buffer", rv, CKR_ATTRIBUTE_TYPE_INVALID);
		P11T_CHECK_ULONG("Should set invalid attribute to -1", attrs[0].ulValueLen, CK_INVALID);
		P11T_CHECK_ULONG("Should set valid attribute to size", attrs[1].ulValueLen, sizeof(CK_OBJECT_CLASS));

		attrs[1].pValue = buffer;
		rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, attrs, 2);
		P11T_CHECK_RV("With one invalid attribute, with buffer", rv, CKR_ATTRIBUTE_TYPE_INVALID);
		P11T_CHECK_ULONG("Should set invalid attribute to -1", attrs[0].ulValueLen, CK_INVALID);
		P11T_CHECK_ULONG("Should set attribute size, in partially valid set", attrs[1].ulValueLen, sizeof(CK_OBJECT_CLASS));
		P11T_CHECK_ULONG("Should set valid attribute, in partially valid set", *((CK_OBJECT_CLASS*)attrs[1].pValue), klass);
	}

	/*
	 * Now we do a bit of testing of other attributes and retrieving sets.
	 * Since we can't list which attributes an object has, we only
	 * work with 'storage' type attributes and their standard 4 attributes.
	 */

	switch(klass)
	{
	case CKO_CERTIFICATE:
	case CKO_DATA:
	case CKO_PRIVATE_KEY:
	case CKO_PUBLIC_KEY:
	case CKO_SECRET_KEY:
		break;
	default:
		return CONTINUE;
	}

	memset(attrs, 0, sizeof(attrs));
	attrs[0].type = CKA_CLASS;
	attrs[1].type = CKA_TOKEN;
	attrs[2].type = CKA_PRIVATE;
	attrs[3].type = CKA_MODIFIABLE;
	attrs[4].type = CKA_LABEL;

	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, attrs, 5);
	P11T_CHECK_RV("Multiple attributes, no buffer", rv, CKR_OK);
	P11T_CHECK_ULONG("Size of CKA_CLASS", attrs[0].ulValueLen, sizeof(CK_OBJECT_CLASS));
	P11T_CHECK_ULONG("Size of CKA_TOKEN", attrs[1].ulValueLen, sizeof(CK_BBOOL));
	P11T_CHECK_ULONG("Size of CKA_PRIVATE", attrs[2].ulValueLen, sizeof(CK_BBOOL));
	P11T_CHECK_ULONG("Size of CKA_MODIFIABLE", attrs[3].ulValueLen, sizeof(CK_BBOOL));

	if(attrs[4].ulValueLen)
		attrs[4].pValue = malloc(attrs[4].ulValueLen);
	attrs[3].pValue = &bools[0];
	attrs[2].pValue = &bools[1];
	attrs[1].pValue = &bools[2];

	if(p11t_test_unexpected)
	{
		rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, attrs, 5);
		P11T_CHECK_RV("Multiple attributes, some buffers", rv, CKR_OK);
		P11T_CHECK_ULONG("Size of CKA_CLASS", attrs[0].ulValueLen, sizeof(CK_OBJECT_CLASS));
		P11T_CHECK_BOOL("CKA_TOKEN boolean value", *((CK_BBOOL*)attrs[1].pValue));
		P11T_CHECK_BOOL("CKA_PRIVATE boolean value", *((CK_BBOOL*)attrs[2].pValue));
		P11T_CHECK_BOOL("CKA_MODIFIABLE boolean value", *((CK_BBOOL*)attrs[3].pValue));

		attrs[0].pValue = buffer;
		attrs[0].ulValueLen = 2;

		rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, attrs, 5);
		P11T_CHECK_RV("Multiple attributes, one small buffer", rv, CKR_BUFFER_TOO_SMALL);
		P11T_CHECK_ULONG("Invalid attribute size", attrs[0].ulValueLen, (CK_ULONG)-1);
		P11T_CHECK_BOOL("CKA_TOKEN boolean value", *((CK_BBOOL*)attrs[1].pValue));
		P11T_CHECK_BOOL("CKA_PRIVATE boolean value", *((CK_BBOOL*)attrs[2].pValue));
		P11T_CHECK_BOOL("CKA_MODIFIABLE boolean value", *((CK_BBOOL*)attrs[3].pValue));
	}

	attrs[0].pValue = buffer;
	attrs[0].ulValueLen = sizeof(CK_OBJECT_CLASS);

	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, attrs, 5);
	P11T_CHECK_RV("Multiple attributes, with buffers", rv, CKR_OK);
	P11T_CHECK_ULONG("Size of CKA_CLASS", attrs[0].ulValueLen, sizeof(CK_OBJECT_CLASS));
	P11T_CHECK_ULONG("CKA_CLASS value", *((CK_ULONG*)attrs[0].pValue), klass);
	P11T_CHECK_BOOL("CKA_TOKEN boolean value", *((CK_BBOOL*)attrs[1].pValue));
	P11T_CHECK_BOOL("CKA_PRIVATE boolean value", *((CK_BBOOL*)attrs[2].pValue));
	P11T_CHECK_BOOL("CKA_MODIFIABLE boolean value", *((CK_BBOOL*)attrs[3].pValue));
	P11T_CHECK_STRING("CKA_LABEL value", attrs[4].pValue, attrs[4].ulValueLen);

	if(attrs[4].ulValueLen)
		free(attrs[4].pValue);

	return CONTINUE;
}

static int
test_objects(CK_SESSION_HANDLE session)
{
	CK_OBJECT_HANDLE_PTR objects;
	CK_ULONG n_objects;
	CK_ULONG i;

	assert(p11t_module_funcs);

	objects = p11t_object_find(session, NULL, 0, &n_objects);
	if(!objects || !n_objects)
		return CONTINUE;

	for(i = 0; i < n_objects; ++i)
	{
		P11T_SECTION("C_GetObjectSize");
		P11T_CHECK_NOTE("Not Tested");

		test_get_attribute_value(session, objects[i]);
	}

	free (objects);
	return CONTINUE;
}

void
p11t_object_tests(void)
{
	CK_SESSION_HANDLE session_rw;
	CK_SESSION_HANDLE session_ro;
	CK_SLOT_ID slot;
	CK_ULONG i;

	for(i = 0; i < p11t_slot_count; ++i)
	{
		slot = p11t_slot_get_id(i);
		session_rw = p11t_session_open(slot, 1);
		session_ro = p11t_session_open(slot, 0);

		/* We need to have the readonly session and log in for tests to proceed */
		if(session_ro == CK_INVALID)
			continue;

		test_find_objects(session_ro);
		test_find_consistent(session_ro);
		test_objects(session_ro);

		if(!p11t_session_login(session_ro))
			continue;

		test_find_objects(session_ro);
		test_objects(session_ro);

		if(session_rw != CK_INVALID)
		{
			test_find_objects(session_rw);
			test_find_consistent(session_rw);
			test_objects(session_rw);
			test_create_object(session_rw, session_ro);
			test_copy_object(session_rw, session_ro);
			test_destroy_object(session_rw, session_ro);
		}

		p11t_session_close_all(slot);
	}
}

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
 */

CK_OBJECT_HANDLE_PTR
p11t_object_find(CK_SESSION_HANDLE session, CK_ATTRIBUTE_PTR attrs,
                 CK_ULONG n_attrs, CK_ULONG_PTR n_objects)
{
	CK_OBJECT_HANDLE_PTR objects;
	CK_ATTRIBUTE dummy_attr;

	assert(p11t_module_funcs);

	if(!attrs)
		attrs = &dummy_attr;

	if((p11t_module_funcs->C_FindObjectsInit)(session, attrs, n_attrs) != CKR_OK)
		return NULL;

	objects = calloc(1024, sizeof(CK_OBJECT_HANDLE));
	if((p11t_module_funcs->C_FindObjects)(session, objects, 1024, n_objects) != CKR_OK)
	{
		free(objects);
		return NULL;
	}

	if((p11t_module_funcs->C_FindObjectsFinal)(session) != CKR_OK)
	{
		free(objects);
		return NULL;
	}

	if(!*n_objects)
	{
		free(objects);
		return NULL;
	}

	return objects;
}

CK_OBJECT_HANDLE
p11t_object_find_one(CK_SESSION_HANDLE session, CK_ATTRIBUTE_PTR attrs,
                     CK_ULONG n_attrs)
{
	CK_OBJECT_HANDLE object;
	CK_ULONG count;

	assert(p11t_module_funcs);

	if((p11t_module_funcs->C_FindObjectsInit)(session, attrs, n_attrs) != CKR_OK)
		return CK_INVALID;

	if((p11t_module_funcs->C_FindObjects)(session, &object, 1, &count) != CKR_OK)
		return CK_INVALID;

	if((p11t_module_funcs->C_FindObjectsFinal)(session) != CKR_OK)
		return CK_INVALID;

	if(count != 1)
		return CK_INVALID;

	return object;
}


int
p11t_object_get(CK_SESSION_HANDLE session, CK_OBJECT_HANDLE object,
                CK_ATTRIBUTE_PTR attrs, CK_ULONG count)
{
	CK_RV rv;

	assert(p11t_module_funcs);

	if(!count)
		return 1;

	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, attrs, count);
	switch(rv)
	{
	case CKR_OK:
	case CKR_BUFFER_TOO_SMALL:
	case CKR_ATTRIBUTE_TYPE_INVALID:
	case CKR_ATTRIBUTE_SENSITIVE:
		return 1;
	}

	return 0;
}