Add single use support

Author: will-v-pi

src/rp2_common/pico_aes/CMakeLists.txt

@@ -16,5 +16,9 @@ pico_mirrored_target_link_libraries(pico_aes INTERFACE
 function(pico_aes_configure_binary TARGET)
     pico_include_in_generated_section(${TARGET} post_data section_aes_scratch.incl)
     pico_include_in_generated_section(${TARGET} post_end section_aes_end.incl)
-    target_compile_definitions(${TARGET} PRIVATE PICO_AES_CONFIGURED=1)
+    target_compile_definitions(${TARGET} PRIVATE PICO_AES_CONFIGURED=1 PICO_AES_SINGLE_USE=0)
+endfunction()
+
+function(pico_aes_configure_single_use_binary TARGET)
+    target_compile_definitions(${TARGET} PRIVATE PICO_AES_CONFIGURED=1 PICO_AES_SINGLE_USE=1)
 endfunction()

src/rp2_common/pico_aes/aes.S

@@ -175,7 +175,11 @@ rcp_iequal_nodelay r1, r0
 @ Put workspace in the second scratch area
 @ The "a"=allocatable attribute (and possibly the %progbits attribute) are necessary to store the murmur3 constants,
 @ otherwise they may end up silently replaced with 0 or 0xffffffff
+#if PICO_AES_SINGLE_USE
+.section .scratch_y.aes,"aw",%progbits
+#else
 .section .aes_scratch_y.aes,"aw",%progbits
+#endif
 
 aes_workspace_start:
 

src/rp2_common/pico_aes/aes.c

@@ -38,9 +38,10 @@ static pico_aes_context_t aes_context = {};
 })
 
 
+#if PICO_AES_SINGLE_USE
 // The function lock_key() is called from pico_aes_decrypt_internal() after key initialisation is complete and before decryption begins.
 // That is a suitable point to lock the OTP area where key information is stored.
-__weak void pico_aes_lock_key() {
+void pico_aes_lock_key() {
     io_rw_32 *sw_lock = __get_opaque_ptr(&otp_hw->sw_lock[0]) + aes_context.otp_key_page;
 #if PICO_AES_HARDENING
     // prevent compiler from re-using sw_lock pointer
@@ -79,7 +80,15 @@ __weak void pico_aes_lock_key() {
 #endif
 }
 
-__weak void pico_aes_lock_all(void) {
+#else // not PICO_AES_SINGLE_USE
+void pico_aes_lock_key(void) {
+#if RC_COUNT
+    rcp_count_check_nodelay(31 + PICO_AES256_RCP_COUNT_DELTA);
+#endif
+}
+#endif // PICO_AES_SINGLE_USE
+
+void pico_aes_lock_all(void) {
     io_rw_32 *sw_lock = __get_opaque_ptr(&otp_hw->sw_lock[0]) + aes_context.otp_key_page;
 #if PICO_AES_HARDENING
     // we only actually need to lock page 2 but we lock and check 0, 1, 2 anyway
@@ -116,21 +125,23 @@ __weak void pico_aes_lock_all(void) {
 #endif
 }
 
-#include <stdio.h>
+
 int pico_aes_try_decrypt(uint8_t(*data)[16], size_t data_size, uint32_t otp_key_page, uint8_t* iv_public) {
     if (!bootrom_try_acquire_lock(BOOTROM_LOCK_SHA_256)) return PICO_ERROR_RESOURCE_IN_USE;
 
     aes_context.otp_key_page = otp_key_page;
 
     uint16_t* otp_data = (uint16_t*)OTP_DATA_GUARDED_BASE;
 
+#if !PICO_AES_SINGLE_USE
     // Restore aes_scratch_y from stored location in RAM_STORE
     extern uint32_t __aes_scratch_y_source__;
     extern uint32_t __aes_scratch_y_start__;
     extern uint32_t __aes_scratch_y_end__;
 
     uint32_t stored_words = (uint32_t)(&__aes_scratch_y_end__ - &__aes_scratch_y_start__);
     memcpy(&__aes_scratch_y_start__, &__aes_scratch_y_source__, stored_words * sizeof(uint32_t));
+#endif // !PICO_AES_SINGLE_USE
 
     pico_aes_decrypt_internal(
         (uint8_t*)&(otp_data[otp_key_page * 0x40]),

src/rp2_common/pico_aes/include/pico/aes.h

@@ -46,6 +46,8 @@
 
 int pico_aes_try_decrypt(uint8_t(*data)[16], size_t data_size, uint32_t otp_key_page, uint8_t* iv_public);
 
+void pico_aes_lock_all(void);
+
 #endif
 
 #endif

src/rp2_common/pico_standard_link/script_include/section_aes_end.incl

@@ -0,0 +1,7 @@
+/* Contains some assertions for use with pico_aes */
+
+SECTIONS
+{
+    ASSERT(chaff==0x20081000, "Chaff array must be located at 0x20081000")
+    ASSERT(__scratch_y_start__ == __scratch_y_end__, "Cannot use scratch_y when using pico_aes")
+}

src/rp2_common/pico_standard_link/script_include/section_aes_scratch.incl

@@ -0,0 +1,14 @@
+/* Defines the following symbols for use by code:
+    __aes_scratch_y_start__, __aes_scratch_y_end__, __aes_scratch_y_source__
+*/
+
+SECTIONS
+{
+    .aes_scratch_y : {
+        __aes_scratch_y_start__ = .;
+        *(.aes_scratch_y.*)
+        . = ALIGN(4);
+        __aes_scratch_y_end__ = .;
+    } > SCRATCH_Y AT> RAM_STORE
+    __aes_scratch_y_source__ = LOADADDR(.aes_scratch_y);
+}

test/pico_aes_test/BUILD.bazel

@@ -1,4 +1,3 @@
-
 # pico_aes is not supported in Bazel, as it requires encryption
 
 exports_files(

test/pico_aes_test/CMakeLists.txt

@@ -22,3 +22,24 @@ target_compile_definitions(pico_aes_test PRIVATE
 pico_aes_configure_binary(pico_aes_test)
 pico_sign_binary(pico_aes_test ${PICO_SDK_PATH}/tools/example_keys/private.pem)
 pico_add_extra_outputs(pico_aes_test)
+
+
+add_executable(pico_aes_test_single
+        pico_aes_test.c
+        )
+target_link_libraries(pico_aes_test_single
+        pico_stdlib
+        pico_aes
+        pico_rand
+)
+target_include_directories(pico_aes_test_single PRIVATE
+        ${CMAKE_CURRENT_LIST_DIR}
+)
+target_compile_definitions(pico_aes_test_single PRIVATE
+        PICO_USE_STACK_GUARDS=1
+        PICO_XOSC_STARTUP_DELAY_MULTIPLIER=64
+        PICO_STACK_SIZE=0x9e0 # maximum stack size that fits
+)
+pico_aes_configure_single_use_binary(pico_aes_test_single)
+pico_sign_binary(pico_aes_test_single ${PICO_SDK_PATH}/tools/example_keys/private.pem)
+pico_add_extra_outputs(pico_aes_test_single)

test/pico_aes_test/pico_aes_test.c

@@ -21,18 +21,25 @@
 #define BUFFER_WORD_SIZE (BUFFER_SIZE/sizeof(uint32_t))
 
 
-void pico_aes_lock_key(void) {
-#if RC_COUNT
-    rcp_count_check_nodelay(31 + PICO_AES256_RCP_COUNT_DELTA);
-#endif
-}
-void pico_aes_lock_all(void) {}
-
 int main() {
     stdio_init_all();
     printf("AES Test Starting\n");
 
     bool passing = true;
+#if PICO_AES_SINGLE_USE
+    uint32_t *data = malloc(BUFFER_SIZE);
+    hard_assert(data);
+    for (int i=0; i < BUFFER_WORD_SIZE; i++) data[i] = get_rand_32();
+    printf("Buffer start   %08x %08x %08x\n", data[0], data[1], data[2]);
+
+    uint32_t* iv_public = malloc(16);
+    hard_assert(iv_public);
+    for (int i=0; i < 16/sizeof(uint32_t); i++) iv_public[i] = get_rand_32();
+
+    pico_aes_try_decrypt((void*)data, BUFFER_SIZE, 29, (uint8_t*)iv_public);
+
+    printf("Buffer encrypt %08x %08x %08x\n", data[0], data[1], data[2]);
+#else
     for (int n=0; n < 10; n++) {
         uint32_t *data = malloc(BUFFER_SIZE);
         uint32_t *data_orig = malloc(BUFFER_SIZE);
@@ -65,6 +72,10 @@ int main() {
         if (!passing) break;
         else printf("run %d passed\n", n);
     }
+#endif
+
+    // Lock all keys now decryption is finished
+    pico_aes_lock_all();
 
     if (passing) printf("PASSED\n");
     else printf("FAILED\n");