1 files changed, 80 insertions, 0 deletions
diff --git a/security/nss/gtests/pk11_gtest/pk11_key_unittest.cc b/security/nss/gtests/pk11_gtest/pk11_key_unittest.cc
new file mode 100644
index 0000000000..9d5c96f3b2
--- /dev/null
+++ b/security/nss/gtests/pk11_gtest/pk11_key_unittest.cc
@@ -0,0 +1,80 @@
+/* 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 <memory>
+#include "nss.h"
+#include "pk11pub.h"
+#include "pk11pqg.h"
+#include "prerror.h"
+#include "secoid.h"
+
+#include "gtest/gtest.h"
+#include "nss_scoped_ptrs.h"
+#include "pk11_keygen.h"
+
+namespace nss_test {
+
+class Pkcs11NullKeyTestBase : public ::testing::Test {
+ protected:
+  // This constructs a key pair, then erases the public value from the public
+  // key.  NSS should reject this.
+  void Test(const Pkcs11KeyPairGenerator& generator,
+            CK_MECHANISM_TYPE dh_mech) {
+    ScopedSECKEYPrivateKey priv;
+    ScopedSECKEYPublicKey pub;
+    generator.GenerateKey(&priv, &pub);
+    ASSERT_TRUE(priv);
+
+    // These don't leak because they are allocated to the arena associated with
+    // the public key.
+    SECItem* pub_val = nullptr;
+    switch (SECKEY_GetPublicKeyType(pub.get())) {
+      case rsaKey:
+        pub_val = &pub->u.rsa.modulus;
+        break;
+
+      case dsaKey:
+        pub_val = &pub->u.dsa.publicValue;
+        break;
+
+      case dhKey:
+        pub_val = &pub->u.dh.publicValue;
+        break;
+
+      case ecKey:
+        pub_val = &pub->u.ec.publicValue;
+        break;
+
+      default:
+        FAIL() << "Unknown key type " << SECKEY_GetPublicKeyType(pub.get());
+    }
+    pub_val->data = nullptr;
+    pub_val->len = 0;
+
+    ScopedPK11SymKey symKey(PK11_PubDeriveWithKDF(
+        priv.get(), pub.get(), false, nullptr, nullptr, dh_mech,
+        CKM_SHA512_HMAC, CKA_DERIVE, 0, CKD_NULL, nullptr, nullptr));
+    ASSERT_FALSE(symKey);
+  }
+};
+
+class Pkcs11DhNullKeyTest : public Pkcs11NullKeyTestBase {};
+TEST_F(Pkcs11DhNullKeyTest, UseNullPublicValue) {
+  Test(Pkcs11KeyPairGenerator(CKM_DH_PKCS_KEY_PAIR_GEN), CKM_DH_PKCS_DERIVE);
+}
+
+class Pkcs11EcdhNullKeyTest : public Pkcs11NullKeyTestBase,
+                              public ::testing::WithParamInterface<SECOidTag> {
+};
+TEST_P(Pkcs11EcdhNullKeyTest, UseNullPublicValue) {
+  Test(Pkcs11KeyPairGenerator(CKM_EC_KEY_PAIR_GEN, GetParam()),
+       CKM_ECDH1_DERIVE);
+}
+INSTANTIATE_TEST_SUITE_P(Pkcs11EcdhNullKeyTest, Pkcs11EcdhNullKeyTest,
+                         ::testing::Values(SEC_OID_SECG_EC_SECP256R1,
+                                           SEC_OID_SECG_EC_SECP384R1,
+                                           SEC_OID_SECG_EC_SECP521R1,
+                                           SEC_OID_CURVE25519));
+
+}  // namespace nss_test