/* This Source Code Form is subject to the terms of the Mozilla Public

* License, v. 2.0. If a copy of the MPL was not distributed with this file,

* You can obtain one at http://mozilla.org/MPL/2.0/. */

#include "pk11pub.h"

#include "nssutil.h"

#include <stdio.h>

#include "prerror.h"

#include "nss.h"

#include "gtest/gtest.h"

#include "scoped_ptrs.h"

#include "cpputil.h"

#include "databuffer.h"

#include "util.h"

namespace nss_test {

static const uint8_t kIv[] = {0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77,

0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff,

0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77};

static const uint8_t kInput[] = {0x77, 0x66, 0x55, 0x44, 0x33, 0x22, 0x11, 0x00,

0xff, 0xee, 0xdd, 0xcc, 0xbb, 0xaa, 0x99, 0x88,

0x77, 0x66, 0x55, 0x44, 0x33, 0x22, 0x11, 0x00};

class EncryptDeriveTest

: public ::testing::Test,

public ::testing::WithParamInterface<CK_MECHANISM_TYPE> {

public:

void TestEncryptDerive() {

ScopedPK11SymKey derived_key(PK11_Derive(key_.get(), derive_mech(),

derive_param(), encrypt_mech(),

CKA_DECRYPT, keysize()));

ASSERT_TRUE(derived_key);

uint8_t derived_key_data[keysize()];

GetKeyData(derived_key, derived_key_data, sizeof(derived_key_data));

RemoveChecksum(derived_key_data);

uint8_t reference_key_data[keysize()];

unsigned int reference_len = 0;

SECStatus rv = PK11_Encrypt(key_.get(), encrypt_mech(), encrypt_param(),

reference_key_data, &reference_len, keysize(),

kInput, 16);

ASSERT_EQ(SECSuccess, rv);

ASSERT_EQ(keysize(), static_cast<size_t>(reference_len));

RemoveChecksum(reference_key_data);

EXPECT_EQ(DataBuffer(reference_key_data, keysize()),

DataBuffer(derived_key_data, keysize()));

}

protected:

unsigned int keysize() const { return 16; }

private:

CK_MECHANISM_TYPE encrypt_mech() const { return GetParam(); }

CK_MECHANISM_TYPE derive_mech() const {

switch (encrypt_mech()) {

case CKM_DES3_ECB:

return CKM_DES3_ECB_ENCRYPT_DATA;

case CKM_DES3_CBC:

return CKM_DES3_CBC_ENCRYPT_DATA;

case CKM_AES_ECB:

return CKM_AES_ECB_ENCRYPT_DATA;

case CKM_AES_CBC:

return CKM_AES_CBC_ENCRYPT_DATA;

case CKM_CAMELLIA_ECB:

return CKM_CAMELLIA_ECB_ENCRYPT_DATA;

case CKM_CAMELLIA_CBC:

return CKM_CAMELLIA_CBC_ENCRYPT_DATA;

case CKM_SEED_ECB:

return CKM_SEED_ECB_ENCRYPT_DATA;

case CKM_SEED_CBC:

return CKM_SEED_CBC_ENCRYPT_DATA;

default:

ADD_FAILURE() << "Unknown mechanism";

break;

}

return CKM_INVALID_MECHANISM;

}

SECItem* derive_param() const {

static CK_AES_CBC_ENCRYPT_DATA_PARAMS aes_data;

static CK_DES_CBC_ENCRYPT_DATA_PARAMS des_data;

static CK_KEY_DERIVATION_STRING_DATA string_data;

static SECItem param = {siBuffer, NULL, 0};

switch (encrypt_mech()) {

case CKM_DES3_ECB:

case CKM_AES_ECB:

case CKM_CAMELLIA_ECB:

case CKM_SEED_ECB:

string_data.pData = toUcharPtr(kInput);

string_data.ulLen = 16;

param.data = reinterpret_cast<uint8_t*>(&string_data);

param.len = sizeof(string_data);

break;

case CKM_DES3_CBC:

des_data.pData = toUcharPtr(kInput);

des_data.length = 16;

PORT_Memcpy(des_data.iv, kIv, 8);

param.data = reinterpret_cast<uint8_t*>(&des_data);

param.len = sizeof(des_data);

break;

case CKM_AES_CBC:

case CKM_CAMELLIA_CBC:

case CKM_SEED_CBC:

aes_data.pData = toUcharPtr(kInput);

aes_data.length = 16;

PORT_Memcpy(aes_data.iv, kIv, 16);

param.data = reinterpret_cast<uint8_t*>(&aes_data);

param.len = sizeof(aes_data);

break;

default:

ADD_FAILURE() << "Unknown mechanism";

break;

}

return &param;

}

SECItem* encrypt_param() const {

static SECItem param = {siBuffer, NULL, 0};

switch (encrypt_mech()) {

case CKM_DES3_ECB:

case CKM_AES_ECB:

case CKM_CAMELLIA_ECB:

case CKM_SEED_ECB:

// No parameter needed here.

break;

case CKM_DES3_CBC:

case CKM_AES_CBC:

case CKM_CAMELLIA_CBC:

case CKM_SEED_CBC:

param.data = toUcharPtr(kIv);

param.len = 16;

break;

default:

ADD_FAILURE() << "Unknown mechanism";

break;

}

return &param;

}

virtual void SetUp() {

slot_.reset(PK11_GetBestSlot(derive_mech(), NULL));

ASSERT_TRUE(slot_);

key_.reset(PK11_TokenKeyGenWithFlags(slot_.get(), encrypt_mech(), NULL,

keysize(), NULL,

CKF_ENCRYPT | CKF_DERIVE, 0, NULL));

ASSERT_TRUE(key_);

}

void GetKeyData(ScopedPK11SymKey& key, uint8_t* buf, size_t max_len) const {

ASSERT_EQ(SECSuccess, PK11_ExtractKeyValue(key.get()));

SECItem* data = PK11_GetKeyData(key.get());

ASSERT_TRUE(data);

ASSERT_EQ(max_len, static_cast<size_t>(data->len));

PORT_Memcpy(buf, data->data, data->len);

}

// Remove checksum if the key is a 3DES key.

void RemoveChecksum(uint8_t* key_data) const {

if (encrypt_mech() != CKM_DES3_CBC && encrypt_mech() != CKM_DES3_ECB) {

return;

}

for (size_t i = 0; i < keysize(); ++i) {

key_data[i] &= 0xfe;

}

}

ScopedPK11SlotInfo slot_;

ScopedPK11SymKey key_;

};

TEST_P(EncryptDeriveTest, Test) {

TestEncryptDerive();

}

static const CK_MECHANISM_TYPE kEncryptDeriveMechanisms[] = {

CKM_DES3_ECB, CKM_DES3_CBC, CKM_AES_ECB, CKM_AES_ECB, CKM_AES_CBC,

CKM_CAMELLIA_ECB, CKM_CAMELLIA_CBC, CKM_SEED_ECB, CKM_SEED_CBC};

INSTANTIATE_TEST_CASE_P(EncryptDeriveTests, EncryptDeriveTest,

::testing::ValuesIn(kEncryptDeriveMechanisms));

// This class handles the case where 3DES takes a 192-bit key

// where all 24 octets will be used.

class EncryptDerive3Test : public EncryptDeriveTest {

protected:

unsigned int keysize() const { return 24; }

};

TEST_P(EncryptDerive3Test, Test) {

TestEncryptDerive();

}

static const CK_MECHANISM_TYPE kDES3EncryptDeriveMechanisms[] = {CKM_DES3_ECB,

CKM_DES3_CBC};

INSTANTIATE_TEST_CASE_P(Encrypt3DeriveTests, EncryptDerive3Test,

::testing::ValuesIn(kDES3EncryptDeriveMechanisms));

} // namespace nss_test