Check set_padding has been called in mbedtls_cipher_finish

Check set_padding has been called in mbedtls_cipher_finish
in modes that require padding.

Signed-off-by: Waleed Elmelegy <waleed.elmelegy@arm.com>
diff --git a/library/cipher.c b/library/cipher.c
index de7f837..995d3d2 100644
--- a/library/cipher.c
+++ b/library/cipher.c
@@ -268,17 +268,6 @@
 
     ctx->cipher_info = cipher_info;
 
-#if defined(MBEDTLS_CIPHER_MODE_WITH_PADDING)
-    /*
-     * Ignore possible errors caused by a cipher mode that doesn't use padding
-     */
-#if defined(MBEDTLS_CIPHER_PADDING_PKCS7)
-    (void) mbedtls_cipher_set_padding_mode(ctx, MBEDTLS_PADDING_PKCS7);
-#else
-    (void) mbedtls_cipher_set_padding_mode(ctx, MBEDTLS_PADDING_NONE);
-#endif
-#endif /* MBEDTLS_CIPHER_MODE_WITH_PADDING */
-
     return 0;
 }
 
@@ -1027,6 +1016,16 @@
 
     *olen = 0;
 
+#if defined(MBEDTLS_CIPHER_MODE_WITH_PADDING)
+    /* CBC mode requires padding so we make sure a call to
+     * mbedtls_cipher_set_padding_mode has been done successfully. */
+    if (MBEDTLS_MODE_CBC == ((mbedtls_cipher_mode_t) ctx->cipher_info->mode)) {
+        if (ctx->get_padding == NULL) {
+            return MBEDTLS_ERR_CIPHER_BAD_INPUT_DATA;
+        }
+    }
+#endif
+
     if (MBEDTLS_MODE_CFB == ((mbedtls_cipher_mode_t) ctx->cipher_info->mode) ||
         MBEDTLS_MODE_OFB == ((mbedtls_cipher_mode_t) ctx->cipher_info->mode) ||
         MBEDTLS_MODE_CTR == ((mbedtls_cipher_mode_t) ctx->cipher_info->mode) ||
diff --git a/tests/suites/test_suite_cipher.aes.data b/tests/suites/test_suite_cipher.aes.data
index 134970f..10bb93e 100644
--- a/tests/suites/test_suite_cipher.aes.data
+++ b/tests/suites/test_suite_cipher.aes.data
@@ -2257,3 +2257,7 @@
 Cipher Corner Case behaviours
 depends_on:MBEDTLS_AES_C
 cipher_special_behaviours:
+
+Check set padding
+depends_on:MBEDTLS_AES_C:MBEDTLS_CIPHER_MODE_CBC
+check_set_padding:
diff --git a/tests/suites/test_suite_cipher.function b/tests/suites/test_suite_cipher.function
index 192e498..4fda4f9 100644
--- a/tests/suites/test_suite_cipher.function
+++ b/tests/suites/test_suite_cipher.function
@@ -2,6 +2,7 @@
 #include "mbedtls/cipher.h"
 #include "mbedtls/aes.h"
 
+
 #if defined(MBEDTLS_GCM_C)
 #include "mbedtls/gcm.h"
 #endif
@@ -400,6 +401,16 @@
     if (-1 != pad_mode) {
         TEST_ASSERT(0 == mbedtls_cipher_set_padding_mode(&ctx_dec, pad_mode));
         TEST_ASSERT(0 == mbedtls_cipher_set_padding_mode(&ctx_enc, pad_mode));
+    } else {
+        if (ctx_dec.cipher_info->mode == MBEDTLS_MODE_CBC) {
+#if defined(MBEDTLS_CIPHER_PADDING_PKCS7)
+            TEST_ASSERT(0 == mbedtls_cipher_set_padding_mode(&ctx_dec, MBEDTLS_PADDING_PKCS7));
+            TEST_ASSERT(0 == mbedtls_cipher_set_padding_mode(&ctx_enc, MBEDTLS_PADDING_PKCS7));
+#else
+            TEST_ASSERT(0 == mbedtls_cipher_set_padding_mode(&ctx_dec, MBEDTLS_PADDING_NONE));
+            TEST_ASSERT(0 == mbedtls_cipher_set_padding_mode(&ctx_enc, MBEDTLS_PADDING_NONE));
+#endif
+        }
     }
 #else
     (void) pad_mode;
@@ -601,6 +612,16 @@
 
     TEST_ASSERT(0 == mbedtls_cipher_reset(&ctx_dec));
 
+#if defined(MBEDTLS_CIPHER_MODE_WITH_PADDING)
+    if (ctx_dec.cipher_info->mode == MBEDTLS_MODE_CBC) {
+#if defined(MBEDTLS_CIPHER_PADDING_PKCS7)
+        TEST_ASSERT(0 == mbedtls_cipher_set_padding_mode(&ctx_dec, MBEDTLS_PADDING_PKCS7));
+#else
+        TEST_ASSERT(0 == mbedtls_cipher_set_padding_mode(&ctx_dec, MBEDTLS_PADDING_NONE));
+#endif
+    }
+#endif /* MBEDTLS_CIPHER_MODE_WITH_PADDING */
+
 #if defined(MBEDTLS_GCM_C) || defined(MBEDTLS_CHACHAPOLY_C)
     int expected = (cipher_info->mode == MBEDTLS_MODE_GCM ||
                     cipher_info->type == MBEDTLS_CIPHER_CHACHA20_POLY1305) ?
@@ -684,6 +705,16 @@
     if (-1 != pad_mode) {
         TEST_ASSERT(0 == mbedtls_cipher_set_padding_mode(&ctx_dec, pad_mode));
         TEST_ASSERT(0 == mbedtls_cipher_set_padding_mode(&ctx_enc, pad_mode));
+    } else {
+        if (ctx_dec.cipher_info->mode == MBEDTLS_MODE_CBC) {
+#if defined(MBEDTLS_CIPHER_PADDING_PKCS7)
+            TEST_ASSERT(0 == mbedtls_cipher_set_padding_mode(&ctx_dec, MBEDTLS_PADDING_PKCS7));
+            TEST_ASSERT(0 == mbedtls_cipher_set_padding_mode(&ctx_enc, MBEDTLS_PADDING_PKCS7));
+#else
+            TEST_ASSERT(0 == mbedtls_cipher_set_padding_mode(&ctx_dec, MBEDTLS_PADDING_NONE));
+            TEST_ASSERT(0 == mbedtls_cipher_set_padding_mode(&ctx_enc, MBEDTLS_PADDING_NONE));
+#endif
+        }
     }
 #else
     (void) pad_mode;
@@ -1198,3 +1229,37 @@
     mbedtls_cipher_free(&ctx_enc);
 }
 /* END_CASE */
+
+/* BEGIN_CASE */
+void check_set_padding()
+{
+    mbedtls_cipher_context_t ctx;
+    unsigned char key[16] = { 0 };
+    unsigned char iv[16] = { 0 };
+    unsigned char input[16] = { 0 };
+    unsigned char output[32] = { 0 };
+    size_t outlen = 0;
+    int result = MBEDTLS_ERR_CIPHER_BAD_INPUT_DATA;
+    const mbedtls_cipher_info_t *cipher_info;
+
+    cipher_info = mbedtls_cipher_info_from_type(MBEDTLS_CIPHER_AES_128_CBC);
+
+    mbedtls_cipher_init(&ctx);
+
+    TEST_ASSERT(0 == mbedtls_cipher_setup(&ctx, cipher_info));
+
+    TEST_ASSERT(0 == mbedtls_cipher_setkey(&ctx, key, 128, MBEDTLS_ENCRYPT));
+
+    TEST_ASSERT(result == mbedtls_cipher_crypt(&ctx, iv, sizeof(iv), input,
+                                               sizeof(input), output, &outlen));
+
+#if defined(MBEDTLS_CIPHER_MODE_WITH_PADDING)
+    TEST_ASSERT(0 == mbedtls_cipher_set_padding_mode(&ctx, MBEDTLS_PADDING_NONE));
+    TEST_ASSERT(0 == mbedtls_cipher_crypt(&ctx, iv, sizeof(iv), input,
+                                          sizeof(input), output, &outlen));
+#endif
+
+exit:
+    mbedtls_cipher_free(&ctx);
+}
+/* END_CASE */