chaps: Prevent interspersal of incremental and one-shot operations.

PKCS#11 spec has variants of the following (from C_Encrypt):
   "C_Encrypt can not be user to terminate a multi-part operation, and
    must be called after C_EncryptInit without intervening
    C_EncryptUpdate calls."

So track whether a one-shot or multi-part operation has been started, and
reject the other style if so.

BUG=None
TEST=Chaps unit tests (with ASAN) plus PKCS11 tests

Change-Id: I484dca635a12246a30c461e6ea7113697e085d26
Reviewed-on: https://chromium-review.googlesource.com/221937
Reviewed-by: Darren Krahn <dkrahn@chromium.org>
Commit-Queue: David Drysdale <drysdale@google.com>
Tested-by: David Drysdale <drysdale@google.com>
diff --git a/chaps/session_impl.cc b/chaps/session_impl.cc
index 76ee65f..033fe31 100644
--- a/chaps/session_impl.cc
+++ b/chaps/session_impl.cc
@@ -280,6 +280,23 @@
     LOG(ERROR) << "Operation is not initialized.";
     return CKR_OPERATION_NOT_INITIALIZED;
   }
+  if (context->is_finished_) {
+    LOG(ERROR) << "Operation is finished.";
+    return CKR_OPERATION_ACTIVE;
+  }
+  context->is_incremental_ = true;
+  return OperationUpdateInternal(operation,
+                                 data_in,
+                                 required_out_length,
+                                 data_out);
+}
+
+CK_RV SessionImpl::OperationUpdateInternal(OperationType operation,
+                                           const string& data_in,
+                                           int* required_out_length,
+                                           string* data_out) {
+  CHECK(operation < kNumOperationTypes);
+  OperationContext* context = &operation_context_[operation];
   if (context->is_cipher_) {
     return CipherUpdate(context, data_in, required_out_length, data_out);
   } else if (context->is_digest_) {
@@ -310,6 +327,19 @@
     LOG(ERROR) << "Operation is not initialized.";
     return CKR_OPERATION_NOT_INITIALIZED;
   }
+  if (!context->is_incremental_ && context->is_finished_) {
+    LOG(ERROR) << "Operation is not incremental.";
+    return CKR_OPERATION_ACTIVE;
+  }
+  context->is_incremental_ = true;
+  return OperationFinalInternal(operation, required_out_length, data_out);
+}
+
+CK_RV SessionImpl::OperationFinalInternal(OperationType operation,
+                                          int* required_out_length,
+                                          string* data_out) {
+  CHECK(operation < kNumOperationTypes);
+  OperationContext* context = &operation_context_[operation];
   context->is_valid_ = false;
   // Complete the operation if it has not already been done.
   if (!context->is_finished_) {
@@ -387,15 +417,23 @@
                                        string* data_out) {
   CHECK(operation < kNumOperationTypes);
   OperationContext* context = &operation_context_[operation];
+  if (!context->is_valid_) {
+    LOG(ERROR) << "Operation is not initialized.";
+    return CKR_OPERATION_NOT_INITIALIZED;
+  }
+  if (context->is_incremental_) {
+    LOG(ERROR) << "Operation is incremental.";
+    return CKR_OPERATION_ACTIVE;
+  }
   CK_RV result = CKR_OK;
   if (!context->is_finished_) {
     string update, final;
     int max = INT_MAX;
-    result = OperationUpdate(operation, data_in, &max, &update);
+    result = OperationUpdateInternal(operation, data_in, &max, &update);
     if (result != CKR_OK)
       return result;
     max = INT_MAX;
-    result = OperationFinal(operation, &max, &final);
+    result = OperationFinalInternal(operation, &max, &final);
     if (result != CKR_OK)
       return result;
     context->data_ = update + final;
@@ -1332,6 +1370,7 @@
   is_cipher_ = false;
   is_digest_ = false;
   is_hmac_ = false;
+  is_incremental_ = false;
   is_finished_ = false;
   key_ = NULL;
   data_.clear();
diff --git a/chaps/session_impl.h b/chaps/session_impl.h
index 37307ce..a725831 100644
--- a/chaps/session_impl.h
+++ b/chaps/session_impl.h
@@ -111,6 +111,7 @@
     bool is_cipher_;  // Set to true when cipher_context_ is valid.
     bool is_digest_;  // Set to true when digest_context_ is valid.
     bool is_hmac_;  // Set to true when hmac_context_ is valid.
+    bool is_incremental_;  // Set when an incremental operation is performed.
     bool is_finished_;  // Set to true when the operation completes.
     union {
       EVP_CIPHER_CTX cipher_context_;
@@ -133,6 +134,13 @@
                       CK_OBJECT_CLASS object_class,
                       CK_KEY_TYPE key_type);
   bool IsValidMechanism(OperationType operation, CK_MECHANISM_TYPE mechanism);
+  CK_RV OperationUpdateInternal(OperationType operation,
+                                const std::string& data_in,
+                                int* required_out_length,
+                                std::string* data_out);
+  CK_RV OperationFinalInternal(OperationType operation,
+                               int* required_out_length,
+                               std::string* data_out);
   CK_RV CipherInit(bool is_encrypt,
                    CK_MECHANISM_TYPE mechanism,
                    const std::string& mechanism_parameter,
diff --git a/chaps/session_test.cc b/chaps/session_test.cc
index 23015b6..4a197b4 100644
--- a/chaps/session_test.cc
+++ b/chaps/session_test.cc
@@ -440,6 +440,62 @@
   EXPECT_EQ(CKR_OK, session_->VerifyFinal(out));
 }
 
+// Test empty multi-part operation.
+TEST_F(TestSession, FinalWithNoUpdate) {
+  EXPECT_EQ(CKR_OK, session_->OperationInit(kDigest, CKM_SHA_1, "", NULL));
+  int len = 20;
+  string out;
+  EXPECT_EQ(CKR_OK, session_->OperationFinal(kDigest, &len, &out));
+  EXPECT_EQ(20, len);
+}
+
+// Test multi-part and single-part operations inhibit each other.
+TEST_F(TestSession, UpdateOperationPreventsSinglePart) {
+  string in(30, 'A');
+  EXPECT_EQ(CKR_OK, session_->OperationInit(kDigest, CKM_SHA_1, "", NULL));
+  EXPECT_EQ(CKR_OK, session_->OperationUpdate(kDigest,
+                                              in.substr(0, 10),
+                                              NULL,
+                                              NULL));
+  int len = 0;
+  string out;
+  EXPECT_EQ(CKR_OPERATION_ACTIVE,
+            session_->OperationSinglePart(kDigest,
+                                          in.substr(10, 20),
+                                          &len,
+                                          &out));
+}
+
+TEST_F(TestSession, SinglePartOperationPreventsUpdate) {
+  string in(30, 'A');
+  EXPECT_EQ(CKR_OK, session_->OperationInit(kDigest, CKM_SHA_1, "", NULL));
+  string out;
+  int len = 0;
+  // Perform a single part operation but leave the output to be collected.
+  EXPECT_EQ(CKR_BUFFER_TOO_SMALL,
+            session_->OperationSinglePart(kDigest, in, &len, &out));
+
+  EXPECT_EQ(CKR_OPERATION_ACTIVE,
+            session_->OperationUpdate(kDigest,
+                                      in.substr(10, 10),
+                                      NULL,
+                                      NULL));
+}
+
+TEST_F(TestSession, SinglePartOperationPreventsFinal) {
+  string in(30, 'A');
+  EXPECT_EQ(CKR_OK, session_->OperationInit(kDigest, CKM_SHA_1, "", NULL));
+  string out;
+  int len = 0;
+  // Perform a single part operation but leave the output to be collected.
+  EXPECT_EQ(CKR_BUFFER_TOO_SMALL,
+            session_->OperationSinglePart(kDigest, in, &len, &out));
+
+  len = 0;
+  EXPECT_EQ(CKR_OPERATION_ACTIVE,
+            session_->OperationFinal(kDigest, &len, &out));
+}
+
 // Test RSA PKCS #1 encryption.
 TEST_F(TestSession, RSAEncrypt) {
   const Object* pub = NULL;