diff --git a/library/aes.c b/library/aes.c
index bf5d432..af13d63 100644
--- a/library/aes.c
+++ b/library/aes.c
@@ -555,14 +555,15 @@
         aes_padlock_ace = mbedtls_padlock_has_support( MBEDTLS_PADLOCK_ACE );
 
     if( aes_padlock_ace )
-        ctx->rk = RK = MBEDTLS_PADLOCK_ALIGN16( ctx->buf );
+        ctx->rk_offset = MBEDTLS_PADLOCK_ALIGN16( ctx->buf ) - ctx->buf;
     else
 #endif
-    ctx->rk = RK = ctx->buf;
+    ctx->rk_offset = 0;
+    RK = ctx->buf + ctx->rk_offset;
 
 #if defined(MBEDTLS_AESNI_C) && defined(MBEDTLS_HAVE_X86_64)
     if( mbedtls_aesni_has_support( MBEDTLS_AESNI_AES ) )
-        return( mbedtls_aesni_setkey_enc( (unsigned char *) ctx->rk, key, keybits ) );
+        return( mbedtls_aesni_setkey_enc( (unsigned char *) RK, key, keybits ) );
 #endif
 
     for( i = 0; i < ( keybits >> 5 ); i++ )
@@ -659,10 +660,11 @@
         aes_padlock_ace = mbedtls_padlock_has_support( MBEDTLS_PADLOCK_ACE );
 
     if( aes_padlock_ace )
-        ctx->rk = RK = MBEDTLS_PADLOCK_ALIGN16( ctx->buf );
+        ctx->rk_offset = MBEDTLS_PADLOCK_ALIGN16( ctx->buf ) - ctx->buf;
     else
 #endif
-    ctx->rk = RK = ctx->buf;
+    ctx->rk_offset = 0;
+    RK = ctx->buf + ctx->rk_offset;
 
     /* Also checks keybits */
     if( ( ret = mbedtls_aes_setkey_enc( &cty, key, keybits ) ) != 0 )
@@ -673,13 +675,13 @@
 #if defined(MBEDTLS_AESNI_C) && defined(MBEDTLS_HAVE_X86_64)
     if( mbedtls_aesni_has_support( MBEDTLS_AESNI_AES ) )
     {
-        mbedtls_aesni_inverse_key( (unsigned char *) ctx->rk,
-                           (const unsigned char *) cty.rk, ctx->nr );
+        mbedtls_aesni_inverse_key( (unsigned char *) RK,
+                           (const unsigned char *) ( cty.buf + cty.rk_offset ), ctx->nr );
         goto exit;
     }
 #endif
 
-    SK = cty.rk + cty.nr * 4;
+    SK = cty.buf + cty.rk_offset + cty.nr * 4;
 
     *RK++ = *SK++;
     *RK++ = *SK++;
@@ -843,7 +845,7 @@
                                   unsigned char output[16] )
 {
     int i;
-    uint32_t *RK = ctx->rk;
+    uint32_t *RK = ctx->buf + ctx->rk_offset;
     struct
     {
         uint32_t X[4];
@@ -907,7 +909,7 @@
                                   unsigned char output[16] )
 {
     int i;
-    uint32_t *RK = ctx->rk;
+    uint32_t *RK = ctx->buf + ctx->rk_offset;
     struct
     {
         uint32_t X[4];
@@ -971,7 +973,6 @@
                            unsigned char output[16] )
 {
     AES_VALIDATE_RET( ctx != NULL );
-    AES_VALIDATE_RET( ctx->rk != NULL );
     AES_VALIDATE_RET( input != NULL );
     AES_VALIDATE_RET( output != NULL );
     AES_VALIDATE_RET( mode == MBEDTLS_AES_ENCRYPT ||