security: various simplifications in KeyChain

 * Remove/consolidate redundant sanity checks
 * Use verifySignature() in importSafeBag()
 * Declare a few functions static
 * Expand unit test coverage

Change-Id: Ia676653b660a6e4093b9a60b32d7b7d756c5dd2b
diff --git a/ndn-cxx/encoding/buffer-stream.hpp b/ndn-cxx/encoding/buffer-stream.hpp
index d31d6d2..6716221 100644
--- a/ndn-cxx/encoding/buffer-stream.hpp
+++ b/ndn-cxx/encoding/buffer-stream.hpp
@@ -1,6 +1,6 @@
 /* -*- Mode:C++; c-file-style:"gnu"; indent-tabs-mode:nil; -*- */
 /*
- * Copyright (c) 2013-2021 Regents of the University of California.
+ * Copyright (c) 2013-2022 Regents of the University of California.
  *
  * This file is part of ndn-cxx library (NDN C++ library with eXperimental eXtensions).
  *
@@ -72,7 +72,7 @@
 public:
   OBufferStream();
 
-  ~OBufferStream();
+  ~OBufferStream() override;
 
   /**
    * Flush written data to the stream and return shared pointer to the underlying buffer
diff --git a/ndn-cxx/security/key-chain.cpp b/ndn-cxx/security/key-chain.cpp
index 573f5c1..327bc99 100644
--- a/ndn-cxx/security/key-chain.cpp
+++ b/ndn-cxx/security/key-chain.cpp
@@ -21,10 +21,7 @@
 
 #include "ndn-cxx/security/key-chain.hpp"
 #include "ndn-cxx/security/signing-helpers.hpp"
-
-#include "ndn-cxx/encoding/buffer-stream.hpp"
-#include "ndn-cxx/util/config-file.hpp"
-#include "ndn-cxx/util/logger.hpp"
+#include "ndn-cxx/security/verification-helpers.hpp"
 
 #include "ndn-cxx/security/pib/impl/pib-memory.hpp"
 #include "ndn-cxx/security/pib/impl/pib-sqlite3.hpp"
@@ -35,13 +32,16 @@
 #include "ndn-cxx/security/tpm/impl/back-end-osx.hpp"
 #endif // NDN_CXX_HAVE_OSX_FRAMEWORKS
 
-#include "ndn-cxx/security/transform/bool-sink.hpp"
 #include "ndn-cxx/security/transform/buffer-source.hpp"
 #include "ndn-cxx/security/transform/digest-filter.hpp"
 #include "ndn-cxx/security/transform/private-key.hpp"
 #include "ndn-cxx/security/transform/public-key.hpp"
 #include "ndn-cxx/security/transform/stream-sink.hpp"
-#include "ndn-cxx/security/transform/verifier-filter.hpp"
+
+#include "ndn-cxx/encoding/buffer-stream.hpp"
+#include "ndn-cxx/util/config-file.hpp"
+#include "ndn-cxx/util/logger.hpp"
+#include "ndn-cxx/util/random.hpp"
 
 #include <boost/lexical_cast.hpp>
 #include <cstdlib>  // for std::getenv()
@@ -200,6 +200,7 @@
 Identity
 KeyChain::createIdentity(const Name& identityName, const KeyParams& params)
 {
+  NDN_LOG_DEBUG("Requesting creation of identity " << identityName);
   Identity id = m_pib->addIdentity(identityName);
 
   Key key;
@@ -214,7 +215,7 @@
     key.getDefaultCertificate();
   }
   catch (const Pib::Error&) {
-    NDN_LOG_DEBUG("No default cert for " << key.getName() << ", requesting self-signing");
+    NDN_LOG_DEBUG("No default certificate for " << key << ", requesting self-signing");
     selfSign(key);
   }
 
@@ -224,9 +225,12 @@
 void
 KeyChain::deleteIdentity(const Identity& identity)
 {
-  BOOST_ASSERT(static_cast<bool>(identity));
+  if (!identity) {
+    return;
+  }
 
   Name identityName = identity.getName();
+  NDN_LOG_DEBUG("Requesting deletion of identity " << identityName);
 
   for (const auto& key : identity.getKeys()) {
     m_tpm->deleteKey(key.getName());
@@ -238,7 +242,7 @@
 void
 KeyChain::setDefaultIdentity(const Identity& identity)
 {
-  BOOST_ASSERT(static_cast<bool>(identity));
+  BOOST_ASSERT(identity);
 
   m_pib->setDefaultIdentity(identity.getName());
 }
@@ -246,7 +250,7 @@
 Key
 KeyChain::createKey(const Identity& identity, const KeyParams& params)
 {
-  BOOST_ASSERT(static_cast<bool>(identity));
+  BOOST_ASSERT(identity);
 
   // create key in TPM
   Name keyName = m_tpm->createKey(identity.getName(), params);
@@ -254,7 +258,7 @@
   // set up key info in PIB
   Key key = identity.addKey(*m_tpm->getPublicKey(keyName), keyName);
 
-  NDN_LOG_DEBUG("Requesting self-signing for newly created key " << key.getName());
+  NDN_LOG_DEBUG("Requesting self-signing for newly created key " << key);
   selfSign(key);
 
   return key;
@@ -269,15 +273,12 @@
 void
 KeyChain::deleteKey(const Identity& identity, const Key& key)
 {
-  BOOST_ASSERT(static_cast<bool>(identity));
-  BOOST_ASSERT(static_cast<bool>(key));
-
-  Name keyName = key.getName();
-  if (identity.getName() != key.getIdentity()) {
-    NDN_THROW(std::invalid_argument("Identity `" + identity.getName().toUri() + "` "
-                                    "does not match key `" + keyName.toUri() + "`"));
+  BOOST_ASSERT(identity);
+  if (!key) {
+    return;
   }
 
+  Name keyName = key.getName();
   identity.removeKey(keyName);
   m_tpm->deleteKey(keyName);
 }
@@ -285,12 +286,8 @@
 void
 KeyChain::setDefaultKey(const Identity& identity, const Key& key)
 {
-  BOOST_ASSERT(static_cast<bool>(identity));
-  BOOST_ASSERT(static_cast<bool>(key));
-
-  if (identity.getName() != key.getIdentity())
-    NDN_THROW(std::invalid_argument("Identity `" + identity.getName().toUri() + "` "
-                                    "does not match key `" + key.getName().toUri() + "`"));
+  BOOST_ASSERT(identity);
+  BOOST_ASSERT(key);
 
   identity.setDefaultKey(key.getName());
 }
@@ -298,44 +295,30 @@
 void
 KeyChain::addCertificate(const Key& key, const Certificate& certificate)
 {
-  BOOST_ASSERT(static_cast<bool>(key));
-
-  auto pk = key.getPublicKey();
-  auto pkCert = certificate.getPublicKey();
-  if (key.getName() != certificate.getKeyName() ||
-      !std::equal(pk.begin(), pk.end(), pkCert.begin(), pkCert.end())) {
-    NDN_THROW(std::invalid_argument("Key `" + key.getName().toUri() + "` "
-                                    "does not match certificate `" + certificate.getName().toUri() + "`"));
-  }
+  BOOST_ASSERT(key);
 
   key.addCertificate(certificate);
 }
 
 void
-KeyChain::deleteCertificate(const Key& key, const Name& certificateName)
+KeyChain::deleteCertificate(const Key& key, const Name& certName)
 {
-  BOOST_ASSERT(static_cast<bool>(key));
+  BOOST_ASSERT(key);
 
-  if (!Certificate::isValidName(certificateName)) {
-    NDN_THROW(std::invalid_argument("Wrong certificate name `" + certificateName.toUri() + "`"));
-  }
-
-  key.removeCertificate(certificateName);
+  key.removeCertificate(certName);
 }
 
 void
 KeyChain::setDefaultCertificate(const Key& key, const Certificate& cert)
 {
-  BOOST_ASSERT(static_cast<bool>(key));
+  BOOST_ASSERT(key);
 
-  addCertificate(key, cert);
-  key.setDefaultCertificate(cert.getName());
+  key.setDefaultCertificate(cert);
 }
 
 shared_ptr<SafeBag>
 KeyChain::exportSafeBag(const Certificate& certificate, const char* pw, size_t pwLen)
 {
-  Name identity = certificate.getIdentity();
   Name keyName = certificate.getKeyName();
 
   ConstBufferPtr encryptedKey;
@@ -352,18 +335,18 @@
 void
 KeyChain::importSafeBag(const SafeBag& safeBag, const char* pw, size_t pwLen)
 {
-  Data certData = safeBag.getCertificate();
-  Certificate cert(std::move(certData));
+  Certificate cert(safeBag.getCertificate());
   Name identity = cert.getIdentity();
   Name keyName = cert.getKeyName();
 
+  // check if private key already exists
   if (m_tpm->hasKey(keyName)) {
     NDN_THROW(Error("Private key `" + keyName.toUri() + "` already exists"));
   }
 
+  // check if public key already exists
   try {
-    Identity existingId = m_pib->getIdentity(identity);
-    existingId.getKey(keyName);
+    m_pib->getIdentity(identity).getKey(keyName);
     NDN_THROW(Error("Public key `" + keyName.toUri() + "` already exists"));
   }
   catch (const Pib::Error&) {
@@ -377,25 +360,18 @@
     NDN_THROW_NESTED(Error("Failed to import private key `" + keyName.toUri() + "`"));
   }
 
-  // check the consistency of private key and certificate
-  const uint8_t content[] = {0x01, 0x02, 0x03, 0x04};
+  // check the consistency of private key and certificate (sign/verify a random message)
+  const auto r = random::generateWord64();
+  const auto msg = make_span(reinterpret_cast<const uint8_t*>(&r), sizeof(r));
   ConstBufferPtr sigBits;
   try {
-    sigBits = m_tpm->sign({content}, keyName, DigestAlgorithm::SHA256);
+    sigBits = m_tpm->sign({msg}, keyName, DigestAlgorithm::SHA256);
   }
   catch (const std::runtime_error&) {
     m_tpm->deleteKey(keyName);
     NDN_THROW(Error("Invalid private key `" + keyName.toUri() + "`"));
   }
-  bool isVerified = false;
-  {
-    using namespace transform;
-    PublicKey publicKey;
-    publicKey.loadPkcs8(cert.getPublicKey());
-    bufferSource(content) >> verifierFilter(DigestAlgorithm::SHA256, publicKey, *sigBits)
-                          >> boolSink(isVerified);
-  }
-  if (!isVerified) {
+  if (!verifySignature({msg}, *sigBits, cert.getPublicKey())) {
     m_tpm->deleteKey(keyName);
     NDN_THROW(Error("Certificate `" + cert.getName().toUri() + "` "
                     "and private key `" + keyName.toUri() + "` do not match"));
@@ -708,7 +684,7 @@
       return prepareSignatureInfoSha256(params);
     }
     case SigningInfo::SIGNER_TYPE_HMAC: {
-      return prepareSignatureInfoHmac(params);
+      return prepareSignatureInfoHmac(params, *m_tpm);
     }
   }
   NDN_THROW(InvalidSigningInfoError("Unrecognized signer type " +
@@ -722,15 +698,15 @@
   sigInfo.setSignatureType(tlv::DigestSha256);
 
   NDN_LOG_TRACE("Prepared signature info: " << sigInfo);
-  return std::make_tuple(SigningInfo::getDigestSha256Identity(), sigInfo);
+  return {SigningInfo::getDigestSha256Identity(), sigInfo};
 }
 
 std::tuple<Name, SignatureInfo>
-KeyChain::prepareSignatureInfoHmac(const SigningInfo& params)
+KeyChain::prepareSignatureInfoHmac(const SigningInfo& params, Tpm& tpm)
 {
   const Name& keyName = params.getSignerName();
-  if (!m_tpm->hasKey(keyName)) {
-    m_tpm->importPrivateKey(keyName, params.getHmacKey());
+  if (!tpm.hasKey(keyName)) {
+    tpm.importPrivateKey(keyName, params.getHmacKey());
   }
 
   auto sigInfo = params.getSignatureInfo();
@@ -738,7 +714,7 @@
   sigInfo.setKeyLocator(keyName);
 
   NDN_LOG_TRACE("Prepared signature info: " << sigInfo);
-  return std::make_tuple(keyName, sigInfo);
+  return {keyName, sigInfo};
 }
 
 std::tuple<Name, SignatureInfo>
@@ -756,7 +732,8 @@
 }
 
 std::tuple<Name, SignatureInfo>
-KeyChain::prepareSignatureInfoWithKey(const SigningInfo& params, const pib::Key& key, optional<Name> certName)
+KeyChain::prepareSignatureInfoWithKey(const SigningInfo& params, const pib::Key& key,
+                                      const optional<Name>& certName)
 {
   auto sigInfo = params.getSignatureInfo();
   sigInfo.setSignatureType(getSignatureType(key.getKeyType(), params.getDigestAlgorithm()));
@@ -776,7 +753,7 @@
   }
 
   NDN_LOG_TRACE("Prepared signature info: " << sigInfo);
-  return std::make_tuple(key.getName(), sigInfo);
+  return {key.getName(), sigInfo};
 }
 
 ConstBufferPtr
diff --git a/ndn-cxx/security/key-chain.hpp b/ndn-cxx/security/key-chain.hpp
index 12123bd..a91e498 100644
--- a/ndn-cxx/security/key-chain.hpp
+++ b/ndn-cxx/security/key-chain.hpp
@@ -76,12 +76,12 @@
 inline namespace v2 {
 
 /**
- * @brief The interface of signing key management.
+ * @brief The main interface for signing key management.
  *
  * The KeyChain class provides an interface to manage entities related to packet signing,
  * such as Identity, Key, and Certificates.  It consists of two parts: a private key module
- * (TPM) and a public key information base (PIB).  Managing signing key and its related
- * entities through KeyChain interface guarantees the consistency between TPM and PIB.
+ * (TPM) and a public key information base (PIB).  Managing signing keys and their related
+ * entities through the KeyChain interface guarantees the consistency between TPM and PIB.
  */
 class KeyChain : noncopyable
 {
@@ -153,34 +153,34 @@
   /**
    * @brief Create an identity @p identityName.
    *
-   * This method will check if the identity exists in PIB and whether the identity has a
-   * default key and default certificate.  If the identity does not exist, this method will
-   * create the identity in PIB.  If the identity's default key does not exist, this method
-   * will create a key pair and set it as the identity's default key.  If the key's default
-   * certificate is missing, this method will create a self-signed certificate for the key.
+   * This method will check if the identity exists in the PIB and whether the identity has
+   * a default key and default certificate.  If the identity does not exist, this method
+   * will create it.  If the identity's default key does not exist, this method will create
+   * a key pair and set it as the identity's default key.  If the key's default certificate
+   * is missing, this method will create a self-signed certificate for the key.
    *
    * If @p identityName did not exist and no default identity was selected before, the created
-   * identity will be set as the default identity
+   * identity will be set as the default identity.
    *
    * @param identityName The name of the identity.
    * @param params The key parameters if a key needs to be created for the identity (default:
-   *               EC key with random key id)
+   *               EC key with random key id).
    * @return The created Identity instance.
    */
   Identity
   createIdentity(const Name& identityName, const KeyParams& params = getDefaultKeyParams());
 
   /**
-   * @brief delete @p identity.
+   * @brief Delete @p identity from this KeyChain.
    *
-   * @pre @p identity must be valid.
-   * @post @p identity becomes invalid.
+   * Attempting to delete an invalid identity has no effect.
    */
   void
   deleteIdentity(const Identity& identity);
 
   /**
    * @brief Set @p identity as the default identity.
+   *
    * @pre @p identity must be valid.
    */
   void
@@ -190,14 +190,14 @@
   /**
    * @brief Create a new key for @p identity.
    *
-   * @param identity Reference to a valid Identity object
-   * @param params Key creation parameters (default: EC key with random key id)
-   * @pre @p identity must be valid.
-   *
    * If @p identity had no default key selected, the created key will be set as the default for
    * this identity.
    *
    * This method will also create a self-signed certificate for the created key.
+   *
+   * @param identity The identity that will own the created key; must be valid.
+   * @param params Key creation parameters (default: EC key with random key id).
+   * @return The created Key instance.
    */
   Key
   createKey(const Identity& identity, const KeyParams& params = getDefaultKeyParams());
@@ -205,24 +205,24 @@
    /**
    * @brief Create a new HMAC key.
    *
-   * @param prefix Prefix used to construct the key name (default: `/localhost/identity/hmac`);
-   *               the full key name will include additional components according to @p params
-   * @param params Key creation parameters
-   * @return A name that can be subsequently used to reference the created key.
-   *
    * The newly created key will be inserted in the TPM. HMAC keys don't have any PIB entries.
+   *
+   * @param prefix Prefix used to construct the key name (default: `/localhost/identity/hmac`);
+   *               the full key name will include additional components according to @p params.
+   * @param params Key creation parameters.
+   * @return A name that can be subsequently used to reference the created key.
    */
   Name
   createHmacKey(const Name& prefix = SigningInfo::getHmacIdentity(),
                 const HmacKeyParams& params = HmacKeyParams());
 
   /**
-   * @brief Delete a key @p key of @p identity.
+   * @brief Delete @p key from @p identity.
+   *
+   * Attempting to delete an invalid key has no effect.
    *
    * @pre @p identity must be valid.
-   * @pre @p key must be valid.
-   * @post @p key becomes invalid.
-   * @throw std::invalid_argument @p key does not belong to @p identity
+   * @throw std::invalid_argument @p key does not belong to @p identity.
    */
   void
   deleteKey(const Identity& identity, const Key& key);
@@ -232,49 +232,49 @@
    *
    * @pre @p identity must be valid.
    * @pre @p key must be valid.
-   * @throw std::invalid_argument @p key does not belong to @p identity
+   * @throw std::invalid_argument @p key does not belong to @p identity.
    */
   void
   setDefaultKey(const Identity& identity, const Key& key);
 
 public: // Certificate management
   /**
-   * @brief Add a certificate @p certificate for @p key
+   * @brief Add a certificate @p cert for @p key.
    *
    * If @p key had no default certificate selected, the added certificate will be set as the
    * default certificate for this key.
    *
-   * @note This method overwrites certificate with the same name, without considering the
-   *       implicit digest.
+   * This method will overwrite a certificate with the same name (without considering the
+   * implicit digest).
    *
    * @pre @p key must be valid.
-   * @throw std::invalid_argument @p key does not match @p certificate
+   * @throw std::invalid_argument The certificate does not match the key.
    */
   void
-  addCertificate(const Key& key, const Certificate& certificate);
+  addCertificate(const Key& key, const Certificate& cert);
 
   /**
-   * @brief delete a certificate with name @p certificateName of @p key.
+   * @brief Delete a certificate with name @p certName from @p key.
    *
-   * If the certificate @p certificateName does not exist, this method has no effect.
+   * If the certificate does not exist, this method has no effect.
    *
    * @pre @p key must be valid.
-   * @throw std::invalid_argument @p certificateName does not follow certificate naming convention.
+   * @throw std::invalid_argument The certificate name is invalid or does not match the key name.
    */
   void
-  deleteCertificate(const Key& key, const Name& certificateName);
+  deleteCertificate(const Key& key, const Name& certName);
 
   /**
    * @brief Set @p cert as the default certificate of @p key.
    *
-   * The certificate @p cert will be added to the @p key, potentially overriding existing
-   * certificate if it has the same name (without considering implicit digest).
+   * The certificate @p cert will be added to @p key, potentially overwriting an existing
+   * certificate with the same name (without considering the implicit digest).
    *
    * @pre @p key must be valid.
-   * @throw std::invalid_argument @p key does not match @p certificate
+   * @throw std::invalid_argument The certificate does not match the key.
    */
   void
-  setDefaultCertificate(const Key& key, const Certificate& certificate);
+  setDefaultCertificate(const Key& key, const Certificate& cert);
 
 public: // signing
   /**
@@ -487,18 +487,18 @@
   std::tuple<Name, SignatureInfo>
   prepareSignatureInfo(const SigningInfo& params);
 
-  std::tuple<Name, SignatureInfo>
+  static std::tuple<Name, SignatureInfo>
   prepareSignatureInfoSha256(const SigningInfo& params);
 
-  std::tuple<Name, SignatureInfo>
-  prepareSignatureInfoHmac(const SigningInfo& params);
+  static std::tuple<Name, SignatureInfo>
+  prepareSignatureInfoHmac(const SigningInfo& params, Tpm& tpm);
 
-  std::tuple<Name, SignatureInfo>
+  static std::tuple<Name, SignatureInfo>
   prepareSignatureInfoWithIdentity(const SigningInfo& params, const pib::Identity& identity);
 
-  std::tuple<Name, SignatureInfo>
+  static std::tuple<Name, SignatureInfo>
   prepareSignatureInfoWithKey(const SigningInfo& params, const pib::Key& key,
-                              optional<Name> certName = nullopt);
+                              const optional<Name>& certName = nullopt);
 
   /**
    * @brief Generate and return a raw signature for the byte ranges in @p bufs using
diff --git a/ndn-cxx/security/pib/impl/key-impl.cpp b/ndn-cxx/security/pib/impl/key-impl.cpp
index 70b0036..f308099 100644
--- a/ndn-cxx/security/pib/impl/key-impl.cpp
+++ b/ndn-cxx/security/pib/impl/key-impl.cpp
@@ -55,6 +55,12 @@
 KeyImpl::addCertificate(const Certificate& cert)
 {
   BOOST_ASSERT(m_certificates.isConsistent());
+
+  auto certKey = cert.getPublicKey();
+  if (!std::equal(m_key.begin(), m_key.end(), certKey.begin(), certKey.end())) {
+    NDN_THROW(std::invalid_argument("Public key of certificate `" + cert.getName().toUri() + "` "
+                                    "does not match key `" + m_keyName.toUri() + "`"));
+  }
   m_certificates.add(cert);
 }
 
diff --git a/tests/unit/security/key-chain.t.cpp b/tests/unit/security/key-chain.t.cpp
index f01aab9..ddee6db 100644
--- a/tests/unit/security/key-chain.t.cpp
+++ b/tests/unit/security/key-chain.t.cpp
@@ -215,57 +215,59 @@
 
 BOOST_FIXTURE_TEST_CASE(Management, KeyChainFixture)
 {
-  Name identityName("/test/id");
-  Name identity2Name("/test/id2");
-
   BOOST_CHECK_EQUAL(m_keyChain.getPib().getIdentities().size(), 0);
-  BOOST_REQUIRE_THROW(m_keyChain.getPib().getDefaultIdentity(), Pib::Error);
+  BOOST_CHECK_THROW(m_keyChain.getPib().getDefaultIdentity(), Pib::Error);
 
-  // Create identity
-  Identity id = m_keyChain.createIdentity(identityName);
-  BOOST_CHECK(id);
-  BOOST_CHECK(m_keyChain.getPib().getIdentities().find(identityName) != m_keyChain.getPib().getIdentities().end());
+  // Create an identity
+  Name idName("/test/id");
+  Identity id = m_keyChain.createIdentity(idName);
+  BOOST_REQUIRE(id);
+  BOOST_CHECK_EQUAL(m_keyChain.getPib().getIdentities().size(), 1);
+  BOOST_CHECK_EQUAL(m_keyChain.getPib().getIdentity(idName), id);
+
   // The first added identity becomes the default identity
-  BOOST_CHECK_NO_THROW(m_keyChain.getPib().getDefaultIdentity());
+  BOOST_CHECK_EQUAL(m_keyChain.getPib().getDefaultIdentity(), id);
   // The default key of the added identity must exist
   Key key = id.getDefaultKey();
+  BOOST_REQUIRE(key);
+  BOOST_CHECK_EQUAL(id.getKeys().size(), 1);
   // The default certificate of the default key must exist
   BOOST_CHECK_NO_THROW(key.getDefaultCertificate());
+  BOOST_CHECK_EQUAL(key.getCertificates().size(), 1);
 
   // Delete key
   Name key1Name = key.getName();
-  BOOST_CHECK_NO_THROW(id.getKey(key1Name));
-  BOOST_CHECK_EQUAL(id.getKeys().size(), 1);
   m_keyChain.deleteKey(id, key);
   // The key instance should not be valid any more
   BOOST_CHECK(!key);
+  BOOST_CHECK_THROW(id.getDefaultKey(), Pib::Error);
   BOOST_CHECK_THROW(id.getKey(key1Name), Pib::Error);
   BOOST_CHECK_EQUAL(id.getKeys().size(), 0);
 
   // Create another key
   m_keyChain.createKey(id);
-  // The added key becomes the default key.
+  // The added key becomes the default key
   Key key2 = id.getDefaultKey();
   BOOST_REQUIRE(key2);
   BOOST_CHECK_NE(key2.getName(), key1Name);
   BOOST_CHECK_EQUAL(id.getKeys().size(), 1);
   BOOST_CHECK_NO_THROW(key2.getDefaultCertificate());
+  BOOST_CHECK_EQUAL(key2.getCertificates().size(), 1);
 
-  // Create the third key
+  // Create a third key
   Key key3 = m_keyChain.createKey(id);
   BOOST_CHECK_NE(key3.getName(), key2.getName());
-  // The added key will not be the default key, because the default key already exists
-  BOOST_CHECK_EQUAL(id.getDefaultKey().getName(), key2.getName());
   BOOST_CHECK_EQUAL(id.getKeys().size(), 2);
-  BOOST_CHECK_NO_THROW(key3.getDefaultCertificate());
+  // The added key will not be the default key, because the default was already set
+  BOOST_CHECK_EQUAL(id.getDefaultKey().getName(), key2.getName());
 
   // Delete cert
-  BOOST_CHECK_EQUAL(key3.getCertificates().size(), 1);
+  BOOST_REQUIRE_EQUAL(key3.getCertificates().size(), 1);
   Certificate key3Cert1 = *key3.getCertificates().begin();
   Name key3CertName = key3Cert1.getName();
   m_keyChain.deleteCertificate(key3, key3CertName);
   BOOST_CHECK_EQUAL(key3.getCertificates().size(), 0);
-  BOOST_REQUIRE_THROW(key3.getDefaultCertificate(), Pib::Error);
+  BOOST_CHECK_THROW(key3.getDefaultCertificate(), Pib::Error);
 
   // Add cert
   m_keyChain.addCertificate(key3, key3Cert1);
@@ -273,41 +275,58 @@
   BOOST_CHECK_NO_THROW(key3.getDefaultCertificate());
   m_keyChain.addCertificate(key3, key3Cert1); // overwriting the cert should work
   BOOST_CHECK_EQUAL(key3.getCertificates().size(), 1);
+
   // Add another cert
   Certificate key3Cert2 = key3Cert1;
   Name key3Cert2Name = key3.getName();
-  key3Cert2Name.append("Self");
-  key3Cert2Name.appendVersion();
+  key3Cert2Name.append("Self").appendVersion();
   key3Cert2.setName(key3Cert2Name);
   m_keyChain.addCertificate(key3, key3Cert2);
   BOOST_CHECK_EQUAL(key3.getCertificates().size(), 2);
+
   // Add empty cert
   Certificate key3Cert3 = key3Cert1;
   key3Cert3.unsetContent();
   BOOST_CHECK_THROW(m_keyChain.addCertificate(key3, key3Cert3), std::invalid_argument);
 
-  // Default certificate setting
-  BOOST_CHECK_EQUAL(key3.getDefaultCertificate().getName(), key3CertName);
-  m_keyChain.setDefaultCertificate(key3, key3Cert2);
-  BOOST_CHECK_EQUAL(key3.getDefaultCertificate().getName(), key3Cert2Name);
+  // Create another identity
+  Name id2Name("/test/id2");
+  Identity id2 = m_keyChain.createIdentity(id2Name);
+  BOOST_REQUIRE(id2);
+  BOOST_CHECK_EQUAL(m_keyChain.getPib().getIdentities().size(), 2);
+  BOOST_CHECK_EQUAL(m_keyChain.getPib().getIdentity(id2Name), id2);
+
+  // Default identity setting
+  BOOST_CHECK_EQUAL(m_keyChain.getPib().getDefaultIdentity().getName(), id.getName());
+  m_keyChain.setDefaultIdentity(id2);
+  BOOST_CHECK_EQUAL(m_keyChain.getPib().getDefaultIdentity().getName(), id2.getName());
 
   // Default key setting
   BOOST_CHECK_EQUAL(id.getDefaultKey().getName(), key2.getName());
   m_keyChain.setDefaultKey(id, key3);
   BOOST_CHECK_EQUAL(id.getDefaultKey().getName(), key3.getName());
+  // Set key of a different identity as default
+  BOOST_CHECK_THROW(m_keyChain.setDefaultKey(id, id2.getDefaultKey()), std::invalid_argument);
 
-  // Default identity setting
-  Identity id2 = m_keyChain.createIdentity(identity2Name);
-  BOOST_CHECK_EQUAL(m_keyChain.getPib().getDefaultIdentity().getName(), id.getName());
-  m_keyChain.setDefaultIdentity(id2);
-  BOOST_CHECK_EQUAL(m_keyChain.getPib().getDefaultIdentity().getName(), id2.getName());
+  // Default certificate setting
+  BOOST_CHECK_EQUAL(key3.getDefaultCertificate().getName(), key3CertName);
+  m_keyChain.setDefaultCertificate(key3, key3Cert2);
+  BOOST_CHECK_EQUAL(key3.getDefaultCertificate().getName(), key3Cert2Name);
+  // Set certificate of a different key as default
+  BOOST_CHECK_THROW(m_keyChain.setDefaultCertificate(key3, key2.getDefaultCertificate()),
+                    std::invalid_argument);
+
+  // Delete certificate name mismatch
+  BOOST_CHECK_THROW(m_keyChain.deleteCertificate(key2, key3CertName), std::invalid_argument);
+  // Delete key name mismatch
+  BOOST_CHECK_THROW(m_keyChain.deleteKey(id, id2.getDefaultKey()), std::invalid_argument);
 
   // Delete identity
   m_keyChain.deleteIdentity(id);
   // The identity instance should not be valid any more
   BOOST_CHECK(!id);
-  BOOST_CHECK_THROW(m_keyChain.getPib().getIdentity(identityName), Pib::Error);
-  BOOST_CHECK(m_keyChain.getPib().getIdentities().find(identityName) == m_keyChain.getPib().getIdentities().end());
+  BOOST_CHECK_THROW(m_keyChain.getPib().getIdentity(idName), Pib::Error);
+  BOOST_CHECK_EQUAL(m_keyChain.getPib().getIdentities().size(), 1);
 }
 
 struct DataPkt
@@ -718,18 +737,21 @@
   Certificate cert = id.getDefaultKey().getDefaultCertificate();
 
   shared_ptr<SafeBag> exported = m_keyChain.exportSafeBag(cert, "1234", 4);
-  Block block = exported->wireEncode();
+  const auto& block = exported->wireEncode();
 
   m_keyChain.deleteIdentity(id);
-  BOOST_CHECK_THROW(m_keyChain.exportSafeBag(cert, "1234", 4), KeyChain::Error);
+  BOOST_CHECK_EXCEPTION(m_keyChain.exportSafeBag(cert, "1234", 4), KeyChain::Error, [] (const auto& e) {
+    return std::string(e.what()).find("Failed to export private key") == 0;
+  });
 
   BOOST_CHECK_EQUAL(m_keyChain.getTpm().hasKey(cert.getKeyName()), false);
   BOOST_CHECK_EQUAL(m_keyChain.getPib().getIdentities().size(), 0);
 
-  SafeBag imported;
-  imported.wireDecode(block);
+  SafeBag imported(block);
   m_keyChain.importSafeBag(imported, "1234", 4);
-  BOOST_CHECK_THROW(m_keyChain.importSafeBag(imported, "1234", 4), KeyChain::Error);
+  BOOST_CHECK_EXCEPTION(m_keyChain.importSafeBag(imported, "1234", 4), KeyChain::Error, [] (const auto& e) {
+    return std::string(e.what()).find("already exists") != std::string::npos;
+  });
 
   BOOST_CHECK_EQUAL(m_keyChain.getTpm().hasKey(cert.getKeyName()), true);
   BOOST_CHECK_EQUAL(m_keyChain.getPib().getIdentities().size(), 1);
@@ -737,11 +759,31 @@
   BOOST_CHECK_EQUAL(newId.getKeys().size(), 1);
   Key newKey = newId.getKey(cert.getKeyName());
   BOOST_CHECK_EQUAL(newKey.getCertificates().size(), 1);
-  BOOST_CHECK_NO_THROW(newKey.getCertificate(cert.getName()));
+  Certificate newCert = newKey.getCertificate(cert.getName());
+  BOOST_CHECK_EQUAL(newCert, cert);
+}
 
-  m_keyChain.deleteIdentity(newId);
-  BOOST_CHECK_EQUAL(m_keyChain.getPib().getIdentities().size(), 0);
-  BOOST_CHECK_EQUAL(m_keyChain.getTpm().hasKey(cert.getKeyName()), false);
+BOOST_FIXTURE_TEST_CASE(ImportInvalid, KeyChainFixture)
+{
+  Identity id = m_keyChain.createIdentity("/TestKeyChain/ExportIdentity");
+  Certificate cert = id.getDefaultKey().getDefaultCertificate();
+
+  auto exported = m_keyChain.exportSafeBag(cert, "1234", 4);
+  m_keyChain.deleteIdentity(id);
+
+  Identity id2 = m_keyChain.createIdentity("/TestKeyChain/AnotherIdentity");
+  Certificate cert2 = id2.getDefaultKey().getDefaultCertificate();
+  m_keyChain.deleteIdentity(id2);
+
+  SafeBag mismatch(cert2, exported->getEncryptedKey());
+  BOOST_CHECK_EXCEPTION(m_keyChain.importSafeBag(mismatch, "1234", 4), KeyChain::Error, [] (const auto& e) {
+    return std::string(e.what()).find("do not match") != std::string::npos;
+  });
+
+  SafeBag invalidPriv(cert2, {0xCA, 0xFE});
+  BOOST_CHECK_EXCEPTION(m_keyChain.importSafeBag(invalidPriv, "1234", 4), KeyChain::Error, [] (const auto& e) {
+    return std::string(e.what()).find("Failed to import private key") == 0;
+  });
 }
 
 BOOST_FIXTURE_TEST_CASE(SelfSignedCertValidity, KeyChainFixture)