core: simplify shm cookie handling

Simplifies SHM cookie handling by storing the cookie in the mobj instead
of putting the burden on the caller. The cookie parameter is dropped
from the thread_rpc_*_payload() functions. All callers of those
functions are also updated and unused cookie members of related structs
are removed too.

Acked-by: Jerome Forissier <jerome.forissier@linaro.org>
Signed-off-by: Jens Wiklander <jens.wiklander@linaro.org>

17 files changed, 134 insertions, 163 deletions

diff --git a/core/arch/arm/include/kernel/thread.h b/core/arch/arm/include/kernel/thread.h
index 6e3dade6..add9ae00 100644
--- a/core/arch/arm/include/kernel/thread.h
+++ b/core/arch/arm/include/kernel/thread.h
@@ -70,7 +70,6 @@ struct thread_specific_data {
 	struct pgt_cache pgt_cache;
 	void *rpc_fs_payload;
 	struct mobj *rpc_fs_payload_mobj;
-	uint64_t rpc_fs_payload_cookie;
 	size_t rpc_fs_payload_size;
 };
 
@@ -584,19 +583,17 @@ bool thread_enable_prealloc_rpc_cache(void);
  * Allocates data for payload buffers.
  *
  * @size:	size in bytes of payload buffer
- * @cookie:	returned cookie used when freeing the buffer
  *
  * @returns	mobj that describes allocated buffer or NULL on error
  */
-struct mobj *thread_rpc_alloc_payload(size_t size, uint64_t *cookie);
+struct mobj *thread_rpc_alloc_payload(size_t size);
 
 /**
  * Free physical memory previously allocated with thread_rpc_alloc_payload()
  *
- * @cookie:	cookie received when allocating the buffer
  * @mobj:	mobj that describes the buffer
  */
-void thread_rpc_free_payload(uint64_t cookie, struct mobj *mobj);
+void thread_rpc_free_payload(struct mobj *mobj);
 
 /**
  * Does an RPC using a preallocated argument buffer
@@ -616,20 +613,18 @@ unsigned long thread_smc(unsigned long func_id, unsigned long a1,
  * Buffer is exported to user mode applications.
  *
  * @size:	size in bytes of payload buffer
- * @cookie:	returned cookie used when freeing the buffer
  *
  * @returns	mobj that describes allocated buffer or NULL on error
  */
-struct mobj *thread_rpc_alloc_global_payload(size_t size, uint64_t *cookie);
+struct mobj *thread_rpc_alloc_global_payload(size_t size);
 
 /**
  * Free physical memory previously allocated with
  * thread_rpc_alloc_global_payload()
  *
- * @cookie:	cookie received when allocating the buffer
  * @mobj:	mobj that describes the buffer
  */
-void thread_rpc_free_global_payload(uint64_t cookie, struct mobj *mobj);
+void thread_rpc_free_global_payload(struct mobj *mobj);
 
 #endif /*ASM*/
 
diff --git a/core/arch/arm/include/mm/mobj.h b/core/arch/arm/include/mm/mobj.h
index f2389a61..66bf567a 100644
--- a/core/arch/arm/include/mm/mobj.h
+++ b/core/arch/arm/include/mm/mobj.h
@@ -13,6 +13,7 @@
 #include <tee_api_types.h>
 #include <types_ext.h>
 
+#define MOBJ_INVALID_COOKIE	0xffffffffffffffff
 
 struct mobj {
 	const struct mobj_ops *ops;
@@ -30,6 +31,7 @@ struct mobj_ops {
 	void (*free)(struct mobj *mobj);
 	void (*update_mapping)(struct mobj *mobj, struct user_ta_ctx *utc,
 			       vaddr_t va);
+	uint64_t (*get_cookie)(struct mobj *mobj);
 };
 
 extern struct mobj mobj_virt;
@@ -85,6 +87,14 @@ static inline void mobj_update_mapping(struct mobj *mobj,
 		mobj->ops->update_mapping(mobj, utc, va);
 }
 
+static inline uint64_t mobj_get_cookie(struct mobj *mobj)
+{
+	if (mobj && mobj->ops && mobj->ops->get_cookie)
+		return mobj->ops->get_cookie(mobj);
+
+	return MOBJ_INVALID_COOKIE;
+}
+
 static inline bool mobj_is_nonsec(struct mobj *mobj)
 {
 	return mobj_matches(mobj, CORE_MEM_NON_SEC);
@@ -182,7 +192,7 @@ void mobj_reg_shm_unguard(struct mobj *mobj);
 struct mobj *mobj_mapped_shm_alloc(paddr_t *pages, size_t num_pages,
 				   paddr_t page_offset, uint64_t cookie);
 
-struct mobj *mobj_shm_alloc(paddr_t pa, size_t size);
+struct mobj *mobj_shm_alloc(paddr_t pa, size_t size, uint64_t cookie);
 
 struct mobj *mobj_paged_alloc(size_t size);
 
diff --git a/core/arch/arm/kernel/ree_fs_ta.c b/core/arch/arm/kernel/ree_fs_ta.c
index 5ac027c7..299b92c2 100644
--- a/core/arch/arm/kernel/ree_fs_ta.c
+++ b/core/arch/arm/kernel/ree_fs_ta.c
@@ -23,7 +23,6 @@
 struct user_ta_store_handle {
 	struct shdr *nw_ta; /* Non-secure (shared memory) */
 	size_t nw_ta_size;
-	uint64_t cookie;
 	struct mobj *mobj;
 	size_t offs;
 	struct shdr *shdr; /* Verified secure copy of @nw_ta's signed header */
@@ -36,14 +35,12 @@ struct user_ta_store_handle {
  * address of the raw TA binary is received in out parameter @ta.
  */
 static TEE_Result rpc_load(const TEE_UUID *uuid, struct shdr **ta,
-			   uint64_t *cookie_ta, size_t *ta_size,
-			   struct mobj **mobj)
+			   size_t *ta_size, struct mobj **mobj)
 {
 	TEE_Result res;
 	struct optee_msg_param params[2];
-	uint64_t cta = 0;
 
-	if (!uuid || !ta || !cookie_ta || !mobj || !ta_size)
+	if (!uuid || !ta || !mobj || !ta_size)
 		return TEE_ERROR_BAD_PARAMETERS;
 
 	memset(params, 0, sizeof(params));
@@ -58,24 +55,23 @@ static TEE_Result rpc_load(const TEE_UUID *uuid, struct shdr **ta,
 	if (res != TEE_SUCCESS)
 		return res;
 
-	*mobj = thread_rpc_alloc_payload(params[1].u.tmem.size, &cta);
+	*mobj = thread_rpc_alloc_payload(params[1].u.tmem.size);
 	if (!*mobj)
 		return TEE_ERROR_OUT_OF_MEMORY;
 
 	*ta = mobj_get_va(*mobj, 0);
 	/* We don't expect NULL as thread_rpc_alloc_payload() was successful */
 	assert(*ta);
-	*cookie_ta = cta;
 	*ta_size = params[1].u.tmem.size;
 
 	params[0].attr = OPTEE_MSG_ATTR_TYPE_VALUE_INPUT;
 	tee_uuid_to_octets((void *)&params[0].u.value, uuid);
 	msg_param_init_memparam(params + 1, *mobj, 0, params[1].u.tmem.size,
-				cta, MSG_PARAM_MEM_DIR_OUT);
+				MSG_PARAM_MEM_DIR_OUT);
 
 	res = thread_rpc_cmd(OPTEE_MSG_RPC_CMD_LOAD_TA, 2, params);
 	if (res != TEE_SUCCESS)
-		thread_rpc_free_payload(cta, *mobj);
+		thread_rpc_free_payload(*mobj);
 	return res;
 }
 
@@ -89,7 +85,6 @@ static TEE_Result ta_open(const TEE_UUID *uuid,
 	uint32_t hash_algo = 0;
 	struct shdr *ta = NULL;
 	size_t ta_size = 0;
-	uint64_t cookie = 0;
 	TEE_Result res;
 	size_t offs;
 
@@ -98,7 +93,7 @@ static TEE_Result ta_open(const TEE_UUID *uuid,
 		return TEE_ERROR_OUT_OF_MEMORY;
 
 	/* Request TA from tee-supplicant */
-	res = rpc_load(uuid, &ta, &cookie, &ta_size, &mobj);
+	res = rpc_load(uuid, &ta, &ta_size, &mobj);
 	if (res != TEE_SUCCESS)
 		goto error;
 
@@ -170,7 +165,6 @@ static TEE_Result ta_open(const TEE_UUID *uuid,
 
 	handle->nw_ta = ta;
 	handle->nw_ta_size = ta_size;
-	handle->cookie = cookie;
 	handle->offs = offs;
 	handle->hash_algo = hash_algo;
 	handle->hash_ctx = hash_ctx;
@@ -182,7 +176,7 @@ static TEE_Result ta_open(const TEE_UUID *uuid,
 error_free_hash:
 	crypto_hash_free_ctx(hash_ctx, hash_algo);
 error_free_payload:
-	thread_rpc_free_payload(cookie, mobj);
+	thread_rpc_free_payload(mobj);
 error:
 	shdr_free(shdr);
 	free(handle);
@@ -248,7 +242,7 @@ static void ta_close(struct user_ta_store_handle *h)
 {
 	if (!h)
 		return;
-	thread_rpc_free_payload(h->cookie, h->mobj);
+	thread_rpc_free_payload(h->mobj);
 	crypto_hash_free_ctx(h->hash_ctx, h->hash_algo);
 	free(h->shdr);
 	free(h);
diff --git a/core/arch/arm/kernel/thread.c b/core/arch/arm/kernel/thread.c
index 67560bd4..a3d71333 100644
--- a/core/arch/arm/kernel/thread.c
+++ b/core/arch/arm/kernel/thread.c
@@ -610,9 +610,8 @@ void __weak __thread_std_smc_entry(struct thread_smc_args *args)
 
 		tee_fs_rpc_cache_clear(&thr->tsd);
 		if (!thread_prealloc_rpc_cache) {
-			thread_rpc_free_arg(thr->rpc_carg);
+			thread_rpc_free_arg(mobj_get_cookie(thr->rpc_mobj));
 			mobj_free(thr->rpc_mobj);
-			thr->rpc_carg = 0;
 			thr->rpc_arg = 0;
 			thr->rpc_mobj = NULL;
 		}
@@ -1331,9 +1330,8 @@ bool thread_disable_prealloc_rpc_cache(uint64_t *cookie)
 	rv = true;
 	for (n = 0; n < CFG_NUM_THREADS; n++) {
 		if (threads[n].rpc_arg) {
+			*cookie = mobj_get_cookie(threads[n].rpc_mobj);
 			mobj_free(threads[n].rpc_mobj);
-			*cookie = threads[n].rpc_carg;
-			threads[n].rpc_carg = 0;
 			threads[n].rpc_arg = NULL;
 			goto out;
 		}
@@ -1374,11 +1372,10 @@ out:
  * Allocates data for struct optee_msg_arg.
  *
  * @size:	size in bytes of struct optee_msg_arg
- * @cookie:	returned cookie used when freeing the buffer
  *
  * @returns	mobj that describes allocated buffer or NULL on error
  */
-static struct mobj *thread_rpc_alloc_arg(size_t size, uint64_t *cookie)
+static struct mobj *thread_rpc_alloc_arg(size_t size)
 {
 	paddr_t pa;
 	uint64_t co;
@@ -1397,19 +1394,17 @@ static struct mobj *thread_rpc_alloc_arg(size_t size, uint64_t *cookie)
 
 	/* Check if this region is in static shared space */
 	if (core_pbuf_is(CORE_MEM_NSEC_SHM, pa, size))
-		mobj = mobj_shm_alloc(pa, size);
+		mobj = mobj_shm_alloc(pa, size, co);
 	else if ((!(pa & SMALL_PAGE_MASK)) && size <= SMALL_PAGE_SIZE)
 		mobj = mobj_mapped_shm_alloc(&pa, 1, 0, co);
 
 	if (!mobj)
 		goto err;
 
-	*cookie = co;
 	return mobj;
 err:
 	thread_rpc_free_arg(co);
 	mobj_free(mobj);
-	*cookie = 0;
 	return NULL;
 }
 
@@ -1420,22 +1415,22 @@ static bool get_rpc_arg(uint32_t cmd, size_t num_params,
 	struct optee_msg_arg *arg = thr->rpc_arg;
 	struct mobj *mobj;
 	size_t sz = OPTEE_MSG_GET_ARG_SIZE(THREAD_RPC_MAX_NUM_PARAMS);
-	uint64_t c;
 
 	if (num_params > THREAD_RPC_MAX_NUM_PARAMS)
 		return false;
 
 	if (!arg) {
-		mobj = thread_rpc_alloc_arg(sz, &c);
+		mobj = thread_rpc_alloc_arg(sz);
 		if (!mobj)
 			return false;
 
 		arg = mobj_get_va(mobj, 0);
-		if (!arg)
-			goto bad;
+		if (!arg) {
+			thread_rpc_free_arg(mobj_get_cookie(mobj));
+			return false;
+		}
 
 		thr->rpc_arg = arg;
-		thr->rpc_carg = c;
 		thr->rpc_mobj = mobj;
 	}
 
@@ -1445,12 +1440,8 @@ static bool get_rpc_arg(uint32_t cmd, size_t num_params,
 	arg->ret = TEE_ERROR_GENERIC; /* in case value isn't updated */
 
 	*arg_ret = arg;
-	*carg_ret = thr->rpc_carg;
+	*carg_ret = mobj_get_cookie(thr->rpc_mobj);
 	return true;
-
-bad:
-	thread_rpc_free_arg(c);
-	return false;
 }
 
 uint32_t thread_rpc_cmd(uint32_t cmd, size_t num_params,
@@ -1528,16 +1519,16 @@ static void thread_rpc_free(unsigned int bt, uint64_t cookie, struct mobj *mobj)
  *		failed.
  * @cookie:	returned cookie used when freeing the buffer
  */
-static struct mobj *thread_rpc_alloc(size_t size, size_t align, unsigned int bt,
-				     uint64_t *cookie)
+static struct mobj *thread_rpc_alloc(size_t size, size_t align, unsigned int bt)
 {
 	uint32_t rpc_args[THREAD_RPC_NUM_ARGS] = { OPTEE_SMC_RETURN_RPC_CMD };
 	struct optee_msg_arg *arg;
 	uint64_t carg;
 	struct mobj *mobj = NULL;
+	uint64_t cookie;
 
 	if (!get_rpc_arg(OPTEE_MSG_RPC_CMD_SHM_ALLOC, 1, &arg, &carg))
-		goto fail;
+		return NULL;
 
 	arg->params[0].attr = OPTEE_MSG_ATTR_TYPE_VALUE_INPUT;
 	arg->params[0].u.value.a = bt;
@@ -1548,55 +1539,56 @@ static struct mobj *thread_rpc_alloc(size_t size, size_t align, unsigned int bt,
 	thread_rpc(rpc_args);
 
 	if (arg->ret != TEE_SUCCESS)
-		goto fail;
+		return NULL;
 
 	if (arg->num_params != 1)
-		goto fail;
+		return NULL;
 
 	if (arg->params[0].attr == OPTEE_MSG_ATTR_TYPE_TMEM_OUTPUT) {
-		*cookie = arg->params[0].u.tmem.shm_ref;
+		cookie = arg->params[0].u.tmem.shm_ref;
 		mobj = mobj_shm_alloc(arg->params[0].u.tmem.buf_ptr,
-				      arg->params[0].u.tmem.size);
+				      arg->params[0].u.tmem.size,
+				      cookie);
 	} else if (arg->params[0].attr == (OPTEE_MSG_ATTR_TYPE_TMEM_OUTPUT |
 					   OPTEE_MSG_ATTR_NONCONTIG)) {
-		*cookie = arg->params[0].u.tmem.shm_ref;
+		cookie = arg->params[0].u.tmem.shm_ref;
 		mobj = msg_param_mobj_from_noncontig(
 			arg->params[0].u.tmem.buf_ptr,
 			arg->params[0].u.tmem.size,
-			*cookie,
+			cookie,
 			true);
-	} else
-		goto fail;
+	} else {
+		return NULL;
+	}
 
-	if (!mobj)
-		goto free_first;
+	if (!mobj) {
+		thread_rpc_free(bt, cookie, mobj);
+		return NULL;
+	}
 
 	assert(mobj_is_nonsec(mobj));
-	return mobj;
 
-free_first:
-	thread_rpc_free(bt, *cookie, mobj);
-fail:
-	*cookie = 0;
-	return NULL;
+	return mobj;
 }
 
-struct mobj *thread_rpc_alloc_payload(size_t size, uint64_t *cookie)
+struct mobj *thread_rpc_alloc_payload(size_t size)
 {
-	return thread_rpc_alloc(size, 8, OPTEE_MSG_RPC_SHM_TYPE_APPL, cookie);
+	return thread_rpc_alloc(size, 8, OPTEE_MSG_RPC_SHM_TYPE_APPL);
 }
 
-void thread_rpc_free_payload(uint64_t cookie, struct mobj *mobj)
+void thread_rpc_free_payload(struct mobj *mobj)
 {
-	thread_rpc_free(OPTEE_MSG_RPC_SHM_TYPE_APPL, cookie, mobj);
+	thread_rpc_free(OPTEE_MSG_RPC_SHM_TYPE_APPL, mobj_get_cookie(mobj),
+			mobj);
 }
 
-struct mobj *thread_rpc_alloc_global_payload(size_t size, uint64_t *cookie)
+struct mobj *thread_rpc_alloc_global_payload(size_t size)
 {
-	return thread_rpc_alloc(size, 8, OPTEE_MSG_RPC_SHM_TYPE_GLOBAL, cookie);
+	return thread_rpc_alloc(size, 8, OPTEE_MSG_RPC_SHM_TYPE_GLOBAL);
 }
 
-void thread_rpc_free_global_payload(uint64_t cookie, struct mobj *mobj)
+void thread_rpc_free_global_payload(struct mobj *mobj)
 {
-	thread_rpc_free(OPTEE_MSG_RPC_SHM_TYPE_GLOBAL, cookie, mobj);
+	thread_rpc_free(OPTEE_MSG_RPC_SHM_TYPE_GLOBAL, mobj_get_cookie(mobj),
+			mobj);
 }
diff --git a/core/arch/arm/kernel/thread_private.h b/core/arch/arm/kernel/thread_private.h
index fe0a1055..148faee5 100644
--- a/core/arch/arm/kernel/thread_private.h
+++ b/core/arch/arm/kernel/thread_private.h
@@ -91,7 +91,6 @@ struct thread_ctx {
 	struct thread_vfp_state vfp_state;
 #endif
 	void *rpc_arg;
-	uint64_t rpc_carg;
 	struct mobj *rpc_mobj;
 	struct mutex_head mutexes;
 	struct thread_specific_data tsd;
diff --git a/core/arch/arm/mm/mobj.c b/core/arch/arm/mm/mobj.c
index b2bdf6bb..d936ff52 100644
--- a/core/arch/arm/mm/mobj.c
+++ b/core/arch/arm/mm/mobj.c
@@ -416,6 +416,11 @@ static TEE_Result mobj_reg_shm_get_cattr(struct mobj *mobj __unused,
 
 static bool mobj_reg_shm_matches(struct mobj *mobj, enum buf_is_attr attr);
 
+static uint64_t mobj_reg_shm_get_cookie(struct mobj *mobj)
+{
+	return to_mobj_reg_shm(mobj)->cookie;
+}
+
 static const struct mobj_ops mobj_reg_shm_ops __rodata_unpaged = {
 	.get_pa = mobj_reg_shm_get_pa,
 	.get_phys_offs = mobj_reg_shm_get_phys_offs,
@@ -423,6 +428,7 @@ static const struct mobj_ops mobj_reg_shm_ops __rodata_unpaged = {
 	.get_cattr = mobj_reg_shm_get_cattr,
 	.matches = mobj_reg_shm_matches,
 	.free = mobj_reg_shm_free,
+	.get_cookie = mobj_reg_shm_get_cookie,
 };
 
 static bool mobj_reg_shm_matches(struct mobj *mobj __maybe_unused,
@@ -716,6 +722,7 @@ service_init(mobj_mapped_shm_init);
 struct mobj_shm {
 	struct mobj mobj;
 	paddr_t pa;
+	uint64_t cookie;
 };
 
 static struct mobj_shm *to_mobj_shm(struct mobj *mobj);
@@ -771,12 +778,18 @@ static void mobj_shm_free(struct mobj *mobj)
 	free(m);
 }
 
+static uint64_t mobj_shm_get_cookie(struct mobj *mobj)
+{
+	return to_mobj_shm(mobj)->cookie;
+}
+
 static const struct mobj_ops mobj_shm_ops __rodata_unpaged = {
 	.get_va = mobj_shm_get_va,
 	.get_pa = mobj_shm_get_pa,
 	.get_phys_offs = mobj_shm_get_phys_offs,
 	.matches = mobj_shm_matches,
 	.free = mobj_shm_free,
+	.get_cookie = mobj_shm_get_cookie,
 };
 
 static struct mobj_shm *to_mobj_shm(struct mobj *mobj)
@@ -785,7 +798,7 @@ static struct mobj_shm *to_mobj_shm(struct mobj *mobj)
 	return container_of(mobj, struct mobj_shm, mobj);
 }
 
-struct mobj *mobj_shm_alloc(paddr_t pa, size_t size)
+struct mobj *mobj_shm_alloc(paddr_t pa, size_t size, uint64_t cookie)
 {
 	struct mobj_shm *m;
 
@@ -799,6 +812,7 @@ struct mobj *mobj_shm_alloc(paddr_t pa, size_t size)
 	m->mobj.size = size;
 	m->mobj.ops = &mobj_shm_ops;
 	m->pa = pa;
+	m->cookie = cookie;
 
 	return &m->mobj;
 }
diff --git a/core/arch/arm/pta/benchmark.c b/core/arch/arm/pta/benchmark.c
index 492ce1d9..5d1c4d91 100644
--- a/core/arch/arm/pta/benchmark.c
+++ b/core/arch/arm/pta/benchmark.c
@@ -27,7 +27,6 @@ static size_t bench_ts_size;
 
 static struct mutex bench_reg_mu = MUTEX_INITIALIZER;
 static struct mobj *bench_mobj;
-static uint64_t bench_cookie;
 
 static TEE_Result rpc_reg_global_buf(uint64_t type, paddr_t phta, size_t size)
 {
@@ -73,8 +72,7 @@ static TEE_Result alloc_benchmark_buffer(uint32_t type,
 		return TEE_ERROR_BAD_STATE;
 	}
 
-	bench_mobj = thread_rpc_alloc_global_payload(bench_ts_size,
-						     &bench_cookie);
+	bench_mobj = thread_rpc_alloc_global_payload(bench_ts_size);
 	if (!bench_mobj) {
 		EMSG(TA_PRINT_PREFIX
 			"can't create mobj for timestamp buffer");
@@ -84,7 +82,7 @@ static TEE_Result alloc_benchmark_buffer(uint32_t type,
 
 	bench_ts_global = (struct tee_ts_global *)mobj_get_va(bench_mobj, 0);
 	if (!bench_ts_global) {
-		thread_rpc_free_global_payload(bench_cookie, bench_mobj);
+		thread_rpc_free_global_payload(bench_mobj);
 
 		mutex_unlock(&bench_reg_mu);
 		return TEE_ERROR_BAD_STATE;
@@ -159,7 +157,7 @@ static TEE_Result unregister_benchmark(uint32_t type,
 
 	res = rpc_reg_global_buf(OPTEE_MSG_RPC_CMD_BENCH_REG_DEL, 0, 0);
 
-	thread_rpc_free_global_payload(bench_cookie, bench_mobj);
+	thread_rpc_free_global_payload(bench_mobj);
 	return res;
 }
 
diff --git a/core/arch/arm/pta/gprof.c b/core/arch/arm/pta/gprof.c
index 35588464..de270245 100644
--- a/core/arch/arm/pta/gprof.c
+++ b/core/arch/arm/pta/gprof.c
@@ -22,10 +22,9 @@ static TEE_Result gprof_send_rpc(TEE_UUID *uuid, void *buf, size_t len,
 	struct optee_msg_param params[3];
 	struct mobj *mobj;
 	TEE_Result res = TEE_ERROR_GENERIC;
-	uint64_t c = 0;
 	char *va;
 
-	mobj = thread_rpc_alloc_payload(sizeof(*uuid) + len, &c);
+	mobj = thread_rpc_alloc_payload(sizeof(*uuid) + len);
 	if (!mobj)
 		return TEE_ERROR_OUT_OF_MEMORY;
 
@@ -40,9 +39,9 @@ static TEE_Result gprof_send_rpc(TEE_UUID *uuid, void *buf, size_t len,
 	params[0].attr = OPTEE_MSG_ATTR_TYPE_VALUE_INOUT;
 	params[0].u.value.a = *id;
 
-	msg_param_init_memparam(params + 1, mobj, 0, sizeof(*uuid), c,
+	msg_param_init_memparam(params + 1, mobj, 0, sizeof(*uuid),
 				MSG_PARAM_MEM_DIR_IN);
-	msg_param_init_memparam(params + 2, mobj, sizeof(*uuid), len, c,
+	msg_param_init_memparam(params + 2, mobj, sizeof(*uuid), len,
 				MSG_PARAM_MEM_DIR_IN);
 
 	res = thread_rpc_cmd(OPTEE_MSG_RPC_CMD_GPROF, 3, params);
@@ -51,7 +50,7 @@ static TEE_Result gprof_send_rpc(TEE_UUID *uuid, void *buf, size_t len,
 
 	*id = (uint32_t)params[0].u.value.a;
 exit:
-	thread_rpc_free_payload(c, mobj);
+	thread_rpc_free_payload(mobj);
 	return res;
 }
 
diff --git a/core/arch/arm/tee/entry_std.c b/core/arch/arm/tee/entry_std.c
index 92fdf0fc..0b60ed5a 100644
--- a/core/arch/arm/tee/entry_std.c
+++ b/core/arch/arm/tee/entry_std.c
@@ -498,7 +498,7 @@ static struct mobj *get_cmd_buffer(paddr_t parg, uint32_t *num_params)
 	*num_params = READ_ONCE(arg->num_params);
 	args_size = OPTEE_MSG_GET_ARG_SIZE(*num_params);
 
-	return mobj_shm_alloc(parg, args_size);
+	return mobj_shm_alloc(parg, args_size, MOBJ_INVALID_COOKIE);
 }
 
 /*
diff --git a/core/arch/arm/tee/pta_socket.c b/core/arch/arm/tee/pta_socket.c
index 3f2d5a3d..d9a5ac67 100644
--- a/core/arch/arm/tee/pta_socket.c
+++ b/core/arch/arm/tee/pta_socket.c
@@ -23,7 +23,6 @@ static TEE_Result socket_open(uint32_t instance_id, uint32_t param_types,
 {
 	struct mobj *mobj;
 	TEE_Result res;
-	uint64_t cookie;
 	void *va;
 	struct optee_msg_param msg_params[4];
 	uint32_t exp_pt = TEE_PARAM_TYPES(TEE_PARAM_TYPE_VALUE_INPUT,
@@ -39,7 +38,7 @@ static TEE_Result socket_open(uint32_t instance_id, uint32_t param_types,
 
 	memset(msg_params, 0, sizeof(msg_params));
 
-	va = tee_fs_rpc_cache_alloc(params[1].memref.size, &mobj, &cookie);
+	va = tee_fs_rpc_cache_alloc(params[1].memref.size, &mobj);
 	if (!va)
 		return TEE_ERROR_OUT_OF_MEMORY;
 
@@ -54,7 +53,7 @@ static TEE_Result socket_open(uint32_t instance_id, uint32_t param_types,
 
 	/* server address */
 	if (!msg_param_init_memparam(msg_params + 2, mobj, 0,
-				     params[1].memref.size, cookie,
+				     params[1].memref.size,
 				     MSG_PARAM_MEM_DIR_IN))
 		return TEE_ERROR_BAD_STATE;
 	memcpy(va, params[1].memref.buffer, params[1].memref.size);
@@ -100,7 +99,6 @@ static TEE_Result socket_send(uint32_t instance_id, uint32_t param_types,
 {
 	struct mobj *mobj;
 	TEE_Result res;
-	uint64_t cookie;
 	void *va;
 	struct optee_msg_param msg_params[3];
 	uint32_t exp_pt = TEE_PARAM_TYPES(TEE_PARAM_TYPE_VALUE_INPUT,
@@ -116,7 +114,7 @@ static TEE_Result socket_send(uint32_t instance_id, uint32_t param_types,
 
 	memset(msg_params, 0, sizeof(msg_params));
 
-	va = tee_fs_rpc_cache_alloc(params[1].memref.size, &mobj, &cookie);
+	va = tee_fs_rpc_cache_alloc(params[1].memref.size, &mobj);
 	if (!va)
 		return TEE_ERROR_OUT_OF_MEMORY;
 
@@ -127,7 +125,7 @@ static TEE_Result socket_send(uint32_t instance_id, uint32_t param_types,
 
 	/* buffer */
 	if (!msg_param_init_memparam(msg_params + 1, mobj, 0,
-				     params[1].memref.size, cookie,
+				     params[1].memref.size,
 				     MSG_PARAM_MEM_DIR_IN))
 		return TEE_ERROR_BAD_STATE;
 
@@ -147,7 +145,6 @@ static TEE_Result socket_recv(uint32_t instance_id, uint32_t param_types,
 {
 	struct mobj *mobj;
 	TEE_Result res;
-	uint64_t cookie;
 	void *va;
 	struct optee_msg_param msg_params[3];
 	uint32_t exp_pt = TEE_PARAM_TYPES(TEE_PARAM_TYPE_VALUE_INPUT,
@@ -163,7 +160,7 @@ static TEE_Result socket_recv(uint32_t instance_id, uint32_t param_types,
 
 	memset(msg_params, 0, sizeof(msg_params));
 
-	va = tee_fs_rpc_cache_alloc(params[1].memref.size, &mobj, &cookie);
+	va = tee_fs_rpc_cache_alloc(params[1].memref.size, &mobj);
 	if (!va)
 		return TEE_ERROR_OUT_OF_MEMORY;
 
@@ -174,7 +171,7 @@ static TEE_Result socket_recv(uint32_t instance_id, uint32_t param_types,
 
 	/* buffer */
 	if (!msg_param_init_memparam(msg_params + 1, mobj, 0,
-				     params[1].memref.size, cookie,
+				     params[1].memref.size,
 				     MSG_PARAM_MEM_DIR_OUT))
 		return TEE_ERROR_BAD_STATE;
 
@@ -194,7 +191,6 @@ static TEE_Result socket_ioctl(uint32_t instance_id, uint32_t param_types,
 {
 	struct mobj *mobj;
 	TEE_Result res;
-	uint64_t cookie;
 	void *va;
 	struct optee_msg_param msg_params[3];
 	uint32_t exp_pt = TEE_PARAM_TYPES(TEE_PARAM_TYPE_VALUE_INPUT,
@@ -210,7 +206,7 @@ static TEE_Result socket_ioctl(uint32_t instance_id, uint32_t param_types,
 
 	memset(msg_params, 0, sizeof(msg_params));
 
-	va = tee_fs_rpc_cache_alloc(params[1].memref.size, &mobj, &cookie);
+	va = tee_fs_rpc_cache_alloc(params[1].memref.size, &mobj);
 	if (!va)
 		return TEE_ERROR_OUT_OF_MEMORY;
 
@@ -221,7 +217,7 @@ static TEE_Result socket_ioctl(uint32_t instance_id, uint32_t param_types,
 
 	/* buffer */
 	if (!msg_param_init_memparam(msg_params + 1, mobj, 0,
-				     params[1].memref.size, cookie,
+				     params[1].memref.size,
 				     MSG_PARAM_MEM_DIR_INOUT))
 		return TEE_ERROR_BAD_STATE;
 
diff --git a/core/include/kernel/msg_param.h b/core/include/kernel/msg_param.h
index e8e9c899..4c9e5ae6 100644
--- a/core/include/kernel/msg_param.h
+++ b/core/include/kernel/msg_param.h
@@ -66,7 +66,6 @@ struct mobj *msg_param_mobj_from_noncontig(paddr_t buf_ptr, size_t size,
  * @mobj	- mobj describing the shared memory buffer
  * @offset	- offset of the buffer
  * @size	- size of the buffer
- * @cookie	- NW cookie of the shared buffer
  * @dir		- data direction
  *
  * Idea behind this function is that thread_rpc_alloc() can return
@@ -81,7 +80,7 @@ struct mobj *msg_param_mobj_from_noncontig(paddr_t buf_ptr, size_t size,
  */
 bool msg_param_init_memparam(struct optee_msg_param *param, struct mobj *mobj,
 			     size_t offset, size_t size,
-			     uint64_t cookie, enum msg_param_mem_dir dir);
+			     enum msg_param_mem_dir dir);
 /**
  * msg_param_get_buf_size() - helper functions that reads [T/R]MEM
  *			      parameter size
diff --git a/core/include/tee/tee_fs_rpc.h b/core/include/tee/tee_fs_rpc.h
index 1cab39b7..2af5ccda 100644
--- a/core/include/tee/tee_fs_rpc.h
+++ b/core/include/tee/tee_fs_rpc.h
@@ -74,6 +74,6 @@ static inline void tee_fs_rpc_cache_clear(
  * cache. The pointer is guaranteed to point to a large enough area or to
  * be NULL.
  */
-void *tee_fs_rpc_cache_alloc(size_t size, struct mobj **mobj, uint64_t *cookie);
+void *tee_fs_rpc_cache_alloc(size_t size, struct mobj **mobj);
 
 #endif /* TEE_FS_RPC_H */
diff --git a/core/kernel/msg_param.c b/core/kernel/msg_param.c
index 000e0f5a..43f6df10 100644
--- a/core/kernel/msg_param.c
+++ b/core/kernel/msg_param.c
@@ -143,7 +143,7 @@ out:
 
 bool msg_param_init_memparam(struct optee_msg_param *param, struct mobj *mobj,
 			     size_t offset, size_t size,
-			     uint64_t cookie, enum msg_param_mem_dir dir)
+			     enum msg_param_mem_dir dir)
 {
 	if (mobj_matches(mobj, CORE_MEM_REG_SHM)) {
 		/* Registered SHM mobj */
@@ -163,7 +163,7 @@ bool msg_param_init_memparam(struct optee_msg_param *param, struct mobj *mobj,
 
 		param->u.rmem.size = size;
 		param->u.rmem.offs = offset;
-		param->u.rmem.shm_ref = cookie;
+		param->u.rmem.shm_ref = mobj_get_cookie(mobj);
 	} else if (mobj_matches(mobj, CORE_MEM_NSEC_SHM)) {
 		/* MOBJ from from predefined pool */
 		paddr_t pa;
@@ -186,7 +186,7 @@ bool msg_param_init_memparam(struct optee_msg_param *param, struct mobj *mobj,
 		}
 
 		param->u.tmem.buf_ptr = pa + offset;
-		param->u.tmem.shm_ref = cookie;
+		param->u.tmem.shm_ref = mobj_get_cookie(mobj);
 		param->u.tmem.size = size;
 	} else
 		return false;
diff --git a/core/tee/tadb.c b/core/tee/tadb.c
index 5b7f2407..125ac833 100644
--- a/core/tee/tadb.c
+++ b/core/tee/tadb.c
@@ -60,7 +60,6 @@ struct tee_tadb_ta_read {
 	struct tadb_entry entry;
 	size_t pos;
 	void *ctx;
-	uint64_t ta_cookie;
 	struct mobj *ta_mobj;
 	uint8_t *ta_buf;
 };
@@ -88,19 +87,18 @@ static TEE_Result ta_operation_open(unsigned int cmd, uint32_t file_number,
 	struct mobj *mobj;
 	TEE_Result res;
 	void *va;
-	uint64_t cookie;
 	struct optee_msg_param params[] = {
 		[0] = { .attr = OPTEE_MSG_ATTR_TYPE_VALUE_INPUT,
 			.u.value.a = cmd },
 		[2] = { .attr = OPTEE_MSG_ATTR_TYPE_VALUE_OUTPUT }
 	};
 
-	va = tee_fs_rpc_cache_alloc(TEE_FS_NAME_MAX, &mobj, &cookie);
+	va = tee_fs_rpc_cache_alloc(TEE_FS_NAME_MAX, &mobj);
 	if (!va)
 		return TEE_ERROR_OUT_OF_MEMORY;
 
 	if (!msg_param_init_memparam(params + 1, mobj, 0, TEE_FS_NAME_MAX,
-				     cookie, MSG_PARAM_MEM_DIR_IN))
+				     MSG_PARAM_MEM_DIR_IN))
 		return TEE_ERROR_BAD_STATE;
 
 	file_num_to_str(va, TEE_FS_NAME_MAX, file_number);
@@ -116,18 +114,17 @@ static TEE_Result ta_operation_remove(uint32_t file_number)
 {
 	struct mobj *mobj;
 	void *va;
-	uint64_t cookie;
 	struct optee_msg_param params[2] = {
 		[0] = { .attr = OPTEE_MSG_ATTR_TYPE_VALUE_INPUT,
 			.u.value.a = OPTEE_MRF_REMOVE },
 	};
 
-	va = tee_fs_rpc_cache_alloc(TEE_FS_NAME_MAX, &mobj, &cookie);
+	va = tee_fs_rpc_cache_alloc(TEE_FS_NAME_MAX, &mobj);
 	if (!va)
 		return TEE_ERROR_OUT_OF_MEMORY;
 
 	if (!msg_param_init_memparam(params + 1, mobj, 0, TEE_FS_NAME_MAX,
-				     cookie, MSG_PARAM_MEM_DIR_IN))
+				     MSG_PARAM_MEM_DIR_IN))
 		return TEE_ERROR_BAD_STATE;
 
 	file_num_to_str(va, TEE_FS_NAME_MAX, file_number);
@@ -703,7 +700,7 @@ static TEE_Result ta_load(struct tee_tadb_ta_read *ta)
 	if (ta->ta_mobj)
 		return TEE_SUCCESS;
 
-	ta->ta_mobj = thread_rpc_alloc_payload(sz, &ta->ta_cookie);
+	ta->ta_mobj = thread_rpc_alloc_payload(sz);
 	if (!ta->ta_mobj)
 		return TEE_ERROR_OUT_OF_MEMORY;
 
@@ -711,12 +708,12 @@ static TEE_Result ta_load(struct tee_tadb_ta_read *ta)
 	assert(ta->ta_buf);
 
 	if (!msg_param_init_memparam(params + 1, ta->ta_mobj, 0, sz,
-				     ta->ta_cookie, MSG_PARAM_MEM_DIR_OUT))
+				     MSG_PARAM_MEM_DIR_OUT))
 		return TEE_ERROR_BAD_STATE;
 
 	res = thread_rpc_cmd(OPTEE_MSG_RPC_CMD_FS, ARRAY_SIZE(params), params);
 	if (res) {
-		thread_rpc_free_payload(ta->ta_cookie, ta->ta_mobj);
+		thread_rpc_free_payload(ta->ta_mobj);
 		ta->ta_mobj = NULL;
 	}
 	return res;
@@ -780,7 +777,7 @@ void tee_tadb_ta_close(struct tee_tadb_ta_read *ta)
 	crypto_authenc_final(ta->ctx, TADB_AUTH_ENC_ALG);
 	crypto_authenc_free_ctx(ta->ctx, TADB_AUTH_ENC_ALG);
 	if (ta->ta_mobj)
-		thread_rpc_free_payload(ta->ta_cookie, ta->ta_mobj);
+		thread_rpc_free_payload(ta->ta_mobj);
 	tee_fs_rpc_close(OPTEE_MSG_RPC_CMD_FS, ta->fd);
 	tadb_put(ta->db);
 	free(ta);
diff --git a/core/tee/tee_fs_rpc.c b/core/tee/tee_fs_rpc.c
index 21ebc586..5c97617c 100644
--- a/core/tee/tee_fs_rpc.c
+++ b/core/tee/tee_fs_rpc.c
@@ -36,9 +36,8 @@ static TEE_Result operation_open(uint32_t id, unsigned int cmd,
 	struct mobj *mobj;
 	TEE_Result res;
 	void *va;
-	uint64_t cookie;
 
-	va = tee_fs_rpc_cache_alloc(TEE_FS_NAME_MAX, &mobj, &cookie);
+	va = tee_fs_rpc_cache_alloc(TEE_FS_NAME_MAX, &mobj);
 	if (!va)
 		return TEE_ERROR_OUT_OF_MEMORY;
 
@@ -46,7 +45,7 @@ static TEE_Result operation_open(uint32_t id, unsigned int cmd,
 	op.params[0].u.value.a = cmd;
 
 	if (!msg_param_init_memparam(op.params + 1, mobj, 0, TEE_FS_NAME_MAX,
-				     cookie, MSG_PARAM_MEM_DIR_IN))
+				     MSG_PARAM_MEM_DIR_IN))
 		return TEE_ERROR_BAD_STATE;
 
 	res = tee_svc_storage_create_filename(va, TEE_FS_NAME_MAX,
@@ -81,9 +80,8 @@ static TEE_Result operation_open_dfh(uint32_t id, unsigned int cmd,
 	struct mobj *mobj;
 	TEE_Result res;
 	void *va;
-	uint64_t cookie;
 
-	va = tee_fs_rpc_cache_alloc(TEE_FS_NAME_MAX, &mobj, &cookie);
+	va = tee_fs_rpc_cache_alloc(TEE_FS_NAME_MAX, &mobj);
 	if (!va)
 		return TEE_ERROR_OUT_OF_MEMORY;
 
@@ -91,7 +89,7 @@ static TEE_Result operation_open_dfh(uint32_t id, unsigned int cmd,
 	op.params[0].u.value.a = cmd;
 
 	if (!msg_param_init_memparam(op.params + 1, mobj, 0, TEE_FS_NAME_MAX,
-				     cookie, MSG_PARAM_MEM_DIR_IN))
+				     MSG_PARAM_MEM_DIR_IN))
 		return TEE_ERROR_BAD_STATE;
 
 	res = tee_svc_storage_create_filename_dfh(va, TEE_FS_NAME_MAX, dfh);
@@ -139,12 +137,11 @@ TEE_Result tee_fs_rpc_read_init(struct tee_fs_rpc_operation *op,
 {
 	struct mobj *mobj;
 	uint8_t *va;
-	uint64_t cookie;
 
 	if (offset < 0)
 		return TEE_ERROR_BAD_PARAMETERS;
 
-	va = tee_fs_rpc_cache_alloc(data_len, &mobj, &cookie);
+	va = tee_fs_rpc_cache_alloc(data_len, &mobj);
 	if (!va)
 		return TEE_ERROR_OUT_OF_MEMORY;
 
@@ -157,7 +154,7 @@ TEE_Result tee_fs_rpc_read_init(struct tee_fs_rpc_operation *op,
 	op->params[0].u.value.b = fd;
 	op->params[0].u.value.c = offset;
 
-	if (!msg_param_init_memparam(op->params + 1, mobj, 0, data_len, cookie,
+	if (!msg_param_init_memparam(op->params + 1, mobj, 0, data_len,
 				     MSG_PARAM_MEM_DIR_OUT))
 		return TEE_ERROR_BAD_STATE;
 
@@ -182,12 +179,11 @@ TEE_Result tee_fs_rpc_write_init(struct tee_fs_rpc_operation *op,
 {
 	struct mobj *mobj;
 	uint8_t *va;
-	uint64_t cookie;
 
 	if (offset < 0)
 		return TEE_ERROR_BAD_PARAMETERS;
 
-	va = tee_fs_rpc_cache_alloc(data_len, &mobj, &cookie);
+	va = tee_fs_rpc_cache_alloc(data_len, &mobj);
 	if (!va)
 		return TEE_ERROR_OUT_OF_MEMORY;
 
@@ -200,7 +196,7 @@ TEE_Result tee_fs_rpc_write_init(struct tee_fs_rpc_operation *op,
 	op->params[0].u.value.b = fd;
 	op->params[0].u.value.c = offset;
 
-	if (!msg_param_init_memparam(op->params + 1, mobj, 0, data_len, cookie,
+	if (!msg_param_init_memparam(op->params + 1, mobj, 0, data_len,
 				     MSG_PARAM_MEM_DIR_IN))
 		return TEE_ERROR_BAD_STATE;
 
@@ -232,9 +228,8 @@ TEE_Result tee_fs_rpc_remove(uint32_t id, struct tee_pobj *po)
 	struct tee_fs_rpc_operation op = { .id = id, .num_params = 2 };
 	struct mobj *mobj;
 	void *va;
-	uint64_t cookie;
 
-	va = tee_fs_rpc_cache_alloc(TEE_FS_NAME_MAX, &mobj, &cookie);
+	va = tee_fs_rpc_cache_alloc(TEE_FS_NAME_MAX, &mobj);
 	if (!va)
 		return TEE_ERROR_OUT_OF_MEMORY;
 
@@ -242,7 +237,7 @@ TEE_Result tee_fs_rpc_remove(uint32_t id, struct tee_pobj *po)
 	op.params[0].u.value.a = OPTEE_MRF_REMOVE;
 
 	if (!msg_param_init_memparam(op.params + 1, mobj, 0, TEE_FS_NAME_MAX,
-				     cookie, MSG_PARAM_MEM_DIR_IN))
+				     MSG_PARAM_MEM_DIR_IN))
 		return TEE_ERROR_BAD_STATE;
 
 	res = tee_svc_storage_create_filename(va, TEE_FS_NAME_MAX,
@@ -260,9 +255,8 @@ TEE_Result tee_fs_rpc_remove_dfh(uint32_t id,
 	struct tee_fs_rpc_operation op = { .id = id, .num_params = 2 };
 	struct mobj *mobj;
 	void *va;
-	uint64_t cookie;
 
-	va = tee_fs_rpc_cache_alloc(TEE_FS_NAME_MAX, &mobj, &cookie);
+	va = tee_fs_rpc_cache_alloc(TEE_FS_NAME_MAX, &mobj);
 	if (!va)
 		return TEE_ERROR_OUT_OF_MEMORY;
 
@@ -270,7 +264,7 @@ TEE_Result tee_fs_rpc_remove_dfh(uint32_t id,
 	op.params[0].u.value.a = OPTEE_MRF_REMOVE;
 
 	if (!msg_param_init_memparam(op.params + 1, mobj, 0, TEE_FS_NAME_MAX,
-				     cookie, MSG_PARAM_MEM_DIR_IN))
+				     MSG_PARAM_MEM_DIR_IN))
 		return TEE_ERROR_BAD_STATE;
 
 	res = tee_svc_storage_create_filename_dfh(va, TEE_FS_NAME_MAX, dfh);
@@ -287,10 +281,9 @@ TEE_Result tee_fs_rpc_rename(uint32_t id, struct tee_pobj *old,
 	struct tee_fs_rpc_operation op = { .id = id, .num_params = 3 };
 	struct mobj *mobj;
 	char *va;
-	uint64_t cookie;
 	bool temp;
 
-	va = tee_fs_rpc_cache_alloc(TEE_FS_NAME_MAX * 2, &mobj, &cookie);
+	va = tee_fs_rpc_cache_alloc(TEE_FS_NAME_MAX * 2, &mobj);
 	if (!va)
 		return TEE_ERROR_OUT_OF_MEMORY;
 
@@ -299,7 +292,7 @@ TEE_Result tee_fs_rpc_rename(uint32_t id, struct tee_pobj *old,
 	op.params[0].u.value.b = overwrite;
 
 	if (!msg_param_init_memparam(op.params + 1, mobj, 0, TEE_FS_NAME_MAX,
-				     cookie, MSG_PARAM_MEM_DIR_IN))
+				     MSG_PARAM_MEM_DIR_IN))
 		return TEE_ERROR_BAD_STATE;
 
 	if (new)
@@ -312,8 +305,7 @@ TEE_Result tee_fs_rpc_rename(uint32_t id, struct tee_pobj *old,
 		return res;
 
 	if (!msg_param_init_memparam(op.params + 2, mobj, TEE_FS_NAME_MAX,
-				     TEE_FS_NAME_MAX, cookie,
-				     MSG_PARAM_MEM_DIR_IN))
+				     TEE_FS_NAME_MAX, MSG_PARAM_MEM_DIR_IN))
 		return TEE_ERROR_BAD_STATE;
 
 	if (new) {
@@ -338,13 +330,12 @@ TEE_Result tee_fs_rpc_opendir(uint32_t id, const TEE_UUID *uuid,
 	struct tee_fs_rpc_operation op = { .id = id, .num_params = 3 };
 	struct mobj *mobj;
 	void *va;
-	uint64_t cookie;
 	struct tee_fs_dir *dir = calloc(1, sizeof(*dir));
 
 	if (!dir)
 		return TEE_ERROR_OUT_OF_MEMORY;
 
-	va = tee_fs_rpc_cache_alloc(TEE_FS_NAME_MAX, &mobj, &cookie);
+	va = tee_fs_rpc_cache_alloc(TEE_FS_NAME_MAX, &mobj);
 	if (!va) {
 		res = TEE_ERROR_OUT_OF_MEMORY;
 		goto err_exit;
@@ -354,7 +345,7 @@ TEE_Result tee_fs_rpc_opendir(uint32_t id, const TEE_UUID *uuid,
 	op.params[0].u.value.a = OPTEE_MRF_OPENDIR;
 
 	if (!msg_param_init_memparam(op.params + 1, mobj, 0, TEE_FS_NAME_MAX,
-				     cookie, MSG_PARAM_MEM_DIR_IN)) {
+				     MSG_PARAM_MEM_DIR_IN)) {
 		res = TEE_ERROR_BAD_STATE;
 		goto err_exit;
 	}
@@ -399,13 +390,12 @@ TEE_Result tee_fs_rpc_readdir(uint32_t id, struct tee_fs_dir *d,
 	struct tee_fs_rpc_operation op = { .id = id, .num_params = 2 };
 	struct mobj *mobj;
 	void *va;
-	uint64_t cookie;
 	const size_t max_name_len = TEE_FS_NAME_MAX + 1;
 
 	if (!d)
 		return TEE_ERROR_ITEM_NOT_FOUND;
 
-	va = tee_fs_rpc_cache_alloc(max_name_len, &mobj, &cookie);
+	va = tee_fs_rpc_cache_alloc(max_name_len, &mobj);
 	if (!va)
 		return TEE_ERROR_OUT_OF_MEMORY;
 
@@ -414,7 +404,7 @@ TEE_Result tee_fs_rpc_readdir(uint32_t id, struct tee_fs_dir *d,
 	op.params[0].u.value.b = d->nw_dir;
 
 	if (!msg_param_init_memparam(op.params + 1, mobj, 0, max_name_len,
-				     cookie, MSG_PARAM_MEM_DIR_OUT))
+				     MSG_PARAM_MEM_DIR_OUT))
 		return TEE_ERROR_BAD_STATE;
 
 	res = operation_commit(&op);
diff --git a/core/tee/tee_fs_rpc_cache.c b/core/tee/tee_fs_rpc_cache.c
index d1fc51c0..3148aea2 100644
--- a/core/tee/tee_fs_rpc_cache.c
+++ b/core/tee/tee_fs_rpc_cache.c
@@ -12,20 +12,17 @@
 void tee_fs_rpc_cache_clear(struct thread_specific_data *tsd)
 {
 	if (tsd->rpc_fs_payload) {
-		thread_rpc_free_payload(tsd->rpc_fs_payload_cookie,
-					tsd->rpc_fs_payload_mobj);
+		thread_rpc_free_payload(tsd->rpc_fs_payload_mobj);
 		tsd->rpc_fs_payload = NULL;
-		tsd->rpc_fs_payload_cookie = 0;
 		tsd->rpc_fs_payload_size = 0;
 		tsd->rpc_fs_payload_mobj = NULL;
 	}
 }
 
-void *tee_fs_rpc_cache_alloc(size_t size, struct mobj **mobj, uint64_t *cookie)
+void *tee_fs_rpc_cache_alloc(size_t size, struct mobj **mobj)
 {
 	struct thread_specific_data *tsd = thread_get_tsd();
 	size_t sz = size;
-	uint64_t c = 0;
 	paddr_t p;
 	void *va;
 
@@ -41,7 +38,7 @@ void *tee_fs_rpc_cache_alloc(size_t size, struct mobj **mobj, uint64_t *cookie)
 	if (sz > tsd->rpc_fs_payload_size) {
 		tee_fs_rpc_cache_clear(tsd);
 
-		*mobj = thread_rpc_alloc_payload(sz,  &c);
+		*mobj = thread_rpc_alloc_payload(sz);
 		if (!*mobj)
 			return NULL;
 
@@ -57,14 +54,12 @@ void *tee_fs_rpc_cache_alloc(size_t size, struct mobj **mobj, uint64_t *cookie)
 
 		tsd->rpc_fs_payload = va;
 		tsd->rpc_fs_payload_mobj = *mobj;
-		tsd->rpc_fs_payload_cookie = c;
 		tsd->rpc_fs_payload_size = sz;
 	} else
 		*mobj = tsd->rpc_fs_payload_mobj;
 
-	*cookie = tsd->rpc_fs_payload_cookie;
 	return tsd->rpc_fs_payload;
 err:
-	thread_rpc_free_payload(c, *mobj);
+	thread_rpc_free_payload(*mobj);
 	return NULL;
 }
diff --git a/core/tee/tee_rpmb_fs.c b/core/tee/tee_rpmb_fs.c
index 7ae604d6..e100fee0 100644
--- a/core/tee/tee_rpmb_fs.c
+++ b/core/tee/tee_rpmb_fs.c
@@ -410,9 +410,7 @@ func_exit:
 
 struct tee_rpmb_mem {
 	struct mobj *phreq_mobj;
-	uint64_t phreq_cookie;
 	struct mobj *phresp_mobj;
-	uint64_t phresp_cookie;
 	size_t req_size;
 	size_t resp_size;
 };
@@ -423,13 +421,11 @@ static void tee_rpmb_free(struct tee_rpmb_mem *mem)
 		return;
 
 	if (mem->phreq_mobj) {
-		thread_rpc_free_payload(mem->phreq_cookie, mem->phreq_mobj);
-		mem->phreq_cookie = 0;
+		thread_rpc_free_payload(mem->phreq_mobj);
 		mem->phreq_mobj = NULL;
 	}
 	if (mem->phresp_mobj) {
-		thread_rpc_free_payload(mem->phresp_cookie, mem->phresp_mobj);
-		mem->phresp_cookie = 0;
+		thread_rpc_free_payload(mem->phresp_mobj);
 		mem->phresp_mobj = NULL;
 	}
 }
@@ -447,9 +443,8 @@ static TEE_Result tee_rpmb_alloc(size_t req_size, size_t resp_size,
 
 	memset(mem, 0, sizeof(*mem));
 
-	mem->phreq_mobj = thread_rpc_alloc_payload(req_s, &mem->phreq_cookie);
-	mem->phresp_mobj =
-		thread_rpc_alloc_payload(resp_s, &mem->phresp_cookie);
+	mem->phreq_mobj = thread_rpc_alloc_payload(req_s);
+	mem->phresp_mobj = thread_rpc_alloc_payload(resp_s);
 
 	if (!mem->phreq_mobj || !mem->phresp_mobj) {
 		res = TEE_ERROR_OUT_OF_MEMORY;
@@ -479,13 +474,11 @@ static TEE_Result tee_rpmb_invoke(struct tee_rpmb_mem *mem)
 	memset(params, 0, sizeof(params));
 
 	if (!msg_param_init_memparam(params + 0, mem->phreq_mobj, 0,
-				     mem->req_size, mem->phreq_cookie,
-				     MSG_PARAM_MEM_DIR_IN))
+				     mem->req_size, MSG_PARAM_MEM_DIR_IN))
 		return TEE_ERROR_BAD_STATE;
 
 	if (!msg_param_init_memparam(params + 1, mem->phresp_mobj, 0,
-				     mem->resp_size, mem->phresp_cookie,
-				     MSG_PARAM_MEM_DIR_OUT))
+				     mem->resp_size, MSG_PARAM_MEM_DIR_OUT))
 		return TEE_ERROR_BAD_STATE;
 
 	return thread_rpc_cmd(OPTEE_MSG_RPC_CMD_RPMB, 2, params);