ux500: Tee: Use hwmem for memrefs

- Instead of kmalloc, hwmem is used when copying memrefs from user space
  to kernel space.
- Fix for a invalid free when calls to secure world fails.
- Remove free to memrefs etc in tee_release since they should be handled
  completely in tee_write.

ST-Ericsson ID: 375595
ST-Ericsson FOSS-OUT ID: NA
ST-Ericsson Linux next: NA

Depends-On: I2067c34223ce49515c6b7ee8fcc4dcecb9119300

Change-Id: Ibab9edd618d3efe6d6d6302cbb9fb4bde987b99a
Signed-off-by: Joakim Bech <joakim.xx.bech@stericsson.com>
Reviewed-on: http://gerrit.lud.stericsson.com/gerrit/42167
Reviewed-by: QATOOLS
Reviewed-by: Jonas ABERG <jonas.aberg@stericsson.com>
Reviewed-by: Berne HEBARK <berne.hebark@stericsson.com>

2 files changed, 159 insertions, 118 deletions

diff --git a/drivers/tee/tee_driver.c b/drivers/tee/tee_driver.c
index feabe3615c1..442dec5fe06 100644
--- a/drivers/tee/tee_driver.c
+++ b/drivers/tee/tee_driver.c
@@ -15,8 +15,10 @@
 #include <linux/uaccess.h>
 #include <linux/tee.h>
 #include <linux/slab.h>
+#include <linux/hwmem.h>
 
 #define TEED_NAME "tee"
+#define TEED_PFX "TEE: "
 
 #define TEED_STATE_OPEN_DEV 0
 #define TEED_STATE_OPEN_SESSION 1
@@ -30,18 +32,23 @@ static int tee_read(struct file *filp, char __user *buffer,
 static int tee_write(struct file *filp, const char __user *buffer,
 		     size_t length, loff_t *offset);
 
-static inline void set_emsg(struct tee_session *ts, u32 msg)
+static inline void set_emsg(struct tee_session *ts, u32 msg, int line)
 {
+	pr_err(TEED_PFX "msg: 0x%08x at line: %d\n", msg, line);
 	ts->err = msg;
 	ts->origin = TEED_ORIGIN_DRIVER;
 }
 
 static void reset_session(struct tee_session *ts)
 {
+	int i;
+
 	ts->state = TEED_STATE_OPEN_DEV;
 	ts->err = TEED_SUCCESS;
 	ts->origin = TEED_ORIGIN_DRIVER;
 	ts->id = 0;
+	for (i = 0; i < TEEC_CONFIG_PAYLOAD_REF_COUNT; i++)
+		ts->vaddr[i] = NULL;
 	ts->ta = NULL;
 	ts->uuid = NULL;
 	ts->cmd = 0;
@@ -55,9 +62,9 @@ static int copy_ta(struct tee_session *ts,
 {
 	ts->ta = kmalloc(ku_buffer->ta_size, GFP_KERNEL);
 	if (ts->ta == NULL) {
-		pr_err("[%s] error, out of memory (ta)\n",
+		pr_err(TEED_PFX "[%s] error, out of memory (ta)\n",
 		       __func__);
-		set_emsg(ts, TEED_ERROR_OUT_OF_MEMORY);
+		set_emsg(ts, TEED_ERROR_OUT_OF_MEMORY, __LINE__);
 		return -ENOMEM;
 	}
 
@@ -73,9 +80,9 @@ static int copy_uuid(struct tee_session *ts,
 	ts->uuid = kmalloc(sizeof(struct tee_uuid), GFP_KERNEL);
 
 	if (ts->uuid == NULL) {
-		pr_err("[%s] error, out of memory (uuid)\n",
+		pr_err(TEED_PFX "[%s] error, out of memory (uuid)\n",
 		       __func__);
-		set_emsg(ts, TEED_ERROR_OUT_OF_MEMORY);
+		set_emsg(ts, TEED_ERROR_OUT_OF_MEMORY, __LINE__);
 		return -ENOMEM;
 	}
 
@@ -84,13 +91,22 @@ static int copy_uuid(struct tee_session *ts,
 	return 0;
 }
 
-static inline void free_operation(struct tee_session *ts)
+static inline void free_operation(struct tee_session *ts,
+				  struct hwmem_alloc **alloc,
+				  int memrefs_allocated)
 {
 	int i;
 
-	for (i = 0; i < 4; ++i) {
-		kfree(ts->op->shm[i].buffer);
-		ts->op->shm[i].buffer = NULL;
+	for (i = 0; i < memrefs_allocated; ++i) {
+		if (ts->op->shm[i].buffer) {
+			hwmem_kunmap(alloc[i]);
+			hwmem_unpin(alloc[i]);
+			hwmem_release(alloc[i]);
+			ts->op->shm[i].buffer = NULL;
+		}
+
+		if (ts->vaddr[i])
+			ts->vaddr[i] = NULL;
 	}
 
 	kfree(ts->op);
@@ -101,7 +117,7 @@ static inline void memrefs_phys_to_virt(struct tee_session *ts)
 {
 	int i;
 
-	for (i = 0; i < 4; ++i) {
+	for (i = 0; i < TEEC_CONFIG_PAYLOAD_REF_COUNT; ++i) {
 		if (ts->op->flags & (1 << i)) {
 			ts->op->shm[i].buffer =
 				phys_to_virt((unsigned long)
@@ -114,7 +130,7 @@ static inline void memrefs_virt_to_phys(struct tee_session *ts)
 {
 	int i;
 
-	for (i = 0; i < 4; ++i) {
+	for (i = 0; i < TEEC_CONFIG_PAYLOAD_REF_COUNT; ++i) {
 		if (ts->op->flags & (1 << i)) {
 			ts->op->shm[i].buffer =
 				(void *)virt_to_phys(ts->op->shm[i].buffer);
@@ -122,34 +138,35 @@ static inline void memrefs_virt_to_phys(struct tee_session *ts)
 	}
 }
 
-static int copy_memref_to_user(struct tee_operation *op,
-			       struct tee_operation *ubuf_op,
+static int copy_memref_to_user(struct tee_session *ts,
+			       struct tee_operation __user *ubuf_op,
 			       int memref)
 {
 	unsigned long bytes_left;
 
 	bytes_left = copy_to_user(ubuf_op->shm[memref].buffer,
-				  op->shm[memref].buffer,
-				  op->shm[memref].size);
+				  ts->vaddr[memref],
+				  ts->op->shm[memref].size);
 
 	if (bytes_left != 0) {
-		pr_err("[%s] Failed to copy result to user space (%lu "
+		pr_err(TEED_PFX "[%s] failed to copy result to user space (%lu "
 		       "bytes left of buffer).\n", __func__, bytes_left);
 		return bytes_left;
 	}
 
-	bytes_left = put_user(op->shm[memref].size, &ubuf_op->shm[memref].size);
+	bytes_left = put_user(ts->op->shm[memref].size,
+			      &ubuf_op->shm[memref].size);
 
 	if (bytes_left != 0) {
-		pr_err("[%s] Failed to copy result to user space (%lu "
+		pr_err(TEED_PFX "[%s] failed to copy result to user space (%lu "
 		       "bytes left of size).\n", __func__, bytes_left);
 		return -EINVAL;
 	}
 
-	bytes_left = put_user(op->shm[memref].flags,
+	bytes_left = put_user(ts->op->shm[memref].flags,
 			      &ubuf_op->shm[memref].flags);
 	if (bytes_left != 0) {
-		pr_err("[%s] Failed to copy result to user space (%lu "
+		pr_err(TEED_PFX "[%s] failed to copy result to user space (%lu "
 		       "bytes left of flags).\n", __func__, bytes_left);
 		return -EINVAL;
 	}
@@ -157,30 +174,66 @@ static int copy_memref_to_user(struct tee_operation *op,
 	return 0;
 }
 
-static int copy_memref_to_kernel(struct tee_operation *op,
-				 struct tee_operation *kbuf_op,
+static int copy_memref_to_kernel(struct tee_session *ts,
+				 struct tee_session *ku_buffer,
+				 struct hwmem_alloc **alloc,
 				 int memref)
 {
-	/* Buffer freed in invoke_command if this function fails */
-	op->shm[memref].buffer = kmalloc(kbuf_op->shm[memref].size, GFP_KERNEL);
+	int ret = -EINVAL;
+	size_t mem_chunks_length = 1;
+	struct hwmem_mem_chunk mem_chunks;
 
-	if (!op->shm[memref].buffer) {
-		pr_err("[%s] out of memory\n", __func__);
-		return -ENOMEM;
+	if (ku_buffer->op->shm[memref].size == 0) {
+		pr_err(TEED_PFX "[%s] error, size of memref is zero "
+		       "(memref: %d)\n", __func__, memref);
+		return ret;
+	}
+
+	alloc[memref] = hwmem_alloc(ku_buffer->op->shm[memref].size,
+				    (HWMEM_ALLOC_HINT_WRITE_COMBINE |
+				     HWMEM_ALLOC_HINT_CACHED |
+				     HWMEM_ALLOC_HINT_CACHE_WB |
+				     HWMEM_ALLOC_HINT_CACHE_AOW |
+				     HWMEM_ALLOC_HINT_INNER_AND_OUTER_CACHE),
+				    (HWMEM_ACCESS_READ | HWMEM_ACCESS_WRITE |
+				     HWMEM_ACCESS_IMPORT),
+				    HWMEM_MEM_CONTIGUOUS_SYS);
+
+	if (IS_ERR(alloc[memref])) {
+		pr_err(TEED_PFX "[%s] couldn't alloc hwmem_alloc (memref: %d)"
+		       "\n", __func__, memref);
+		return PTR_ERR(alloc[memref]);
+	}
+
+	ret = hwmem_pin(alloc[memref], &mem_chunks, &mem_chunks_length);
+	if (ret) {
+		pr_err(TEED_PFX "[%s] couldn't pin buffer (memref: %d)\n",
+		       __func__, memref);
+		return ret;
 	}
 
 	/*
-	 * Copy shared memory operations to a local kernel
-	 * buffer if they are of type input.
+	 * Since phys_to_virt is not working for hwmem memory we are storing the
+	 * virtual addresses in separate array in tee_session and we keep the
+	 * address of the physical pointers in the memref buffer.
 	 */
-	if (kbuf_op->shm[memref].flags & TEEC_MEM_INPUT) {
-		memcpy(op->shm[memref].buffer,
-		       kbuf_op->shm[memref].buffer,
-		       kbuf_op->shm[memref].size);
+	ts->op->shm[memref].buffer = (void *)mem_chunks.paddr;
+	ts->vaddr[memref] = hwmem_kmap(alloc[memref]);
+
+	/* Buffer unmapped/freed in invoke_command if this function fails. */
+	if (!ts->op->shm[memref].buffer || !ts->vaddr[memref]) {
+		pr_err(TEED_PFX "[%s] out of memory (memref: %d)\n",
+		       __func__, memref);
+		return -ENOMEM;
 	}
 
-	op->shm[memref].size = kbuf_op->shm[memref].size;
-	op->shm[memref].flags = kbuf_op->shm[memref].flags;
+	if (ku_buffer->op->shm[memref].flags & TEEC_MEM_INPUT)
+		memcpy(ts->vaddr[memref],
+		       ku_buffer->op->shm[memref].buffer,
+		       ku_buffer->op->shm[memref].size);
+
+	ts->op->shm[memref].size = ku_buffer->op->shm[memref].size;
+	ts->op->shm[memref].flags = ku_buffer->op->shm[memref].flags;
 
 	return 0;
 }
@@ -191,7 +244,7 @@ static int open_tee_device(struct tee_session *ts,
 	int ret;
 
 	if (ku_buffer->driver_cmd != TEED_OPEN_SESSION) {
-		set_emsg(ts, TEED_ERROR_BAD_STATE);
+		set_emsg(ts, TEED_ERROR_BAD_STATE, __LINE__);
 		return -EINVAL;
 	}
 
@@ -200,7 +253,7 @@ static int open_tee_device(struct tee_session *ts,
 	} else if (ku_buffer->uuid) {
 		ret = copy_uuid(ts, ku_buffer);
 	} else {
-		set_emsg(ts, TEED_ERROR_COMMUNICATION);
+		set_emsg(ts, TEED_ERROR_COMMUNICATION, __LINE__);
 		return -EINVAL;
 	}
 
@@ -215,63 +268,61 @@ static int invoke_command(struct tee_session *ts,
 {
 	int i;
 	int ret = 0;
-	struct tee_operation *kbuf_op =
-		(struct tee_operation *)ku_buffer->op;
+	/* To keep track of which memrefs to free when failure occurs. */
+	int memrefs_allocated = 0;
+	struct hwmem_alloc *alloc[TEEC_CONFIG_PAYLOAD_REF_COUNT];
 
 	ts->op = kmalloc(sizeof(struct tee_operation), GFP_KERNEL);
 
 	if (!ts->op) {
 		if (ts->op == NULL) {
-			pr_err("[%s] error, out of memory "
+			pr_err(TEED_PFX "[%s] error, out of memory "
 			       "(op)\n", __func__);
-			set_emsg(ts, TEED_ERROR_OUT_OF_MEMORY);
+			set_emsg(ts, TEED_ERROR_OUT_OF_MEMORY, __LINE__);
 			return -ENOMEM;
 		}
 	}
 
-	/* Copy memrefs to kernel space. */
-	ts->op->flags = kbuf_op->flags;
+	ts->op->flags = ku_buffer->op->flags;
 	ts->cmd = ku_buffer->cmd;
 
-	for (i = 0; i < 4; ++i) {
-		/* We only want to copy memrefs in use. */
-		if (kbuf_op->flags & (1 << i)) {
-			ret = copy_memref_to_kernel(ts->op, kbuf_op, i);
-
-			if (ret)
+	for (i = 0; i < TEEC_CONFIG_PAYLOAD_REF_COUNT; ++i) {
+		ts->op->shm[i].buffer = NULL;
+		memrefs_allocated++;
+
+		/* We only want to copy memrefs in use to kernel space. */
+		if (ku_buffer->op->flags & (1 << i)) {
+			ret = copy_memref_to_kernel(ts, ku_buffer, alloc, i);
+			if (ret) {
+				pr_err(TEED_PFX "[%s] failed copy memref[%d] "
+				       "to kernel", __func__, i);
 				goto err;
+			}
 		} else {
-			ts->op->shm[i].buffer = NULL;
 			ts->op->shm[i].size = 0;
 			ts->op->shm[i].flags = 0;
 		}
 	}
 
-	/* Secure world expects physical addresses. */
-	memrefs_virt_to_phys(ts);
-
 	if (call_sec_world(ts, TEED_INVOKE)) {
+		set_emsg(ts, TEED_ERROR_COMMUNICATION, __LINE__);
 		ret = -EINVAL;
 		goto err;
 	}
 
-	/*
-	 * Convert physical addresses back to virtual address so the
-	 * kernel can free the buffers when closing the session.
-	 */
-	memrefs_phys_to_virt(ts);
-
-	for (i = 0; i < 4; ++i) {
-		if ((kbuf_op->flags & (1 << i)) &&
-		    (kbuf_op->shm[i].flags & TEEC_MEM_OUTPUT)) {
-			struct tee_operation *ubuf_op =
-				(struct tee_operation *)u_buffer->op;
-
-			ret = copy_memref_to_user(ts->op, ubuf_op, i);
+	for (i = 0; i < TEEC_CONFIG_PAYLOAD_REF_COUNT; ++i) {
+		if ((ku_buffer->op->flags & (1 << i)) &&
+		    (ku_buffer->op->shm[i].flags & TEEC_MEM_OUTPUT)) {
+			ret = copy_memref_to_user(ts, u_buffer->op, i);
+			if (ret) {
+				pr_err(TEED_PFX "[%s] failed copy memref[%d] "
+				       "to user", __func__, i);
+				goto err;
+			}
 		}
 	}
 err:
-	free_operation(ts);
+	free_operation(ts, alloc, memrefs_allocated);
 
 	return ret;
 }
@@ -279,15 +330,15 @@ err:
 static int tee_open(struct inode *inode, struct file *filp)
 {
 	struct tee_session *ts;
-
 	filp->private_data = kmalloc(sizeof(struct tee_session),
 				     GFP_KERNEL);
 
-	if (filp->private_data == NULL)
+	if (filp->private_data == NULL) {
+		pr_err(TEED_PFX "[%s] allocation failed", __func__);
 		return -ENOMEM;
+	}
 
-	ts = (struct tee_session *) (filp->private_data);
-
+	ts = (struct tee_session *)(filp->private_data);
 	reset_session(ts);
 
 	return 0;
@@ -295,28 +346,6 @@ static int tee_open(struct inode *inode, struct file *filp)
 
 static int tee_release(struct inode *inode, struct file *filp)
 {
-	struct tee_session *ts;
-	int i;
-
-	ts = (struct tee_session *) (filp->private_data);
-
-	if (ts == NULL)
-			goto no_ts;
-
-	if (ts->op) {
-		for (i = 0; i < 4; ++i) {
-			kfree(ts->op->shm[i].buffer);
-			ts->op->shm[i].buffer = NULL;
-		}
-	}
-
-	kfree(ts->op);
-	ts->op = NULL;
-
-	kfree(ts->ta);
-	ts->ta = NULL;
-
-no_ts:
 	kfree(filp->private_data);
 	filp->private_data = NULL;
 
@@ -334,15 +363,15 @@ static int tee_read(struct file *filp, char __user *buffer,
 	struct tee_session *ts;
 
 	if (length != sizeof(struct tee_read)) {
-		pr_err("[%s] error, incorrect input length\n",
+		pr_err(TEED_PFX "[%s] error, incorrect input length\n",
 		       __func__);
 		return -EINVAL;
 	}
 
-	ts = (struct tee_session *) (filp->private_data);
+	ts = (struct tee_session *)(filp->private_data);
 
 	if (ts == NULL) {
-		pr_err("[%s] error, private_data not "
+		pr_err(TEED_PFX "[%s] error, private_data not "
 		       "initialized\n", __func__);
 		return -EINVAL;
 	}
@@ -355,7 +384,7 @@ static int tee_read(struct file *filp, char __user *buffer,
 	mutex_unlock(&sync);
 
 	if (copy_to_user(buffer, &buf, length)) {
-		pr_err("[%s] error, copy_to_user failed!\n",
+		pr_err(TEED_PFX "[%s] error, copy_to_user failed!\n",
 		       __func__);
 		return -EINVAL;
 	}
@@ -364,31 +393,31 @@ static int tee_read(struct file *filp, char __user *buffer,
 }
 
 /*
- * Called when a process writes to a dev file
+ * Called when a process writes to a dev file.
  */
 static int tee_write(struct file *filp, const char __user *buffer,
 		     size_t length, loff_t *offset)
 {
 	struct tee_session ku_buffer;
 	struct tee_session *ts;
-	int ret = length;
+	int ret = 0;
 
 	if (length != sizeof(struct tee_session)) {
-		pr_err("[%s] error, incorrect input length\n",
+		pr_err(TEED_PFX "[%s] error, incorrect input length\n",
 		       __func__);
 		return -EINVAL;
 	}
 
 	if (copy_from_user(&ku_buffer, buffer, length)) {
-		pr_err("[%s] error, tee_session "
+		pr_err(TEED_PFX "[%s] error, tee_session "
 		       "copy_from_user failed\n", __func__);
 		return -EINVAL;
 	}
 
-	ts = (struct tee_session *) (filp->private_data);
+	ts = (struct tee_session *)(filp->private_data);
 
 	if (ts == NULL) {
-		pr_err("[%s] error, private_data not "
+		pr_err(TEED_PFX "[%s] error, private_data not "
 		       "initialized\n", __func__);
 		return -EINVAL;
 	}
@@ -409,8 +438,11 @@ static int tee_write(struct file *filp, const char __user *buffer,
 
 		case TEED_CLOSE_SESSION:
 			/* no caching implemented yet... */
-			if (call_sec_world(ts, TEED_CLOSE_SESSION))
+			if (call_sec_world(ts, TEED_CLOSE_SESSION)) {
+				set_emsg(ts, TEED_ERROR_COMMUNICATION,
+					 __LINE__);
 				ret = -EINVAL;
+			}
 
 			kfree(ts->ta);
 			ts->ta = NULL;
@@ -419,13 +451,13 @@ static int tee_write(struct file *filp, const char __user *buffer,
 			break;
 
 		default:
-			set_emsg(ts, TEED_ERROR_BAD_PARAMETERS);
+			set_emsg(ts, TEED_ERROR_BAD_PARAMETERS, __LINE__);
 			ret = -EINVAL;
 		}
 		break;
 	default:
-		pr_err("[%s] unknown state\n", __func__);
-		set_emsg(ts, TEED_ERROR_BAD_STATE);
+		pr_err(TEED_PFX "[%s] unknown state\n", __func__);
+		set_emsg(ts, TEED_ERROR_BAD_STATE, __LINE__);
 		ret = -EINVAL;
 	}
 
@@ -433,10 +465,12 @@ static int tee_write(struct file *filp, const char __user *buffer,
 	 * We expect that ret has value zero when reaching the end here.
 	 * If it has any other value some error must have occured.
 	 */
-	if (!ret)
+	if (!ret) {
 		ret = length;
-	else
+	} else {
+		pr_err(TEED_PFX "[%s], forcing error to -EINVAL\n", __func__);
 		ret = -EINVAL;
+	}
 
 	mutex_unlock(&sync);
 
@@ -465,7 +499,8 @@ int teec_open_session(struct tee_context *context,
 	int res = TEED_SUCCESS;
 
 	if (session == NULL || destination == NULL) {
-		pr_err("[%s] session or destination == NULL\n", __func__);
+		pr_err(TEED_PFX "[%s] session or destination == NULL\n",
+		       __func__);
 		if (error_origin != NULL)
 			*error_origin = TEED_ORIGIN_DRIVER;
 		res = TEED_ERROR_BAD_PARAMETERS;
@@ -476,12 +511,12 @@ int teec_open_session(struct tee_context *context,
 
 	/*
 	 * Open a session towards an application already loaded inside
-	 * the TEE
+	 * the TEE.
 	 */
 	session->uuid = kmalloc(sizeof(struct tee_uuid), GFP_KERNEL);
 
 	if (session->uuid == NULL) {
-		pr_err("[%s] error, out of memory (uuid)\n",
+		pr_err(TEED_PFX "[%s] error, out of memory (uuid)\n",
 		       __func__);
 		if (error_origin != NULL)
 			*error_origin = TEED_ORIGIN_DRIVER;
@@ -506,13 +541,14 @@ int teec_close_session(struct tee_session *session)
 	mutex_lock(&sync);
 
 	if (session == NULL) {
-		pr_err("[%s] error, session == NULL\n", __func__);
+		pr_err(TEED_PFX "[%s] error, session == NULL\n", __func__);
 		res = TEED_ERROR_BAD_PARAMETERS;
 		goto exit;
 	}
 
 	if (call_sec_world(session, TEED_CLOSE_SESSION)) {
-		pr_err("[%s] error, call_sec_world failed\n", __func__);
+		pr_err(TEED_PFX "[%s] error, call_sec_world failed\n",
+		       __func__);
 		res = TEED_ERROR_GENERIC;
 		goto exit;
 	}
@@ -539,7 +575,8 @@ int teec_invoke_command(
 	mutex_lock(&sync);
 
 	if (session == NULL || operation == NULL || error_origin == NULL) {
-		pr_err("[%s] error, input parameters == NULL\n", __func__);
+		pr_err(TEED_PFX "[%s] error, input parameters == NULL\n",
+		       __func__);
 		if (error_origin != NULL)
 			*error_origin = TEED_ORIGIN_DRIVER;
 		res = TEED_ERROR_BAD_PARAMETERS;
@@ -561,13 +598,15 @@ int teec_invoke_command(
 	 * Call secure world
 	 */
 	if (call_sec_world(session, TEED_INVOKE)) {
-		pr_err("[%s] error, call_sec_world failed\n", __func__);
+		pr_err(TEED_PFX "[%s] error, call_sec_world failed\n",
+		       __func__);
 		if (error_origin != NULL)
 			*error_origin = TEED_ORIGIN_DRIVER;
 		res = TEED_ERROR_GENERIC;
 	}
 	if (session->err != TEED_SUCCESS) {
-		pr_err("[%s] error, call_sec_world failed\n", __func__);
+		pr_err(TEED_PFX "[%s] error, call_sec_world failed\n",
+		       __func__);
 		if (error_origin != NULL)
 			*error_origin = session->origin;
 		res = session->err;
@@ -632,7 +671,7 @@ static int __init tee_init(void)
 	err = misc_register(&tee_dev);
 
 	if (err) {
-		pr_err("[%s] error %d adding character device "
+		pr_err(TEED_PFX "[%s] error %d adding character device "
 		       "TEE\n", __func__, err);
 	}
 
diff --git a/include/linux/tee.h b/include/linux/tee.h
index 8b71224ac77..c0f8f11d58d 100644
--- a/include/linux/tee.h
+++ b/include/linux/tee.h
@@ -160,6 +160,7 @@ struct tee_context {};
  * @err: Error code (as in Global Platform TEE Client API spec)
  * @origin: Origin for the error code (also from spec).
  * @id: Implementation defined type, 0 if not used.
+ * @vaddr: Virtual address for the memrefs.
  * @ta: The trusted application.
  * @uuid: The uuid for the trusted application.
  * @cmd: The command to be executed in the trusted application.
@@ -179,6 +180,7 @@ struct tee_session {
 	uint32_t err;
 	uint32_t origin;
 	uint32_t id;
+	uint32_t *vaddr[TEEC_CONFIG_PAYLOAD_REF_COUNT];
 	void *ta;
 	struct tee_uuid *uuid;
 	unsigned int cmd;