Add USE_MBEDTLS variant

This variant runs from xip_sram, and uses the MbedTLS AES code instead. This is much faster, but not secure against side channel attacks.

Currently no way to enable it, but that can be added if it's going to be useful

Author: will-v-pi

enc_bootloader/CMakeLists.txt

@@ -33,19 +33,39 @@ if (NOT USE_PRECOMPILED)
     # Encrypted Bootloader
     add_executable(enc_bootloader
         enc_bootloader.c
-        aes.S
     )
 
     target_link_libraries(enc_bootloader
         pico_stdlib
     )
 
+    if (USE_MBEDTLS)
+        target_sources(enc_bootloader PRIVATE mbedtls_aes.c)
+
+        target_link_libraries(enc_bootloader pico_mbedtls)
+
+        target_compile_definitions(enc_bootloader PRIVATE
+            PICO_STACK_SIZE=0x800
+            # 0x20080000 -> 0x20081000 doesn't overlap the stack
+            ROM_CHAIN_WORKSPACE=0x20080000)
+
+        target_include_directories(enc_bootloader PRIVATE ${CMAKE_CURRENT_LIST_DIR})
+
+        pico_set_linker_script(enc_bootloader ${CMAKE_CURRENT_LIST_DIR}/memmap_mbedtls.ld)
+    else()
+        target_sources(enc_bootloader PRIVATE aes.S)
+
+        target_compile_definitions(enc_bootloader PRIVATE
+            PICO_STACK_SIZE=0x180
+            # AES Code & workspace from 0x20080044 -> 0x20081604, so 0x20080200 -> 0x20081200 is inside that
+            ROM_CHAIN_WORKSPACE=0x20080200)
+
+        pico_set_linker_script(enc_bootloader ${CMAKE_CURRENT_LIST_DIR}/memmap_enc_bootloader.ld)
+    endif()
+
     target_compile_definitions(enc_bootloader PRIVATE
-        # increase startup delay for stability
-        PICO_XOSC_STARTUP_DELAY_MULTIPLIER=64
         # use stack guards, as AES variables are written near the stack
         PICO_USE_STACK_GUARDS=1
-        PICO_STACK_SIZE=0x180
         # The following are to reduce the size of the binary
         PICO_NO_PROGRAM_INFO=1
         # No spinlocks used
@@ -66,7 +86,6 @@ if (NOT USE_PRECOMPILED)
     pico_minimize_runtime(enc_bootloader)
 
     pico_set_binary_type(enc_bootloader no_flash)
-    pico_set_linker_script(enc_bootloader ${CMAKE_CURRENT_LIST_DIR}/memmap_enc_bootloader.ld)
     pico_add_dis_output(enc_bootloader)
 else()
     project(enc_bootloader C CXX ASM)

enc_bootloader/enc_bootloader.c

@@ -192,7 +192,7 @@ int main() {
 
     // Chain into decrypted image
     rom_chain_image(
-        (uint8_t*)0x20080200, // AES Code & workspace from 0x20080030 -> 0x20081500
+        (uint8_t*)ROM_CHAIN_WORKSPACE,
         4 * 1024,
         data_start_addr,
         data_size

enc_bootloader/mbedtls_aes.c

@@ -0,0 +1,73 @@
+#include <mbedtls/aes.h>
+#include "pico/stdlib.h"
+
+extern void lock_key();
+
+int mb_aes_crypt_ctr_xor(mbedtls_aes_context *ctx,
+    size_t length,
+    unsigned char iv0[16],
+    unsigned char nonce_xor[16],
+    unsigned char stream_block[16],
+    const unsigned char *input,
+    unsigned char *output)
+{
+    int c;
+    int ret = 0;
+    size_t n = 0;
+    uint32_t counter = 0;
+
+    assert(length == (uint32_t)length);
+
+    while (length--) {
+        if (n == 0) {
+            for (int i = 16; i > 0; i--) {
+                nonce_xor[i-1] = iv0[i-1];
+                if (i - (int)(16 - sizeof(counter)) > (int)0) {
+                    nonce_xor[i-1] ^= (unsigned char)(counter >> ((16-i)*8));
+                }
+            }
+
+            ret = mbedtls_aes_crypt_ecb(ctx, MBEDTLS_AES_ENCRYPT, nonce_xor, stream_block);
+            if (ret != 0) {
+                break;
+            }
+            counter++;
+        }
+        c = *input++;
+        *output++ = (unsigned char) (c ^ stream_block[n]);
+
+        n = (n + 1) & 0x0F;
+    }
+
+    return ret;
+}
+
+void decrypt(uint8_t* key4way, uint8_t* IV_OTPsalt, uint8_t* IV_public, uint8_t(*buf)[16], int nblk) {
+    mbedtls_aes_context aes;
+
+    uint32_t aes_key[8];
+    uint32_t* key4waywords = (uint32_t*)key4way;
+    // Key is stored as a 4-way share of each word, ie X[0] = A[0] ^ B[0] ^ C[0] ^ D[0], stored as A[0], B[0], C[0], D[0]
+    for (int i=0; i < count_of(aes_key); i++) {
+        aes_key[i] = key4waywords[i*4]
+                         ^ key4waywords[i*4 + 1]
+                         ^ key4waywords[i*4 + 2]
+                         ^ key4waywords[i*4 + 3];
+    }
+
+    uint8_t iv[16];
+    for (int i=0; i < sizeof(iv); i++) {
+        iv[i] = IV_OTPsalt[i] ^ IV_public[i];
+    }
+
+    int len = nblk * 16;
+
+    mbedtls_aes_setkey_enc(&aes, (uint8_t*)aes_key, 256);
+
+    lock_key();
+
+    uint8_t xor_working_block[16] = {0};
+    uint8_t stream_block[16] = {0};
+    size_t nc_off = 0;
+    mb_aes_crypt_ctr_xor(&aes, len, (uint8_t*)iv, xor_working_block, stream_block, (uint8_t*)buf, (uint8_t*)buf);
+}

enc_bootloader/mbedtls_config.h

@@ -0,0 +1,10 @@
+#ifndef _MBEDTLS_CONFIG_H
+#define _MBEDTLS_CONFIG_H
+
+#define MBEDTLS_HAVE_ASM
+#define MBEDTLS_AES_C
+#define MBEDTLS_AES_ROM_TABLES
+// #define MBEDTLS_AES_FEWER_TABLES
+#define MBEDTLS_CIPHER_MODE_CTR
+
+#endif

enc_bootloader/memmap_enc_bootloader.ld

@@ -257,4 +257,8 @@ SECTIONS
 
     /* todo assert on extra code */
     ASSERT(chaff==0x20081000, "Chaff array must be located at 0x20081000")
+
+    /* Provide symbols for picotool */
+    __enc_bootloader_start = ORIGIN(RAM_START);
+    __enc_bootloader_end = ORIGIN(RAM) + LENGTH(RAM);
 }