fix(s2tt): fix bugs on s2tt related to NS memory assignment/unassignment

This patch fixes a number of bugs found on smc_rtt_fold(),
smc_rtt_create() and map_upmap_ns() related to NS memory assigment
and unassignment.

- SMC_RTT_FOLD: ASSIGNED_NS case did not propagate attributes to the
                parent S2TTE.
- SMC_RTT_CREATE: ASSIGNED_NS case did not propagate attributes to the
                child S2TT (unfold operation) and the refcount was
		not updated properly.

Signed-off-by: Javier Almansa Sobrino <javier.almansasobrino@arm.com>
Change-Id: I240d16ca8e17485fb09183d3ee8172af689563d0
diff --git a/lib/realm/src/s2tt.c b/lib/realm/src/s2tt.c
index c36655c..b66dde6 100644
--- a/lib/realm/src/s2tt.c
+++ b/lib/realm/src/s2tt.c
@@ -65,6 +65,8 @@
 
 #define S2TTE_ATTRS	(S2TTE_MEMATTR_FWB_NORMAL_WB | S2TTE_AP_RW | \
 			S2TTE_SH_IS | S2TTE_AF)
+#define S2TTE_NS_ATTR_MASK (S2TTE_MEMATTR_MASK | S2TTE_AP_MASK | \
+			    S2TTE_SH_MASK)
 
 #define S2TTE_TABLE	S2TTE_L012_TABLE
 #define S2TTE_BLOCK	(S2TTE_ATTRS | S2TTE_L012_BLOCK)
@@ -497,14 +499,17 @@
  * - MemAttr
  * - S2AP
  * - Shareability
+ * Any other field on @s2tte is masked out.
  */
-unsigned long s2tte_create_assigned_ns(unsigned long pa, long level)
+unsigned long s2tte_create_assigned_ns(unsigned long s2tte, long level)
 {
+	unsigned long new_s2tte = s2tte & ~DESC_TYPE_MASK;
+
 	assert(level >= RTT_MIN_BLOCK_LEVEL);
 	if (level == RTT_PAGE_LEVEL) {
-		return (pa | S2TTE_PAGE_NS);
+		return (new_s2tte | S2TTE_PAGE_NS);
 	}
-	return (pa | S2TTE_BLOCK_NS);
+	return (new_s2tte | S2TTE_BLOCK_NS);
 }
 
 /*
@@ -513,9 +518,7 @@
 bool host_ns_s2tte_is_valid(unsigned long s2tte, long level)
 {
 	unsigned long mask = addr_level_mask(~0UL, level) |
-			     S2TTE_MEMATTR_MASK |
-			     S2TTE_AP_MASK |
-			     S2TTE_SH_MASK;
+						S2TTE_NS_ATTR_MASK;
 
 	/*
 	 * Test that all fields that are not controlled by the host are zero
@@ -552,9 +555,7 @@
 unsigned long host_ns_s2tte(unsigned long s2tte, long level)
 {
 	unsigned long mask = addr_level_mask(~0UL, level) |
-			     S2TTE_MEMATTR_MASK |
-			     S2TTE_AP_MASK |
-			     S2TTE_SH_MASK;
+						S2TTE_NS_ATTR_MASK;
 	return (s2tte & mask);
 }
 
@@ -874,20 +875,26 @@
 }
 
 /*
- * Populates @s2tt with assigned_ns s2ttes that refer to a
+ * Populates @s2tt with NS attributes @attrs that refer to a
  * contiguous memory block starting at @pa, and mapped at level @level.
  *
  * The granule is populated before it is made a table,
  * hence, don't use s2tte_write for access.
  */
-void s2tt_init_assigned_ns(unsigned long *s2tt, unsigned long pa, long level)
+void s2tt_init_assigned_ns(unsigned long *s2tt, unsigned long attrs,
+			   unsigned long pa, long level)
 {
 	const unsigned long map_size = s2tte_map_size(level);
 
+	assert(addr_is_level_aligned(pa, level));
+
 	for (unsigned int i = 0U; i < S2TTES_PER_S2TT; i++) {
-		s2tt[i] = s2tte_create_assigned_ns(pa, level);
+		unsigned long s2tte = attrs & S2TTE_NS_ATTR_MASK;
+
+		s2tt[i] = s2tte_create_assigned_ns(s2tte | pa, level);
 		pa += map_size;
 	}
+
 	dsb(ish);
 }
 
@@ -998,7 +1005,8 @@
 
 static bool __table_maps_block(unsigned long *table,
 			       long level,
-			       s2tte_type_level_checker s2tte_is_x)
+			       s2tte_type_level_checker s2tte_is_x,
+			       bool check_ns_attrs)
 {
 	unsigned long base_pa;
 	unsigned long map_size = s2tte_map_size(level);
@@ -1026,6 +1034,18 @@
 		if (s2tte_pa(s2tte, level) != expected_pa) {
 			return false;
 		}
+
+		if (check_ns_attrs) {
+			unsigned long ns_attrs = s2tte & S2TTE_NS_ATTR_MASK;
+
+			/*
+			 * We match all the attributes in the S2TTE
+			 * except for the AF bit.
+			 */
+			if ((s2tte & S2TTE_NS_ATTR_MASK) != ns_attrs) {
+				return false;
+			}
+		}
 	}
 
 	return true;
@@ -1037,7 +1057,8 @@
  */
 bool table_maps_assigned_empty_block(unsigned long *table, long level)
 {
-	return __table_maps_block(table, level, s2tte_is_assigned_empty);
+	return __table_maps_block(table, level, s2tte_is_assigned_empty,
+				  false);
 }
 
 /*
@@ -1046,18 +1067,20 @@
  */
 bool table_maps_assigned_ram_block(unsigned long *table, long level)
 {
-	return __table_maps_block(table, level, s2tte_is_assigned_ram);
+	return __table_maps_block(table, level, s2tte_is_assigned_ram, false);
 }
 
 /*
- * Returns true if all s2ttes in @table are assigned_ns s2ttes and
- * refer to a contiguous block of granules aligned to @level - 1.
+ * Returns true if
+ *    - all s2ttes in @table are assigned_ns s2ttes and
+ *    - they refer to a contiguous block of granules aligned to @level - 1 and
+ *    - all the s2tte attributes in @table controlled by the host are identical
  *
  * @pre: @table maps IPA outside PAR.
  */
 bool table_maps_assigned_ns_block(unsigned long *table, long level)
 {
-	return __table_maps_block(table, level, s2tte_is_assigned_ns);
+	return __table_maps_block(table, level, s2tte_is_assigned_ns, true);
 }
 
 /*
@@ -1066,7 +1089,8 @@
  */
 bool table_maps_assigned_destroyed_block(unsigned long *table, long level)
 {
-	return __table_maps_block(table, level, s2tte_is_assigned_destroyed);
+	return __table_maps_block(table, level, s2tte_is_assigned_destroyed,
+				  false);
 }
 
 /*