diff --git a/library/aesce.c b/library/aesce.c
index 356d0a3..e47665a 100644
--- a/library/aesce.c
+++ b/library/aesce.c
@@ -185,15 +185,17 @@
                              const unsigned char *key,
                              const size_t key_bit_length)
 {
-
-    const uint32_t key_len_in_words = key_bit_length / 32;
-    const size_t round_key_len_in_words = 4;
     static uint8_t const rcon[] = { 0x01, 0x02, 0x04, 0x08, 0x10,
                                     0x20, 0x40, 0x80, 0x1b, 0x36 };
-    /* Require max(key_len_in_words, round_key_len_len_in_words) + 7 */
-    const size_t round_keys_needed = key_len_in_words + 7;
+    /* See https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.197.pdf
+     *   - Section 5, Nr = Nk + 6
+     *   - Section 5.2, the key expansion size is Nb*(Nr+1)
+     */
+    const uint32_t key_len_in_words = key_bit_length / 32;  /* Nk */
+    const size_t round_key_len_in_words = 4;                /* Nb */
+    const size_t round_keys_needed = key_len_in_words + 6;  /* Nr */
     const size_t key_expansion_size_in_words =
-        round_keys_needed * round_key_len_in_words;
+        round_key_len_in_words * (round_keys_needed + 1);   /* Nb*(Nr+1) */
     const uint32_t *rko_end = (uint32_t *) rk + key_expansion_size_in_words;
 
     memcpy(rk, key, key_len_in_words * 4);