diff --git a/library/aes.c b/library/aes.c
index 3de571e..7a364a0 100644
--- a/library/aes.c
+++ b/library/aes.c
@@ -511,6 +511,8 @@
 
 void mbedtls_aes_init( mbedtls_aes_context *ctx )
 {
+    MBEDTLS_AES_VALIDATE( ctx != NULL );
+
     memset( ctx, 0, sizeof( mbedtls_aes_context ) );
 }
 
@@ -525,12 +527,16 @@
 #if defined(MBEDTLS_CIPHER_MODE_XTS)
 void mbedtls_aes_xts_init( mbedtls_aes_xts_context *ctx )
 {
+    MBEDTLS_AES_VALIDATE( ctx != NULL );
+
     mbedtls_aes_init( &ctx->crypt );
     mbedtls_aes_init( &ctx->tweak );
 }
 
 void mbedtls_aes_xts_free( mbedtls_aes_xts_context *ctx )
 {
+    MBEDTLS_AES_VALIDATE( ctx != NULL );
+
     mbedtls_aes_free( &ctx->crypt );
     mbedtls_aes_free( &ctx->tweak );
 }
@@ -546,14 +552,7 @@
     unsigned int i;
     uint32_t *RK;
 
-#if !defined(MBEDTLS_AES_ROM_TABLES)
-    if( aes_init_done == 0 )
-    {
-        aes_gen_tables();
-        aes_init_done = 1;
-
-    }
-#endif
+    MBEDTLS_AES_VALIDATE_RET( ctx != NULL && key != NULL );
 
     switch( keybits )
     {
@@ -563,6 +562,15 @@
         default : return( MBEDTLS_ERR_AES_INVALID_KEY_LENGTH );
     }
 
+#if !defined(MBEDTLS_AES_ROM_TABLES)
+    if( aes_init_done == 0 )
+    {
+        aes_gen_tables();
+        aes_init_done = 1;
+
+    }
+#endif
+
 #if defined(MBEDTLS_PADLOCK_C) && defined(MBEDTLS_PADLOCK_ALIGN16)
     if( aes_padlock_ace == -1 )
         aes_padlock_ace = mbedtls_padlock_has_support( MBEDTLS_PADLOCK_ACE );
@@ -662,6 +670,8 @@
     uint32_t *RK;
     uint32_t *SK;
 
+    MBEDTLS_AES_VALIDATE_RET( ctx != NULL && key != NULL );
+
     mbedtls_aes_init( &cty );
 
 #if defined(MBEDTLS_PADLOCK_C) && defined(MBEDTLS_PADLOCK_ALIGN16)