feat(rmm): unify design of RSI/PSCI handlers

This patch modifies RSI handlers to use consistent design
pattern. It eliminates the split of responsibility between
the 'handle_realm_rsi' dispatcher function and ther handler
function varied between commands, e.g.:

- Among commands which only have a single return value (in X0),
  some set X1-X3 to zero (e.g. RSI_REALM_CONFIG), while
  others leave caller values in X1-X3 (e.g. RSI_HOST_CALL).

- Among commands which cause a REC exit to the Host, the exit
  reason is in some cases set by the dispatcher (e.g.
  RSI_ATTESTATION_TOKEN_CONTINUE exiting due to a pending IRQ)
  and in others (e.g. RSI_IPA_STATE_SET) is set by the handler.

- In some cases all of the command logic is contained within
  the handler (e.g. RSI_VERSION, RSI_ATTESTATION_TOKEN_INIT,
  RSI_MEASUREMENT_READ, RSI_MEASUREMENT_EXTEND), while in
  others a significant part of the command logic is implemented
  by the handler (in particular RSI_ATTESTATION_TOKEN_CONTINUE
  and RSI_HOST_CALL).

These discrepancies made it difficult to follow the program
flow, and added complexity for implementation of new commands
or modification of the existing ones.

This commit introduces the following design pattern:

- All RSI and PSCI commands handlers have "void" return value
 and pointer to 'struct rsi_result' type parameter. It contains
 an "action" enumeration, which informs the dispatcher to do
 one of the following:
 - Return to the Realm, passing GPRs provided by the handler.
 - Exit to the Host, with 'rec_exit' fields (including
   exit_reason) populated by the handler.
 - Exit to the Host due to a Stage 2 translation fault.

- When GPRs are returned to the Realm, X0-X3 are always updated
  (with non-result registers being zero).

- The dispatcher contains only a switch over FID, and the logic
  necessary to perform the actions listed above. All other
  logic for a given command is contained within the handler itself.

This commit includes no changes to functionality.

Signed-off-by: AlexeiFedorov <Alexei.Fedorov@arm.com>
Change-Id: Iba80b679ca5aa9d83ba6c876e9d48b7b0b0d2579
diff --git a/runtime/core/exit.c b/runtime/core/exit.c
index 1bc5c22..ab78add 100644
--- a/runtime/core/exit.c
+++ b/runtime/core/exit.c
@@ -296,8 +296,7 @@
  * Handle FPU or SVE exceptions.
  * Returns: true if the exception is handled.
  */
-static bool
-handle_simd_exception(simd_t exp_type, struct rec *rec)
+static bool handle_simd_exception(simd_t exp_type, struct rec *rec)
 {
 	/*
 	 * If the REC wants to use SVE and if SVE is not enabled for this REC
@@ -326,19 +325,6 @@
 	return true;
 }
 
-/*
- * Return 'false' if no IRQ is pending,
- * return 'true' if there is an IRQ pending, and need to return to host.
- */
-static bool check_pending_irq(void)
-{
-	unsigned long pending_irq;
-
-	pending_irq = read_isr_el1();
-
-	return (pending_irq != 0UL);
-}
-
 static void advance_pc(void)
 {
 	unsigned long pc = read_elr_el2();
@@ -346,19 +332,11 @@
 	write_elr_el2(pc + 4UL);
 }
 
-static void return_result_to_realm(struct rec *rec, struct smc_result result)
-{
-	rec->regs[0] = result.x[0];
-	rec->regs[1] = result.x[1];
-	rec->regs[2] = result.x[2];
-	rec->regs[3] = result.x[3];
-}
-
 static inline bool rsi_handler_needs_fpu(unsigned int id)
 {
 #ifdef RMM_FPU_USE_AT_REL2
-	if (id == SMC_RSI_ATTEST_TOKEN_CONTINUE ||
-	    id == SMC_RSI_MEASUREMENT_EXTEND) {
+	if ((id == SMC_RSI_ATTEST_TOKEN_CONTINUE) ||
+	    (id == SMC_RSI_MEASUREMENT_EXTEND)) {
 		return true;
 	}
 #endif
@@ -371,16 +349,18 @@
  */
 static bool handle_realm_rsi(struct rec *rec, struct rmi_rec_exit *rec_exit)
 {
-	bool ret_to_rec = true;	/* Return to Realm */
+	struct rsi_result res = { 0 };
 	unsigned int function_id = (unsigned int)rec->regs[0];
 	bool restore_rec_simd_state = false;
+	bool needs_fpu, ret_to_rec;
 
 	RSI_LOG_SET(rec->regs);
 
 	/* Ignore SVE hint bit, until RMM supports SVE hint bit */
 	function_id &= ~MASK(SMC_SVE_HINT);
 
-	if (rsi_handler_needs_fpu(function_id) == true) {
+	needs_fpu = rsi_handler_needs_fpu(function_id);
+	if (needs_fpu) {
 		/*
 		 * RSI handler uses FPU at REL2, so actively save REC SIMD state
 		 * if REC is using SIMD or NS SIMD state. Restore the same before
@@ -403,172 +383,47 @@
 
 	switch (function_id) {
 	case SMCCC_VERSION:
-		rec->regs[0] = SMCCC_VERSION_NUMBER;
+		res.action = UPDATE_REC_RETURN_TO_REALM;
+		res.smc_res.x[0] = SMCCC_VERSION_NUMBER;
 		break;
 	case SMC_RSI_ABI_VERSION:
-		rec->regs[0] = handle_rsi_version();
+		handle_rsi_version(&res);
 		break;
 	case SMC32_PSCI_FID_MIN ... SMC32_PSCI_FID_MAX:
-	case SMC64_PSCI_FID_MIN ... SMC64_PSCI_FID_MAX: {
-		struct psci_result res;
-
-		res = psci_rsi(rec,
-			       function_id,
-			       rec->regs[1],
-			       rec->regs[2],
-			       rec->regs[3]);
-
-		if (!rec->psci_info.pending) {
-			rec->regs[0] = res.smc_res.x[0];
-			rec->regs[1] = res.smc_res.x[1];
-			rec->regs[2] = res.smc_res.x[2];
-			rec->regs[3] = res.smc_res.x[3];
-		}
-
-		if (res.hvc_forward.forward_psci_call) {
-			unsigned int i;
-
-			rec_exit->exit_reason = RMI_EXIT_PSCI;
-			rec_exit->gprs[0] = function_id;
-			rec_exit->gprs[1] = res.hvc_forward.x1;
-			rec_exit->gprs[2] = res.hvc_forward.x2;
-			rec_exit->gprs[3] = res.hvc_forward.x3;
-
-			for (i = 4U; i < REC_EXIT_NR_GPRS; i++) {
-				rec_exit->gprs[i] = 0UL;
-			}
-
-			advance_pc();
-			ret_to_rec = false;
-		}
+	case SMC64_PSCI_FID_MIN ... SMC64_PSCI_FID_MAX:
+		handle_psci(rec, rec_exit, &res);
 		break;
-	}
 	case SMC_RSI_ATTEST_TOKEN_INIT:
-		rec->regs[0] = handle_rsi_attest_token_init(rec);
+		handle_rsi_attest_token_init(rec, &res);
 		break;
-	case SMC_RSI_ATTEST_TOKEN_CONTINUE: {
-		struct attest_result res;
-		while (true) {
-			/*
-			 * Possible outcomes:
-			 *     if res.incomplete is true
-			 *         if IRQ pending
-			 *             check for pending IRQ and return to host
-			 *         else try a new iteration
-			 *     else
-			 *         if RTT table walk has failed,
-			 *             emulate data abort back to host
-			 *         otherwise
-			 *             return to realm because the token
-			 *             creation is complete or input parameter
-			 *             validation failed.
-			 */
-			handle_rsi_attest_token_continue(rec, &res);
-
-			if (res.incomplete) {
-				if (check_pending_irq()) {
-					rec_exit->exit_reason = RMI_EXIT_IRQ;
-
-					/* Copy the result to rec prior to return to host */
-					return_result_to_realm(rec, res.smc_res);
-					advance_pc();
-
-					/* Return to NS host to handle IRQ. */
-					ret_to_rec = false;
-					break;
-				}
-			} else {
-				if (res.walk_result.abort) {
-					emulate_stage2_data_abort(
-						rec, rec_exit,
-						res.walk_result.rtt_level);
-					ret_to_rec = false; /* Exit to Host */
-					break;
-				}
-
-				/* Return to Realm */
-				return_result_to_realm(rec, res.smc_res);
-				break;
-			}
-		}
+	case SMC_RSI_ATTEST_TOKEN_CONTINUE:
+		handle_rsi_attest_token_continue(rec, rec_exit, &res);
 		break;
-	}
 	case SMC_RSI_MEASUREMENT_READ:
-		rec->regs[0] = handle_rsi_read_measurement(rec);
+		handle_rsi_measurement_read(rec, &res);
 		break;
 	case SMC_RSI_MEASUREMENT_EXTEND:
-		rec->regs[0] = handle_rsi_extend_measurement(rec);
+		handle_rsi_measurement_extend(rec, &res);
 		break;
-	case SMC_RSI_REALM_CONFIG: {
-		struct rsi_result res;
-
-		res = handle_rsi_realm_config(rec);
-		if (res.walk_result.abort) {
-			emulate_stage2_data_abort(rec, rec_exit,
-						  res.walk_result.rtt_level);
-			ret_to_rec = false; /* Exit to Host */
-		} else {
-			/* Return to Realm */
-			return_result_to_realm(rec, res.smc_res);
-		}
+	case SMC_RSI_REALM_CONFIG:
+		handle_rsi_realm_config(rec, &res);
 		break;
-	}
 	case SMC_RSI_IPA_STATE_SET:
-		if (handle_rsi_ipa_state_set(rec, rec_exit)) {
-			rec->regs[0] = RSI_ERROR_INPUT;
-		} else {
-			advance_pc();
-			ret_to_rec = false; /* Return to Host */
-		}
+		handle_rsi_ipa_state_set(rec, rec_exit, &res);
 		break;
-	case SMC_RSI_IPA_STATE_GET: {
-		struct rsi_walk_smc_result res;
-
-		res = handle_rsi_ipa_state_get(rec);
-		if (res.walk_result.abort) {
-			emulate_stage2_data_abort(rec, rec_exit,
-						  res.walk_result.rtt_level);
-			/* Exit to Host */
-			ret_to_rec = false;
-		} else {
-			/* Exit to Realm */
-			return_result_to_realm(rec, res.smc_res);
-		}
+	case SMC_RSI_IPA_STATE_GET:
+		handle_rsi_ipa_state_get(rec, &res);
 		break;
-	}
-	case SMC_RSI_HOST_CALL: {
-		struct rsi_result res;
-
-		res = handle_rsi_host_call(rec, rec_exit);
-		if (res.walk_result.abort) {
-			emulate_stage2_data_abort(rec, rec_exit,
-						  res.walk_result.rtt_level);
-			/* Exit to Host */
-			ret_to_rec = false;
-		} else {
-			rec->regs[0] = res.smc_res.x[0];
-
-			/*
-			 * Return to Realm in case of error,
-			 * parent function calls advance_pc()
-			 */
-			if (rec->regs[0] == RSI_SUCCESS) {
-				advance_pc();
-
-				/* Exit to Host */
-				rec->host_call = true;
-				rec_exit->exit_reason = RMI_EXIT_HOST_CALL;
-				ret_to_rec = false;
-			}
-		}
+	case SMC_RSI_HOST_CALL:
+		handle_rsi_host_call(rec, rec_exit, &res);
 		break;
-	}
 	default:
-		rec->regs[0] = SMC_UNKNOWN;
+		res.action = UPDATE_REC_RETURN_TO_REALM;
+		res.smc_res.x[0] = SMC_UNKNOWN;
 		break;
 	}
 
-	if (rsi_handler_needs_fpu(function_id) == true) {
+	if (needs_fpu) {
 		if (restore_rec_simd_state == true) {
 			rec_simd_enable_restore(rec);
 		} else {
@@ -578,6 +433,20 @@
 		simd_enable(rec_simd_type(rec));
 	}
 
+	if ((res.action & FLAG_UPDATE_REC) != 0) {
+		for (unsigned int i = 0U; i < SMC_RESULT_REGS; ++i) {
+			rec->regs[i] = res.smc_res.x[i];
+		}
+	}
+
+	if ((res.action & FLAG_STAGE_2_ABORT) != 0) {
+		emulate_stage2_data_abort(rec, rec_exit, res.rtt_level);
+	} else {
+		advance_pc();
+	}
+
+	ret_to_rec = ((res.action & FLAG_EXIT_TO_HOST) == 0);
+
 	/* Log RSI call */
 	RSI_LOG_EXIT(function_id, rec->regs, ret_to_rec);
 	return ret_to_rec;
@@ -600,18 +469,7 @@
 		realm_inject_undef_abort();
 		return true;
 	case ESR_EL2_EC_SMC:
-		if (!handle_realm_rsi(rec, rec_exit)) {
-			return false;
-		}
-		/*
-		 * Advance PC.
-		 * HCR_EL2.TSC traps execution of the SMC instruction.
-		 * It is not a routing control for the SMC exception.
-		 * Trap exceptions and SMC exceptions have different
-		 * preferred return addresses.
-		 */
-		advance_pc();
-		return true;
+		return handle_realm_rsi(rec, rec_exit);
 	case ESR_EL2_EC_SYSREG: {
 		bool ret = handle_sysreg_access_trap(rec, rec_exit, esr);
 
@@ -654,7 +512,7 @@
 {
 	const unsigned long esr = read_esr_el2();
 
-	if (esr & ESR_EL2_SERROR_IDS_BIT) {
+	if ((esr & ESR_EL2_SERROR_IDS_BIT) != 0UL) {
 		/*
 		 * Implementation defined content of the esr.
 		 */
@@ -765,7 +623,7 @@
 	default:
 		INFO("Unrecognized exit reason: %d\n", exception);
 		break;
-	};
+	}
 
 	return false;
 }
diff --git a/runtime/include/psci.h b/runtime/include/psci.h
index e11b5e8..b3b33da 100644
--- a/runtime/include/psci.h
+++ b/runtime/include/psci.h
@@ -97,22 +97,6 @@
 
 struct rec;
 
-struct psci_result {
-	struct {
-		bool forward_psci_call;
-		unsigned long x1;
-		unsigned long x2;
-		unsigned long x3;
-	} hvc_forward;
-	struct smc_result smc_res;
-};
-
-struct psci_result psci_rsi(struct rec *rec,
-			    unsigned int function_id,
-			    unsigned long arg0,
-			    unsigned long arg1,
-			    unsigned long arg2);
-
 unsigned long psci_complete_request(struct rec *calling_rec,
 				    struct rec *target_rec);
 
diff --git a/runtime/include/rsi-handler.h b/runtime/include/rsi-handler.h
index cc604d1..5bf1787 100644
--- a/runtime/include/rsi-handler.h
+++ b/runtime/include/rsi-handler.h
@@ -9,54 +9,89 @@
 #include <rsi-walk.h>
 #include <smc.h>
 
+struct rec;
+struct rmi_rec_exit;
+
 /*
- * Result of an RSI command which performed an RTT walk.
+ * If set, update REC registers to values provided by the handler.
+ */
+#define FLAG_UPDATE_REC		1
+
+/*
+ * If set, exit to Host.  Otherwise, return to Realm.
+ */
+#define FLAG_EXIT_TO_HOST	2
+
+/*
+ * If set, present emulated Stage 2 abort to Host.
+ */
+#define FLAG_STAGE_2_ABORT	4
+
+enum rsi_action {
+	/*
+	 * Update REC registers to values provided by the handler,
+	 * and return to Realm.
+	 */
+	UPDATE_REC_RETURN_TO_REALM	= FLAG_UPDATE_REC,
+
+	/*
+	 * Leave REC registers unchanged, and exit to Host,
+	 * with rec_exit fields populated by the handler.
+	 */
+	EXIT_TO_HOST			= FLAG_EXIT_TO_HOST,
+
+	/*
+	 * Update REC registers to values provided by the handler,
+	 * with rec_exit fields and exit to Host, with rec_exit
+	 * fields populated by the handler.
+	 */
+	UPDATE_REC_EXIT_TO_HOST		= FLAG_UPDATE_REC |
+					  FLAG_EXIT_TO_HOST,
+
+	/*
+	 * Exit to Host, indicating a Stage 2 translation fault
+	 * encountered by the handler.
+	 */
+	STAGE_2_TRANSLATION_FAULT	= FLAG_EXIT_TO_HOST |
+					  FLAG_STAGE_2_ABORT
+};
+
+/*
+ * Result of RSI command handler
  */
 struct rsi_result {
 	/*
-	 * Result of RTT walk performed by RSI command.
+	 * Action which should be taken following execution of the handler.
 	 */
-	struct rsi_walk_result walk_result;
+	enum rsi_action action;
 
 	/*
-	 * If @walk_result.abort is false,
-	 * @smc_res contains GPR values to be returned to the Realm.
+	 * If the handler performed an RTT walk,
+	 * @rtt_level is the level at which the walk terminated.
 	 */
-	struct smc_result smc_res;
-};
-
-struct attest_result {
-	/*
-	 * If true, RMM should repeat the operation.
-	 *
-	 * If false, contents of @access are valid.
-	 */
-	bool incomplete;
+	unsigned long rtt_level;
 
 	/*
-	 * Result of RTT walk performed by RSI command.
-	 */
-	struct rsi_walk_result walk_result;
-
-	/*
-	 * If @incomplete is false and @walk_result.abort is false,
+	 * If @action is RETURN_TO_REALM,
 	 * @smc_result contains GPR values to be returned to the Realm.
 	 */
 	struct smc_result smc_res;
 };
 
-unsigned long handle_rsi_version(void);
-struct rsi_result handle_rsi_realm_config(struct rec *rec);
-struct rsi_result handle_rsi_host_call(struct rec *rec,
-				       struct rmi_rec_exit *rec_exit);
-bool handle_rsi_ipa_state_set(struct rec *rec, struct rmi_rec_exit *rec_exit);
-struct rsi_walk_smc_result handle_rsi_ipa_state_get(struct rec *rec);
-unsigned long handle_rsi_read_measurement(struct rec *rec);
-unsigned long handle_rsi_extend_measurement(struct rec *rec);
-unsigned long handle_rsi_attest_token_init(struct rec *rec);
-void attest_realm_token_sign_continue_start(void);
+void handle_rsi_version(struct rsi_result *res);
+void handle_rsi_realm_config(struct rec *rec, struct rsi_result *res);
+void handle_rsi_host_call(struct rec *rec, struct rmi_rec_exit *rec_exit,
+			  struct rsi_result *res);
+void handle_rsi_ipa_state_set(struct rec *rec, struct rmi_rec_exit *rec_exit,
+			      struct rsi_result *res);
+void handle_rsi_ipa_state_get(struct rec *rec, struct rsi_result *res);
+void handle_rsi_measurement_read(struct rec *rec, struct rsi_result *res);
+void handle_rsi_measurement_extend(struct rec *rec, struct rsi_result *res);
+void handle_rsi_attest_token_init(struct rec *rec, struct rsi_result *res);
 void handle_rsi_attest_token_continue(struct rec *rec,
-				      struct attest_result *res);
-void attest_realm_token_sign_continue_finish(void);
+				      struct rmi_rec_exit *rec_exit,
+				      struct rsi_result *res);
+void handle_psci(struct rec *rec, struct rmi_rec_exit *rec_exit,
+		 struct rsi_result *res);
 
 #endif /* RSI_HANDLER_H */
diff --git a/runtime/include/rsi-walk.h b/runtime/include/rsi-walk.h
index 2294e88..9f6c7cc 100644
--- a/runtime/include/rsi-walk.h
+++ b/runtime/include/rsi-walk.h
@@ -20,15 +20,4 @@
 	unsigned long rtt_level;
 };
 
-struct rsi_walk_smc_result {
-	/* Result of RTT walk performed by RSI command */
-	struct rsi_walk_result walk_result;
-
-	/*
-	 * If @walk_result.abort is false, @smc_res contains GPR values to be
-	 * returned to the Realm.
-	 */
-	struct smc_result smc_res;
-};
-
 #endif /* RSI_WALK_H */
diff --git a/runtime/rmi/system.c b/runtime/rmi/system.c
index 24e53ef..024d419 100644
--- a/runtime/rmi/system.c
+++ b/runtime/rmi/system.c
@@ -3,7 +3,6 @@
  * SPDX-FileCopyrightText: Copyright TF-RMM Contributors.
  */
 #include <assert.h>
-#include <debug.h>
 #include <smc-handler.h>
 #include <smc-rmi.h>
 
diff --git a/runtime/rsi/config.c b/runtime/rsi/config.c
index 8282d19..63ec51c 100644
--- a/runtime/rsi/config.c
+++ b/runtime/rsi/config.c
@@ -8,37 +8,36 @@
 #include <rsi-handler.h>
 #include <smc-rsi.h>
 
-struct rsi_result handle_rsi_realm_config(struct rec *rec)
+void handle_rsi_realm_config(struct rec *rec, struct rsi_result *res)
 {
-	struct rsi_result res = { 0 };
 	unsigned long ipa = rec->regs[1];
 	enum s2_walk_status walk_status;
 	struct s2_walk_result walk_res;
 	struct granule *gr;
 	struct rsi_realm_config *config;
 
+	res->action = UPDATE_REC_RETURN_TO_REALM;
+
 	if (!GRANULE_ALIGNED(ipa) || !addr_in_rec_par(rec, ipa)) {
-		res.smc_res.x[0] = RSI_ERROR_INPUT;
-		return res;
+		res->smc_res.x[0] = RSI_ERROR_INPUT;
+		return;
 	}
 
 	walk_status = realm_ipa_to_pa(rec, ipa, &walk_res);
 
 	if (walk_status == WALK_FAIL) {
 		if (walk_res.ripas_val == RIPAS_EMPTY) {
-			res.smc_res.x[0] = RSI_ERROR_INPUT;
+			res->smc_res.x[0] = RSI_ERROR_INPUT;
 		} else {
-			/* Exit to Host */
-			res.walk_result.abort = true;
-			res.walk_result.rtt_level = walk_res.rtt_level;
+			res->action = STAGE_2_TRANSLATION_FAULT;
+			res->rtt_level = walk_res.rtt_level;
 		}
-		return res;
+		return;
 	}
 
 	if (walk_status == WALK_INVALID_PARAMS) {
-		/* Return error to Realm */
-		res.smc_res.x[0] = RSI_ERROR_INPUT;
-		return res;
+		res->smc_res.x[0] = RSI_ERROR_INPUT;
+		return;
 	}
 
 	/* Map Realm data granule to RMM address space */
@@ -55,6 +54,5 @@
 	granule_unlock(walk_res.llt);
 
 	/* Write output values */
-	res.smc_res.x[0] = RSI_SUCCESS;
-	return res;
+	res->smc_res.x[0] = RSI_SUCCESS;
 }
diff --git a/runtime/rsi/host_call.c b/runtime/rsi/host_call.c
index 3a008ef..24ce759 100644
--- a/runtime/rsi/host_call.c
+++ b/runtime/rsi/host_call.c
@@ -26,10 +26,9 @@
  *     results to host call data structure (in Realm memory).
  *   - Return value is RSI_SUCCESS.
  */
-static unsigned int do_host_call(struct rec *rec,
-				 struct rmi_rec_exit *rec_exit,
-				 struct rmi_rec_entry *rec_entry,
-				 struct rsi_walk_result *rsi_walk_result)
+static void do_host_call(struct rec *rec, struct rmi_rec_exit *rec_exit,
+			 struct rmi_rec_entry *rec_entry,
+			 struct rsi_result *res)
 {
 	enum s2_walk_status walk_status;
 	struct s2_walk_result walk_result;
@@ -39,7 +38,6 @@
 	unsigned char *data;
 	struct rsi_host_call *host_call;
 	unsigned int i;
-	unsigned int ret = RSI_SUCCESS;
 
 	assert(addr_in_rec_par(rec, ipa));
 
@@ -54,12 +52,12 @@
 		break;
 	case WALK_FAIL:
 		if (walk_result.ripas_val == RIPAS_EMPTY) {
-			ret = RSI_ERROR_INPUT;
+			res->smc_res.x[0] = RSI_ERROR_INPUT;
 		} else {
-			rsi_walk_result->abort = true;
-			rsi_walk_result->rtt_level = walk_result.rtt_level;
+			res->action = STAGE_2_TRANSLATION_FAULT;
+			res->rtt_level = walk_result.rtt_level;
 		}
-		return ret;
+		return;
 	case WALK_INVALID_PARAMS:
 		assert(false);
 		break;
@@ -76,11 +74,24 @@
 		for (i = 0U; i < RSI_HOST_CALL_NR_GPRS; i++) {
 			rec_exit->gprs[i] = host_call->gprs[i];
 		}
+
+		/* Record that a host call is pending */
+		rec->host_call = true;
+
+		/*
+		 * Notify the Host.
+		 * Leave REC registers unchanged,
+		 * these will be read and updated by complete_rsi_host_call.
+		 */
+		res->action = EXIT_TO_HOST;
+		rec_exit->exit_reason = RMI_EXIT_HOST_CALL;
 	} else {
 		/* Copy host call results to host call data structure */
 		for (i = 0U; i < RSI_HOST_CALL_NR_GPRS; i++) {
 			host_call->gprs[i] = rec_entry->gprs[i];
 		}
+
+		rec->regs[0] = RSI_SUCCESS;
 	}
 
 	/* Unmap Realm data granule */
@@ -88,41 +99,39 @@
 
 	/* Unlock last level RTT */
 	granule_unlock(walk_result.llt);
-
-	return ret;
 }
 
-struct rsi_result handle_rsi_host_call(struct rec *rec,
-					struct rmi_rec_exit *rec_exit)
+void handle_rsi_host_call(struct rec *rec, struct rmi_rec_exit *rec_exit,
+			  struct rsi_result *res)
 {
-	struct rsi_result res = { 0 };
 	unsigned long ipa = rec->regs[1];
 
+	res->action = UPDATE_REC_RETURN_TO_REALM;
+
 	if (!ALIGNED(ipa, sizeof(struct rsi_host_call))) {
-		res.smc_res.x[0] = RSI_ERROR_INPUT;
-		return res;
+		res->smc_res.x[0] = RSI_ERROR_INPUT;
+		return;
 	}
 
 	if ((ipa / GRANULE_SIZE) !=
 		((ipa + sizeof(struct rsi_host_call) - 1UL) / GRANULE_SIZE)) {
-		res.smc_res.x[0] = RSI_ERROR_INPUT;
-		return res;
+		res->smc_res.x[0] = RSI_ERROR_INPUT;
+		return;
 	}
 
 	if (!addr_in_rec_par(rec, ipa)) {
-		res.smc_res.x[0] = RSI_ERROR_INPUT;
-		return res;
+		res->smc_res.x[0] = RSI_ERROR_INPUT;
+		return;
 	}
 
-	res.smc_res.x[0] = do_host_call(rec, rec_exit, NULL, &res.walk_result);
-
-	return res;
+	do_host_call(rec, rec_exit, NULL, res);
 }
 
 struct rsi_walk_result complete_rsi_host_call(struct rec *rec,
 					      struct rmi_rec_entry *rec_entry)
 {
-	struct rsi_walk_result res = { false, 0UL };
+	struct rsi_result res = { 0UL };
+	struct rsi_walk_result walk_res = { false, 0UL };
 
 	/*
 	 * Do the necessary to walk the S2 RTTs and copy args from NS Host
@@ -133,7 +142,12 @@
 	 * the RIPAS of IPA at that time). This is a situation which can be
 	 * controlled from Realm and Realm should avoid this.
 	 */
-	(void)do_host_call(rec, NULL, rec_entry, &res);
+	do_host_call(rec, NULL, rec_entry, &res);
 
-	return res;
+	if (res.action == STAGE_2_TRANSLATION_FAULT) {
+		walk_res.abort = true;
+		walk_res.rtt_level = res.rtt_level;
+	}
+
+	return walk_res;
 }
diff --git a/runtime/rsi/memory.c b/runtime/rsi/memory.c
index 2324fcf..1affda7 100644
--- a/runtime/rsi/memory.c
+++ b/runtime/rsi/memory.c
@@ -10,32 +10,22 @@
 #include <smc-rsi.h>
 #include <status.h>
 
-bool handle_rsi_ipa_state_set(struct rec *rec, struct rmi_rec_exit *rec_exit)
+void handle_rsi_ipa_state_set(struct rec *rec,
+			      struct rmi_rec_exit *rec_exit,
+			      struct rsi_result *res)
 {
 	unsigned long start = rec->regs[1];
 	unsigned long size = rec->regs[2];
 	unsigned long end = start + size;
 	enum ripas ripas_val = (enum ripas)rec->regs[3];
 
-	if (ripas_val > RIPAS_RAM) {
-		return true;
-	}
-
-	if (!GRANULE_ALIGNED(start)) {
-		return true;
-	}
-
-	if (!GRANULE_ALIGNED(size)) {
-		return true;
-	}
-
-	if (end <= start) {
-		/* Size is zero, or range overflows */
-		return true;
-	}
-
-	if (!region_in_rec_par(rec, start, end)) {
-		return true;
+	if ((ripas_val > RIPAS_RAM)	    ||
+	    !GRANULE_ALIGNED(start) || !GRANULE_ALIGNED(size) ||
+	    (end <= start)	    || /* Size is zero, or range overflows */
+	    !region_in_rec_par(rec, start, end)) {
+		res->action = UPDATE_REC_RETURN_TO_REALM;
+		res->smc_res.x[0] = RSI_ERROR_INPUT;
+		return;
 	}
 
 	rec->set_ripas.base = start;
@@ -48,35 +38,32 @@
 	rec_exit->ripas_size = size;
 	rec_exit->ripas_value = (unsigned int)ripas_val;
 
-	return false;
+	res->action = UPDATE_REC_EXIT_TO_HOST;
 }
 
-struct rsi_walk_smc_result handle_rsi_ipa_state_get(struct rec *rec)
+void handle_rsi_ipa_state_get(struct rec *rec,
+			      struct rsi_result *res)
 {
-	struct rsi_walk_smc_result res = { 0 };
+	unsigned long ipa = rec->regs[1];
+	unsigned long rtt_level;
 	enum s2_walk_status ws;
-	unsigned long rtt_level, ipa;
 	enum ripas ripas_val;
 
-	ipa = rec->regs[1];
-
-	/* Exit to realm */
-	res.walk_result.abort = false;
+	res->action = UPDATE_REC_RETURN_TO_REALM;
 
 	if (!GRANULE_ALIGNED(ipa) || !addr_in_rec_par(rec, ipa)) {
-		res.smc_res.x[0] = RSI_ERROR_INPUT;
-		return res;
+		res->smc_res.x[0] = RSI_ERROR_INPUT;
+		return;
 	}
 
 	ws = realm_ipa_get_ripas(rec, ipa, &ripas_val, &rtt_level);
 	if (ws == WALK_SUCCESS) {
-		res.smc_res.x[0] = RSI_SUCCESS;
-		res.smc_res.x[1] = ripas_val;
+		res->smc_res.x[0] = RSI_SUCCESS;
+		res->smc_res.x[1] = ripas_val;
 	} else {
 		/* Exit to Host */
-		res.walk_result.abort = true;
-		res.walk_result.rtt_level = rtt_level;
+		res->action = STAGE_2_TRANSLATION_FAULT;
+		res->rtt_level = rtt_level;
+		res->smc_res.x[0] = RSI_ERROR_INPUT;
 	}
-
-	return res;
 }
diff --git a/runtime/rsi/psci.c b/runtime/rsi/psci.c
index 4a62685..4fcae13 100644
--- a/runtime/rsi/psci.c
+++ b/runtime/rsi/psci.c
@@ -7,41 +7,64 @@
 #include <psci.h>
 #include <realm.h>
 #include <rec.h>
+#include <rsi-handler.h>
 #include <smc-rmi.h>
 #include <smc.h>
 #include <stdint.h>
 
-static struct psci_result psci_version(struct rec *rec)
+/*
+ * Copy @count GPRs from @rec to @rec_exit.
+ * The remaining @rec_exit.gprs[] values are zero filled.
+ */
+static void forward_args_to_host(unsigned int count, struct rec *rec,
+				 struct rmi_rec_exit *rec_exit)
 {
-	struct psci_result result = { 0 };
-	unsigned int version_1_1 = (1U << 16) | 1U;
+	unsigned int i;
 
-	result.smc_res.x[0] = (unsigned long)version_1_1;
-	return result;
+	assert(count <= 4U);
+
+	for (i = 0U; i < count; ++i) {
+		rec_exit->gprs[i] = rec->regs[i];
+	}
+
+	for (i = count; i < REC_EXIT_NR_GPRS; ++i) {
+		rec_exit->gprs[i] = 0UL;
+	}
 }
 
-static struct psci_result psci_cpu_suspend(struct rec *rec,
-					  unsigned long entry_point_address,
-					  unsigned long context_id)
+static void psci_version(struct rsi_result *res)
 {
-	struct psci_result result = { 0 };
+	const unsigned long version_1_1 = (1UL << 16) | 1UL;
+
+	res->action = UPDATE_REC_RETURN_TO_REALM;
+	res->smc_res.x[0] = version_1_1;
+}
+
+static void psci_cpu_suspend(struct rec *rec, struct rmi_rec_exit *rec_exit,
+			     struct rsi_result *res)
+{
+	res->action = UPDATE_REC_EXIT_TO_HOST;
 
 	/*
-	 * We treat all target power states as suspend requests, so all we
-	 * need to do is inform that NS hypervisor and we can ignore all the
-	 * parameters.
+	 * We treat all target power states as suspend requests,
+	 * so all we need to do is forward the FID to the NS hypervisor,
+	 * and we can ignore all the parameters.
 	 */
-	result.hvc_forward.forward_psci_call = true;
+	forward_args_to_host(1U, rec, rec_exit);
 
-	result.smc_res.x[0] = PSCI_RETURN_SUCCESS;
-	return result;
+	/*
+	 * The exit to the Host is just a notification; the Host does not need
+	 * to complete a PSCI request before the next call to RMI_REC_ENTER.
+	 * We therefore update the REC immediately with the results of the PSCI
+	 * command.
+	 */
+	res->smc_res.x[0] = PSCI_RETURN_SUCCESS;
 }
 
-static struct psci_result psci_cpu_off(struct rec *rec)
+static void psci_cpu_off(struct rec *rec, struct rmi_rec_exit *rec_exit,
+			 struct rsi_result *res)
 {
-	struct psci_result result = { 0 };
-
-	result.hvc_forward.forward_psci_call = true;
+	res->action = UPDATE_REC_EXIT_TO_HOST;
 
 	/*
 	 * It should be fine to set this flag without holding a lock on the
@@ -55,8 +78,16 @@
 	 */
 	rec->runnable = false;
 
-	result.smc_res.x[0] = PSCI_RETURN_SUCCESS;
-	return result;
+	/* Notify the Host, passing the FID only. */
+	forward_args_to_host(1U, rec, rec_exit);
+
+	/*
+	 * The exit to the Host is just a notification; the Host does not need
+	 * to complete a PSCI request before the next call to RMI_REC_ENTER.
+	 * We therefore update the REC immediately with the results of the PSCI
+	 * command.
+	 */
+	res->smc_res.x[0] = PSCI_RETURN_SUCCESS;
 }
 
 static void psci_reset_rec(struct rec *rec, unsigned long caller_sctlr_el1)
@@ -86,18 +117,19 @@
 	return rec_count;
 }
 
-static struct psci_result psci_cpu_on(struct rec *rec,
-				      unsigned long target_cpu,
-				      unsigned long entry_point_address,
-				      unsigned long context_id)
+static void psci_cpu_on(struct rec *rec, struct rmi_rec_exit *rec_exit,
+			struct rsi_result *res)
 {
-	struct psci_result result = { 0 };
+	unsigned long target_cpu = rec->regs[1];
+	unsigned long entry_point_address = rec->regs[2];
 	unsigned long target_rec_idx;
 
+	res->action = UPDATE_REC_RETURN_TO_REALM;
+
 	/* Check that entry_point_address is a Protected Realm Address */
 	if (!addr_in_rec_par(rec, entry_point_address)) {
-		result.smc_res.x[0] = PSCI_RETURN_INVALID_ADDRESS;
-		return result;
+		res->smc_res.x[0] = PSCI_RETURN_INVALID_ADDRESS;
+		return;
 	}
 
 	/* Get REC index from MPIDR */
@@ -109,33 +141,40 @@
 	 * consecutively increasing indexes starting from zero.
 	 */
 	if (target_rec_idx >= rd_map_read_rec_count(rec->realm_info.g_rd)) {
-		result.smc_res.x[0] = PSCI_RETURN_INVALID_PARAMS;
-		return result;
+		res->smc_res.x[0] = PSCI_RETURN_INVALID_PARAMS;
+		return;
 	}
 
 	/* Check if we're trying to turn ourselves on */
 	if (target_rec_idx == rec->rec_idx) {
-		result.smc_res.x[0] = PSCI_RETURN_ALREADY_ON;
-		return result;
+		res->smc_res.x[0] = PSCI_RETURN_ALREADY_ON;
+		return;
 	}
 
+	/* Record that a PSCI request is outstanding */
 	rec->psci_info.pending = true;
 
-	result.hvc_forward.forward_psci_call = true;
-	result.hvc_forward.x1 = target_cpu;
-	return result;
+	/*
+	 * Notify the Host, passing the FID and MPIDR arguments.
+	 * Leave REC registers unchanged; these will be read and updated by
+	 * psci_complete_request.
+	 */
+	forward_args_to_host(2U, rec, rec_exit);
+	res->action = EXIT_TO_HOST;
 }
 
-static struct psci_result psci_affinity_info(struct rec *rec,
-					     unsigned long target_affinity,
-					     unsigned long lowest_affinity_level)
+static void psci_affinity_info(struct rec *rec, struct rmi_rec_exit *rec_exit,
+			       struct rsi_result *res)
 {
-	struct psci_result result = { 0 };
+	unsigned long target_affinity = rec->regs[1];
+	unsigned long lowest_affinity_level = rec->regs[2];
 	unsigned long target_rec_idx;
 
+	res->action = UPDATE_REC_RETURN_TO_REALM;
+
 	if (lowest_affinity_level != 0UL) {
-		result.smc_res.x[0] = PSCI_RETURN_INVALID_PARAMS;
-		return result;
+		res->smc_res.x[0] = PSCI_RETURN_INVALID_PARAMS;
+		return;
 	}
 
 	/* Get REC index from MPIDR */
@@ -147,21 +186,27 @@
 	 * consecutively increasing indexes starting from zero.
 	 */
 	if (target_rec_idx >= rd_map_read_rec_count(rec->realm_info.g_rd)) {
-		result.smc_res.x[0] = PSCI_RETURN_INVALID_PARAMS;
-		return result;
+		res->smc_res.x[0] = PSCI_RETURN_INVALID_PARAMS;
+		return;
 	}
 
 	/* Check if the vCPU targets itself */
 	if (target_rec_idx == rec->rec_idx) {
-		result.smc_res.x[0] = PSCI_AFFINITY_INFO_ON;
-		return result;
+		res->smc_res.x[0] = PSCI_AFFINITY_INFO_ON;
+		return;
 	}
 
+	/* Record that a PSCI request is outstanding */
 	rec->psci_info.pending = true;
 
-	result.hvc_forward.forward_psci_call = true;
-	result.hvc_forward.x1 = target_affinity;
-	return result;
+	/*
+	 * Notify the Host, passing the FID and MPIDR arguments.
+	 * Leave REC registers unchanged; these will be read and updated
+	 * by psci_complete_request.
+	 */
+	forward_args_to_host(2U, rec, rec_exit);
+
+	res->action = EXIT_TO_HOST;
 }
 
 /*
@@ -191,31 +236,21 @@
 	/* TODO: Invalidate all stage 2 entris to ensure REC exits */
 }
 
-static struct psci_result psci_system_off(struct rec *rec)
+static void psci_system_off_reset(struct rec *rec,
+				  struct rmi_rec_exit *rec_exit,
+				  struct rsi_result *res)
 {
-	struct psci_result result = { 0 };
-
 	system_off_reboot(rec);
 
-	result.hvc_forward.forward_psci_call = true;
-	return result;
+	/* Notify the Host, passing the FID only */
+	forward_args_to_host(1U, rec, rec_exit);
+
+	res->action = EXIT_TO_HOST;
 }
 
-static struct psci_result psci_system_reset(struct rec *rec)
+static void psci_features(struct rec *rec, struct rsi_result *res)
 {
-	struct psci_result result = { 0 };
-
-	system_off_reboot(rec);
-
-	result.hvc_forward.forward_psci_call = true;
-	return result;
-}
-
-static struct psci_result psci_features(struct rec *rec,
-				       unsigned int psci_func_id)
-{
-	struct psci_result result = { 0 };
-	unsigned long ret;
+	unsigned int psci_func_id = (unsigned int)rec->regs[1];
 
 	switch (psci_func_id) {
 	case SMC32_PSCI_CPU_SUSPEND:
@@ -229,66 +264,56 @@
 	case SMC32_PSCI_SYSTEM_RESET:
 	case SMC32_PSCI_FEATURES:
 	case SMCCC_VERSION:
-		ret = 0UL;
+		res->smc_res.x[0] = PSCI_RETURN_SUCCESS;
 		break;
 	default:
-		ret = PSCI_RETURN_NOT_SUPPORTED;
+		res->smc_res.x[0] = PSCI_RETURN_NOT_SUPPORTED;
 	}
 
-	result.smc_res.x[0] = ret;
-	return result;
+	res->action = UPDATE_REC_RETURN_TO_REALM;
 }
 
-struct psci_result psci_rsi(struct rec *rec,
-			    unsigned int function_id,
-			    unsigned long arg0,
-			    unsigned long arg1,
-			    unsigned long arg2)
+void handle_psci(struct rec *rec,
+		 struct rmi_rec_exit *rec_exit,
+		 struct rsi_result *res)
 {
-	struct psci_result result = { false };
+	unsigned int function_id = (unsigned int)rec->regs[0];
 
 	switch (function_id) {
 	case SMC32_PSCI_VERSION:
-		result = psci_version(rec);
+		psci_version(res);
 		break;
 	case SMC32_PSCI_CPU_SUSPEND:
 	case SMC64_PSCI_CPU_SUSPEND:
-		result = psci_cpu_suspend(rec, arg0, arg1);
+		psci_cpu_suspend(rec, rec_exit, res);
 		break;
 	case SMC32_PSCI_CPU_OFF:
-		result = psci_cpu_off(rec);
+		psci_cpu_off(rec, rec_exit, res);
 		break;
 	case SMC32_PSCI_CPU_ON:
-		arg0 = (unsigned int)arg0;
-		arg1 = (unsigned int)arg1;
-		arg2 = (unsigned int)arg2;
-		FALLTHROUGH;
 	case SMC64_PSCI_CPU_ON:
-		result = psci_cpu_on(rec, arg0, arg1, arg2);
+		psci_cpu_on(rec, rec_exit, res);
 		break;
 	case SMC32_PSCI_AFFINITY_INFO:
-		arg0 = (unsigned int)arg0;
-		arg1 = (unsigned int)arg1;
-		FALLTHROUGH;
 	case SMC64_PSCI_AFFINITY_INFO:
-		result = psci_affinity_info(rec, arg0, arg1);
+		psci_affinity_info(rec, rec_exit, res);
 		break;
 	case SMC32_PSCI_SYSTEM_OFF:
-		result = psci_system_off(rec);
-		break;
 	case SMC32_PSCI_SYSTEM_RESET:
-		result = psci_system_reset(rec);
+		psci_system_off_reset(rec, rec_exit, res);
 		break;
 	case SMC32_PSCI_FEATURES:
-		result = psci_features(rec, arg0);
+		psci_features(rec, res);
 		break;
 	default:
-		result.smc_res.x[0] = PSCI_RETURN_NOT_SUPPORTED;
-		result.hvc_forward.forward_psci_call = false;
+		res->action = UPDATE_REC_RETURN_TO_REALM;
+		res->smc_res.x[0] = PSCI_RETURN_NOT_SUPPORTED;
 		break;
 	}
 
-	return result;
+	if ((res->action & FLAG_EXIT_TO_HOST) != 0) {
+		rec_exit->exit_reason = RMI_EXIT_PSCI;
+	}
 }
 
 /*
diff --git a/runtime/rsi/realm_attest.c b/runtime/rsi/realm_attest.c
index 4fde391..c83888c 100644
--- a/runtime/rsi/realm_attest.c
+++ b/runtime/rsi/realm_attest.c
@@ -51,12 +51,10 @@
 }
 
 /*
- * Function to continue with the sign operation.
- * It returns void as the result will be updated in the
- * struct attest_result passed as argument.
+ * Function to continue with the sign operation
  */
 static void attest_token_continue_sign_state(struct rec *rec,
-					     struct attest_result *res)
+					     struct rsi_result *res)
 {
 	/*
 	 * Sign and finish creating the token.
@@ -72,7 +70,6 @@
 		 * to check is there anything else to do (pending IRQ)
 		 * or next signing iteration can be executed.
 		 */
-		res->incomplete = true;
 		res->smc_res.x[0] = RSI_INCOMPLETE;
 
 		/* If this was the last signing cycle */
@@ -88,12 +85,10 @@
 }
 
 /*
- * Function to continue with the token write operation.
- * It returns void as the result will be updated in the
- * struct attest_result passed as argument.
+ * Function to continue with the token write operation
  */
 static void attest_token_continue_write_state(struct rec *rec,
-					      struct attest_result *res)
+					      struct rsi_result *res)
 {
 	struct granule *gr;
 	uint8_t *realm_att_token;
@@ -117,12 +112,11 @@
 			res->smc_res.x[0] = RSI_ERROR_INPUT;
 		} else {
 			/*
-			 * Translation failed, IPA is not mapped. Return to NS host to
-			 * fix the issue.
+			 * Translation failed, IPA is not mapped.
+			 * Return to NS host to fix the issue.
 			 */
-			res->walk_result.abort = true;
-			res->walk_result.rtt_level = walk_res.rtt_level;
-			res->smc_res.x[0] = RSI_INCOMPLETE;
+			res->action = STAGE_2_TRANSLATION_FAULT;
+			res->rtt_level = walk_res.rtt_level;
 		}
 		return;
 	}
@@ -154,10 +148,9 @@
 	rec->token_sign_ctx.state = ATTEST_SIGN_NOT_STARTED;
 }
 
-unsigned long handle_rsi_attest_token_init(struct rec *rec)
+void handle_rsi_attest_token_init(struct rec *rec, struct rsi_result *res)
 {
 	struct rd *rd = NULL;
-	unsigned long ret;
 	unsigned long realm_buf_ipa = rec->regs[1];
 	void *rpv_ptr;
 	size_t rpv_len;
@@ -165,6 +158,8 @@
 
 	assert(rec != NULL);
 
+	res->action = UPDATE_REC_RETURN_TO_REALM;
+
 	/*
 	 * Calling RSI_ATTESTATION_TOKEN_INIT any time aborts any ongoing
 	 * operation.
@@ -183,7 +178,8 @@
 	}
 
 	if (!GRANULE_ALIGNED(realm_buf_ipa)) {
-		return RSI_ERROR_INPUT;
+		res->smc_res.x[0] = RSI_ERROR_INPUT;
+		return;
 	}
 
 	/*
@@ -193,7 +189,7 @@
 	granule_lock(rec->realm_info.g_rd, GRANULE_STATE_RD);
 	rd = granule_map(rec->realm_info.g_rd, SLOT_RD);
 	if (!addr_in_par(rd, realm_buf_ipa)) {
-		ret = RSI_ERROR_INPUT;
+		res->smc_res.x[0] = RSI_ERROR_INPUT;
 		goto out_unmap_rd;
 	}
 
@@ -212,67 +208,93 @@
 					    rec->rmm_realm_token_buf,
 					    sizeof(rec->rmm_realm_token_buf));
 	if (att_ret != 0) {
-		ERROR("FATAL_ERROR: Realm token creation failed,\n");
+		ERROR("FATAL_ERROR: Realm token creation failed\n");
 		panic();
 	}
 
 	rec->token_sign_ctx.state = ATTEST_SIGN_IN_PROGRESS;
-	ret = RSI_SUCCESS;
+	res->smc_res.x[0] = RSI_SUCCESS;
 
 out_unmap_rd:
 	buffer_unmap(rd);
 	granule_unlock(rec->realm_info.g_rd);
-	return ret;
+}
+
+/*
+ * Return 'false' if no IRQ is pending,
+ * return 'true' if there is an IRQ pending, and need to return to Host.
+ */
+static bool check_pending_irq(void)
+{
+	return (read_isr_el1() != 0UL);
 }
 
 void handle_rsi_attest_token_continue(struct rec *rec,
-				      struct attest_result *res)
+				      struct rmi_rec_exit *rec_exit,
+				      struct rsi_result *res)
 {
 	assert(rec != NULL);
-	assert(res != NULL);
+	assert(rec_exit != NULL);
 
-	/* Initialize attest_result */
-	res->incomplete = false;
-	res->walk_result.abort = false;
+	res->action = UPDATE_REC_RETURN_TO_REALM;
 
 	if (!verify_input_parameters_consistency(rec)) {
 		res->smc_res.x[0] = RSI_ERROR_INPUT;
 		return;
 	}
 
-	switch (rec->token_sign_ctx.state) {
-	case ATTEST_SIGN_NOT_STARTED:
-		/*
-		 * Before this call the initial attestation token call
-		 * (SMC_RSI_ATTEST_TOKEN_INIT) must have been executed
-		 * successfully.
-		 */
-		res->smc_res.x[0] = RSI_ERROR_STATE;
-		break;
-	case ATTEST_SIGN_IN_PROGRESS:
-		attest_token_continue_sign_state(rec, res);
-		break;
-	case ATTEST_SIGN_TOKEN_WRITE_IN_PROGRESS:
-		attest_token_continue_write_state(rec, res);
-		break;
-	default:
-		/* Any other state is considered an error. */
-		assert(false);
+	while (true) {
+		switch (rec->token_sign_ctx.state) {
+		case ATTEST_SIGN_NOT_STARTED:
+			/*
+			 * Before this call the initial attestation token call
+			 * (SMC_RSI_ATTEST_TOKEN_INIT) must have been executed
+			 * successfully.
+			 */
+			res->smc_res.x[0] = RSI_ERROR_STATE;
+			break;
+		case ATTEST_SIGN_IN_PROGRESS:
+			attest_token_continue_sign_state(rec, res);
+			break;
+		case ATTEST_SIGN_TOKEN_WRITE_IN_PROGRESS:
+			attest_token_continue_write_state(rec, res);
+			break;
+		default:
+			/* Any other state is considered an error */
+			panic();
+		}
+
+		if (res->smc_res.x[0] == RSI_INCOMPLETE) {
+			if (check_pending_irq()) {
+				res->action = UPDATE_REC_EXIT_TO_HOST;
+				res->smc_res.x[0] = RSI_INCOMPLETE;
+				rec_exit->exit_reason = RMI_EXIT_IRQ;
+				break;
+			}
+			if ((res->action & FLAG_EXIT_TO_HOST) != 0) {
+				break;
+			}
+		} else {
+			break;
+		}
 	}
 }
 
-unsigned long handle_rsi_extend_measurement(struct rec *rec)
+void handle_rsi_measurement_extend(struct rec *rec, struct rsi_result *res)
 {
 	struct granule *g_rd;
 	struct rd *rd;
 	unsigned long index;
 	unsigned long rd_addr;
 	size_t size;
-	unsigned long ret;
 	void *extend_measurement;
 	unsigned char *current_measurement;
 	int __unused meas_ret;
 
+	assert(rec != NULL);
+
+	res->action = UPDATE_REC_RETURN_TO_REALM;
+
 	/*
 	 * rd lock is acquired so that measurement cannot be updated
 	 * simultaneously by another rec
@@ -293,14 +315,14 @@
 
 	if ((index == RIM_MEASUREMENT_SLOT) ||
 	    (index >= MEASUREMENT_SLOT_NR)) {
-		ret = RSI_ERROR_INPUT;
+		res->smc_res.x[0] = RSI_ERROR_INPUT;
 		goto out_unmap_rd;
 	}
 
 	size  = rec->regs[2];
 
 	if (size > MAX_EXTENDED_SIZE) {
-		ret = RSI_ERROR_INPUT;
+		res->smc_res.x[0] = RSI_ERROR_INPUT;
 		goto out_unmap_rd;
 	}
 
@@ -313,15 +335,14 @@
 			   size,
 			   current_measurement);
 
-	ret = RSI_SUCCESS;
+	res->smc_res.x[0] = RSI_SUCCESS;
 
 out_unmap_rd:
 	buffer_unmap(rd);
 	granule_unlock(g_rd);
-	return ret;
 }
 
-unsigned long handle_rsi_read_measurement(struct rec *rec)
+void handle_rsi_measurement_read(struct rec *rec, struct rsi_result *res)
 {
 	struct rd *rd;
 	unsigned long idx;
@@ -329,11 +350,14 @@
 
 	assert(rec != NULL);
 
+	res->action = UPDATE_REC_RETURN_TO_REALM;
+
 	/* X1: Index */
 	idx = rec->regs[1];
 
 	if (idx >= MEASUREMENT_SLOT_NR) {
-		return RSI_ERROR_INPUT;
+		res->smc_res.x[0] = RSI_ERROR_INPUT;
+		return;
 	}
 
 	/*
@@ -356,5 +380,5 @@
 	buffer_unmap(rd);
 	granule_unlock(rec->realm_info.g_rd);
 
-	return RSI_SUCCESS;
+	res->smc_res.x[0] = RSI_SUCCESS;
 }
diff --git a/runtime/rsi/system.c b/runtime/rsi/system.c
index c224dc5..1b5d1bb 100644
--- a/runtime/rsi/system.c
+++ b/runtime/rsi/system.c
@@ -3,12 +3,14 @@
  * SPDX-FileCopyrightText: Copyright TF-RMM Contributors.
  */
 #include <assert.h>
+#include <rsi-handler.h>
 #include <smc-rsi.h>
 
 COMPILER_ASSERT(RSI_ABI_VERSION_MAJOR <= 0x7FFF);
 COMPILER_ASSERT(RSI_ABI_VERSION_MINOR <= 0xFFFF);
 
-unsigned long handle_rsi_version(void)
+void handle_rsi_version(struct rsi_result *res)
 {
-	return RSI_ABI_VERSION;
+	res->action = UPDATE_REC_RETURN_TO_REALM;
+	res->smc_res.x[0] = RSI_ABI_VERSION;
 }