Move memory training code into it's own stage (fusee-mtc)

sept/sept-secondary/src/main.c

@@ -90,6 +90,20 @@ static void exfiltrate_keys_and_reboot_if_needed(uint32_t version) {
     }
 }
 
+static void display_splash_screen(void) {
+    /* Draw splash. */
+    draw_splash((volatile uint32_t *)g_framebuffer);
+    
+    /* Turn on the backlight. */
+    display_backlight(true);
+    
+    /* Ensure the splash screen is displayed for at least one second. */
+    mdelay(1000);
+    
+    /* Turn off the backlight. */
+    display_backlight(false);
+}
+
 static void setup_env(void) {
     g_framebuffer = (void *)0xC0000000;
 
@@ -104,14 +118,7 @@ static void setup_env(void) {
 
     /* Set the framebuffer. */
     display_init_framebuffer(g_framebuffer);
-
-    /* Draw splash. */
-    draw_splash((volatile uint32_t *)g_framebuffer);
-
-    /* Turn on the backlight after initializing the lfb */
-    /* to avoid flickering. */
-    display_backlight(true);
-
+    
     /* Set up the exception handlers. */
     setup_exception_handlers();
 
@@ -122,9 +129,6 @@ static void setup_env(void) {
 static void cleanup_env(void) {
     /* Unmount the SD card. */
     unmount_sd();
-
-    /* Turn off the backlight. */
-    display_backlight(false);
     
     /* Terminate the display. */
     display_end();
@@ -160,9 +164,11 @@ int sept_main(uint32_t version) {
 
     /* Load the loader payload into DRAM. */
     load_stage2();
+    
+    /* Display the splash screen. */
+    display_splash_screen();
 
     /* Setup argument data. */
-    log_level = SCREEN_LOG_LEVEL_MANDATORY;
     stage2_path = stage2_get_program_path();
     strcpy(g_chainloader_arg_data, stage2_path);
     stage2_args = (stage2_args_t *)(g_chainloader_arg_data + strlen(stage2_path) + 1); /* May be unaligned. */

sept/sept-secondary/src/stage2.c

@@ -25,6 +25,42 @@ const char *stage2_get_program_path(void) {
     return g_stage2_path;
 }
 
+static bool run_mtc(const char *mtc_path, uintptr_t mtc_address) {
+    FILINFO info;
+    size_t size;
+    
+    /* Check if the MTC binary is present. */
+    if (f_stat(mtc_path, &info) != FR_OK) {
+        print(SCREEN_LOG_LEVEL_WARNING, "Stage2's MTC binary not found!\n");
+        return false;
+    }
+    
+    size = (size_t)info.fsize;
+    
+    /* Try to read the MTC binary. */
+    if (read_from_file((void *)mtc_address, size, mtc_path) != size) {
+        print(SCREEN_LOG_LEVEL_WARNING, "Failed to read stage2's MTC binary (%s)!\n", mtc_path);
+        return false;
+    }
+    
+    ScreenLogLevel mtc_log_level = log_get_log_level();
+    bool mtc_res = false;
+    int mtc_argc = 1;
+    char mtc_arg_data[CHAINLOADER_ARG_DATA_MAX_SIZE] = {0};
+    stage2_mtc_args_t *mtc_args = (stage2_mtc_args_t *)mtc_arg_data;
+
+    /* Setup argument data. */
+    memcpy(&mtc_args->log_level, &mtc_log_level, sizeof(mtc_log_level));
+    
+    /* Run the MTC binary. */
+    mtc_res = (((int (*)(int, void *))mtc_address)(mtc_argc, mtc_arg_data) == 0);
+    
+    /* Cleanup right away. */
+    memset((void *)mtc_address, 0, size);
+    
+    return mtc_res;
+}
+
 /* We get the luxury of assuming a constant filename/load address. */
 void load_stage2(void) {
     FILINFO info;
@@ -32,15 +68,22 @@ void load_stage2(void) {
     uintptr_t tmp_addr;
     stage2_config_t config = {
         .path = "sept/payload.bin",
+        .mtc_path = "atmosphere/fusee-mtc.bin",
         .load_address = 0xF0000000,
         .entrypoint = 0xF0000000,
     };
     
     print(SCREEN_LOG_LEVEL_DEBUG, "Stage 2 Config:\n");
     print(SCREEN_LOG_LEVEL_DEBUG | SCREEN_LOG_LEVEL_NO_PREFIX, "    File Path:    %s\n", config.path);
+    print(SCREEN_LOG_LEVEL_DEBUG | SCREEN_LOG_LEVEL_NO_PREFIX, "    MTC File Path:    %s\n", config.mtc_path);
     print(SCREEN_LOG_LEVEL_DEBUG | SCREEN_LOG_LEVEL_NO_PREFIX, "    Load Address: 0x%08x\n", config.load_address);
     print(SCREEN_LOG_LEVEL_DEBUG | SCREEN_LOG_LEVEL_NO_PREFIX, "    Entrypoint:   0x%p\n", config.entrypoint);
 
+    /* Run the MTC binary. */
+    if (!run_mtc(config.mtc_path, config.load_address)) {
+        print(SCREEN_LOG_LEVEL_WARNING, "DRAM training failed! Continuing with untrained DRAM.\n");
+    }
+    
     if (f_stat(config.path, &info) != FR_OK) {
         fatal_error("Failed to stat stage2 (%s)!\n", config.path);
     }
@@ -66,6 +109,7 @@ void load_stage2(void) {
         tmp_addr = config.load_address;
     }
     
+    /* Try to read stage2. */
     if (read_from_file((void *)tmp_addr, size, config.path) != size) {
         fatal_error("Failed to read stage2 (%s)!\n", config.path);
     }

sept/sept-secondary/src/stage2.h

@@ -32,6 +32,7 @@
 #define STAGE2_ARGC 2
 
 #define STAGE2_NAME_KEY "stage2_path"
+#define STAGE2_MTC_NAME_KEY "stage2_mtc_path"
 #define STAGE2_ADDRESS_KEY "stage2_addr"
 #define STAGE2_ENTRYPOINT_KEY "stage2_entrypoint"
 
@@ -39,6 +40,7 @@
 
 typedef struct {
     char path[0x100];
+    char mtc_path[0x100];
     uintptr_t load_address;
     uintptr_t entrypoint;
 } stage2_config_t;
@@ -49,7 +51,11 @@ typedef struct {
     char bct0[BCTO_MAX_SIZE];
 } stage2_args_t;
 
+typedef struct {
+    ScreenLogLevel log_level;
+} stage2_mtc_args_t;
+
 const char *stage2_get_program_path(void);
 void load_stage2(void);
 
-#endif
+#endif