[PATCH 1/2] lib: sbi_domain: introduce atomic hart assignment helper

[Huamao Wu]

The per-hart domain assignment is represented by two fields: a
per-hart pointer (in sbi_scratch) and a per-domain bitmap
(assigned_harts).  These must stay consistent, but
switch_to_next_domain_context() and sbi_domain_register() update
them in separate locked steps with no surrounding critical
section, leaving a window where a concurrent reader can observe
a hart that belongs to no domain.

Introduce sbi_domain_assign_hart() which atomically removes a hart
from its current domain and adds it to the target domain under a
new domain_assignment_lock.

Add sbi_domain_check_hart_assignment() and its locked variant
sbi_domain_check_hart_assignment_locked() for callers that need a
consistent snapshot of both fields.  The unlocked variant is for
one-shot checks where the caller does not need to hold
domain_assignment_lock across a multi-step operation, such as
platform code on systems with an HSM that must validate domain
ownership of a hart before starting it.

Signed-off-by: Huamao Wu <huamao.wu at linux.spacemit.com>
---
 include/sbi/sbi_domain.h | 19 ++++++++++
 lib/sbi/sbi_domain.c     | 81 +++++++++++++++++++++++++++++++++++++---
 2 files changed, 94 insertions(+), 6 deletions(-)

diff --git a/include/sbi/sbi_domain.h b/include/sbi/sbi_domain.h
index 16edd4ce..a8528bd1 100644
--- a/include/sbi/sbi_domain.h
+++ b/include/sbi/sbi_domain.h
@@ -228,6 +228,25 @@ struct sbi_domain *sbi_hartindex_to_domain(u32 hartindex);
 /** Update HART local pointer to point to specified domain */
 void sbi_update_hartindex_to_domain(u32 hartindex, struct sbi_domain *dom);
 
+/** Atomically assign HART to specified domain */
+void sbi_domain_assign_hart(u32 hartindex, struct sbi_domain *dom);
+
+/** Lock domain assignment state for a compound operation */
+void sbi_domain_assignment_lock(void);
+
+/** Unlock domain assignment state after a compound operation */
+void sbi_domain_assignment_unlock(void);
+
+/** Check HART assignment while domain assignment lock is held */
+bool sbi_domain_check_hart_assignment_locked(const struct sbi_domain *dom,
+					     u32 hartindex,
+					     struct sbi_domain **mapped_dom);
+
+/** Atomically check HART assignment and return its mapped domain */
+bool sbi_domain_check_hart_assignment(const struct sbi_domain *dom,
+				      u32 hartindex,
+				      struct sbi_domain **mapped_dom);
+
 /** Get pointer to sbi_domain for current HART */
 #define sbi_domain_thishart_ptr() \
 	sbi_hartindex_to_domain(current_hartindex())
diff --git a/lib/sbi/sbi_domain.c b/lib/sbi/sbi_domain.c
index fa69170b..9ea7b7cb 100644
--- a/lib/sbi/sbi_domain.c
+++ b/lib/sbi/sbi_domain.c
@@ -37,6 +37,8 @@ struct sbi_domain root = {
 
 static unsigned long domain_hart_ptr_offset;
 
+static spinlock_t domain_assignment_lock = SPIN_LOCK_INITIALIZER;
+
 struct sbi_domain *sbi_hartindex_to_domain(u32 hartindex)
 {
 	struct sbi_scratch *scratch;
@@ -59,6 +61,78 @@ void sbi_update_hartindex_to_domain(u32 hartindex, struct sbi_domain *dom)
 	sbi_scratch_write_type(scratch, void *, domain_hart_ptr_offset, dom);
 }
 
+void sbi_domain_assignment_lock(void)
+{
+	spin_lock(&domain_assignment_lock);
+}
+
+void sbi_domain_assignment_unlock(void)
+{
+	spin_unlock(&domain_assignment_lock);
+}
+
+static void __sbi_domain_assign_hart_locked(u32 hartindex,
+					    struct sbi_domain *dom)
+{
+	struct sbi_domain *current_dom;
+
+	current_dom = sbi_hartindex_to_domain(hartindex);
+	if (current_dom) {
+		spin_lock(&current_dom->assigned_harts_lock);
+		sbi_hartmask_clear_hartindex(hartindex,
+						&current_dom->assigned_harts);
+		spin_unlock(&current_dom->assigned_harts_lock);
+	}
+	sbi_update_hartindex_to_domain(hartindex, dom);
+	if (dom) {
+		spin_lock(&dom->assigned_harts_lock);
+		sbi_hartmask_set_hartindex(hartindex, &dom->assigned_harts);
+		spin_unlock(&dom->assigned_harts_lock);
+	}
+}
+
+void sbi_domain_assign_hart(u32 hartindex, struct sbi_domain *dom)
+{
+	sbi_domain_assignment_lock();
+	__sbi_domain_assign_hart_locked(hartindex, dom);
+	sbi_domain_assignment_unlock();
+}
+
+bool sbi_domain_check_hart_assignment_locked(const struct sbi_domain *dom,
+						     u32 hartindex,
+						     struct sbi_domain **mapped_dom)
+{
+	bool assigned;
+	struct sbi_domain *mapped;
+
+	mapped = sbi_hartindex_to_domain(hartindex);
+	assigned = false;
+	if (dom && mapped == dom) {
+		spin_lock(&mapped->assigned_harts_lock);
+		assigned = sbi_hartmask_test_hartindex(hartindex,
+						      &mapped->assigned_harts);
+		spin_unlock(&mapped->assigned_harts_lock);
+	}
+	if (mapped_dom)
+		*mapped_dom = mapped;
+
+	return assigned;
+}
+
+bool sbi_domain_check_hart_assignment(const struct sbi_domain *dom,
+					      u32 hartindex,
+					      struct sbi_domain **mapped_dom)
+{
+	bool assigned;
+
+	sbi_domain_assignment_lock();
+	assigned = sbi_domain_check_hart_assignment_locked(dom, hartindex,
+						   mapped_dom);
+	sbi_domain_assignment_unlock();
+
+	return assigned;
+}
+
 bool sbi_domain_is_assigned_hart(const struct sbi_domain *dom, u32 hartindex)
 {
 	bool ret;
@@ -668,12 +742,7 @@ int sbi_domain_register(struct sbi_domain *dom,
 		if (!sbi_hartmask_test_hartindex(i, dom->possible_harts))
 			continue;
 
-		tdom = sbi_hartindex_to_domain(i);
-		if (tdom)
-			sbi_hartmask_clear_hartindex(i,
-					&tdom->assigned_harts);
-		sbi_update_hartindex_to_domain(i, dom);
-		sbi_hartmask_set_hartindex(i, &dom->assigned_harts);
+		sbi_domain_assign_hart(i, dom);
 
 		/*
 		 * If cold boot HART is assigned to this domain then
-- 
2.43.0