Catch failures of AES or DES operations

Declare all AES and DES functions that return int as needing to have
their result checked, and do check the result in our code.

A DES or AES block operation can fail in alternative implementations of
mbedtls_internal_aes_encrypt() (under MBEDTLS_AES_ENCRYPT_ALT),
mbedtls_internal_aes_decrypt() (under MBEDTLS_AES_DECRYPT_ALT),
mbedtls_des_crypt_ecb() (under MBEDTLS_DES_CRYPT_ECB_ALT),
mbedtls_des3_crypt_ecb() (under MBEDTLS_DES3_CRYPT_ECB_ALT).
A failure can happen if the accelerator peripheral is in a bad state.
Several block modes were not catching the error.

This commit does the following code changes, grouped together to avoid
having an intermediate commit where the build fails:

* Add MBEDTLS_CHECK_RETURN to all functions returning int in aes.h and des.h.
* Fix all places where this causes a GCC warning, indicating that our code
  was not properly checking the result of an AES operation:
    * In library code: on failure, goto exit and return ret.
    * In pkey programs: goto exit.
    * In the benchmark program: exit (not ideal since there's no error
      message, but it's what the code currently does for failures).
    * In test code: TEST_ASSERT.
* Changelog entry.

Signed-off-by: Gilles Peskine <Gilles.Peskine@arm.com>
diff --git a/ChangeLog.d/check-return.txt b/ChangeLog.d/check-return.txt
new file mode 100644
index 0000000..062b8a1
--- /dev/null
+++ b/ChangeLog.d/check-return.txt
@@ -0,0 +1,12 @@
+Bugfix
+   * Failures of alternative implementations of AES or DES single-block
+     functions enabled with MBEDTLS_AES_ENCRYPT_ALT, MBEDTLS_AES_DECRYPT_ALT,
+     MBEDTLS_DES_CRYPT_ECB_ALT or MBEDTLS_DES3_CRYPT_ECB_ALT were ignored.
+     This does not concern the implementation provided with Mbed TLS,
+     where this function cannot fail, or full-module replacements with
+     MBEDTLS_AES_ALT or MBEDTLS_DES_ALT. Reported by Armelle Duboc in #1092.
+
+Changes
+   * Warn if errors from AES or DES functions are ignored. This is currently
+     supported on GCC-like compilers and on MSVC and can be configured by
+     setting MBEDTLS_CHECK_RETURN in config.h.
diff --git a/include/mbedtls/aes.h b/include/mbedtls/aes.h
index a74e413..4662260 100644
--- a/include/mbedtls/aes.h
+++ b/include/mbedtls/aes.h
@@ -175,6 +175,7 @@
  * \return         \c 0 on success.
  * \return         #MBEDTLS_ERR_AES_INVALID_KEY_LENGTH on failure.
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_aes_setkey_enc( mbedtls_aes_context *ctx, const unsigned char *key,
                     unsigned int keybits );
 
@@ -193,6 +194,7 @@
  * \return         \c 0 on success.
  * \return         #MBEDTLS_ERR_AES_INVALID_KEY_LENGTH on failure.
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_aes_setkey_dec( mbedtls_aes_context *ctx, const unsigned char *key,
                     unsigned int keybits );
 
@@ -213,6 +215,7 @@
  * \return         \c 0 on success.
  * \return         #MBEDTLS_ERR_AES_INVALID_KEY_LENGTH on failure.
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_aes_xts_setkey_enc( mbedtls_aes_xts_context *ctx,
                                 const unsigned char *key,
                                 unsigned int keybits );
@@ -233,6 +236,7 @@
  * \return         \c 0 on success.
  * \return         #MBEDTLS_ERR_AES_INVALID_KEY_LENGTH on failure.
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_aes_xts_setkey_dec( mbedtls_aes_xts_context *ctx,
                                 const unsigned char *key,
                                 unsigned int keybits );
@@ -261,6 +265,7 @@
 
  * \return         \c 0 on success.
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_aes_crypt_ecb( mbedtls_aes_context *ctx,
                     int mode,
                     const unsigned char input[16],
@@ -308,6 +313,7 @@
  * \return         #MBEDTLS_ERR_AES_INVALID_INPUT_LENGTH
  *                 on failure.
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_aes_crypt_cbc( mbedtls_aes_context *ctx,
                     int mode,
                     size_t length,
@@ -352,6 +358,7 @@
  *                     smaller than an AES block in size (16 Bytes) or if \p
  *                     length is larger than 2^20 blocks (16 MiB).
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_aes_crypt_xts( mbedtls_aes_xts_context *ctx,
                            int mode,
                            size_t length,
@@ -400,6 +407,7 @@
  *
  * \return         \c 0 on success.
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_aes_crypt_cfb128( mbedtls_aes_context *ctx,
                        int mode,
                        size_t length,
@@ -444,6 +452,7 @@
  *
  * \return         \c 0 on success.
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_aes_crypt_cfb8( mbedtls_aes_context *ctx,
                     int mode,
                     size_t length,
@@ -498,6 +507,7 @@
  *
  * \return         \c 0 on success.
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_aes_crypt_ofb( mbedtls_aes_context *ctx,
                        size_t length,
                        size_t *iv_off,
@@ -584,6 +594,7 @@
  *
  * \return                 \c 0 on success.
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_aes_crypt_ctr( mbedtls_aes_context *ctx,
                        size_t length,
                        size_t *nc_off,
@@ -604,6 +615,7 @@
  *
  * \return          \c 0 on success.
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_internal_aes_encrypt( mbedtls_aes_context *ctx,
                                   const unsigned char input[16],
                                   unsigned char output[16] );
@@ -619,6 +631,7 @@
  *
  * \return          \c 0 on success.
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_internal_aes_decrypt( mbedtls_aes_context *ctx,
                                   const unsigned char input[16],
                                   unsigned char output[16] );
@@ -668,6 +681,7 @@
  * \return         \c 0 on success.
  * \return         \c 1 on failure.
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_aes_self_test( int verbose );
 
 #endif /* MBEDTLS_SELF_TEST */
diff --git a/include/mbedtls/des.h b/include/mbedtls/des.h
index 29b96db..a82ecf8 100644
--- a/include/mbedtls/des.h
+++ b/include/mbedtls/des.h
@@ -147,6 +147,7 @@
  *                 security risk. We recommend considering stronger ciphers
  *                 instead.
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_des_key_check_key_parity( const unsigned char key[MBEDTLS_DES_KEY_SIZE] );
 
 /**
@@ -160,6 +161,7 @@
  *                 security risk. We recommend considering stronger ciphers
  *                 instead.
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_des_key_check_weak( const unsigned char key[MBEDTLS_DES_KEY_SIZE] );
 
 /**
@@ -174,6 +176,7 @@
  *                 security risk. We recommend considering stronger ciphers
  *                 instead.
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_des_setkey_enc( mbedtls_des_context *ctx, const unsigned char key[MBEDTLS_DES_KEY_SIZE] );
 
 /**
@@ -188,6 +191,7 @@
  *                 security risk. We recommend considering stronger ciphers
  *                 instead.
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_des_setkey_dec( mbedtls_des_context *ctx, const unsigned char key[MBEDTLS_DES_KEY_SIZE] );
 
 /**
@@ -198,6 +202,7 @@
  *
  * \return         0
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_des3_set2key_enc( mbedtls_des3_context *ctx,
                       const unsigned char key[MBEDTLS_DES_KEY_SIZE * 2] );
 
@@ -209,6 +214,7 @@
  *
  * \return         0
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_des3_set2key_dec( mbedtls_des3_context *ctx,
                       const unsigned char key[MBEDTLS_DES_KEY_SIZE * 2] );
 
@@ -220,6 +226,7 @@
  *
  * \return         0
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_des3_set3key_enc( mbedtls_des3_context *ctx,
                       const unsigned char key[MBEDTLS_DES_KEY_SIZE * 3] );
 
@@ -231,6 +238,7 @@
  *
  * \return         0
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_des3_set3key_dec( mbedtls_des3_context *ctx,
                       const unsigned char key[MBEDTLS_DES_KEY_SIZE * 3] );
 
@@ -247,6 +255,7 @@
  *                 security risk. We recommend considering stronger ciphers
  *                 instead.
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_des_crypt_ecb( mbedtls_des_context *ctx,
                     const unsigned char input[8],
                     unsigned char output[8] );
@@ -274,6 +283,7 @@
  *                 security risk. We recommend considering stronger ciphers
  *                 instead.
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_des_crypt_cbc( mbedtls_des_context *ctx,
                     int mode,
                     size_t length,
@@ -291,6 +301,7 @@
  *
  * \return         0 if successful
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_des3_crypt_ecb( mbedtls_des3_context *ctx,
                      const unsigned char input[8],
                      unsigned char output[8] );
@@ -316,6 +327,7 @@
  *
  * \return         0 if successful, or MBEDTLS_ERR_DES_INVALID_INPUT_LENGTH
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_des3_crypt_cbc( mbedtls_des3_context *ctx,
                      int mode,
                      size_t length,
@@ -346,6 +358,7 @@
  *
  * \return         0 if successful, or 1 if the test failed
  */
+MBEDTLS_CHECK_RETURN
 int mbedtls_des_self_test( int verbose );
 
 #endif /* MBEDTLS_SELF_TEST */
diff --git a/library/aes.c b/library/aes.c
index 3f61642..e4bcfa0 100644
--- a/library/aes.c
+++ b/library/aes.c
@@ -1052,6 +1052,7 @@
                     unsigned char *output )
 {
     int i;
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     unsigned char temp[16];
 
     AES_VALIDATE_RET( ctx != NULL );
@@ -1081,7 +1082,9 @@
         while( length > 0 )
         {
             memcpy( temp, input, 16 );
-            mbedtls_aes_crypt_ecb( ctx, mode, input, output );
+            ret = mbedtls_aes_crypt_ecb( ctx, mode, input, output );
+            if( ret != 0 )
+                goto exit;
 
             for( i = 0; i < 16; i++ )
                 output[i] = (unsigned char)( output[i] ^ iv[i] );
@@ -1100,7 +1103,9 @@
             for( i = 0; i < 16; i++ )
                 output[i] = (unsigned char)( input[i] ^ iv[i] );
 
-            mbedtls_aes_crypt_ecb( ctx, mode, output, output );
+            ret = mbedtls_aes_crypt_ecb( ctx, mode, output, output );
+            if( ret != 0 )
+                goto exit;
             memcpy( iv, output, 16 );
 
             input  += 16;
@@ -1108,8 +1113,10 @@
             length -= 16;
         }
     }
+    ret = 0;
 
-    return( 0 );
+exit:
+    return( ret );
 }
 #endif /* MBEDTLS_CIPHER_MODE_CBC */
 
@@ -1292,6 +1299,7 @@
                        unsigned char *output )
 {
     int c;
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     size_t n;
 
     AES_VALIDATE_RET( ctx != NULL );
@@ -1312,7 +1320,11 @@
         while( length-- )
         {
             if( n == 0 )
-                mbedtls_aes_crypt_ecb( ctx, MBEDTLS_AES_ENCRYPT, iv, iv );
+            {
+                ret = mbedtls_aes_crypt_ecb( ctx, MBEDTLS_AES_ENCRYPT, iv, iv );
+                if( ret != 0 )
+                    goto exit;
+            }
 
             c = *input++;
             *output++ = (unsigned char)( c ^ iv[n] );
@@ -1326,7 +1338,11 @@
         while( length-- )
         {
             if( n == 0 )
-                mbedtls_aes_crypt_ecb( ctx, MBEDTLS_AES_ENCRYPT, iv, iv );
+            {
+                ret = mbedtls_aes_crypt_ecb( ctx, MBEDTLS_AES_ENCRYPT, iv, iv );
+                if( ret != 0 )
+                    goto exit;
+            }
 
             iv[n] = *output++ = (unsigned char)( iv[n] ^ *input++ );
 
@@ -1335,8 +1351,10 @@
     }
 
     *iv_off = n;
+    ret = 0;
 
-    return( 0 );
+exit:
+    return( ret );
 }
 
 /*
@@ -1349,6 +1367,7 @@
                             const unsigned char *input,
                             unsigned char *output )
 {
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     unsigned char c;
     unsigned char ov[17];
 
@@ -1361,7 +1380,9 @@
     while( length-- )
     {
         memcpy( ov, iv, 16 );
-        mbedtls_aes_crypt_ecb( ctx, MBEDTLS_AES_ENCRYPT, iv, iv );
+        ret = mbedtls_aes_crypt_ecb( ctx, MBEDTLS_AES_ENCRYPT, iv, iv );
+        if( ret != 0 )
+            goto exit;
 
         if( mode == MBEDTLS_AES_DECRYPT )
             ov[16] = *input;
@@ -1373,8 +1394,10 @@
 
         memcpy( iv, ov + 1, 16 );
     }
+    ret = 0;
 
-    return( 0 );
+exit:
+    return( ret );
 }
 #endif /* MBEDTLS_CIPHER_MODE_CFB */
 
@@ -1436,6 +1459,7 @@
                        unsigned char *output )
 {
     int c, i;
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     size_t n;
 
     AES_VALIDATE_RET( ctx != NULL );
@@ -1453,7 +1477,9 @@
     while( length-- )
     {
         if( n == 0 ) {
-            mbedtls_aes_crypt_ecb( ctx, MBEDTLS_AES_ENCRYPT, nonce_counter, stream_block );
+            ret = mbedtls_aes_crypt_ecb( ctx, MBEDTLS_AES_ENCRYPT, nonce_counter, stream_block );
+            if( ret != 0 )
+                goto exit;
 
             for( i = 16; i > 0; i-- )
                 if( ++nonce_counter[i - 1] != 0 )
@@ -1466,8 +1492,10 @@
     }
 
     *nc_off = n;
+    ret = 0;
 
-    return( 0 );
+exit:
+    return( ret );
 }
 #endif /* MBEDTLS_CIPHER_MODE_CTR */
 
diff --git a/library/des.c b/library/des.c
index eddf55e..9ff4190 100644
--- a/library/des.c
+++ b/library/des.c
@@ -28,6 +28,7 @@
 #if defined(MBEDTLS_DES_C)
 
 #include "mbedtls/des.h"
+#include "mbedtls/error.h"
 #include "mbedtls/platform_util.h"
 
 #include <string.h>
@@ -665,6 +666,7 @@
                     unsigned char *output )
 {
     int i;
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     unsigned char temp[8];
 
     if( length % 8 )
@@ -677,7 +679,9 @@
             for( i = 0; i < 8; i++ )
                 output[i] = (unsigned char)( input[i] ^ iv[i] );
 
-            mbedtls_des_crypt_ecb( ctx, output, output );
+            ret = mbedtls_des_crypt_ecb( ctx, output, output );
+            if( ret != 0 )
+                goto exit;
             memcpy( iv, output, 8 );
 
             input  += 8;
@@ -690,7 +694,9 @@
         while( length > 0 )
         {
             memcpy( temp, input, 8 );
-            mbedtls_des_crypt_ecb( ctx, input, output );
+            ret = mbedtls_des_crypt_ecb( ctx, input, output );
+            if( ret != 0 )
+                goto exit;
 
             for( i = 0; i < 8; i++ )
                 output[i] = (unsigned char)( output[i] ^ iv[i] );
@@ -702,8 +708,10 @@
             length -= 8;
         }
     }
+    ret = 0;
 
-    return( 0 );
+exit:
+    return( ret );
 }
 #endif /* MBEDTLS_CIPHER_MODE_CBC */
 
@@ -764,6 +772,7 @@
                      unsigned char *output )
 {
     int i;
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     unsigned char temp[8];
 
     if( length % 8 )
@@ -776,7 +785,9 @@
             for( i = 0; i < 8; i++ )
                 output[i] = (unsigned char)( input[i] ^ iv[i] );
 
-            mbedtls_des3_crypt_ecb( ctx, output, output );
+            ret = mbedtls_des3_crypt_ecb( ctx, output, output );
+            if( ret != 0 )
+                goto exit;
             memcpy( iv, output, 8 );
 
             input  += 8;
@@ -789,7 +800,9 @@
         while( length > 0 )
         {
             memcpy( temp, input, 8 );
-            mbedtls_des3_crypt_ecb( ctx, input, output );
+            ret = mbedtls_des3_crypt_ecb( ctx, input, output );
+            if( ret != 0 )
+                goto exit;
 
             for( i = 0; i < 8; i++ )
                 output[i] = (unsigned char)( output[i] ^ iv[i] );
@@ -801,8 +814,10 @@
             length -= 8;
         }
     }
+    ret = 0;
 
-    return( 0 );
+exit:
+    return( ret );
 }
 #endif /* MBEDTLS_CIPHER_MODE_CBC */
 
@@ -895,39 +910,43 @@
         switch( i )
         {
         case 0:
-            mbedtls_des_setkey_dec( &ctx, des3_test_keys );
+            ret = mbedtls_des_setkey_dec( &ctx, des3_test_keys );
             break;
 
         case 1:
-            mbedtls_des_setkey_enc( &ctx, des3_test_keys );
+            ret = mbedtls_des_setkey_enc( &ctx, des3_test_keys );
             break;
 
         case 2:
-            mbedtls_des3_set2key_dec( &ctx3, des3_test_keys );
+            ret = mbedtls_des3_set2key_dec( &ctx3, des3_test_keys );
             break;
 
         case 3:
-            mbedtls_des3_set2key_enc( &ctx3, des3_test_keys );
+            ret = mbedtls_des3_set2key_enc( &ctx3, des3_test_keys );
             break;
 
         case 4:
-            mbedtls_des3_set3key_dec( &ctx3, des3_test_keys );
+            ret = mbedtls_des3_set3key_dec( &ctx3, des3_test_keys );
             break;
 
         case 5:
-            mbedtls_des3_set3key_enc( &ctx3, des3_test_keys );
+            ret = mbedtls_des3_set3key_enc( &ctx3, des3_test_keys );
             break;
 
         default:
             return( 1 );
         }
+        if( ret != 0 )
+            goto exit;
 
         for( j = 0; j < 100; j++ )
         {
             if( u == 0 )
-                mbedtls_des_crypt_ecb( &ctx, buf, buf );
+                ret = mbedtls_des_crypt_ecb( &ctx, buf, buf );
             else
-                mbedtls_des3_crypt_ecb( &ctx3, buf, buf );
+                ret = mbedtls_des3_crypt_ecb( &ctx3, buf, buf );
+            if( ret != 0 )
+                goto exit;
         }
 
         if( ( v == MBEDTLS_DES_DECRYPT &&
@@ -970,41 +989,45 @@
         switch( i )
         {
         case 0:
-            mbedtls_des_setkey_dec( &ctx, des3_test_keys );
+            ret = mbedtls_des_setkey_dec( &ctx, des3_test_keys );
             break;
 
         case 1:
-            mbedtls_des_setkey_enc( &ctx, des3_test_keys );
+            ret = mbedtls_des_setkey_enc( &ctx, des3_test_keys );
             break;
 
         case 2:
-            mbedtls_des3_set2key_dec( &ctx3, des3_test_keys );
+            ret = mbedtls_des3_set2key_dec( &ctx3, des3_test_keys );
             break;
 
         case 3:
-            mbedtls_des3_set2key_enc( &ctx3, des3_test_keys );
+            ret = mbedtls_des3_set2key_enc( &ctx3, des3_test_keys );
             break;
 
         case 4:
-            mbedtls_des3_set3key_dec( &ctx3, des3_test_keys );
+            ret = mbedtls_des3_set3key_dec( &ctx3, des3_test_keys );
             break;
 
         case 5:
-            mbedtls_des3_set3key_enc( &ctx3, des3_test_keys );
+            ret = mbedtls_des3_set3key_enc( &ctx3, des3_test_keys );
             break;
 
         default:
             return( 1 );
         }
+        if( ret != 0 )
+            goto exit;
 
         if( v == MBEDTLS_DES_DECRYPT )
         {
             for( j = 0; j < 100; j++ )
             {
                 if( u == 0 )
-                    mbedtls_des_crypt_cbc( &ctx, v, 8, iv, buf, buf );
+                    ret = mbedtls_des_crypt_cbc( &ctx, v, 8, iv, buf, buf );
                 else
-                    mbedtls_des3_crypt_cbc( &ctx3, v, 8, iv, buf, buf );
+                    ret = mbedtls_des3_crypt_cbc( &ctx3, v, 8, iv, buf, buf );
+                if( ret != 0 )
+                    goto exit;
             }
         }
         else
@@ -1014,9 +1037,11 @@
                 unsigned char tmp[8];
 
                 if( u == 0 )
-                    mbedtls_des_crypt_cbc( &ctx, v, 8, iv, buf, buf );
+                    ret = mbedtls_des_crypt_cbc( &ctx, v, 8, iv, buf, buf );
                 else
-                    mbedtls_des3_crypt_cbc( &ctx3, v, 8, iv, buf, buf );
+                    ret = mbedtls_des3_crypt_cbc( &ctx3, v, 8, iv, buf, buf );
+                if( ret != 0 )
+                    goto exit;
 
                 memcpy( tmp, prv, 8 );
                 memcpy( prv, buf, 8 );
@@ -1050,6 +1075,8 @@
     mbedtls_des_free( &ctx );
     mbedtls_des3_free( &ctx3 );
 
+    if( ret != 0 )
+        ret = 1;
     return( ret );
 }
 
diff --git a/programs/pkey/dh_client.c b/programs/pkey/dh_client.c
index d6e4990..f6c6226 100644
--- a/programs/pkey/dh_client.c
+++ b/programs/pkey/dh_client.c
@@ -274,7 +274,9 @@
     mbedtls_printf( "...\n  . Receiving and decrypting the ciphertext" );
     fflush( stdout );
 
-    mbedtls_aes_setkey_dec( &aes, buf, 256 );
+    ret = mbedtls_aes_setkey_dec( &aes, buf, 256 );
+    if( ret != 0 )
+        goto exit;
 
     memset( buf, 0, sizeof( buf ) );
 
@@ -284,7 +286,9 @@
         goto exit;
     }
 
-    mbedtls_aes_crypt_ecb( &aes, MBEDTLS_AES_DECRYPT, buf, buf );
+    ret = mbedtls_aes_crypt_ecb( &aes, MBEDTLS_AES_DECRYPT, buf, buf );
+    if( ret != 0 )
+        goto exit;
     buf[16] = '\0';
     mbedtls_printf( "\n  . Plaintext is \"%s\"\n\n", (char *) buf );
 
diff --git a/programs/pkey/dh_server.c b/programs/pkey/dh_server.c
index dccf095..8f032a3 100644
--- a/programs/pkey/dh_server.c
+++ b/programs/pkey/dh_server.c
@@ -295,9 +295,13 @@
     mbedtls_printf( "...\n  . Encrypting and sending the ciphertext" );
     fflush( stdout );
 
-    mbedtls_aes_setkey_enc( &aes, buf, 256 );
+    ret = mbedtls_aes_setkey_enc( &aes, buf, 256 );
+    if( ret != 0 )
+        goto exit;
     memcpy( buf, PLAINTEXT, 16 );
-    mbedtls_aes_crypt_ecb( &aes, MBEDTLS_AES_ENCRYPT, buf, buf );
+    ret = mbedtls_aes_crypt_ecb( &aes, MBEDTLS_AES_ENCRYPT, buf, buf );
+    if( ret != 0 )
+        goto exit;
 
     if( ( ret = mbedtls_net_send( &client_fd, buf, 16 ) ) != 16 )
     {
diff --git a/programs/test/benchmark.c b/programs/test/benchmark.c
index 9c5911b..88c9e65 100644
--- a/programs/test/benchmark.c
+++ b/programs/test/benchmark.c
@@ -449,7 +449,8 @@
     {
         mbedtls_des3_context des3;
         mbedtls_des3_init( &des3 );
-        mbedtls_des3_set3key_enc( &des3, tmp );
+        if( mbedtls_des3_set3key_enc( &des3, tmp ) != 0 )
+            mbedtls_exit( 1 );
         TIME_AND_TSC( "3DES",
                 mbedtls_des3_crypt_cbc( &des3, MBEDTLS_DES_ENCRYPT, BUFSIZE, tmp, buf, buf ) );
         mbedtls_des3_free( &des3 );
@@ -459,7 +460,8 @@
     {
         mbedtls_des_context des;
         mbedtls_des_init( &des );
-        mbedtls_des_setkey_enc( &des, tmp );
+        if( mbedtls_des_setkey_enc( &des, tmp ) != 0 )
+            mbedtls_exit( 1 );
         TIME_AND_TSC( "DES",
                 mbedtls_des_crypt_cbc( &des, MBEDTLS_DES_ENCRYPT, BUFSIZE, tmp, buf, buf ) );
         mbedtls_des_free( &des );
@@ -497,7 +499,7 @@
 
             memset( buf, 0, sizeof( buf ) );
             memset( tmp, 0, sizeof( tmp ) );
-            mbedtls_aes_setkey_enc( &aes, tmp, keysize );
+            CHECK_AND_CONTINUE( mbedtls_aes_setkey_enc( &aes, tmp, keysize ) );
 
             TIME_AND_TSC( title,
                 mbedtls_aes_crypt_cbc( &aes, MBEDTLS_AES_ENCRYPT, BUFSIZE, tmp, buf, buf ) );
@@ -518,7 +520,7 @@
 
             memset( buf, 0, sizeof( buf ) );
             memset( tmp, 0, sizeof( tmp ) );
-            mbedtls_aes_xts_setkey_enc( &ctx, tmp, keysize * 2 );
+            CHECK_AND_CONTINUE( mbedtls_aes_xts_setkey_enc( &ctx, tmp, keysize * 2 ) );
 
             TIME_AND_TSC( title,
                     mbedtls_aes_crypt_xts( &ctx, MBEDTLS_AES_ENCRYPT, BUFSIZE,
diff --git a/tests/suites/test_suite_aes.function b/tests/suites/test_suite_aes.function
index a90277d..5a64099 100644
--- a/tests/suites/test_suite_aes.function
+++ b/tests/suites/test_suite_aes.function
@@ -67,7 +67,7 @@
 
     mbedtls_aes_init( &ctx );
 
-    mbedtls_aes_setkey_enc( &ctx, key_str->x, key_str->len * 8 );
+    TEST_ASSERT( mbedtls_aes_setkey_enc( &ctx, key_str->x, key_str->len * 8 ) == 0 );
     TEST_ASSERT( mbedtls_aes_crypt_cbc( &ctx, MBEDTLS_AES_ENCRYPT, src_str->len, iv_str->x, src_str->x, output ) == cbc_result );
     if( cbc_result == 0 )
     {
@@ -92,7 +92,7 @@
     memset(output, 0x00, 100);
     mbedtls_aes_init( &ctx );
 
-    mbedtls_aes_setkey_dec( &ctx, key_str->x, key_str->len * 8 );
+    TEST_ASSERT( mbedtls_aes_setkey_dec( &ctx, key_str->x, key_str->len * 8 ) == 0 );
     TEST_ASSERT( mbedtls_aes_crypt_cbc( &ctx, MBEDTLS_AES_DECRYPT, src_str->len, iv_str->x, src_str->x, output ) == cbc_result );
     if( cbc_result == 0)
     {
@@ -244,7 +244,7 @@
     mbedtls_aes_init( &ctx );
 
 
-    mbedtls_aes_setkey_enc( &ctx, key_str->x, key_str->len * 8 );
+    TEST_ASSERT( mbedtls_aes_setkey_enc( &ctx, key_str->x, key_str->len * 8 ) == 0 );
     TEST_ASSERT( mbedtls_aes_crypt_cfb128( &ctx, MBEDTLS_AES_ENCRYPT, 16, &iv_offset, iv_str->x, src_str->x, output ) == 0 );
 
     TEST_ASSERT( mbedtls_test_hexcmp( output, dst->x, 16, dst->len ) == 0 );
@@ -266,7 +266,7 @@
     mbedtls_aes_init( &ctx );
 
 
-    mbedtls_aes_setkey_enc( &ctx, key_str->x, key_str->len * 8 );
+    TEST_ASSERT( mbedtls_aes_setkey_enc( &ctx, key_str->x, key_str->len * 8 ) == 0 );
     TEST_ASSERT( mbedtls_aes_crypt_cfb128( &ctx, MBEDTLS_AES_DECRYPT, 16, &iv_offset, iv_str->x, src_str->x, output ) == 0 );
 
     TEST_ASSERT( mbedtls_test_hexcmp( output, dst->x, 16, dst->len ) == 0 );
@@ -287,7 +287,7 @@
     mbedtls_aes_init( &ctx );
 
 
-    mbedtls_aes_setkey_enc( &ctx, key_str->x, key_str->len * 8 );
+    TEST_ASSERT( mbedtls_aes_setkey_enc( &ctx, key_str->x, key_str->len * 8 ) == 0 );
     TEST_ASSERT( mbedtls_aes_crypt_cfb8( &ctx, MBEDTLS_AES_ENCRYPT, src_str->len, iv_str->x, src_str->x, output ) == 0 );
 
     TEST_ASSERT( mbedtls_test_hexcmp( output, dst->x,
@@ -309,7 +309,7 @@
     mbedtls_aes_init( &ctx );
 
 
-    mbedtls_aes_setkey_enc( &ctx, key_str->x, key_str->len * 8 );
+    TEST_ASSERT( mbedtls_aes_setkey_enc( &ctx, key_str->x, key_str->len * 8 ) == 0 );
     TEST_ASSERT( mbedtls_aes_crypt_cfb8( &ctx, MBEDTLS_AES_DECRYPT, src_str->len, iv_str->x, src_str->x, output ) == 0 );
 
     TEST_ASSERT( mbedtls_test_hexcmp( output, dst->x,
diff --git a/tests/suites/test_suite_des.function b/tests/suites/test_suite_des.function
index 5b24935..7256fb5 100644
--- a/tests/suites/test_suite_des.function
+++ b/tests/suites/test_suite_des.function
@@ -24,7 +24,7 @@
     mbedtls_des_init( &ctx );
 
 
-    mbedtls_des_setkey_enc( &ctx, key_str->x );
+    TEST_ASSERT( mbedtls_des_setkey_enc( &ctx, key_str->x ) == 0 );
     TEST_ASSERT( mbedtls_des_crypt_ecb( &ctx, src_str->x, output ) == 0 );
 
     TEST_ASSERT( mbedtls_test_hexcmp( output, dst->x, 8, dst->len ) == 0 );
@@ -44,7 +44,7 @@
     mbedtls_des_init( &ctx );
 
 
-    mbedtls_des_setkey_dec( &ctx, key_str->x );
+    TEST_ASSERT( mbedtls_des_setkey_dec( &ctx, key_str->x ) == 0 );
     TEST_ASSERT( mbedtls_des_crypt_ecb( &ctx, src_str->x, output ) == 0 );
 
     TEST_ASSERT( mbedtls_test_hexcmp( output, dst->x, 8, dst->len ) == 0 );
@@ -65,7 +65,7 @@
     mbedtls_des_init( &ctx );
 
 
-    mbedtls_des_setkey_enc( &ctx, key_str->x );
+    TEST_ASSERT( mbedtls_des_setkey_enc( &ctx, key_str->x ) == 0 );
     TEST_ASSERT( mbedtls_des_crypt_cbc( &ctx, MBEDTLS_DES_ENCRYPT, src_str->len, iv_str->x, src_str->x, output ) == cbc_result );
     if( cbc_result == 0 )
     {
@@ -91,7 +91,7 @@
     mbedtls_des_init( &ctx );
 
 
-    mbedtls_des_setkey_dec( &ctx, key_str->x );
+    TEST_ASSERT( mbedtls_des_setkey_dec( &ctx, key_str->x ) == 0 );
     TEST_ASSERT( mbedtls_des_crypt_cbc( &ctx, MBEDTLS_DES_DECRYPT, src_str->len, iv_str->x, src_str->x, output ) == cbc_result );
     if( cbc_result == 0 )
     {
@@ -117,9 +117,9 @@
 
 
     if( key_count == 2 )
-        mbedtls_des3_set2key_enc( &ctx, key_str->x );
+        TEST_ASSERT( mbedtls_des3_set2key_enc( &ctx, key_str->x ) == 0 );
     else if( key_count == 3 )
-        mbedtls_des3_set3key_enc( &ctx, key_str->x );
+        TEST_ASSERT( mbedtls_des3_set3key_enc( &ctx, key_str->x ) == 0 );
     else
         TEST_ASSERT( 0 );
 
@@ -144,9 +144,9 @@
 
 
     if( key_count == 2 )
-        mbedtls_des3_set2key_dec( &ctx, key_str->x );
+        TEST_ASSERT( mbedtls_des3_set2key_dec( &ctx, key_str->x ) == 0 );
     else if( key_count == 3 )
-        mbedtls_des3_set3key_dec( &ctx, key_str->x );
+        TEST_ASSERT( mbedtls_des3_set3key_dec( &ctx, key_str->x ) == 0 );
     else
         TEST_ASSERT( 0 );
 
@@ -172,9 +172,9 @@
 
 
     if( key_count == 2 )
-        mbedtls_des3_set2key_enc( &ctx, key_str->x );
+        TEST_ASSERT( mbedtls_des3_set2key_enc( &ctx, key_str->x ) == 0 );
     else if( key_count == 3 )
-        mbedtls_des3_set3key_enc( &ctx, key_str->x );
+        TEST_ASSERT( mbedtls_des3_set3key_enc( &ctx, key_str->x ) == 0 );
     else
         TEST_ASSERT( 0 );
 
@@ -205,9 +205,9 @@
 
 
     if( key_count == 2 )
-        mbedtls_des3_set2key_dec( &ctx, key_str->x );
+        TEST_ASSERT( mbedtls_des3_set2key_dec( &ctx, key_str->x ) == 0 );
     else if( key_count == 3 )
-        mbedtls_des3_set3key_dec( &ctx, key_str->x );
+        TEST_ASSERT( mbedtls_des3_set3key_dec( &ctx, key_str->x ) == 0 );
     else
         TEST_ASSERT( 0 );