From 5ed875ce90d532d0148f7f78cf62458ca04af334 Mon Sep 17 00:00:00 2001
From: Kostas Missos <ctcaer@gmail.com>
Date: Sun, 16 Dec 2018 16:55:56 +0200
Subject: [PATCH] Make tsec checks faster and cover unluckiest cases

+ changes for readability
---
 bootloader/hos/hos.c  |  6 +++++-
 bootloader/hos/hos.h  |  2 ++
 bootloader/sec/tsec.c | 11 +++++++++--
 bootloader/soc/smmu.c |  2 +-
 4 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/bootloader/hos/hos.c b/bootloader/hos/hos.c
index ed3fadf..3009ccf 100644
--- a/bootloader/hos/hos.c
+++ b/bootloader/hos/hos.c
@@ -184,8 +184,12 @@ int keygen(u8 *keyblob, u32 kb, tsec_ctxt_t *tsec_ctxt)
 		memset(tmp, 0x00, 0x20);
 		retries++;
 
-		if (retries > 3)
+		// We rely on racing conditions, make sure we cover even the unluckiest cases.
+		if (retries > 15)
+		{
+			gfx_printf(&gfx_con, "%k\nFailed to get TSEC keys. Please try again.%k\n\n", 0xFFFF0000, 0xFFCCCCCC);
 			return 0;
+		}
 	}
 
 	if (kb >= KB_FIRMWARE_VERSION_620)
diff --git a/bootloader/hos/hos.h b/bootloader/hos/hos.h
index 9f206ee..e962e6f 100644
--- a/bootloader/hos/hos.h
+++ b/bootloader/hos/hos.h
@@ -32,6 +32,8 @@
 #define KB_FIRMWARE_VERSION_620 6
 #define KB_FIRMWARE_VERSION_MAX KB_FIRMWARE_VERSION_620
 
+#define HOS_PKG11_MAGIC 0x31314B50
+
 typedef struct _launch_ctxt_t
 {
 	void *keyblob;
diff --git a/bootloader/sec/tsec.c b/bootloader/sec/tsec.c
index 4de5cc1..b2e6922 100644
--- a/bootloader/sec/tsec.c
+++ b/bootloader/sec/tsec.c
@@ -64,6 +64,7 @@ int tsec_query(u8 *tsec_keys, u8 kb, tsec_ctxt_t *tsec_ctxt)
 	int res = 0;
 	u8 *fwbuf = NULL;
 	u32 *pdir, *car, *fuse, *pmc, *flowctrl, *se, *mc, *iram, *evec;
+	u32 *pkg11_magic_off;
 
 	//Enable clocks.
 	clock_enable_host1x();
@@ -161,6 +162,8 @@ int tsec_query(u8 *tsec_keys, u8 kb, tsec_ctxt_t *tsec_ctxt)
 		// IRAM
 		iram = page_alloc(0x30);
 		memcpy(iram, tsec_ctxt->pkg1, 0x30000);
+		// PKG1.1 magic offset.
+		pkg11_magic_off = (u32 *)(iram + ((tsec_ctxt->pkg11_off + 0x20) / 4));
 		smmu_map(pdir, 0x40010000, (u32)iram, 0x30, _READABLE | _WRITABLE | _NONSECURE);
 
 		// Exception vectors
@@ -216,16 +219,17 @@ int tsec_query(u8 *tsec_keys, u8 kb, tsec_ctxt_t *tsec_ctxt)
 		u32 key[16] = {0};
 		u32 kidx = 0;
 
-		while (memcmp((u8 *)(iram + ((tsec_ctxt->pkg11_off + 0x20) / 4)), "PK11", 4))
+		while (*pkg11_magic_off != HOS_PKG11_MAGIC)
 		{
 			smmu_flush_all();
+
 			if (k == se[SE_KEYTABLE_DATA0_REG_OFFSET / 4])
 				continue;
 			k = se[SE_KEYTABLE_DATA0_REG_OFFSET / 4];
 			key[kidx++] = k;
 
 			// Failsafe.
-			if ((u32)get_tmr_us() - start > 500000)
+			if ((u32)get_tmr_us() - start > 125000)
 				break;
 		}
 
@@ -237,6 +241,9 @@ int tsec_query(u8 *tsec_keys, u8 kb, tsec_ctxt_t *tsec_ctxt)
 			goto out;
 		}
 
+		// Give some extra time to make sure PKG1.1 is decrypted.
+		msleep(50);
+
 		memcpy(tsec_keys, &key, 0x20);
 		memcpy(tsec_ctxt->pkg1, iram, 0x30000);
 		
diff --git a/bootloader/soc/smmu.c b/bootloader/soc/smmu.c
index 85a1deb..5aceaa4 100644
--- a/bootloader/soc/smmu.c
+++ b/bootloader/soc/smmu.c
@@ -94,7 +94,7 @@ void smmu_enable()
 
 	cluster_boot_cpu0((u32)smmu_payload);
 	smmu_used = true;
-	msleep(100);
+	msleep(150);
 
 	smmu_flush_all();
 }