Tegra: smmu: make the context save sequence robust

This patch sanity checks the SMMU context created by the platform
code. The first entry contains the size of the array; which the
driver now verifies before moving on with the save.

This patch also fixes an error in the calculation of the size of
the context that gets copied to TZDRAM.

Signed-off-by: Varun Wadekar <vwadekar@nvidia.com>
diff --git a/plat/nvidia/tegra/common/drivers/smmu/smmu.c b/plat/nvidia/tegra/common/drivers/smmu/smmu.c
index e8b0d0b..a985532 100644
--- a/plat/nvidia/tegra/common/drivers/smmu/smmu.c
+++ b/plat/nvidia/tegra/common/drivers/smmu/smmu.c
@@ -87,7 +87,7 @@
  */
 void tegra_smmu_save_context(uint64_t smmu_ctx_addr)
 {
-	uint32_t i;
+	uint32_t i, num_entries = 0;
 	smmu_regs_t *smmu_ctx_regs;
 #if DEBUG
 	plat_params_from_bl2_t *params_from_bl2 = bl31_get_plat_params();
@@ -110,13 +110,29 @@
 	smmu_ctx_regs = plat_get_smmu_ctx();
 	assert(smmu_ctx_regs);
 
+	/*
+	 * smmu_ctx_regs[0].val contains the size of the context table minus
+	 * the last entry. Sanity check the table size before we start with
+	 * the context save operation.
+	 */
+	while (smmu_ctx_regs[num_entries].val != 0xFFFFFFFFU) {
+		num_entries++;
+	}
+
+	/* panic if the sizes do not match */
+	if (num_entries != smmu_ctx_regs[0].val)
+		panic();
+
 	/* save SMMU register values */
-	for (i = 1; i < smmu_ctx_regs[0].val; i++)
+	for (i = 1; i < num_entries; i++)
 		smmu_ctx_regs[i].val = mmio_read_32(smmu_ctx_regs[i].reg);
 
+	/* increment by 1 to take care of the last entry */
+	num_entries++;
+
 	/* Save SMMU config settings */
 	memcpy16((void *)(uintptr_t)smmu_ctx_addr, (void *)smmu_ctx_regs,
-		 sizeof(smmu_regs_t));
+		 (sizeof(smmu_regs_t) * num_entries));
 
 	/* save the SMMU table address */
 	mmio_write_32(TEGRA_SCRATCH_BASE + SECURE_SCRATCH_RSV11_LO,
diff --git a/plat/nvidia/tegra/soc/t186/plat_psci_handlers.c b/plat/nvidia/tegra/soc/t186/plat_psci_handlers.c
index 66a5999..9790b81 100644
--- a/plat/nvidia/tegra/soc/t186/plat_psci_handlers.c
+++ b/plat/nvidia/tegra/soc/t186/plat_psci_handlers.c
@@ -46,11 +46,8 @@
 
 extern void prepare_cpu_pwr_dwn(void);
 extern void tegra186_cpu_reset_handler(void);
-extern uint32_t __tegra186_cpu_reset_handler_data,
-		__tegra186_cpu_reset_handler_end;
-
-/* TZDRAM offset for saving SMMU context */
-#define TEGRA186_SMMU_CTX_OFFSET	16
+extern uint32_t __tegra186_cpu_reset_handler_end,
+		__tegra186_smmu_context;
 
 /* state id mask */
 #define TEGRA186_STATE_ID_MASK		0xF
@@ -151,9 +148,8 @@
 
 		/* save SMMU context to TZDRAM */
 		smmu_ctx_base = params_from_bl2->tzdram_base +
-			((uintptr_t)&__tegra186_cpu_reset_handler_data -
-			 (uintptr_t)tegra186_cpu_reset_handler) +
-			TEGRA186_SMMU_CTX_OFFSET;
+			((uintptr_t)&__tegra186_smmu_context -
+			 (uintptr_t)tegra186_cpu_reset_handler);
 		tegra_smmu_save_context((uintptr_t)smmu_ctx_base);
 
 		/* Prepare for system suspend */
diff --git a/plat/nvidia/tegra/soc/t186/plat_trampoline.S b/plat/nvidia/tegra/soc/t186/plat_trampoline.S
index 21393d9..ba696f3 100644
--- a/plat/nvidia/tegra/soc/t186/plat_trampoline.S
+++ b/plat/nvidia/tegra/soc/t186/plat_trampoline.S
@@ -94,6 +94,8 @@
 __tegra186_cpu_reset_handler_data:
 	.quad	tegra_secure_entrypoint
 	.quad	__BL31_END__ - BL31_BASE
+	.globl	__tegra186_smmu_context
+__tegra186_smmu_context:
 	.rept	TEGRA186_SMMU_CTX_SIZE
 	.quad	0
 	.endr