From 6fb3bead9e0eea3640ad1209347691c2b40512a2 Mon Sep 17 00:00:00 2001
From: Imre Kis <imre.kis@arm.com>
Date: Wed, 5 Feb 2025 14:27:45 +0100
Subject: [PATCH 2/8] Load initial image state in PSA FWU M update agent

Set initial image state based on the image state returned by
psa_fwu_query. This way the update agent has the correct view of images
after reboot and it can accept or reject them.

Upstream-Status: Submitted [https://review.trustedfirmware.org/c/TS/trusted-services/+/35155]
Signed-off-by: Imre Kis <imre.kis@arm.com>
Change-Id: I150e4fdb4584c8d07f5f1f642ee88197f9cff49b
Signed-off-by: Harsimran Singh Tungal <harsimransingh.tungal@arm.com>
---
 .../psa_fwu_m/agent/psa_fwu_m_update_agent.c  | 23 +++++++--
 .../test/test_psa_fwu_m_update_agent.cpp      | 49 ++++++++++++++++++-
 docs/services/fwu/psa-fwu-m.rst               | 14 +++++-
 3 files changed, 80 insertions(+), 6 deletions(-)

diff --git a/components/service/fwu/psa_fwu_m/agent/psa_fwu_m_update_agent.c b/components/service/fwu/psa_fwu_m/agent/psa_fwu_m_update_agent.c
index 6de9ba71..48b86f6e 100644
--- a/components/service/fwu/psa_fwu_m/agent/psa_fwu_m_update_agent.c
+++ b/components/service/fwu/psa_fwu_m/agent/psa_fwu_m_update_agent.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, Arm Limited. All rights reserved.
+ * Copyright (c) 2024-2025, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -609,8 +609,11 @@ struct update_agent *psa_fwu_m_update_agent_init(
 	const struct psa_fwu_m_image_mapping image_mapping[], size_t image_count,
 	uint32_t max_payload_size)
 {
+	psa_status_t psa_status = PSA_ERROR_GENERIC_ERROR;
 	struct psa_fwu_m_update_agent *context = NULL;
+	struct psa_fwu_component_info_t info = { 0 };
 	struct psa_fwu_m_image *images = NULL;
+	enum psa_fwu_m_state state = regular;
 	struct update_agent *agent = NULL;
 	size_t i = 0;
 
@@ -637,9 +640,23 @@ struct update_agent *psa_fwu_m_update_agent_init(
 	}
 
 	for (i = 0; i < image_count; i++) {
+		psa_status = psa_fwu_query(image_mapping[i].component, &info);
+		if (psa_status != PSA_SUCCESS) {
+			free(images);
+			free(context);
+			free(agent);
+			return NULL;
+		}
+
 		images[i].uuid = image_mapping[i].uuid;
 		images[i].component = image_mapping[i].component;
-		images[i].selected_for_staging = false;
+		if (info.state == PSA_FWU_TRIAL) {
+			images[i].selected_for_staging = true;
+			state = trial;
+		} else {
+			images[i].selected_for_staging = false;
+		}
+
 		images[i].read = NULL; /* Cannot read images */
 		images[i].write = image_write;
 	}
@@ -654,7 +671,7 @@ struct update_agent *psa_fwu_m_update_agent_init(
 	context->images = images;
 	context->image_count = image_count + 1;
 	context->max_payload_size = max_payload_size;
-	context->state = regular;
+	context->state = state;
 
 	agent->context = context;
 	agent->interface = &interface;
diff --git a/components/service/fwu/psa_fwu_m/agent/test/test_psa_fwu_m_update_agent.cpp b/components/service/fwu/psa_fwu_m/agent/test/test_psa_fwu_m_update_agent.cpp
index de289fff..3805d182 100644
--- a/components/service/fwu/psa_fwu_m/agent/test/test_psa_fwu_m_update_agent.cpp
+++ b/components/service/fwu/psa_fwu_m/agent/test/test_psa_fwu_m_update_agent.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, Arm Limited. All rights reserved.
+ * Copyright (c) 2024-2025, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -14,6 +14,9 @@
 
 TEST_GROUP(psa_fwu_m_update_agent) {
 	TEST_SETUP() {
+		psa_fwu_component_info_t info = {0};
+		expect_mock_psa_fwu_query(mapping[0].component, &info, PSA_SUCCESS);
+		expect_mock_psa_fwu_query(mapping[1].component, &info, PSA_SUCCESS);
 		agent = psa_fwu_m_update_agent_init(mapping, 2, 4096);
 		handle = 0;
 		progress = 0;
@@ -667,4 +670,46 @@ TEST(psa_fwu_m_update_agent, select_previous)
 
 	expect_mock_psa_fwu_reject(0, PSA_SUCCESS);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, update_agent_select_previous(agent));
-}
\ No newline at end of file
+}
+
+TEST(psa_fwu_m_update_agent, boot_in_trial_mode_query_fail) {
+	psa_fwu_component_info_t info = {0};
+
+	expect_mock_psa_fwu_query(mapping[0].component, &info, PSA_ERROR_GENERIC_ERROR);
+	POINTERS_EQUAL(NULL, psa_fwu_m_update_agent_init(mapping, 2, 4096));
+}
+
+TEST(psa_fwu_m_update_agent, boot_in_trial_mode_select_previous) {
+	psa_fwu_component_info_t info0 = {0};
+	psa_fwu_component_info_t info1 = {0};
+
+	info1.state = PSA_FWU_TRIAL;
+
+	expect_mock_psa_fwu_query(mapping[0].component, &info0, PSA_SUCCESS);
+	expect_mock_psa_fwu_query(mapping[1].component, &info1, PSA_SUCCESS);
+
+	update_agent *agent = psa_fwu_m_update_agent_init(mapping, 2, 4096);
+
+	expect_mock_psa_fwu_reject(0, PSA_SUCCESS);
+	LONGS_EQUAL(FWU_STATUS_SUCCESS, update_agent_select_previous(agent));
+
+	psa_fwu_m_update_agent_deinit(agent);
+}
+
+TEST(psa_fwu_m_update_agent, boot_in_trial_mode_accept) {
+	psa_fwu_component_info_t info0 = {0};
+	psa_fwu_component_info_t info1 = {0};
+
+	info1.state = PSA_FWU_TRIAL;
+
+	expect_mock_psa_fwu_query(mapping[0].component, &info0, PSA_SUCCESS);
+	expect_mock_psa_fwu_query(mapping[1].component, &info1, PSA_SUCCESS);
+
+	update_agent *agent = psa_fwu_m_update_agent_init(mapping, 2, 4096);
+
+	expect_mock_psa_fwu_accept(PSA_SUCCESS);
+	LONGS_EQUAL(FWU_STATUS_DENIED, update_agent_accept_image(agent, &mapping[0].uuid));
+	LONGS_EQUAL(FWU_STATUS_SUCCESS, update_agent_accept_image(agent, &mapping[1].uuid));
+
+	psa_fwu_m_update_agent_deinit(agent);
+}
diff --git a/docs/services/fwu/psa-fwu-m.rst b/docs/services/fwu/psa-fwu-m.rst
index 26ffed09..1358015f 100644
--- a/docs/services/fwu/psa-fwu-m.rst
+++ b/docs/services/fwu/psa-fwu-m.rst
@@ -44,6 +44,11 @@ The solutions to these differences:
 * Convert the image query result returned by FWU-M to FWU-A format. There are similar field, but this imposes some
   limitations.
 
+Initialization
+```````````````
+
+The initial image and agent state is determined based on the image state returned by ``psa_fwu_query()``.
+
 
 ``fwu_discover()``
 ``````````````````
@@ -71,6 +76,10 @@ agent switches to trial state, so the client can validate the new set of images
 
 On calling ``fwu_end_staging()`` the agent calls ``psa_fwu_finish()`` on each selected image, then calls
 ``psa_fwu_install()``. If all images have been accepted (see ``fwu_commit()``) it also calls ``psa_fwu_accept()``.
+The implementation treats ``PSA_SUCCESS_REBOOT`` and ``PSA_SUCCESS_RESTART`` status values as error. In an A+M system the M
+class side shouldn't restart the system, so calling ``psa_fwu_request_reboot()`` does not fit the system. There's also no
+PSA FWU A return code for inidicating the restart request to the normal world. If the normal world has to restart the
+system after ending the staging phase, it has to do it in an implementation defined way.
 
 .. uml:: ../uml/psa_fwu_m_update_agent/fwu_end_staging.puml
 
@@ -136,7 +145,10 @@ calls ``psa_fwu_accept()`` when all images have been accepted. This results in a
 `````````````````````````
 
 Selects previous working state (i.e. rejects the firmware update) and transitions back to regular state after calling
-``psa_fwu_reject()``.
+``psa_fwu_reject()``. The implementation treats ``PSA_SUCCESS_REBOOT`` and ``PSA_SUCCESS_RESTART`` status values as error.
+In an A+M system the M class side shouldn't restart the system, so calling ``psa_fwu_request_reboot()`` does not fit the
+system. There's also no PSA FWU A return code for inidicating the restart request to the normal world. If the normal
+world has to restart the system when rejecting the installed firmware, it has to do it in an implementation defined way.
 
 .. uml:: ../uml/psa_fwu_m_update_agent/fwu_select_previous.puml
 
-- 
2.25.1

