Update to PSA FWU 1.0 protocol (DEN0118)

Update FWU provider and clients to follow the 1.0 version specification.

Signed-off-by: Imre Kis <imre.kis@arm.com>
Change-Id: I0710b0df388609a7c3f2ce4a5e98a9fe5c2e2b99
diff --git a/components/service/fwu/agent/update_agent.c b/components/service/fwu/agent/update_agent.c
index 77037b1..733a012 100644
--- a/components/service/fwu/agent/update_agent.c
+++ b/components/service/fwu/agent/update_agent.c
@@ -54,6 +54,18 @@
 
 static int cancel_staging(void *context);
 
+static int discover(void *context, struct fwu_discovery_result *result)
+{
+	result->service_status = 0;
+	result->version_major = FWU_PROTOCOL_VERSION_MAJOR;
+	result->version_minor = FWU_PROTOCOL_VERSION_MINOR;
+	result->max_payload_size = 0;
+	result->flags = 0;
+	result->vendor_specific_flags = 0;
+
+	return FWU_STATUS_SUCCESS;
+}
+
 static int begin_staging(void *context, uint32_t vendor_flags, uint32_t partial_update_count,
 			 const struct uuid_octets *update_guid)
 {
@@ -295,7 +307,7 @@
 
 
 static const struct update_agent_interface interface = {
-	.discover = NULL,
+	.discover = discover,
 	.begin_staging = begin_staging,
 	.end_staging = end_staging,
 	.cancel_staging = cancel_staging,
diff --git a/components/service/fwu/provider/fwu_provider.c b/components/service/fwu/provider/fwu_provider.c
index 33b744f..ded6db7 100644
--- a/components/service/fwu/provider/fwu_provider.c
+++ b/components/service/fwu/provider/fwu_provider.c
@@ -15,7 +15,12 @@
 #include "service/fwu/provider/serializer/fwu_provider_serializer.h"
 #include "fwu_uuid.h"
 
+#ifndef FWU_PROVIDER_MAX_PARTIAL_UPDATE_COUNT
+#define FWU_PROVIDER_MAX_PARTIAL_UPDATE_COUNT	(4)
+#endif /* FWU_PROVIDER_MAX_PARTIAL_UPDATE_COUNT */
+
 /* Service request handlers */
+static rpc_status_t discover_handler(void *context, struct rpc_request *req);
 static rpc_status_t begin_staging_handler(void *context, struct rpc_request *req);
 static rpc_status_t end_staging_handler(void *context, struct rpc_request *req);
 static rpc_status_t cancel_staging_handler(void *context, struct rpc_request *req);
@@ -28,15 +33,16 @@
 
 /* Handler mapping table for service */
 static const struct service_handler handler_table[] = {
-	{ TS_FWU_OPCODE_BEGIN_STAGING, begin_staging_handler },
-	{ TS_FWU_OPCODE_END_STAGING, end_staging_handler },
-	{ TS_FWU_OPCODE_CANCEL_STAGING, cancel_staging_handler },
-	{ TS_FWU_OPCODE_OPEN, open_handler },
-	{ TS_FWU_OPCODE_WRITE_STREAM, write_stream_handler },
-	{ TS_FWU_OPCODE_READ_STREAM, read_stream_handler },
-	{ TS_FWU_OPCODE_COMMIT, commit_handler },
-	{ TS_FWU_OPCODE_ACCEPT_IMAGE, accept_image_handler },
-	{ TS_FWU_OPCODE_SELECT_PREVIOUS, select_previous_handler }
+	{ FWU_FUNC_ID_DISCOVER, discover_handler },
+	{ FWU_FUNC_ID_BEGIN_STAGING, begin_staging_handler },
+	{ FWU_FUNC_ID_END_STAGING, end_staging_handler },
+	{ FWU_FUNC_ID_CANCEL_STAGING, cancel_staging_handler },
+	{ FWU_FUNC_ID_OPEN, open_handler },
+	{ FWU_FUNC_ID_WRITE_STREAM, write_stream_handler },
+	{ FWU_FUNC_ID_READ_STREAM, read_stream_handler },
+	{ FWU_FUNC_ID_COMMIT, commit_handler },
+	{ FWU_FUNC_ID_ACCEPT_IMAGE, accept_image_handler },
+	{ FWU_FUNC_ID_SELECT_PREVIOUS, select_previous_handler }
 };
 
 struct rpc_service_interface *fwu_provider_init(struct fwu_provider *context,
@@ -80,13 +86,89 @@
 	return serializer;
 }
 
+static uint16_t generate_function_presence(const struct update_agent *agent,
+					   uint8_t function_presence[FWU_FUNC_ID_COUNT])
+{
+	uint16_t num_func = 0;
+
+#define ADD_FUNC_IF_PRESENT(func, id) \
+do { \
+	if (agent->interface->func != NULL) \
+		function_presence[num_func++] = (id); \
+} while (0)
+
+	ADD_FUNC_IF_PRESENT(discover, FWU_FUNC_ID_DISCOVER);
+	ADD_FUNC_IF_PRESENT(begin_staging, FWU_FUNC_ID_BEGIN_STAGING);
+	ADD_FUNC_IF_PRESENT(end_staging, FWU_FUNC_ID_END_STAGING);
+	ADD_FUNC_IF_PRESENT(cancel_staging, FWU_FUNC_ID_CANCEL_STAGING);
+	ADD_FUNC_IF_PRESENT(open, FWU_FUNC_ID_OPEN);
+	ADD_FUNC_IF_PRESENT(write_stream, FWU_FUNC_ID_WRITE_STREAM);
+	ADD_FUNC_IF_PRESENT(read_stream, FWU_FUNC_ID_READ_STREAM);
+	ADD_FUNC_IF_PRESENT(commit, FWU_FUNC_ID_COMMIT);
+	ADD_FUNC_IF_PRESENT(accept_image, FWU_FUNC_ID_ACCEPT_IMAGE);
+	ADD_FUNC_IF_PRESENT(select_previous, FWU_FUNC_ID_SELECT_PREVIOUS);
+
+#undef ADD_FUNC_IF_PRESENT
+
+	return num_func;
+}
+
+static rpc_status_t discover_handler(void *context, struct rpc_request *req)
+{
+	rpc_status_t rpc_status = RPC_ERROR_INTERNAL;
+	struct fwu_provider *this_instance = (struct fwu_provider *)context;
+	const struct fwu_provider_serializer *serializer = get_fwu_serializer(this_instance, req);
+	struct fwu_discovery_result discovery_result = { 0 };
+	struct rpc_buffer *resp_buf = &req->response;
+
+	if (!serializer)
+		return rpc_status;
+
+	req->service_status = update_agent_discover(this_instance->update_agent, &discovery_result);
+
+	if (!req->service_status) {
+		uint16_t num_func = 0;
+		uint8_t function_presence[FWU_FUNC_ID_COUNT] = { 0 };
+
+		num_func = generate_function_presence(this_instance->update_agent,
+						      function_presence);
+
+		rpc_status = serializer->serialize_discover_resp(
+			resp_buf, discovery_result.service_status, discovery_result.version_major,
+			discovery_result.version_minor, num_func, discovery_result.max_payload_size,
+			discovery_result.flags, discovery_result.vendor_specific_flags,
+			function_presence);
+	} else {
+		/*
+		 * The actual service call failed, but the request was successful on the RPC level
+		 */
+		rpc_status = RPC_SUCCESS;
+	}
+
+	return rpc_status;
+}
+
 static rpc_status_t begin_staging_handler(void *context, struct rpc_request *req)
 {
+	rpc_status_t rpc_status = RPC_ERROR_INTERNAL;
+	struct rpc_buffer *req_buf = &req->request;
 	struct fwu_provider *this_instance = (struct fwu_provider *)context;
+	const struct fwu_provider_serializer *serializer = get_fwu_serializer(this_instance, req);
+	uint32_t vendor_flags = 0;
+	uint32_t partial_update_count = 0;
+	struct uuid_octets update_guid[FWU_PROVIDER_MAX_PARTIAL_UPDATE_COUNT];
 
-	req->service_status = update_agent_begin_staging(this_instance->update_agent, 0, 0, NULL);
+	if (serializer)
+		rpc_status = serializer->deserialize_begin_staging_req(
+			req_buf, &vendor_flags, &partial_update_count,
+			FWU_PROVIDER_MAX_PARTIAL_UPDATE_COUNT, update_guid);
 
-	return RPC_SUCCESS;
+	if (rpc_status == RPC_SUCCESS)
+		req->service_status = update_agent_begin_staging(
+			this_instance->update_agent, vendor_flags, partial_update_count,
+			update_guid);
+
+	return rpc_status;
 }
 
 static rpc_status_t end_staging_handler(void *context, struct rpc_request *req)
@@ -113,15 +195,16 @@
 	struct rpc_buffer *req_buf = &req->request;
 	struct fwu_provider *this_instance = (struct fwu_provider *)context;
 	const struct fwu_provider_serializer *serializer = get_fwu_serializer(this_instance, req);
-	struct uuid_octets image_type_uuid;
+	struct uuid_octets image_type_uuid = { 0 };
+	uint8_t op_type = 0;
 
 	if (serializer)
-		rpc_status = serializer->deserialize_open_req(req_buf, &image_type_uuid);
+		rpc_status = serializer->deserialize_open_req(req_buf, &image_type_uuid, &op_type);
 
 	if (rpc_status == RPC_SUCCESS) {
 		uint32_t handle = 0;
 		req->service_status =
-			update_agent_open(this_instance->update_agent, &image_type_uuid, 0,
+			update_agent_open(this_instance->update_agent, &image_type_uuid, op_type,
 					  &handle);
 
 		if (!req->service_status) {
diff --git a/components/service/fwu/provider/serializer/fwu_provider_serializer.h b/components/service/fwu/provider/serializer/fwu_provider_serializer.h
index c9df194..ee44abe 100644
--- a/components/service/fwu/provider/serializer/fwu_provider_serializer.h
+++ b/components/service/fwu/provider/serializer/fwu_provider_serializer.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2022-2024, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -21,9 +21,25 @@
  * implement this interface.
  */
 struct fwu_provider_serializer {
+	/* Operation: discover */
+	rpc_status_t (*serialize_discover_resp)(const struct rpc_buffer *resp_buf,
+						int16_t service_status, uint8_t version_major,
+						uint8_t version_minor, uint16_t num_func,
+						uint64_t max_payload_size, uint32_t flags,
+						uint32_t vendor_specific_flags,
+						uint8_t *function_presence);
+
+	/* Operation: begin staging */
+	rpc_status_t (*deserialize_begin_staging_req)(const struct rpc_buffer *req_buf,
+						      uint32_t *vendor_flags,
+						      uint32_t *partial_update_count,
+						      uint32_t max_update_count,
+						      struct uuid_octets *update_guid);
+
 	/* Operation: open */
 	rpc_status_t (*deserialize_open_req)(const struct rpc_buffer *req_buf,
-					     struct uuid_octets *image_type_uuid);
+					     struct uuid_octets *image_type_uuid,
+					     uint8_t *op_type);
 
 	rpc_status_t (*serialize_open_resp)(struct rpc_buffer *resp_buf, uint32_t handle);
 
diff --git a/components/service/fwu/provider/serializer/packed-c/packedc_fwu_provider_serializer.c b/components/service/fwu/provider/serializer/packed-c/packedc_fwu_provider_serializer.c
index cb013d4..e45160f 100644
--- a/components/service/fwu/provider/serializer/packed-c/packedc_fwu_provider_serializer.c
+++ b/components/service/fwu/provider/serializer/packed-c/packedc_fwu_provider_serializer.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2022-2024, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -9,9 +9,73 @@
 
 #include "protocols/rpc/common/packed-c/status.h"
 #include "protocols/service/fwu/packed-c/fwu_proto.h"
+#include "util.h"
+
+rpc_status_t serialize_discover_resp(const struct rpc_buffer *resp_buf, int16_t service_status,
+				     uint8_t version_major, uint8_t version_minor,
+				     uint16_t num_func, uint64_t max_payload_size, uint32_t flags,
+				     uint32_t vendor_specific_flags, uint8_t *function_presence)
+{
+	rpc_status_t rpc_status = RPC_ERROR_INVALID_RESPONSE_BODY;
+	struct ts_fwu_discover_out *resp_msg = NULL;
+	size_t len = 0;
+
+	if (ADD_OVERFLOW(sizeof(*resp_msg), num_func, &len))
+		return RPC_ERROR_INVALID_RESPONSE_BODY;
+
+	if (len <= resp_buf->size) {
+		resp_msg = (struct ts_fwu_discover_out *)resp_buf->data;
+
+		resp_msg->service_status = service_status;
+		resp_msg->version_major = version_major;
+		resp_msg->version_minor = version_minor;
+		resp_msg->off_function_presence =
+			offsetof(struct ts_fwu_discover_out, function_presence);
+		resp_msg->num_func = num_func;
+		resp_msg->max_payload_size = max_payload_size;
+		resp_msg->flags = flags;
+		resp_msg->vendor_specific_flags = vendor_specific_flags;
+		memcpy(resp_msg->function_presence, function_presence, num_func);
+
+		rpc_status = RPC_SUCCESS;
+	}
+
+	return rpc_status;
+}
+
+rpc_status_t deserialize_begin_staging_req(const struct rpc_buffer *req_buf, uint32_t *vendor_flags,
+					   uint32_t *partial_update_count,
+					   uint32_t max_update_count,
+					   struct uuid_octets *update_guid)
+{
+	rpc_status_t rpc_status = RPC_ERROR_INVALID_REQUEST_BODY;
+	size_t expected_fixed_len = sizeof(struct ts_fwu_begin_staging_in);
+
+	if (expected_fixed_len <= req_buf->data_length) {
+		const struct ts_fwu_begin_staging_in *recv_msg =
+			(const struct ts_fwu_begin_staging_in *)req_buf->data;
+		size_t full_len = 0;
+
+		if (ADD_OVERFLOW(expected_fixed_len, recv_msg->partial_update_count, &full_len))
+			return RPC_ERROR_INVALID_REQUEST_BODY;
+
+		if (recv_msg->partial_update_count > max_update_count)
+			return RPC_ERROR_INTERNAL;
+
+		*vendor_flags = recv_msg->vendor_flags;
+		*partial_update_count = recv_msg->partial_update_count;
+
+		memcpy(update_guid, recv_msg->update_guid,
+		       UUID_OCTETS_LEN * recv_msg->partial_update_count);
+
+		rpc_status = RPC_SUCCESS;
+	}
+
+	return rpc_status;
+}
 
 static rpc_status_t deserialize_open_req(const struct rpc_buffer *req_buf,
-					 struct uuid_octets *image_type_uuid)
+					 struct uuid_octets *image_type_uuid, uint8_t *op_type)
 {
 	rpc_status_t rpc_status = RPC_ERROR_INVALID_REQUEST_BODY;
 	size_t expected_fixed_len = sizeof(struct ts_fwu_open_in);
@@ -20,9 +84,14 @@
 		const struct ts_fwu_open_in *recv_msg =
 			(const struct ts_fwu_open_in *)req_buf->data;
 
-		memcpy(image_type_uuid->octets, recv_msg->image_type_uuid, UUID_OCTETS_LEN);
-
-		rpc_status = RPC_SUCCESS;
+		if (recv_msg->op_type == FWU_OPEN_OP_TYPE_READ ||
+		    recv_msg->op_type == FWU_OPEN_OP_TYPE_WRITE) {
+			memcpy(image_type_uuid->octets, recv_msg->image_type_uuid, UUID_OCTETS_LEN);
+			*op_type = recv_msg->op_type;
+			rpc_status = RPC_SUCCESS;
+		} else {
+			rpc_status = RPC_ERROR_INVALID_REQUEST_BODY;
+		}
 	}
 
 	return rpc_status;
@@ -181,9 +250,17 @@
 const struct fwu_provider_serializer *packedc_fwu_provider_serializer_instance(void)
 {
 	static const struct fwu_provider_serializer instance = {
-		deserialize_open_req,	     serialize_open_resp,      deserialize_write_stream_req,
-		deserialize_read_stream_req, read_stream_resp_payload, serialize_read_stream_resp,
-		deserialize_commit_req,	     serialize_commit_resp,    deserialize_accept_req
+		serialize_discover_resp,
+		deserialize_begin_staging_req,
+		deserialize_open_req,
+		serialize_open_resp,
+		deserialize_write_stream_req,
+		deserialize_read_stream_req,
+		read_stream_resp_payload,
+		serialize_read_stream_resp,
+		deserialize_commit_req,
+		serialize_commit_resp,
+		deserialize_accept_req
 	};
 
 	return &instance;
diff --git a/components/service/fwu/test/fwu_client/direct/direct_fwu_client.cpp b/components/service/fwu/test/fwu_client/direct/direct_fwu_client.cpp
index ecca254..89f870c 100644
--- a/components/service/fwu/test/fwu_client/direct/direct_fwu_client.cpp
+++ b/components/service/fwu/test/fwu_client/direct/direct_fwu_client.cpp
@@ -9,6 +9,7 @@
 #include <cstring>
 
 #include "service/fwu/common/update_agent_interface.h"
+#include "protocols/service/fwu/packed-c/fwu_proto.h"
 
 direct_fwu_client::direct_fwu_client(struct update_agent **update_agent)
 	: fwu_client()
@@ -25,9 +26,54 @@
 {
 }
 
-int direct_fwu_client::begin_staging(void)
+int direct_fwu_client::discover(int16_t *service_status, uint8_t *version_major,
+				uint8_t *version_minor, uint16_t *num_func,
+				uint64_t *max_payload_size, uint32_t *flags,
+				uint32_t *vendor_specific_flags, uint8_t *function_presence)
 {
-	return update_agent_begin_staging(*m_update_agent, 0, 0, NULL);
+	struct fwu_discovery_result discovery_result = { 0 };
+	int res = 0;
+
+	res = update_agent_discover(*m_update_agent, &discovery_result);
+	if (res)
+		return res;
+
+	*service_status = discovery_result.service_status;
+	*version_major = discovery_result.version_major;
+	*version_minor = discovery_result.version_minor;
+	*max_payload_size = discovery_result.max_payload_size;
+	*flags = discovery_result.flags;
+	*vendor_specific_flags = discovery_result.vendor_specific_flags;
+
+	*num_func = 0;
+
+#define ADD_FUNC_IF_PRESENT(func, id) \
+do { \
+	if ((*m_update_agent)->interface->func != NULL) \
+		function_presence[(*num_func)++] = (id); \
+} while (0)
+
+	ADD_FUNC_IF_PRESENT(discover, FWU_FUNC_ID_DISCOVER);
+	ADD_FUNC_IF_PRESENT(begin_staging, FWU_FUNC_ID_BEGIN_STAGING);
+	ADD_FUNC_IF_PRESENT(end_staging, FWU_FUNC_ID_END_STAGING);
+	ADD_FUNC_IF_PRESENT(cancel_staging, FWU_FUNC_ID_CANCEL_STAGING);
+	ADD_FUNC_IF_PRESENT(open, FWU_FUNC_ID_OPEN);
+	ADD_FUNC_IF_PRESENT(write_stream, FWU_FUNC_ID_WRITE_STREAM);
+	ADD_FUNC_IF_PRESENT(read_stream, FWU_FUNC_ID_READ_STREAM);
+	ADD_FUNC_IF_PRESENT(commit, FWU_FUNC_ID_COMMIT);
+	ADD_FUNC_IF_PRESENT(accept_image, FWU_FUNC_ID_ACCEPT_IMAGE);
+	ADD_FUNC_IF_PRESENT(select_previous, FWU_FUNC_ID_SELECT_PREVIOUS);
+
+#undef ADD_FUNC_IF_PRESENT
+
+	return res;
+}
+
+int direct_fwu_client::begin_staging(uint32_t vendor_flags, uint32_t partial_update_count,
+				     struct uuid_octets update_guid[])
+{
+	return update_agent_begin_staging(*m_update_agent, vendor_flags, partial_update_count,
+					  update_guid);
 }
 
 int direct_fwu_client::end_staging(void)
@@ -50,9 +96,9 @@
 	return update_agent_select_previous(*m_update_agent);
 }
 
-int direct_fwu_client::open(const struct uuid_octets *uuid, uint32_t *handle)
+int direct_fwu_client::open(const struct uuid_octets *uuid, op_type op_type, uint32_t *handle)
 {
-	return update_agent_open(*m_update_agent, uuid, 0, handle);
+	return update_agent_open(*m_update_agent, uuid, (uint8_t)op_type, handle);
 }
 
 int direct_fwu_client::commit(uint32_t handle, bool accepted)
diff --git a/components/service/fwu/test/fwu_client/direct/direct_fwu_client.h b/components/service/fwu/test/fwu_client/direct/direct_fwu_client.h
index 3cd332d..744c824 100644
--- a/components/service/fwu/test/fwu_client/direct/direct_fwu_client.h
+++ b/components/service/fwu/test/fwu_client/direct/direct_fwu_client.h
@@ -22,7 +22,12 @@
 	explicit direct_fwu_client(struct update_agent **update_agent);
 	~direct_fwu_client();
 
-	int begin_staging(void);
+	int discover(int16_t *service_status, uint8_t *version_major, uint8_t *version_minor,
+		     uint16_t *num_func, uint64_t *max_payload_size, uint32_t *flags,
+		     uint32_t *vendor_specific_flags, uint8_t *function_presence);
+
+	int begin_staging(uint32_t vendor_flags, uint32_t partial_update_count,
+			  struct uuid_octets update_guid[]);
 
 	int end_staging(void);
 
@@ -32,7 +37,7 @@
 
 	int select_previous(void);
 
-	int open(const struct uuid_octets *uuid, uint32_t *handle);
+	int open(const struct uuid_octets *uuid, op_type op_type, uint32_t *handle);
 
 	int commit(uint32_t handle, bool accepted);
 
diff --git a/components/service/fwu/test/fwu_client/fwu_client.h b/components/service/fwu/test/fwu_client/fwu_client.h
index 55f4c87..072d647 100644
--- a/components/service/fwu/test/fwu_client/fwu_client.h
+++ b/components/service/fwu/test/fwu_client/fwu_client.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2022-2024, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -19,6 +19,11 @@
  */
 class fwu_client {
 public:
+	enum class op_type {
+		READ = 0,
+		WRITE = 1,
+	};
+
 	fwu_client()
 	{
 	}
@@ -27,7 +32,13 @@
 	{
 	}
 
-	virtual int begin_staging(void) = 0;
+	virtual int discover(int16_t *service_status, uint8_t *version_major,
+			     uint8_t *version_minor, uint16_t *num_func, uint64_t *max_payload_size,
+			     uint32_t *flags, uint32_t *vendor_specific_flags,
+			     uint8_t *function_presence) = 0;
+
+	virtual int begin_staging(uint32_t vendor_flags, uint32_t partial_update_count,
+				  struct uuid_octets update_guid[]) = 0;
 
 	virtual int end_staging(void) = 0;
 
@@ -37,7 +48,7 @@
 
 	virtual int select_previous(void) = 0;
 
-	virtual int open(const struct uuid_octets *uuid, uint32_t *handle) = 0;
+	virtual int open(const struct uuid_octets *uuid, op_type op_type, uint32_t *handle) = 0;
 
 	virtual int commit(uint32_t handle, bool accepted) = 0;
 
diff --git a/components/service/fwu/test/fwu_client/remote/remote_fwu_client.cpp b/components/service/fwu/test/fwu_client/remote/remote_fwu_client.cpp
index 803dfe6..2ad8ce1 100644
--- a/components/service/fwu/test/fwu_client/remote/remote_fwu_client.cpp
+++ b/components/service/fwu/test/fwu_client/remote/remote_fwu_client.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2022-2024, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -15,6 +15,7 @@
 #include "protocols/service/fwu/packed-c/fwu_proto.h"
 #include "protocols/service/fwu/packed-c/opcodes.h"
 #include "protocols/service/fwu/packed-c/status.h"
+#include "util.h"
 
 remote_fwu_client::remote_fwu_client()
 	: fwu_client()
@@ -97,19 +98,114 @@
 	return fwu_status;
 }
 
-int remote_fwu_client::begin_staging(void)
+int remote_fwu_client::discover(int16_t *service_status, uint8_t *version_major,
+				uint8_t *version_minor, uint16_t *num_func,
+				uint64_t *max_payload_size, uint32_t *flags,
+				uint32_t *vendor_specific_flags, uint8_t *function_presence)
 {
-	return invoke_no_param(TS_FWU_OPCODE_BEGIN_STAGING);
+	int fwu_status = FWU_STATUS_NOT_AVAILABLE;
+	struct ts_fwu_discover_out *resp_msg = NULL;
+	size_t max_resp_len = 0;
+
+	if (!m_service_context)
+		return fwu_status;
+
+	if (ADD_OVERFLOW(FWU_FUNC_ID_COUNT, sizeof(*resp_msg), &max_resp_len))
+		return FWU_STATUS_OUT_OF_BOUNDS;
+
+	rpc_call_handle call_handle = NULL;
+	uint8_t *req_buf = NULL;
+
+	call_handle = rpc_caller_session_begin(m_rpc_session, &req_buf, 0, max_resp_len);
+
+	if (call_handle) {
+		uint8_t *resp_buf = NULL;
+		size_t resp_len = 0;
+		service_status_t req_service_status = 0;
+
+		m_client.rpc_status = rpc_caller_session_invoke(call_handle, FWU_FUNC_ID_DISCOVER,
+								&resp_buf, &resp_len,
+								&req_service_status);
+
+		if (m_client.rpc_status == TS_RPC_CALL_ACCEPTED) {
+			fwu_status = req_service_status;
+
+			if (fwu_status == FWU_STATUS_SUCCESS) {
+				if (resp_len <= max_resp_len) {
+					resp_msg = (struct ts_fwu_discover_out *)resp_buf;
+
+					*service_status = resp_msg->service_status;
+					*version_major = resp_msg->version_major;
+					*num_func = resp_msg->num_func;
+					*max_payload_size = resp_msg->max_payload_size;
+					*flags = resp_msg->flags;
+					*vendor_specific_flags = resp_msg->vendor_specific_flags;
+					memcpy(function_presence, resp_msg->function_presence,
+					resp_msg->num_func);
+				} else {
+					fwu_status = FWU_STATUS_OUT_OF_BOUNDS;
+				}
+			}
+		}
+
+		rpc_caller_session_end(call_handle);
+	}
+
+	return fwu_status;
+}
+
+int remote_fwu_client::begin_staging(uint32_t vendor_flags, uint32_t partial_update_count,
+				     struct uuid_octets update_guid[])
+{
+	int fwu_status = FWU_STATUS_NOT_AVAILABLE;
+	struct ts_fwu_begin_staging_in *req_msg = NULL;
+	size_t req_len = 0;
+	size_t uuids_size = 0;
+
+	if (!m_service_context)
+		return fwu_status;
+
+	if (MUL_OVERFLOW(partial_update_count, sizeof(*update_guid), &uuids_size) ||
+	    ADD_OVERFLOW(uuids_size, sizeof(*req_msg), &req_len))
+		return fwu_status;
+
+	rpc_call_handle call_handle = NULL;
+	uint8_t *req_buf = NULL;
+
+	call_handle = rpc_caller_session_begin(m_rpc_session, &req_buf, req_len, 0);
+
+	if (call_handle) {
+		uint8_t *resp_buf = NULL;
+		size_t resp_len = 0;
+		service_status_t service_status = 0;
+
+		req_msg = (struct ts_fwu_begin_staging_in *)req_buf;
+		req_msg->reserved = 0;
+		req_msg->vendor_flags = vendor_flags;
+		req_msg->partial_update_count = partial_update_count;
+		memcpy(req_msg->update_guid, update_guid, uuids_size);
+
+		m_client.rpc_status = rpc_caller_session_invoke(call_handle,
+							FWU_FUNC_ID_BEGIN_STAGING,
+							&resp_buf, &resp_len, &service_status);
+
+		if (m_client.rpc_status == TS_RPC_CALL_ACCEPTED)
+			fwu_status = service_status;
+
+		rpc_caller_session_end(call_handle);
+	}
+
+	return fwu_status;
 }
 
 int remote_fwu_client::end_staging(void)
 {
-	return invoke_no_param(TS_FWU_OPCODE_END_STAGING);
+	return invoke_no_param(FWU_FUNC_ID_END_STAGING);
 }
 
 int remote_fwu_client::cancel_staging(void)
 {
-	return invoke_no_param(TS_FWU_OPCODE_CANCEL_STAGING);
+	return invoke_no_param(FWU_FUNC_ID_CANCEL_STAGING);
 }
 
 int remote_fwu_client::accept(const struct uuid_octets *image_type_uuid)
@@ -136,7 +232,7 @@
 		memcpy(req_buf, &req_msg, req_len);
 
 		m_client.rpc_status = rpc_caller_session_invoke(call_handle,
-							TS_FWU_OPCODE_ACCEPT_IMAGE,
+							FWU_FUNC_ID_ACCEPT_IMAGE,
 							&resp_buf, &resp_len, &service_status);
 
 		if (m_client.rpc_status == TS_RPC_CALL_ACCEPTED)
@@ -150,10 +246,10 @@
 
 int remote_fwu_client::select_previous(void)
 {
-	return invoke_no_param(TS_FWU_OPCODE_SELECT_PREVIOUS);
+	return invoke_no_param(FWU_FUNC_ID_SELECT_PREVIOUS);
 }
 
-int remote_fwu_client::open(const struct uuid_octets *uuid, uint32_t *handle)
+int remote_fwu_client::open(const struct uuid_octets *uuid, op_type op_type, uint32_t *handle)
 {
 	int fwu_status = FWU_STATUS_NOT_AVAILABLE;
 	struct ts_fwu_open_in req_msg = { 0 };
@@ -163,6 +259,7 @@
 		return fwu_status;
 
 	memcpy(req_msg.image_type_uuid, uuid->octets, OSF_UUID_OCTET_LEN);
+	req_msg.op_type = static_cast<uint8_t>(op_type);
 
 	rpc_call_handle call_handle;
 	uint8_t *req_buf;
@@ -178,7 +275,7 @@
 		memcpy(req_buf, &req_msg, req_len);
 
 		m_client.rpc_status = rpc_caller_session_invoke(call_handle,
-							TS_FWU_OPCODE_OPEN, &resp_buf,
+							FWU_FUNC_ID_OPEN, &resp_buf,
 							&resp_len, &service_status);
 
 		if (m_client.rpc_status == TS_RPC_CALL_ACCEPTED) {
@@ -224,7 +321,7 @@
 		memcpy(req_buf, &req_msg, req_len);
 
 		m_client.rpc_status = rpc_caller_session_invoke(call_handle,
-							TS_FWU_OPCODE_COMMIT, &resp_buf,
+							FWU_FUNC_ID_COMMIT, &resp_buf,
 							&resp_len, &service_status);
 
 		if (m_client.rpc_status == TS_RPC_CALL_ACCEPTED)
@@ -280,7 +377,7 @@
 			total_written += write_len;
 
 			m_client.rpc_status = rpc_caller_session_invoke(call_handle,
-								TS_FWU_OPCODE_WRITE_STREAM,
+								FWU_FUNC_ID_WRITE_STREAM,
 								&resp_buf, &resp_len, &service_status);
 
 			rpc_caller_session_end(call_handle);
@@ -327,7 +424,7 @@
 		memcpy(req_buf, &req_msg, req_len);
 
 		m_client.rpc_status = rpc_caller_session_invoke(call_handle,
-							TS_FWU_OPCODE_READ_STREAM,
+							FWU_FUNC_ID_READ_STREAM,
 							&resp_buf, &resp_len, &service_status);
 
 		size_t proto_overhead = offsetof(ts_fwu_read_stream_out, payload);
diff --git a/components/service/fwu/test/fwu_client/remote/remote_fwu_client.h b/components/service/fwu/test/fwu_client/remote/remote_fwu_client.h
index ce0511c..1370e9f 100644
--- a/components/service/fwu/test/fwu_client/remote/remote_fwu_client.h
+++ b/components/service/fwu/test/fwu_client/remote/remote_fwu_client.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2022-2024, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -21,7 +21,12 @@
 	remote_fwu_client();
 	~remote_fwu_client();
 
-	int begin_staging(void);
+	int discover(int16_t *service_status, uint8_t *version_major, uint8_t *version_minor,
+		     uint16_t *num_func, uint64_t *max_payload_size, uint32_t *flags,
+		     uint32_t *vendor_specific_flags, uint8_t *function_presence);
+
+	int begin_staging(uint32_t vendor_flags, uint32_t partial_update_count,
+			  struct uuid_octets update_guid[]);
 
 	int end_staging(void);
 
@@ -31,7 +36,7 @@
 
 	int select_previous(void);
 
-	int open(const struct uuid_octets *uuid, uint32_t *handle);
+	int open(const struct uuid_octets *uuid, op_type op_type, uint32_t *handle);
 
 	int commit(uint32_t handle, bool accepted);
 
diff --git a/components/service/fwu/test/image_directory_checker/image_directory_checker.cpp b/components/service/fwu/test/image_directory_checker/image_directory_checker.cpp
index 49610c1..8aaa321 100644
--- a/components/service/fwu/test/image_directory_checker/image_directory_checker.cpp
+++ b/components/service/fwu/test/image_directory_checker/image_directory_checker.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2022-2024, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -34,7 +34,7 @@
 
 	uuid_guid_octets_from_canonical(&uuid, FWU_DIRECTORY_CANONICAL_UUID);
 
-	status = fwu_client->open(&uuid, &stream_handle);
+	status = fwu_client->open(&uuid, fwu_client::op_type::READ, &stream_handle);
 	if (status)
 		return status;
 
diff --git a/components/service/fwu/test/metadata_fetcher/client/client_metadata_fetcher.cpp b/components/service/fwu/test/metadata_fetcher/client/client_metadata_fetcher.cpp
index 77d211c..8a64637 100644
--- a/components/service/fwu/test/metadata_fetcher/client/client_metadata_fetcher.cpp
+++ b/components/service/fwu/test/metadata_fetcher/client/client_metadata_fetcher.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2022-2024, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -40,7 +40,7 @@
 
 	uuid_guid_octets_from_canonical(&uuid, FWU_METADATA_CANONICAL_UUID);
 
-	status = m_fwu_client->open(&uuid, &stream_handle);
+	status = m_fwu_client->open(&uuid, fwu_client::op_type::READ, &stream_handle);
 	LONGS_EQUAL(0, status);
 
 	size_t read_len = 0;
diff --git a/components/service/fwu/test/ref_scenarios/image_directory_tests.cpp b/components/service/fwu/test/ref_scenarios/image_directory_tests.cpp
index 453525b..126f3fa 100644
--- a/components/service/fwu/test/ref_scenarios/image_directory_tests.cpp
+++ b/components/service/fwu/test/ref_scenarios/image_directory_tests.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2022-2024, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -91,7 +91,7 @@
 		int status = 0;
 		uint32_t stream_handle = 0;
 
-		status = m_fwu_client->open(&uuid, &stream_handle);
+		status = m_fwu_client->open(&uuid, fwu_client::op_type::WRITE, &stream_handle);
 		LONGS_EQUAL(0, status);
 
 		stream_handles.push_back(stream_handle);
diff --git a/components/service/fwu/test/ref_scenarios/invalid_behaviour_tests.cpp b/components/service/fwu/test/ref_scenarios/invalid_behaviour_tests.cpp
index 100940a..ec77406 100644
--- a/components/service/fwu/test/ref_scenarios/invalid_behaviour_tests.cpp
+++ b/components/service/fwu/test/ref_scenarios/invalid_behaviour_tests.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2022-2024, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -64,7 +64,7 @@
 
 	/* Open a fw image when not STAGING */
 	m_dut->whole_volume_image_type_uuid(0, &uuid);
-	status = m_fwu_client->open(&uuid, &stream_handle);
+	status = m_fwu_client->open(&uuid, fwu_client::op_type::WRITE, &stream_handle);
 	LONGS_EQUAL(FWU_STATUS_DENIED, status);
 
 	/* Write to a stream when not STAGING. Note also that it's not possible
@@ -103,11 +103,11 @@
 	m_dut->boot();
 
 	/* Expect to be able to transition to STAGING */
-	status = m_fwu_client->begin_staging();
+	status = m_fwu_client->begin_staging(0, 0, NULL);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 	/* And re-enter STAGING (implicit cancel) */
-	status = m_fwu_client->begin_staging();
+	status = m_fwu_client->begin_staging(0, 0, NULL);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 	/* Opening a couple of streams for installing images associated
@@ -117,11 +117,11 @@
 	uint32_t stream_handle2 = 0;
 
 	m_dut->whole_volume_image_type_uuid(0, &uuid);
-	status = m_fwu_client->open(&uuid, &stream_handle1);
+	status = m_fwu_client->open(&uuid, fwu_client::op_type::WRITE, &stream_handle1);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 	m_dut->whole_volume_image_type_uuid(1, &uuid);
-	status = m_fwu_client->open(&uuid, &stream_handle2);
+	status = m_fwu_client->open(&uuid, fwu_client::op_type::WRITE, &stream_handle2);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 	/* Attempting to end staging with open install streams should fail */
@@ -184,14 +184,14 @@
 	m_dut->boot();
 
 	/* Expect to be able to transition to STAGING */
-	status = m_fwu_client->begin_staging();
+	status = m_fwu_client->begin_staging(0, 0, NULL);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 	/* Install an image into location 0 */
 	uint32_t stream_handle = 0;
 
 	m_dut->whole_volume_image_type_uuid(0, &uuid);
-	status = m_fwu_client->open(&uuid, &stream_handle);
+	status = m_fwu_client->open(&uuid, fwu_client::op_type::WRITE, &stream_handle);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 	std::vector<uint8_t> image_data;
@@ -214,7 +214,7 @@
 	/* If we've transitioned to TRAIL, attempting to begin staging
 	 * again should be denied.
 	 */
-	status = m_fwu_client->begin_staging();
+	status = m_fwu_client->begin_staging(0, 0, NULL);
 	LONGS_EQUAL(FWU_STATUS_DENIED, status);
 
 	/* Activate the update. We'd expect the update to have been installed
@@ -226,7 +226,7 @@
 	/* If all's well, the DUT should have rebooted to TRIAL. Confirm this by
 	 * trying to begin staging - this should be denied.
 	 */
-	status = m_fwu_client->begin_staging();
+	status = m_fwu_client->begin_staging(0, 0, NULL);
 	LONGS_EQUAL(FWU_STATUS_DENIED, status);
 
 	/* All other staging related operations should also be denied */
@@ -238,7 +238,7 @@
 
 	/* Attempting to install images should also be denied */
 	m_dut->whole_volume_image_type_uuid(0, &uuid);
-	status = m_fwu_client->open(&uuid, &stream_handle);
+	status = m_fwu_client->open(&uuid, fwu_client::op_type::WRITE, &stream_handle);
 	LONGS_EQUAL(FWU_STATUS_DENIED, status);
 
 	/* Reading the image directory should be ok though */
@@ -250,10 +250,10 @@
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 	/* Should have transitioned back to REGULAR */
-	status = m_fwu_client->begin_staging();
+	status = m_fwu_client->begin_staging(0, 0, NULL);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
-	status = m_fwu_client->begin_staging();
+	status = m_fwu_client->begin_staging(0, 0, NULL);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 	status = m_fwu_client->cancel_staging();
diff --git a/components/service/fwu/test/ref_scenarios/oversize_image_tests.cpp b/components/service/fwu/test/ref_scenarios/oversize_image_tests.cpp
index c8d91a1..65b7ebb 100644
--- a/components/service/fwu/test/ref_scenarios/oversize_image_tests.cpp
+++ b/components/service/fwu/test/ref_scenarios/oversize_image_tests.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2023-2024, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -70,10 +70,10 @@
 	m_dut->generate_image_data(&image_data, img_size);
 
 	/* Install the image */
-	status = m_fwu_client->begin_staging();
+	status = m_fwu_client->begin_staging(0, 0, NULL);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
-	status = m_fwu_client->open(&uuid, &stream_handle);
+	status = m_fwu_client->open(&uuid, fwu_client::op_type::WRITE, &stream_handle);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 	status = m_fwu_client->write_stream(stream_handle,
@@ -114,10 +114,10 @@
 	m_dut->generate_image_data(&image_data, img_size);
 
 	/* Install the image */
-	status = m_fwu_client->begin_staging();
+	status = m_fwu_client->begin_staging(0, 0, NULL);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
-	status = m_fwu_client->open(&uuid, &stream_handle);
+	status = m_fwu_client->open(&uuid, fwu_client::op_type::WRITE, &stream_handle);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 	status = m_fwu_client->write_stream(stream_handle,
@@ -156,10 +156,10 @@
 	m_dut->generate_image_data(&image_data, img_size);
 
 	/* Install the image */
-	status = m_fwu_client->begin_staging();
+	status = m_fwu_client->begin_staging(0, 0, NULL);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
-	status = m_fwu_client->open(&uuid, &stream_handle);
+	status = m_fwu_client->open(&uuid, fwu_client::op_type::WRITE, &stream_handle);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 	status = m_fwu_client->write_stream(stream_handle,
@@ -201,7 +201,7 @@
 	m_dut->boot();
 
 	/* Perform multi-image update transaction */
-	status = m_fwu_client->begin_staging();
+	status = m_fwu_client->begin_staging(0, 0, NULL);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 	/* Install good image for location 0 */
@@ -210,7 +210,7 @@
 	size_t img_size = max_image_size(&uuid);
 	m_dut->generate_image_data(&image_data, img_size);
 
-	status = m_fwu_client->open(&uuid, &stream_handle);
+	status = m_fwu_client->open(&uuid, fwu_client::op_type::WRITE, &stream_handle);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 	status = m_fwu_client->write_stream(stream_handle,
@@ -226,7 +226,7 @@
 	img_size = max_image_size(&uuid) + 1;
 	m_dut->generate_image_data(&image_data, img_size);
 
-	status = m_fwu_client->open(&uuid, &stream_handle);
+	status = m_fwu_client->open(&uuid, fwu_client::op_type::WRITE, &stream_handle);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 	status = m_fwu_client->write_stream(stream_handle,
@@ -242,7 +242,7 @@
 	img_size = max_image_size(&uuid);
 	m_dut->generate_image_data(&image_data, img_size);
 
-	status = m_fwu_client->open(&uuid, &stream_handle);
+	status = m_fwu_client->open(&uuid, fwu_client::op_type::WRITE, &stream_handle);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 	status = m_fwu_client->write_stream(stream_handle,
diff --git a/components/service/fwu/test/ref_scenarios/power_failure_tests.cpp b/components/service/fwu/test/ref_scenarios/power_failure_tests.cpp
index a009e8c..47fb812 100644
--- a/components/service/fwu/test/ref_scenarios/power_failure_tests.cpp
+++ b/components/service/fwu/test/ref_scenarios/power_failure_tests.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2022-2024, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -62,7 +62,7 @@
 	m_metadata_checker->check_regular(boot_info.boot_index);
 
 	/* Begin staging */
-	status = m_fwu_client->begin_staging();
+	status = m_fwu_client->begin_staging(0, 0, NULL);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 	m_metadata_checker->check_ready_for_staging(boot_info.boot_index);
 
@@ -75,13 +75,13 @@
 	m_metadata_checker->check_regular(boot_info.boot_index);
 
 	/* Begin staging again */
-	status = m_fwu_client->begin_staging();
+	status = m_fwu_client->begin_staging(0, 0, NULL);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 	m_metadata_checker->check_ready_for_staging(boot_info.boot_index);
 
 	/* Start installing an image but don't commit it */
 	m_dut->whole_volume_image_type_uuid(0, &uuid);
-	status = m_fwu_client->open(&uuid, &stream_handle);
+	status = m_fwu_client->open(&uuid, fwu_client::op_type::WRITE, &stream_handle);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 	m_dut->generate_image_data(&image_data);
@@ -100,13 +100,13 @@
 	m_metadata_checker->check_regular(boot_info.boot_index);
 
 	/* Begin staging again */
-	status = m_fwu_client->begin_staging();
+	status = m_fwu_client->begin_staging(0, 0, NULL);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 	m_metadata_checker->check_ready_for_staging(boot_info.boot_index);
 
 	/* Start installing an image but this time commit it without ending staging */
 	m_dut->whole_volume_image_type_uuid(1, &uuid);
-	status = m_fwu_client->open(&uuid, &stream_handle);
+	status = m_fwu_client->open(&uuid, fwu_client::op_type::WRITE, &stream_handle);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 	m_dut->generate_image_data(&image_data);
@@ -147,13 +147,13 @@
 	m_metadata_checker->check_regular(boot_info.boot_index);
 
 	/* Begin staging */
-	status = m_fwu_client->begin_staging();
+	status = m_fwu_client->begin_staging(0, 0, NULL);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 	m_metadata_checker->check_ready_for_staging(boot_info.boot_index);
 
 	/* Install a partial update */
 	m_dut->whole_volume_image_type_uuid(2, &uuid);
-	status = m_fwu_client->open(&uuid, &stream_handle);
+	status = m_fwu_client->open(&uuid, fwu_client::op_type::WRITE, &stream_handle);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 	m_dut->generate_image_data(&image_data);
diff --git a/components/service/fwu/test/ref_scenarios/rollback_tests.cpp b/components/service/fwu/test/ref_scenarios/rollback_tests.cpp
index db77a3f..50005b4 100644
--- a/components/service/fwu/test/ref_scenarios/rollback_tests.cpp
+++ b/components/service/fwu/test/ref_scenarios/rollback_tests.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2022-2024, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -67,12 +67,12 @@
 	unsigned int pre_update_bank_index = boot_info.boot_index;
 
 	/* Install the update */
-	status = m_fwu_client->begin_staging();
+	status = m_fwu_client->begin_staging(0, 0, NULL);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 	m_metadata_checker->check_ready_for_staging(boot_info.boot_index);
 
 	m_dut->whole_volume_image_type_uuid(0, &uuid);
-	status = m_fwu_client->open(&uuid, &stream_handle);
+	status = m_fwu_client->open(&uuid, fwu_client::op_type::WRITE, &stream_handle);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 	std::vector<uint8_t> image_data;
@@ -132,12 +132,12 @@
 	unsigned int pre_update_bank_index = boot_info.boot_index;
 
 	/* Install the update */
-	status = m_fwu_client->begin_staging();
+	status = m_fwu_client->begin_staging(0, 0, NULL);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 	m_metadata_checker->check_ready_for_staging(boot_info.boot_index);
 
 	m_dut->whole_volume_image_type_uuid(0, &uuid);
-	status = m_fwu_client->open(&uuid, &stream_handle);
+	status = m_fwu_client->open(&uuid, fwu_client::op_type::WRITE, &stream_handle);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 	std::vector<uint8_t> image_data;
@@ -204,12 +204,12 @@
 	unsigned int pre_update_bank_index = boot_info.boot_index;
 
 	/* Install the update */
-	status = m_fwu_client->begin_staging();
+	status = m_fwu_client->begin_staging(0, 0, NULL);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 	m_metadata_checker->check_ready_for_staging(boot_info.boot_index);
 
 	m_dut->whole_volume_image_type_uuid(0, &uuid);
-	status = m_fwu_client->open(&uuid, &stream_handle);
+	status = m_fwu_client->open(&uuid, fwu_client::op_type::WRITE, &stream_handle);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 	std::vector<uint8_t> image_data;
diff --git a/components/service/fwu/test/ref_scenarios/update_fmp_tests.cpp b/components/service/fwu/test/ref_scenarios/update_fmp_tests.cpp
index 310899a..0a37363 100644
--- a/components/service/fwu/test/ref_scenarios/update_fmp_tests.cpp
+++ b/components/service/fwu/test/ref_scenarios/update_fmp_tests.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2023-2024, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -92,12 +92,12 @@
 		struct uuid_octets *uuid = &(img_info[ImageIndex].ImageTypeId);
 
 		if (!is_staging) {
-			status = client->begin_staging();
+			status = client->begin_staging(0, 0, NULL);
 			LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 			is_staging = true;
 		}
 
-		status = client->open(uuid, &stream_handle);
+		status = client->open(uuid, fwu_client::op_type::WRITE, &stream_handle);
 		LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 		status = client->write_stream(stream_handle, static_cast<const uint8_t *>(Image),
@@ -127,7 +127,7 @@
 
 		uuid_guid_octets_from_canonical(&uuid, FWU_DIRECTORY_CANONICAL_UUID);
 
-		status = client->open(&uuid, &stream_handle);
+		status = client->open(&uuid, fwu_client::op_type::READ, &stream_handle);
 		LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 		// Determine the size of the FW directory without reading any info.
@@ -139,7 +139,7 @@
 		// to reset the read seek.
 		status = client->commit(stream_handle, false);
 		LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
-		status = client->open(&uuid, &stream_handle);
+		status = client->open(&uuid, fwu_client::op_type::READ, &stream_handle);
 		LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 		img_dir = (ts_fwu_image_directory *)new uint8_t[reported_total_len];
diff --git a/components/service/fwu/test/ref_scenarios/update_scenario_tests.cpp b/components/service/fwu/test/ref_scenarios/update_scenario_tests.cpp
index b5c8341..57bc900 100644
--- a/components/service/fwu/test/ref_scenarios/update_scenario_tests.cpp
+++ b/components/service/fwu/test/ref_scenarios/update_scenario_tests.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2022-2024, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -40,6 +40,39 @@
 	fwu_client *m_fwu_client;
 };
 
+TEST(FwuUpdateScenarioTests, discover)
+{
+	int16_t service_status = 0;
+	uint8_t version_major = 0;
+	uint8_t version_minor = 0;
+	uint16_t num_func = 0;
+	uint64_t max_payload_size = 0;
+	uint32_t flags = 0;
+	uint32_t vendor_specific_flags = 0;
+	uint8_t function_presence[10] = { 0 };
+	uint8_t expected_functions[10] = { 0, 16, 17, 18, 19, 20, 21, 22, 23, 24 };
+
+	m_dut = fwu_dut_factory::create(1, false);
+	m_fwu_client = m_dut->create_fwu_client();
+	m_metadata_checker = m_dut->create_metadata_checker();
+
+	m_dut->boot();
+
+	int res = m_fwu_client->discover(&service_status, &version_major, &version_minor, &num_func,
+					 &max_payload_size, &flags, &vendor_specific_flags,
+					 function_presence);
+	LONGS_EQUAL(0, res);
+
+	UNSIGNED_LONGS_EQUAL(0, service_status);
+	UNSIGNED_LONGS_EQUAL(1, version_major);
+	UNSIGNED_LONGS_EQUAL(0, version_minor);
+	UNSIGNED_LONGS_EQUAL(10, sizeof(expected_functions));
+	UNSIGNED_LONGS_EQUAL(0, max_payload_size);
+	UNSIGNED_LONGS_EQUAL(0, flags);
+	UNSIGNED_LONGS_EQUAL(0, vendor_specific_flags);
+	MEMCMP_EQUAL(expected_functions, function_presence, sizeof(expected_functions));
+}
+
 TEST(FwuUpdateScenarioTests, wholeFirmwareUpdateFlow)
 {
 	int status = 0;
@@ -67,12 +100,12 @@
 	unsigned int pre_update_bank_index = boot_info.boot_index;
 
 	/* Perform staging steps where a single image is installed */
-	status = m_fwu_client->begin_staging();
+	status = m_fwu_client->begin_staging(0, 0, NULL);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 	m_metadata_checker->check_ready_for_staging(boot_info.boot_index);
 
 	m_dut->whole_volume_image_type_uuid(0, &uuid);
-	status = m_fwu_client->open(&uuid, &stream_handle);
+	status = m_fwu_client->open(&uuid, fwu_client::op_type::WRITE, &stream_handle);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 	std::vector<uint8_t> image_data;
@@ -152,13 +185,13 @@
 
 	/* Perform staging steps where multiple images are installed */
 	std::vector<uint8_t> image_data;
-	status = m_fwu_client->begin_staging();
+	status = m_fwu_client->begin_staging(0, 0, NULL);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 	m_metadata_checker->check_ready_for_staging(boot_info.boot_index);
 
 	/* Install whole image for location 0 */
 	m_dut->whole_volume_image_type_uuid(0, &uuid);
-	status = m_fwu_client->open(&uuid, &stream_handle);
+	status = m_fwu_client->open(&uuid, fwu_client::op_type::WRITE, &stream_handle);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 	m_dut->generate_image_data(&image_data);
@@ -174,7 +207,7 @@
 
 	/* Install whole image for location 1 */
 	m_dut->whole_volume_image_type_uuid(1, &uuid);
-	status = m_fwu_client->open(&uuid, &stream_handle);
+	status = m_fwu_client->open(&uuid, fwu_client::op_type::WRITE, &stream_handle);
 	LONGS_EQUAL(FWU_STATUS_SUCCESS, status);
 
 	m_dut->generate_image_data(&image_data);
diff --git a/protocols/service/fwu/packed-c/fwu_proto.h b/protocols/service/fwu/packed-c/fwu_proto.h
index 6cd691a..2537444 100644
--- a/protocols/service/fwu/packed-c/fwu_proto.h
+++ b/protocols/service/fwu/packed-c/fwu_proto.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2022-2024, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -17,7 +17,7 @@
  * The major version number of the FWU-A access protocol. It will be incremented
  * on significant updates that may include breaking changes.
  */
-#define FWU_PROTOCOL_VERSION_MAJOR 2
+#define FWU_PROTOCOL_VERSION_MAJOR 1
 
 /**
  * The minor version number  It will be incremented in
@@ -32,11 +32,13 @@
 #define FWU_DIRECTORY_CANONICAL_UUID	"deee58d9-5147-4ad3-a290-77666e2341a5"
 #define FWU_METADATA_CANONICAL_UUID	"8a7a84a0-8387-40f6-ab41-a8b9a5a60d23"
 
+#define FWU_OPEN_OP_TYPE_READ  (0)
+#define FWU_OPEN_OP_TYPE_WRITE (1)
+
 /**
  * Image directory
  */
-#define FWU_READ_PERM  (1u << 1)
-#define FWU_WRITE_PERM (1u << 0)
+#define FWU_IMAGE_DIRECTORY_VERSION	(2)
 
 struct __attribute__((__packed__)) ts_fwu_image_info_entry {
 	uint8_t img_type_uuid[OSF_UUID_OCTET_LEN];
@@ -63,14 +65,27 @@
  */
 
 struct __attribute__((__packed__)) ts_fwu_discover_out {
+	int16_t service_status;
 	uint8_t version_major;
 	uint8_t version_minor;
+	uint16_t off_function_presence;
 	uint16_t num_func;
+	uint64_t max_payload_size;
+	uint32_t flags;
+	uint32_t vendor_specific_flags;
 	uint8_t function_presence[];
 };
 
+struct __attribute__((__packed__)) ts_fwu_begin_staging_in {
+	uint32_t reserved;
+	uint32_t vendor_flags;
+	uint32_t partial_update_count;
+	uint8_t update_guid[];
+};
+
 struct __attribute__((__packed__)) ts_fwu_open_in {
 	uint8_t image_type_uuid[OSF_UUID_OCTET_LEN];
+	uint8_t op_type;
 };
 
 struct __attribute__((__packed__)) ts_fwu_open_out {
diff --git a/protocols/service/fwu/packed-c/opcodes.h b/protocols/service/fwu/packed-c/opcodes.h
index 1905a65..d6bbf3d 100644
--- a/protocols/service/fwu/packed-c/opcodes.h
+++ b/protocols/service/fwu/packed-c/opcodes.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2022-2024, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -10,15 +10,17 @@
 /**
  * Service-level opcodes
  */
-#define TS_FWU_OPCODE_BASE	      (0x10)
-#define TS_FWU_OPCODE_BEGIN_STAGING   (TS_FWU_OPCODE_BASE + 0)
-#define TS_FWU_OPCODE_END_STAGING     (TS_FWU_OPCODE_BASE + 1)
-#define TS_FWU_OPCODE_CANCEL_STAGING  (TS_FWU_OPCODE_BASE + 2)
-#define TS_FWU_OPCODE_OPEN	      (TS_FWU_OPCODE_BASE + 3)
-#define TS_FWU_OPCODE_WRITE_STREAM    (TS_FWU_OPCODE_BASE + 4)
-#define TS_FWU_OPCODE_READ_STREAM     (TS_FWU_OPCODE_BASE + 5)
-#define TS_FWU_OPCODE_COMMIT	      (TS_FWU_OPCODE_BASE + 6)
-#define TS_FWU_OPCODE_ACCEPT_IMAGE    (TS_FWU_OPCODE_BASE + 7)
-#define TS_FWU_OPCODE_SELECT_PREVIOUS (TS_FWU_OPCODE_BASE + 8)
+#define FWU_FUNC_ID_DISCOVER		(0)
+#define FWU_FUNC_ID_BEGIN_STAGING	(16)
+#define FWU_FUNC_ID_END_STAGING	(17)
+#define FWU_FUNC_ID_CANCEL_STAGING	(18)
+#define FWU_FUNC_ID_OPEN		(19)
+#define FWU_FUNC_ID_WRITE_STREAM	(20)
+#define FWU_FUNC_ID_READ_STREAM	(21)
+#define FWU_FUNC_ID_COMMIT		(22)
+#define FWU_FUNC_ID_ACCEPT_IMAGE	(23)
+#define FWU_FUNC_ID_SELECT_PREVIOUS	(24)
+
+#define FWU_FUNC_ID_COUNT		(10)
 
 #endif /* FWU_PROTO_OPCODES_H */