#include "p11-tests-lib.h"

#include <string.h>

static const char *login_user_pin = NULL;
static const char *login_so_pin = NULL;
static const char *login_context_pin = NULL;

void
p11t_session_config(const char *name, const char *value)
{
	if(strcmp(name, "login-user-pin") == 0)
		login_user_pin = value;
	if(strcmp(name, "login-so-pin") == 0)
		login_so_pin = value;
	if(strcmp(name, "login-context-pin") == 0)
		login_context_pin = value;
}

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

static int
test_session_info(CK_SESSION_HANDLE session, CK_SLOT_ID slot, CK_FLAGS flags, CK_STATE state)
{
	CK_SESSION_INFO info;
	CK_RV rv;

	if(session == CK_INVALID)
		return CONTINUE;

	P11T_SECTION("C_GetSessionInfo");

	if(p11t_test_unexpected)
	{
		rv = (p11t_module_funcs->C_GetSessionInfo)((CK_SESSION_HANDLE)-33, &info);
		P11T_CHECK_RV("Invalid session", rv, CKR_SESSION_HANDLE_INVALID);

		rv = (p11t_module_funcs->C_GetSessionInfo)(session, NULL);
		P11T_CHECK_RV("Null arguments", rv, CKR_ARGUMENTS_BAD);
	}

	rv = (p11t_module_funcs->C_GetSessionInfo)(session, &info);
	P11T_CHECK_RV("Valid call", rv, CKR_OK);

	P11T_SECTION("CK_SESSION_INFO");

	P11T_CHECK_ULONG("slotID", info.slotID, slot);
	P11T_CHECK_ULONG("state", info.state, state);
	P11T_CHECK_FLAG("flags", info.flags, flags);

	return CONTINUE;
}

static int
test_session_user_login(CK_SLOT_ID slot)
{
	CK_SESSION_HANDLE session_ro = CK_INVALID;
	CK_SESSION_HANDLE session_rw = CK_INVALID;
	CK_UTF8CHAR_PTR pin;
	CK_ULONG n_pin;
	CK_RV rv;

	pin = p11t_session_get_pin (slot, CKU_USER, &n_pin);

	rv = (p11t_module_funcs->C_OpenSession)(slot, CKF_SERIAL_SESSION, NULL, NULL, &session_ro);
	P11T_CHECK_RV("C_Login: open RO session", rv, CKR_OK);
	rv = (p11t_module_funcs->C_OpenSession)(slot, CKF_SERIAL_SESSION | CKF_RW_SESSION, NULL, NULL, &session_rw);
	P11T_CHECK_RV("C_Login: open RW session", rv, CKR_OK);

	/* Check the session state before login */
	test_session_info(session_ro, slot, CKF_SERIAL_SESSION, CKS_RO_PUBLIC_SESSION);
	test_session_info(session_rw, slot, CKF_SERIAL_SESSION | CKF_RW_SESSION, CKS_RW_PUBLIC_SESSION);

	P11T_SECTION("C_Login");

	if(p11t_test_unexpected)
	{
		rv = (p11t_module_funcs->C_Login)((CK_ULONG)-55, CKU_USER, pin, n_pin);
		P11T_CHECK_RV("Invalid session", rv, CKR_SESSION_HANDLE_INVALID);

		rv = (p11t_module_funcs->C_Login)(session_ro, CK_INVALID, pin, n_pin);
		P11T_CHECK_RV("Invalid user type", rv, CKR_USER_TYPE_INVALID);
	}

	rv = (p11t_module_funcs->C_Login)(session_ro, CKU_USER, pin, n_pin);
	if(rv == CKR_USER_TYPE_INVALID || rv == CKR_USER_PIN_NOT_INITIALIZED)
		return CONTINUE;
	P11T_CHECK_RV("Normal login", rv, CKR_OK);

	P11T_CHECK_NOTE("Login changes all session's state");
	test_session_info(session_ro, slot, CKF_SERIAL_SESSION, CKS_RO_USER_FUNCTIONS);
	test_session_info(session_rw, slot, CKF_SERIAL_SESSION | CKF_RW_SESSION, CKS_RW_USER_FUNCTIONS);

	P11T_SECTION("C_Logout");

	if(p11t_test_unexpected)
	{
		rv = (p11t_module_funcs->C_Logout)((CK_ULONG)-55);
		P11T_CHECK_RV("Invalid session", rv, CKR_SESSION_HANDLE_INVALID);
	}

	rv = (p11t_module_funcs->C_Logout)(session_rw);
	P11T_CHECK_RV("Normal logout", rv, CKR_OK);

	P11T_CHECK_NOTE("Logout changes all session state");
	test_session_info(session_ro, slot, CKF_SERIAL_SESSION, CKS_RO_PUBLIC_SESSION);
	test_session_info(session_rw, slot, CKF_SERIAL_SESSION | CKF_RW_SESSION, CKS_RW_PUBLIC_SESSION);

	if(p11t_test_unexpected)
	{
		rv = (p11t_module_funcs->C_Logout)(session_rw);
		P11T_CHECK_RV("Double logout", rv, CKR_USER_NOT_LOGGED_IN);
	}

	return CONTINUE;
}

static int
test_session_so_login(CK_SLOT_ID slot)
{
	CK_SESSION_HANDLE session_ro = CK_INVALID;
	CK_SESSION_HANDLE session_rw = CK_INVALID;
	CK_UTF8CHAR_PTR pin;
	CK_ULONG n_pin;
	CK_RV rv;

	(p11t_module_funcs->C_CloseAllSessions) (slot);

	pin = p11t_session_get_pin (slot, CKU_SO, &n_pin);

	rv = (p11t_module_funcs->C_OpenSession)(slot, CKF_SERIAL_SESSION | CKF_RW_SESSION, NULL, NULL, &session_rw);
	P11T_CHECK_RV("C_OpenSession: open RW session", rv, CKR_OK);

	/* Check the session state before login */
	test_session_info(session_rw, slot, CKF_SERIAL_SESSION | CKF_RW_SESSION, CKS_RW_PUBLIC_SESSION);

	P11T_SECTION("C_Login");

	if(p11t_test_unexpected)
	{
		rv = (p11t_module_funcs->C_OpenSession)(slot, CKF_SERIAL_SESSION, NULL, NULL, &session_ro);
		P11T_CHECK_RV("Test opening RO session before login as SO", rv, CKR_OK);

		rv = (p11t_module_funcs->C_Login)(session_ro, CKU_SO, pin, n_pin);
		if(rv == CKR_USER_TYPE_INVALID)
			return CONTINUE;
		P11T_CHECK_RV("Login as SO not allowed with RO session", rv, CKR_SESSION_READ_ONLY_EXISTS);

		rv = (p11t_module_funcs->C_CloseSession)(session_ro);
		P11T_CHECK_RV("Test closing RO session before login as SO", rv, CKR_OK);
	}

	rv = (p11t_module_funcs->C_Login)(session_rw, CKU_SO, pin, n_pin);
	if(rv == CKR_USER_TYPE_INVALID)
		return CONTINUE;
	P11T_CHECK_RV("Login as SO", rv, CKR_OK);

	P11T_CHECK_NOTE("SO login changes all session state");
	test_session_info(session_rw, slot, CKF_SERIAL_SESSION | CKF_RW_SESSION, CKS_RW_SO_FUNCTIONS);

	if(p11t_test_unexpected)
	{
		P11T_SECTION("C_OpenSession");

		rv = (p11t_module_funcs->C_OpenSession)(slot, CKF_SERIAL_SESSION, NULL, NULL, &session_ro);
		P11T_CHECK_RV("Can't open RO session when logged in as SO", rv, CKR_SESSION_READ_WRITE_SO_EXISTS);
	}

	P11T_SECTION("C_Logout");

	rv = (p11t_module_funcs->C_Logout)(session_rw);
	P11T_CHECK_RV("Logout from SO", rv, CKR_OK);

	P11T_CHECK_NOTE("Logout SO changes goes back to public");
	test_session_info(session_rw, slot, CKF_SERIAL_SESSION | CKF_RW_SESSION, CKS_RW_PUBLIC_SESSION);

	return CONTINUE;
}

static int
test_session_main(CK_SLOT_ID slot)
{
	CK_SESSION_HANDLE session_ro = CK_INVALID;
	CK_SESSION_HANDLE session_rw = CK_INVALID;
	CK_SESSION_HANDLE session_ro2 = CK_INVALID;
	CK_SESSION_INFO info;
	CK_RV rv;

	assert(p11t_module_funcs);

	P11T_SECTION("C_OpenSession");

	if(p11t_test_unexpected)
	{
		if (!p11t_slot_virtual)
		{
			rv = (p11t_module_funcs->C_OpenSession)((CK_SLOT_ID)-5, CKF_SERIAL_SESSION, NULL, NULL, &session_ro);
			P11T_CHECK_RV("Invalid slot", rv, CKR_SLOT_ID_INVALID);
		}

		rv = (p11t_module_funcs->C_OpenSession)(slot, CKF_SERIAL_SESSION, NULL, NULL, NULL);
		P11T_CHECK_RV("Null arguments", rv, CKR_ARGUMENTS_BAD);

		rv = (p11t_module_funcs->C_OpenSession)(slot, 0, NULL, NULL, &session_ro);
		P11T_CHECK_RV("No flags", rv, CKR_SESSION_PARALLEL_NOT_SUPPORTED);

		rv = (p11t_module_funcs->C_OpenSession)(slot, CKF_RW_SESSION, NULL, NULL, &session_ro);
		P11T_CHECK_RV("Without serial flag", rv, CKR_SESSION_PARALLEL_NOT_SUPPORTED);
	}

	rv = (p11t_module_funcs->C_OpenSession)(slot, CKF_SERIAL_SESSION, NULL, NULL, &session_ro);
	P11T_CHECK_RV("Valid flags", rv, CKR_OK);
	rv = (p11t_module_funcs->C_OpenSession)(slot, CKF_SERIAL_SESSION, NULL, NULL, &session_ro2);
	P11T_CHECK_RV("Valid flags", rv, CKR_OK);

	rv = (p11t_module_funcs->C_OpenSession)(slot, CKF_SERIAL_SESSION | CKF_RW_SESSION, NULL, NULL, &session_rw);
	if(rv == CKR_TOKEN_WRITE_PROTECTED)
		session_rw = CK_INVALID;
	else
		P11T_CHECK_RV("Read write session", rv, CKR_OK);

	/* Test all the sessions and validate their state */
	test_session_info(session_ro, slot, CKF_SERIAL_SESSION, CKS_RO_PUBLIC_SESSION);
	test_session_info(session_ro2, slot, CKF_SERIAL_SESSION, CKS_RO_PUBLIC_SESSION);
	test_session_info(session_rw, slot, CKF_SERIAL_SESSION | CKF_RW_SESSION, CKS_RW_PUBLIC_SESSION);

	/* These functions should return errors */
	P11T_SECTION("C_GetFunctionStatus");
	rv = (p11t_module_funcs->C_GetFunctionStatus) (session_ro);
	P11T_CHECK_RV("Should return error", rv, CKR_FUNCTION_NOT_PARALLEL);

	P11T_SECTION("C_CancelFunction");
	rv = (p11t_module_funcs->C_CancelFunction) (session_ro);
	P11T_CHECK_RV("Should return error", rv, CKR_FUNCTION_NOT_PARALLEL);


	P11T_SECTION("C_CloseSession");

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

	if(session_ro != CK_INVALID)
	{
		rv = (p11t_module_funcs->C_CloseSession)(session_ro);
		P11T_CHECK_RV("Normal call", rv, CKR_OK);

		if(p11t_test_unexpected)
		{
			rv = (p11t_module_funcs->C_GetSessionInfo)(session_ro, &info);
			P11T_CHECK_RV("Check open session was closed", rv, CKR_SESSION_HANDLE_INVALID);

			/*
			 * Note that CKR_SESSION_CLOSED is a valid return in this case.
			 * That should only be returned in the rare case when a session
			 * was closed during the execution of a function. A corner case.
			 */

			rv = (p11t_module_funcs->C_CloseSession)(session_ro);
			P11T_CHECK_RV("Close twice", rv, CKR_SESSION_HANDLE_INVALID);
		}
	}

	if(session_rw != CK_INVALID)
	{
		rv = (p11t_module_funcs->C_CloseSession)(session_rw);
		P11T_CHECK_RV("Normal call", rv, CKR_OK);
	}

	P11T_SECTION("C_CloseAllSessions");

	if(p11t_test_unexpected && !p11t_slot_virtual)
	{
		rv = (p11t_module_funcs->C_CloseAllSessions)((CK_SLOT_ID)-34);
		P11T_CHECK_RV("Invalid slot id", rv, CKR_SLOT_ID_INVALID);
	}

	rv = (p11t_module_funcs->C_CloseAllSessions)(slot);
	P11T_CHECK_RV("Normal call", rv, CKR_OK);

	if(p11t_test_unexpected)
	{
		rv = (p11t_module_funcs->C_GetSessionInfo)(session_ro2, &info);
		P11T_CHECK_RV("Check open session was closed", rv, CKR_SESSION_HANDLE_INVALID);

		rv = (p11t_module_funcs->C_CloseAllSessions)(slot);
		P11T_CHECK_RV("Call when no sessions open", rv, CKR_OK);
	}

	return CONTINUE;
}

static int
test_session_pin(CK_SLOT_ID slot)
{
	P11T_SECTION("C_InitPIN");
	P11T_CHECK_NOTE("Not Tested");

	P11T_SECTION("C_SetPIN");
	P11T_CHECK_NOTE("Not Tested");

	return CONTINUE;
}

static int
test_session_operation_state(CK_SLOT_ID slot)
{
	P11T_SECTION("C_GetOperationState");
	P11T_CHECK_NOTE("Not Tested");

	P11T_SECTION("C_SetOperationState");
	P11T_CHECK_NOTE("Not Tested");

	return CONTINUE;
}

void
p11t_session_tests()
{
	CK_ULONG i;

	for(i = 0; i < p11t_slot_count; ++i)
	{
		CK_SLOT_ID slot = p11t_slot_get_id(i);
		test_session_pin(slot);
		test_session_main(slot);
		test_session_user_login(slot);
		test_session_so_login(slot);
		test_session_operation_state(slot);
	}
}

/* ----------------------------------------------------------------------------------
 * METHODS
 */

CK_SESSION_HANDLE
p11t_session_open(CK_SLOT_ID slot, int readwrite)
{
	CK_SESSION_HANDLE session;
	CK_FLAGS flags;

	assert(p11t_module_funcs);

	flags = CKF_SERIAL_SESSION;
	if(readwrite)
		flags |= CKF_RW_SESSION;

	if((p11t_module_funcs->C_OpenSession)(slot, flags, NULL, NULL, &session) != CKR_OK)
		return CK_INVALID;

	return session;
}

CK_UTF8CHAR_PTR
p11t_session_get_pin(CK_SLOT_ID slot, CK_USER_TYPE user, CK_ULONG_PTR n_pin)
{
	CK_TOKEN_INFO_PTR info;
	const char *pin;

	*n_pin = 0;

	info = p11t_slot_get_token_info(slot);
	if(info && info->flags & CKF_PROTECTED_AUTHENTICATION_PATH)
		return NULL;

	if(user == CKU_USER)
		pin = login_user_pin;
	else if(user == CKU_SO)
		pin = login_so_pin;
	else if(user == CKU_CONTEXT_SPECIFIC)
		pin = login_context_pin;
	else
		pin = NULL;

	if(pin)
		*n_pin = (CK_ULONG)strlen(pin);

	return (CK_UTF8CHAR_PTR)pin;
}


int
p11t_session_login(CK_SESSION_HANDLE session)
{
	CK_SESSION_INFO info;
	CK_UTF8CHAR_PTR pin;
	CK_ULONG n_pin;

	assert(session != CK_INVALID);
	assert(p11t_module_funcs);

	if((p11t_module_funcs->C_GetSessionInfo)(session, &info) != CKR_OK)
		return 0;

	switch(info.state)
	{
	case CKS_RO_USER_FUNCTIONS:
	case CKS_RW_USER_FUNCTIONS:
	case CKS_RW_SO_FUNCTIONS:
		return 1;
	}

	pin = p11t_session_get_pin(info.slotID, CKU_USER, &n_pin);
	if((p11t_module_funcs->C_Login)(session, CKU_USER, pin, n_pin) != CKR_OK)
		return 0;

	return 1;
}

int
p11t_session_logout(CK_SESSION_HANDLE session)
{
	CK_SESSION_INFO info;

	assert(session != CK_INVALID);
	assert(p11t_module_funcs);

	if((p11t_module_funcs->C_GetSessionInfo)(session, &info) != CKR_OK)
		return 0;

	switch(info.state)
	{
	case CKS_RO_PUBLIC_SESSION:
	case CKS_RW_PUBLIC_SESSION:
		return 1;
	}

	if((p11t_module_funcs->C_Logout)(session) != CKR_OK)
		return 0;

	return 1;
}

int
p11t_session_close(CK_SESSION_HANDLE session)
{
	CK_RV rv;

	assert(session != CK_INVALID);
	assert(p11t_module_funcs);

	rv = (p11t_module_funcs->C_CloseSession)(session);
	switch(rv)
	{
	case CKR_SESSION_HANDLE_INVALID:
	case CKR_SESSION_CLOSED:
	case CKR_OK:
		return 1;
	default:
		return 0;
	}
}

int
p11t_session_close_all(CK_SLOT_ID slot)
{
	assert(p11t_module_funcs);

	if((p11t_module_funcs->C_CloseAllSessions)(slot) == CKR_OK)
		return 1;
	return 0;
}