I am trying to understand $AES$ and I wondered what happens in the case of a $128$ bit block size and a $256$ bit key. All documents I read about $AES$ told me the key expansions algorithm for $128$ bit keys and $128$ bit blocks. But I have a $256$ bit key ($8$ words) and $128$ bit block ($4$ words) and I am not sure how the key expansion works now. Should I expand it how usual but instead beginning in the second round beginning in the third round? Like this graphic shows:
Well, as SEJPM and Mok-Kong Shen pointed out, it's described in FIPS 197, but let's explain it in details:
Here's the Key Expansion algorithm pseudo-code from FIPS 197:
- $Nk$ is the number of $32$-bit words that composed the key
- $Nr$ is the number of rounds
- $Nb$ is the number of columns in the state block (which is always $4$).
- $word$ is a $32$-bit word (can be seen as an array of $4$ bytes).
- $w$ denotes a $4$-byte word array where the number of elements is equal to $Nr + 1$ (because AES uses a different key at each round)
For a $256$-bit key, we have $Nk = 8$ and $Nr = 14$ according to FIPS 197.
When you are asking
Should I expand it how usual but instead beginning in the second round beginning in the third round?
Don't you mean the fourth round? Let me explain.
As you can see with the above pseudo-code, the first loop just copy the secret key at the beginning of the array $w$. Then, we get into the second loop to compute rounds key which are derived from the given one.
So for a $128$-bit key, you start to derive round keys as soon as $i=4$ but only when $i=8$ for a $256$-bit key (because the key is longer).
Also, as mentionned in FIPS 197, note that the routine isn't the same regarding to the key length (which is $Nk$) because there are arithmetic conditions between $i$ and $Nk$ to get into different branches of the algorithm. For example, only $256$-bit keys will satisfy the condition else if (Nk >6 and i mod Nk = 4).
Hope it's more clear now.