1 files changed, 37 insertions, 28 deletions

diff --git a/gtests/pk11_gtest/pk11_seed_cbc_unittest.cc b/gtests/pk11_gtest/pk11_seed_cbc_unittest.cc
index dd90b7853..45f77278e 100644
--- a/gtests/pk11_gtest/pk11_seed_cbc_unittest.cc
+++ b/gtests/pk11_gtest/pk11_seed_cbc_unittest.cc
@@ -14,35 +14,40 @@
 #include "util.h"
 
 namespace nss_test {
-class Pkcs11SeedCbcTest : public ::testing::Test {
+class Pkcs11SeedTest : public ::testing::Test {
  protected:
-  enum class Action { Encrypt, Decrypt };
-
-  SECStatus EncryptDecryptSeed(Action action, unsigned int input_size,
-                               unsigned int output_size) {
+  void EncryptDecryptSeed(SECStatus expected, unsigned int input_size,
+                          unsigned int output_size,
+                          CK_MECHANISM_TYPE mech = CKM_SEED_CBC) {
     // Generate a random key.
     ScopedPK11SlotInfo slot(PK11_GetInternalSlot());
     ScopedPK11SymKey sym_key(
-        PK11_KeyGen(slot.get(), kMech, nullptr, 16, nullptr));
+        PK11_KeyGen(slot.get(), mech, nullptr, 16, nullptr));
     EXPECT_TRUE(!!sym_key);
 
-    std::vector<uint8_t> data(input_size);
+    std::vector<uint8_t> plaintext(input_size, 0xFF);
     std::vector<uint8_t> init_vector(16);
-    std::vector<uint8_t> output(output_size);
-    SECItem params = {siBuffer, init_vector.data(),
-                      (unsigned int)init_vector.size()};
+    std::vector<uint8_t> ciphertext(output_size, 0);
+    SECItem iv_param = {siBuffer, init_vector.data(),
+                        (unsigned int)init_vector.size()};
+    std::vector<uint8_t> decrypted(output_size, 0);
 
-    // Try to encrypt/decrypt.
+    // Try to encrypt, decrypt if positive test.
     unsigned int output_len = 0;
-    if (action == Action::Encrypt) {
-      return PK11_Encrypt(sym_key.get(), kMech, &params, output.data(),
-                          &output_len, output_size, data.data(), data.size());
-    } else {
-      return PK11_Decrypt(sym_key.get(), kMech, &params, output.data(),
-                          &output_len, output_size, data.data(), data.size());
+    EXPECT_EQ(expected,
+              PK11_Encrypt(sym_key.get(), mech, &iv_param, ciphertext.data(),
+                           &output_len, output_size, plaintext.data(),
+                           plaintext.size()));
+
+    if (expected == SECSuccess) {
+      EXPECT_EQ(expected,
+                PK11_Decrypt(sym_key.get(), mech, &iv_param, decrypted.data(),
+                             &output_len, output_size, ciphertext.data(),
+                             output_len));
+      decrypted.resize(output_len);
+      EXPECT_EQ(plaintext, decrypted);
     }
   }
-  const CK_MECHANISM_TYPE kMech = CKM_SEED_CBC;
 };
 
 // The intention here is to test the arguments of these functions
@@ -51,21 +56,25 @@ class Pkcs11SeedCbcTest : public ::testing::Test {
 // The input data size must be multiple of 16.
 // If not, some padding should be added.
 // The output size must be at least the size of input data.
-TEST_F(Pkcs11SeedCbcTest, SeedCBC_ValidArgs) {
-  EXPECT_EQ(SECSuccess, EncryptDecryptSeed(Action::Encrypt, 16, 16));
-  EXPECT_EQ(SECSuccess, EncryptDecryptSeed(Action::Decrypt, 16, 16));
+TEST_F(Pkcs11SeedTest, CBC_ValidArgs) {
+  EncryptDecryptSeed(SECSuccess, 16, 16);
   // No problem if maxLen is bigger than input data.
-  EXPECT_EQ(SECSuccess, EncryptDecryptSeed(Action::Encrypt, 16, 32));
-  EXPECT_EQ(SECSuccess, EncryptDecryptSeed(Action::Decrypt, 16, 32));
+  EncryptDecryptSeed(SECSuccess, 16, 32);
 }
 
-TEST_F(Pkcs11SeedCbcTest, SeedCBC_InvalidArgs) {
+TEST_F(Pkcs11SeedTest, CBC_InvalidArgs) {
   // maxLen lower than input data.
-  EXPECT_EQ(SECFailure, EncryptDecryptSeed(Action::Encrypt, 16, 10));
-  EXPECT_EQ(SECFailure, EncryptDecryptSeed(Action::Decrypt, 16, 10));
+  EncryptDecryptSeed(SECFailure, 16, 10);
   // input data not multiple of SEED_BLOCK_SIZE (16)
-  EXPECT_EQ(SECFailure, EncryptDecryptSeed(Action::Encrypt, 17, 32));
-  EXPECT_EQ(SECFailure, EncryptDecryptSeed(Action::Decrypt, 17, 32));
+  EncryptDecryptSeed(SECFailure, 17, 32);
+}
+
+TEST_F(Pkcs11SeedTest, ECB_Singleblock) {
+  EncryptDecryptSeed(SECSuccess, 16, 16, CKM_SEED_ECB);
+}
+
+TEST_F(Pkcs11SeedTest, ECB_Multiblock) {
+  EncryptDecryptSeed(SECSuccess, 64, 64, CKM_SEED_ECB);
 }
 
 }  // namespace nss_test
\ No newline at end of file