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.

[wtc, 2011/06/21 00:33:32]

Please explain this is necessary because NSS doesn't support
CKM_AES_CTR yet.

[Alpha Left Google, 2011/06/22 21:25:17]

Done.

+      return CKM_AES_ECB;
+    default:
+      NOTREACHED() << "Unsupported mode of operation";
+      break;
+  }
+  return 0;
+}
+
+}  // 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();
+  if (mode == 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));
+  } else {
+    param_.reset(PK11_ParamFromIV(GetMechanism(mode), NULL));
+  }

[wtc, 2011/06/21 00:33:32]

I would probably use a "switch" statement here so that we
can add new modes of operation more easily in the future.

[Alpha Left Google, 2011/06/22 21:25:17]

Done.

+
   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;
+  return mode_ == CTR ? CryptCTR(context.get(), plaintext, ciphertext) :
+      Crypt(context.get(), plaintext, ciphertext);

[wtc, 2011/06/21 00:33:32]

Nit: please use an "if" or "switch" statement here and in
the Decrypt method on lines 96-97 below.

[Alpha Left Google, 2011/06/22 21:25:17]

Done.

 }
 
 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),
+      key_->key(), param_.get()));
   if (!context.get())
     return false;
 
-  size_t plaintext_len = ciphertext.size();
-  std::vector<unsigned char> buffer(plaintext_len);
+  return mode_ == CTR ? CryptCTR(context.get(), ciphertext, plaintext) :
+      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;

[Ryan Sleevi, 2011/06/20 23:47:23]

nit: Is it necessary to pad |output_len| ? The underlying cipher is ECB, and the
mask code is smart enough to round up, so I'm not sure why.

Was this supposed to pad output_len to a multiple of AES_BLOCK_SIZE for
PK11_CipherOp?

[Alpha Left Google, 2011/06/20 23:59:44]

This is just for convenience, it's easier to check make sure this is correct if
we spend a few more bytes.

+  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 input_len = input.size();
+  uint8* input_data =
+      reinterpret_cast<uint8*>(const_cast<char*>(input.data()));
+
+  scoped_array<uint8> ctr_mask;
+  GenerateCounterMask(input_len, &ctr_mask, &input_len);

[wtc, 2011/06/21 00:33:32]

Since PK11_CipherOp can encrypt in place, you should be
able to use output_data to receive the CTR mask and pass
it to PK11_CipherOp for in-place encryption.

This will avoid one memory allocation (of ctr_mask).

+  CHECK(input_len <= output_len);
+  input_data = ctr_mask.get();
+
+  int op_len;
+  SECStatus rv = PK11_CipherOp(context,
+                               output_data,
+                               &op_len,
+                               output_len,
+                               input_data,
+                               input_len);

[Ryan Sleevi, 2011/06/20 23:47:23]

BUG: There is no guarantee that |op_len| will == |input_len| or |output_len|,
AFAICT. In the CBC mode/Encryptor::Crypt(), you see that beyond this point,
|op_len| is used for the result length (lines 126, 128, and 132).

The risk here is that the buffer past |op_len| (either due to overflow of int vs
size_t, or due to the PKCS#11 module only doing a subset of |op_len|) will
contain the mask value, rather than the encrypted mask value.

[Alpha Left Google, 2011/06/20 23:59:44]

Since this is using ECB and input_len is a multiple of 16 bytes I would like to
enforce op_len == input_len. I'll make a check here.

+  if (SECSuccess != rv)
+    return false;
+
+  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