[PATCH 2/2] lib: sbi: serialize HSM hart_start ownership check with assignment lock

[Huamao Wu]

sbi_hsm_hart_start() checks domain ownership, acquires the start
ticket, and transitions the hart to START_PENDING in three unlocked
steps.  A concurrent domain assignment can reassign the hart
between the check and the state transition, causing HART_START to
proceed for a domain that no longer owns the target hart.

Hold domain_assignment_lock across the ownership check, ticket
acquisition, and STOPPED -> START_PENDING cmpxchg so that no
assignment can occur between check and commit.  Release the lock
before hardware start or IPI dispatch.

Convert switch_to_next_domain_context() to use
sbi_domain_assign_hart() instead of its open-coded three-step
sequence, eliminating the same TOCTOU window on the context switch
path.

Signed-off-by: Huamao Wu <huamao.wu at linux.spacemit.com>
---
 lib/sbi/sbi_domain_context.c | 13 ++-----------
 lib/sbi/sbi_hsm.c            | 15 ++++++++++++---
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/lib/sbi/sbi_domain_context.c b/lib/sbi/sbi_domain_context.c
index 0861d541..f71dd94f 100644
--- a/lib/sbi/sbi_domain_context.c
+++ b/lib/sbi/sbi_domain_context.c
@@ -110,24 +110,15 @@ static int switch_to_next_domain_context(struct hart_context *ctx,
 {
 	u32 hartindex = current_hartindex();
 	struct sbi_trap_context *trap_ctx;
-	struct sbi_domain *current_dom, *target_dom;
+	struct sbi_domain *target_dom;
 	struct sbi_scratch *scratch = sbi_scratch_thishart_ptr();
 
 	if (!ctx || !dom_ctx || ctx == dom_ctx)
 		return SBI_EINVAL;
 
-	current_dom = ctx->dom;
 	target_dom = dom_ctx->dom;
 	/* Assign current hart to target domain */
-	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, target_dom);
-
-	spin_lock(&target_dom->assigned_harts_lock);
-	sbi_hartmask_set_hartindex(hartindex, &target_dom->assigned_harts);
-	spin_unlock(&target_dom->assigned_harts_lock);
+	sbi_domain_assign_hart(hartindex, target_dom);
 
 	/* Save current CSR context and restore target domain's CSR context */
 	ctx->sstatus	= csr_swap(CSR_SSTATUS, dom_ctx->sstatus);
diff --git a/lib/sbi/sbi_hsm.c b/lib/sbi/sbi_hsm.c
index 0a355f9c..5293f995 100644
--- a/lib/sbi/sbi_hsm.c
+++ b/lib/sbi/sbi_hsm.c
@@ -318,8 +318,6 @@ int sbi_hsm_hart_start(struct sbi_scratch *scratch,
 	/* For now, we only allow start mode to be S-mode or U-mode. */
 	if (smode != PRV_S && smode != PRV_U)
 		return SBI_EINVAL;
-	if (dom && !sbi_domain_is_assigned_hart(dom, hartindex))
-		return SBI_EINVAL;
 	if (dom && !sbi_domain_check_addr(dom, saddr, smode,
 					  SBI_DOMAIN_EXECUTE))
 		return SBI_EINVALID_ADDR;
@@ -329,8 +327,18 @@ int sbi_hsm_hart_start(struct sbi_scratch *scratch,
 		return SBI_EINVAL;
 
 	hdata = sbi_scratch_offset_ptr(rscratch, hart_data_offset);
-	if (!hsm_start_ticket_acquire(hdata))
+
+	/* Serialize ownership validation with acceptance of this HART_START. */
+	sbi_domain_assignment_lock();
+	if (dom && !sbi_domain_check_hart_assignment_locked(dom, hartindex,
+							    NULL)) {
+		sbi_domain_assignment_unlock();
 		return SBI_EINVAL;
+	}
+	if (!hsm_start_ticket_acquire(hdata)) {
+		sbi_domain_assignment_unlock();
+		return SBI_EINVAL;
+	}
 
 	init_count = sbi_init_count(hartindex);
 	entry_count = sbi_entry_count(hartindex);
@@ -346,6 +354,7 @@ int sbi_hsm_hart_start(struct sbi_scratch *scratch,
 	 */
 	hstate = atomic_cmpxchg(&hdata->state, SBI_HSM_STATE_STOPPED,
 				SBI_HSM_STATE_START_PENDING);
+	sbi_domain_assignment_unlock();
 	if (hstate == SBI_HSM_STATE_STARTED) {
 		rc = SBI_EALREADY;
 		goto err;
-- 
2.43.0