5
\$\begingroup\$

I am a hobbyist and do not work in the industry, but I do have an interest in programming an cryptography. I decided to have a go at writing a 128bit implementation of AES in C++. This is my first ever C++ program and I learn as I went. So I'm sure there are a lot of bad habits or things that are simply wrong.

But it does give the correct results as described by FIPS 197 :) Appreciate any feedback that I can get.

#include <iostream>
#include <array>

using namespace std;

void printBlock(const array<array<int,4>,4>& block){
    //Function to print a block to terminal
    for (size_t i = 0; i < 4; ++i){
        printf("%.2x, %.2x, %.2x, %.2x\n", block[i][0], block[i][1], block[i][2], block[i][3]);
    }
} 
void printRow(const array<int,4>& row){
    //Function to print a row to terminal 
    printf("%.2x, %.2x, %.2x, %.2x\n", row[0], row[1], row[2], row[3]);
}
void print_1d_State(const array<int,16>& e){
    //Function to print a 1d version of the state to terminal 
    printf("%.2x, %.2x, %.2x, %.2x, %.2x, %.2x, %.2x, %.2x, %.2x, %.2x, %.2x, %.2x, %.2x, %.2x, %.2x, %.2x\n", 
    e[0], e[1], e[2], e[3], e[4], e[5], e[6], e[7], e[8], e[9], e[10], e[11], e[12], e[13], e[14], e[15]);
}
void printKeySchedule(const array<array<int,16>,11>& keySchedual){
    //Function to print keys for debugging
    for (size_t i=0; i<11; ++i){
        print_1d_State(keySchedual[i]);
    }
}

array<array<int,4>,4> inputToState(const array<int,16>& input){
    //Function turns a 128bit/ 16 element array and minipulates into a 4x4 matrix
    array<array<int,4>,4> state = {};

    //Loop the input and assign to state
    for(size_t i=0; i<4; ++i){
        for(size_t j=0; j<4; ++j){
            state[i][j] = input[i+4*j];
        }
    }
    
    return state;
}
array<int,16> stateToOutput(const array<array<int,4>,4>& state){
    //Function takes a 128bit 4x4 block/ matrix and minipulates into a 1d array of words  
    array<int,16> output = {};

    //Loop the state and assign to output
    for(size_t i=0; i<4; ++i){
        for(size_t j=0; j<4; ++j){
            output[i+4*j] = state[i][j];
        }
    }

    return output;
}
array<array<int,4>,4> keyToWords(const array<int,16>& roundKey){
    //Function will take a rounds key and minipulate it into a 4x4 array
    //With each row being a word
    array<array<int,4>,4> wordBlock = {};

    for (size_t i= 0; i < 4; ++i){
        //Get words from key 
        wordBlock[i] = {roundKey[4*i], roundKey[1+i*4], roundKey[2+i*4], roundKey[3+i*4]};
    }
    return wordBlock;
}
array<array<int,4>,4> transposeState(const array<array<int,4>,4>& state){
    //Function will transpose the state into a 4x4 matrix of words 
    array<array<int,4>,4> tState = {};

    for(size_t i = 0; i < 4; ++i){
        for(size_t j = 0; j < 4; ++j){
            tState[i][j] = state[j][i];
        }
    }

    return tState;
}

int multiply_by_2(const int v){
/* 
Function to impliment multiplication by 2 of Galois Field GF(2^8).
Shift bits left by 1, if high bit is 0 return the value.
If the high bit is one XOR the value with 0x1B.  
(0x1B comes from the field representation)
*/

    int s = v << 1; //Shift bits left by 1

    //If high bit = 1 (0x80 = 10000000bin)
    if (v & 0x80){
        s &= 0xff; //Bitwise AND (ff = 11111111bin) This effectivly takes first 8 bits
        s = s ^ 0x1b;
    }

    return s;
}
int multiply_by_3(int v){
/* 
Function to impliment multiplication by 3 of Galois Field GF(2^8).
This is simply the XOR of multiply by 2, with the original value.
*/
    return multiply_by_2(v) ^ v;
}

int lookupByte(int &byte){
    //This method takes a byte, performs a lookup in the AES SBOX and returns the corresponding value

    int x = byte >> 4; //Shifts 4 bits right i.e. takes first 4 bits and discards the rest 
    int y = byte & 0x0f; // 0x0f = 15 = 00001111(bin). Effectivly takes last 4 bits and dicards the rest

    const int sbox[16][16] = {{0x63, 0x7c, 0x77, 0x7b, 0xf2, 0x6b, 0x6f, 0xc5,0x30, 0x01, 0x67, 0x2b, 0xfe, 0xd7, 0xab, 0x76},
                          {0xca, 0x82, 0xc9, 0x7d, 0xfa, 0x59, 0x47, 0xf0,0xad, 0xd4, 0xa2, 0xaf, 0x9c, 0xa4, 0x72, 0xc0},
                          {0xb7, 0xfd, 0x93, 0x26, 0x36, 0x3f, 0xf7, 0xcc,0x34, 0xa5, 0xe5, 0xf1, 0x71, 0xd8, 0x31, 0x15},
                          {0x04, 0xc7, 0x23, 0xc3, 0x18, 0x96, 0x05, 0x9a,0x07, 0x12, 0x80, 0xe2, 0xeb, 0x27, 0xb2, 0x75},
                          {0x09, 0x83, 0x2c, 0x1a, 0x1b, 0x6e, 0x5a, 0xa0,0x52, 0x3b, 0xd6, 0xb3, 0x29, 0xe3, 0x2f, 0x84},
                          {0x53, 0xd1, 0x00, 0xed, 0x20, 0xfc, 0xb1, 0x5b,0x6a, 0xcb, 0xbe, 0x39, 0x4a, 0x4c, 0x58, 0xcf},
                          {0xd0, 0xef, 0xaa, 0xfb, 0x43, 0x4d, 0x33, 0x85,0x45, 0xf9, 0x02, 0x7f, 0x50, 0x3c, 0x9f, 0xa8},
                          {0x51, 0xa3, 0x40, 0x8f, 0x92, 0x9d, 0x38, 0xf5,0xbc, 0xb6, 0xda, 0x21, 0x10, 0xff, 0xf3, 0xd2},
                          {0xcd, 0x0c, 0x13, 0xec, 0x5f, 0x97, 0x44, 0x17,0xc4, 0xa7, 0x7e, 0x3d, 0x64, 0x5d, 0x19, 0x73},
                          {0x60, 0x81, 0x4f, 0xdc, 0x22, 0x2a, 0x90, 0x88,0x46, 0xee, 0xb8, 0x14, 0xde, 0x5e, 0x0b, 0xdb},
                          {0xe0, 0x32, 0x3a, 0x0a, 0x49, 0x06, 0x24, 0x5c,0xc2, 0xd3, 0xac, 0x62, 0x91, 0x95, 0xe4, 0x79},
                          {0xe7, 0xc8, 0x37, 0x6d, 0x8d, 0xd5, 0x4e, 0xa9,0x6c, 0x56, 0xf4, 0xea, 0x65, 0x7a, 0xae, 0x08},
                          {0xba, 0x78, 0x25, 0x2e, 0x1c, 0xa6, 0xb4, 0xc6,0xe8, 0xdd, 0x74, 0x1f, 0x4b, 0xbd, 0x8b, 0x8a},
                          {0x70, 0x3e, 0xb5, 0x66, 0x48, 0x03, 0xf6, 0x0e,0x61, 0x35, 0x57, 0xb9, 0x86, 0xc1, 0x1d, 0x9e},
                          {0xe1, 0xf8, 0x98, 0x11, 0x69, 0xd9, 0x8e, 0x94,0x9b, 0x1e, 0x87, 0xe9, 0xce, 0x55, 0x28, 0xdf},
                          {0x8c, 0xa1, 0x89, 0x0d, 0xbf, 0xe6, 0x42, 0x68,0x41, 0x99, 0x2d, 0x0f, 0xb0, 0x54, 0xbb, 0x16}};

     return sbox[x][y];
    
}
int invLookupByte(int &byte){
    //This method takes a byte, performs a lookup in the inverse AES SBOX and returns the corresponding value

    int x = byte >> 4; //Shifts 4 bits right i.e. takes first 4 bits and discards the rest 
    int y = byte & 0x0f; // 0x0f = 15 = 00001111(bin). Effectivly takes last 4 bits and dicards the rest

    const int sboxinv[16][16] = {{0x52, 0x09, 0x6a, 0xd5, 0x30, 0x36, 0xa5, 0x38,0xbf, 0x40, 0xa3, 0x9e, 0x81, 0xf3, 0xd7, 0xfb},
                            {0x7c, 0xe3, 0x39, 0x82, 0x9b, 0x2f, 0xff, 0x87,0x34, 0x8e, 0x43, 0x44, 0xc4, 0xde, 0xe9, 0xcb},
                            {0x54, 0x7b, 0x94, 0x32, 0xa6, 0xc2, 0x23, 0x3d,0xee, 0x4c, 0x95, 0x0b, 0x42, 0xfa, 0xc3, 0x4e},
                            {0x08, 0x2e, 0xa1, 0x66, 0x28, 0xd9, 0x24, 0xb2,0x76, 0x5b, 0xa2, 0x49, 0x6d, 0x8b, 0xd1, 0x25},
                            {0x72, 0xf8, 0xf6, 0x64, 0x86, 0x68, 0x98, 0x16,0xd4, 0xa4, 0x5c, 0xcc, 0x5d, 0x65, 0xb6, 0x92},
                            {0x6c, 0x70, 0x48, 0x50, 0xfd, 0xed, 0xb9, 0xda,0x5e, 0x15, 0x46, 0x57, 0xa7, 0x8d, 0x9d, 0x84},
                            {0x90, 0xd8, 0xab, 0x00, 0x8c, 0xbc, 0xd3, 0x0a,0xf7, 0xe4, 0x58, 0x05, 0xb8, 0xb3, 0x45, 0x06},
                            {0xd0, 0x2c, 0x1e, 0x8f, 0xca, 0x3f, 0x0f, 0x02,0xc1, 0xaf, 0xbd, 0x03, 0x01, 0x13, 0x8a, 0x6b},
                            {0x3a, 0x91, 0x11, 0x41, 0x4f, 0x67, 0xdc, 0xea,0x97, 0xf2, 0xcf, 0xce, 0xf0, 0xb4, 0xe6, 0x73},
                            {0x96, 0xac, 0x74, 0x22, 0xe7, 0xad, 0x35, 0x85,0xe2, 0xf9, 0x37, 0xe8, 0x1c, 0x75, 0xdf, 0x6e},
                            {0x47, 0xf1, 0x1a, 0x71, 0x1d, 0x29, 0xc5, 0x89,0x6f, 0xb7, 0x62, 0x0e, 0xaa, 0x18, 0xbe, 0x1b},
                            {0xfc, 0x56, 0x3e, 0x4b, 0xc6, 0xd2, 0x79, 0x20,0x9a, 0xdb, 0xc0, 0xfe, 0x78, 0xcd, 0x5a, 0xf4},
                            {0x1f, 0xdd, 0xa8, 0x33, 0x88, 0x07, 0xc7, 0x31,0xb1, 0x12, 0x10, 0x59, 0x27, 0x80, 0xec, 0x5f},
                            {0x60, 0x51, 0x7f, 0xa9, 0x19, 0xb5, 0x4a, 0x0d,0x2d, 0xe5, 0x7a, 0x9f, 0x93, 0xc9, 0x9c, 0xef},
                            {0xa0, 0xe0, 0x3b, 0x4d, 0xae, 0x2a, 0xf5, 0xb0,0xc8, 0xeb, 0xbb, 0x3c, 0x83, 0x53, 0x99, 0x61},
                            {0x17, 0x2b, 0x04, 0x7e, 0xba, 0x77, 0xd6, 0x26,0xe1, 0x69, 0x14, 0x63, 0x55, 0x21, 0x0c, 0x7d}};

    return sboxinv[x][y];
}

auto subBytes(const array<array<int,4>,4>& state){
    //Function takes the current state and subs each byte using sbox look up 
    //the sboxLookup variable indicates the use of sbox or sboxInv i.e. forward or reverse lookup. 
    array<array<int,4>,4> result = {};
    int byte; 

    for (size_t i=0; i<4; ++i){
        for (size_t j=0; j<4; ++j){
            byte = state[i][j];
            result[i][j] = lookupByte(byte);
        }
    }

    return result;
}
auto unSubBytes(const array<array<int,4>,4>& state){
    //Function takes the current states and subs each byte using sbox look up 
    //the sboxLookup variable indicates the use of sbox or sboxInv i.e. forward or reverse lookup. 
    array<array<int,4>,4> result = {};
    int byte; 

    for (size_t i=0; i<4; ++i){
        for (size_t j=0; j<4; ++j){
            byte = state[i][j];
            result[i][j] = invLookupByte(byte);
        }
    }

    return result;
}

array<int,4> shiftRow(const array<int,4>& row, const int shift){
    //Recursive function to shft a row left by given value 
    array<int,4> result = {};
    result = row;

    if(shift){
        //Shift by 1
        int temp = result[0]; 

        for (size_t i=0; i<3; ++i){
            result[i] = result[i+1];
        }
        result[3] = temp; 
        //reduce shift and perform again
        result = shiftRow(result, shift -1);
    }
    else{ 
        return result;
    }
}
auto shiftRows(const array<array<int,4>,4>& state){
    //Function to shift rows of the state left by the value of their index
    array<array<int,4>,4> result = {};

    for (size_t i = 0; i < 4; i++){
        result[i] = shiftRow(state[i], i);
    }

    return result;
}
array<int,4> unShiftRow(const array<int,4>& row, const int shift){
    //Recursive function to shift a row right by given value 
    array<int,4> result = {};
    result = row;

    if(shift){
        //Shift by 1
        int temp = result[3]; 

        for (int i=3; i>0; --i){
            result[i] = result[i-1];
        }
        result[0] = temp; 
        //reduce shift and perform again
        result = unShiftRow(result, shift -1);
    }
    else{ 
        return result;
    }
}
auto unShiftRows(const array<array<int,4>,4>& state){
    //Function to shift rows of the state right by the value of their index
    array<array<int,4>,4> result = {};

    for (size_t i = 0; i < 4; i++){
        result[i] = unShiftRow(state[i], i);
    }

    return result;
}

auto mixColumn(const array<int,4>& stateColumn){
    //Function takes a column of the state and 'mixes' it according to the matrix multiplication defined in FIPS 197
    array<int,4> result = {};
    
    result[0] = multiply_by_2(stateColumn[0]) ^ multiply_by_3(stateColumn[1]) ^ stateColumn[2] ^ stateColumn[3];
    result[1] = multiply_by_2(stateColumn[1]) ^ multiply_by_3(stateColumn[2]) ^ stateColumn[3] ^ stateColumn[0];
    result[2] = multiply_by_2(stateColumn[2]) ^ multiply_by_3(stateColumn[3]) ^ stateColumn[0] ^ stateColumn[1];
    result[3] = multiply_by_2(stateColumn[3]) ^ multiply_by_3(stateColumn[0]) ^ stateColumn[1] ^ stateColumn[2];

    return result;
}
auto mixColumns(const array<array<int,4>,4>& state){
    //Function takes the current states and applies the 'mixColumn' function column by column
    array<array<int,4>,4> result = {};
    array<int,4> column = {};
    
    //Grab the columns and mix
    for(size_t i = 0; i < 4; ++i){
        for(size_t j = 0; j < 4; ++j){
            column[j] = state[j][i];
        }

        column = mixColumn(column);   

        //Transpose the columns back into rows
        for(size_t j = 0; j < 4; ++j){
            result[j][i] = column[j];
        }
    }

    return result;
}
auto unMixColumns(const array<array<int,4>,4>& state){
    //Function performs the inverse of the mixColumns function by using the 'cheat' method of 
    //applying mixColumns 3 times
    array<array<int,4>,4> result = {};

    result = mixColumns(state);
    result = mixColumns(result);
    result = mixColumns(result);

    return result;
}

array<int,4> xorWords(const array<int,4>& wordA, const array<int,4>& wordB){
    //Function takes 2 words, XORs each element and returns the result 
    array<int,4> result = {};
    
    for (size_t i = 0; i < 4; i++){
        result[i] = wordA[i] ^ wordB[i];
    }
    return result;
}
array<int,4> subWord(const array<int,4>& word){
    //Function takes the a word, performs an sbox lookup on each element and returns the result 
    array<int,4> result = {};
    int byte; 

    for (size_t i=0; i<4; ++i){
        byte = word[i];
        result[i] = lookupByte(byte);    
    }

    return result;
}
array<int,4> rotWord(const array<int,4>& word){
    //Function takes the a word, performs a rightward cyclic shift and returns the result
    array<int,4> result = {};
    
    result = shiftRow(word, 1);

    return result;
}

auto addRoundKey(const array<array<int,4>,4>& state, const array<int,16>& roundKey){
    //Function takes 2 4x4 matricies, xors the corresponding columns and returns the result
    array<array<int,4>,4> result = {};
    array<array<int,4>,4> stateWords = {};
    array<array<int,4>,4> keyWords = {};

    //Convert state and key into rows of words 
    stateWords = transposeState(state);
    keyWords = keyToWords(roundKey);

    //Loop columns of the state and xor with the corresponding key column 
    for (size_t i = 0; i < 4; ++i){
        result[i] = xorWords(stateWords[i], keyWords[i]);
    }


    //Transpose words back into columns when returning 
    return transposeState(result);
}

array<array<int,16>,11> generateKeys(const array<int,16>& key){
    //Function performs the key expansion algorithm to expand a single 128 bit key into 
    //10 round keys know as the keySchedule. Note this is coded specifically for 128bit keys and
    //will NOT work for other AES varients 
    array<array<int,16>,11> keySchedule = {};
    array<array<int,4>,4> roundKey = {};
    array<array<int,4>,4> pRoundKey = {};
    array<int,16> temp = {};

    const array<array<int,4>,10> RCON = {{{0x01, 0x00, 0x00, 0x00},
                                        {0x02, 0x00, 0x00, 0x00},
                                        {0x04, 0x00, 0x00, 0x00},
                                        {0x08, 0x00, 0x00, 0x00},
                                        {0x10, 0x00, 0x00, 0x00},
                                        {0x20, 0x00, 0x00, 0x00},
                                        {0x40, 0x00, 0x00, 0x00},
                                        {0x80, 0x00, 0x00, 0x00},
                                        {0x1b, 0x00, 0x00, 0x00},
                                        {0x36, 0x00, 0x00, 0x00}}};

    keySchedule[0] = key; 

    //Extract the words from the key and arrange into a block
    for (size_t i=0; i<4; ++i){
        pRoundKey[i] = {key[4*i], key[4*i+1], key[4*i+2], key[4*i+3]};
    }

    //Loop rounds
    for (size_t i=0; i<10; ++i){
        //Rotate last word of PRK
        roundKey[0] = rotWord(pRoundKey[3]);
        //SubBytes of last word in PRK
        roundKey[0] = subWord(roundKey[0]);
        //XOR : first word PRK, current RK state, rcon(round)
        roundKey[0] = xorWords(xorWords(pRoundKey[0], roundKey[0]), RCON[i]);
        //XOR the other words in sequential order
        for (size_t j=1; j<4; ++j){
            roundKey[j] = xorWords(roundKey[j-1], pRoundKey[j]);
        }

        //Arrange block back into 1d key
        for (size_t r=0; r<4; ++r){
            for (size_t c=0; c<4; ++c){
                temp[4*r+c] = roundKey[r][c];
            }
        } 
        //Add new round key to key schedule 
        keySchedule[i+1] = temp;

        //Set new PRK 
        pRoundKey = roundKey;
    }

    return keySchedule;
}

array<array<int,4>,4> encryptBlock(const array<array<int,4>,4>& block, const array<array<int,16>,11> keySchedual){
    //Function will take a plain text 128bit block, perform the Rijndael algorithm with a 
    //128bit key and return the cipher text vesion of the block
    array<array<int,4>,4> state = block;

    //First round
    //AddRoundKey 
    state = addRoundKey(state, keySchedual[0]);

    //9 main rounds
    for (size_t i=1; i<10; ++i){
        //SubBytes
        state = subBytes(state);
        //ShiftRows
        state = shiftRows(state);
        //MixColumns
        state = mixColumns(state);
        //AddRoundKey
        state = addRoundKey(state, keySchedual[i]);
    }

    //Final round
    //SubBytes
    state = subBytes(state);
    //ShiftRows
    state = shiftRows(state);
    //AddRoundKey
    state = addRoundKey(state, keySchedual[10]);

    return state;
}
array<array<int,4>,4> decryptBlock(const array<array<int,4>,4>& block, const array<array<int,16>,11> keySchedual){
    //Function will take a cipher text 128bit block, perform the inverse Rijndael algorithm with a 
    //128bit key and return the plain text vesion of the block
    array<array<int,4>,4> state = {};

    state = block;

    //Undo final round
    //undo AddRoundKey
    state = addRoundKey(state, keySchedual[10]);
    //undo ShiftRows
    state = unShiftRows(state);
    //undo SubBytes
    state = unSubBytes(state);
    
    //Undo 9 main rounds
    for (size_t i=9; i>0; --i){
        //Undo addRoundKey
        state = addRoundKey(state, keySchedual[i]);
        //Undo mixColumns
        state = unMixColumns(state);
        //Undo shiftRows
        state = unShiftRows(state);
        //Undo SubBytes
        state = unSubBytes(state);  
    }
    
    //undo first round
    //Undo AddRoundKey 
    state = addRoundKey(state, keySchedual[0]);

    return state; 
}

int main(){
    /*
    Example values given in FIPS 197
    Input = 32 43 f6 a8 88 5a 30 8d 31 31 98 a2 e0 37 07 34 
    Cipher Key = 2b 7e 15 16 28 ae d2 a6 ab f7 15 88 09 cf 4f 3c
    */

    array<int,16> input = {0x32,0x43,0xf6,0xa8,0x88,0x5a,0x30,0x8d,
                        0x31,0x31,0x98,0xa2,0xe0,0x37,0x07,0x34};;
    array<int,16> key = {0x2b,0x7e,0x15,0x16,0x28,0xae,0xd2,0xa6,
                        0xab,0xf7,0x15,0x88, 0x09,0xcf,0x4f,0x3c};

    array<array<int,16>,11> keySchedual = {};

    array<array<int,4>,4> pTextBlock = {};
    array<array<int,4>,4> cTextBlock = {}; 
    array<int,16> output = {};

    //Generate keys 
    keySchedual = generateKeys(key);

    //Perform encryption 
    cout << "INITIAL INPUT:\n";
    print_1d_State(input);

    pTextBlock = inputToState(input);
    cTextBlock = encryptBlock(pTextBlock, keySchedual);

    cout << "CIPHER TEXT:\n";
    output = stateToOutput(cTextBlock);
    print_1d_State(output);

    //Perform decryption 
    pTextBlock = decryptBlock(cTextBlock, keySchedual);

    cout << "DECRYPTED OUTPUT:\n";
    output = stateToOutput(pTextBlock);
    print_1d_State(output);

}
\$\endgroup\$
0

1 Answer 1

4
\$\begingroup\$

Overall, this is an excellent first attempt. It’s good enough that I would accept it as okay even in a project I was in charge of. It’s not perfect, of course—nothing really is, and there’s always room to improve even experts’ code—but it is solid, and most of the areas of improvement are things that are nontrivial… things I wouldn’t even expect from an experienced C++ programmer.

There is one small bug—technically two, but it’s the exact same bug duplicated in two places—but it’s a doozy. I’m surprised the code even compiles. Certainly I would have expected it to throw up a few warnings, even if warnings are not enabled (but they always should be!).

I’ll get into a line-by-line review of the code, but before I do, I want to cover a few issues that are endemic. Since they come up over and over again, there’s no point mentioning them repeatedly, so I’ll just get them out of the way.

Types

There are no types in the code. For example, for a block, you just use array<int, 16> (or array<array<int, 4>, 4>, depending on the context).

C++ is a strongly-typed language. In fact, it is probably the most strongly-typed language you’ll ever use. It’s all about the types. If you get the types right, everything else Just Works.

It’s hard to appreciate just how magical it gets when the types are perfect: You’ll find that writing code is almost effortless, and everything just works perfectly, automatically. You’ll find that you can’t make mistakes; they’re all caught immediately, and the fix is also immediately obvious.

For example, let’s consider your main() function. There are ~30 lines of code there to set up a cipher, encrypt a block, print it, decrypt it, and print it again. But… that’s just four “things” to do… so, why couldn’t it be done in ~4 lines?

auto main() -> int
{
    auto key = aes_128_key{0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6,
                           0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c};
    auto cipher = aes_128{key};

    auto input = aes_128_block{0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d,
                               0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37, 0x07, 0x34};
    std::cout << "INITIAL INPUT:\n" << input;

    auto output = cipher.encrypt(input);
    std::cout << "CIPHER TEXT:\n" << output;

    output = cipher.decrypt(output);
    std::cout << "DECRYPTED OUTPUT:\n" << output;
}

I mean, do you need any more than that? The key expansion could be computed in the constructor for the aes_128 class; it doesn’t need to be done separately. And with a dedicated block type, you don’t need to transform flat 16-element arrays into 4×4 arrays, or have a separate print function; the block type would know how to do all that itself. And everything could be checked rigorously to detect errors: if you mistakenly provide only 120 bits to the key—you accidentally forgot an octet—that could be detected immediately at compile time (whereas an array would simply silently replace the missing 8 bytes with zero).

And that’s just the tip of the iceberg. With a dedicated type, you could add more convenience functionality. For example, let’s say you want to encrypt an entire file. That could be as easy as:

auto file = std::ifstream{"file.dat", std::ios_base::binary};
auto out = std::ofstream{"file.enc", std::ios_base::binary};

cipher.encrypt(file, out);

The moral here is that when coding in C++, you should look for the “things” and the “operations”… and every “thing” should be an object with a meaningful type, and every “operation” should be a function. With AES, the “things” are things like “key”, “block”, and so on. The “operations” are “encrypt” and “decrypt”, but there are also lower-level operations like “shift rows” and “mix columns”. You’ve done an excellent job of picking out the operations… but you’re missing out on a lot of the power of C++ by not also extracting the “things”.

Unsigned numbers

Throughout your code, you use signed integers—int—for all values. Normally that would exactly the right thing to do. Normally you should always use signed integers—I often see newbies using unsigned integers for things they figure can never be negative… but that’s misguided.

There are only two situations where you should use unsigned integers:

  1. When you are bit-twiddling.
  2. When you are using modular arithmetic.

So, normally, using int would be exactly the right thing to do.

However

This just happens to be a rare case where using unsigned integers is proper:

  1. You are doing bit-twiddling. For example, the final step in each round is XORing the round key and the current state.
  2. More importantly, you are using modular arithmetic. Modular arithmetic (and/or elliptic curves) is one of the key mathematical tricks that makes encryption cheap to do, and expensive to undo.

So, this is a very rare situation where I have to reverse my usual rule, and say you shouldn’t use signed integers, you should use unsigned integers.

Now, this kinda connects to the note about types above. If you had used specific types throughout, then there would only be one or two places where you’d have to edit the code to switch from signed int to unsigned int. For example, if you had a block type, there would be maybe one or two lines in that type where you would have to make the switch… and then everywhere that used the block type would be automatically fixed. But because you used ad hoc types like array<int, 16>, you’d have to make the fix all over the place.

Testing

It’s good that you checked that your code actually works! You’d be surprised, and probably horrified, at how often I see beginners write code, then submit it for review, and they haven’t even checked that it works—sometimes they haven’t even checked that it compiles.

However, one of the best practices you can learn as a programmer—and not just a C++ programmer—is testing. A good C++ developer always tests their code; in fact, I’d say that anyone who doesn’t test their code is simply not a good C++ developer. I actually recommend writing the tests before you start writing the actual code; this can be a great way to make sure your interface is good, and it can help you get a better understanding of what you need, and what you don’t.

You should choose a testing framework, then learn it, love it, live it. Personally, I swear by Boost.Test, but that’s not for the faint of heart; it can be a nightmare to set up, but once it’s up and going, it’s wonderful. A much simpler, but still excellent option is Catch2. I’ll use Catch2 to illustrate the idea of testing here.

Now, AES has a “proper” testing procedure specified by NIST—the Advanced Encryption Standard Algorithm Validation System (AESAVS). This procedure is mega complex, with 3 different test categories, and 7 different encryption modes. For now, let’s just focus on the easiest: the known answer test (KAT), in electronic code book (ECB) mode. You can get a copy of the KAT test vectors from the Internet Archive.

If you download the KAT zip file and extract it, you’ll see 144 files. The ones we want are the ones that start with ECB, with 128 in them (for the 128-bit key size). That narrows the list to 8 files:

  • ECBGFSbox128e.txt
  • ECBKeySbox128e.txt
  • ECBVarKey128e.txt
  • ECBVarTxt128e.txt
  • ECBGFSbox128d.txt
  • ECBKeySbox128d.txt
  • ECBVarKey128d.txt
  • ECBVarTxt128d.txt

The ECB is the mode, then the test type, then 128 for the key size, then e or d for encryption or decryption. The four test types are:

  • GFSbox. This sets the key to zero, to focus on testing the substitution box in the rounds.
  • KeySbox. This sets the text to zero, to focus on testing the key expansion.
  • VarKey for variable key. Basically, set the key and the text to zero, and then set one bit at a time in the key.
  • VarText for variable text. Basically, set the key and the text to zero, and then set one bit at a time in the text.

So let’s pick the GFSbox type, and open up ECBGFSbox128e.txt, and this is what we see:

[ENCRYPT]

COUNT = 0
KEY = 00000000000000000000000000000000
PLAINTEXT = f34481ec3cc627bacd5dc3fb08f273e6
CIPHERTEXT = 0336763e966d92595a567cc9ce537f5e

COUNT = 1
KEY = 00000000000000000000000000000000
PLAINTEXT = 9798c4640bad75c7c3227db910174e72
CIPHERTEXT = a9a1631bf4996954ebc093957b234589

COUNT = 2
KEY = 00000000000000000000000000000000
PLAINTEXT = 96ab5c2ff612d9dfaae8c31f30c42168
CIPHERTEXT = ff4f8391a6a40ca5b25d23bedd44a597

COUNT = 3
KEY = 00000000000000000000000000000000
PLAINTEXT = 6a118a874519e64e9963798a503f1d35
CIPHERTEXT = dc43be40be0e53712f7e2bf5ca707209

COUNT = 4
KEY = 00000000000000000000000000000000
PLAINTEXT = cb9fceec81286ca3e989bd979b0cb284
CIPHERTEXT = 92beedab1895a94faa69b632e5cc47ce

COUNT = 5
KEY = 00000000000000000000000000000000
PLAINTEXT = b26aeb1874e47ca8358ff22378f09144
CIPHERTEXT = 459264f4798f6a78bacb89c15ed3d601

COUNT = 6
KEY = 00000000000000000000000000000000
PLAINTEXT = 58c8e00b2631686d54eab84b91f0aca1
CIPHERTEXT = 08a4e2efec8a8e3312ca7460b9040bbf

In theory, you’re supposed to write a whole testing suite that reads this file (and the others), parses it, and then automatically generates the tests. But, let’s just do it the simple way for now. There are 7 tests in that file, each with a key of zero, a different input (PLAINTEXT), and the expected output (CIPHERTEXT). Let’s just copy out the inputs and outputs:

TEST_CASE("ECBGFSbox128e")
{
    auto const data = GENERATE(table<aes_128_block, aes_128_block>>({
        {{0xf3, 0x44, 0x81, 0xec, 0x3c, 0xc6, 0x27, 0xba, 0xcd, 0x5d, 0xc3, 0xfb, 0x08, 0xf2, 0x73, 0xe6}, {0x03, 0x36, 0x76, 0x3e, 0x96, 0x6d, 0x92, 0x59, 0x5a, 0x56, 0x7c, 0xc9, 0xce, 0x53, 0x7f, 0x5e}},
        {{0x97, 0x98, 0xc4, 0x64, 0x0b, 0xad, 0x75, 0xc7, 0xc3, 0x22, 0x7d, 0xb9, 0x10, 0x17, 0x4e, 0x72}, {0xa9, 0xa1, 0x63, 0x1b, 0xf4, 0x99, 0x69, 0x54, 0xeb, 0xc0, 0x93, 0x95, 0x7b, 0x23, 0x45, 0x89}},
        // ... and so on ...
    }));

    auto key = aes_128_key{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
    auto cipher = aes_128{key};

    auto input = std::get<0>(data);
    auto output = std::get<1>(data);

    REQUIRE(cipher.encrypt(input) == output);
}

That will generate 7 tests, that each use a zero key to encrypt the given input, then checks if it equals the expected output. If the encryption algorithm is correct, those 7 tests will all pass.

(Note that I think this technically not the correct way to do the AESAVS ECB tests. Doing it this way will recreate the cipher for each of the 7 tests. I think the proper way is to make the cipher only once and reuse it. But, whatever.)

If you do basically the same thing for the other three test types, and also do decryption tests as well, you will end up with dozens of tests. If your algorithm passes all of those, then you can be pretty damn confident it’s okay.

And if you want to get crazy, you can also try writing tests for the other, chained and counter modes. And if you want get really crazy, you can try implementing the multi-block and Monte Carlo tests. If you go that far, you’ve basically passed the AESAVS, and could submit your implementation for certification.

But let’s not get crazy here. Generally, when you are implementing an algorithm that transforms one thing into another, you will want to write a test that basically looks like:

TEST_CASE("algorithm test")
{
    auto const data = GENERATE(table<input_type, output_type>>({
        {input_1, expected_1},
        {input_2, expected_2},
        {input_3, expected_3},
        // ... and so on ...
    }));

    REQUIRE(algorithm(std::get<0>(data)) == std::get<1>(data));
}

The whole AESAVS dance was just about getting those input and expected values.

I took the liberty of setting up a Godbolt that shows the test in action. I had to make some small fixes to get it to compile (which I’ll mention in the code review), but as you can see, all the tests pass.

This is how you properly test an algorithm. Doing it this way makes the tests automated, and does the checking in code—you don’t need to eyeball output to make sure it’s good. In practice, you should implement the tests first, which means they will (and should!) all fail at first, and then you start implementing the algorithm to start making tests pass. This should be your standard development cycle: run tests, write code to fix failures, repeat.

Code review

#include <iostream>
#include <array>

You are missing quite a few includes that you need. For example, you use printf(), but don’t include <cstdio>; you use size_t, but don’t include anything that defines that (there are several options).

using namespace std;

Never, never do this. It’s not worth all the potential problems, which can be both subtle and infuriating. It’s just a few more characters to write the std:: wherever needed, and it helps identify when you are using standard library types, versus when you are using your own types or other library’s types.

Note that if you’re actually using your own types, you’ll be using a lot fewer std:: types anyway.

void printBlock(const array<array<int,4>,4>& block){
    //Function to print a block to terminal
    for (size_t i = 0; i < 4; ++i){
        printf("%.2x, %.2x, %.2x, %.2x\n", block[i][0], block[i][1], block[i][2], block[i][3]);
    }
} 
void printRow(const array<int,4>& row){
    //Function to print a row to terminal 
    printf("%.2x, %.2x, %.2x, %.2x\n", row[0], row[1], row[2], row[3]);
}
void print_1d_State(const array<int,16>& e){
    //Function to print a 1d version of the state to terminal 
    printf("%.2x, %.2x, %.2x, %.2x, %.2x, %.2x, %.2x, %.2x, %.2x, %.2x, %.2x, %.2x, %.2x, %.2x, %.2x, %.2x\n", 
    e[0], e[1], e[2], e[3], e[4], e[5], e[6], e[7], e[8], e[9], e[10], e[11], e[12], e[13], e[14], e[15]);
}
void printKeySchedule(const array<array<int,16>,11>& keySchedual){
    //Function to print keys for debugging
    for (size_t i=0; i<11; ++i){
        print_1d_State(keySchedual[i]);
    }
}

As I mentioned in the preamble, if you write proper types, you can give those types their own stream inserters, and you won’t need these functions. You could just do std::cout << ___.

However, let’s assume we need these functions. In that case, you shouldn’t use the C library functions; you should use IOstreams. Ideally you should take the stream as a parameter, too, so you can print to any stream.

For maximum efficiency, you should probably write to a buffer, and then stream the whole buffer. For example (with hacky and untested code):

template <typename CharT, typename Traits>
auto print_1d_state(std::basic_ostream<CharT, Traits>& out, std::array<int, 16> const& state)
{
    constexpr auto digits = std::array{
        '0', '1', '2', '3', '4', '5', '6', '7',
        '8', '9', 'a', 'b', 'c', 'd', 'e', 'f'};

    auto buffer = std::array<char, 64>{};

    auto p = buffer.data();
    for (auto byte : state)
    {
        *p++ = digits[(static_cast<unsigned int>(byte) & 0xF0u) >> 4];
        *p++ = digits[static_cast<unsigned int>(byte) & 0xFu];
        *p++ = ',';
        *p++ = ' ';
    }

    p -= 2;
    *p++ = '\n';
    *p++ = '\0';

    out << buffer.data();
}

// usage:
//  print_1d_state(std::cout, state);

But again, this would be better as an inserter for an actual state type.

array<array<int,4>,4> inputToState(const array<int,16>& input){
    //Function turns a 128bit/ 16 element array and minipulates into a 4x4 matrix
    array<array<int,4>,4> state = {};

    //Loop the input and assign to state
    for(size_t i=0; i<4; ++i){
        for(size_t j=0; j<4; ++j){
            state[i][j] = input[i+4*j];
        }
    }
    
    return state;
}
array<int,16> stateToOutput(const array<array<int,4>,4>& state){
    //Function takes a 128bit 4x4 block/ matrix and minipulates into a 1d array of words  
    array<int,16> output = {};

    //Loop the state and assign to output
    for(size_t i=0; i<4; ++i){
        for(size_t j=0; j<4; ++j){
            output[i+4*j] = state[i][j];
        }
    }

    return output;
}

The only reason these two functions are necessary is because you are using raw arrays, and need to convert a flat array to a square “matrix”.

But suppose you had an actual state type instead. And suppose it looked like this:

class state
{
    std::array<std::byte, 16> _bytes;

public:
    // ... [snip] ...

    constexpr auto at(std::size_t n) const { return _bytes[n]; }
    constexpr auto at(std::size_t i, std::size_t j) const { return _bytes[i + (4 * j)]; }

    // ... [snip] ...
};

You see? You don’t need two different types to have the best of both worlds, and there’s no need to convert back and forth. Want to treat the state as a flat array of 16 octets? Use .at(0) to .at(15). Want to treat the state as a 4×4 matrix? Use .at(0, 0) to .at(3, 3).

(Note that in a future version of C++, you will be allowed to use commas in subscripts. So you will one day be able to write s[0] for a 1D flat view, and s[0, 0] for a 2D matrix view.)

array<array<int,4>,4> keyToWords(const array<int,16>& roundKey){
    //Function will take a rounds key and minipulate it into a 4x4 array
    //With each row being a word
    array<array<int,4>,4> wordBlock = {};

    for (size_t i= 0; i < 4; ++i){
        //Get words from key 
        wordBlock[i] = {roundKey[4*i], roundKey[1+i*4], roundKey[2+i*4], roundKey[3+i*4]};
    }
    return wordBlock;
}

This is another transformation function that would be unnecessary with a custom state type.

The next few functions are fine.

int lookupByte(int &byte){
    //This method takes a byte, performs a lookup in the AES SBOX and returns the corresponding value

    int x = byte >> 4; //Shifts 4 bits right i.e. takes first 4 bits and discards the rest 
    int y = byte & 0x0f; // 0x0f = 15 = 00001111(bin). Effectivly takes last 4 bits and dicards the rest

    const int sbox[16][16] = {{0x63, 0x7c, 0x77, 0x7b, 0xf2, 0x6b, 0x6f, 0xc5,0x30, 0x01, 0x67, 0x2b, 0xfe, 0xd7, 0xab, 0x76},
                          {0xca, 0x82, 0xc9, 0x7d, 0xfa, 0x59, 0x47, 0xf0,0xad, 0xd4, 0xa2, 0xaf, 0x9c, 0xa4, 0x72, 0xc0},
                          {0xb7, 0xfd, 0x93, 0x26, 0x36, 0x3f, 0xf7, 0xcc,0x34, 0xa5, 0xe5, 0xf1, 0x71, 0xd8, 0x31, 0x15},
                          {0x04, 0xc7, 0x23, 0xc3, 0x18, 0x96, 0x05, 0x9a,0x07, 0x12, 0x80, 0xe2, 0xeb, 0x27, 0xb2, 0x75},
                          {0x09, 0x83, 0x2c, 0x1a, 0x1b, 0x6e, 0x5a, 0xa0,0x52, 0x3b, 0xd6, 0xb3, 0x29, 0xe3, 0x2f, 0x84},
                          {0x53, 0xd1, 0x00, 0xed, 0x20, 0xfc, 0xb1, 0x5b,0x6a, 0xcb, 0xbe, 0x39, 0x4a, 0x4c, 0x58, 0xcf},
                          {0xd0, 0xef, 0xaa, 0xfb, 0x43, 0x4d, 0x33, 0x85,0x45, 0xf9, 0x02, 0x7f, 0x50, 0x3c, 0x9f, 0xa8},
                          {0x51, 0xa3, 0x40, 0x8f, 0x92, 0x9d, 0x38, 0xf5,0xbc, 0xb6, 0xda, 0x21, 0x10, 0xff, 0xf3, 0xd2},
                          {0xcd, 0x0c, 0x13, 0xec, 0x5f, 0x97, 0x44, 0x17,0xc4, 0xa7, 0x7e, 0x3d, 0x64, 0x5d, 0x19, 0x73},
                          {0x60, 0x81, 0x4f, 0xdc, 0x22, 0x2a, 0x90, 0x88,0x46, 0xee, 0xb8, 0x14, 0xde, 0x5e, 0x0b, 0xdb},
                          {0xe0, 0x32, 0x3a, 0x0a, 0x49, 0x06, 0x24, 0x5c,0xc2, 0xd3, 0xac, 0x62, 0x91, 0x95, 0xe4, 0x79},
                          {0xe7, 0xc8, 0x37, 0x6d, 0x8d, 0xd5, 0x4e, 0xa9,0x6c, 0x56, 0xf4, 0xea, 0x65, 0x7a, 0xae, 0x08},
                          {0xba, 0x78, 0x25, 0x2e, 0x1c, 0xa6, 0xb4, 0xc6,0xe8, 0xdd, 0x74, 0x1f, 0x4b, 0xbd, 0x8b, 0x8a},
                          {0x70, 0x3e, 0xb5, 0x66, 0x48, 0x03, 0xf6, 0x0e,0x61, 0x35, 0x57, 0xb9, 0x86, 0xc1, 0x1d, 0x9e},
                          {0xe1, 0xf8, 0x98, 0x11, 0x69, 0xd9, 0x8e, 0x94,0x9b, 0x1e, 0x87, 0xe9, 0xce, 0x55, 0x28, 0xdf},
                          {0x8c, 0xa1, 0x89, 0x0d, 0xbf, 0xe6, 0x42, 0x68,0x41, 0x99, 0x2d, 0x0f, 0xb0, 0x54, 0xbb, 0x16}};

     return sbox[x][y];
    
}

Now, I mentioned at the top that you should be using unsigned types, but there is another issue with using int: AES is a byte-wise algorithm (where bytes are 8 bits). Basically everywhere you have int, you should have unsigned char or std::byte. Normally that doesn’t matter much, because if arguments are going to be passed in registers, they’ll be basically expanded to int size regardless. However, it does have some costs, and some that are quite significant. For example, if all your operands are int sized, then you’ll fit a lot fewer of them at a time when doing vectorized operations… and vectorization could speed things up a lot, so you really want as much as that to happen as possible.

But even here there is a huge cost. That substitution box is 256 elements. If those elements were bytes, it would be 256 bytes. But you’ve used ints. Given that ints are often 2, 4, and sometimes even 8 bytes, this array will actually be 512, 1024, or 2048 bytes. An array of 256 bytes would have a much better chance of being able to fit in cache memory, and stay there, than a 1 kiB array. And if your substitution box is being bumped in and out of cache every round… your encryption will be slow.

The solution is to use bytes. And in fact, you should use bytes everywhere, rather than ints. This will not only allow more stuff to fit in cache during encryption/decryption, it will also allow for more vectorization.

Now, the question is whether to use unsigned char… which is the old school way of writing “byte”… or std::byte. The latter is more explicit, and safer… butstd::byte still has still a lot of missing functionality, and using it can be a little clunky. I’ll show you what code looks like with both so you can decide:

// With unsigned char:
constexpr auto lookupByte(unsigned char b)
{
    constexpr auto sbox = std::array<unsigned char, 256>{
        0x63u, 0x7cu, 0x77u, 0x7bu, 0xf2u, 0x6bu, 0x6fu, 0xc5u,
        0x30u, 0x01u, 0x67u, 0x2bu, 0xfeu, 0xd7u, 0xabu, 0x76u,
        // ...
        0x8cu, 0xa1u, 0x89u, 0x0du, 0xbfu, 0xe6u, 0x42u, 0x68u,
        0x41u, 0x99u, 0x2du, 0x0fu, 0xb0u, 0x54u, 0xbbu, 0x16u
    };

    return sbox[b];
}

// With std::byte:
constexpr auto lookupByte(std::byte b)
{
    // Unfortunately, there are no byte literals (though there should be!),
    // and (unsigned) integer literals won't implicitly convert (because they
    // might overflow - this is a good thing), so we need to *EXPLICITLY*
    // convert every value to byte.
    //
    // We could do std::byte{0x??} for every single value… but we can save
    // some typing and some space if we use an alias for std::byte.

    using x = std::byte;

    constexpr auto sbox = std::array<std::byte, 256>{
        x{0x63u}, x{0x7cu}, x{0x77u}, x{0x7bu}, x{0xf2u}, x{0x6bu}, x{0x6fu}, x{0xc5u},
        x{0x30u}, x{0x01u}, x{0x67u}, x{0x2bu}, x{0xfeu}, x{0xd7u}, x{0xabu}, x{0x76u},
        // ...
        x{0x8cu}, x{0xa1u}, x{0x89u}, x{0x0du}, x{0xbfu}, x{0xe6u}, x{0x42u}, x{0x68u},
        x{0x41u}, x{0x99u}, x{0x2du}, x{0x0fu}, x{0xb0u}, x{0x54u}, x{0xbbu}, x{0x16u}
    };

    // And, std::byte doesn’t automatically convert to a number (which is
    // another good thing), so we need to manually do that.

    return sbox[std::to_integer<unsigned int>(b)];
    // or, in C++23:
    //  return sbox[std::to_underlying(b)];
}

std::byte is wordier because you need to be more explicit, which is really a good thing, but it can be aggravating.

Whether you use unsigned int or std::byte, what you really should do is use an alias.

Here’s how: First, you should start with having everything in your own namespace. For me, I’d use namespace indi, and then I’d probably put all the encryption stuff in a encryption namespace. For you, I’ll assume namespace wire::encryption.

Once you have a namespace, you’d just make an alias for byte:

namespace wire::encryption {

using byte = unsigned char; // or std::byte

} // namespace wire::encryption

And then you’d use byte everywhere, instead of int.

I’d also add a couple of convenience features. I would add a user-defined literal for bytes, and a conversion function to numbers:

namespace wire::encryption {

using byte = unsigned char; // or std::byte

consteval auto operator""_b(unsigned long long int value)
{
    if (value > 0xFFu)
        throw std::out_of_range{"byte value out of range"};

    return static_cast<byte>(value);
}

constexpr auto to_number(byte b) noexcept
{
    return static_cast<unsigned char>(b);
}

} // namespace wire::encryption

The UDL allows you to do:

    constexpr auto sbox = std::array{
        0x63_b, 0x7c_b, 0x77_b, 0x7b_b, 0xf2_b, 0x6b_b, 0x6f_b, 0xc5_b,
        0x30_b, 0x01_b, 0x67_b, 0x2b_b, 0xfe_b, 0xd7_b, 0xab_b, 0x76_b,
        // ...
        0x8c_b, 0xa1_b, 0x89_b, 0x0d_b, 0xbf_b, 0xe6_b, 0x42_b, 0x68_b,
        0x41_b, 0x99_b, 0x2d_b, 0x0f_b, 0xb0_b, 0x54_b, 0xbb_b, 0x16_b
    };

And the conversion helper allows you to do:

    return sbox[to_number(b)];

And everything will work whether you use unsigned char or std::byte.

Anywho, going back to the function in question:

int lookupByte(int &byte){
    //This method takes a byte, performs a lookup in the AES SBOX and returns the corresponding value

    int x = byte >> 4; //Shifts 4 bits right i.e. takes first 4 bits and discards the rest 
    int y = byte & 0x0f; // 0x0f = 15 = 00001111(bin). Effectivly takes last 4 bits and dicards the rest

    const int sbox[16][16] = {{0x63, 0x7c, 0x77, 0x7b, 0xf2, 0x6b, 0x6f, 0xc5,0x30, 0x01, 0x67, 0x2b, 0xfe, 0xd7, 0xab, 0x76},
                          {0xca, 0x82, 0xc9, 0x7d, 0xfa, 0x59, 0x47, 0xf0,0xad, 0xd4, 0xa2, 0xaf, 0x9c, 0xa4, 0x72, 0xc0},
                          {0xb7, 0xfd, 0x93, 0x26, 0x36, 0x3f, 0xf7, 0xcc,0x34, 0xa5, 0xe5, 0xf1, 0x71, 0xd8, 0x31, 0x15},
                          {0x04, 0xc7, 0x23, 0xc3, 0x18, 0x96, 0x05, 0x9a,0x07, 0x12, 0x80, 0xe2, 0xeb, 0x27, 0xb2, 0x75},
                          {0x09, 0x83, 0x2c, 0x1a, 0x1b, 0x6e, 0x5a, 0xa0,0x52, 0x3b, 0xd6, 0xb3, 0x29, 0xe3, 0x2f, 0x84},
                          {0x53, 0xd1, 0x00, 0xed, 0x20, 0xfc, 0xb1, 0x5b,0x6a, 0xcb, 0xbe, 0x39, 0x4a, 0x4c, 0x58, 0xcf},
                          {0xd0, 0xef, 0xaa, 0xfb, 0x43, 0x4d, 0x33, 0x85,0x45, 0xf9, 0x02, 0x7f, 0x50, 0x3c, 0x9f, 0xa8},
                          {0x51, 0xa3, 0x40, 0x8f, 0x92, 0x9d, 0x38, 0xf5,0xbc, 0xb6, 0xda, 0x21, 0x10, 0xff, 0xf3, 0xd2},
                          {0xcd, 0x0c, 0x13, 0xec, 0x5f, 0x97, 0x44, 0x17,0xc4, 0xa7, 0x7e, 0x3d, 0x64, 0x5d, 0x19, 0x73},
                          {0x60, 0x81, 0x4f, 0xdc, 0x22, 0x2a, 0x90, 0x88,0x46, 0xee, 0xb8, 0x14, 0xde, 0x5e, 0x0b, 0xdb},
                          {0xe0, 0x32, 0x3a, 0x0a, 0x49, 0x06, 0x24, 0x5c,0xc2, 0xd3, 0xac, 0x62, 0x91, 0x95, 0xe4, 0x79},
                          {0xe7, 0xc8, 0x37, 0x6d, 0x8d, 0xd5, 0x4e, 0xa9,0x6c, 0x56, 0xf4, 0xea, 0x65, 0x7a, 0xae, 0x08},
                          {0xba, 0x78, 0x25, 0x2e, 0x1c, 0xa6, 0xb4, 0xc6,0xe8, 0xdd, 0x74, 0x1f, 0x4b, 0xbd, 0x8b, 0x8a},
                          {0x70, 0x3e, 0xb5, 0x66, 0x48, 0x03, 0xf6, 0x0e,0x61, 0x35, 0x57, 0xb9, 0x86, 0xc1, 0x1d, 0x9e},
                          {0xe1, 0xf8, 0x98, 0x11, 0x69, 0xd9, 0x8e, 0x94,0x9b, 0x1e, 0x87, 0xe9, 0xce, 0x55, 0x28, 0xdf},
                          {0x8c, 0xa1, 0x89, 0x0d, 0xbf, 0xe6, 0x42, 0x68,0x41, 0x99, 0x2d, 0x0f, 0xb0, 0x54, 0xbb, 0x16}};

     return sbox[x][y];
    
}

First, I don’t see the point of taking the argument by reference. I assume that’s a typo.

Making the s-box a 2D array and then using bit-twiddling to extract the two nibbles from the byte seems a bit pointless. Why not make the s-box a 1D array, and just use the byte value directly?

I would also make the function constexpr, and I would consider making it noexcept.

Making it noexcept is a bit tricky, because a byte may not be 8-bits. If bytes are 8-bits, then the function cannot fail: all possible 8-bit values are covered in the s-box. However… if a byte is greater than 8 bits, then it is possible to have a value outside the range of 0–255… which means the function could fail. What you could do is use a conditional noexcept specifier, and only make it noexcept if a byte is 8 bits:

constexpr auto lookupByte(byte b) noexcept(std::numeric_limits<unsigned char>::digits == 8)
{
    // ...
}

There’s one more thing I’d suggest. When filling out a big block of numbers like this, it is easy to screw up and add too many or too few values. If you specify the array size, and add too many values, the compiler will catch it, which is cool. But if you add too few values, the compiler won’t complain… it will just add zeroes at the end. To avoid this problem, I would recommend letting the compiler deduce the number of values, and then double check that the count is correct:

    constexpr auto sbox = std::array{
        0x63_b, 0x7c_b, 0x77_b, 0x7b_b, 0xf2_b, 0x6b_b, 0x6f_b, 0xc5_b,
        0x30_b, 0x01_b, 0x67_b, 0x2b_b, 0xfe_b, 0xd7_b, 0xab_b, 0x76_b,
        // ...
        0x8c_b, 0xa1_b, 0x89_b, 0x0d_b, 0xbf_b, 0xe6_b, 0x42_b, 0x68_b,
        0x41_b, 0x99_b, 0x2d_b, 0x0f_b, 0xb0_b, 0x54_b, 0xbb_b, 0x16_b
    };

    static_assert(sbox.size() == 256);

The inverse function is essentially the same, so let’s skip to the actual substitution function:

auto subBytes(const array<array<int,4>,4>& state){
    //Function takes the current state and subs each byte using sbox look up 
    //the sboxLookup variable indicates the use of sbox or sboxInv i.e. forward or reverse lookup. 
    array<array<int,4>,4> result = {};
    int byte; 

    for (size_t i=0; i<4; ++i){
        for (size_t j=0; j<4; ++j){
            byte = state[i][j];
            result[i][j] = lookupByte(byte);
        }
    }

    return result;
}

Generally, writing naked for loops in modern C++ is frowned upon. Instead, you should use algorithms.

Let’s assume the state is a flat array<int, 16> instead of a 2D array. In that case, this function could be:

constexpr auto subBytes(std::array<int, 16> state)
{
    std::ranges::transform(state, state.begin(), lookupByte);

    return state;
}

Obviously, things are more complicated with a 2D array… but that’s really a hint that you shouldn’t be using 2D arrays.

On to the row shift:

array<int,4> shiftRow(const array<int,4>& row, const int shift){
    //Recursive function to shft a row left by given value 
    array<int,4> result = {};
    result = row;

    if(shift){
        //Shift by 1
        int temp = result[0]; 

        for (size_t i=0; i<3; ++i){
            result[i] = result[i+1];
        }
        result[3] = temp; 
        //reduce shift and perform again
        result = shiftRow(result, shift -1);
    }
    else{ 
        return result;
    }
}

This is a place where using an algorithm would really help, for a couple of reasons.

First, the simplicity. Check it out:

constexpr auto shiftRow(array<int, 4> row, int const shift) noexcept
{
    std::ranges::rotate(row, row.begin() + shift);

    return row;
}

Second, and more importantly… you have a serious bug. In fact, I’m surprised your code even compiles. Surely it must give some warnings, at least.

The problem is that if shift is not zero… what does the function return? Answer: it returns nothing; it just falls off the end and boom, undefined behaviour. (In practice, it looks like you may be getting lucky, in that the UB is just that the last return value that was left in the return slot gets returned again… which happens to make things work.)

The fix is simple. Instead of the last few lines being:

    else{ 
        return result;
    }

They should just be:

    return result;

(This is basically what I had to do to get the Godbolt tests to compile.)

But again, the smart thing to do is just use the built-in algorithm. It’s simpler, it’s already rigorously tested, and it just works.

The same applies to the next function:

auto shiftRows(const array<array<int,4>,4>& state){
    //Function to shift rows of the state left by the value of their index
    array<array<int,4>,4> result = {};

    for (size_t i = 0; i < 4; i++){
        result[i] = shiftRow(state[i], i);
    }

    return result;
}

This could be:

constexpr auto shiftRows(std::array<std::array<byte, 4>, 4> state)
{
    std::ranges::transform(state, state.begin(), shiftRow);

    return state;
}

Note that I take the argument by value, not by const&, because we’re making a copy anyway. You could also take it by const& as you do, then copy it into result, again, as you do; that’s fine, nothing wrong with that at all.

The un-shift functions are basically identical, so I’ll skip ’em.

auto mixColumn(const array<int,4>& stateColumn){
    //Function takes a column of the state and 'mixes' it according to the matrix multiplication defined in FIPS 197
    array<int,4> result = {};
    
    result[0] = multiply_by_2(stateColumn[0]) ^ multiply_by_3(stateColumn[1]) ^ stateColumn[2] ^ stateColumn[3];
    result[1] = multiply_by_2(stateColumn[1]) ^ multiply_by_3(stateColumn[2]) ^ stateColumn[3] ^ stateColumn[0];
    result[2] = multiply_by_2(stateColumn[2]) ^ multiply_by_3(stateColumn[3]) ^ stateColumn[0] ^ stateColumn[1];
    result[3] = multiply_by_2(stateColumn[3]) ^ multiply_by_3(stateColumn[0]) ^ stateColumn[1] ^ stateColumn[2];

    return result;
}

This is fine, but it could be simplified. There’s no need to create the result array, and then later fill it in—you could do it all in one step—but realistically, that unnecessary step will be optimized away by any modern compiler. Still, you could just as easily write:

constexpr auto mixColumn(std::array<byte, 4> const& stateColumn)
{
    return std::array<byte, 4>{
        multiply_by_2(stateColumn[0]) ^ multiply_by_3(stateColumn[1]) ^ stateColumn[2] ^ stateColumn[3],
        multiply_by_2(stateColumn[1]) ^ multiply_by_3(stateColumn[2]) ^ stateColumn[3] ^ stateColumn[0],
        multiply_by_2(stateColumn[2]) ^ multiply_by_3(stateColumn[3]) ^ stateColumn[0] ^ stateColumn[1],
        multiply_by_2(stateColumn[3]) ^ multiply_by_3(stateColumn[0]) ^ stateColumn[1] ^ stateColumn[2]
    };
}

The next function could also be simplified quite a bit:

auto mixColumns(const array<array<int,4>,4>& state){
    //Function takes the current states and applies the 'mixColumn' function column by column
    array<array<int,4>,4> result = {};
    array<int,4> column = {};
    
    //Grab the columns and mix
    for(size_t i = 0; i < 4; ++i){
        for(size_t j = 0; j < 4; ++j){
            column[j] = state[j][i];
        }

        column = mixColumn(column);   

        //Transpose the columns back into rows
        for(size_t j = 0; j < 4; ++j){
            result[j][i] = column[j];
        }
    }

    return result;
}

You see the two inner loops? They do the same thing, in the same order, using column as an intermediate to bridge between them. You could simplify that to:

constexpr auto mixColumns(std::array<std::array<byte, 4>, 4> const& state)
{
    auto result = std::array<std::array<byte, 4>, 4>{};

    for(auto i = 0; i < 4; ++i)
    {
        for(int j = 0; j < 4; ++j)
            result[j][i] = mixColumn(state[j][i]);
    }

    return result;
}

Let’s skip unMixColums().

array<int,4> xorWords(const array<int,4>& wordA, const array<int,4>& wordB){
    //Function takes 2 words, XORs each element and returns the result 
    array<int,4> result = {};
    
    for (size_t i = 0; i < 4; i++){
        result[i] = wordA[i] ^ wordB[i];
    }
    return result;
}

Another function that could be massively simplified and made safer with standard algorithms:

auto xorWords(std::array<byte, 4> const& wordA, std::array<byte, 4> const& wordB)
{
    auto result = std::array<byte, 4>{};

    std::ranges::transform(wordA, wordB, result.begin(), std::bit_xor<>{});

    return result;
}

I’ll skip the next few functions because it’s mostly more of the same.

array<array<int,16>,11> generateKeys(const array<int,16>& key){
    //Function performs the key expansion algorithm to expand a single 128 bit key into 
    //10 round keys know as the keySchedule. Note this is coded specifically for 128bit keys and
    //will NOT work for other AES varients 
    array<array<int,16>,11> keySchedule = {};
    array<array<int,4>,4> roundKey = {};
    array<array<int,4>,4> pRoundKey = {};
    array<int,16> temp = {};

    const array<array<int,4>,10> RCON = {{{0x01, 0x00, 0x00, 0x00},
                                        {0x02, 0x00, 0x00, 0x00},
                                        {0x04, 0x00, 0x00, 0x00},
                                        {0x08, 0x00, 0x00, 0x00},
                                        {0x10, 0x00, 0x00, 0x00},
                                        {0x20, 0x00, 0x00, 0x00},
                                        {0x40, 0x00, 0x00, 0x00},
                                        {0x80, 0x00, 0x00, 0x00},
                                        {0x1b, 0x00, 0x00, 0x00},
                                        {0x36, 0x00, 0x00, 0x00}}};

    keySchedule[0] = key; 

    //Extract the words from the key and arrange into a block
    for (size_t i=0; i<4; ++i){
        pRoundKey[i] = {key[4*i], key[4*i+1], key[4*i+2], key[4*i+3]};
    }

    //Loop rounds
    for (size_t i=0; i<10; ++i){
        //Rotate last word of PRK
        roundKey[0] = rotWord(pRoundKey[3]);
        //SubBytes of last word in PRK
        roundKey[0] = subWord(roundKey[0]);
        //XOR : first word PRK, current RK state, rcon(round)
        roundKey[0] = xorWords(xorWords(pRoundKey[0], roundKey[0]), RCON[i]);
        //XOR the other words in sequential order
        for (size_t j=1; j<4; ++j){
            roundKey[j] = xorWords(roundKey[j-1], pRoundKey[j]);
        }

        //Arrange block back into 1d key
        for (size_t r=0; r<4; ++r){
            for (size_t c=0; c<4; ++c){
                temp[4*r+c] = roundKey[r][c];
            }
        } 
        //Add new round key to key schedule 
        keySchedule[i+1] = temp;

        //Set new PRK 
        pRoundKey = roundKey;
    }

    return keySchedule;
}

I note that you list all or most of the functions’ variables at the top of the function. That is an archaic practice going back to the primordial days of C—nowadays, even C programmers don’t do that anymore.

Declaring all the variables at the top creates both confusion, and the possibility of bugs when values get jiggered with in surprising places. For example, I think that you intend for pRoundKey to be carried over from loop iteration to loop iteration, while roundKey exists only within the loop. (Giving variables such similar names is also asking for trouble.)

Instead, declare variables right where they’re needed.

The name RCON is a dangerous idea. In C++, identifiers in ALL_CAPS (also known as SCREAMING_SNAKE_CASE), are traditionally reserved for the preprocessor. If you use such an identifier, you risk having it stomped on by someone’s macro.

Summary

Aside from that one bug (which, again, I’m really surprised the compiler didn’t choke on, or at least warn about), the implementation obviously works; you checked it yourself and saw that it gives the correct result, and I ran a small battery of tests, and it worked.

The biggest problems with the code are:

  1. a lack of custom types; and
  2. not using standard algorithms, and instead rolling your own loops all over the place.

Making proper types can make the code both much simpler, and much safer, because your custom types can implement any useful operations, and prevent any undesired ones. Getting the types right is the secret to good C++ code; otherwise, you’re just writing jazzed-up C.

One noteworthy issue with the code that is related to not using custom types is the fact that you spend a lot of time converting back and forth between flat 16 element arrays to 4×4 2D arrays. Indeed, close to half the implementation is just converting back and forth, which is both a fantastic waste of time, and a lot of unnecessary code.

Even if you’re not going to use custom types, it would probably make more sense to just use flat arrays throughout—never 2D arrays—and wherever it would be more convenient to view the flat array as a 4×4 matrix, use something like the proposed mdspan, or just a simple function to convert a 2D (i, j) pair to a flat index, like so:

template <std::size_t M, std::size_t N>
constexpr auto to_index(std::size_t i, std::size_t j)
{
    // An optional check, if you want:
    if constexpr (not NDEBUG)
    {
        if (i >= M or j >= N)
            throw std::out_of_range{"coordinate out of range"};
    }

    return i * M + j;
}

////////////////////////////////////////
// Usage:
auto block = std::array<byte, 16>{};

// A 1D view of the block:
for (auto i = 0; i != 16; ++i)
    std::cout << block[i] << " ";

// A 2D view of the block:
for (auto i = 0; i != 4; ++i)
{
    for (auto j = 0; j != 4; ++j)
        std::cout << block[to_index<4, 4>(i, j)] << " ";
    std::cout << "\n";
}

2D arrays are almost never worth the hassle. But especially here, where you’re just jumping back and forth between a flat view and a 2D view. Removing the need to do that will massively simplify your code.

As for standard algorithms, they are superior to hand-rolled loops in at least three ways:

  1. Hand-rolled loops are too easy to fuck up. That actually happened in this code. (Well, technically, the fuck-up wasn’t in a loop, but rather a recursive function… but loops and recursive functions are basically interchangeable, so, yeah.)
  2. They are pretty much self documenting. When you see a naked for loop, you have literally no idea what it’s really doing until you sit down and reason through the code. But a standard algorithm? rotate() is pretty self-explanatory. transform(something, operation) is also pretty clear. And so on.
  3. They are rigorously tested, and optimized. You might be able to write a safer and/or faster version… but almost certainly not in the general case.

Plus, they offer intriguing possibilities for acceleration… not so much today (though, maybe!), but definitely in the future. I can’t think of any place where you might benefit from it offhand, but imagine you were doing an operation of 4 bytes like so:

// Naked loop version (bad)
for (auto i = 0; i != 4; ++i)
    b[i] = some_op(b[i]);

// Algorithm version
std::transform(b.begin(), b.end(), b.begin(), some_op);

Now, if it’s possible, a good compiler will probably be smart enough to unroll that loop and vectorize the operation, so that some_op() will be done to all 4 bytes in a single operation… at least a 4× speedup. (And the same would happen for the algorithm.) But why rely on the compiler’s good graces? Why not specify that you want the whole operation vectorized?

std::transform(std::execution::unseq, b.begin(), b.end(), b.begin(), some_op);

The difference is between silently praying that the compiler will vectorize the operation, and explicily asking the compiler to do that, while also documenting for readers of the code that you know it is safe and desirable for vectorization to happen here.

I wouldn’t bother doing this today, especially since the constrained algorithms don’t support execution policies yet. But it is a possibility for the future. Once you have everything written using standard algorithms, adding the execution policy later is trivial.

Another issue with your code is the interface. It’s not particularly user-friendly. You should consider how a user would want to do AES encryption/decryption. They’d want to set up an encrypter/decrypter with a key, then just blast bytes into it, getting encrypted bytes out. Something like this:

auto const key = aes_128_key{/* write key here, or load it from a file or whatever */};

auto encrypter = aes_128{key};  // key expansion happens here

auto const input = /* ... */;   // the input could be a byte array or vector,
                                // or an input stream from a file, or
                                // whatever.

auto output = /* ... */;    // the output could be a vector we back-insert
                            // into, or an output stream to a file, or
                            // whatever.

encrypter(std::ranges::begin(input), std::ranges::end(input), std::ranges::begin(output));
// or:
//  encryper(input, std::ranges::begin(output));

// Internally, the encrypter reads 128 bits at a time, padding if necessary.

That’s what a good interface looks like: Easy to use, hard to misuse.

Finally, look into testing, with a proper testing framework. NIST specifies a complete testing program for AES implementations… but you don’t need to go that far (unless you want to!). For an algorithm that just transforms data from one form to another—which would include encryption and decryption—all you really need is a set of pairs of input and expected output. Then it’s just:

TEST_CASE("algorithm test")
{
    auto const data = GENERATE(table<input_type, output_type>>({
        {input_1, expected_1},
        {input_2, expected_2},
        {input_3, expected_3},
        // ... and so on ...
    }));

    REQUIRE(algorithm(std::get<0>(data)) == std::get<1>(data));
}

Sure you can do more testing if your algorithm has interesting or special cases, or particular usage patterns—for example, you could do one set of tests with input that is exactly the size of a block, and another set with input that is smaller (to make sure that padding works), and other set with input that is larger (with various chaining strategies). But this is the basic pattern you’ll be using.

One last thing I would suggest, and that is parameterizing your implementation. AES comes in 3 flavours, after all, and the 3 flavours are more or less identical. It is possible to make a single implementation that can do AES 128, AES 192, and AES 256. You could make the actual version a template parameter, so that, using the interface I suggested above:

// To change to AES 192 or 256, simply change the template parameter (and
// write/load a different sized key, of course).
auto const key = aes<128>::key{/* key here, or load it from wherever */};
auto encrypter = aes{key};  // deduces aes<128>, from the key

auto const input = /* ... */;

auto output = /* ... */;

encrypter(std::ranges::begin(input), std::ranges::end(input), std::ranges::begin(output));

With some simple constraints (the key size/template parameter can only be 128, 192, or 256), and a few simple specialized traits (like, for example, see how to parameterize the number of rounds below), this should be pretty easy to make a single implementation for all versions of AES (or, indeed, for all versions of Rijndael, if you want to get crazy).

template <std::size_t KeySize>
    requires (KeySize == 128)
        or (KeySize == 192)
        or (KeySize == 256)
consteval auto _number_of_rounds() noexcept
{
    if constexpr (KeySize == 128)
        return 10uz;
    if constexpr (KeySize == 192)
        return 12uz;
    if constexpr (KeySize == 256)
        return 14uz;
}

template <std::size_t KeySize>
inline constexpr auto number_of_rounds = _number_of_rounds<KeySize>();

// Usage:
std::cout << number_of_rounds<128>; // prints 10

That’s it!

\$\endgroup\$
1
  • \$\begingroup\$ Wow, I wasn't expecting this level of feedback. As I was writing parts of this program I was thinking "there must be a better way to do this", I just didn't know what they were. Thank you for taking the time to read and review my code! I learnt a lot from your answer and my next program will, for sure, be better :) \$\endgroup\$ Commented May 24, 2022 at 8:19

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.