feat(runtime/rmi): implement RMI_PDEV_STOP

RMI:
This RMI handler set the PDEV state to STOPPING, this allows NS Host to
invoke PDEV communicate call to do deinit the device connection and
transition the PDEV to STOPPED state.

DEV_ASSIGN App:
Added command to deinit the connection. This terminates the secure
session. This command also unset the public key context and certificate
chain hash in the libspdm connection state.

Signed-off-by: Arunachalam Ganapathy <arunachalam.ganapathy@arm.com>
Change-Id: I7d0d9b1654137f739793005cf290b8bb4c74957a
diff --git a/app/device_assignment/el0_app/src/dev_assign_cmds.c b/app/device_assignment/el0_app/src/dev_assign_cmds.c
index 28c091a..8927469 100644
--- a/app/device_assignment/el0_app/src/dev_assign_cmds.c
+++ b/app/device_assignment/el0_app/src/dev_assign_cmds.c
@@ -45,7 +45,7 @@
 	return DEV_ASSIGN_STATUS_SUCCESS;
 }
 
-/* dsm_cmd_init_connection_main */
+/* dev_assign_cmd_init_connection_main */
 int dev_assign_cmd_init_connection_main(struct dev_assign_info *info)
 {
 	void *spdm_context = info->libspdm_ctx;
@@ -75,3 +75,18 @@
 
 	return (int)LIBSPDM_STATUS_SUCCESS;
 }
+
+/* dev_assign_cmd_stop_connection_main */
+int dev_assign_cmd_stop_connection_main(struct dev_assign_info *info)
+{
+	/* Send GET_VERSION, this resets the SPDM connection */
+	(void)libspdm_init_connection(info->libspdm_ctx, true);
+
+	/*
+	 * Unset public key context in libspdm connection once the connection is
+	 * terminated.
+	 */
+	dev_assign_unset_pubkey(info);
+
+	return DEV_ASSIGN_STATUS_SUCCESS;
+}
diff --git a/app/device_assignment/el0_app/src/dev_assign_el0_app.c b/app/device_assignment/el0_app/src/dev_assign_el0_app.c
index 9e5662b..f862994 100644
--- a/app/device_assignment/el0_app/src/dev_assign_el0_app.c
+++ b/app/device_assignment/el0_app/src/dev_assign_el0_app.c
@@ -203,6 +203,16 @@
 	/* Nothing to do */
 }
 
+void dev_assign_unset_pubkey(struct dev_assign_info *info)
+{
+	if ((info->key_sig_algo == RMI_SIGNATURE_ALGORITHM_ECDSA_P256) ||
+	    (info->key_sig_algo == RMI_SIGNATURE_ALGORITHM_ECDSA_P384)) {
+		mbedtls_ecdh_free(&info->pk_ctx.ecdh);
+	} else {
+		mbedtls_rsa_free(&info->pk_ctx.rsa);
+	}
+}
+
 /*
  * Returns the min heap size. This include libspdm context, libspdm secured
  * message context, libspdm scratch space, libspdm send recv buffer and
@@ -511,6 +521,9 @@
 	case DEVICE_ASSIGN_APP_FUNC_ID_CONNECT_INIT:
 		ret = (unsigned long)dev_assign_cmd_init_connection_main(info);
 		break;
+	case DEVICE_ASSIGN_APP_FUNC_ID_STOP_CONNECTION:
+		ret = (unsigned long)dev_assign_cmd_stop_connection_main(info);
+		break;
 	default:
 		assert(false);
 		return (unsigned long)DEV_ASSIGN_STATUS_ERROR;
@@ -549,6 +562,7 @@
 			(struct dev_assign_params *)shared);
 	}
 	case DEVICE_ASSIGN_APP_FUNC_ID_CONNECT_INIT:
+	case DEVICE_ASSIGN_APP_FUNC_ID_STOP_CONNECTION:
 		return dev_assign_communicate_cmd_cmn(func_id, heap);
 	case DEVICE_ASSIGN_APP_FUNC_ID_DEINIT:
 		return (unsigned long)dev_assign_deinit(heap);
diff --git a/app/device_assignment/el0_app/src/dev_assign_private.h b/app/device_assignment/el0_app/src/dev_assign_private.h
index eda7e3c..a485c18 100644
--- a/app/device_assignment/el0_app/src/dev_assign_private.h
+++ b/app/device_assignment/el0_app/src/dev_assign_private.h
@@ -12,7 +12,9 @@
 #include <industry_standard/spdm_secured_message.h>
 #include <library/spdm_requester_lib.h>
 #include <library/spdm_secured_message_lib.h>
+#include <mbedtls/ecdh.h>
 #include <mbedtls/memory_buffer_alloc.h>
+#include <mbedtls/rsa.h>
 #include <psa/crypto.h>
 #include <sizes.h>
 #include <utils_def.h>
@@ -243,10 +245,20 @@
 
 	buffer_alloc_ctx mbedtls_heap_ctx;
 
+	/* Public key context */
+	uint32_t key_sig_algo;
+	union {
+		mbedtls_ecdh_context ecdh;
+		mbedtls_rsa_context rsa;
+	} pk_ctx;
+
 	/* Exit and Entry args for dev_communicate cmds */
 	struct rmi_dev_comm_enter enter_args;
 	struct rmi_dev_comm_exit exit_args;
 
+	/* ID of the SPDM session started by libspdm_start_session */
+	uint32_t session_id;
+
 	/*
 	 * The PSA equivalent of the 'rmi_hash_algo'. This value is used by PSA
 	 * crypto calls to calculate hash of cached device objects.
@@ -262,5 +274,8 @@
 };
 
 int dev_assign_cmd_init_connection_main(struct dev_assign_info *info);
+int dev_assign_cmd_stop_connection_main(struct dev_assign_info *info);
+
+void dev_assign_unset_pubkey(struct dev_assign_info *info);
 
 #endif /* DEV_ASSIGN_PRIVATE_H */
diff --git a/app/device_assignment/rmm_stub/include/dev_assign_structs.h b/app/device_assignment/rmm_stub/include/dev_assign_structs.h
index 1411393..e2c623e 100644
--- a/app/device_assignment/rmm_stub/include/dev_assign_structs.h
+++ b/app/device_assignment/rmm_stub/include/dev_assign_structs.h
@@ -95,4 +95,13 @@
  */
 #define DEVICE_ASSIGN_APP_FUNC_ID_DEINIT		4
 
+/*
+ * App function ID to stop the libspdm session that is associated with this app
+ * instance.
+ *
+ * ret0 == DEV_ASSIGN_STATUS_SUCCESS if the session is stopped successfully.
+ *         DEV_ASSIGN_STATUS_ERROR if libspdm returned error.
+ */
+#define DEVICE_ASSIGN_APP_FUNC_ID_STOP_CONNECTION	0x80
+
 #endif /* DEV_ASSIGN_STRUCTS_H */
diff --git a/app/device_assignment/rmm_stub/src/dev_assign_app_stub.c b/app/device_assignment/rmm_stub/src/dev_assign_app_stub.c
index 9486950..4c021a2 100644
--- a/app/device_assignment/rmm_stub/src/dev_assign_app_stub.c
+++ b/app/device_assignment/rmm_stub/src/dev_assign_app_stub.c
@@ -57,7 +57,8 @@
 	struct rmi_dev_comm_exit *shared_ret;
 
 	assert((dev_cmd == DEVICE_ASSIGN_APP_FUNC_ID_RESUME) ||
-		(dev_cmd == DEVICE_ASSIGN_APP_FUNC_ID_CONNECT_INIT));
+		(dev_cmd == DEVICE_ASSIGN_APP_FUNC_ID_CONNECT_INIT) ||
+		(dev_cmd == DEVICE_ASSIGN_APP_FUNC_ID_STOP_CONNECTION));
 
 	app_map_shared_page(app_data);
 	shared = app_data->el2_shared_page;
diff --git a/runtime/core/handler.c b/runtime/core/handler.c
index 140a746..4076efb 100644
--- a/runtime/core/handler.c
+++ b/runtime/core/handler.c
@@ -169,7 +169,7 @@
 	HANDLER(PDEV_IDE_RESET,		0, 0, NULL,			 true, true),
 	HANDLER(PDEV_NOTIFY,		0, 0, NULL,			 true, true),
 	HANDLER(PDEV_SET_PUBKEY,	4, 0, NULL,			 true, true),
-	HANDLER(PDEV_STOP,		0, 0, NULL,			 true, true),
+	HANDLER(PDEV_STOP,		1, 0, smc_pdev_stop,		 true, true),
 	HANDLER(RTT_AUX_CREATE,		0, 0, NULL,			 true, true),
 	HANDLER(RTT_AUX_DESTROY,	0, 0, NULL,			 true, true),
 	HANDLER(RTT_AUX_FOLD,		0, 0, NULL,			 true, true),
diff --git a/runtime/include/smc-handler.h b/runtime/include/smc-handler.h
index 75bb669..0ea36f5 100644
--- a/runtime/include/smc-handler.h
+++ b/runtime/include/smc-handler.h
@@ -118,6 +118,8 @@
 
 unsigned long smc_pdev_abort(unsigned long pdev_ptr);
 
+unsigned long smc_pdev_stop(unsigned long pdev_ptr);
+
 unsigned long smc_pdev_destroy(unsigned long pdev_ptr);
 
 #endif /* SMC_HANDLER_H */
diff --git a/runtime/rmi/pdev.c b/runtime/rmi/pdev.c
index f2afd06..a52bc20 100644
--- a/runtime/rmi/pdev.c
+++ b/runtime/rmi/pdev.c
@@ -434,6 +434,10 @@
 		rc = dev_assign_dev_communicate(&pd->da_app_data, enter_args,
 			exit_args, DEVICE_ASSIGN_APP_FUNC_ID_CONNECT_INIT);
 		break;
+	case RMI_PDEV_STATE_STOPPING:
+		rc = dev_assign_dev_communicate(&pd->da_app_data, enter_args,
+			exit_args, DEVICE_ASSIGN_APP_FUNC_ID_STOP_CONNECTION);
+		break;
 	default:
 		assert(false);
 		rc = -1;
@@ -546,11 +550,17 @@
 	 */
 	switch (pd->dev_comm_state) {
 	case DEV_COMM_ERROR:
-		pd->rmi_state = RMI_PDEV_STATE_ERROR;
+		if (pd->rmi_state == RMI_PDEV_STATE_STOPPING) {
+			pd->rmi_state = RMI_PDEV_STATE_STOPPED;
+		} else {
+			pd->rmi_state = RMI_PDEV_STATE_ERROR;
+		}
 		break;
 	case DEV_COMM_IDLE:
 		if (pd->rmi_state == RMI_PDEV_STATE_NEW) {
 			pd->rmi_state = RMI_PDEV_STATE_NEEDS_KEY;
+		} else if (pd->rmi_state == RMI_PDEV_STATE_STOPPING) {
+			pd->rmi_state = RMI_PDEV_STATE_STOPPED;
 		} else {
 			pd->rmi_state = RMI_PDEV_STATE_ERROR;
 		}
@@ -618,6 +628,62 @@
 }
 
 /*
+ * Stop the PDEV. This sets the PDEV state to STOPPING, and a PDEV communicate
+ * call can do device specific cleanup like terminating a secure session.
+ *
+ * pdev_ptr	- PA of the PDEV
+ */
+unsigned long smc_pdev_stop(unsigned long pdev_ptr)
+{
+	struct granule *g_pdev;
+	unsigned long smc_rc;
+	struct pdev *pd;
+
+	if (!is_rmi_feat_da_enabled()) {
+		return SMC_NOT_SUPPORTED;
+	}
+
+	if (!GRANULE_ALIGNED(pdev_ptr)) {
+		return RMI_ERROR_INPUT;
+	}
+
+	/* Lock pdev granule and map it */
+	g_pdev = find_lock_granule(pdev_ptr, GRANULE_STATE_PDEV);
+	if (g_pdev == NULL) {
+		return RMI_ERROR_INPUT;
+	}
+
+	pd = buffer_granule_map(g_pdev, SLOT_PDEV);
+	if (pd == NULL) {
+		granule_unlock(g_pdev);
+		return RMI_ERROR_INPUT;
+	}
+
+	if ((pd->rmi_state == RMI_PDEV_STATE_COMMUNICATING) ||
+	    (pd->rmi_state == RMI_PDEV_STATE_STOPPING) ||
+	    (pd->rmi_state == RMI_PDEV_STATE_STOPPED)) {
+		smc_rc = RMI_ERROR_DEVICE;
+		goto out_pdev_buf_unmap;
+	}
+
+	if (pd->num_vdevs != 0U) {
+		smc_rc = RMI_ERROR_DEVICE;
+		goto out_pdev_buf_unmap;
+	}
+
+	pd->rmi_state = RMI_PDEV_STATE_STOPPING;
+	pd->dev_comm_state = DEV_COMM_PENDING;
+	smc_rc = RMI_SUCCESS;
+
+out_pdev_buf_unmap:
+	buffer_unmap(pd);
+
+	granule_unlock(g_pdev);
+
+	return smc_rc;
+}
+
+/*
  * Abort device communication associated with a PDEV.
  *
  * pdev_ptr	- PA of the PDEV