From ea1f89ed4614a28d756d0b5c9a6963fc64b511cb Mon Sep 17 00:00:00 2001
From: xulu <xulu@allwinnertech.com>
Date: Wed, 22 Sep 2021 17:13:08 +0800
Subject: [PATCH] optee: add optee mbedtls rsa demo.

Signed-off-by: xulu <xulu@allwinnertech.com>
Change-Id: Iac38b8c583da4ad1254695749ffa8753c8d47223
---
 security/optee-rsa-demo/Makefile              | 128 ++++++++
 security/optee-rsa-demo/src/Makefile          |  70 +++++
 security/optee-rsa-demo/src/na/Makefile       | 102 +++++++
 security/optee-rsa-demo/src/na/liboptee_rsa.c | 207 +++++++++++++
 security/optee-rsa-demo/src/na/liboptee_rsa.h |  28 ++
 .../optee-rsa-demo/src/na/optee-rsa-demo.c    |  64 ++++
 security/optee-rsa-demo/src/scripts/common.mk |   9 +
 security/optee-rsa-demo/src/ta/Makefile       |   2 +
 .../src/ta/include/rsa-demo_ta.h              |  12 +
 .../src/ta/include/user_ta_header_defines.h   |  16 +
 .../optee-rsa-demo/src/ta/optee_rsa-demo_ta.c | 275 ++++++++++++++++++
 .../src/ta/optee_rsa-demo_ta.c-bak            | 272 +++++++++++++++++
 security/optee-rsa-demo/src/ta/sub.mk         |   2 +
 security/optee-rsa-demo/src/ta/ta_common.mk   |   3 +
 14 files changed, 1190 insertions(+)
 create mode 100755 security/optee-rsa-demo/Makefile
 create mode 100644 security/optee-rsa-demo/src/Makefile
 create mode 100644 security/optee-rsa-demo/src/na/Makefile
 create mode 100644 security/optee-rsa-demo/src/na/liboptee_rsa.c
 create mode 100644 security/optee-rsa-demo/src/na/liboptee_rsa.h
 create mode 100644 security/optee-rsa-demo/src/na/optee-rsa-demo.c
 create mode 100644 security/optee-rsa-demo/src/scripts/common.mk
 create mode 100644 security/optee-rsa-demo/src/ta/Makefile
 create mode 100644 security/optee-rsa-demo/src/ta/include/rsa-demo_ta.h
 create mode 100644 security/optee-rsa-demo/src/ta/include/user_ta_header_defines.h
 create mode 100644 security/optee-rsa-demo/src/ta/optee_rsa-demo_ta.c
 create mode 100644 security/optee-rsa-demo/src/ta/optee_rsa-demo_ta.c-bak
 create mode 100644 security/optee-rsa-demo/src/ta/sub.mk
 create mode 100644 security/optee-rsa-demo/src/ta/ta_common.mk

diff --git a/security/optee-rsa-demo/Makefile b/security/optee-rsa-demo/Makefile
new file mode 100755
index 000000000..5fba4d36a
--- /dev/null
+++ b/security/optee-rsa-demo/Makefile
@@ -0,0 +1,128 @@
+##############################################
+# OpenWrt Makefile for rsa-demo program
+#
+#
+# Most of the variables used here are defined in
+# the include directives below. We just need to
+# specify a basic description of the package,
+# where to build our program, where to find
+# the source files, and where to install the
+# compiled program on the router.
+#
+# Be very careful of spacing in this file.
+# Indents should be tabs, not spaces, and
+# there should be no trailing whitespace in
+# lines that are not commented.
+#
+##############################################
+include $(TOPDIR)/rules.mk
+include $(BUILD_DIR)/kernel.mk
+
+# Name and release number of this package
+PKG_NAME:=optee-rsa-demo
+PKG_RELEASE:=1
+
+# This specifies the directory where we're going to build the program.
+# The root build directory, $(BUILD_DIR), is by default the build_mipsel
+# directory in your OpenWrt SDK directory
+PKG_BUILD_DIR := $(COMPILE_DIR)/$(PKG_NAME)
+
+PKG_CONFIG_DEPENDS += CONFIG_OPTEE_ENCRYPT_TA CONFIG_OPTEE_ENCRYPT_TA_SSK_KEY CONFIG_OPTEE_ENCRYPT_TA_ROTPK_DERIVE_KEY
+
+PKG_FILE_DEPENDS += $(STAGING_DIR)/usr/dev_kit/arm-plat-$(TARGET_CHIP)/
+
+include $(BUILD_DIR)/package.mk
+
+# Specify package information for this program.
+# The variables defined here should be self explanatory.
+# If you are running Kamikaze, delete the DESCRIPTION
+# variable below and uncomment the Kamikaze define
+# directive for the description below
+define Package/$(PKG_NAME)
+  SECTION:=utils
+  CATEGORY:=Security
+  SUBMENU:=OPTEE
+  TITLE:=OPTEE rsa-demo
+  DEPENDS:=+optee-client +liballwinner-base +optee-os-dev-kit
+endef
+
+
+# Uncomment portion below for Kamikaze and delete DESCRIPTION variable above
+define Package/$(PKG_NAME)/description
+	rsa-demo TA & CA demo source code!
+endef
+
+# Specify what needs to be done to prepare for building the package.
+# In our case, we need to copy the source files to the build directory.
+# This is NOT the default.  The default uses the PKG_SOURCE_URL and the
+# PKG_SOURCE which is not defined here to download the source from the web.
+# In order to just build a simple program that we have just written, it is
+# much easier to do it this way.
+define Build/Prepare
+	mkdir -p $(PKG_BUILD_DIR)
+	$(CP) -r ./src/* $(PKG_BUILD_DIR)/
+endef
+
+define Build/Configure
+endef
+
+# TA MUST be compiled with 32bit cross toolchain.
+ifeq ($(CONFIG_ARCH_64BIT),y)
+TARGET_CROSS_TA=$(subst aarch64,arm,$(CONFIG_TOOLCHAIN_ROOT)/bin/arm-openwrt-linux-)
+else
+TARGET_CROSS_TA:=$(TARGET_CROSS)
+endif
+
+define Build/Compile/Source
+	$(MAKE) -C $(PKG_BUILD_DIR)/ \
+		ARCH="$(TARGET_ARCH)" \
+		AR="$(TARGET_AR)" \
+		CC="$(TARGET_CC)" \
+		CROSS_COMPILE_TA="$(TARGET_CROSS_TA)" \
+		CROSS_COMPILE_HOST="$(TARGET_CROSS)" \
+		PLATFORM="$(TARGET_CHIP)" \
+		TA_LDFLAGS="$(TARGET_LDFLAGS)"\
+		DEV_KIT_DIR="$(STAGING_DIR)/usr/dev_kit"\
+		CA_DEV_KIT_DIR="$(STAGING_DIR)/usr"
+endef
+
+define Build/Compile
+	$(call Build/Compile/Source)
+endef
+
+# We do not need to define Build/Configure or Build/Compile directives
+# The defaults are appropriate for compiling a simple program such as this one
+
+
+# Specify where and how to install the program. Since we only have one file,
+# the rsa-demo executable, install it by copying it to the /bin directory on
+# the router. The $(1) variable represents the root directory on the router running
+# OpenWrt. The $(INSTALL_DIR) variable contains a command to prepare the install
+# directory if it does not already exist.  Likewise $(INSTALL_BIN) contains the
+# command to copy the binary file from its current location (in our case the build
+# directory) to the install directory.
+define Build/InstallDev
+	$(INSTALL_DIR) $(1)/usr/include/
+	$(INSTALL_DIR) $(1)/usr/lib
+
+	$(CP) $(PKG_BUILD_DIR)/na/liboptee_rsa.h $(1)/usr/include
+	$(CP) $(PKG_BUILD_DIR)/out/na/lib/liboptee_rsa.so $(1)/usr/lib
+endef
+
+define Package/$(PKG_NAME)/install
+	$(INSTALL_DIR) $(1)/usr/bin
+	$(INSTALL_DIR) $(1)/usr/lib
+	$(INSTALL_DIR) $(1)/lib/optee_armtz
+
+	$(CP) $(PKG_BUILD_DIR)/out/na/lib/liboptee_rsa.so $(1)/usr/lib
+	$(CP) $(PKG_BUILD_DIR)/out/na/*_na $(1)/usr/bin
+	$(CP) $(PKG_BUILD_DIR)/out/ta/*.ta $(1)/lib/optee_armtz
+endef
+
+
+# This line executes the necessary commands to compile our program.
+# The above define directives specify all the information needed, but this
+# line calls BuildPackage which in turn actually uses this information to
+# build a package.
+
+$(eval $(call BuildPackage,$(PKG_NAME)))
diff --git a/security/optee-rsa-demo/src/Makefile b/security/optee-rsa-demo/src/Makefile
new file mode 100644
index 000000000..c16ed7a95
--- /dev/null
+++ b/security/optee-rsa-demo/src/Makefile
@@ -0,0 +1,70 @@
+out-dir := $(CURDIR)/out
+
+DEV_KIT_DIR ?= $(CURDIR)/../../dev_kit
+#include $(DEV_KIT_DIR)/../platform_config.mk
+CA_DEV_KIT_DIR ?= $(DEV_KIT_DIR)/arm-plat-$(PLATFORM)/export-ca
+TA_DEV_KIT_DIR ?= $(DEV_KIT_DIR)/arm-plat-$(PLATFORM)/export-ta_arm32
+
+ifeq ($(ATF_EXIST),y)
+CROSS_COMPILE_HOST ?= $(DEV_KIT_DIR)/../tools/toolchain/gcc-aarch64/bin/aarch64-linux-gnu-
+CROSS_COMPILE_TA   ?= $(DEV_KIT_DIR)/../tools/toolchain/gcc-arm-gnueabi/bin/arm-linux-gnueabi-
+else
+CROSS_COMPILE_HOST ?= $(DEV_KIT_DIR)/../tools/toolchain/gcc-arm-gnueabi/bin/arm-linux-gnueabi-
+CROSS_COMPILE_TA   ?= $(DEV_KIT_DIR)/../tools/toolchain/gcc-arm-gnueabi/bin/arm-linux-gnueabi-
+endif
+
+export CA_DEV_KIT_DIR
+export TA_DEV_KIT_DIR
+
+ifneq ($V,1)
+	q := @
+	echo := @echo
+else
+	q :=
+	echo := @:
+endif
+# export 'q', used by sub-makefiles.
+export q
+
+# If _HOST or _TA specific compilers are not specified, then use CROSS_COMPILE
+CROSS_COMPILE_HOST ?= $(CROSS_COMPILE)
+CROSS_COMPILE_TA ?= $(CROSS_COMPILE)
+
+CFLAGS+=$(TA_CFLAGS)
+LDFLAGS+=$(TA_LDFLAGS)
+
+.PHONY: all
+ifneq ($(wildcard $(TA_DEV_KIT_DIR)/host_include/conf.mk),)
+all: na ta
+else
+all:
+	$(q)echo "TA_DEV_KIT_DIR is not correctly defined" && false
+endif
+
+.PHONY: na
+na:
+	$(q)$(MAKE) -C na CROSS_COMPILE="$(CROSS_COMPILE_HOST)" \
+			     --no-builtin-variables \
+			     O=$(out-dir)/na
+
+.PHONY: ta
+ta:
+	$(q)$(MAKE) -C ta CROSS_COMPILE="$(CROSS_COMPILE_TA)" \
+			     O=$(out-dir)/ta
+
+.PHONY: clean
+RMDIR := rmdir --ignore-fail-on-non-empty
+define rm-build-subdir
+echo `[ -d $1 ] && { cd $1; $(RMDIR) $(out-dir)/ta $(out-dir) 2> /dev/null; true; }` > /dev/null
+endef
+
+ifneq ($(wildcard $(TA_DEV_KIT_DIR)/host_include/conf.mk),)
+clean:
+	-$(q)$(MAKE) -C na O=$(out-dir)/na $@
+	-$(q)$(MAKE) -C ta O=$(out-dir)/ta $@
+	-$(q)$(call rm-build-subdir,ta)
+else
+clean:
+	$(q)echo "TA_DEV_KIT_DIR is not correctly defined"
+	$(q)echo "You can remove manually $(out-dir)"
+endif
diff --git a/security/optee-rsa-demo/src/na/Makefile b/security/optee-rsa-demo/src/na/Makefile
new file mode 100644
index 000000000..e1401e8c1
--- /dev/null
+++ b/security/optee-rsa-demo/src/na/Makefile
@@ -0,0 +1,102 @@
+# Normally this makefile shouldn't be called directly and we expect the output
+# path to be on a certain location to fit together with the other OP-TEE
+# gits and helper scripts.
+
+include ../scripts/common.mk
+out-dir := $(call strip-trailing-slashes-and-dots,$(O))
+ifeq ($(out-dir),)
+$(error invalid output directory (O=$(O)))
+endif
+
+#include $(TA_DEV_KIT_DIR)/host_include/conf.mk
+
+# By default we expect optee_client exported folder to be on a certain relative
+# path, but if the client specifies the OPTEE_CLIENT_EXPORT then that path will
+# be used instead.
+OPTEE_CLIENT_EXPORT ?= $(CA_DEV_KIT_DIR)
+
+CC		?= $(CROSS_COMPILE)gcc
+CPP		?= $(CROSS_COMPILE)cpp
+LD		?= $(CROSS_COMPILE)ld
+AR		?= $(CROSS_COMPILE)ar
+NM		?= $(CROSS_COMPILE)nm
+OBJCOPY		?= $(CROSS_COMPILE)objcopy
+OBJDUMP		?= $(CROSS_COMPILE)objdump
+READELF		?= $(CROSS_COMPILE)readelf
+
+
+srcs +=	optee-rsa-demo.c
+
+objs 	:= $(patsubst %.c,$(out-dir)/%.o, $(srcs))
+
+CFLAGS += -I./
+CFLAGS += -I./../ta/include
+
+CFLAGS += -I$(OPTEE_CLIENT_EXPORT)/include
+
+ifndef CFG_GP_PACKAGE_PATH
+CFLAGS += -Wall -Wcast-align -Werror \
+	  -Werror-implicit-function-declaration -Wextra -Wfloat-equal \
+	  -Wformat-nonliteral -Wformat-security -Wformat=2 -Winit-self \
+	  -Wmissing-declarations -Wmissing-format-attribute \
+	  -Wmissing-include-dirs -Wmissing-noreturn \
+	  -Wmissing-prototypes -Wnested-externs -Wpointer-arith \
+	  -Wshadow -Wstrict-prototypes -Wswitch-default \
+	  -Wwrite-strings \
+	  -Wno-missing-field-initializers -Wno-format-zero-length
+endif
+
+CFLAGS += -g3 -fPIC
+LDFLAGS += -L$(OPTEE_CLIENT_EXPORT)/exportlib -lteec
+LDFLAGS += -lpthread -lm
+
+################################ LIBROTPK
+lib_srcs +=	liboptee_rsa.c
+lib_objs 	:= $(patsubst %.c,$(out-dir)/%.o, $(lib_srcs))
+
+LIB_OUT_DIR = $(out-dir)/lib
+LIB_CFLAGS = $(CFLAGS)
+LIB_CFLAGS += -I./
+LIB_CFLAGS += -I./../ta/include
+LIB_CFLAGS += -I$(OPTEE_CLIENT_EXPORT)/include
+LIB_LDFLAGS = $(LDFLAGS)
+LIB_LDFLAGS += -L$(OPTEE_CLIENT_EXPORT)/exportlib -lteec
+
+LIB_NAME = liboptee_rsa.so
+LIB_CFLAGS += --no-builtin-variables
+LIB_CFLAGS += -fPIC
+LIB_LDFLAGS += -lpthread -lm
+LIB_LDFLAGS += -shared
+################################ LIBROTPK
+
+.PHONY: all
+all: $(LIB_NAME) rsa-demo_na
+
+$(LIB_NAME):
+	mkdir -p $(LIB_OUT_DIR)
+	$(CC) -c $(LIB_CFLAGS) -o $(lib_objs) $(lib_srcs)
+	$(CC) $(lib_objs) $(LIB_CFLAGS) $(LIB_LDFLAGS) -o $@
+	cp -rf $@ $(LIB_OUT_DIR)/$(LIB_NAME)
+
+rsa-demo_na: $(objs)
+	@echo "  LD      $(out-dir)/$@"
+	$(q)@$(CC) -o $(out-dir)/$@ $+ $(LDFLAGS) -L$(LIB_OUT_DIR) -loptee_rsa
+
+$(out-dir)/%.o: $(CURDIR)/%.c
+	$(q)mkdir -p $(out-dir)/
+	@echo '  CC      $<'
+	$(q)$(CC) $(CFLAGS) -c $< -o $@
+
+RMDIR := rmdir --ignore-fail-on-non-empty
+define rm-build-dirs
+	$(q)$(RMDIR) $(out-dir) 2> /dev/null; true
+	$(q)$(RMDIR) $(out-dir) 2> /dev/null; true
+endef
+
+.PHONY: clean
+clean:
+	@echo '  CLEAN $(out-dir)'
+	$(q)rm -rf $(out-dir)/
+	$(q)$(foreach obj,$(lib_objs), rm -f $(obj))
+	$(q)$(foreach obj,$(objs), rm -f $(obj))
+	$(call rm-build-dirs)
diff --git a/security/optee-rsa-demo/src/na/liboptee_rsa.c b/security/optee-rsa-demo/src/na/liboptee_rsa.c
new file mode 100644
index 000000000..5ff6e3be6
--- /dev/null
+++ b/security/optee-rsa-demo/src/na/liboptee_rsa.c
@@ -0,0 +1,207 @@
+#include <tee_client_api.h>
+#include <rsa-demo_ta.h>
+#include <stdio.h>
+#include <string.h>
+#include <ctype.h>
+
+#include "liboptee_rsa.h"
+
+
+//#define DBG_ON
+
+#ifdef DBG_ON
+#define PR_ERR(...)	printf(__VA_ARGS__)
+#define PR_DBG(...)	printf(__VA_ARGS__)
+#else
+#define PR_ERR(...)	printf(__VA_ARGS__)
+#define PR_DBG(...)	(void)0
+#endif
+
+#ifdef DBG_ON
+void sunxi_dump(uint8_t *buf, uint32_t ttl_len)
+{
+	uint32_t len;
+
+	for (len = 0; len < ttl_len; len++) {
+		printf("0x%02x ", ((char *)buf)[len]);
+		if (len % 8 == 7) {
+			printf("\n");
+		}
+	}
+	printf("\n");
+}
+
+#else
+void sunxi_dump(uint8_t *buf, uint32_t ttl_len)
+{
+	(void)buf;
+	(void)ttl_len;
+}
+#endif
+
+TEEC_Context ctx;
+TEEC_Session sess;
+static const TEEC_UUID ta_uuid = RSA_DEMO_UUID;
+
+int rsa_demo_init(void)
+{
+	TEEC_Result teecErr = TEEC_ERROR_GENERIC;
+
+	int ret = 0;
+
+	/* init a context connecting to the TEE */
+	PR_DBG("NA: init context\n");
+	teecErr = TEEC_InitializeContext(NULL, &ctx);
+	if (teecErr != TEEC_SUCCESS) {
+		PR_ERR("NA: failed to init context: 0x%x\n", teecErr);
+		ret = -1;
+		return ret;
+	}
+
+	/* open a session to the ta */
+	PR_DBG("NA: open session\n");
+	teecErr = TEEC_OpenSession(&ctx, &sess, &ta_uuid,
+				TEEC_LOGIN_PUBLIC, NULL, NULL, NULL);
+	if (teecErr != TEEC_SUCCESS) {
+		PR_ERR("NA: failed to open session: 0x%x\n", teecErr);
+		ret = -1;
+		goto exit;
+	}
+
+	return ret;
+
+exit:
+	TEEC_FinalizeContext(&ctx);
+
+	return ret;
+}
+
+int rsa_demo_encrypt(uint8_t *in_buf, uint32_t in_len,
+		     uint8_t *out_buf, uint32_t *out_len)
+{
+	TEEC_Result teecErr = TEEC_ERROR_GENERIC;
+	TEEC_Operation operation;
+
+	int command;
+
+	TEEC_SharedMemory tee_params_p;
+	tee_params_p.size = in_len;
+	tee_params_p.flags = TEEC_MEM_INPUT;
+	teecErr = TEEC_AllocateSharedMemory(&ctx, &tee_params_p);
+	if (teecErr != TEEC_SUCCESS)
+		goto out;
+
+	TEEC_SharedMemory tee_params_q;
+	tee_params_q.size = RSA_KEY_SIZE / 8;
+	tee_params_q.flags = TEEC_MEM_INPUT | TEEC_MEM_OUTPUT;
+	teecErr = TEEC_AllocateSharedMemory(&ctx, &tee_params_q);
+	if (teecErr != TEEC_SUCCESS)
+		goto out1;
+
+	memset(&operation, 0, sizeof(TEEC_Operation));
+	operation.paramTypes =
+	TEEC_PARAM_TYPES(TEEC_MEMREF_WHOLE, TEEC_MEMREF_WHOLE,
+					TEEC_VALUE_INOUT, TEEC_NONE);
+	operation.started = 1;
+
+	operation.params[0].memref.parent = &tee_params_p;
+	operation.params[0].memref.offset = 0;
+	operation.params[0].memref.size   = 0;
+
+	operation.params[1].memref.parent = &tee_params_q;
+	operation.params[1].memref.offset = 0;
+	operation.params[1].memref.size   = 0;
+
+	memcpy(tee_params_p.buffer, in_buf, in_len);
+	operation.params[2].value.a = in_len;
+	command = CMD_RSA_DEMO_ENCRYPT;
+
+	PR_DBG("NA: invoke command: %d\n", command);
+	teecErr = TEEC_InvokeCommand(&sess, command, &operation, NULL);
+	if (teecErr == TEEC_SUCCESS) {
+		*out_len = operation.params[2].value.b;
+		memcpy(out_buf, tee_params_q.buffer, *out_len);
+	} else {
+		PR_ERR("NA: invoke command(%d) failed: %d\n", command, teecErr);
+	}
+
+	TEEC_ReleaseSharedMemory(&tee_params_q);
+
+out1:
+	TEEC_ReleaseSharedMemory(&tee_params_p);
+out:
+	if (teecErr)
+		PR_ERR("NA: %s finish with %d\n", __func__, teecErr);
+
+	return teecErr;
+}
+
+int rsa_demo_decrypt(uint8_t *in_buf, uint32_t in_len,
+		     uint8_t *out_buf, uint32_t *out_len)
+{
+	TEEC_Result teecErr = TEEC_ERROR_GENERIC;
+	TEEC_Operation operation;
+
+	int command;
+
+	TEEC_SharedMemory tee_params_p;
+	tee_params_p.size = in_len;
+	tee_params_p.flags = TEEC_MEM_INPUT;
+	teecErr = TEEC_AllocateSharedMemory(&ctx, &tee_params_p);
+	if (teecErr != TEEC_SUCCESS)
+		goto out;
+
+	TEEC_SharedMemory tee_params_q;
+	tee_params_q.size = RSA_KEY_SIZE / 8;
+	tee_params_q.flags = TEEC_MEM_INPUT | TEEC_MEM_OUTPUT;
+	teecErr = TEEC_AllocateSharedMemory(&ctx, &tee_params_q);
+	if (teecErr != TEEC_SUCCESS)
+		goto out1;
+
+	memset(&operation, 0, sizeof(TEEC_Operation));
+	operation.paramTypes =
+	TEEC_PARAM_TYPES(TEEC_MEMREF_WHOLE, TEEC_MEMREF_WHOLE,
+					TEEC_VALUE_INOUT, TEEC_NONE);
+	operation.started = 1;
+
+	operation.params[0].memref.parent = &tee_params_p;
+	operation.params[0].memref.offset = 0;
+	operation.params[0].memref.size   = 0;
+
+	operation.params[1].memref.parent = &tee_params_q;
+	operation.params[1].memref.offset = 0;
+	operation.params[1].memref.size   = 0;
+
+	memcpy(tee_params_p.buffer, in_buf, in_len);
+	operation.params[2].value.a = in_len;
+	command = CMD_RSA_DEMO_DECRYPT;
+
+	PR_DBG("NA: invoke command: %d\n", command);
+	teecErr = TEEC_InvokeCommand(&sess, command, &operation, NULL);
+	if (teecErr == TEEC_SUCCESS) {
+		*out_len = operation.params[2].value.b;
+		memcpy(out_buf, tee_params_q.buffer, *out_len);
+	} else {
+		PR_ERR("NA: invoke command(%d) failed: %d\n", command, teecErr);
+	}
+
+	TEEC_ReleaseSharedMemory(&tee_params_q);
+
+out1:
+	TEEC_ReleaseSharedMemory(&tee_params_p);
+out:
+	if (teecErr)
+		PR_ERR("NA: %s finish with %d\n", __func__, teecErr);
+
+	return teecErr;
+}
+
+void rsa_demo_uninit(void)
+{
+	PR_DBG("NA: close session\n");
+	TEEC_CloseSession(&sess);
+
+	PR_DBG("NA: finalize context\n");
+	TEEC_FinalizeContext(&ctx);
+}
+
diff --git a/security/optee-rsa-demo/src/na/liboptee_rsa.h b/security/optee-rsa-demo/src/na/liboptee_rsa.h
new file mode 100644
index 000000000..af8fcac52
--- /dev/null
+++ b/security/optee-rsa-demo/src/na/liboptee_rsa.h
@@ -0,0 +1,28 @@
+#ifndef _LIBROTPK_H__
+#define _LIBROTPK_H__
+
+#include <string.h>
+#include <stdio.h>
+#include <stdint.h>
+
+/**
+ * sunxi_dump() - dump buf.
+ * @buf: input buf.
+ * @ttl_len: dump length.
+ *
+ * return value: zero, write success; non-zero, write failed.
+ */
+void sunxi_dump(uint8_t *buf, uint32_t ttl_len);
+
+int rsa_demo_init(void);
+
+int rsa_demo_encrypt(uint8_t *in_buf, uint32_t in_len,
+		     uint8_t *out_buf, uint32_t *out_len);
+
+int rsa_demo_decrypt(uint8_t *in_buf, uint32_t in_len,
+		     uint8_t *out_buf, uint32_t *out_len);
+
+void rsa_demo_uninit(void);
+
+#endif
+
diff --git a/security/optee-rsa-demo/src/na/optee-rsa-demo.c b/security/optee-rsa-demo/src/na/optee-rsa-demo.c
new file mode 100644
index 000000000..6f6c4f6ea
--- /dev/null
+++ b/security/optee-rsa-demo/src/na/optee-rsa-demo.c
@@ -0,0 +1,64 @@
+#include <stdio.h>
+#include <string.h>
+#include "liboptee_rsa.h"
+
+static void usage(char *app)
+{
+	printf("usage: %s [string]\n", app);
+	printf(" [string]: string len should be less than 128 bytes\n");
+}
+
+int main(int argc, char *argv[])
+{
+	if (argc != 2) {
+		usage(argv[0]);
+		return -1;
+	}
+
+	uint8_t *src_buf = (uint8_t *)argv[1];
+	uint32_t src_len = strlen(argv[1]);
+
+	uint8_t enc_buf[1024/8] = {0};
+	uint32_t enc_len = 0;
+
+	uint8_t dec_buf[1024/8] = {0};
+	uint32_t dec_len = 0;
+
+	uint32_t ret = 0;
+
+	// connect with TA and generate an 1024 bits rsa key in TA
+	ret = rsa_demo_init();
+	if (ret) {
+		printf("rsa_demo_init failed: %d\n", ret);
+		return ret;
+	}
+
+	printf("src_len: %d, src_buf: %s\n", src_len, src_buf);
+
+	// encrypt src_buf with rsa publich key
+	ret = rsa_demo_encrypt(src_buf, src_len, enc_buf, &enc_len);
+	if (ret) {
+		printf("rsa_demo_encrypt failed: %d\n", ret);
+		return ret;
+	}
+
+	// decrypt enc_buf with rsa private key
+	ret = rsa_demo_decrypt(enc_buf, enc_len, dec_buf, &dec_len);
+	if (ret) {
+		printf("rsa_demo_decryp failed: %d\n", ret);
+		return ret;
+	}
+
+	// disconnect with TA
+	rsa_demo_uninit();
+
+	printf("dec_len: %d, dec_buf: %s\n", dec_len, dec_buf);
+
+	if (dec_len != src_len)
+		printf("ERROR: decrypt error, length dismatch!\n");
+
+	if(memcmp(src_buf, dec_buf, dec_len))
+		printf("ERROR: decrypt error, data dismatch!\n");
+
+	return 0;
+}
diff --git a/security/optee-rsa-demo/src/scripts/common.mk b/security/optee-rsa-demo/src/scripts/common.mk
new file mode 100644
index 000000000..103a4155a
--- /dev/null
+++ b/security/optee-rsa-demo/src/scripts/common.mk
@@ -0,0 +1,9 @@
+# Makefile script used in various OP-TEE test components
+
+# Since 'clean' make rules rely on 'rmdir', removed directories shall
+# not end with 'current directory' (./ or .) subpath information.
+# This macro remove trailing './' info from input path.
+define strip-trailing-slashes-and-dots
+$(eval _o := $(patsubst %/,%,$(patsubst %/.,%,$(1))))$(if \
+	$(filter-out $(1),$(_o)),$(call strip-trailing-slashes-and-dots,$(_o)),$(_o))
+endef
diff --git a/security/optee-rsa-demo/src/ta/Makefile b/security/optee-rsa-demo/src/ta/Makefile
new file mode 100644
index 000000000..4b96d8048
--- /dev/null
+++ b/security/optee-rsa-demo/src/ta/Makefile
@@ -0,0 +1,2 @@
+BINARY = e0c2b919-5b19-4150-a078-d8258374f4d9
+include ta_common.mk
diff --git a/security/optee-rsa-demo/src/ta/include/rsa-demo_ta.h b/security/optee-rsa-demo/src/ta/include/rsa-demo_ta.h
new file mode 100644
index 000000000..b99db2d63
--- /dev/null
+++ b/security/optee-rsa-demo/src/ta/include/rsa-demo_ta.h
@@ -0,0 +1,12 @@
+#ifndef __RSA_DEMO_H
+#define __RSA_DEMO_H
+
+#define RSA_DEMO_UUID  {0xe0c2b919, 0x5b19, 0x4150, \
+	{ 0xa0, 0x78, 0xd8, 0x25, 0x83, 0x74, 0xf4, 0xd9 }}
+
+#define CMD_RSA_DEMO_ENCRYPT 220
+#define CMD_RSA_DEMO_DECRYPT 221
+
+#define RSA_KEY_SIZE 1024
+
+#endif
diff --git a/security/optee-rsa-demo/src/ta/include/user_ta_header_defines.h b/security/optee-rsa-demo/src/ta/include/user_ta_header_defines.h
new file mode 100644
index 000000000..91e19e110
--- /dev/null
+++ b/security/optee-rsa-demo/src/ta/include/user_ta_header_defines.h
@@ -0,0 +1,16 @@
+#ifndef USER_TA_HEADER_DEFINES_H
+#define USER_TA_HEADER_DEFINES_H
+
+#include "rsa-demo_ta.h"
+#define TA_UUID RSA_DEMO_UUID
+
+/*
+ * This is important to have TA_FLAG_SINGLE_INSTANCE && !TA_FLAG_MULTI_SESSION
+ * as it is used by the ytest
+ */
+#define TA_FLAGS		(TA_FLAG_USER_MODE | TA_FLAG_EXEC_DDR | \
+				TA_FLAG_MULTI_SESSION)
+#define TA_STACK_SIZE		(2 * 1024)
+#define TA_DATA_SIZE		(32 * 1024)
+
+#endif
diff --git a/security/optee-rsa-demo/src/ta/optee_rsa-demo_ta.c b/security/optee-rsa-demo/src/ta/optee_rsa-demo_ta.c
new file mode 100644
index 000000000..408be6190
--- /dev/null
+++ b/security/optee-rsa-demo/src/ta/optee_rsa-demo_ta.c
@@ -0,0 +1,275 @@
+#include <tee_ta_api.h>
+#include <stdio.h>
+#include <utee_syscalls.h>
+#include <string.h>
+#include <tee_api.h>
+
+#include <rsa-demo_ta.h>
+#include <mbedtls/rsa.h>
+
+#include <mbedtls/entropy.h>
+#include <mbedtls/ctr_drbg.h>
+
+//#define DBG_ON
+#define USE_PERSISTENT_KEY
+
+#ifdef DBG_ON
+#define PR_ERR(...)	printf(__VA_ARGS__)
+#define PR_DBG(...)	printf(__VA_ARGS__)
+#else
+#define PR_ERR(...)	printf(__VA_ARGS__)
+#define PR_DBG(...)	(void)0
+#endif
+
+#ifdef USE_PERSISTENT_KEY
+#define KEY_LEN 128
+
+#define RSA_N   "9292758453063D803DD603D5E777D788" \
+                "8ED1D5BF35786190FA2F23EBC0848AEA" \
+                "DDA92CA6C3D80B32C4D109BE0F36D6AE" \
+                "7130B9CED7ACDF54CFC7555AC14EEBAB" \
+                "93A89813FBF3C4F8066D2D800F7C38A8" \
+                "1AE31942917403FF4946B0A83D3D3E05" \
+                "EE57C6F5F5606FB5D4BC6CD34EE0801A" \
+                "5E94BB77B07507233A0BC7BAC8F90F79"
+
+#define RSA_E   "10001"
+
+#define RSA_D   "24BF6185468786FDD303083D25E64EFC" \
+                "66CA472BC44D253102F8B4A9D3BFA750" \
+                "91386C0077937FE33FA3252D28855837" \
+                "AE1B484A8A9A45F7EE8C0C634F99E8CD" \
+                "DF79C5CE07EE72C7F123142198164234" \
+                "CABB724CF78B8173B9F880FC86322407" \
+                "AF1FEDFDDE2BEB674CA15F3E81A1521E" \
+                "071513A1E85B5DFA031F21ECAE91A34D"
+
+#define RSA_P   "C36D0EB7FCD285223CFB5AABA5BDA3D8" \
+                "2C01CAD19EA484A87EA4377637E75500" \
+                "FCB2005C5C7DD6EC4AC023CDA285D796" \
+                "C3D9E75E1EFC42488BB4F1D13AC30A57"
+
+#define RSA_Q   "C000DF51A7C77AE8D7C7370C1FF55B69" \
+                "E211C2B9E5DB1ED0BF61D0D9899620F4" \
+                "910E4168387E3C30AA1E00C339A79508" \
+                "8452DD96A9A5EA5D9DCA68DA636032AF"
+
+#define RSA_DP  "C1ACF567564274FB07A0BBAD5D26E298" \
+                "3C94D22288ACD763FD8E5600ED4A702D" \
+                "F84198A5F06C2E72236AE490C93F07F8" \
+                "3CC559CD27BC2D1CA488811730BB5725"
+
+#define RSA_DQ  "4959CBF6F8FEF750AEE6977C155579C7" \
+                "D8AAEA56749EA28623272E4F7D0592AF" \
+                "7C1F1313CAC9471B5C523BFE592F517B" \
+                "407A1BD76C164B93DA2D32A383E58357"
+
+#define RSA_QP  "9AE7FBC99546432DF71896FC239EADAE" \
+                "F38D18D2B2F0E2DD275AA977E2BF4411" \
+                "F5A3B2A5D33605AEBBCCBA7FEB9F2D2F" \
+                "A74206CEC169D74BF5A8C50D6F48EA08"
+#endif
+
+#ifdef DBG_ON
+static void dump(uint8_t *buf, int ttl_len)
+{
+	int len;
+
+	for (len = 0; len < ttl_len; len++) {
+		printf("0x%02x ", ((char *)buf)[len]);
+		if (len % 8 == 7) {
+			printf("\n");
+		}
+	}
+	printf("\n");
+}
+
+static void dump_rsa_key(mbedtls_rsa_context *ctx)
+{
+	size_t olen;
+	uint8_t buf[516];
+
+	printf("\n  +++++++++++++++ rsa keypair +++++++++++++++\n\n");
+	mbedtls_mpi_write_string(&ctx->N , 16, (char *)buf, sizeof(buf), &olen);
+	printf("N: %s\n", buf);
+
+	mbedtls_mpi_write_string(&ctx->E , 16, (char *)buf, sizeof(buf), &olen);
+	printf("E: %s\n", buf);
+
+	mbedtls_mpi_write_string(&ctx->D , 16, (char *)buf, sizeof(buf), &olen);
+	printf("D: %s\n", buf);
+
+	mbedtls_mpi_write_string(&ctx->P , 16, (char *)buf, sizeof(buf), &olen);
+	printf("P: %s\n", buf);
+
+	mbedtls_mpi_write_string(&ctx->Q , 16, (char *)buf, sizeof(buf), &olen);
+	printf("Q: %s\n", buf);
+
+	mbedtls_mpi_write_string(&ctx->DP, 16, (char *)buf, sizeof(buf), &olen);
+	printf("DP: %s\n", buf);
+
+	mbedtls_mpi_write_string(&ctx->DQ, 16, (char *)buf, sizeof(buf), &olen);
+	printf("DQ: %s\n", buf);
+
+	mbedtls_mpi_write_string(&ctx->QP, 16, (char *)buf, sizeof(buf), &olen);
+	printf("QP: %s\n", buf);
+	printf("\n  +++++++++++++++ rsa keypair +++++++++++++++\n\n");
+}
+
+#else
+static void dump(uint8_t *buf, int ttl_len)
+{
+	(void)buf;
+	(void)ttl_len;
+}
+
+static void dump_rsa_key(mbedtls_rsa_context *ctx)
+{
+	(void)ctx;
+}
+#endif
+
+mbedtls_rsa_context rsa;
+
+static int f_rng(void *rng __unused, unsigned char *output, size_t output_len)
+{
+	TEE_GenerateRandom(output, output_len);
+	return 0;
+}
+
+/*
+ * Trusted Application Entry Points
+ */
+
+/* Called each time a new instance is created */
+TEE_Result TA_CreateEntryPoint(void)
+{
+	PR_DBG("TA: create entry!\n");
+	return TEE_SUCCESS;
+}
+
+/* Called each time an instance is destroyed */
+void TA_DestroyEntryPoint(void)
+{
+}
+
+/* Called each time a session is opened */
+TEE_Result TA_OpenSessionEntryPoint(uint32_t nParamTypes,
+				    TEE_Param pParams[4],
+				    void **ppSessionContext)
+{
+	(void)nParamTypes;
+	(void)pParams;
+	(void)ppSessionContext;
+
+#ifdef USE_PERSISTENT_KEY
+	mbedtls_rsa_init( &rsa, MBEDTLS_RSA_PKCS_V15, 0 );
+#else
+	TEE_Result res = TEE_ERROR_GENERIC;
+
+	PR_DBG("TA: generating rsa key...\n");
+
+	mbedtls_rsa_init( &rsa, MBEDTLS_RSA_PKCS_V15, 0 );
+	res = mbedtls_rsa_gen_key( &rsa, f_rng, NULL, RSA_KEY_SIZE, 65537 );
+	if( res != 0 ) {
+		PR_ERR("TA: gen rsa key failed: %d!\n", res);
+		PR_ERR("TA: open session failed!\n");
+		return res;
+	} else if ((size_t)mbedtls_mpi_bitlen(&rsa.N) != RSA_KEY_SIZE) {
+		PR_ERR("TA: gen rsa key failed!\n");
+		PR_ERR("TA: open session failed!\n");
+		return TEE_ERROR_BAD_PARAMETERS;
+	}
+
+	dump_rsa_key(&rsa);
+#endif
+	PR_DBG("TA: open session!\n");
+	return TEE_SUCCESS;
+}
+
+/* Called each time a session is closed */
+void TA_CloseSessionEntryPoint(void *pSessionContext)
+{
+	(void)pSessionContext;
+
+	mbedtls_rsa_free(&rsa);
+}
+
+/* Called when a command is invoked */
+TEE_Result TA_InvokeCommandEntryPoint(void *pSessionContext,
+				      uint32_t nCommandID, uint32_t nParamTypes,
+				      TEE_Param pParams[4])
+{
+	(void)pSessionContext;
+	(void)nParamTypes;
+	(void)pParams;
+
+	TEE_Result res = 0;
+
+	uint8_t *buf_in = (unsigned char *)pParams[0].memref.buffer;
+	uint32_t in_len = pParams[2].value.a;
+
+	uint8_t buf_out[RSA_KEY_SIZE/8] = {0};
+	uint32_t out_len = 0;
+
+	switch (nCommandID) {
+	case CMD_RSA_DEMO_ENCRYPT:
+#ifdef USE_PERSISTENT_KEY
+		if ((res = mbedtls_mpi_read_string(&rsa.N, 16, RSA_N)) != 0 ||
+		    (res = mbedtls_mpi_read_string(&rsa.E, 16, RSA_E)) != 0) {
+			PR_ERR("TA: get rsa pub key failed: 0x%x\n", res);
+			return res;
+		}
+
+		dump_rsa_key(&rsa);
+		rsa.len = RSA_KEY_SIZE / 8;
+#endif
+		res = mbedtls_rsa_pkcs1_encrypt(&rsa, f_rng, NULL,
+				MBEDTLS_RSA_PUBLIC, in_len, buf_in, buf_out);
+		if (res) {
+			PR_ERR("TA: rsa encrypt failed, return: 0x%x\n", res);
+		} else {
+			PR_DBG("TA: rsa encrypt success. Encrypt data:\n");
+			dump(buf_out, RSA_KEY_SIZE/8);
+		}
+
+		memcpy(pParams[1].memref.buffer, buf_out, RSA_KEY_SIZE/8);
+		pParams[2].value.b = RSA_KEY_SIZE / 8;
+
+		return res;
+	case CMD_RSA_DEMO_DECRYPT:
+#ifdef USE_PERSISTENT_KEY
+		if ((res = mbedtls_mpi_read_string(&rsa.N, 16, RSA_N)) != 0 ||
+		   (res = mbedtls_mpi_read_string(&rsa.E, 16, RSA_E)) != 0 ||
+		   (res = mbedtls_mpi_read_string(&rsa.D, 16, RSA_D)) != 0 ||
+		   (res = mbedtls_mpi_read_string(&rsa.P, 16, RSA_P)) != 0 ||
+		   (res = mbedtls_mpi_read_string(&rsa.Q, 16, RSA_Q)) != 0 ||
+		   (res = mbedtls_mpi_read_string(&rsa.DP, 16, RSA_DP)) != 0 ||
+		   (res = mbedtls_mpi_read_string(&rsa.DQ, 16, RSA_DQ)) != 0 ||
+		   (res = mbedtls_mpi_read_string(&rsa.QP, 16, RSA_QP)) != 0 ) {
+			PR_ERR("TA: get rsa pri key failed: 0x%x\n", res);
+			return res;
+		}
+
+		dump_rsa_key(&rsa);
+		rsa.len = RSA_KEY_SIZE / 8;
+#endif
+		res = mbedtls_rsa_pkcs1_decrypt(&rsa, f_rng, NULL,
+					MBEDTLS_RSA_PRIVATE, &out_len,
+					buf_in, buf_out, RSA_KEY_SIZE/8);
+		if (res) {
+			PR_ERR("TA: rsa decrypt failed, return: 0x%x\n", res);
+		} else {
+			PR_DBG("TA: rsa decrypt success. Decrypt data:\n");
+			dump(buf_out, out_len);
+		}
+
+		memcpy(pParams[1].memref.buffer, buf_out, out_len);
+		pParams[2].value.b = out_len;
+
+		return res;
+	default:
+		return TEE_ERROR_BAD_PARAMETERS;
+	}
+}
+
diff --git a/security/optee-rsa-demo/src/ta/optee_rsa-demo_ta.c-bak b/security/optee-rsa-demo/src/ta/optee_rsa-demo_ta.c-bak
new file mode 100644
index 000000000..4e685373f
--- /dev/null
+++ b/security/optee-rsa-demo/src/ta/optee_rsa-demo_ta.c-bak
@@ -0,0 +1,272 @@
+#include <tee_ta_api.h>
+#include <stdio.h>
+#include <utee_syscalls.h>
+#include <string.h>
+#include <tee_api.h>
+
+#include <rsa-demo_ta.h>
+#include <mbedtls/rsa.h>
+
+//#define DBG_ON
+#define USE_PERSISTENT_KEY
+
+#ifdef DBG_ON
+#define PR_ERR(...)	printf(__VA_ARGS__)
+#define PR_DBG(...)	printf(__VA_ARGS__)
+#else
+#define PR_ERR(...)	printf(__VA_ARGS__)
+#define PR_DBG(...)	(void)0
+#endif
+
+#ifdef USE_PERSISTENT_KEY
+#define KEY_LEN 128
+
+#define RSA_N   "9292758453063D803DD603D5E777D788" \
+                "8ED1D5BF35786190FA2F23EBC0848AEA" \
+                "DDA92CA6C3D80B32C4D109BE0F36D6AE" \
+                "7130B9CED7ACDF54CFC7555AC14EEBAB" \
+                "93A89813FBF3C4F8066D2D800F7C38A8" \
+                "1AE31942917403FF4946B0A83D3D3E05" \
+                "EE57C6F5F5606FB5D4BC6CD34EE0801A" \
+                "5E94BB77B07507233A0BC7BAC8F90F79"
+
+#define RSA_E   "10001"
+
+#define RSA_D   "24BF6185468786FDD303083D25E64EFC" \
+                "66CA472BC44D253102F8B4A9D3BFA750" \
+                "91386C0077937FE33FA3252D28855837" \
+                "AE1B484A8A9A45F7EE8C0C634F99E8CD" \
+                "DF79C5CE07EE72C7F123142198164234" \
+                "CABB724CF78B8173B9F880FC86322407" \
+                "AF1FEDFDDE2BEB674CA15F3E81A1521E" \
+                "071513A1E85B5DFA031F21ECAE91A34D"
+
+#define RSA_P   "C36D0EB7FCD285223CFB5AABA5BDA3D8" \
+                "2C01CAD19EA484A87EA4377637E75500" \
+                "FCB2005C5C7DD6EC4AC023CDA285D796" \
+                "C3D9E75E1EFC42488BB4F1D13AC30A57"
+
+#define RSA_Q   "C000DF51A7C77AE8D7C7370C1FF55B69" \
+                "E211C2B9E5DB1ED0BF61D0D9899620F4" \
+                "910E4168387E3C30AA1E00C339A79508" \
+                "8452DD96A9A5EA5D9DCA68DA636032AF"
+
+#define RSA_DP  "C1ACF567564274FB07A0BBAD5D26E298" \
+                "3C94D22288ACD763FD8E5600ED4A702D" \
+                "F84198A5F06C2E72236AE490C93F07F8" \
+                "3CC559CD27BC2D1CA488811730BB5725"
+
+#define RSA_DQ  "4959CBF6F8FEF750AEE6977C155579C7" \
+                "D8AAEA56749EA28623272E4F7D0592AF" \
+                "7C1F1313CAC9471B5C523BFE592F517B" \
+                "407A1BD76C164B93DA2D32A383E58357"
+
+#define RSA_QP  "9AE7FBC99546432DF71896FC239EADAE" \
+                "F38D18D2B2F0E2DD275AA977E2BF4411" \
+                "F5A3B2A5D33605AEBBCCBA7FEB9F2D2F" \
+                "A74206CEC169D74BF5A8C50D6F48EA08"
+#endif
+
+#ifdef DBG_ON
+static void dump(uint8_t *buf, int ttl_len)
+{
+	int len;
+
+	for (len = 0; len < ttl_len; len++) {
+		printf("0x%02x ", ((char *)buf)[len]);
+		if (len % 8 == 7) {
+			printf("\n");
+		}
+	}
+	printf("\n");
+}
+
+static void dump_rsa_key(mbedtls_rsa_context *ctx)
+{
+	size_t olen;
+	uint8_t buf[516];
+
+	printf("\n  +++++++++++++++ rsa keypair +++++++++++++++\n\n");
+	mbedtls_mpi_write_string(&ctx->N , 16, (char *)buf, sizeof(buf), &olen);
+	printf("N: %s\n", buf);
+
+	mbedtls_mpi_write_string(&ctx->E , 16, (char *)buf, sizeof(buf), &olen);
+	printf("E: %s\n", buf);
+
+	mbedtls_mpi_write_string(&ctx->D , 16, (char *)buf, sizeof(buf), &olen);
+	printf("D: %s\n", buf);
+
+	mbedtls_mpi_write_string(&ctx->P , 16, (char *)buf, sizeof(buf), &olen);
+	printf("P: %s\n", buf);
+
+	mbedtls_mpi_write_string(&ctx->Q , 16, (char *)buf, sizeof(buf), &olen);
+	printf("Q: %s\n", buf);
+
+	mbedtls_mpi_write_string(&ctx->DP, 16, (char *)buf, sizeof(buf), &olen);
+	printf("DP: %s\n", buf);
+
+	mbedtls_mpi_write_string(&ctx->DQ, 16, (char *)buf, sizeof(buf), &olen);
+	printf("DQ: %s\n", buf);
+
+	mbedtls_mpi_write_string(&ctx->QP, 16, (char *)buf, sizeof(buf), &olen);
+	printf("QP: %s\n", buf);
+	printf("\n  +++++++++++++++ rsa keypair +++++++++++++++\n\n");
+}
+
+#else
+static void dump(uint8_t *buf, int ttl_len)
+{
+	(void)buf;
+	(void)ttl_len;
+}
+
+static void dump_rsa_key(mbedtls_rsa_context *ctx)
+{
+	(void)ctx;
+}
+#endif
+
+mbedtls_rsa_context rsa;
+
+static int f_rng(void *rng __unused, unsigned char *output, size_t output_len)
+{
+	TEE_GenerateRandom(output, output_len);
+	return 0;
+}
+
+/*
+ * Trusted Application Entry Points
+ */
+
+/* Called each time a new instance is created */
+TEE_Result TA_CreateEntryPoint(void)
+{
+	PR_DBG("TA: create entry!\n");
+	return TEE_SUCCESS;
+}
+
+/* Called each time an instance is destroyed */
+void TA_DestroyEntryPoint(void)
+{
+}
+
+/* Called each time a session is opened */
+TEE_Result TA_OpenSessionEntryPoint(uint32_t nParamTypes,
+				    TEE_Param pParams[4],
+				    void **ppSessionContext)
+{
+	(void)nParamTypes;
+	(void)pParams;
+	(void)ppSessionContext;
+
+#ifdef USE_PERSISTENT_KEY
+	mbedtls_rsa_init( &rsa, MBEDTLS_RSA_PKCS_V15, 0 );
+#else
+	TEE_Result res = TEE_ERROR_GENERIC;
+
+	PR_DBG("TA: generating rsa key...\n");
+
+	mbedtls_rsa_init( &rsa, MBEDTLS_RSA_PKCS_V15, 0 );
+	res = mbedtls_rsa_gen_key( &rsa, f_rng, NULL, RSA_KEY_SIZE, 65537 );
+	if( res != 0 ) {
+		PR_ERR("TA: gen rsa key failed: %d!\n", res);
+		PR_ERR("TA: open session failed!\n");
+		return res;
+	} else if ((size_t)mbedtls_mpi_bitlen(&rsa.N) != RSA_KEY_SIZE) {
+		PR_ERR("TA: gen rsa key failed!\n");
+		PR_ERR("TA: open session failed!\n");
+		return TEE_ERROR_BAD_PARAMETERS;
+	}
+
+	dump_rsa_key(&rsa);
+#endif
+	PR_DBG("TA: open session!\n");
+	return TEE_SUCCESS;
+}
+
+/* Called each time a session is closed */
+void TA_CloseSessionEntryPoint(void *pSessionContext)
+{
+	(void)pSessionContext;
+
+	mbedtls_rsa_free(&rsa);
+}
+
+/* Called when a command is invoked */
+TEE_Result TA_InvokeCommandEntryPoint(void *pSessionContext,
+				      uint32_t nCommandID, uint32_t nParamTypes,
+				      TEE_Param pParams[4])
+{
+	(void)pSessionContext;
+	(void)nParamTypes;
+	(void)pParams;
+
+	TEE_Result res = 0;
+
+	uint8_t *buf_in = (unsigned char *)pParams[0].memref.buffer;
+	uint32_t in_len = pParams[2].value.a;
+
+	uint8_t buf_out[RSA_KEY_SIZE/8] = {0};
+	uint32_t out_len = 0;
+
+	switch (nCommandID) {
+	case CMD_RSA_DEMO_ENCRYPT:
+#ifdef USE_PERSISTENT_KEY
+		if ((res = mbedtls_mpi_read_string(&rsa.N, 16, RSA_N)) != 0 ||
+		    (res = mbedtls_mpi_read_string(&rsa.E, 16, RSA_E)) != 0) {
+			PR_ERR("TA: get rsa pub key failed: 0x%x\n", res);
+			return res;
+		}
+
+		dump_rsa_key(&rsa);
+		rsa.len = RSA_KEY_SIZE / 8;
+#endif
+		res = mbedtls_rsa_pkcs1_encrypt(&rsa, f_rng, NULL,
+				MBEDTLS_RSA_PUBLIC, in_len, buf_in, buf_out);
+		if (res) {
+			PR_ERR("TA: rsa encrypt failed, return: 0x%x\n", res);
+		} else {
+			PR_DBG("TA: rsa encrypt success. Encrypt data:\n");
+			dump(buf_out, RSA_KEY_SIZE/8);
+		}
+
+		memcpy(pParams[1].memref.buffer, buf_out, RSA_KEY_SIZE/8);
+		pParams[2].value.b = RSA_KEY_SIZE / 8;
+
+		return res;
+	case CMD_RSA_DEMO_DECRYPT:
+#ifdef USE_PERSISTENT_KEY
+		if ((res = mbedtls_mpi_read_string(&rsa.N, 16, RSA_N)) != 0 ||
+		   (res = mbedtls_mpi_read_string(&rsa.E, 16, RSA_E)) != 0 ||
+		   (res = mbedtls_mpi_read_string(&rsa.D, 16, RSA_D)) != 0 ||
+		   (res = mbedtls_mpi_read_string(&rsa.P, 16, RSA_P)) != 0 ||
+		   (res = mbedtls_mpi_read_string(&rsa.Q, 16, RSA_Q)) != 0 ||
+		   (res = mbedtls_mpi_read_string(&rsa.DP, 16, RSA_DP)) != 0 ||
+		   (res = mbedtls_mpi_read_string(&rsa.DQ, 16, RSA_DQ)) != 0 ||
+		   (res = mbedtls_mpi_read_string(&rsa.QP, 16, RSA_QP)) != 0 ) {
+			PR_ERR("TA: get rsa pri key failed: 0x%x\n", res);
+			return res;
+		}
+
+		dump_rsa_key(&rsa);
+		rsa.len = RSA_KEY_SIZE / 8;
+#endif
+		res = mbedtls_rsa_pkcs1_decrypt(&rsa, f_rng, NULL,
+					MBEDTLS_RSA_PRIVATE, &out_len,
+					buf_in, buf_out, RSA_KEY_SIZE/8);
+		if (res) {
+			PR_ERR("TA: rsa decrypt failed, return: 0x%x\n", res);
+		} else {
+			PR_DBG("TA: rsa decrypt success. Decrypt data:\n");
+			dump(buf_out, out_len);
+		}
+
+		memcpy(pParams[1].memref.buffer, buf_out, out_len);
+		pParams[2].value.b = out_len;
+
+		return res;
+	default:
+		return TEE_ERROR_BAD_PARAMETERS;
+	}
+}
+
diff --git a/security/optee-rsa-demo/src/ta/sub.mk b/security/optee-rsa-demo/src/ta/sub.mk
new file mode 100644
index 000000000..35ae18de0
--- /dev/null
+++ b/security/optee-rsa-demo/src/ta/sub.mk
@@ -0,0 +1,2 @@
+global-incdirs-y += include
+srcs-y += optee_rsa-demo_ta.c
diff --git a/security/optee-rsa-demo/src/ta/ta_common.mk b/security/optee-rsa-demo/src/ta/ta_common.mk
new file mode 100644
index 000000000..b8ce142bc
--- /dev/null
+++ b/security/optee-rsa-demo/src/ta/ta_common.mk
@@ -0,0 +1,3 @@
+# Note that it is important that that $(BINARY) is set before we include
+# ta_dev_kit.mk. In the optee_test this is done in the individual TA Makefiles.
+include $(TA_DEV_KIT_DIR)/mk/ta_dev_kit.mk
-- 
2.29.0

