Add AES OCB decryption.

Also, refactor to extract functionality that will be common to all AEAD modes.

Change-Id: I4bcf12c9d2d464ab1af559c69031904ffae45e25
diff --git a/aes_operation.cpp b/aes_operation.cpp
index 346c857..0c7a86c 100644
--- a/aes_operation.cpp
+++ b/aes_operation.cpp
@@ -14,6 +14,8 @@
  * limitations under the License.
  */
 
+#include <stdio.h>
+
 #include <openssl/aes.h>
 #include <openssl/rand.h>
 
@@ -21,85 +23,58 @@
 
 namespace keymaster {
 
-keymaster_error_t AesOcbEncryptOperation::Begin() {
-    chunk_.reset(new uint8_t[chunk_length_]);
-    if (!chunk_.get())
-        return KM_ERROR_MEMORY_ALLOCATION_FAILED;
+keymaster_error_t AesOcbOperation::Initialize(uint8_t* key, size_t key_size, size_t nonce_length,
+                                              size_t tag_length) {
+    if (tag_length > MAX_TAG_LENGTH ||nonce_length > MAX_NONCE_LENGTH)
+        return KM_ERROR_INVALID_KEY_BLOB;
 
-    if (!RAND_bytes(nonce_, NONCE_LENGTH))
-        return KM_ERROR_UNKNOWN_ERROR;
-
-    if (ae_init(ctx_.get(), key_, key_size_, array_size(nonce_), tag_length_) != AE_SUCCESS) {
-        memset_s(ctx_.get(), 0, ae_ctx_sizeof());
+    if (ae_init(ctx(), key, key_size, nonce_length, tag_length) != AE_SUCCESS) {
+        memset_s(ctx(), 0, ae_ctx_sizeof());
         return KM_ERROR_UNKNOWN_ERROR;
     }
-
     return KM_ERROR_OK;
 }
 
-keymaster_error_t AesOcbEncryptOperation::Update(const Buffer& input, Buffer* output,
-                                                 size_t* input_consumed) {
-    const uint8_t* plaintext = input.peek_read();
-    const uint8_t* plaintext_end = plaintext + input.available_read();
+keymaster_error_t AesOcbOperation::EncryptChunk(const uint8_t* nonce, size_t /* nonce_length */,
+                                                size_t tag_length,
+                                                const keymaster_blob_t additional_data,
+                                                uint8_t* chunk, size_t chunk_size, Buffer* output) {
+    if (!ctx())
+        return KM_ERROR_UNKNOWN_ERROR;
+    uint8_t __attribute__((aligned(16))) tag[MAX_TAG_LENGTH];
 
-    while (plaintext + chunk_unfilled_space() < plaintext_end) {
-        size_t to_process = chunk_unfilled_space();
-        memcpy(chunk_.get() + chunk_offset_, plaintext, to_process);
-        chunk_offset_ += to_process;
-        assert(chunk_offset_ == chunk_length_);
+    // Encrypt chunk in place.
+    int ae_err = ae_encrypt(ctx(), nonce, chunk, chunk_size, additional_data.data,
+                            additional_data.data_length, chunk, tag, AE_FINALIZE);
 
-        keymaster_error_t error = ProcessChunk(output);
-        if (error != KM_ERROR_OK)
-            return error;
-        plaintext += to_process;
-    }
-
-    // Copy remaining data into chunk_.
-    assert(plaintext_end - plaintext < chunk_unfilled_space());
-    memcpy(chunk_.get() + chunk_offset_, plaintext, plaintext_end - plaintext);
-    chunk_offset_ += (plaintext_end - plaintext);
-
-    *input_consumed = input.available_read();
-    return KM_ERROR_OK;
-}
-
-keymaster_error_t AesOcbEncryptOperation::Finish(const Buffer& /* signature */, Buffer* output) {
-    keymaster_error_t error = KM_ERROR_OK;
-    if (chunk_offset_ > 0)
-        error = ProcessChunk(output);
-    return error;
-}
-
-keymaster_error_t AesOcbEncryptOperation::ProcessChunk(Buffer* output) {
-    if (!nonce_written_) {
-        if (!output->reserve(NONCE_LENGTH + chunk_length_ + tag_length_))
-            return KM_ERROR_MEMORY_ALLOCATION_FAILED;
-        output->write(nonce_, NONCE_LENGTH);
-        nonce_written_ = true;
-    } else {
-        IncrementNonce();
-    }
-
-    if (!output->reserve(output->available_read() + chunk_offset_ + tag_length_))
-        return KM_ERROR_MEMORY_ALLOCATION_FAILED;
-
-    int ae_err = ae_encrypt(ctx_.get(), nonce_, chunk_.get(), chunk_offset_, additional_data_.data,
-                            additional_data_.data_length, output->peek_write(), tag_, AE_FINALIZE);
     if (ae_err < 0)
         return KM_ERROR_UNKNOWN_ERROR;
-    output->advance_write(chunk_offset_);
-    chunk_offset_ = 0;
+    assert(ae_err == (int)buffered_data_length());
 
-    // Output the tag.
-    output->write(tag_, tag_length_);
+    output->write(chunk, buffered_data_length());
+    output->write(tag, tag_length);
 
     return KM_ERROR_OK;
 }
 
-void AesOcbEncryptOperation::IncrementNonce() {
-    for (int i = NONCE_LENGTH - 1; i > 0; --i)
-        if (++nonce_[i])
-            break;
+keymaster_error_t AesOcbOperation::DecryptChunk(const uint8_t* nonce, size_t /* nonce_length */,
+                                                const uint8_t* tag, size_t /* tag_length */,
+                                                const keymaster_blob_t additional_data,
+                                                uint8_t* chunk, size_t chunk_size, Buffer* output) {
+    if (!ctx())
+        return KM_ERROR_UNKNOWN_ERROR;
+
+    // Decrypt chunk in place
+    int ae_err = ae_decrypt(ctx(), nonce, chunk, chunk_size, additional_data.data,
+                            additional_data.data_length, chunk, tag, AE_FINALIZE);
+    if (ae_err == AE_INVALID)
+        return KM_ERROR_VERIFICATION_FAILED;
+    else if (ae_err < 0)
+        return KM_ERROR_UNKNOWN_ERROR;
+    assert(ae_err == (int)buffered_data_length());
+    output->write(chunk, chunk_size);
+
+    return KM_ERROR_OK;
 }
 
 }  // namespace keymaster