#include "p11-tests-lib.h"

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

#include <openssl/dh.h>

static void
test_dh_key_pair_gen (CK_SLOT_ID slot, CK_MECHANISM_TYPE mech, CK_MECHANISM_INFO_PTR info)
{
	CK_SESSION_HANDLE session;

	session = p11t_session_open(slot, 0);
	if(session == CK_INVALID || !p11t_session_login (session))
		return;

	p11t_dh_test_generate_pair (session);
}

void
p11t_dh_tests (void)
{
	p11t_slot_for_each_mech (CKM_DH_PKCS_KEY_PAIR_GEN, test_dh_key_pair_gen);
}

int
p11t_dh_test_generate_pair (CK_SESSION_HANDLE session)
{
	CK_ATTRIBUTE pub_attrs[3];
	CK_ATTRIBUTE priv_attrs[3];
	CK_MECHANISM mech;
	CK_OBJECT_HANDLE pub_key;
	CK_OBJECT_HANDLE priv_key;
	CK_ULONG bits;
	CK_RV rv;
	DH *dh;

	P11T_SECTION ("C_GenerateKeyPair");

	dh = DH_generate_parameters (256, 2, NULL, NULL);
	assert (dh);

	pub_attrs[0].type = CKA_PRIME;
	pub_attrs[0].ulValueLen = BN_num_bytes (dh->p);
	pub_attrs[0].pValue = alloca (pub_attrs[0].ulValueLen);
	BN_bn2bin (dh->p, (unsigned char*)pub_attrs[0].pValue);

	pub_attrs[1].type = CKA_BASE;
	pub_attrs[1].ulValueLen = BN_num_bytes (dh->g);
	pub_attrs[1].pValue = alloca (pub_attrs[1].ulValueLen);
	BN_bn2bin (dh->g, (unsigned char*)pub_attrs[1].pValue);

	priv_attrs[0].type = CKA_PRIME;
	priv_attrs[0].ulValueLen = BN_num_bytes (dh->p);
	priv_attrs[0].pValue = alloca (priv_attrs[0].ulValueLen);
	BN_bn2bin (dh->p, (unsigned char*)priv_attrs[0].pValue);

	priv_attrs[1].type = CKA_BASE;
	priv_attrs[1].ulValueLen = BN_num_bytes (dh->g);
	priv_attrs[1].pValue = alloca (priv_attrs[1].ulValueLen);
	BN_bn2bin (dh->g, (unsigned char*)priv_attrs[1].pValue);

	priv_attrs[2].type = CKA_VALUE_BITS;
	priv_attrs[2].ulValueLen = sizeof (bits);
	priv_attrs[2].pValue = &bits;

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

	if (p11t_test_unexpected) {
		rv = (p11t_module_funcs->C_GenerateKeyPair) (session, &mech, pub_attrs, 0, priv_attrs, 0, &pub_key, &priv_key);
		P11T_CHECK_RV ("DH Key Pair without CKA_PRIME", rv, CKR_TEMPLATE_INCOMPLETE);

		bits = 1024;
		rv = (p11t_module_funcs->C_GenerateKeyPair) (session, &mech, pub_attrs, 2, priv_attrs, 3, &pub_key, &priv_key);
		P11T_CHECK_RV ("DH Key Pair with CKA_VALUE_BITS larger than CKA_PRIME", rv, CKR_TEMPLATE_INCONSISTENT);
	}

	bits = 256;
	rv = (p11t_module_funcs->C_GenerateKeyPair) (session, &mech, pub_attrs, 2, priv_attrs, 3, &pub_key, &priv_key);
	P11T_CHECK_RV ("DH Key Pair", rv, CKR_OK);

	p11t_key_test (session, pub_key, CKO_PUBLIC_KEY);
	p11t_key_test (session, priv_key, CKO_PRIVATE_KEY);

	/* Test corner cases */
	return p11t_object_generate_pair_bad (session, CKM_DH_PKCS_KEY_PAIR_GEN);
}