[PATCH v3] xen/riscv: add p2m context switch handling for VSATP and HGATP

[Oleksii Kurochko <oleksii.kurochko@xxxxxxxxx>]

Introduce helpers to manage VS-stage and G-stage translation state during
vCPU context switches.

As VSATP and HGATP cannot be updated atomically, clear VSATP on context
switch-out to prevent speculative VS-stage translations from being associated
with an incorrect VMID. On context switch-in, restore HGATP and VSATP in the
required order.

Add p2m_handle_vmenter() to perform VMID management and issue TLB flushes
only when required (e.g. on VMID reuse or generation change).

This provides the necessary infrastructure for correct p2m context switching
on RISC-V.

Signed-off-by: Oleksii Kurochko <oleksii.kurochko@xxxxxxxxx>
---
Changes in v3:
 - Add comment above p2m_ctxt_switch_{to, from}().
 - Code style fixes.
 - Refactor p2m_ctxt_switch_to().
 - Update the comment at the end of p2m_ctxt_switch_from().
 - Refactor the code of p2m_handle_vmenter().
---
Changes in v2:
 - Add vsatp field declaration to arch_vcpu.
 - s/p2m_ctx_switch_{from,to}/p2m_ctxt_switch_{from,to}.
 - Introduce p2m_handle_vmenter() for proper handling of VMID,
   hgatp and vsatp updates.
 - Introduce is_p2m_switch_finished and init it inisde
   p2m_ctx_switch_to() for furhter handling in p2m_handle_vmenter().
 - Code style fixes.
 - Add is_idle_vcpu() check in p2m_ctxt_switch_from().
 - use csr_swap() in p2m_ctxt_switch_from().
 - move flush_tlb_guest_local() to the end if p2m_handle_vmenter() and
   drop unnessary anymore comments.
 - Correct printk()'s arguments in p2m_handle_vmenter().
---
 xen/arch/riscv/include/asm/domain.h |  1 +
 xen/arch/riscv/include/asm/p2m.h    |  4 ++
 xen/arch/riscv/p2m.c                | 79 +++++++++++++++++++++++++++++
 xen/arch/riscv/traps.c              |  2 +
 4 files changed, 86 insertions(+)

diff --git a/xen/arch/riscv/include/asm/domain.h 
b/xen/arch/riscv/include/asm/domain.h
index 3da2387cb197..42bb678fcbf9 100644
--- a/xen/arch/riscv/include/asm/domain.h
+++ b/xen/arch/riscv/include/asm/domain.h
@@ -59,6 +59,7 @@ struct arch_vcpu {
     register_t hstateen0;
     register_t hvip;
 
+    register_t vsatp;
     register_t vsie;
 
     /*
diff --git a/xen/arch/riscv/include/asm/p2m.h b/xen/arch/riscv/include/asm/p2m.h
index f63b5dec99b1..60f27f9b347e 100644
--- a/xen/arch/riscv/include/asm/p2m.h
+++ b/xen/arch/riscv/include/asm/p2m.h
@@ -255,6 +255,10 @@ static inline bool p2m_is_locked(const struct p2m_domain 
*p2m)
 struct page_info *p2m_get_page_from_gfn(struct p2m_domain *p2m, gfn_t gfn,
                                         p2m_type_t *t);
 
+void p2m_ctxt_switch_from(struct vcpu *p);
+void p2m_ctxt_switch_to(struct vcpu *n);
+void p2m_handle_vmenter(void);
+
 #endif /* ASM__RISCV__P2M_H */
 
 /*
diff --git a/xen/arch/riscv/p2m.c b/xen/arch/riscv/p2m.c
index 0abeb374c110..7ae854707174 100644
--- a/xen/arch/riscv/p2m.c
+++ b/xen/arch/riscv/p2m.c
@@ -1434,3 +1434,82 @@ struct page_info *p2m_get_page_from_gfn(struct 
p2m_domain *p2m, gfn_t gfn,
 
     return get_page(page, p2m->domain) ? page : NULL;
 }
+
+/* Should be called before other CSRs are stored to avoid speculation */
+void p2m_ctxt_switch_from(struct vcpu *p)
+{
+    if ( is_idle_vcpu(p) )
+        return;
+
+    /*
+     * No mechanism is provided to atomically change vsatp and hgatp
+     * together. Hence, to prevent speculative execution causing one
+     * guest’s VS-stage translations to be cached under another guest’s
+     * VMID, world-switch code should zero vsatp, then swap hgatp, then
+     * finally write the new vsatp value what will be done in
+     * p2m_handle_vmenter().
+     */
+    p->arch.vsatp = csr_swap(CSR_VSATP, 0);
+
+    /*
+     * Nothing to do with HGATP as it will be update in p2m_ctxt_switch_to()
+     * or/and in p2m_handle_vmenter().
+     */
+}
+
+/* Should be called after other CSRs are restored to avoid speculation */
+void p2m_ctxt_switch_to(struct vcpu *n)
+{
+    struct p2m_domain *p2m = p2m_get_hostp2m(n->domain);
+
+    if ( is_idle_vcpu(n) )
+        return;
+
+    csr_write(CSR_HGATP, construct_hgatp(p2m, n->arch.vmid.vmid));
+    /*
+     * As VMID is unique per vCPU and just re-used here thereby there is no
+     * need for G-stage TLB flush here.
+     */
+
+    csr_write(CSR_VSATP, n->arch.vsatp);
+    /*
+     * As at the start of context switch VSATP were set to 0, so no speculation
+     * could happen thereby there is no need for VS TLB flush here.
+     */
+}
+
+void p2m_handle_vmenter(void)
+{
+    struct vcpu *c = current;
+    struct p2m_domain *p2m = p2m_get_hostp2m(c->domain);
+    struct vcpu_vmid *p_vmid = &c->arch.vmid;
+    uint16_t old_vmid, new_vmid;
+    bool need_flush;
+
+    BUG_ON(is_idle_vcpu(current));
+
+    old_vmid = p_vmid->vmid;
+    need_flush = vmid_handle_vmenter(p_vmid);
+    new_vmid = p_vmid->vmid;
+
+#ifdef P2M_DEBUG
+    printk("%pv: oldvmid(%d) new_vmid(%d), need_flush(%d)\n",
+           c, old_vmid, new_vmid, need_flush);
+#endif
+
+    if ( old_vmid != new_vmid )
+        csr_write(CSR_HGATP, construct_hgatp(p2m, p_vmid->vmid));
+
+    if ( unlikely(need_flush) )
+    {
+        local_hfence_gvma_all();
+        flush_tlb_guest_local();
+    }
+
+    /*
+     * There is no need to set VSATP to 0 to stop speculation before updating
+     * HGATP, as VSATP is not modified here. There is also no need to flush
+     * the VS-stage TLB: even if speculation occurs, it will use the old VMID,
+     * which will not be reused until need_flush is set to true.
+     */
+}
diff --git a/xen/arch/riscv/traps.c b/xen/arch/riscv/traps.c
index ce8d346a14d2..7ef089809734 100644
--- a/xen/arch/riscv/traps.c
+++ b/xen/arch/riscv/traps.c
@@ -177,6 +177,8 @@ static void check_for_pcpu_work(void)
     vcpu_sync_interrupts(c);
 
     vcpu_flush_interrupts(c);
+
+    p2m_handle_vmenter();
 }
 
 static void timer_interrupt(void)
-- 
2.52.0