Implement AES-CTR for NSS.

crypto/encryptor_nss.cc

 // Copyright (c) 2011 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "crypto/encryptor.h"
 
 #include <cryptohi.h>
 #include <vector>
 
 #include "base/logging.h"
 #include "crypto/nss_util.h"
 #include "crypto/symmetric_key.h"
 
 namespace crypto {
 
+namespace {
+
+inline CK_MECHANISM_TYPE GetMechanism(Encryptor::Mode mode) {
+  switch (mode) {
+    case Encryptor::CBC:
+      return CKM_AES_CBC_PAD;
+    case Encryptor::CTR:
+      // AES-CTR encryption uses ECB encryptor as a building block since
+      // NSS doesn't support CTR encryption mode.
+      return CKM_AES_ECB;
+    default:
+      NOTREACHED() << "Unsupported mode of operation";
+      break;
+  }
+  return 0;

[wtc, 2011/06/24 18:06:06]

Let's use
  return static_cast<CK_MECHANISM_TYPE>(-1);

0 is actually a valid mechanism (CKM_RSA_PKCS_KEY_PAIR_GEN).

[Chris Palmer, 2011/10/03 19:50:09]

+1

+}
+
+}  // namespace
+
 Encryptor::Encryptor()
     : key_(NULL),
       mode_(CBC) {
   EnsureNSSInit();
 }
 
 Encryptor::~Encryptor() {
 }
 
 bool Encryptor::Init(SymmetricKey* key, Mode mode, const std::string& iv) {
   DCHECK(key);
-  DCHECK_EQ(CBC, mode);
+  DCHECK(CBC == mode || CTR == mode) << "Unsupported mode of operation";
 
   key_ = key;
   mode_ = mode;
 
-  if (iv.size() != AES_BLOCK_SIZE)
+  if (mode == CBC && iv.size() != AES_BLOCK_SIZE)
     return false;
 
-  slot_.reset(PK11_GetBestSlot(CKM_AES_CBC_PAD, NULL));
+  slot_.reset(PK11_GetBestSlot(GetMechanism(mode), NULL));
   if (!slot_.get())
     return false;
 
-  SECItem iv_item;
-  iv_item.type = siBuffer;
-  iv_item.data = reinterpret_cast<unsigned char*>(
-      const_cast<char *>(iv.data()));
-  iv_item.len = iv.size();
+  switch (mode) {
+    case CBC:
+      SECItem iv_item;
+      iv_item.type = siBuffer;
+      iv_item.data = reinterpret_cast<unsigned char*>(
+          const_cast<char *>(iv.data()));
+      iv_item.len = iv.size();
 
-  param_.reset(PK11_ParamFromIV(CKM_AES_CBC_PAD, &iv_item));
+      param_.reset(PK11_ParamFromIV(GetMechanism(mode), &iv_item));
+      break;
+    case CTR:
+      param_.reset(PK11_ParamFromIV(GetMechanism(mode), NULL));
+      break;
+  }
+
   if (!param_.get())
     return false;
-
   return true;
 }
 
 bool Encryptor::Encrypt(const std::string& plaintext, std::string* ciphertext) {
-  ScopedPK11Context context(PK11_CreateContextBySymKey(CKM_AES_CBC_PAD,
+  ScopedPK11Context context(PK11_CreateContextBySymKey(GetMechanism(mode_),
                                                        CKA_ENCRYPT,
                                                        key_->key(),
                                                        param_.get()));
   if (!context.get())
     return false;
 
-  size_t ciphertext_len = plaintext.size() + AES_BLOCK_SIZE;
-  std::vector<unsigned char> buffer(ciphertext_len);
-
-  int op_len;
-  SECStatus rv = PK11_CipherOp(context.get(),
-                               &buffer[0],
-                               &op_len,
-                               ciphertext_len,
-                               reinterpret_cast<unsigned char*>(
-                                   const_cast<char*>(plaintext.data())),
-                               plaintext.size());
-  if (SECSuccess != rv)
-    return false;
-
-  unsigned int digest_len;
-  rv = PK11_DigestFinal(context.get(),
-                        &buffer[op_len],
-                        &digest_len,
-                        ciphertext_len - op_len);
-  if (SECSuccess != rv)
-    return false;
-
-  ciphertext->assign(reinterpret_cast<char *>(&buffer[0]),
-                     op_len + digest_len);
-  return true;
+  if (mode_ == CTR)
+    return CryptCTR(context.get(), plaintext, ciphertext);
+  else
+    return Crypt(context.get(), plaintext, ciphertext);
 }
 
 bool Encryptor::Decrypt(const std::string& ciphertext, std::string* plaintext) {
   if (ciphertext.empty())
     return false;
 
-  ScopedPK11Context context(PK11_CreateContextBySymKey(CKM_AES_CBC_PAD,
-                                                       CKA_DECRYPT,
-                                                       key_->key(),
-                                                       param_.get()));
+  ScopedPK11Context context(PK11_CreateContextBySymKey(
+      GetMechanism(mode_), (mode_ == CTR ? CKA_ENCRYPT : CKA_DECRYPT),

[wtc, 2011/06/24 18:06:06]

Are you sure we have to use CKA_ENCRYPT for CTR mode?

[Alpha Left Google, 2011/06/24 18:52:27]

Yup, decrypt does't work since we need to encrypt the counter to decrypt.

+      key_->key(), param_.get()));
   if (!context.get())
     return false;
 
-  size_t plaintext_len = ciphertext.size();
-  std::vector<unsigned char> buffer(plaintext_len);
+  if (mode_ == CTR)
+    return CryptCTR(context.get(), ciphertext, plaintext);
+  else
+    return Crypt(context.get(), ciphertext, plaintext);
+}
+
+bool Encryptor::Crypt(PK11Context* context, const std::string& input,
+                      std::string* output) {
+  size_t output_len = input.size() + AES_BLOCK_SIZE;
+  CHECK(output_len > input.size()) << "Output size overflow";
+
+  output->resize(output_len);
+  uint8* output_data =
+      reinterpret_cast<uint8*>(const_cast<char*>(output->data()));
+
+  int input_len = input.size();
+  uint8* input_data =
+      reinterpret_cast<uint8*>(const_cast<char*>(input.data()));
 
   int op_len;
-  SECStatus rv = PK11_CipherOp(context.get(),
-                               &buffer[0],
+  SECStatus rv = PK11_CipherOp(context,
+                               output_data,
                                &op_len,
-                               plaintext_len,
-                               reinterpret_cast<unsigned char*>(
-                                   const_cast<char*>(ciphertext.data())),
-                               ciphertext.size());
+                               output_len,
+                               input_data,
+                               input_len);
+
   if (SECSuccess != rv)
     return false;
 
   unsigned int digest_len;
-  rv = PK11_DigestFinal(context.get(),
-                        &buffer[op_len],
+  rv = PK11_DigestFinal(context,
+                        output_data + op_len,
                         &digest_len,
-                        plaintext_len - op_len);
+                        output_len - op_len);
   if (SECSuccess != rv)
     return false;
 
-  plaintext->assign(reinterpret_cast<char *>(&buffer[0]),
-                    op_len + digest_len);
+  output->resize(op_len + digest_len);
   return true;
 }
 
+bool Encryptor::CryptCTR(PK11Context* context, const std::string& input,
+                         std::string* output) {
+  if (!counter_.get()) {
+    LOG(ERROR) << "Counter value not set in CTR mode.";
+    return false;
+  }
+
+  size_t output_len = input.size() + AES_BLOCK_SIZE;

[wtc, 2011/06/24 18:06:06]

You can use the same formula as in GenerateCounterMask to
estimate the output_len:

  size_t output_len = ((input.size() + AES_BLOCK_SIZE - 1) / AES_BLOCK_SIZE) *
AES_BLOCK_SIZE;

[Alpha Left Google, 2011/06/24 18:52:27]

Done.

[Chris Palmer, 2011/10/03 19:50:09]

No, CTR mode gives you a stream cipher. Unlike CBC mode, the ciphertext size
need not be a multiple of the block size. You could pad it, of course, as with
CBC mode, if you wanted to "blind" or "hide" the size of the plaintext somewhat.
But it's ok and normal for the ciphertext to be the same length as the
plaintext.

+  CHECK(output_len > input.size()) << "Output size overflow";
+  output->resize(output_len);
+  uint8* output_data =
+      reinterpret_cast<uint8*>(const_cast<char*>(output->data()));
+
+  size_t mask_len;
+  GenerateCounterMask(input.size(), output_data, &mask_len);
+  CHECK(mask_len <= output_len);
+
+  int op_len;
+  SECStatus rv = PK11_CipherOp(context,
+                               output_data,
+                               &op_len,
+                               output_len,
+                               output_data,
+                               mask_len);
+  if (SECSuccess != rv)
+    return false;
+  CHECK(op_len == static_cast<int>(mask_len));
+
+  unsigned int digest_len;
+  rv = PK11_DigestFinal(context,
+                        NULL,
+                        &digest_len,
+                        0);
+  if (SECSuccess != rv)
+    return false;
+  CHECK(!digest_len);
+
+  // Use |output_data| to mask |input|.
+  MaskMessage(
+      reinterpret_cast<uint8*>(const_cast<char*>(input.data())),
+      input.length(), output_data, output_data);
+  output->resize(input.length());
+  return true;
+}
+
 }  // namespace crypto