Updated aes.c / aes.h

... thanks to @gemarcano

source/godmode.c

@@ -7,7 +7,7 @@
 #include "virtual.h"
 #include "image.h"
 
-#define VERSION "0.6.5"
+#define VERSION "0.6.6"
 
 #define N_PANES 2
 #define IMG_DRV "789I"

source/nand/aes.c

@@ -1,61 +1,66 @@
 /* original version by megazig */
 #include "aes.h"
 
-void setup_aeskeyX(u8 keyslot, void* keyx)
+//FIXME some things make assumptions about alignemnts!
+
+void setup_aeskeyX(uint8_t keyslot, void* keyx)
 {
-    u32 * _keyx = (u32*)keyx;
+    uint32_t * _keyx = (uint32_t*)keyx;
     *REG_AESKEYCNT = (*REG_AESKEYCNT >> 6 << 6) | keyslot | 0x80;
     if (keyslot > 3) {
-        *REG_AESCNT = (*REG_AESCNT | (AES_CNT_INPUT_ENDIAN | AES_CNT_INPUT_ORDER));
         *REG_AESKEYXFIFO = _keyx[0];
         *REG_AESKEYXFIFO = _keyx[1];
         *REG_AESKEYXFIFO = _keyx[2];
         *REG_AESKEYXFIFO = _keyx[3];
     } else {
-        vu32* reg_aeskeyx = REG_AESKEY0123 + (((0x30*keyslot) + 0x10)/4);
+        uint32_t old_aescnt = *REG_AESCNT;
+        volatile uint32_t* reg_aeskeyx = REG_AESKEY0123 + (((0x30u * keyslot) + 0x10u)/4u);
         *REG_AESCNT = (*REG_AESCNT & ~(AES_CNT_INPUT_ENDIAN | AES_CNT_INPUT_ORDER));
-        for (u32 i = 0; i < 4; i++)
+        for (uint32_t i = 0; i < 4u; i++)
             reg_aeskeyx[i] = _keyx[i];
+        *REG_AESCNT = old_aescnt;
     }
 }
 
-void setup_aeskeyY(u8 keyslot, void* keyy)
+void setup_aeskeyY(uint8_t keyslot, void* keyy)
 {
-    u32 * _keyy = (u32*)keyy;
+    uint32_t * _keyy = (uint32_t*)keyy;
     *REG_AESKEYCNT = (*REG_AESKEYCNT >> 6 << 6) | keyslot | 0x80;
     if (keyslot > 3) {
-        *REG_AESCNT = (*REG_AESCNT | (AES_CNT_INPUT_ENDIAN | AES_CNT_INPUT_ORDER));
         *REG_AESKEYYFIFO = _keyy[0];
         *REG_AESKEYYFIFO = _keyy[1];
         *REG_AESKEYYFIFO = _keyy[2];
         *REG_AESKEYYFIFO = _keyy[3];
     } else {
-        vu32* reg_aeskeyy = REG_AESKEY0123 + (((0x30*keyslot) + 0x20)/4);
+        uint32_t old_aescnt = *REG_AESCNT;
+        volatile uint32_t* reg_aeskeyy = REG_AESKEY0123 + (((0x30u * keyslot) + 0x20u)/4u);
         *REG_AESCNT = (*REG_AESCNT & ~(AES_CNT_INPUT_ENDIAN | AES_CNT_INPUT_ORDER));
-        for (u32 i = 0; i < 4; i++)
+        for (uint32_t i = 0; i < 4u; i++)
             reg_aeskeyy[i] = _keyy[i];
+        *REG_AESCNT = old_aescnt;
     }
 }
 
-void setup_aeskey(u8 keyslot, void* key)
+void setup_aeskey(uint8_t keyslot, void* key)
 {
-    u32 * _key = (u32*)key;
+    uint32_t * _key = (uint32_t*)key;
     *REG_AESKEYCNT = (*REG_AESKEYCNT >> 6 << 6) | keyslot | 0x80;
     if (keyslot > 3) {
-        *REG_AESCNT = (*REG_AESCNT | (AES_CNT_INPUT_ENDIAN | AES_CNT_INPUT_ORDER));
         *REG_AESKEYFIFO = _key[0];
         *REG_AESKEYFIFO = _key[1];
         *REG_AESKEYFIFO = _key[2];
         *REG_AESKEYFIFO = _key[3];
     } else {
-        vu32* reg_aeskey = REG_AESKEY0123 + ((0x30*keyslot)/4);
+        uint32_t old_aescnt = *REG_AESCNT;
+        volatile uint32_t* reg_aeskey = REG_AESKEY0123 + ((0x30u * keyslot)/4u);
         *REG_AESCNT = (*REG_AESCNT & ~(AES_CNT_INPUT_ENDIAN | AES_CNT_INPUT_ORDER));
-        for (u32 i = 0; i < 4; i++)
+        for (uint32_t i = 0; i < 4u; i++)
             reg_aeskey[i] = _key[i];
+        *REG_AESCNT = old_aescnt;
     }
 }
 
-void use_aeskey(u32 keyno)
+void use_aeskey(uint32_t keyno)
 {
     if (keyno > 0x3F)
         return;
@@ -65,7 +70,7 @@ void use_aeskey(u32 keyno)
 
 void set_ctr(void* iv)
 {
-    u32 * _iv = (u32*)iv;
+    uint32_t * _iv = (uint32_t*)iv;
     *REG_AESCNT = (*REG_AESCNT) | AES_CNT_INPUT_ENDIAN | AES_CNT_INPUT_ORDER;
     *(REG_AESCTR + 0) = _iv[3];
     *(REG_AESCTR + 1) = _iv[2];
@@ -73,15 +78,15 @@ void set_ctr(void* iv)
     *(REG_AESCTR + 3) = _iv[0];
 }
 
-void add_ctr(void* ctr, u32 carry)
+void add_ctr(void* ctr, uint32_t carry)
 {
-    u32 counter[4];
+    uint32_t counter[4];
-    u8 *outctr = (u8 *) ctr;
+    uint8_t *outctr = (uint8_t *) ctr;
-    u32 sum;
+    uint32_t sum;
     int32_t i;
 
     for(i = 0; i < 4; i++) {
-        counter[i] = (outctr[i*4+0]<<24) | (outctr[i*4+1]<<16) | (outctr[i*4+2]<<8) | (outctr[i*4+3]<<0);
+        counter[i] = ((uint32_t)outctr[i*4+0]<<24) | ((uint32_t)outctr[i*4+1]<<16) | ((uint32_t)outctr[i*4+2]<<8) | ((uint32_t)outctr[i*4+3]<<0);
     }
 
     for(i=3; i>=0; i--)
@@ -105,10 +110,29 @@ void add_ctr(void* ctr, u32 carry)
     }
 }
 
-void aes_decrypt(void* inbuf, void* outbuf, size_t size, u32 mode)
+void ctr_decrypt(void *inbuf, void *outbuf, size_t size, uint32_t mode, uint8_t *ctr)
 {
-    u32 in  = (u32)inbuf;
+    size_t blocks_left = size;
-    u32 out = (u32)outbuf;
+    size_t blocks;
+    uint8_t *in  = inbuf;
+    uint8_t *out = outbuf;
+
+    while (blocks_left)
+    {
+        set_ctr(ctr);
+        blocks = (blocks_left >= 0xFFFF) ? 0xFFFF : blocks_left;
+        aes_decrypt(in, out, blocks, mode);
+        add_ctr(ctr, blocks);
+        in += blocks * AES_BLOCK_SIZE;
+        out += blocks * AES_BLOCK_SIZE;
+        blocks_left -= blocks;
+    }
+}
+
+void aes_decrypt(void* inbuf, void* outbuf, size_t size, uint32_t mode)
+{
+    uint8_t *in  = inbuf;
+    uint8_t *out = outbuf;
     size_t block_count = size;
     size_t blocks;
     while (block_count != 0)
@@ -127,62 +151,115 @@ void aes_decrypt(void* inbuf, void* outbuf, size_t size, u32 mode)
     }
 }
 
+void aes_cmac(void* inbuf, void* outbuf, size_t size)
+{
+    // only works for full blocks
+    uint32_t zeroes[4] = { 0 };
+    uint32_t xorpad[4] = { 0 };
+    uint32_t mode = AES_CBC_ENCRYPT_MODE | AES_CNT_INPUT_ORDER | AES_CNT_OUTPUT_ORDER |
+        AES_CNT_INPUT_ENDIAN | AES_CNT_OUTPUT_ENDIAN;
+    uint32_t* out = (uint32_t*) outbuf;
+    uint32_t* in  = (uint32_t*) inbuf;
+        
+    // create xorpad for last block
+    set_ctr(zeroes);
+    aes_decrypt(xorpad, xorpad, 1, mode);
+    for (uint32_t i = 0; i < 4; i++) {
+        if (i && (xorpad[i] >> 31))
+            xorpad[i-i] |= 1;
+        xorpad[i] <<= 1;
+    }   
+    
+    // process blocks
+    for (uint32_t i = 0; i < 4; i++)
+        out[i] = 0;
+    while (size-- > 0) {
+        for (uint32_t i = 0; i < 4; i++)
+			out[i] ^= *(in++);
+        if (!size) { // last block
+            for (uint32_t i = 0; i < 4; i++)
+                out[i] ^= xorpad[i];
+        }
+        set_ctr(zeroes);
+        aes_decrypt(out, out, 1, mode);
+    }
+}
+
 void aes_fifos(void* inbuf, void* outbuf, size_t blocks)
 {
-    u32 in  = (u32)inbuf;
+    if (!inbuf || !outbuf) return;
-    if (!in) return;
+
+    uint8_t *in = inbuf;
+    uint8_t *out = outbuf;
 
-    u32 out = (u32)outbuf;
     size_t curblock = 0;
     while (curblock != blocks)
     {
         while (aescnt_checkwrite());
 
-        u32 ii = 0;
+        size_t blocks_to_read = blocks - curblock > 4 ? 4 : blocks - curblock;
-        for (ii = in; ii != in + AES_BLOCK_SIZE; ii += 4)
+
+        for (size_t wblocks = 0; wblocks < blocks_to_read; ++wblocks)
+        for (uint8_t *ii = in + AES_BLOCK_SIZE * wblocks; ii != in + (AES_BLOCK_SIZE * (wblocks + 1)); ii += 4)
         {
-            set_aeswrfifo( *(u32*)(ii) );
+            uint32_t data = ii[0];
+            data |= (uint32_t)(ii[1]) << 8;
+            data |= (uint32_t)(ii[2]) << 16;
+            data |= (uint32_t)(ii[3]) << 24;
+            set_aeswrfifo(data);
         }
+
         if (out)
         {
-            while (aescnt_checkread()) ;
+            for (size_t rblocks = 0; rblocks < blocks_to_read; ++rblocks)
-            for (ii = out; ii != out + AES_BLOCK_SIZE; ii += 4)
             {
-                *(u32*)ii = read_aesrdfifo();
+                while (aescnt_checkread()) ;
+                for (uint8_t *ii = out + AES_BLOCK_SIZE * rblocks; ii != out + (AES_BLOCK_SIZE * (rblocks + 1)); ii += 4)
+                {
+                    uint32_t data = read_aesrdfifo();
+                    ii[0] = data;
+                    ii[1] = data >> 8;
+                    ii[2] = data >> 16;
+                    ii[3] = data >> 24;
                 }
             }
-        curblock++;
-    }
         }
 
-void set_aeswrfifo(u32 value)
+        in += blocks_to_read * AES_BLOCK_SIZE;
+        out += blocks_to_read * AES_BLOCK_SIZE;
+        curblock += blocks_to_read;
+    }
+}
+
+void set_aeswrfifo(uint32_t value)
 {
     *REG_AESWRFIFO = value;
 }
 
-u32 read_aesrdfifo(void)
+uint32_t read_aesrdfifo(void)
 {
     return *REG_AESRDFIFO;
 }
 
-u32 aes_getwritecount()
+uint32_t aes_getwritecount()
 {
     return *REG_AESCNT & 0x1F;
 }
 
-u32 aes_getreadcount()
+uint32_t aes_getreadcount()
 {
     return (*REG_AESCNT >> 5) & 0x1F;
 }
 
-u32 aescnt_checkwrite()
+uint32_t aescnt_checkwrite()
 {
     size_t ret = aes_getwritecount();
     return (ret > 0xF);
 }
 
-u32 aescnt_checkread()
+uint32_t aescnt_checkread()
 {
     size_t ret = aes_getreadcount();
     return (ret <= 3);
 }
+

source/nand/aes.h

@@ -1,37 +1,43 @@
 #pragma once
 
-#include "common.h"
+#include <stdint.h>
+#include <stddef.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
 
 #define AES_BLOCK_SIZE 0x10
 
-#define AES_CCM_DECRYPT_MODE (0 << 27)
+#define AES_CCM_DECRYPT_MODE (0u << 27)
-#define AES_CCM_ENCRYPT_MODE (1 << 27)
+#define AES_CCM_ENCRYPT_MODE (1u << 27)
-#define AES_CTR_MODE         (2 << 27)
+#define AES_CTR_MODE         (2u << 27)
-#define AES_CBC_DECRYPT_MODE (4 << 27)
+#define AES_CBC_DECRYPT_MODE (4u << 27)
-#define AES_CBC_ENCRYPT_MODE (5 << 27)
+#define AES_CBC_ENCRYPT_MODE (5u << 27)
-#define AES_ECB_DECRYPT_MODE (6 << 27)
+#define AES_ECB_DECRYPT_MODE (6u << 27)
-#define AES_ECB_ENCRYPT_MODE (7 << 27)
+#define AES_ECB_ENCRYPT_MODE (7u << 27)
 
-#define REG_AESCNT     ((volatile u32*)0x10009000)
+#define REG_AESCNT     ((volatile uint32_t*)0x10009000)
-#define REG_AESBLKCNT  ((volatile u32*)0x10009004)
+#define REG_AESBLKCNT  ((volatile uint32_t*)0x10009004)
-#define REG_AESWRFIFO  ((volatile u32*)0x10009008)
+#define REG_AESWRFIFO  ((volatile uint32_t*)0x10009008)
-#define REG_AESRDFIFO  ((volatile u32*)0x1000900C)
+#define REG_AESRDFIFO  ((volatile uint32_t*)0x1000900C)
-#define REG_AESKEYSEL  ((volatile u8 *)0x10009010)
+#define REG_AESKEYSEL  ((volatile uint8_t *)0x10009010)
-#define REG_AESKEYCNT  ((volatile u8 *)0x10009011)
+#define REG_AESKEYCNT  ((volatile uint8_t *)0x10009011)
-#define REG_AESCTR     ((volatile u32*)0x10009020)
+#define REG_AESCTR     ((volatile uint32_t*)0x10009020)
-#define REG_AESKEYFIFO ((volatile u32*)0x10009100)
+#define REG_AESKEYFIFO ((volatile uint32_t*)0x10009100)
-#define REG_AESKEYXFIFO ((volatile u32*)0x10009104)
+#define REG_AESKEYXFIFO ((volatile uint32_t*)0x10009104)
-#define REG_AESKEYYFIFO ((volatile u32*)0x10009108)
+#define REG_AESKEYYFIFO ((volatile uint32_t*)0x10009108)
+#define REG_AESMAC      ((volatile uint32_t*)0x10009030)
 // see https://www.3dbrew.org/wiki/AES_Registers#AES_KEY0.2F1.2F2.2F3
-#define REG_AESKEY0123 ((volatile u32*)0x10009040)
+#define REG_AESKEY0123 ((volatile uint32_t*)0x10009040)
 
-#define AES_CNT_START         0x80000000
+#define AES_CNT_START         0x80000000u
-#define AES_CNT_INPUT_ORDER   0x02000000
+#define AES_CNT_INPUT_ORDER   0x02000000u
-#define AES_CNT_OUTPUT_ORDER  0x01000000
+#define AES_CNT_OUTPUT_ORDER  0x01000000u
-#define AES_CNT_INPUT_ENDIAN  0x00800000
+#define AES_CNT_INPUT_ENDIAN  0x00800000u
-#define AES_CNT_OUTPUT_ENDIAN 0x00400000
+#define AES_CNT_OUTPUT_ENDIAN 0x00400000u
-#define AES_CNT_FLUSH_READ    0x00000800
+#define AES_CNT_FLUSH_READ    0x00000800u
-#define AES_CNT_FLUSH_WRITE   0x00000400
+#define AES_CNT_FLUSH_WRITE   0x00000400u
 
 #define AES_CNT_CTRNAND_MODE (AES_CTR_MODE | AES_CNT_INPUT_ORDER | AES_CNT_OUTPUT_ORDER | AES_CNT_INPUT_ENDIAN | AES_CNT_OUTPUT_ENDIAN)
 #define AES_CNT_TWLNAND_MODE AES_CTR_MODE
@@ -40,18 +46,23 @@
 #define AES_CNT_ECB_DECRYPT_MODE (AES_ECB_DECRYPT_MODE | AES_CNT_INPUT_ORDER | AES_CNT_OUTPUT_ORDER | AES_CNT_INPUT_ENDIAN | AES_CNT_OUTPUT_ENDIAN)
 #define AES_CNT_ECB_ENCRYPT_MODE (AES_ECB_ENCRYPT_MODE | AES_CNT_INPUT_ORDER | AES_CNT_OUTPUT_ORDER | AES_CNT_INPUT_ENDIAN | AES_CNT_OUTPUT_ENDIAN)
 
-
+void setup_aeskeyX(uint8_t keyslot, void* keyx);
-void setup_aeskeyX(u8 keyslot, void* keyx);
+void setup_aeskeyY(uint8_t keyslot, void* keyy);
-void setup_aeskeyY(u8 keyslot, void* keyy);
+void setup_aeskey(uint8_t keyslot, void* keyy);
-void setup_aeskey(u8 keyslot, void* keyy);
+void use_aeskey(uint32_t keyno);
-void use_aeskey(u32 keyno);
 void set_ctr(void* iv);
-void add_ctr(void* ctr, u32 carry);
+void add_ctr(void* ctr, uint32_t carry);
-void aes_decrypt(void* inbuf, void* outbuf, size_t size, u32 mode);
+void aes_decrypt(void* inbuf, void* outbuf, size_t size, uint32_t mode);
+void ctr_decrypt(void* inbuf, void* outbuf, size_t size, uint32_t mode, uint8_t *ctr);
 void aes_fifos(void* inbuf, void* outbuf, size_t blocks);
-void set_aeswrfifo(u32 value);
+void set_aeswrfifo(uint32_t value);
-u32 read_aesrdfifo(void);
+uint32_t read_aesrdfifo(void);
-u32 aes_getwritecount();
+uint32_t aes_getwritecount(void);
-u32 aes_getreadcount();
+uint32_t aes_getreadcount(void);
-u32 aescnt_checkwrite();
+uint32_t aescnt_checkwrite(void);
-u32 aescnt_checkread();
+uint32_t aescnt_checkread(void);
+
+#ifdef __cplusplus
+}
+#endif
+

source/nand/nand.c

@@ -247,6 +247,7 @@ void CryptNand(u8* buffer, u32 sector, u32 count, u32 keyslot)
 {
     u32 mode = (sector >= (0x0B100000 / 0x200)) ? AES_CNT_CTRNAND_MODE : AES_CNT_TWLNAND_MODE;
     u8 ctr[16] __attribute__((aligned(32)));
+    u32 blocks = count * (0x200 / 0x10);
     
     // copy NAND CTR and increment it
     memcpy(ctr, (sector >= (0x0B100000 / 0x200)) ? CtrNandCtr : TwlNandCtr, 16);
@@ -254,13 +255,7 @@ void CryptNand(u8* buffer, u32 sector, u32 count, u32 keyslot)
     
     // decrypt the data
     use_aeskey(keyslot);
-    for (u32 s = 0; s < count; s++) {
+    ctr_decrypt((void*) buffer, (void*) buffer, blocks, mode, ctr);
-        for (u32 b = 0x0; b < 0x200; b += 0x10, buffer += 0x10) {
-            set_ctr(ctr);
-            aes_decrypt((void*) buffer, (void*) buffer, 1, mode);
-            add_ctr(ctr, 0x1);
-        }
-    }
 }
 
 void CryptSector0x96(u8* buffer, bool encrypt)