#include "p11-tests.h"

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

/* An very likey to be invalid attribute */
#define CKA_INVALID (CKA_VENDOR_DEFINED | 0x00FFFFFF)

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

static void
create_object(CK_SESSION_HANDLE session_rw, CK_SESSION_HANDLE session_ro)
{
	/** C_CreateObject */

	/** - Not Implemented */
}

static void
copy_object(CK_SESSION_HANDLE session_rw, CK_SESSION_HANDLE session_ro)
{
	/** C_CopyObject */

	/** - Not Implemented */
}

static void
destroy_object(CK_SESSION_HANDLE session_rw, CK_SESSION_HANDLE session_ro)
{
	/** C_DestroyObject */

	/** - Not Implemented */
}

static void
find_objects(CK_SESSION_HANDLE session)
{
	CK_OBJECT_HANDLE objects[1024];
	CK_OBJECT_HANDLE extra;
	CK_ATTRIBUTE dummy;
	CK_ULONG object_count, count;
	CK_RV rv;

	assert(p11t_module_funcs);

	/** C_FindObjectsInit */

	if(p11t_test_unexpected)
	{
		/** - Invalid session */
		rv = (p11t_module_funcs->C_FindObjectsInit)((CK_SESSION_HANDLE)-99, &dummy, 0);
		p11t_check_returns("C_FindObjectsInit: invalid session", rv, CKR_SESSION_HANDLE_INVALID);

		/** - Attribute count without buffer */
		rv = (p11t_module_funcs->C_FindObjectsInit)((CK_SESSION_HANDLE)-99, NULL, 1);
		p11t_check_returns("C_FindObjectsInit: count without buffer", rv, CKR_ARGUMENTS_BAD);

	}

	rv = (p11t_module_funcs->C_FindObjectsInit)(session, &dummy, 0);
	p11t_check_returns("C_FindObjectsInit: all", rv, CKR_OK);

	if(p11t_test_unexpected)
	{
		/** - Double call. */
		rv = (p11t_module_funcs->C_FindObjectsInit)(session, &dummy, 0);
		p11t_check_returns("C_FindObjectsInit: all", rv, CKR_OPERATION_ACTIVE);

	}

	/** C_FindObjects */

	if(p11t_test_unexpected)
	{
		/** - Invalid session */
		rv = (p11t_module_funcs->C_FindObjects)((CK_SESSION_HANDLE)-99, NULL, 0, &object_count);
		p11t_check_returns("C_FindObjects: invalid session", rv, CKR_SESSION_HANDLE_INVALID);

		/** - Null object count */
		rv = (p11t_module_funcs->C_FindObjects)(session, NULL, 0, NULL);
		p11t_check_returns("C_FindObjects: invalid session", rv, CKR_ARGUMENTS_BAD);

		/** - Retrieve a single object before remainder */
		rv = (p11t_module_funcs->C_FindObjects)(session, objects, 1, &count);
		p11t_check_returns("C_FindObjects: single object", rv, CKR_OK);
		if(count != 0)
			p11t_check_ulong("C_FindObjects: should return one object", count, 1);
	}

	do
	{
		/** - Retrieve remaining objects */
		rv = (p11t_module_funcs->C_FindObjects)(session, objects, 1024, &count);
		p11t_check_returns("C_FindObjects: all objects", rv, CKR_OK);
	}
	while(count == 1024);

	if(p11t_test_unexpected)
	{
		/** - Extra call after retrieving all objects. */
		rv = (p11t_module_funcs->C_FindObjects(session, &extra, 1, &count));
		p11t_check_returns("C_FindObjects: extra call", rv, CKR_OK);
		p11t_check_ulong("C_FindObjects: should return no objects in extra call", count, 0);
	}

	/** C_FindObjectsFinal */

	if(p11t_test_unexpected)
	{
		/** - Invalid session */
		rv = (p11t_module_funcs->C_FindObjectsFinal)((CK_SESSION_HANDLE)-88);
		p11t_check_returns("C_FindObjectsFinal: invalid session", rv, CKR_SESSION_HANDLE_INVALID);
	}

	/** - Normal call */
	rv = (p11t_module_funcs->C_FindObjectsFinal)(session);
	p11t_check_returns("C_FindObjectsFinal", rv, CKR_OK);

	if(p11t_test_unexpected)
	{
		/** - Extra call */
		rv = (p11t_module_funcs->C_FindObjectsFinal)(session);
		p11t_check_returns("C_FindObjectsFinal: extra call", rv, CKR_OPERATION_NOT_INITIALIZED);

		/** C_FindObjects */

		/** - Out of order call */
		rv = (p11t_module_funcs->C_FindObjects)(session, &extra, 1, &count);
		p11t_check_returns("C_FindObjectsFinal: extra call", rv, CKR_OPERATION_NOT_INITIALIZED);
	}
}

static void
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);

	/** C_GetAttributeValue */

	attrs[0].type = CKA_CLASS;

	if(p11t_test_unexpected)
	{
		/** - Invalid session */
		rv = (p11t_module_funcs->C_GetAttributeValue)((CK_SESSION_HANDLE)-57, object, attrs, 1);
		p11t_check_returns("C_GetAttributeValue: invalid session", rv, CKR_SESSION_HANDLE_INVALID);

		/** - Invalid object */
		rv = (p11t_module_funcs->C_GetAttributeValue)(session, (CK_OBJECT_HANDLE)-58, attrs, 1);
		p11t_check_returns("C_GetAttributeValue: invalid object", rv, CKR_OBJECT_HANDLE_INVALID);

		/** - No template */
		rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, NULL, 0);
		p11t_check_returns("C_GetAttributeValue: no template", rv, CKR_ARGUMENTS_BAD);

		/** - Buffer too small */
		attrs[0].pValue = buffer;
		attrs[0].ulValueLen = 2;
		rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, attrs, 1);
		p11t_check_returns("C_GetAttributeValue: buffer too small", rv, CKR_BUFFER_TOO_SMALL);
		p11t_check_ulong("C_GetAttributeValue: should return size", attrs[0].ulValueLen, sizeof(CK_OBJECT_CLASS));
	}

	/** - Retrieve attribute length */
	attrs[0].pValue = NULL;
	attrs[0].ulValueLen = 0;
	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, attrs, 1);
	p11t_check_returns("C_GetAttributeValue: get size", rv, CKR_OK);
	p11t_check_ulong("C_GetAttributeValue: size of class", attrs[0].ulValueLen, sizeof(CK_OBJECT_CLASS));

	/** - Retrieve single attribute */
	attrs[0].pValue = buffer;
	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, attrs, 1);
	p11t_check_returns("C_GetAttributeValue: get object", 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[0].pValue = NULL;

		/** - With one invalid attribute, no buffer. */
		rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, attrs, 2);
		p11t_check_returns("C_GetAttributeValue: with one invalid attribute", rv, CKR_ATTRIBUTE_TYPE_INVALID);
		p11t_check_ulong("C_GetAttributeValue: should set invalid attribute to -1", attrs[0].ulValueLen, CK_INVALID);
		p11t_check_ulong("C_GetAttributeValue: should set valid attribute to size", attrs[1].ulValueLen, sizeof(CK_OBJECT_CLASS));

		/** - With one invalid attribute, with buffer. */
		attrs[1].pValue = buffer;
		rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, attrs, 2);
		p11t_check_returns("C_GetAttributeValue: with one invalid attribute, buffer", rv, CKR_ATTRIBUTE_TYPE_INVALID);
		p11t_check_ulong("C_GetAttributeValue: should set invalid attribute to -1", attrs[0].ulValueLen, CK_INVALID);
		p11t_check_ulong("C_GetAttributeValue: should set valid attribute to size", attrs[1].ulValueLen, sizeof(CK_OBJECT_CLASS));
		p11t_check_ulong("C_GetAttributeValue: should set valid attribute", *((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;
	}

	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;

	/** - Multiple attributes, no buffer. */
	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, attrs, 5);
	p11t_check_returns("C_GetAttributeValue: multiple attrs, no buffer", rv, CKR_OK);
	p11t_check_ulong("C_GetAttributeValue: CKA_CLASS size", attrs[0].ulValueLen, sizeof(CK_OBJECT_CLASS));
	p11t_check_ulong("C_GetAttributeValue: CKA_TOKEN size", attrs[1].ulValueLen, sizeof(CK_BBOOL));
	p11t_check_ulong("C_GetAttributeValue: CKA_PRIVATE size", attrs[2].ulValueLen, sizeof(CK_BBOOL));
	p11t_check_ulong("C_GetAttributeValue: CKA_MODIFIABLE size", 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)
	{
		/** - Multiple attributes, some buffers */
		rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, attrs, 5);
		p11t_check_returns("C_GetAttributeValue: multiple attrs, some buffer", rv, CKR_OK);
		p11t_check_ulong("C_GetAttributeValue: CKA_CLASS no buffer, size", attrs[0].ulValueLen, sizeof(CK_OBJECT_CLASS));
		p11t_check_bool("C_GetAttributeValue: CKA_TOKEN", *((CK_BBOOL*)attrs[1].pValue));
		p11t_check_bool("C_GetAttributeValue: CKA_PRIVATE", *((CK_BBOOL*)attrs[2].pValue));
		p11t_check_bool("C_GetAttributeValue: CKA_MODIFIABLE", *((CK_BBOOL*)attrs[3].pValue));

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

		/** - Multiple attributes, one small buffer */
		rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, attrs, 5);
		p11t_check_returns("C_GetAttributeValue: multiple attrs, one small buffer", rv, CKR_BUFFER_TOO_SMALL);
		p11t_check_ulong("C_GetAttributeValue: CKA_CLASS no buffer, size", attrs[0].ulValueLen, sizeof(CK_OBJECT_CLASS));
		p11t_check_bool("C_GetAttributeValue: CKA_TOKEN", *((CK_BBOOL*)attrs[1].pValue));
		p11t_check_bool("C_GetAttributeValue: CKA_PRIVATE", *((CK_BBOOL*)attrs[2].pValue));
		p11t_check_bool("C_GetAttributeValue: CKA_MODIFIABLE", *((CK_BBOOL*)attrs[3].pValue));
	}

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

	/** - Multiple attributes, with buffers */
	rv = (p11t_module_funcs->C_GetAttributeValue)(session, object, attrs, 5);
	p11t_check_returns("C_GetAttributeValue: multiple attrs, with buffers", rv, CKR_OK);
	p11t_check_ulong("C_GetAttributeValue: CKA_CLASS no buffer, size", attrs[0].ulValueLen, sizeof(CK_OBJECT_CLASS));
	p11t_check_ulong("C_GetAttributeValue: CKA_CLASS value", *((CK_ULONG*)attrs[0].pValue), klass);
	p11t_check_bool("C_GetAttributeValue: CKA_TOKEN", *((CK_BBOOL*)attrs[1].pValue));
	p11t_check_bool("C_GetAttributeValue: CKA_PRIVATE", *((CK_BBOOL*)attrs[2].pValue));
	p11t_check_bool("C_GetAttributeValue: CKA_MODIFIABLE", *((CK_BBOOL*)attrs[3].pValue));
	p11t_check_string("C_GetAttributeValue: CKA_LABEL", attrs[4].pValue, attrs[4].ulValueLen);
}

static void
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;

	for(i = 0; i < n_objects; ++i)
	{
		/** C_GetObjectSize */

		/** - Not Implemented */

		test_get_attribute_value(session, objects[i]);
	}

	free (objects);
}

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;

		find_objects(session_ro);
		test_objects(session_ro);

		if(!p11t_session_login(session_ro))
			continue;

		find_objects(session_ro);
		test_objects(session_ro);

		if(session_rw != CK_INVALID)
		{
			find_objects(session_rw);
			test_objects(session_rw);
			create_object(session_rw, session_ro);
			copy_object(session_rw, session_ro);
			destroy_object(session_rw, session_ro);
		}

		p11t_session_close_all(slot);
	}
}

/* ----------------------------------------------------------------------------------
 * 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(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;
}