diff --git a/programs/pkey/dh_client.c b/programs/pkey/dh_client.c
index eccb42a..d633e4d 100644
--- a/programs/pkey/dh_client.c
+++ b/programs/pkey/dh_client.c
@@ -270,7 +270,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 ) );
 
@@ -280,7 +282,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 0ddb85c..75713ff 100644
--- a/programs/pkey/dh_server.c
+++ b/programs/pkey/dh_server.c
@@ -290,9 +290,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 )
     {