/*
 * COPYRIGHT (c) International Business Machines Corp. 2017
 *
 * This program is provided under the terms of the Common Public License,
 * version 1.0 (CPL-1.0). Any use, reproduction or distribution for this
 * software constitutes recipient's acceptance of CPL-1.0 terms which can be
 * found in the file LICENSE file or at
 * https://opensource.org/licenses/cpl1.0.php
 */

/*
 * Testcase for
 * C_GetOperationState / C_SetOperationState
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <memory.h>
#include <time.h>

#include "pkcs11types.h"
#include "regress.h"
#include "common.c"


CK_BYTE_PTR alloc_random_buf(CK_SESSION_HANDLE sess, CK_LONG nbytes)
{
	CK_RV rc;
	CK_BYTE_PTR ptr = malloc(nbytes);
	if (ptr == NULL) {
		testcase_error("malloc(%lu) failed", nbytes);
		return NULL;
	}
	rc = funcs->C_GenerateRandom(sess, ptr, nbytes);
	if (rc != CKR_OK) {
		testcase_error("C_GenerateRandom() rc=%s", p11_get_ckr(rc));
		free(ptr);
		return NULL;
	}
	return ptr;
}


int sess_opstate_funcs(int loops)
{
	CK_SESSION_HANDLE s1, s2;
	CK_SLOT_ID	  slot_id = SLOT_ID;
	CK_ULONG	  flags;
	CK_RV		  rc;
	int               i, counter, rbytes;
	CK_BYTE           *rdata = NULL;
	CK_MECHANISM      mech1 = {CKM_SHA256, 0, 0};
	CK_MECHANISM      mech2 = {CKM_SHA_1, 0, 0};
	CK_ULONG          r1hlen, r2hlen, hlen;
	CK_BYTE           r1hash[32], r2hash[32], hash[32];
	CK_ULONG          opstatelen;
	CK_BYTE           *opstate = NULL;

	// open 2 sessions
	flags = CKF_SERIAL_SESSION | CKF_RW_SESSION;
	rc = funcs->C_OpenSession(slot_id, flags, NULL, NULL, &s1);
	if (rc != CKR_OK) {
		testcase_error("C_OpenSession() rc=%s", p11_get_ckr(rc));
		goto out;
	}
	rc = funcs->C_OpenSession(slot_id, flags, NULL, NULL, &s2);
	if (rc != CKR_OK) {
		testcase_error("C_OpenSession() rc=%s", p11_get_ckr(rc));
		goto out;
	}

	// init digest for both sessions
	rc = funcs->C_DigestInit(s1, &mech1);
	if (rc != CKR_OK) {
		testcase_error("C_DigestInit rc=%s", p11_get_ckr(rc));
		goto out;
	}
	rc = funcs->C_DigestInit(s2, &mech1);
	if (rc != CKR_OK) {
		testcase_error("C_DigestInit rc=%s", p11_get_ckr(rc));
		goto out;
	}

	// now loop over some digest updates
	for (counter = 0; counter < loops; counter ++) {

		// create some random data
		rbytes = random() % sizeof(rdata);
		rdata = alloc_random_buf(s1, rbytes);
		if (!rdata) goto out;

		// digest update on session 1
		rc = funcs->C_DigestUpdate(s1, rdata, rbytes);
		if (rc != CKR_OK) {
			testcase_error("C_DigestUpdate rc=%s", p11_get_ckr(rc));
			goto out;
		}

		// restore op state on session 2
		if (opstate != NULL) {
			rc = funcs->C_SetOperationState(s2, opstate, opstatelen, 0, 0);
			if (rc != CKR_OK) {
				testcase_error("C_SetOperationState rc=%s", p11_get_ckr(rc));
				goto out;
			}
			free(opstate);
			opstate = NULL;
		}

		// digest update on session 2
		rc = funcs->C_DigestUpdate(s2, rdata, rbytes);
		if (rc != CKR_OK) {
			testcase_error("C_DigestUpdate rc=%s", p11_get_ckr(rc));
			goto out;
		}

		// fetch op state on session 2
		opstatelen = 0;
		rc = funcs->C_GetOperationState(s2, NULL, &opstatelen);
		if (rc != CKR_OK) {
			testcase_error("C_GetOperationState rc=%s", p11_get_ckr(rc));
			goto out;
		}
		opstate = malloc(opstatelen);
		if (opstate == NULL) {
			testcase_error("malloc(%lu) failed", opstatelen);
			goto out;
		}
		rc = funcs->C_GetOperationState(s2, opstate, &opstatelen);
		if (rc != CKR_OK) {
			testcase_error("C_GetOperationState rc=%s", p11_get_ckr(rc));
			goto out;
		}

		free(rdata);
		rdata = NULL;

		// now do something different on session 2, but first
		// we have to wipe out the started digest operation

		hlen = sizeof(hash);
		rc = funcs->C_DigestFinal(s2, hash, &hlen);
		if (rc != CKR_OK) {
			testcase_error("C_DigestFinal rc=%s", p11_get_ckr(rc));
			goto out;
		}

		// so now let's do a digest init/update/finish
		// to randomize the memory a little

		rc = funcs->C_DigestInit(s2, &mech2);
		if (rc != CKR_OK) {
			testcase_error("C_DigestInit rc=%s", p11_get_ckr(rc));
			goto out;
		}
		for (i = 0; i < loops; i++) {
			rbytes = random() % sizeof(rdata);
			rdata = alloc_random_buf(s1, rbytes);
			if (!rdata) goto out;
			rc = funcs->C_DigestUpdate(s2, rdata, rbytes);
			if (rc != CKR_OK) {
				testcase_error("C_DigestUpdate rc=%s", p11_get_ckr(rc));
				goto out;
			}
			free(rdata);
			rdata = NULL;
		}
		hlen = sizeof(hash);
		rc = funcs->C_DigestFinal(s2, hash, &hlen);
		if (rc != CKR_OK) {
			testcase_error("C_DigestFinal rc=%s", p11_get_ckr(rc));
			goto out;
		}

	}

	// restore op state on session 2
	rc = funcs->C_SetOperationState(s2, opstate, opstatelen, 0, 0);
	if (rc != CKR_OK) {
		testcase_error("C_SetOperationState rc=%s", p11_get_ckr(rc));
		goto out;
	}

	// digest finish
	r1hlen = sizeof(r1hash);
	rc = funcs->C_DigestFinal(s1, r1hash, &r1hlen);
	if (rc != CKR_OK) {
		testcase_error("C_DigestFinal rc=%s", p11_get_ckr(rc));
		goto out;
	}
	r2hlen = sizeof(r2hash);
	rc = funcs->C_DigestFinal(s2, r2hash, &r2hlen);
	if (rc != CKR_OK) {
		testcase_error("C_DigestFinal rc=%s", p11_get_ckr(rc));
		goto out;
	}

	// check both hashes
	if (r1hlen != r2hlen) {
		testcase_fail("hash length differ");
		goto out;
	}
	if (memcmp(r1hash, r2hash, r1hlen) != 0) {
		testcase_fail("hash values differs");
		goto out;
	}

	testcase_pass("Get/SetOperationState digest test");

out:
	if (opstate) free(opstate);
	if (rdata) free(rdata);
	funcs->C_CloseAllSessions(slot_id);
	return rc;
}


int main(int argc, char **argv)
{
	CK_C_INITIALIZE_ARGS cinit_args;
	int rc, i, j, loops = 0;
	CK_RV rv;

	SLOT_ID = 0;
	no_init = FALSE;

	srandom(time(0));

	for (i=0; i < argc; i++) {
		if (strncmp(argv[i], "loops=", 6) == 0) {
			sscanf(argv[i]+6, "%i", &loops);
			for (j=i; j < argc; j++)
				argv[j] = argv[j+1];
			argc--;
		}
	}
	if (loops < 1) loops = 100;

	rc = do_ParseArgs(argc, argv);
	if (rc != 1)
		return rc;

	printf("Using slot #%lu...\n\n", SLOT_ID );
	printf("With option: no_init: %d\n", no_init);
	printf("Running %d loops...\n", loops );

	rc = do_GetFunctionList();
	if (!rc) {
		PRINT_ERR("ERROR do_GetFunctionList() Failed , rc = 0x%0x\n", rc);
		return rc;
	}

	memset( &cinit_args, 0x0, sizeof(cinit_args) );
	cinit_args.flags = CKF_OS_LOCKING_OK;

	funcs->C_Initialize( &cinit_args );

	{
		CK_SESSION_HANDLE  hsess = 0;

		rc = funcs->C_GetFunctionStatus(hsess);
		if (rc  != CKR_FUNCTION_NOT_PARALLEL)
			return rc;

		rc = funcs->C_CancelFunction(hsess);
		if (rc  != CKR_FUNCTION_NOT_PARALLEL)
			return rc;

	}

	rv = sess_opstate_funcs(loops);

	/* make sure we return non-zero if rv is non-zero */
	return ((rv==0) || (rv % 256) ? rv : -1);
}
