exo: implement mariko se/tzram context save

exosphere/program/source/boot/secmon_main.cpp

@@ -42,7 +42,7 @@ namespace ams::secmon {
            i2c::SetRegisterAddress(i2c::Port_5, MemoryRegionVirtualDeviceI2c5.GetAddress());
         pinmux::SetRegisterAddress(MemoryRegionVirtualDeviceApbMisc.GetAddress(), MemoryRegionVirtualDeviceGpio.GetAddress());
            pmc::SetRegisterAddress(MemoryRegionVirtualDevicePmc.GetAddress());
-            se::SetRegisterAddress(MemoryRegionVirtualDeviceSecurityEngine.GetAddress());
+            se::SetRegisterAddress(MemoryRegionVirtualDeviceSecurityEngine.GetAddress(), MemoryRegionVirtualDeviceSecurityEngine2.GetAddress());
           uart::SetRegisterAddress(MemoryRegionVirtualDeviceUart.GetAddress());
            wdt::SetRegisterAddress(MemoryRegionVirtualDeviceTimer.GetAddress());
           util::SetRegisterAddress(MemoryRegionVirtualDeviceTimer.GetAddress());

exosphere/program/source/smc/secmon_smc_power_management.cpp

@@ -325,7 +325,11 @@ namespace ams::secmon::smc {
         }
 
         void SaveSecureContextForMariko() {
-            /* TODO: Implement this when adding ams-on-mariko support. */
+            /* Save security engine context to TZRAM SE carveout (inaccessible to cpu). */
+            se::SaveContextAutomatic();
+
+            /* Save TZRAM to shadow-TZRAM in always-on power domain. */
+            se::SaveTzramAutomatic();
         }
 
         void SaveSecureContext() {

libraries/libexosphere/include/exosphere/se/se_management.hpp

@@ -18,7 +18,7 @@
 
 namespace ams::se {
 
-    void SetRegisterAddress(uintptr_t address);
+    void SetRegisterAddress(uintptr_t address, uintptr_t address2);
 
     void Initialize();
 

libraries/libexosphere/include/exosphere/se/se_suspend.hpp

@@ -53,4 +53,8 @@ namespace ams::se {
     bool ValidateStickyBits(const StickyBits &bits);
     void SaveContext(Context *dst);
 
+    void ConfigureAutomaticContextSave();
+    void SaveContextAutomatic();
+    void SaveTzramAutomatic();
+
 }

libraries/libexosphere/source/se/se_execute.hpp

@@ -19,6 +19,7 @@
 namespace ams::se {
 
     volatile SecurityEngineRegisters *GetRegisters();
+    volatile SecurityEngineRegisters *GetRegisters2();
 
     void ExecuteOperation(volatile SecurityEngineRegisters *SE, SE_OPERATION_OP op, void *dst, size_t dst_size, const void *src, size_t src_size);
     void ExecuteOperationSingleBlock(volatile SecurityEngineRegisters *SE, void *dst, size_t dst_size, const void *src, size_t src_size);

libraries/libexosphere/source/se/se_management.cpp

@@ -21,6 +21,7 @@ namespace ams::se {
     namespace {
 
         constinit uintptr_t g_register_address  = secmon::MemoryRegionPhysicalDeviceSecurityEngine.GetAddress();
+        constinit uintptr_t g_register2_address = secmon::MemoryRegionPhysicalDeviceSecurityEngine2.GetAddress();
         constinit DoneHandler g_done_handler   = nullptr;
 
     }
@@ -29,8 +30,13 @@ namespace ams::se {
         return reinterpret_cast<volatile SecurityEngineRegisters *>(g_register_address);
     }
 
-    void SetRegisterAddress(uintptr_t address) {
+    volatile SecurityEngineRegisters *GetRegisters2() {
+        return reinterpret_cast<volatile SecurityEngineRegisters *>(g_register2_address);
+    }
+
+    void SetRegisterAddress(uintptr_t address, uintptr_t address2) {
         g_register_address  = address;
+        g_register2_address = address2;
     }
 
     void Initialize() {

libraries/libexosphere/source/se/se_registers.hpp

@@ -32,7 +32,8 @@ namespace ams::se {
         u32 SE_OUT_CUR_LL_ID;
         u32 SE_HASH_RESULT[0x10];
         u32 SE_CTX_SAVE_CONFIG;
-        u32 _0x74[0x63];
+        u32 SE_CTX_SAVE_AUTO;
+        u32 _0x78[0x62];
         u32 SE_SHA_CONFIG;
         u32 SE_SHA_MSG_LENGTH[0x4];
         u32 SE_SHA_MSG_LEFT[0x4];
@@ -61,7 +62,9 @@ namespace ams::se {
         u32 SE_RSA_KEYTABLE_ADDR;
         u32 SE_RSA_KEYTABLE_DATA;
         u32 SE_RSA_OUTPUT[0x40];
-        u32 _0x528[0xB6];
+        u32 _0x528[0x6];
+        u32 SE_TZRAM_OPERATION;
+        u32 _0x544[0xAF];
         u32 SE_STATUS;
         u32 SE_ERR_STATUS;
         u32 SE_MISC;
@@ -100,6 +103,7 @@ namespace ams::se {
 
     /* SE_STATUS. */
     DEFINE_SE_REG_TWO_BIT_ENUM(STATUS_STATE, 0, IDLE, BUSY, WAIT_OUT, WAIT_IN);
+    DEFINE_SE_REG_BIT_ENUM(STATUS_MEM_INTERFACE, 2, IDLE, BUSY);
 
     /* SE_SECURITY */
     DEFINE_SE_REG_BIT_ENUM(SECURITY_HARD_SETTING,    0,  SECURE, NONSECURE);
@@ -111,6 +115,12 @@ namespace ams::se {
     DEFINE_SE_REG(TZRAM_SETTING, 0, BITSIZEOF(u32));
     constexpr inline u32 SE_TZRAM_SETTING_SECURE = 0;
 
+    /* SE_TZRAM_OPERATION */
+    DEFINE_SE_REG_BIT_ENUM(TZRAM_OPERATION_REQ,  0, IDLE, INITIATE);
+    DEFINE_SE_REG_BIT_ENUM(TZRAM_OPERATION_MODE, 1, SAVE, RESTORE);
+    DEFINE_SE_REG_BIT_ENUM(TZRAM_OPERATION_BUSY, 2, NO, YES);
+    DEFINE_SE_REG(TZRAM_OPERATION_CURR_ADDR, 16, 16);
+
     /* SE_OPERATION */
     DEFINE_SE_REG_THREE_BIT_ENUM(OPERATION_OP, 0, ABORT, START, RESTART_OUT, CTX_SAVE, RESTART_IN, RESERVED_5, RESERVED_6, RESERVED_7);
 
@@ -168,6 +178,11 @@ namespace ams::se {
     DEFINE_SE_REG_BIT_ENUM(CTX_SAVE_CONFIG_STICKY_WORD_QUAD, 24, WORDS_0_3, WORDS_4_7);
     DEFINE_SE_REG_THREE_BIT_ENUM(CTX_SAVE_CONFIG_SRC, 29, STICKY_BITS, RSA_KEYTABLE, AES_KEYTABLE, PKA1_STICKY_BITS, MEM, RESERVED5, SRK, PKA1_KEYTABLE);
 
+    /* SE_CTX_SAVE_AUTO */
+    DEFINE_SE_REG_BIT_ENUM(CTX_SAVE_AUTO_ENABLE, 0, NO, YES);
+    DEFINE_SE_REG_BIT_ENUM(CTX_SAVE_AUTO_LOCK, 8, NO, YES);
+    DEFINE_SE_REG(CTX_SAVE_AUTO_CURR_CNT, 16, 10);
+
     /* SE_SHA_CONFIG */
     DEFINE_SE_REG(SHA_CONFIG_HW_INIT_HASH, 0, 1);
 

libraries/libexosphere/source/se/se_suspend.cpp

@@ -20,6 +20,10 @@ namespace ams::se {
 
     namespace {
 
+        constexpr inline size_t SE1ContextSaveOperationCount = 133;
+        constexpr inline size_t SE2ContextSaveOperationCount = 646;
+        static_assert(((SE1ContextSaveOperationCount - 2) + 1) * se::AesBlockSize == sizeof(se::Context));
+
         constinit const u8 FixedPattern[AesBlockSize] = {
             0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F
         };
@@ -64,6 +68,44 @@ namespace ams::se {
             ExecuteContextSaveOperation(SE, dst, AesBlockSize, nullptr, 0);
         }
 
+        void ConfigureForAutomaticContextSave(volatile SecurityEngineRegisters *SE) {
+            /* Configure the engine to do RNG encryption. */
+            reg::Write(SE->SE_CONFIG,  SE_REG_BITS_ENUM(CONFIG_ENC_MODE, AESMODE_KEY128),
+                                       SE_REG_BITS_ENUM(CONFIG_DEC_MODE, AESMODE_KEY128),
+                                       SE_REG_BITS_ENUM(CONFIG_ENC_ALG,             RNG),
+                                       SE_REG_BITS_ENUM(CONFIG_DEC_ALG,             NOP),
+                                       SE_REG_BITS_ENUM(CONFIG_DST,              MEMORY));
+
+            reg::Write(SE->SE_CRYPTO_CONFIG, SE_REG_BITS_ENUM (CRYPTO_CONFIG_MEMIF,              AHB),
+                                             SE_REG_BITS_VALUE(CRYPTO_CONFIG_CTR_CNTN,             0),
+                                             SE_REG_BITS_ENUM (CRYPTO_CONFIG_KEYSCH_BYPASS,  DISABLE),
+                                             SE_REG_BITS_ENUM (CRYPTO_CONFIG_CORE_SEL,       ENCRYPT),
+                                             SE_REG_BITS_ENUM (CRYPTO_CONFIG_IV_SELECT,     ORIGINAL),
+                                             SE_REG_BITS_ENUM (CRYPTO_CONFIG_VCTRAM_SEL,      MEMORY),
+                                             SE_REG_BITS_ENUM (CRYPTO_CONFIG_INPUT_SEL,       RANDOM),
+                                             SE_REG_BITS_ENUM (CRYPTO_CONFIG_XOR_POS,         BYPASS),
+                                             SE_REG_BITS_ENUM (CRYPTO_CONFIG_HASH_ENB,       DISABLE));
+        }
+
+        void WaitAutomaticContextSaveDone(volatile SecurityEngineRegisters *SE) {
+            /* Wait for operation. */
+            while (!reg::HasValue(SE->SE_INT_STATUS, SE_REG_BITS_ENUM(INT_STATUS_SE_OP_DONE, ACTIVE))) { /* ... */ }
+
+            /* Wait for the engine to be idle. */
+            while (!reg::HasValue(SE->SE_STATUS, SE_REG_BITS_ENUM(STATUS_STATE, IDLE))) { /* ... */ }
+
+            /* Wait for the memory interface to be idle. */
+            while (!reg::HasValue(SE->SE_STATUS, SE_REG_BITS_ENUM(STATUS_MEM_INTERFACE, IDLE))) { /* ... */ }
+        }
+
+        void ValidateErrStatus(volatile SecurityEngineRegisters *SE) {
+            /* Ensure there is no error status. */
+            AMS_ABORT_UNLESS(reg::Read(SE->SE_ERR_STATUS) == 0);
+
+            /* Ensure no error occurred. */
+            AMS_ABORT_UNLESS(reg::HasValue(SE->SE_INT_STATUS, SE_REG_BITS_ENUM(INT_STATUS_ERR_STAT, CLEAR)));
+        }
+
     }
 
     bool ValidateStickyBits(const StickyBits &bits) {
@@ -237,15 +279,76 @@ namespace ams::se {
         }
     }
 
-    void ValidateErrStatus() {
+    void ConfigureAutomaticContextSave() {
-        /* Get the registers. */
+        /* Get registers. */
+        auto *SE  = GetRegisters();
+        auto *SE2 = GetRegisters2();
+
+        /* Automatic context save is supported only on mariko. */
+        if (fuse::GetSocType() == fuse::SocType_Mariko) {
+            /* Configure SE1 to do automatic context save. */
+            reg::Write(SE->SE_CTX_SAVE_AUTO,  SE_REG_BITS_ENUM(CTX_SAVE_AUTO_ENABLE, YES),
+                                              SE_REG_BITS_ENUM(CTX_SAVE_AUTO_LOCK,   YES));
+
+            /* Configure SE2 to do automatic context save. */
+            reg::Write(SE2->SE_CTX_SAVE_AUTO, SE_REG_BITS_ENUM(CTX_SAVE_AUTO_ENABLE, YES),
+                                              SE_REG_BITS_ENUM(CTX_SAVE_AUTO_LOCK,   YES));
+        }
+    }
+
+    void SaveContextAutomatic() {
+        /* Get registers. */
+        auto *SE  = GetRegisters();
+        auto *SE2 = GetRegisters2();
+
+        /* Ensure there's no error status before or after we save context. */
+        ValidateErrStatus();
+        ON_SCOPE_EXIT { ValidateErrStatus(); };
+
+        /* Perform atomic context save. */
+        {
+            /* Check that context save has not already been performed. */
+            AMS_ABORT_UNLESS(reg::HasValue(SE->SE_CTX_SAVE_AUTO,  SE_REG_BITS_VALUE(CTX_SAVE_AUTO_CURR_CNT, 0)));
+            AMS_ABORT_UNLESS(reg::HasValue(SE2->SE_CTX_SAVE_AUTO, SE_REG_BITS_VALUE(CTX_SAVE_AUTO_CURR_CNT, 0)));
+
+            /* Configure SE1 to do context save. */
+            ConfigureForAutomaticContextSave(SE);
+            ConfigureForAutomaticContextSave(SE2);
+
+            /* Start the context save operation. */
+            reg::Write(SE->SE_OPERATION,  SE_REG_BITS_ENUM(OPERATION_OP, CTX_SAVE));
+            reg::Write(SE2->SE_OPERATION, SE_REG_BITS_ENUM(OPERATION_OP, CTX_SAVE));
+
+            /* Wait for the context save operation to complete. */
+            WaitAutomaticContextSaveDone(SE);
+            WaitAutomaticContextSaveDone(SE2);
+
+            /* Check that the correct sizes were written. */
+            AMS_ABORT_UNLESS(reg::HasValue(SE->SE_CTX_SAVE_AUTO,  SE_REG_BITS_VALUE(CTX_SAVE_AUTO_CURR_CNT, SE1ContextSaveOperationCount)));
+            AMS_ABORT_UNLESS(reg::HasValue(SE2->SE_CTX_SAVE_AUTO, SE_REG_BITS_VALUE(CTX_SAVE_AUTO_CURR_CNT, SE2ContextSaveOperationCount)));
+        }
+    }
+
+    void SaveTzramAutomatic() {
+        /* Get registers. */
         auto *SE  = GetRegisters();
 
-        /* Ensure there is no error status. */
+        /* Begin save-to-shadow-tzram operation. */
-        AMS_ABORT_UNLESS(reg::Read(SE->SE_ERR_STATUS) == 0);
+        reg::Write(SE->SE_TZRAM_OPERATION, SE_REG_BITS_ENUM(TZRAM_OPERATION_MODE,     SAVE),
+                                           SE_REG_BITS_ENUM(TZRAM_OPERATION_REQ,  INITIATE));
 
-        /* Ensure no error occurred. */
+        /* Wait for operation to complete. */
-        AMS_ABORT_UNLESS(reg::HasValue(SE->SE_INT_STATUS, SE_REG_BITS_ENUM(INT_STATUS_ERR_STAT, CLEAR)));
+        while (reg::HasValue(SE->SE_TZRAM_OPERATION, SE_REG_BITS_ENUM(TZRAM_OPERATION_BUSY, YES))) { /* ... */ }
+    }
+
+    void ValidateErrStatus() {
+        /* Ensure SE has no error status. */
+        ValidateErrStatus(GetRegisters());
+
+        /* If on mariko, ensure SE2 has no error status. */
+        if (fuse::GetSocType() == fuse::SocType_Mariko) {
+            ValidateErrStatus(GetRegisters2());
+        }
     }
 
 }