Add a Mediatek secure heap which uses TEE service call to protect
buffer. Currently this secure heap is NULL, Prepare for the later patch.
Mainly there are two changes:
a) Add a heap_init ops since TEE probe late than secure heap, thus
   initialize the heap when we require the buffer the first time.
b) Add a priv_data for each heap, like the special data used by MTK
   (such as "TEE session") can be placed in priv_data.

Currently our heap depends on CMA which could only be bool, thus
depend on "TEE=y".

Signed-off-by: Yong Wu <yong.wu@mediatek.com>
---
 drivers/dma-buf/heaps/Kconfig           |   7 ++
 drivers/dma-buf/heaps/Makefile          |   1 +
 drivers/dma-buf/heaps/secure_heap.c     |  11 +++
 drivers/dma-buf/heaps/secure_heap.h     |   4 +
 drivers/dma-buf/heaps/secure_heap_mtk.c | 114 ++++++++++++++++++++++++
 5 files changed, 137 insertions(+)
 create mode 100644 drivers/dma-buf/heaps/secure_heap_mtk.c

diff --git a/drivers/dma-buf/heaps/Kconfig b/drivers/dma-buf/heaps/Kconfig
index 3a9943e94200..12962189878e 100644
--- a/drivers/dma-buf/heaps/Kconfig
+++ b/drivers/dma-buf/heaps/Kconfig
@@ -18,3 +18,10 @@ config DMABUF_HEAPS_SECURE
 	depends on DMABUF_HEAPS
 	help
 	  Choose this option to enable dma-buf secure heap. If in doubt, say N.
+
+config DMABUF_HEAPS_SECURE_MTK
+	bool "MediaTek DMA-BUF Secure Heap"
+	depends on DMABUF_HEAPS_SECURE && TEE=y
+	help
+	  Enable secure dma-buf heaps for MediaTek platform. This heap is backed by
+	  TEE client interfaces. If in doubt, say N.
diff --git a/drivers/dma-buf/heaps/Makefile b/drivers/dma-buf/heaps/Makefile
index b1ad9d1f2fbe..9751dea345df 100644
--- a/drivers/dma-buf/heaps/Makefile
+++ b/drivers/dma-buf/heaps/Makefile
@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: GPL-2.0
 obj-$(CONFIG_DMABUF_HEAPS_SECURE)	+= secure_heap.o
+obj-$(CONFIG_DMABUF_HEAPS_SECURE_MTK)	+= secure_heap_mtk.o
 obj-$(CONFIG_DMABUF_HEAPS_SYSTEM)	+= system_heap.o
 obj-$(CONFIG_DMABUF_HEAPS_CMA)		+= cma_heap.o
diff --git a/drivers/dma-buf/heaps/secure_heap.c b/drivers/dma-buf/heaps/secure_heap.c
index 7cb4db3e55c2..ca4b433fb3f1 100644
--- a/drivers/dma-buf/heaps/secure_heap.c
+++ b/drivers/dma-buf/heaps/secure_heap.c
@@ -150,11 +150,22 @@ secure_heap_allocate(struct dma_heap *heap, unsigned long size,
 		     unsigned long fd_flags, unsigned long heap_flags)
 {
 	struct secure_heap *sec_heap = dma_heap_get_drvdata(heap);
+	const struct secure_heap_ops *ops = sec_heap->ops;
 	struct secure_buffer *sec_buf;
 	DEFINE_DMA_BUF_EXPORT_INFO(exp_info);
 	struct dma_buf *dmabuf;
 	int ret;
 
+	/*
+	 * In some implements, TEE is required to protect buffer. However TEE probe
+	 * may be late, Thus heap_init is performed when the first buffer is requested.
+	 */
+	if (ops->heap_init) {
+		ret = ops->heap_init(sec_heap);
+		if (ret)
+			return ERR_PTR(ret);
+	}
+
 	sec_buf = kzalloc(sizeof(*sec_buf), GFP_KERNEL);
 	if (!sec_buf)
 		return ERR_PTR(-ENOMEM);
diff --git a/drivers/dma-buf/heaps/secure_heap.h b/drivers/dma-buf/heaps/secure_heap.h
index ec5349cd28d0..1ce9c431d989 100644
--- a/drivers/dma-buf/heaps/secure_heap.h
+++ b/drivers/dma-buf/heaps/secure_heap.h
@@ -17,9 +17,13 @@ struct secure_heap {
 	const char		*name;
 
 	const struct secure_heap_ops *ops;
+
+	void			*priv_data;
 };
 
 struct secure_heap_ops {
+	int	(*heap_init)(struct secure_heap *sec_heap);
+
 	int	(*memory_alloc)(struct secure_heap *sec_heap, struct secure_buffer *sec_buf);
 	void	(*memory_free)(struct secure_heap *sec_heap, struct secure_buffer *sec_buf);
 
diff --git a/drivers/dma-buf/heaps/secure_heap_mtk.c b/drivers/dma-buf/heaps/secure_heap_mtk.c
new file mode 100644
index 000000000000..c7e609dd7bd3
--- /dev/null
+++ b/drivers/dma-buf/heaps/secure_heap_mtk.c
@@ -0,0 +1,114 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * DMABUF MediaTek secure heap exporter
+ *
+ * Copyright (C) 2023 MediaTek Inc.
+ */
+#include <linux/dma-buf.h>
+#include <linux/err.h>
+#include <linux/module.h>
+#include <linux/slab.h>
+#include <linux/tee_drv.h>
+#include <linux/uuid.h>
+
+#include "secure_heap.h"
+
+#define TZ_TA_MEM_UUID_MTK		"4477588a-8476-11e2-ad15-e41f1390d676"
+
+#define TEE_PARAM_NUM			4
+
+enum mtk_secure_mem_type {
+	/*
+	 * MediaTek static chunk memory carved out for TrustZone. The memory
+	 * management is inside the TEE.
+	 */
+	MTK_SECURE_MEMORY_TYPE_CM_TZ	= 1,
+};
+
+struct mtk_secure_heap_data {
+	struct tee_context	*tee_ctx;
+	u32			tee_session;
+
+	const enum mtk_secure_mem_type mem_type;
+
+};
+
+static int mtk_tee_ctx_match(struct tee_ioctl_version_data *ver, const void *data)
+{
+	return ver->impl_id == TEE_IMPL_ID_OPTEE;
+}
+
+static int mtk_tee_session_init(struct mtk_secure_heap_data *data)
+{
+	struct tee_param t_param[TEE_PARAM_NUM] = {0};
+	struct tee_ioctl_open_session_arg arg = {0};
+	uuid_t ta_mem_uuid;
+	int ret;
+
+	data->tee_ctx = tee_client_open_context(NULL, mtk_tee_ctx_match, NULL, NULL);
+	if (IS_ERR(data->tee_ctx)) {
+		pr_err_once("%s: open context failed, ret=%ld\n", __func__,
+			    PTR_ERR(data->tee_ctx));
+		return -ENODEV;
+	}
+
+	arg.num_params = TEE_PARAM_NUM;
+	arg.clnt_login = TEE_IOCTL_LOGIN_PUBLIC;
+	ret = uuid_parse(TZ_TA_MEM_UUID_MTK, &ta_mem_uuid);
+	if (ret)
+		goto close_context;
+	memcpy(&arg.uuid, &ta_mem_uuid.b, sizeof(ta_mem_uuid));
+
+	ret = tee_client_open_session(data->tee_ctx, &arg, t_param);
+	if (ret < 0 || arg.ret) {
+		pr_err_once("%s: open session failed, ret=%d:%d\n",
+			    __func__, ret, arg.ret);
+		ret = -EINVAL;
+		goto close_context;
+	}
+	data->tee_session = arg.session;
+	return 0;
+
+close_context:
+	tee_client_close_context(data->tee_ctx);
+	return ret;
+}
+
+static int mtk_secure_heap_init(struct secure_heap *sec_heap)
+{
+	struct mtk_secure_heap_data *data = sec_heap->priv_data;
+
+	if (!data->tee_ctx)
+		return mtk_tee_session_init(data);
+	return 0;
+}
+
+static const struct secure_heap_ops mtk_sec_mem_ops = {
+	.heap_init		= mtk_secure_heap_init,
+};
+
+static struct mtk_secure_heap_data mtk_sec_heap_data = {
+	.mem_type = MTK_SECURE_MEMORY_TYPE_CM_TZ,
+};
+
+static struct secure_heap mtk_secure_heaps[] = {
+	{
+		.name		= "secure_mtk_cm",
+		.ops		= &mtk_sec_mem_ops,
+		.priv_data	= &mtk_sec_heap_data,
+	},
+};
+
+static int mtk_sec_heap_init(void)
+{
+	struct secure_heap *sec_heap = mtk_secure_heaps;
+	unsigned int i;
+
+	for (i = 0; i < ARRAY_SIZE(mtk_secure_heaps); i++, sec_heap++)
+		secure_heap_add(sec_heap);
+	return 0;
+}
+
+module_init(mtk_sec_heap_init);
+MODULE_DESCRIPTION("MediaTek Secure Heap Driver");
+MODULE_LICENSE("GPL");
-- 
2.25.1