Add MAC length checking for AES OCB.

Change-Id: I9ea564a00af3beb168cd5dc24a52c5f9c8b856a0
diff --git a/aes_key.cpp b/aes_key.cpp
index 00c22c8..5bd1f8a 100644
--- a/aes_key.cpp
+++ b/aes_key.cpp
@@ -51,18 +51,30 @@
         if (!authorizations.GetTagValue(TAG_CHUNK_LENGTH, &chunk_length) ||
             !chunk_length_is_supported(chunk_length)) {
             // TODO(swillden): Add a better error code for this.
-            *error = KM_ERROR_INVALID_INPUT_LENGTH;
+            *error = KM_ERROR_INVALID_ARGUMENT;
             return NULL;
         }
     }
 
     // Padding is optional
     keymaster_padding_t padding;
-    if (authorizations.GetTagValue(TAG_PADDING, &padding) && !padding_is_supported(padding)) {
+    if (authorizations.GetTagValue(TAG_PADDING, &padding) &&
+        !padding_is_supported(block_mode, padding)) {
         *error = KM_ERROR_UNSUPPORTED_PADDING_MODE;
         return NULL;
     }
 
+    // Check required for some modes.
+    uint32_t mac_length;
+    if (mac_length_required(block_mode)) {
+        if (!authorizations.GetTagValue(TAG_MAC_LENGTH, &mac_length) ||
+            !mac_length_is_supported(block_mode, mac_length)) {
+            // TODO(swillden): Add a better error code for this.
+            *error = KM_ERROR_INVALID_ARGUMENT;
+            return NULL;
+        }
+    }
+
     // Verify purpose is compatible with block mode.
     if (!ModeAndPurposesAreCompatible(authorizations, block_mode, logger)) {
         *error = KM_ERROR_INCOMPATIBLE_BLOCK_MODE;
diff --git a/aes_key.h b/aes_key.h
index c505741..87a83e7 100644
--- a/aes_key.h
+++ b/aes_key.h
@@ -28,6 +28,7 @@
 class AesKey : public Key {
   public:
     static const int MAX_KEY_SIZE = 32;
+    static const int MAX_MAC_LENGTH = 16;
 
     AesKey(const UnencryptedKeyBlob& blob, const Logger& logger, keymaster_error_t* error);
     ~AesKey();
@@ -47,7 +48,13 @@
         return (chunk_length <= MAX_AES_CHUNK_LENGTH);
     }
 
-    static bool padding_is_supported(keymaster_padding_t padding) {
+    static bool mac_length_required(keymaster_block_mode_t) { return true; }
+
+    static bool mac_length_is_supported(keymaster_block_mode_t, uint32_t mac_length) {
+        return (mac_length <= MAX_MAC_LENGTH);
+    }
+
+    static bool padding_is_supported(keymaster_block_mode_t, keymaster_padding_t padding) {
         return (padding == KM_PAD_NONE);
     }
 
diff --git a/google_keymaster_test.cpp b/google_keymaster_test.cpp
index c59af35..4f17259 100644
--- a/google_keymaster_test.cpp
+++ b/google_keymaster_test.cpp
@@ -319,7 +319,7 @@
         Authorization(TAG_PURPOSE, KM_PURPOSE_DECRYPT),
         Authorization(TAG_ALGORITHM, KM_ALGORITHM_AES), Authorization(TAG_KEY_SIZE, 128),
         Authorization(TAG_BLOCK_MODE, KM_MODE_OCB), Authorization(TAG_CHUNK_LENGTH, 4096),
-        Authorization(TAG_PADDING, KM_PAD_NONE),
+        Authorization(TAG_MAC_LENGTH, 16), Authorization(TAG_PADDING, KM_PAD_NONE),
     };
     req_.key_description.Reinitialize(params, array_length(params));
     device.GenerateKey(req_, &rsp_);
@@ -332,7 +332,7 @@
         Authorization(TAG_PURPOSE, KM_PURPOSE_DECRYPT),
         Authorization(TAG_ALGORITHM, KM_ALGORITHM_AES), Authorization(TAG_KEY_SIZE, 129),
         Authorization(TAG_BLOCK_MODE, KM_MODE_OCB), Authorization(TAG_CHUNK_LENGTH, 4096),
-        Authorization(TAG_PADDING, KM_PAD_NONE),
+        Authorization(TAG_MAC_LENGTH, 16), Authorization(TAG_PADDING, KM_PAD_NONE),
     };
     req_.key_description.Reinitialize(params, array_length(params));
     device.GenerateKey(req_, &rsp_);
@@ -343,9 +343,8 @@
     keymaster_key_param_t params[] = {
         Authorization(TAG_PURPOSE, KM_PURPOSE_ENCRYPT),
         Authorization(TAG_PURPOSE, KM_PURPOSE_DECRYPT),
-        Authorization(TAG_ALGORITHM, KM_ALGORITHM_AES),
-        Authorization(TAG_BLOCK_MODE, KM_MODE_OCB),
-        Authorization(TAG_CHUNK_LENGTH, 4096),
+        Authorization(TAG_ALGORITHM, KM_ALGORITHM_AES), Authorization(TAG_BLOCK_MODE, KM_MODE_OCB),
+        Authorization(TAG_MAC_LENGTH, 16), Authorization(TAG_CHUNK_LENGTH, 4096),
         Authorization(TAG_PADDING, KM_PAD_NONE),
     };
 
@@ -368,7 +367,7 @@
     };
     req_.key_description.Reinitialize(params, array_length(params));
     device.GenerateKey(req_, &rsp_);
-    EXPECT_EQ(KM_ERROR_INVALID_INPUT_LENGTH, rsp_.error);
+    EXPECT_EQ(KM_ERROR_INVALID_ARGUMENT, rsp_.error);
 }
 
 TEST_F(NewKeyGeneration, AesEcbUnsupported) {
@@ -396,6 +395,19 @@
     EXPECT_EQ(KM_ERROR_UNSUPPORTED_PADDING_MODE, rsp_.error);
 }
 
+TEST_F(NewKeyGeneration, AesOcbInvalidMacLength) {
+    keymaster_key_param_t params[] = {
+        Authorization(TAG_PURPOSE, KM_PURPOSE_ENCRYPT),
+        Authorization(TAG_PURPOSE, KM_PURPOSE_DECRYPT),
+        Authorization(TAG_ALGORITHM, KM_ALGORITHM_AES), Authorization(TAG_KEY_SIZE, 128),
+        Authorization(TAG_BLOCK_MODE, KM_MODE_OCB), Authorization(TAG_CHUNK_LENGTH, 4096),
+        Authorization(TAG_MAC_LENGTH, 17), Authorization(TAG_PADDING, KM_PAD_NONE),
+    };
+    req_.key_description.Reinitialize(params, array_length(params));
+    device.GenerateKey(req_, &rsp_);
+    EXPECT_EQ(KM_ERROR_INVALID_ARGUMENT, rsp_.error);
+}
+
 typedef KeymasterTest GetKeyCharacteristics;
 TEST_F(GetKeyCharacteristics, SimpleRsa) {
     keymaster_key_param_t params[] = {