From 8e5228866c6be84502ef6bde92180e9ba9eab4ff Mon Sep 17 00:00:00 2001
From: Michael Scire <SciresM@gmail.com>
Date: Thu, 8 Mar 2018 04:59:00 -0800
Subject: [PATCH] Fix smcComputeCmac, generally improve userpage cache usage

---
 exosphere/src/se.c       | 24 ++++++++++++++----------
 exosphere/src/smc_user.c |  9 ++++++++-
 exosphere/src/userpage.c |  2 ++
 3 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/exosphere/src/se.c b/exosphere/src/se.c
index e279d3d21..0a63154d0 100644
--- a/exosphere/src/se.c
+++ b/exosphere/src/se.c
@@ -216,7 +216,7 @@ void clear_aes_keyslot_iv(unsigned int keyslot) {
     }
 
     for (size_t i = 0; i < (0x10 >> 2); i++) {
-        SECURITY_ENGINE->AES_KEYTABLE_ADDR = (keyslot << 4) | 8;
+        SECURITY_ENGINE->AES_KEYTABLE_ADDR = (keyslot << 4) | 8 | i;
         SECURITY_ENGINE->AES_KEYTABLE_DATA = 0;
     }
 }
@@ -584,11 +584,14 @@ void se_compute_aes_cmac(unsigned int keyslot, void *cmac, size_t cmac_size, con
     if (keyslot >= KEYSLOT_AES_MAX) {
         generic_panic();
     }
-
+    
+    if (data_size) {
+        flush_dcache_range((uint8_t *)data, (uint8_t *)data + data_size);
+    }
+    
     /* Generate the derived key, to be XOR'd with final output block. */
-    uint8_t derived_key[0x10];
-    memset(derived_key, 0, sizeof(derived_key));
-    se_aes_128_ecb_encrypt_block(keyslot, derived_key, sizeof(derived_key), derived_key, sizeof(derived_key));
+    uint8_t derived_key[0x10] = {0};
+    se_aes_ecb_encrypt_block(keyslot, derived_key, sizeof(derived_key), derived_key, sizeof(derived_key), config_high);
     shift_left_xor_rb(derived_key);
     if (data_size & 0xF) {
         shift_left_xor_rb(derived_key);
@@ -597,6 +600,7 @@ void se_compute_aes_cmac(unsigned int keyslot, void *cmac, size_t cmac_size, con
     SECURITY_ENGINE->CONFIG_REG = (ALG_AES_ENC | DST_HASHREG) | (config_high << 16);
     SECURITY_ENGINE->CRYPTO_REG = (keyslot << 24) | (0x145);
     clear_aes_keyslot_iv(keyslot);
+    
 
     unsigned int num_blocks = (data_size + 0xF) >> 4;
     /* Handle aligned blocks. */
@@ -605,22 +609,22 @@ void se_compute_aes_cmac(unsigned int keyslot, void *cmac, size_t cmac_size, con
         trigger_se_blocking_op(OP_START, NULL, 0, data, data_size);
         SECURITY_ENGINE->CRYPTO_REG |= 0x80;
     }
-
+    
     /* Create final block. */
-    uint8_t last_block[0x10];
-    memset(last_block, 0, sizeof(last_block));
+    uint8_t last_block[0x10] = {0};
     if (data_size & 0xF) {
         memcpy(last_block, data + (data_size & ~0xF), data_size & 0xF);
         last_block[data_size & 0xF] = 0x80; /* Last block = data || 100...0 */
     } else if (data_size >= 0x10) {
         memcpy(last_block, data + data_size - 0x10, 0x10);
     }
-
+    
     for (unsigned int i = 0; i < 0x10; i++) {
         last_block[i] ^= derived_key[i];
     }
-
+    
     /* Perform last operation. */
+    SECURITY_ENGINE->BLOCK_COUNT_REG = 0;
     flush_dcache_range(last_block, last_block + sizeof(last_block));
     trigger_se_blocking_op(OP_START, NULL, 0, last_block, sizeof(last_block));
 
diff --git a/exosphere/src/smc_user.c b/exosphere/src/smc_user.c
index ff6a4034e..d25bf046f 100644
--- a/exosphere/src/smc_user.c
+++ b/exosphere/src/smc_user.c
@@ -355,8 +355,9 @@ uint32_t user_compute_cmac(smc_args_t *args) {
     if (upage_init(&page_ref, user_address) == 0 || user_copy_to_secure(&page_ref, user_data, user_address, size) == 0) {
         return 2;
     }
-
+    
     flush_dcache_range(user_data, user_data + size);
+
     se_compute_aes_128_cmac(keyslot, result_cmac, 0x10, user_data, size);
 
     /* Copy CMAC out. */
@@ -399,6 +400,8 @@ uint32_t user_load_rsa_oaep_key(smc_args_t *args) {
     if (upage_init(&page_ref, user_address) == 0 || user_copy_to_secure(&page_ref, user_data, user_address, size) == 0) {
         return 2;
     }
+    
+    flush_dcache_range(user_data, user_data + size);
 
     /* Ensure that our private key is 0x100 bytes. */
     if (gcm_decrypt_key(user_data, size, user_data, size, sealed_kek, 0x10, wrapped_key, 0x10, CRYPTOUSECASE_RSAOAEP, is_personalized) < 0x100) {
@@ -446,6 +449,8 @@ uint32_t user_decrypt_rsa_private_key(smc_args_t *args) {
     if (upage_init(&page_ref, user_address) == 0 || user_copy_to_secure(&page_ref, user_data, user_address, size) == 0) {
         return 2;
     }
+    
+    flush_dcache_range(user_data, user_data + size);
 
     size_t out_size;
 
@@ -494,6 +499,8 @@ uint32_t user_load_secure_exp_mod_key(smc_args_t *args) {
     if (upage_init(&page_ref, user_address) == 0 || user_copy_to_secure(&page_ref, user_data, user_address, size) == 0) {
         return 2;
     }
+    
+    flush_dcache_range(user_data, user_data + size);
 
     size_t out_size;
 
diff --git a/exosphere/src/userpage.c b/exosphere/src/userpage.c
index d09c57157..01ee0afe2 100644
--- a/exosphere/src/userpage.c
+++ b/exosphere/src/userpage.c
@@ -32,6 +32,8 @@ bool upage_init(upage_ref_t *upage, void *user_address) {
             mmu_map_page(mmu_l3_tbl, USER_PAGE_SECURE_MONITOR_ADDR, upage->user_address, userpage_attributes);
             tlb_invalidate_page_inner_shareable((void *)USER_PAGE_SECURE_MONITOR_ADDR);
             upage->secure_monitor_address = USER_PAGE_SECURE_MONITOR_ADDR;
+        } else {
+            generic_panic();
         }
     }