Add a unit test to check KeygenHandler's thread-safety

We'll want some semblance of thread-safety when we make keygen asynchronous.

R=wtc,mattm
BUG=148
TEST=unit test

Review URL: http://codereview.chromium.org/2838010

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@50903 0039d316-1c4b-4281-b951-d872f2087c98

6 files changed, 169 insertions, 16 deletions

diff --git a/base/nss_util.cc b/base/nss_util.cc
index 7043b8d..091ca05 100644
--- a/base/nss_util.cc
+++ b/base/nss_util.cc
@@ -18,6 +18,11 @@
 #include "base/singleton.h"
 #include "base/string_util.h"
 
+#if defined(USE_NSS)
+#include "base/lock.h"
+#include "base/scoped_ptr.h"
+#endif // defined(USE_NSS)
+
 // On some platforms, we use NSS for SSL only -- we don't use NSS for crypto
 // or certificate verification, and we don't use the NSS certificate and key
 // databases.
@@ -163,11 +168,18 @@ class NSSInitSingleton {
     // log in.
     PK11SlotInfo* slot = PK11_GetInternalKeySlot();
     if (slot) {
+      // PK11_InitPin may write to the keyDB, but no other thread can use NSS
+      // yet, so we don't need to lock.
       if (PK11_NeedUserInit(slot))
         PK11_InitPin(slot, NULL, NULL);
       PK11_FreeSlot(slot);
     }
 
+    // TODO(davidben): When https://bugzilla.mozilla.org/show_bug.cgi?id=564011
+    // is fixed, we will no longer need the lock. We should detect this and not
+    // initialize a Lock here.
+    write_lock_.reset(new Lock());
+
     root_ = InitDefaultRootCerts();
 #endif  // defined(USE_NSS_FOR_SSL_ONLY)
   }
@@ -219,10 +231,19 @@ class NSSInitSingleton {
     return PK11_GetInternalKeySlot();
   }
 
+#if defined(USE_NSS)
+  Lock* write_lock() {
+    return write_lock_.get();
+  }
+#endif // defined(USE_NSS)
+
  private:
   PK11SlotInfo* real_db_slot_;  // Overrides internal key slot if non-NULL.
   SECMODModule *root_;
   bool chromeos_user_logged_in_;
+#if defined(USE_NSS)
+  scoped_ptr<Lock> write_lock_;
+#endif // defined(USE_NSS)
 };
 
 }  // namespace
@@ -237,6 +258,25 @@ void EnsureNSSInit() {
   Singleton<NSSInitSingleton>::get();
 }
 
+#if defined(USE_NSS)
+Lock* GetNSSWriteLock() {
+  return Singleton<NSSInitSingleton>::get()->write_lock();
+}
+
+AutoNSSWriteLock::AutoNSSWriteLock() : lock_(GetNSSWriteLock()) {
+  // May be NULL if the lock is not needed in our version of NSS.
+  if (lock_)
+    lock_->Acquire();
+}
+
+AutoNSSWriteLock::~AutoNSSWriteLock() {
+  if (lock_) {
+    lock_->AssertAcquired();
+    lock_->Release();
+  }
+}
+#endif  // defined(USE_NSS)
+
 #if defined(OS_CHROMEOS)
 void OpenPersistentNSSDB() {
   Singleton<NSSInitSingleton>::get()->OpenPersistentNSSDB();
diff --git a/base/nss_util.h b/base/nss_util.h
index 94a81cb..8de3537 100644
--- a/base/nss_util.h
+++ b/base/nss_util.h
@@ -7,6 +7,10 @@
 
 #include "base/basictypes.h"
 
+#if defined(USE_NSS)
+class Lock;
+#endif  // defined(USE_NSS)
+
 // This file specifically doesn't depend on any NSS or NSPR headers because it
 // is included by various (non-crypto) parts of chrome to call the
 // initialization functions.
@@ -33,6 +37,30 @@ void OpenPersistentNSSDB();
 // We use a int64 instead of PRTime here to avoid depending on NSPR headers.
 Time PRTimeToBaseTime(int64 prtime);
 
+#if defined(USE_NSS)
+// NSS has a bug which can cause a deadlock or stall in some cases when writing
+// to the certDB and keyDB. It also has a bug which causes concurrent key pair
+// generations to scribble over each other. To work around this, we synchronize
+// writes to the NSS databases with a global lock. The lock is hidden beneath a
+// function for easy disabling when the bug is fixed. Callers should allow for
+// it to return NULL in the future.
+//
+// See https://bugzilla.mozilla.org/show_bug.cgi?id=564011
+Lock* GetNSSWriteLock();
+
+// A helper class that acquires the NSS write Lock while the AutoNSSWriteLock
+// is in scope.
+class AutoNSSWriteLock {
+ public:
+  AutoNSSWriteLock();
+  ~AutoNSSWriteLock();
+ private:
+  Lock *lock_;
+  DISALLOW_COPY_AND_ASSIGN(AutoNSSWriteLock);
+};
+
+#endif  // defined(USE_NSS)
+
 }  // namespace base
 
 #endif  // BASE_NSS_UTIL_H_
diff --git a/net/base/cert_database_nss.cc b/net/base/cert_database_nss.cc
index 0701e6a..98930ff 100644
--- a/net/base/cert_database_nss.cc
+++ b/net/base/cert_database_nss.cc
@@ -70,9 +70,13 @@ int CertDatabase::AddUserCert(X509Certificate* cert_obj) {
   }
   nickname = username + "'s " + ca_name + " ID";
 
-  slot = PK11_ImportCertForKey(cert,
-                               const_cast<char*>(nickname.c_str()),
-                               NULL);
+  {
+    base::AutoNSSWriteLock lock;
+    slot = PK11_ImportCertForKey(cert,
+                                 const_cast<char*>(nickname.c_str()),
+                                 NULL);
+  }
+
   if (!slot) {
     LOG(ERROR) << "Couldn't import user certificate.";
     return ERR_ADD_USER_CERT_FAILED;
diff --git a/net/base/keygen_handler_unittest.cc b/net/base/keygen_handler_unittest.cc
index e6b3641..15ec0ce 100644
--- a/net/base/keygen_handler_unittest.cc
+++ b/net/base/keygen_handler_unittest.cc
@@ -4,11 +4,20 @@
 
 #include "net/base/keygen_handler.h"
 
+#include "build/build_config.h" // Needs to be imported early for USE_NSS
+
+#if defined(USE_NSS)
+#include <private/pprthred.h>  // PR_DetachThread
+#endif
+
 #include <string>
 
 #include "base/base64.h"
 #include "base/logging.h"
 #include "base/nss_util.h"
+#include "base/task.h"
+#include "base/waitable_event.h"
+#include "base/worker_pool.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 namespace net {
@@ -27,11 +36,10 @@ class KeygenHandlerTest : public ::testing::Test {
   }
 };
 
-TEST_F(KeygenHandlerTest, FLAKY_SmokeTest) {
-  KeygenHandler handler(2048, "some challenge");
-  handler.set_stores_key(false);  // Don't leave the key-pair behind
-  std::string result = handler.GenKeyAndSignChallenge();
-  LOG(INFO) << "KeygenHandler produced: " << result;
+// Assert that |result| is a valid output for KeygenHandler given challenge
+// string of |challenge|.
+void AssertValidSignedPublicKeyAndChallenge(const std::string& result,
+                                            const std::string& challenge) {
   ASSERT_GT(result.length(), 0U);
 
   // Verify it's valid base64:
@@ -64,6 +72,72 @@ TEST_F(KeygenHandlerTest, FLAKY_SmokeTest) {
   //    openssl asn1parse -inform DER
 }
 
+TEST_F(KeygenHandlerTest, FLAKY_SmokeTest) {
+  KeygenHandler handler(2048, "some challenge");
+  handler.set_stores_key(false);  // Don't leave the key-pair behind
+  std::string result = handler.GenKeyAndSignChallenge();
+  LOG(INFO) << "KeygenHandler produced: " << result;
+  AssertValidSignedPublicKeyAndChallenge(result, "some challenge");
+}
+
+class ConcurrencyTestTask : public Task {
+ public:
+  ConcurrencyTestTask(base::WaitableEvent* event,
+                      const std::string& challenge, std::string* result)
+      : event_(event),
+        challenge_(challenge),
+        result_(result) {
+  }
+
+  virtual void Run() {
+    KeygenHandler handler(2048, "some challenge");
+    handler.set_stores_key(false); // Don't leave the key-pair behind.
+    *result_ = handler.GenKeyAndSignChallenge();
+    event_->Signal();
+#if defined(USE_NSS)
+    // Detach the thread from NSPR.
+    // Calling NSS functions attaches the thread to NSPR, which stores
+    // the NSPR thread ID in thread-specific data.
+    // The threads in our thread pool terminate after we have called
+    // PR_Cleanup.  Unless we detach them from NSPR, net_unittests gets
+    // segfaults on shutdown when the threads' thread-specific data
+    // destructors run.
+    PR_DetachThread();
+#endif
+  }
+
+ private:
+  base::WaitableEvent* event_;
+  std::string challenge_;
+  std::string* result_;
+};
+
+// We asynchronously generate the keys so as not to hang up the IO thread. This
+// test tries to catch concurrency problems in the keygen implementation.
+TEST_F(KeygenHandlerTest, ConcurrencyTest) {
+  const int NUM_HANDLERS = 5;
+  base::WaitableEvent* events[NUM_HANDLERS] = { NULL };
+  std::string results[NUM_HANDLERS];
+  for (int i = 0; i < NUM_HANDLERS; i++) {
+    events[i] = new base::WaitableEvent(false, false);
+    WorkerPool::PostTask(FROM_HERE,
+                         new ConcurrencyTestTask(events[i], "some challenge",
+                                                 &results[i]),
+                         true);
+  }
+
+  for (int i = 0; i < NUM_HANDLERS; i++) {
+    // Make sure the job completed
+    bool signaled = events[i]->Wait();
+    EXPECT_TRUE(signaled);
+    delete events[i];
+    events[i] = NULL;
+
+    LOG(INFO) << "KeygenHandler " << i << " produced: " << results[i];
+    AssertValidSignedPublicKeyAndChallenge(results[i], "some challenge");
+  }
+}
+
 }  // namespace
 
 }  // namespace net
diff --git a/net/net.gyp b/net/net.gyp
index 5a25071..ec5e187 100644
--- a/net/net.gyp
+++ b/net/net.gyp
@@ -770,6 +770,7 @@
         [ 'OS == "linux" or OS == "freebsd" or OS == "openbsd"', {
             'dependencies': [
               '../build/linux/system.gyp:gtk',
+              '../build/linux/system.gyp:nss',
             ],
             'sources!': [
               'base/sdch_filter_unittest.cc',
diff --git a/net/third_party/mozilla_security_manager/nsKeygenHandler.cpp b/net/third_party/mozilla_security_manager/nsKeygenHandler.cpp
index 8be54a4..ffef66d 100644
--- a/net/third_party/mozilla_security_manager/nsKeygenHandler.cpp
+++ b/net/third_party/mozilla_security_manager/nsKeygenHandler.cpp
@@ -41,6 +41,7 @@
 #include "net/third_party/mozilla_security_manager/nsKeygenHandler.h"
 
 #include <pk11pub.h>
+#include <prerror.h>   // PR_GetError()
 #include <secmod.h>
 #include <secder.h>    // DER_Encode()
 #include <cryptohi.h>  // SEC_DerSignData()
@@ -152,13 +153,16 @@ std::string GenKeyAndSignChallenge(int key_size_in_bits,
   }
 
   LOG(INFO) << "Creating key pair...";
-  privateKey = PK11_GenerateKeyPair(slot,
-                                    keyGenMechanism,
-                                    keyGenParams,
-                                    &publicKey,
-                                    PR_TRUE,  // isPermanent?
-                                    PR_TRUE,  // isSensitive?
-                                    NULL);
+  {
+    base::AutoNSSWriteLock lock;
+    privateKey = PK11_GenerateKeyPair(slot,
+                                      keyGenMechanism,
+                                      keyGenParams,
+                                      &publicKey,
+                                      PR_TRUE,  // isPermanent?
+                                      PR_TRUE,  // isSensitive?
+                                      NULL);
+  }
   LOG(INFO) << "done.";
 
   if (!privateKey) {
@@ -227,7 +231,7 @@ std::string GenKeyAndSignChallenge(int key_size_in_bits,
 
  failure:
   if (!isSuccess) {
-    LOG(ERROR) << "SSL Keygen failed!";
+    LOG(ERROR) << "SSL Keygen failed! (NSS error code " << PR_GetError() << ")";
   } else {
     LOG(INFO) << "SSL Keygen succeeded!";
   }
@@ -237,6 +241,7 @@ std::string GenKeyAndSignChallenge(int key_size_in_bits,
     // On successful keygen we need to keep the private key, of course,
     // or we won't be able to use the client certificate.
     if (!isSuccess || !stores_key) {
+      base::AutoNSSWriteLock lock;
       PK11_DestroyTokenObject(privateKey->pkcs11Slot, privateKey->pkcs11ID);
     }
     SECKEY_DestroyPrivateKey(privateKey);
@@ -244,6 +249,7 @@ std::string GenKeyAndSignChallenge(int key_size_in_bits,
 
   if (publicKey) {
     if (!isSuccess || !stores_key) {
+      base::AutoNSSWriteLock lock;
       PK11_DestroyTokenObject(publicKey->pkcs11Slot, publicKey->pkcs11ID);
     }
     SECKEY_DestroyPublicKey(publicKey);