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

[Oleksii Kurochko <oleksii.kurochko@xxxxxxxxx>]

Add helpers to safely update VS-stage and G-stage translation registers
during vCPU context switches.

Because VSATP and HGATP cannot be updated atomically, VSATP is cleared on
switch-out to prevent speculative VS-stage translations being cached under
an incorrect VMID. On VM entry, HGATP is reconstructed, VMID handling is
performed, and VSATP is restored.

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

Signed-off-by: Oleksii Kurochko <oleksii.kurochko@xxxxxxxxx>
---
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    |  11 +++
 xen/arch/riscv/p2m.c                | 123 ++++++++++++++++++++++++++++
 xen/arch/riscv/traps.c              |   2 +
 4 files changed, 137 insertions(+)

diff --git a/xen/arch/riscv/include/asm/domain.h 
b/xen/arch/riscv/include/asm/domain.h
index b5a8a9f711ac..c029e50b96fe 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..0cdd3dc44683 100644
--- a/xen/arch/riscv/include/asm/p2m.h
+++ b/xen/arch/riscv/include/asm/p2m.h
@@ -90,6 +90,13 @@ struct p2m_domain {
      */
     bool clean_dcache;
 
+    /*
+     * Inidicate that context switch is fully finished. It is needed to
+     * detect in p2m_handle_vmenter() to undestand what to write to
+     * CSR_VSATP register.
+     */
+    bool is_ctxt_switch_finished;
+
     /* Highest guest frame that's ever been mapped in the p2m */
     gfn_t max_mapped_gfn;
 
@@ -255,6 +262,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..275c38020ae2 100644
--- a/xen/arch/riscv/p2m.c
+++ b/xen/arch/riscv/p2m.c
@@ -1434,3 +1434,126 @@ struct page_info *p2m_get_page_from_gfn(struct 
p2m_domain *p2m, gfn_t gfn,
 
     return get_page(page, p2m->domain) ? page : NULL;
 }
+
+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 is constructed each time when
+     * p2m_handle_vmenter() is called.
+     */
+}
+
+void p2m_ctxt_switch_to(struct vcpu *n)
+{
+    if ( is_idle_vcpu(n) )
+        return;
+
+    n->domain->arch.p2m.is_ctxt_switch_finished = false;
+
+    /*
+     * Nothing to do with HGATP or VSATP, they will be set in
+     * p2_handle_vmenter()
+     */
+}
+
+void p2m_handle_vmenter(void)
+{
+    struct p2m_domain *p2m = &current->domain->arch.p2m;
+    struct vcpu_vmid *p_vmid = &current->arch.vmid;
+    uint16_t old_vmid, new_vmid;
+    bool need_flush;
+    register_t vsatp_old = 0;
+
+    BUG_ON(is_idle_vcpu(current));
+
+    /*
+     * 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
+     *
+     * CSR_VSATP is already set to 0 in p2m_ctxt_switch_from() in the
+     * case when n->arch.is_p2m_switch_finished = false. Also, there is
+     * BUG_ON() below to verify that.
+     */
+    if ( p2m->is_ctxt_switch_finished )
+        vsatp_old = csr_swap(CSR_VSATP, 0);
+
+    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",
+           current, old_vmid, new_vmid, need_flush);
+#endif
+
+    csr_write(CSR_HGATP, construct_hgatp(p2m_get_hostp2m(current->domain),
+              new_vmid));
+
+    if ( unlikely(need_flush) )
+        local_hfence_gvma_all();
+
+    if ( p2m->is_ctxt_switch_finished )
+        csr_swap(CSR_VSATP, vsatp_old);
+        /*
+         * We are not coming from a context switch here, so the VSATP value is
+         * the same as it was before csr_swap() was executed at the start of
+         * this function. Since VSATP was set to 0, no speculation could occur,
+         * and the VS-stage TLB cannot be polluted.
+         * Therefore, no additional VS TLB flush is required.
+         */
+    else
+    {
+        vsatp_old = csr_swap(CSR_VSATP, current->arch.vsatp);
+
+        /*
+         * vsatp_old should be zero as in the case of context switch it was
+         * set to 0 in p2m_ctxt_switch_from().
+         */
+        BUG_ON(vsatp_old);
+
+        p2m->is_ctxt_switch_finished = true;
+
+        /*
+         * TODO: further investigation is needed here.
+         *
+         *       In my opinion, a VS-stage TLB flush is not always strictly
+         *       necessary.
+         *       If a context switch occurs and VSATP is set to 0 before any
+         *       context-switch-related operations begin, no speculation can
+         *       occur. Therefore, at the time this function executes, the
+         *       VS-stage TLB should not be polluted with incorrect entries
+         *       belonging to a previously running vCPU. Another one reason is
+         *       that about SATP register is mentioned the following in RISC-V
+         *       spec:
+         *         Changing satp.MODE from Bare to other modes and vice versa
+         *         also takes effect immediately, without the need to execute
+         *         an SFENCE.VMA instruction. Likewise, changes to satp.ASID
+         *         take effect immediately.
+         *       I expect the same for VSATP (as it is VS copy of SATP) but it
+         *       isn't mentioned explicitly in the spec.
+         *
+         *       The only case where a VS-stage TLB flush seems necessary is
+         *       when the VASID remains unchanged but VSATP is updated to point
+         *       to a different VS page table. In that case, flushing
+         *       guarantees that the guest observes a clean context switch
+         *       without any possibility of using stale TLB entries.
+         */
+        flush_tlb_guest_local();
+    }
+}
diff --git a/xen/arch/riscv/traps.c b/xen/arch/riscv/traps.c
index e8d9ca902d9c..d6543eac1390 100644
--- a/xen/arch/riscv/traps.c
+++ b/xen/arch/riscv/traps.c
@@ -175,6 +175,8 @@ static void check_for_pcpu_work(void)
     vcpu_sync_interrupts(current);
 
     vcpu_flush_interrupts(current);
+
+    p2m_handle_vmenter();
 }
 
 static void timer_interrupt(void)
-- 
2.52.0