fix(realm): cleanup pdev after setup

Change-Id: Ia2f47fef71ce7336cece5f14b50e0ea4631c0358
Signed-off-by: Soby Mathew <soby.mathew@arm.com>
diff --git a/tftf/tests/runtime_services/host_realm_managment/host_da_helper.c b/tftf/tests/runtime_services/host_realm_managment/host_da_helper.c
index 7d193c8..8bcf370 100644
--- a/tftf/tests/runtime_services/host_realm_managment/host_da_helper.c
+++ b/tftf/tests/runtime_services/host_realm_managment/host_da_helper.c
@@ -497,11 +497,15 @@
 
 	/* Allocate granule for PDEV and delegate */
 	h_pdev->pdev = page_alloc(PAGE_SIZE);
+	if (h_pdev->pdev == NULL) {
+		return -1;
+	}
+
 	memset(h_pdev->pdev, 0, GRANULE_SIZE);
 	ret = host_rmi_granule_delegate((u_register_t)h_pdev->pdev);
 	if (ret != RMI_SUCCESS) {
 		ERROR("PDEV delegate failed 0x%lx\n", ret);
-		return -1;
+		goto err_undelegate_pdev;
 	}
 
 	/*
@@ -524,76 +528,111 @@
 	/* Allocate aux granules for PDEV and delegate */
 	INFO("PDEV create requires %u aux pages\n", h_pdev->pdev_aux_num);
 	for (i = 0; i < h_pdev->pdev_aux_num; i++) {
-		h_pdev->pdev_aux[i] = page_alloc(PAGE_SIZE);
-		ret = host_rmi_granule_delegate((u_register_t)h_pdev->pdev_aux[i]);
+		void *pdev_aux = page_alloc(PAGE_SIZE);
+
+		if (pdev_aux == NULL) {
+			goto err_undelegate_pdev_aux;
+		}
+
+		ret = host_rmi_granule_delegate((u_register_t)pdev_aux);
 		if (ret != RMI_SUCCESS) {
 			ERROR("Aux granule delegate failed 0x%lx\n", ret);
-			return -1;
+			goto err_undelegate_pdev;
 		}
+
+		h_pdev->pdev_aux[i] = pdev_aux;
 	}
 
 	/* Allocate dev_comm_data and send/recv buffer for Dev communication */
 	h_pdev->dev_comm_data = (struct rmi_dev_comm_data *)page_alloc(PAGE_SIZE);
+	if (h_pdev->dev_comm_data == NULL) {
+		goto err_undelegate_pdev_aux;
+	}
+
 	memset(h_pdev->dev_comm_data, 0, sizeof(struct rmi_dev_comm_data));
+
 	h_pdev->dev_comm_data->enter.req_addr = (unsigned long)
 		page_alloc(PAGE_SIZE);
+	if (h_pdev->dev_comm_data->enter.req_addr == 0UL) {
+		goto err_undelegate_pdev_aux;
+	}
+
 	h_pdev->dev_comm_data->enter.resp_addr = (unsigned long)
 		page_alloc(PAGE_SIZE);
+	if (h_pdev->dev_comm_data->enter.resp_addr == 0UL) {
+		goto err_undelegate_pdev_aux;
+	}
 
 	/* Allocate buffer to cache device certificate */
 	h_pdev->cert_slot_id = 0;
 	h_pdev->cert_chain = (uint8_t *)page_alloc(HOST_PDEV_CERT_LEN_MAX);
 	h_pdev->cert_chain_len = 0;
 	if (h_pdev->cert_chain == NULL) {
-		return -1;
+		goto err_undelegate_pdev_aux;
 	}
 
 	/* Allocate buffer to store extracted public key */
 	h_pdev->public_key = (void *)page_alloc(PAGE_SIZE);
 	if (h_pdev->public_key == NULL) {
-		return -1;
+		goto err_undelegate_pdev_aux;
 	}
 	h_pdev->public_key_len = PAGE_SIZE;
 
 	/* Allocate buffer to store public key metadata */
 	h_pdev->public_key_metadata = (void *)page_alloc(PAGE_SIZE);
 	if (h_pdev->public_key_metadata == NULL) {
-		return -1;
+		goto err_undelegate_pdev_aux;
 	}
+
 	h_pdev->public_key_metadata_len = PAGE_SIZE;
 
 	/* Set algorithm to use for device digests */
 	h_pdev->pdev_hash_algo = RMI_HASH_SHA_512;
 
 	return 0;
+
+err_undelegate_pdev_aux:
+	/* Undelegate all the delegated pages */
+	for (int i = 0; i < h_pdev->pdev_aux_num; i++) {
+		if (h_pdev->pdev_aux[i]) {
+			host_rmi_granule_undelegate((u_register_t)
+						    h_pdev->pdev_aux[i]);
+		}
+	}
+
+err_undelegate_pdev:
+	host_rmi_granule_undelegate((u_register_t)h_pdev->pdev);
+
+	return -1;
 }
 
+
 /*
  * Stop PDEV and ternimate secure session and call PDEV destroy
  */
 int host_pdev_reclaim(struct host_pdev *h_pdev)
 {
 	u_register_t ret;
-	int rc;
+	int rc, result = 0;
 
 	/* Move the device to STOPPING state */
 	rc = host_pdev_transition(h_pdev, RMI_PDEV_STATE_STOPPING);
 	if (rc != 0) {
 		ERROR("PDEV transition: to PDEV_STATE_STOPPING failed\n");
-		return rc;
+		result = -1;
 	}
 
 	/* Do pdev_communicate to terminate secure session */
 	rc = host_pdev_transition(h_pdev, RMI_PDEV_STATE_STOPPED);
 	if (rc != 0) {
 		ERROR("PDEV transition: to PDEV_STATE_STOPPED failed\n");
-		return rc;
+		result = -1;
 	}
 
 	rc = host_pdev_destroy(h_pdev);
 	if (rc != 0) {
 		ERROR("PDEV transition: to STATE_NULL failed\n");
-		return rc;
+		result = -1;
 	}
 
 	/* Undelegate all aux granules */
@@ -601,7 +640,7 @@
 		ret = host_rmi_granule_undelegate((u_register_t)h_pdev->pdev_aux[i]);
 		if (ret != RMI_SUCCESS) {
 			ERROR("Aux granule undelegate failed 0x%lx\n", ret);
-			return -1;
+			result = -1;
 		}
 	}
 
@@ -609,10 +648,10 @@
 	ret = host_rmi_granule_undelegate((u_register_t)h_pdev->pdev);
 	if (ret != RMI_SUCCESS) {
 		ERROR("PDEV undelegate failed 0x%lx\n", ret);
-		return -1;
+		result = -1;
 	}
 
-	return 0;
+	return result;
 }
 
 int host_create_realm_with_feat_da(struct realm *realm)
@@ -661,6 +700,10 @@
 
 	/* Allocate granule for VDEV and delegate */
 	h_vdev->vdev_ptr = (void *)page_alloc(PAGE_SIZE);
+	if (h_vdev->vdev_ptr == NULL) {
+		return -1;
+	}
+
 	memset(h_vdev->vdev_ptr, 0, GRANULE_SIZE);
 	ret = host_rmi_granule_delegate((u_register_t)h_vdev->vdev_ptr);
 	if (ret != RMI_SUCCESS) {
@@ -670,27 +713,42 @@
 
 	/* Allocate dev_comm_data and send/recv buffer for Dev communication */
 	h_vdev->dev_comm_data = (struct rmi_dev_comm_data *)page_alloc(PAGE_SIZE);
+	if (h_vdev->dev_comm_data == NULL) {
+		goto err_undel_vdev;
+	}
 	memset(h_vdev->dev_comm_data, 0, sizeof(struct rmi_dev_comm_data));
 	h_vdev->dev_comm_data->enter.req_addr = (unsigned long)
 		page_alloc(PAGE_SIZE);
+	if (h_vdev->dev_comm_data->enter.req_addr == 0UL) {
+		goto err_undel_vdev;
+	}
+
 	h_vdev->dev_comm_data->enter.resp_addr = (unsigned long)
 		page_alloc(PAGE_SIZE);
+	if (h_vdev->dev_comm_data->enter.resp_addr == 0UL) {
+		goto err_undel_vdev;
+	}
 
 	/* Allocate buffer to cache device measurements */
 	h_vdev->meas = (uint8_t *)page_alloc(HOST_VDEV_MEAS_LEN_MAX);
-	h_vdev->meas_len = 0;
 	if (h_vdev->meas == NULL) {
-		return -1;
+		goto err_undel_vdev;
 	}
 
+	h_vdev->meas_len = 0;
+
 	/* Allocate buffer to cache device interface report */
 	h_vdev->ifc_report = (uint8_t *)page_alloc(HOST_VDEV_IFC_REPORT_LEN_MAX);
 	h_vdev->ifc_report_len = 0;
 	if (h_vdev->ifc_report == NULL) {
-		return -1;
+		goto err_undel_vdev;
 	}
 
 	return 0;
+
+err_undel_vdev:
+	host_rmi_granule_undelegate((u_register_t)h_vdev->vdev_ptr);
+	return -1;
 }
 
 int host_assign_vdev_to_realm(struct realm *realm,
diff --git a/tftf/tests/runtime_services/host_realm_managment/host_rmi_da_flow.c b/tftf/tests/runtime_services/host_realm_managment/host_rmi_da_flow.c
index 0452f58..9a140ac 100644
--- a/tftf/tests/runtime_services/host_realm_managment/host_rmi_da_flow.c
+++ b/tftf/tests/runtime_services/host_realm_managment/host_rmi_da_flow.c
@@ -44,7 +44,7 @@
 	struct host_pdev *h_pdev;
 	struct host_vdev *h_vdev;
 	uint8_t public_key_algo;
-	int ret, rc;
+	int rc;
 	bool realm_rc;
 	struct realm realm;
 
@@ -53,18 +53,10 @@
 
 	INFO("DA on bdf: 0x%x, doe_cap_base: 0x%x\n", pdev_bdf, doe_cap_base);
 
-	/* Initialize Host NS heap memory */
-	ret = page_pool_init((u_register_t)PAGE_POOL_BASE,
-				(u_register_t)PAGE_POOL_MAX_SIZE);
-	if (ret != HEAP_INIT_SUCCESS) {
-		ERROR("Failed to init heap pool %d\n", ret);
-		return TEST_RESULT_FAIL;
-	}
-
 	/*
-	 * 2. Create a Realm with DA feature enabled
+	 * Create a Realm with DA feature enabled
 	 *
-	 * todo: creating this after host_pdev_setup cases Realm create to
+	 * todo: creating this after host_pdev_setup causes Realm create to
 	 * fail.
 	 */
 	rc = host_create_realm_with_feat_da(&realm);
@@ -81,8 +73,9 @@
 	/* Allocate granules. Skip DA ABIs if host_pdev_setup fails */
 	rc = host_pdev_setup(h_pdev);
 	if (rc == -1) {
-		INFO("host_pdev_setup failed. skipping DA ABIs...\n");
-		return TEST_RESULT_SKIPPED;
+		INFO("host_pdev_setup failed.\n");
+		(void)host_destroy_realm(&realm);
+		return TEST_RESULT_FAIL;
 	}
 
 	/* todo: move to tdi_pdev_setup */
@@ -93,14 +86,14 @@
 	rc = host_pdev_transition(h_pdev, RMI_PDEV_STATE_NEW);
 	if (rc != 0) {
 		ERROR("PDEV transition: NULL -> STATE_NEW failed\n");
-		return TEST_RESULT_FAIL;
+		goto err_pdev_reclaim;
 	}
 
 	/* Call rmi_pdev_communicate to transition PDEV to NEEDS_KEY */
 	rc = host_pdev_transition(h_pdev, RMI_PDEV_STATE_NEEDS_KEY);
 	if (rc != 0) {
 		ERROR("PDEV transition: PDEV_NEW -> PDEV_NEEDS_KEY failed\n");
-		return TEST_RESULT_FAIL;
+		goto err_pdev_reclaim;
 	}
 
 	/* Get public key. Verifying cert_chain not done by host but by Realm? */
@@ -113,7 +106,7 @@
 						 &public_key_algo);
 	if (rc != 0) {
 		ERROR("Get public key failed\n");
-		return TEST_RESULT_FAIL;
+		goto err_pdev_reclaim;
 	}
 
 	if (public_key_algo == PUBLIC_KEY_ALGO_ECDSA_ECC_NIST_P256) {
@@ -130,14 +123,14 @@
 	rc = host_pdev_transition(h_pdev, RMI_PDEV_STATE_HAS_KEY);
 	if (rc != 0) {
 		INFO("PDEV transition: PDEV_NEEDS_KEY -> PDEV_HAS_KEY failed\n");
-		return TEST_RESULT_FAIL;
+		goto err_pdev_reclaim;
 	}
 
 	/* Call rmi_pdev_comminucate to transition PDEV to READY state */
 	rc = host_pdev_transition(h_pdev, RMI_PDEV_STATE_READY);
 	if (rc != 0) {
 		INFO("PDEV transition: PDEV_HAS_KEY -> PDEV_READY failed\n");
-		return TEST_RESULT_FAIL;
+		goto err_pdev_reclaim;
 	}
 
 
@@ -147,7 +140,7 @@
 	rc = host_assign_vdev_to_realm(&realm, h_pdev, h_vdev);
 	if (rc != 0) {
 		INFO("VDEV assign to realm failed\n");
-		return TEST_RESULT_FAIL;
+		goto err_pdev_reclaim;
 	}
 
 	/*
@@ -157,7 +150,7 @@
 					    RMI_EXIT_HOST_CALL, 0U);
 	if (!realm_rc) {
 		INFO("Realm DA_RSI_CALLS failed\n");
-		return TEST_RESULT_FAIL;
+		goto err_pdev_reclaim;
 	}
 
 	/*
@@ -166,7 +159,7 @@
 	rc = host_unassign_vdev_from_realm(&realm, h_pdev, h_vdev);
 	if (rc != 0) {
 		INFO("VDEV unassign to realm failed\n");
-		return TEST_RESULT_FAIL;
+		goto err_pdev_reclaim;
 	}
 
 	/*
@@ -174,6 +167,7 @@
 	 */
 	if (!host_destroy_realm(&realm)) {
 		INFO("Realm destroy failed\n");
+		(void)host_pdev_reclaim(h_pdev);
 		return TEST_RESULT_FAIL;
 	}
 
@@ -187,6 +181,11 @@
 	}
 
 	return TEST_RESULT_SUCCESS;
+
+err_pdev_reclaim:
+	(void)host_destroy_realm(&realm);
+	(void)host_pdev_reclaim(h_pdev);
+	return TEST_RESULT_FAIL;
 }
 
 /*