1 files changed, 69 insertions, 71 deletions

diff --git a/security/nss/gtests/pk11_gtest/pk11_signature_test.h b/security/nss/gtests/pk11_gtest/pk11_signature_test.h
index e6a0a9c57..b14104371 100644
--- a/security/nss/gtests/pk11_gtest/pk11_signature_test.h
+++ b/security/nss/gtests/pk11_gtest/pk11_signature_test.h
@@ -9,26 +9,37 @@
 
 #include "cpputil.h"
 #include "scoped_ptrs.h"
+#include "databuffer.h"
 
 #include "gtest/gtest.h"
 
 namespace nss_test {
 
+// For test vectors.
+struct Pkcs11SignatureTestParams {
+  const DataBuffer pkcs8_;
+  const DataBuffer spki_;
+  const DataBuffer data_;
+  const DataBuffer signature_;
+};
+
 class Pk11SignatureTest : public ::testing::Test {
  protected:
-  virtual CK_MECHANISM_TYPE mechanism() = 0;
-  virtual SECItem* parameters() = 0;
-  virtual SECOidTag hashOID() = 0;
+  Pk11SignatureTest(CK_MECHANISM_TYPE mechanism, SECOidTag hash_oid)
+      : mechanism_(mechanism), hash_oid_(hash_oid) {}
 
-  ScopedSECKEYPrivateKey ImportPrivateKey(const uint8_t* pkcs8,
-                                          size_t pkcs8_len) {
+  virtual const SECItem* parameters() const { return nullptr; }
+  CK_MECHANISM_TYPE mechanism() const { return mechanism_; }
+
+  ScopedSECKEYPrivateKey ImportPrivateKey(const DataBuffer& pkcs8) {
     ScopedPK11SlotInfo slot(PK11_GetInternalSlot());
     if (!slot) {
+      ADD_FAILURE() << "No slot";
       return nullptr;
     }
 
-    SECItem pkcs8Item = {siBuffer, toUcharPtr(pkcs8),
-                         static_cast<unsigned int>(pkcs8_len)};
+    SECItem pkcs8Item = {siBuffer, toUcharPtr(pkcs8.data()),
+                         static_cast<unsigned int>(pkcs8.len())};
 
     SECKEYPrivateKey* key = nullptr;
     SECStatus rv = PK11_ImportDERPrivateKeyInfoAndReturnKey(
@@ -42,9 +53,9 @@ class Pk11SignatureTest : public ::testing::Test {
     return ScopedSECKEYPrivateKey(key);
   }
 
-  ScopedSECKEYPublicKey ImportPublicKey(const uint8_t* spki, size_t spki_len) {
-    SECItem spkiItem = {siBuffer, toUcharPtr(spki),
-                        static_cast<unsigned int>(spki_len)};
+  ScopedSECKEYPublicKey ImportPublicKey(const DataBuffer& spki) {
+    SECItem spkiItem = {siBuffer, toUcharPtr(spki.data()),
+                        static_cast<unsigned int>(spki.len())};
 
     ScopedCERTSubjectPublicKeyInfo certSpki(
         SECKEY_DecodeDERSubjectPublicKeyInfo(&spkiItem));
@@ -52,87 +63,74 @@ class Pk11SignatureTest : public ::testing::Test {
     return ScopedSECKEYPublicKey(SECKEY_ExtractPublicKey(certSpki.get()));
   }
 
-  ScopedSECItem ComputeHash(const uint8_t* data, size_t len) {
-    unsigned int hLen = HASH_ResultLenByOidTag(hashOID());
-    ScopedSECItem hash(SECITEM_AllocItem(nullptr, nullptr, hLen));
-    if (!hash) {
-      return nullptr;
-    }
-
-    SECStatus rv = PK11_HashBuf(hashOID(), hash->data, data, len);
-    if (rv != SECSuccess) {
-      return nullptr;
-    }
-
-    return hash;
+  bool ComputeHash(const DataBuffer& data, DataBuffer* hash) {
+    hash->Allocate(static_cast<size_t>(HASH_ResultLenByOidTag(hash_oid_)));
+    SECStatus rv =
+        PK11_HashBuf(hash_oid_, hash->data(), data.data(), data.len());
+    return rv == SECSuccess;
   }
 
-  ScopedSECItem SignHashedData(ScopedSECKEYPrivateKey& privKey,
-                               ScopedSECItem& hash) {
-    unsigned int sLen = PK11_SignatureLen(privKey.get());
-    ScopedSECItem sig(SECITEM_AllocItem(nullptr, nullptr, sLen));
-    if (!sig) {
-      return nullptr;
-    }
-
-    SECStatus rv = PK11_SignWithMechanism(privKey.get(), mechanism(),
-                                          parameters(), sig.get(), hash.get());
-    if (rv != SECSuccess) {
-      return nullptr;
-    }
-
-    return sig;
+  bool SignHashedData(ScopedSECKEYPrivateKey& privKey, const DataBuffer& hash,
+                      DataBuffer* sig) {
+    SECItem hashItem = {siBuffer, toUcharPtr(hash.data()),
+                        static_cast<unsigned int>(hash.len())};
+    int sigLen = PK11_SignatureLen(privKey.get());
+    EXPECT_LT(0, sigLen);
+    sig->Allocate(static_cast<size_t>(sigLen));
+    SECItem sigItem = {siBuffer, toUcharPtr(sig->data()),
+                       static_cast<unsigned int>(sig->len())};
+    SECStatus rv = PK11_SignWithMechanism(privKey.get(), mechanism_,
+                                          parameters(), &sigItem, &hashItem);
+    return rv == SECSuccess;
   }
 
-  ScopedSECItem ImportPrivateKeyAndSignHashedData(const uint8_t* pkcs8,
-                                                  size_t pkcs8_len,
-                                                  const uint8_t* data,
-                                                  size_t data_len) {
-    ScopedSECKEYPrivateKey privKey(ImportPrivateKey(pkcs8, pkcs8_len));
+  bool ImportPrivateKeyAndSignHashedData(const DataBuffer& pkcs8,
+                                         const DataBuffer& data,
+                                         DataBuffer* sig) {
+    ScopedSECKEYPrivateKey privKey(ImportPrivateKey(pkcs8));
     if (!privKey) {
-      return nullptr;
+      return false;
     }
 
-    ScopedSECItem hash(ComputeHash(data, data_len));
-    if (!hash) {
-      return nullptr;
+    DataBuffer hash;
+    if (!ComputeHash(data, &hash)) {
+      ADD_FAILURE() << "Failed to compute hash";
+      return false;
     }
-
-    return ScopedSECItem(SignHashedData(privKey, hash));
+    return SignHashedData(privKey, hash, sig);
   }
 
-  void Verify(const uint8_t* spki, size_t spki_len, const uint8_t* data,
-              size_t data_len, const uint8_t* sig, size_t sig_len) {
-    ScopedSECKEYPublicKey pubKey(ImportPublicKey(spki, spki_len));
+  void Verify(const Pkcs11SignatureTestParams& params, const DataBuffer& sig) {
+    ScopedSECKEYPublicKey pubKey(ImportPublicKey(params.spki_));
     ASSERT_TRUE(pubKey);
 
-    ScopedSECItem hash(ComputeHash(data, data_len));
-    ASSERT_TRUE(hash);
-
-    SECItem sigItem = {siBuffer, toUcharPtr(sig),
-                       static_cast<unsigned int>(sig_len)};
+    DataBuffer hash;
+    ASSERT_TRUE(ComputeHash(params.data_, &hash));
 
     // Verify.
+    SECItem hashItem = {siBuffer, toUcharPtr(hash.data()),
+                        static_cast<unsigned int>(hash.len())};
+    SECItem sigItem = {siBuffer, toUcharPtr(sig.data()),
+                       static_cast<unsigned int>(sig.len())};
     SECStatus rv = PK11_VerifyWithMechanism(
-        pubKey.get(), mechanism(), parameters(), &sigItem, hash.get(), nullptr);
+        pubKey.get(), mechanism_, parameters(), &sigItem, &hashItem, nullptr);
     EXPECT_EQ(rv, SECSuccess);
   }
 
-  void SignAndVerify(const uint8_t* pkcs8, size_t pkcs8_len,
-                     const uint8_t* spki, size_t spki_len, const uint8_t* data,
-                     size_t data_len) {
-    ScopedSECItem sig(
-        ImportPrivateKeyAndSignHashedData(pkcs8, pkcs8_len, data, data_len));
-    ASSERT_TRUE(sig);
-
-    Verify(spki, spki_len, data, data_len, sig->data, sig->len);
+  void Verify(const Pkcs11SignatureTestParams& params) {
+    Verify(params, params.signature_);
   }
-};
 
-#define SIG_TEST_VECTOR_VERIFY(spki, data, sig) \
-  Verify(spki, sizeof(spki), data, sizeof(data), sig, sizeof(sig));
+  void SignAndVerify(const Pkcs11SignatureTestParams& params) {
+    DataBuffer sig;
+    ASSERT_TRUE(
+        ImportPrivateKeyAndSignHashedData(params.pkcs8_, params.data_, &sig));
+    Verify(params, sig);
+  }
 
-#define SIG_TEST_VECTOR_SIGN_VERIFY(pkcs8, spki, data) \
-  SignAndVerify(pkcs8, sizeof(pkcs8), spki, sizeof(spki), data, sizeof(data));
+ private:
+  CK_MECHANISM_TYPE mechanism_;
+  SECOidTag hash_oid_;
+};
 
 }  // namespace nss_test