feat(rmm): move attestation data to REC auxiliary granule

This patch moves attestation data from REC to
'rec_attest_data' structure in auxiliary granule.
In 'handle_rsi_attest_token_init' function
initialisation of variables which depend on 'rec'
parameter is moved after
assert(rec != NULL);
to prevent accessing 'rec' in case of NULL pointer.

Signed-off-by: AlexeiFedorov <Alexei.Fedorov@arm.com>
Change-Id: Ia182d66bfb19a2baea5a9d2039dbbf6523e72bdd
diff --git a/lib/allocator/include/memory_alloc.h b/lib/allocator/include/memory_alloc.h
index 852b648..b4a9665 100644
--- a/lib/allocator/include/memory_alloc.h
+++ b/lib/allocator/include/memory_alloc.h
@@ -29,9 +29,17 @@
  * Total size is ~3 Pages (rounded up to page size).
  */
 #define REC_SIMD_PAGES		3
+#define REC_SIMD_SIZE		(REC_SIMD_PAGES * SZ_4K)
 
-/* Number of aux granules pages per REC to be used */
-#define REC_NUM_PAGES		(REC_HEAP_PAGES + REC_PMU_PAGES + REC_SIMD_PAGES)
+/* Number of pages per REC for 'rec_attest_data' structure */
+#define REC_ATTEST_PAGES	1
+#define REC_ATTEST_SIZE		(REC_ATTEST_PAGES * SZ_4K)
+
+/* Number of pages per REC to be allocated */
+#define REC_NUM_PAGES		(REC_HEAP_PAGES	  + \
+				 REC_PMU_PAGES	  + \
+				 REC_SIMD_PAGES	  + \
+				 REC_ATTEST_PAGES)
 
 struct buffer_alloc_ctx {
 	unsigned char		*buf;
diff --git a/lib/realm/include/rec.h b/lib/realm/include/rec.h
index 961bfaf..8140af1 100644
--- a/lib/realm/include/rec.h
+++ b/lib/realm/include/rec.h
@@ -107,13 +107,38 @@
 } __attribute__((aligned(CACHE_WRITEBACK_GRANULE)));
 
 /*
+ * Data used when handling attestation requests
+ */
+struct rec_attest_data {
+	unsigned char rmm_realm_token_buf[SZ_1K];
+	size_t rmm_realm_token_len;
+
+	struct token_sign_ctx token_sign_ctx;
+
+	/* Buffer allocation info used for heap init and management */
+	struct {
+		struct buffer_alloc_ctx ctx;
+		bool ctx_initialised;
+	} alloc_info;
+};
+COMPILER_ASSERT(sizeof(struct rec_attest_data) <= GRANULE_SIZE);
+
+/*
  * This structure contains pointers to data that are allocated
  * in auxilary granules for a REC.
  */
 struct rec_aux_data {
-	uint8_t *attest_heap_buf; /* pointer to the heap buffer */
-	struct pmu_state *pmu;	  /* pointer to PMU state */
-	struct rec_simd_state rec_simd; /* REC SIMD context region */
+	/* Pointer to the heap buffer */
+	uint8_t *attest_heap_buf;
+
+	/* Pointer to PMU state */
+	struct pmu_state *pmu;
+
+	/* SIMD context region */
+	struct rec_simd_state rec_simd;
+
+	/* Pointer to attestation-related data */
+	struct rec_attest_data *attest_data;
 };
 
 struct rec {
@@ -189,18 +214,6 @@
 	/* Addresses of auxiliary granules */
 	struct granule *g_aux[MAX_REC_AUX_GRANULES];
 	struct rec_aux_data aux_data;
-
-	unsigned char rmm_realm_token_buf[SZ_1K];
-	size_t rmm_realm_token_len;
-
-	struct token_sign_ctx token_sign_ctx;
-
-	/* Buffer allocation info used for heap init and management */
-	struct {
-		struct buffer_alloc_ctx ctx;
-		bool ctx_initialised;
-	} alloc_info;
-
 	struct {
 		unsigned long vsesr_el2;
 		bool inject;
diff --git a/runtime/core/run.c b/runtime/core/run.c
index a94d5af..fb5528a 100644
--- a/runtime/core/run.c
+++ b/runtime/core/run.c
@@ -28,17 +28,22 @@
 void init_rec_aux_data(struct rec_aux_data *aux_data, void *rec_aux,
 		       unsigned long num_aux)
 {
-	/* Ensure we have enough aux granules for use by REC */
+	/*
+	 * Ensure we have enough aux granules for use by REC:
+	 * - REC_HEAP_PAGES for MbedTLS heap
+	 * - REC_PMU_PAGES for PMU state
+	 * - REC_SIMD_PAGES for SIMD state
+	 * - REC_ATTEST_PAGES for 'rec_attest_data' structure
+	 */
 	assert(num_aux >= REC_NUM_PAGES);
 
 	aux_data->attest_heap_buf = (uint8_t *)rec_aux;
-
 	aux_data->pmu = (struct pmu_state *)((uint8_t *)rec_aux +
-					     REC_HEAP_SIZE);
-
+				REC_HEAP_SIZE);
 	aux_data->rec_simd.simd = (struct simd_state *)((uint8_t *)rec_aux +
-							REC_HEAP_SIZE +
-							REC_PMU_SIZE);
+				REC_HEAP_SIZE + REC_PMU_SIZE);
+	aux_data->attest_data = (struct rec_attest_data *)((uint8_t *)rec_aux +
+				REC_HEAP_SIZE + REC_PMU_SIZE + REC_SIMD_SIZE);
 }
 
 /*
@@ -354,6 +359,7 @@
 	int realm_exception_code;
 	void *rec_aux;
 	unsigned int cpuid = my_cpuid();
+	struct rec_attest_data *attest_data;
 
 	assert(cpuid < MAX_CPUS);
 	assert(rec->ns == NULL);
@@ -368,23 +374,22 @@
 	/* Map auxiliary granules */
 	rec_aux = map_rec_aux(rec->g_aux, rec->num_rec_aux);
 
-	init_rec_aux_data(&(rec->aux_data), rec_aux, rec->num_rec_aux);
-
 	/*
 	 * The attset heap on the REC aux pages is mapped now. It is time to
 	 * associate it with the current CPU.
 	 * This heap will be used for attestation RSI calls when the
 	 * REC is running.
 	 */
-	attestation_heap_ctx_assign_pe(&rec->alloc_info.ctx);
+	attest_data = rec->aux_data.attest_data;
+	attestation_heap_ctx_assign_pe(&attest_data->alloc_info.ctx);
 
 	/*
 	 * Initialise the heap for attestation if necessary.
 	 */
-	if (!rec->alloc_info.ctx_initialised) {
+	if (!attest_data->alloc_info.ctx_initialised) {
 		(void)attestation_heap_ctx_init(rec->aux_data.attest_heap_buf,
 						REC_HEAP_SIZE);
-		rec->alloc_info.ctx_initialised = true;
+		attest_data->alloc_info.ctx_initialised = true;
 	}
 
 	rec_simd_state_init(rec);
@@ -456,6 +461,7 @@
 
 	/* Undo the heap association */
 	attestation_heap_ctx_unassign_pe();
+
 	/* Unmap auxiliary granules */
 	unmap_rec_aux(rec_aux, rec->num_rec_aux);
 }
diff --git a/runtime/include/run.h b/runtime/include/run.h
index 546e766..a504a51 100644
--- a/runtime/include/run.h
+++ b/runtime/include/run.h
@@ -7,6 +7,7 @@
 #define RUN_H
 
 struct granule;
+struct rec;
 struct rec_aux_data;
 
 void init_rec_aux_data(struct rec_aux_data *aux_data, void *rec_aux,
diff --git a/runtime/rmi/rec.c b/runtime/rmi/rec.c
index b21923e..7226f90 100644
--- a/runtime/rmi/rec.c
+++ b/runtime/rmi/rec.c
@@ -17,6 +17,7 @@
 #include <psci.h>
 #include <realm.h>
 #include <rec.h>
+#include <run.h>
 #include <smc-handler.h>
 #include <smc-rmi.h>
 #include <smc.h>
@@ -201,6 +202,8 @@
 	unsigned long ret;
 	bool ns_access_ok;
 	unsigned int num_rec_aux;
+	void *rec_aux;
+	struct rec_attest_data *attest_data;
 
 	g_rec_params = find_granule(rec_params_addr);
 	if ((g_rec_params == NULL) || (g_rec_params->state != GRANULE_STATE_NS)) {
@@ -297,9 +300,20 @@
 	new_rec_state = GRANULE_STATE_REC;
 	rec->runnable = rec_params.flags & REC_PARAMS_FLAG_RUNNABLE;
 
-	rec->alloc_info.ctx_initialised = false;
+	/*
+	 * The access to rec_aux granule is not protected by the lock
+	 * in its granule but by the lock of the parent REC granule.
+	 */
+	rec_aux = map_rec_aux(rec->g_aux, rec->num_rec_aux);
+	init_rec_aux_data(&(rec->aux_data), rec_aux, rec->num_rec_aux);
+
+	attest_data = rec->aux_data.attest_data;
+	attest_data->alloc_info.ctx_initialised = false;
+
 	/* Initialize attestation state */
-	rec->token_sign_ctx.state = ATTEST_SIGN_NOT_STARTED;
+	attest_data->token_sign_ctx.state = ATTEST_SIGN_NOT_STARTED;
+
+	unmap_rec_aux(rec_aux, rec->num_rec_aux);
 
 	set_rd_rec_count(rd, rec_idx + 1U);
 
diff --git a/runtime/rsi/realm_attest.c b/runtime/rsi/realm_attest.c
index c83888c..b6f39c6 100644
--- a/runtime/rsi/realm_attest.c
+++ b/runtime/rsi/realm_attest.c
@@ -36,8 +36,10 @@
  */
 static void save_input_parameters(struct rec *rec)
 {
-	rec->token_sign_ctx.token_ipa = rec->regs[1];
-	(void)memcpy(rec->token_sign_ctx.challenge, &rec->regs[2],
+	struct rec_attest_data *attest_data = rec->aux_data.attest_data;
+
+	attest_data->token_sign_ctx.token_ipa = rec->regs[1];
+	(void)memcpy(attest_data->token_sign_ctx.challenge, &rec->regs[2],
 		     ATTEST_CHALLENGE_SIZE);
 }
 
@@ -47,21 +49,24 @@
  */
 static bool verify_input_parameters_consistency(struct rec *rec)
 {
-	return rec->token_sign_ctx.token_ipa == rec->regs[1];
+	struct rec_attest_data *attest_data = rec->aux_data.attest_data;
+
+	return attest_data->token_sign_ctx.token_ipa == rec->regs[1];
 }
 
 /*
  * Function to continue with the sign operation
  */
-static void attest_token_continue_sign_state(struct rec *rec,
-					     struct rsi_result *res)
+static void attest_token_continue_sign_state(
+					struct rec_attest_data *attest_data,
+					struct rsi_result *res)
 {
 	/*
 	 * Sign and finish creating the token.
 	 */
 	enum attest_token_err_t ret =
-		attest_realm_token_sign(&(rec->token_sign_ctx.ctx),
-					&(rec->rmm_realm_token_len));
+		attest_realm_token_sign(&(attest_data->token_sign_ctx.ctx),
+					&(attest_data->rmm_realm_token_len));
 
 	if ((ret == ATTEST_TOKEN_ERR_COSE_SIGN_IN_PROGRESS) ||
 		(ret == ATTEST_TOKEN_ERR_SUCCESS)) {
@@ -74,7 +79,7 @@
 
 		/* If this was the last signing cycle */
 		if (ret == ATTEST_TOKEN_ERR_SUCCESS) {
-			rec->token_sign_ctx.state =
+			attest_data->token_sign_ctx.state =
 				ATTEST_SIGN_TOKEN_WRITE_IN_PROGRESS;
 		}
 	} else {
@@ -96,6 +101,7 @@
 	enum s2_walk_status walk_status;
 	struct s2_walk_result walk_res = { 0UL };
 	size_t attest_token_len;
+	struct rec_attest_data *attest_data = rec->aux_data.attest_data;
 
 	/*
 	 * Translate realm granule IPA to PA. If returns with
@@ -126,9 +132,9 @@
 	realm_att_token = granule_map(gr, SLOT_RSI_CALL);
 
 	attest_token_len = attest_cca_token_create(realm_att_token,
-						   ATTEST_TOKEN_BUFFER_SIZE,
-						   (void *)rec->rmm_realm_token_buf,
-						   rec->rmm_realm_token_len);
+						ATTEST_TOKEN_BUFFER_SIZE,
+						&attest_data->rmm_realm_token_buf,
+						attest_data->rmm_realm_token_len);
 
 	/* Unmap realm granule */
 	buffer_unmap(realm_att_token);
@@ -145,19 +151,23 @@
 	}
 
 	/* The signing has either succeeded or failed. Reset the state. */
-	rec->token_sign_ctx.state = ATTEST_SIGN_NOT_STARTED;
+	attest_data->token_sign_ctx.state = ATTEST_SIGN_NOT_STARTED;
 }
 
 void handle_rsi_attest_token_init(struct rec *rec, struct rsi_result *res)
 {
 	struct rd *rd = NULL;
-	unsigned long realm_buf_ipa = rec->regs[1];
+	unsigned long realm_buf_ipa;
+	struct rec_attest_data *attest_data;
 	void *rpv_ptr;
 	size_t rpv_len;
 	int att_ret;
 
 	assert(rec != NULL);
 
+	realm_buf_ipa = rec->regs[1];
+	attest_data = rec->aux_data.attest_data;
+
 	res->action = UPDATE_REC_RETURN_TO_REALM;
 
 	/*
@@ -165,10 +175,10 @@
 	 * operation.
 	 * TODO: This can be moved to attestation lib
 	 */
-	if (rec->token_sign_ctx.state != ATTEST_SIGN_NOT_STARTED) {
+	if (attest_data->token_sign_ctx.state != ATTEST_SIGN_NOT_STARTED) {
 		int restart;
 
-		rec->token_sign_ctx.state = ATTEST_SIGN_NOT_STARTED;
+		attest_data->token_sign_ctx.state = ATTEST_SIGN_NOT_STARTED;
 		restart = attestation_heap_reinit_pe(rec->aux_data.attest_heap_buf,
 							REC_HEAP_SIZE);
 		if (restart != 0) {
@@ -204,15 +214,15 @@
 					    MEASUREMENT_SLOT_NR,
 					    rpv_ptr,
 					    rpv_len,
-					    &rec->token_sign_ctx,
-					    rec->rmm_realm_token_buf,
-					    sizeof(rec->rmm_realm_token_buf));
+					    &attest_data->token_sign_ctx,
+					    attest_data->rmm_realm_token_buf,
+					    sizeof(attest_data->rmm_realm_token_buf));
 	if (att_ret != 0) {
 		ERROR("FATAL_ERROR: Realm token creation failed\n");
 		panic();
 	}
 
-	rec->token_sign_ctx.state = ATTEST_SIGN_IN_PROGRESS;
+	attest_data->token_sign_ctx.state = ATTEST_SIGN_IN_PROGRESS;
 	res->smc_res.x[0] = RSI_SUCCESS;
 
 out_unmap_rd:
@@ -233,9 +243,12 @@
 				      struct rmi_rec_exit *rec_exit,
 				      struct rsi_result *res)
 {
+	struct rec_attest_data *attest_data;
+
 	assert(rec != NULL);
 	assert(rec_exit != NULL);
 
+	attest_data = rec->aux_data.attest_data;
 	res->action = UPDATE_REC_RETURN_TO_REALM;
 
 	if (!verify_input_parameters_consistency(rec)) {
@@ -244,7 +257,7 @@
 	}
 
 	while (true) {
-		switch (rec->token_sign_ctx.state) {
+		switch (attest_data->token_sign_ctx.state) {
 		case ATTEST_SIGN_NOT_STARTED:
 			/*
 			 * Before this call the initial attestation token call
@@ -254,7 +267,7 @@
 			res->smc_res.x[0] = RSI_ERROR_STATE;
 			break;
 		case ATTEST_SIGN_IN_PROGRESS:
-			attest_token_continue_sign_state(rec, res);
+			attest_token_continue_sign_state(attest_data, res);
 			break;
 		case ATTEST_SIGN_TOKEN_WRITE_IN_PROGRESS:
 			attest_token_continue_write_state(rec, res);