diff --git a/security/keys/trusted-keys/trusted_tee.c b/security/keys/trusted-keys/trusted_tee.c
index ac3e270ade69..3085343c489a 100644
--- a/security/keys/trusted-keys/trusted_tee.c
+++ b/security/keys/trusted-keys/trusted_tee.c
@@ -8,6 +8,7 @@
 
 #include <linux/err.h>
 #include <linux/key-type.h>
+#include <linux/minmax.h>
 #include <linux/module.h>
 #include <linux/slab.h>
 #include <linux/string.h>
@@ -65,38 +66,37 @@ static int trusted_tee_seal(struct trusted_key_payload *p, char *datablob)
 	int ret;
 	struct tee_ioctl_invoke_arg inv_arg;
 	struct tee_param param[4];
-	struct tee_shm *reg_shm_in = NULL, *reg_shm_out = NULL;
+	struct tee_shm *shm;
+	uint8_t *buf;
 
 	memset(&inv_arg, 0, sizeof(inv_arg));
 	memset(&param, 0, sizeof(param));
 
-	reg_shm_in = tee_shm_register_kernel_buf(pvt_data.ctx, p->key,
-						 p->key_len);
-	if (IS_ERR(reg_shm_in)) {
-		dev_err(pvt_data.dev, "key shm register failed\n");
-		return PTR_ERR(reg_shm_in);
+	shm = tee_shm_alloc_kernel_buf(pvt_data.ctx,
+				       p->key_len + sizeof(p->blob));
+	if (IS_ERR(shm)) {
+		dev_err(pvt_data.dev, "key shm alloc failed\n");
+		return PTR_ERR(shm);
 	}
-
-	reg_shm_out = tee_shm_register_kernel_buf(pvt_data.ctx, p->blob,
-						  sizeof(p->blob));
-	if (IS_ERR(reg_shm_out)) {
-		dev_err(pvt_data.dev, "blob shm register failed\n");
-		ret = PTR_ERR(reg_shm_out);
+	buf = tee_shm_get_va(shm, 0);
+	if (IS_ERR(buf)) {
+		ret = PTR_ERR(buf);
 		goto out;
 	}
+	memcpy(buf, p->key, p->key_len);
 
 	inv_arg.func = TA_CMD_SEAL;
 	inv_arg.session = pvt_data.session_id;
 	inv_arg.num_params = 4;
 
 	param[0].attr = TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_INPUT;
-	param[0].u.memref.shm = reg_shm_in;
+	param[0].u.memref.shm = shm;
 	param[0].u.memref.size = p->key_len;
 	param[0].u.memref.shm_offs = 0;
 	param[1].attr = TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_OUTPUT;
-	param[1].u.memref.shm = reg_shm_out;
+	param[1].u.memref.shm = shm;
 	param[1].u.memref.size = sizeof(p->blob);
-	param[1].u.memref.shm_offs = 0;
+	param[1].u.memref.shm_offs = p->key_len;
 
 	ret = tee_client_invoke_func(pvt_data.ctx, &inv_arg, param);
 	if ((ret < 0) || (inv_arg.ret != 0)) {
@@ -104,14 +104,13 @@ static int trusted_tee_seal(struct trusted_key_payload *p, char *datablob)
 			inv_arg.ret);
 		ret = -EFAULT;
 	} else {
+		memcpy(p->blob, buf + p->key_len,
+		       min(param[1].u.memref.size, sizeof(p->blob)));
 		p->blob_len = param[1].u.memref.size;
 	}
 
 out:
-	if (reg_shm_out)
-		tee_shm_free(reg_shm_out);
-	if (reg_shm_in)
-		tee_shm_free(reg_shm_in);
+	tee_shm_free(shm);
 
 	return ret;
 }
@@ -166,11 +165,9 @@ static int trusted_tee_unseal(struct trusted_key_payload *p, char *datablob)
 		p->key_len = param[1].u.memref.size;
 	}
 
+	tee_shm_free(reg_shm_out);
 out:
-	if (reg_shm_out)
-		tee_shm_free(reg_shm_out);
-	if (reg_shm_in)
-		tee_shm_free(reg_shm_in);
+	tee_shm_free(reg_shm_in);
 
 	return ret;
 }
@@ -183,15 +180,21 @@ static int trusted_tee_get_random(unsigned char *key, size_t key_len)
 	int ret;
 	struct tee_ioctl_invoke_arg inv_arg;
 	struct tee_param param[4];
-	struct tee_shm *reg_shm = NULL;
+	struct tee_shm *shm;
+	void *buf;
 
 	memset(&inv_arg, 0, sizeof(inv_arg));
 	memset(&param, 0, sizeof(param));
 
-	reg_shm = tee_shm_register_kernel_buf(pvt_data.ctx, key, key_len);
-	if (IS_ERR(reg_shm)) {
-		dev_err(pvt_data.dev, "key shm register failed\n");
-		return PTR_ERR(reg_shm);
+	shm = tee_shm_alloc_kernel_buf(pvt_data.ctx, key_len);
+	if (IS_ERR(shm)) {
+		dev_err(pvt_data.dev, "key shm alloc failed\n");
+		return PTR_ERR(shm);
+	}
+	buf = tee_shm_get_va(shm, 0);
+	if (IS_ERR(buf)) {
+		ret = PTR_ERR(buf);
+		goto out;
 	}
 
 	inv_arg.func = TA_CMD_GET_RANDOM;
@@ -199,7 +202,7 @@ static int trusted_tee_get_random(unsigned char *key, size_t key_len)
 	inv_arg.num_params = 4;
 
 	param[0].attr = TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_OUTPUT;
-	param[0].u.memref.shm = reg_shm;
+	param[0].u.memref.shm = shm;
 	param[0].u.memref.size = key_len;
 	param[0].u.memref.shm_offs = 0;
 
@@ -209,18 +212,19 @@ static int trusted_tee_get_random(unsigned char *key, size_t key_len)
 			inv_arg.ret);
 		ret = -EFAULT;
 	} else {
+		memcpy(key, buf, min(param[0].u.memref.size, key_len));
 		ret = param[0].u.memref.size;
 	}
 
-	tee_shm_free(reg_shm);
+out:
+	tee_shm_free(shm);
 
 	return ret;
 }
 
 static int optee_ctx_match(struct tee_ioctl_version_data *ver, const void *data)
 {
-	if (ver->impl_id == TEE_IMPL_ID_OPTEE &&
-	    ver->gen_caps & TEE_GEN_CAP_REG_MEM)
+	if (ver->impl_id == TEE_IMPL_ID_OPTEE)
 		return 1;
 	else
 		return 0;
