[PATCH v7 10/11] firmware: imx: enclave-fw: add handling for save/restore IMEM region

[Pankaj Gupta]

Some IMEM region is lost during kernel power down. Due to this,
firmware's functionaity cannot work correctly.

Saving encrypted IMEM region in kernel memory during power down,
and restore IMEM region on resume.

Signed-off-by: Gaurav Jain <gaurav.jain@xxxxxxx>
Signed-off-by: Pankaj Gupta <pankaj.gupta@xxxxxxx>
---
 drivers/firmware/imx/ele_base_msg.c       |  53 ++++++++++
 drivers/firmware/imx/ele_common.c         | 115 ++++++++++++++++++++++
 drivers/firmware/imx/ele_common.h         |   8 ++
 drivers/firmware/imx/se_fw.c              |  36 +++++++
 include/linux/firmware/imx/ele_base_msg.h |  14 +++
 5 files changed, 226 insertions(+)

diff --git a/drivers/firmware/imx/ele_base_msg.c b/drivers/firmware/imx/ele_base_msg.c
index 3a3af2321f67..844bb0e54c02 100644
--- a/drivers/firmware/imx/ele_base_msg.c
+++ b/drivers/firmware/imx/ele_base_msg.c
@@ -216,3 +216,56 @@ int ele_start_rng(struct device *dev)
 
 	return ret;
 }
+
+int ele_service_swap(struct device *dev,
+		     phys_addr_t addr,
+		     u32 addr_size, u16 flag)
+{
+	struct ele_mu_priv *priv = dev_get_drvdata(dev);
+	int ret;
+	unsigned int status;
+
+	ret = imx_se_alloc_tx_rx_buf(priv);
+	if (ret)
+		return ret;
+
+	ret = plat_fill_cmd_msg_hdr(priv,
+				    (struct mu_hdr *)&priv->tx_msg->header,
+				    ELE_SERVICE_SWAP_REQ,
+				    ELE_SERVICE_SWAP_REQ_MSG_SZ,
+				    true);
+	if (ret)
+		return ret;
+
+	priv->tx_msg->data[0] = flag;
+	priv->tx_msg->data[1] = addr_size;
+	priv->tx_msg->data[2] = ELE_NONE_VAL;
+	priv->tx_msg->data[3] = lower_32_bits(addr);
+	priv->tx_msg->data[4] = plat_add_msg_crc((uint32_t *)&priv->tx_msg[0],
+						 ELE_SERVICE_SWAP_REQ_MSG_SZ);
+	ret = imx_ele_msg_send_rcv(priv);
+	if (ret < 0)
+		return ret;
+
+	ret  = validate_rsp_hdr(priv,
+				priv->rx_msg->header,
+				ELE_SERVICE_SWAP_REQ,
+				ELE_SERVICE_SWAP_RSP_MSG_SZ,
+				true);
+	if (ret)
+		return ret;
+
+	status = RES_STATUS(priv->rx_msg->data[0]);
+	if (status != priv->success_tag) {
+		dev_err(dev, "Command Id[%d], Response Failure = 0x%x",
+			ELE_SERVICE_SWAP_REQ, status);
+		ret = -1;
+	} else {
+		if (flag == ELE_IMEM_EXPORT)
+			ret = priv->rx_msg->data[1];
+		else
+			ret = 0;
+	}
+
+	return ret;
+}
diff --git a/drivers/firmware/imx/ele_common.c b/drivers/firmware/imx/ele_common.c
index d4b829c19133..348f7194778f 100644
--- a/drivers/firmware/imx/ele_common.c
+++ b/drivers/firmware/imx/ele_common.c
@@ -45,6 +45,18 @@ void imx_se_free_tx_rx_buf(struct ele_mu_priv *priv)
 		devm_kfree(priv->dev, priv->rx_msg);
 }
 
+uint32_t plat_add_msg_crc(uint32_t *msg, uint32_t msg_len)
+{
+	uint32_t i;
+	uint32_t crc = 0;
+	uint32_t nb_words = msg_len / (uint32_t)sizeof(uint32_t);
+
+	for (i = 0; i < nb_words - 1; i++)
+		crc ^= *(msg + i);
+
+	return crc;
+}
+
 int imx_ele_msg_send_rcv(struct ele_mu_priv *priv)
 {
 	unsigned int wait;
@@ -177,3 +189,106 @@ int ele_do_start_rng(struct device *dev)
 
 	return 0;
 }
+
+#ifdef CONFIG_PM_SLEEP
+int save_imem(struct device *dev)
+{
+	int ret;
+	struct ele_mu_priv *priv = dev_get_drvdata(dev);
+
+	/* EXPORT command will save encrypted IMEM to given address,
+	 * so later in resume, IMEM can be restored from the given
+	 * address.
+	 *
+	 * Size must be at least 64 kB.
+	 */
+	ret = ele_service_swap(dev,
+			       priv->imem.phyaddr,
+			       ELE_IMEM_SIZE,
+			       ELE_IMEM_EXPORT);
+	if (ret < 0)
+		dev_err(dev, "Failed to export IMEM\n");
+	else
+		dev_info(dev,
+				"Exported %d bytes of encrypted IMEM\n",
+				ret);
+
+	return ret;
+}
+
+int restore_imem(struct device *dev,
+		 uint8_t *pool_name)
+{
+	int ret;
+	u32 imem_state;
+	u32 *get_info_buf = NULL;
+	phys_addr_t get_info_phyaddr = 0;
+	struct ele_mu_priv *priv = dev_get_drvdata(dev);
+
+	get_info_phyaddr
+		= pool_name ? get_phy_buf_mem_pool(dev,
+						   pool_name,
+						   &get_info_buf,
+						   DEVICE_GET_INFO_SZ)
+			    : 0x0;
+
+	if (!get_info_buf) {
+		dev_err(dev, "Unable to alloc sram from sram pool\n");
+		return -ENOMEM;
+	}
+
+	ret = ele_do_start_rng(dev);
+	if (ret)
+		goto exit;
+
+	/* get info from ELE */
+	ret = ele_get_info(dev, get_info_phyaddr, ELE_GET_INFO_READ_SZ);
+	if (ret) {
+		dev_err(dev, "Failed to get info from ELE.\n");
+		goto exit;
+	}
+
+	/* Get IMEM state, if 0xFE then import IMEM */
+	imem_state = (get_info_buf[ELE_IMEM_STATE_WORD]
+			& ELE_IMEM_STATE_MASK) >> 16;
+	if (imem_state == ELE_IMEM_STATE_BAD) {
+		/* IMPORT command will restore IMEM from the given
+		 * address, here size is the actual size returned by ELE
+		 * during the export operation
+		 */
+		ret = ele_service_swap(dev,
+				       priv->imem.phyaddr,
+				       priv->imem.size,
+				       ELE_IMEM_IMPORT);
+		if (ret) {
+			dev_err(dev, "Failed to import IMEM\n");
+			goto exit;
+		}
+	} else
+		goto exit;
+
+	/* After importing IMEM, check if IMEM state is equal to 0xCA
+	 * to ensure IMEM is fully loaded and
+	 * ELE functionality can be used.
+	 */
+	ret = ele_get_info(dev, get_info_phyaddr, ELE_GET_INFO_READ_SZ);
+	if (ret) {
+		dev_err(dev, "Failed to get info from ELE.\n");
+		goto exit;
+	}
+
+	imem_state = (get_info_buf[ELE_IMEM_STATE_WORD]
+			& ELE_IMEM_STATE_MASK) >> 16;
+	if (imem_state == ELE_IMEM_STATE_OK)
+		dev_info(dev, "Successfully restored IMEM\n");
+	else
+		dev_err(dev, "Failed to restore IMEM\n");
+
+exit:
+	if (pool_name && get_info_buf)
+		free_phybuf_mem_pool(dev, pool_name,
+				get_info_buf, DEVICE_GET_INFO_SZ);
+
+	return ret;
+}
+#endif
diff --git a/drivers/firmware/imx/ele_common.h b/drivers/firmware/imx/ele_common.h
index f9e1d949dc6a..e1979b555b58 100644
--- a/drivers/firmware/imx/ele_common.h
+++ b/drivers/firmware/imx/ele_common.h
@@ -9,6 +9,7 @@
 
 #include "se_fw.h"
 
+uint32_t plat_add_msg_crc(uint32_t *msg, uint32_t msg_len);
 int imx_ele_msg_send_rcv(struct ele_mu_priv *priv);
 void imx_se_free_tx_rx_buf(struct ele_mu_priv *priv);
 int imx_se_alloc_tx_rx_buf(struct ele_mu_priv *priv);
@@ -29,4 +30,11 @@ static inline int ele_trng_init(struct device *dev)
 #endif
 
 int ele_do_start_rng(struct device *dev);
+
+#ifdef CONFIG_PM_SLEEP
+int save_imem(struct device *dev);
+int restore_imem(struct device *dev,
+		 uint8_t *pool_name);
 #endif
+
+#endif /*__ELE_COMMON_H__ */
diff --git a/drivers/firmware/imx/se_fw.c b/drivers/firmware/imx/se_fw.c
index 3638a4fed5e7..db992b3cceb9 100644
--- a/drivers/firmware/imx/se_fw.c
+++ b/drivers/firmware/imx/se_fw.c
@@ -52,6 +52,7 @@ struct imx_info {
 	/* platform specific flag to enable/disable the ELE True RNG */
 	bool start_rng;
 	bool enable_ele_trng;
+	bool imem_mgmt;
 };
 
 struct imx_info_list {
@@ -82,6 +83,7 @@ static const struct imx_info_list imx8ulp_info = {
 				.init_fw = false,
 				.start_rng = true,
 				.enable_ele_trng = false,
+				.imem_mgmt = true,
 			},
 	},
 };
@@ -107,6 +109,7 @@ static const struct imx_info_list imx93_info = {
 				.init_fw = true,
 				.start_rng = true,
 				.enable_ele_trng = true,
+				.imem_mgmt = false,
 			},
 	},
 };
@@ -1062,6 +1065,17 @@ static int se_probe_cleanup(struct platform_device *pdev)
 	if (priv->rx_chan)
 		mbox_free_channel(priv->rx_chan);
 
+	/* free the buffer in ele-mu remove, previously allocated
+	 * in ele-mu probe to store encrypted IMEM
+	 */
+	if (priv->imem.buf) {
+		dmam_free_coherent(&pdev->dev,
+				   ELE_IMEM_SIZE,
+				   priv->imem.buf,
+				   priv->imem.phyaddr);
+		priv->imem.buf = NULL;
+	}
+
 	if (priv->flags & RESERVED_DMA_POOL) {
 		of_reserved_mem_device_release(dev);
 		priv->flags &= (~RESERVED_DMA_POOL);
@@ -1281,6 +1295,18 @@ static int se_fw_probe(struct platform_device *pdev)
 			dev_err(dev, "Failed to init ele-trng\n");
 	}
 
+	if (info->imem_mgmt) {
+		/* allocate buffer where ELE store encrypted IMEM */
+		priv->imem.buf = dmam_alloc_coherent(dev, ELE_IMEM_SIZE,
+						     &priv->imem.phyaddr,
+						     GFP_KERNEL);
+		if (!priv->imem.buf) {
+			dev_err(dev,
+				"dmam-alloc-failed: To store encr-IMEM.\n");
+			ret = -ENOMEM;
+			goto exit;
+		}
+	}
 	dev_info(dev, "i.MX secure-enclave: %s interface to firmware, configured.\n",
 		info->pdev_name[1]);
 
@@ -1303,17 +1329,27 @@ static int se_fw_remove(struct platform_device *pdev)
 #ifdef CONFIG_PM_SLEEP
 static int se_fw_suspend(struct device *dev)
 {
+	struct ele_mu_priv *priv = dev_get_drvdata(dev);
+	const struct imx_info *info = priv->info;
+
+	if (info && info->imem_mgmt)
+		priv->imem.size = save_imem(dev);
+
 	return 0;
 }
 
 static int se_fw_resume(struct device *dev)
 {
 	struct ele_mu_priv *priv = dev_get_drvdata(dev);
+	const struct imx_info *info = priv->info;
 	int i;
 
 	for (i = 0; i < priv->max_dev_ctx; i++)
 		wake_up_interruptible(&priv->ctxs[i]->wq);
 
+	if (info && info->imem_mgmt)
+		restore_imem(dev, info->pool_name);
+
 	return 0;
 }
 #endif
diff --git a/include/linux/firmware/imx/ele_base_msg.h b/include/linux/firmware/imx/ele_base_msg.h
index 3ca4b47e4c4e..f027f687de2e 100644
--- a/include/linux/firmware/imx/ele_base_msg.h
+++ b/include/linux/firmware/imx/ele_base_msg.h
@@ -45,9 +45,23 @@
 #define ELE_GET_TRNG_STATE_RETRY_COUNT	0x5
 #define CSAL_TRNG_STATE_MASK		0x0000ffff
 
+#define ELE_SERVICE_SWAP_REQ		0xDF
+#define ELE_SERVICE_SWAP_REQ_MSG_SZ	0x18
+#define ELE_SERVICE_SWAP_RSP_MSG_SZ	0x0C
+#define ELE_IMEM_SIZE			0x10000
+#define ELE_IMEM_STATE_OK		0xCA
+#define ELE_IMEM_STATE_BAD		0xFE
+#define ELE_IMEM_STATE_WORD		0x27
+#define ELE_IMEM_STATE_MASK		0x00ff0000
+#define ELE_IMEM_EXPORT			0x1
+#define ELE_IMEM_IMPORT			0x2
+
 int ele_get_info(struct device *dev, phys_addr_t addr, u32 data_size);
 int ele_ping(struct device *dev);
 int ele_start_rng(struct device *dev);
 int ele_get_trng_state(struct device *dev);
+int ele_service_swap(struct device *dev,
+		     phys_addr_t addr,
+		     u32 addr_size, u16 flag);
 
 #endif
-- 
2.34.1