[PATCH v1 2/2] RISC-V: KVM: fix vcpu vector context handling for kernel-mode vector

[PATCH v1 2/2] RISC-V: KVM: fix vcpu vector context handling for kernel-mode vector

[Andy Chiu]

Running vector workloads like perf + mcf on KVM can result in an
unexpected termination due to a vtype corruption. This happens because
the kernel-mode vector (KMV) misattributes the guest's vcpu context as
the user's context and source from a wrong status.VS.

The simplified call chain that results in this problem is shown as
follow:

__riscv_sys_ioctl()
    kvm_arch_vcpu_ioctl_run()
      kvm_riscv_vcpu_exit()
        kvm_riscv_vcpu_sbi_ecall()
          kvm_riscv_vcpu_pmu_ctr_stop()
            kvm_vcpu_write_guest()
              __copy_to_user()
                enter_vector_usercopy()
                  kernel_vector_begin()

kernel_vector_begin() should use the sstatus.VS from guest's vcpu
context instead of task_pt_reg(current). Also, it should not save
guest's v-reg into the user's context memory.

To resolve this, the vcpu context must be correctly saved when KMV is
serving a guest. However, invoking KVM functions directly from generic
RISC-V architecture code introduces a reverse dependency, breaking
builds when KVM is configured as N or M.

Address this by registering an RCU-protected callback for context
flushing. KVM registers this callback at module initialization and
unregisters it on exit. When KMV starts a kernel context, it can now
safely flush the vector context via the callback.

Fixes: ecd2ada8a5e0 ("riscv: Add support for kernel mode vector")
Signed-off-by: Andy Chiu <tchiu@tenstorrent.com>
---
 arch/riscv/include/asm/kvm_vcpu_vector.h | 24 ++++++++++++
 arch/riscv/include/asm/processor.h       |  1 +
 arch/riscv/kernel/kernel_mode_vector.c   | 48 +++++++++++++++++++-----
 arch/riscv/kvm/main.c                    |  4 ++
 arch/riscv/kvm/vcpu.c                    |  6 +++
 arch/riscv/kvm/vcpu_vector.c             | 18 +++++++++
 6 files changed, 91 insertions(+), 10 deletions(-)

diff --git a/arch/riscv/include/asm/kvm_vcpu_vector.h b/arch/riscv/include/asm/kvm_vcpu_vector.h
index 57a798a4cb0d..e36f46de7429 100644
--- a/arch/riscv/include/asm/kvm_vcpu_vector.h
+++ b/arch/riscv/include/asm/kvm_vcpu_vector.h
@@ -35,6 +35,22 @@ void kvm_riscv_vcpu_host_vector_save(struct kvm_cpu_context *cntx);
 void kvm_riscv_vcpu_host_vector_restore(struct kvm_cpu_context *cntx);
 int kvm_riscv_vcpu_alloc_vector_context(struct kvm_vcpu *vcpu);
 void kvm_riscv_vcpu_free_vector_context(struct kvm_vcpu *vcpu);
+void kvm_riscv_register_vctx_callback(bool (*func)(void));
+void kvm_riscv_unregister_vctx_callback(void);
+bool kvm_riscv_vcpu_flush_vector(void);
+
+static inline void kvm_riscv_v_init(void)
+{
+	if (has_vector())
+		kvm_riscv_register_vctx_callback(&kvm_riscv_vcpu_flush_vector);
+}
+
+static inline void kvm_riscv_v_exit(void)
+{
+	if (has_vector())
+		kvm_riscv_unregister_vctx_callback();
+}
+
 #else
 
 struct kvm_cpu_context;
@@ -69,6 +85,14 @@ static inline int kvm_riscv_vcpu_alloc_vector_context(struct kvm_vcpu *vcpu)
 static inline void kvm_riscv_vcpu_free_vector_context(struct kvm_vcpu *vcpu)
 {
 }
+
+static inline void kvm_riscv_v_init(void)
+{
+}
+
+static inline void kvm_riscv_v_exit(void)
+{
+}
 #endif
 
 int kvm_riscv_vcpu_get_reg_vector(struct kvm_vcpu *vcpu,
diff --git a/arch/riscv/include/asm/processor.h b/arch/riscv/include/asm/processor.h
index 812517b2cec1..e1b1948f4a4c 100644
--- a/arch/riscv/include/asm/processor.h
+++ b/arch/riscv/include/asm/processor.h
@@ -97,6 +97,7 @@ struct pt_regs;
 
 #define RISCV_V_CTX_UNIT_DEPTH		0x00010000
 #define RISCV_KERNEL_MODE_V		0x00000001
+#define RISCV_V_VCPU_NEED_RESTORE	0x00000002
 #define RISCV_PREEMPT_V			0x00000100
 #define RISCV_PREEMPT_V_DIRTY		0x80000000
 #define RISCV_PREEMPT_V_NEED_RESTORE	0x40000000
diff --git a/arch/riscv/kernel/kernel_mode_vector.c b/arch/riscv/kernel/kernel_mode_vector.c
index 307ac369c3d4..a83d5c26d3b2 100644
--- a/arch/riscv/kernel/kernel_mode_vector.c
+++ b/arch/riscv/kernel/kernel_mode_vector.c
@@ -12,12 +12,31 @@
 #include <linux/types.h>
 
 #include <asm/vector.h>
+#include <asm/kvm_vcpu_vector.h>
 #include <asm/switch_to.h>
 #include <asm/simd.h>
 #ifdef CONFIG_RISCV_ISA_V_PREEMPTIVE
 #include <asm/asm-prototypes.h>
 #endif
 
+static bool (* __rcu kvm_flush_vector_ctx_callback)(void);
+
+void kvm_riscv_register_vctx_callback(bool (*func)(void))
+{
+	if (WARN_ON_ONCE(rcu_access_pointer(kvm_flush_vector_ctx_callback)))
+		return;
+
+	rcu_assign_pointer(kvm_flush_vector_ctx_callback, func);
+}
+EXPORT_SYMBOL_GPL(kvm_riscv_register_vctx_callback);
+
+void kvm_riscv_unregister_vctx_callback(void)
+{
+	rcu_assign_pointer(kvm_flush_vector_ctx_callback, NULL);
+	synchronize_rcu();
+}
+EXPORT_SYMBOL_GPL(kvm_riscv_unregister_vctx_callback);
+
 static inline void riscv_v_flags_set(u32 flags)
 {
 	WRITE_ONCE(current->thread.riscv_v_flags, flags);
@@ -80,6 +99,19 @@ void put_cpu_vector_context(void)
 		preempt_enable();
 }
 
+static void __riscv_flush_vector_context(void)
+{
+	bool (*vcpu_flush_v_callback)(void);
+
+	/* rcu_deference is protected by get/put_cpu_vector_context() */
+	vcpu_flush_v_callback = rcu_dereference(kvm_flush_vector_ctx_callback);
+	if (vcpu_flush_v_callback && vcpu_flush_v_callback())
+		return;
+
+	riscv_v_vstate_save(&current->thread.vstate, task_pt_regs(current));
+	riscv_v_vstate_set_restore(current, task_pt_regs(current));
+}
+
 #ifdef CONFIG_RISCV_ISA_V_PREEMPTIVE
 static __always_inline u32 *riscv_v_flags_ptr(void)
 {
@@ -123,7 +155,7 @@ static int riscv_v_stop_kernel_context(void)
 
 static int riscv_v_start_kernel_context(void)
 {
-	struct __riscv_v_ext_state *kvstate, *uvstate;
+	struct __riscv_v_ext_state *kvstate;
 
 	kvstate = &current->thread.kernel_vstate;
 	if (!kvstate->datap)
@@ -141,13 +173,10 @@ static int riscv_v_start_kernel_context(void)
 	}
 
 	/* Transfer the ownership of V from user to kernel, then save */
-	riscv_v_start(RISCV_PREEMPT_V | RISCV_PREEMPT_V_DIRTY);
-	if (__riscv_v_vstate_check(task_pt_regs(current)->status, DIRTY)) {
-		uvstate = &current->thread.vstate;
-		__riscv_v_vstate_save(uvstate, uvstate->datap);
-	}
-	riscv_preempt_v_clear_dirty(current);
-	riscv_v_vstate_set_restore(current, task_pt_regs(current));
+	get_cpu_vector_context();
+	__riscv_flush_vector_context();
+	put_cpu_vector_context();
+	riscv_v_start(RISCV_PREEMPT_V);
 	return 0;
 }
 
@@ -213,8 +242,7 @@ void kernel_vector_begin(void)
 
 	if (riscv_v_start_kernel_context()) {
 		get_cpu_vector_context();
-		riscv_v_vstate_save(&current->thread.vstate, task_pt_regs(current));
-		riscv_v_vstate_set_restore(current, task_pt_regs(current));
+		__riscv_flush_vector_context();
 	}
 
 	riscv_v_enable();
diff --git a/arch/riscv/kvm/main.c b/arch/riscv/kvm/main.c
index 0924c75100a2..da8a5f6d4b4d 100644
--- a/arch/riscv/kvm/main.c
+++ b/arch/riscv/kvm/main.c
@@ -14,6 +14,7 @@
 #include <asm/kvm_mmu.h>
 #include <asm/kvm_nacl.h>
 #include <asm/sbi.h>
+#include <asm/kvm_vcpu_vector.h>
 
 DEFINE_STATIC_KEY_FALSE(kvm_riscv_vsstage_tlb_no_gpa);
 
@@ -76,6 +77,7 @@ static void kvm_riscv_teardown(void)
 {
 	kvm_riscv_aia_exit();
 	kvm_riscv_nacl_exit();
+	kvm_riscv_v_exit();
 	kvm_unregister_perf_callbacks();
 }
 
@@ -170,6 +172,8 @@ static int __init riscv_kvm_init(void)
 
 	kvm_riscv_setup_vendor_features();
 
+	kvm_riscv_v_init();
+
 	kvm_register_perf_callbacks();
 
 	rc = kvm_init(sizeof(struct kvm_vcpu), 0, THIS_MODULE);
diff --git a/arch/riscv/kvm/vcpu.c b/arch/riscv/kvm/vcpu.c
index cf6e231e76e2..60288bf6e9be 100644
--- a/arch/riscv/kvm/vcpu.c
+++ b/arch/riscv/kvm/vcpu.c
@@ -765,6 +765,12 @@ static void noinstr kvm_riscv_vcpu_enter_exit(struct kvm_vcpu *vcpu,
 	kvm_riscv_vcpu_swap_in_guest_state(vcpu);
 	guest_state_enter_irqoff();
 
+	if (current->thread.riscv_v_flags & RISCV_V_VCPU_NEED_RESTORE) {
+		current->thread.riscv_v_flags &= ~RISCV_V_VCPU_NEED_RESTORE;
+		__kvm_riscv_vector_restore(gcntx);
+		gcntx->sstatus = (gcntx->sstatus & ~SR_VS) | SR_VS_CLEAN;
+	}
+
 	if (kvm_riscv_nacl_sync_sret_available()) {
 		nsh = nacl_shmem();
 
diff --git a/arch/riscv/kvm/vcpu_vector.c b/arch/riscv/kvm/vcpu_vector.c
index 62d2fb77bb9b..3b94e6e7fec8 100644
--- a/arch/riscv/kvm/vcpu_vector.c
+++ b/arch/riscv/kvm/vcpu_vector.c
@@ -95,6 +95,24 @@ void kvm_riscv_vcpu_free_vector_context(struct kvm_vcpu *vcpu)
 	kfree(vcpu->arch.guest_context.vector.datap);
 	kfree(vcpu->arch.host_context.vector.datap);
 }
+
+bool kvm_riscv_vcpu_flush_vector(void)
+{
+	struct kvm_vcpu *vcpu;
+
+	vcpu = *this_cpu_ptr(kvm_get_running_vcpus());
+	if (vcpu && riscv_isa_extension_available(vcpu->arch.isa, v)) {
+		kvm_riscv_vcpu_guest_vector_save(&vcpu->arch.guest_context,
+						 vcpu->arch.isa);
+
+		if (vcpu->arch.guest_context.sstatus != SR_VS_OFF)
+			current->thread.riscv_v_flags |= RISCV_V_VCPU_NEED_RESTORE;
+
+		return true;
+	}
+	return false;
+}
+
 #endif
 
 static int kvm_riscv_vcpu_vreg_addr(struct kvm_vcpu *vcpu,
-- 
2.43.0

Re: [PATCH v1 2/2] RISC-V: KVM: fix vcpu vector context handling for kernel-mode vector

[sashiko-bot]

Thank you for your contribution! Sashiko AI review found 4 potential issue(s) to consider:
- [High] Softirqs preempting `vcpu_load()` or `vcpu_put()` can call `kernel_vector_begin()` and catastrophically corrupt both host user and guest vector states.
- [High] Missing `rcu_read_lock()` in `__riscv_flush_vector_context()` causes a lockdep warning and a Use-After-Free on PREEMPT_RT kernels during KVM module unload.
- [High] Missing bitmask when checking guest `sstatus.VS` field unconditionally forces vector enablement in the guest OS.
- [High] A preemption window allows guest vector state to overwrite the host userspace vector state buffer.
--

--- Patch [2]: [PATCH v1 2/2] RISC-V: KVM: fix vcpu vector context handling for kernel-mode vector ---
commit 1f94d99a9de1c321461c33680cf69b26f4efb606
Author: Andy Chiu <tchiu@tenstorrent.com>
Subject: RISC-V: KVM: fix vcpu vector context handling for kernel-mode vector

This patch registers an RCU-protected callback to allow the kernel-mode
vector code to safely flush the KVM guest's vector context. This addresses a
reverse dependency issue when KVM is configured as a module.

Fixes: ecd2ada8a5e0 ("riscv: Add support for kernel mode vector")

> diff --git a/arch/riscv/kernel/kernel_mode_vector.c b/arch/riscv/kernel/kernel_mode_vector.c
> index 307ac369c3d4a..a83d5c26d3b2c 100644
> --- a/arch/riscv/kernel/kernel_mode_vector.c
> +++ b/arch/riscv/kernel/kernel_mode_vector.c
[ ... ]
> @@ -80,6 +99,19 @@ void put_cpu_vector_context(void)
>  		preempt_enable();
>  }
>  
> +static void __riscv_flush_vector_context(void)
> +{
> +	bool (*vcpu_flush_v_callback)(void);
> +
> +	/* rcu_deference is protected by get/put_cpu_vector_context() */
> +	vcpu_flush_v_callback = rcu_dereference(kvm_flush_vector_ctx_callback);

[Severity: High]
Does this missing rcu_read_lock() trigger a lockdep warning and risk a
use-after-free on PREEMPT_RT?

The caller get_cpu_vector_context() disables softirqs (or preemption on
RT), which does not constitute an RCU read-side critical section.

Furthermore, on PREEMPT_RT, preempt_disable() does not block
synchronize_rcu(). Could kvm_riscv_unregister_vctx_callback() return while
the callback is still executing and prematurely free the memory?

> +	if (vcpu_flush_v_callback && vcpu_flush_v_callback())
> +		return;
> +
> +	riscv_v_vstate_save(&current->thread.vstate, task_pt_regs(current));
> +	riscv_v_vstate_set_restore(current, task_pt_regs(current));
> +}
[ ... ]
> @@ -141,13 +173,10 @@ static int riscv_v_start_kernel_context(void)
>  	}
>  
>  	/* Transfer the ownership of V from user to kernel, then save */
> -	riscv_v_start(RISCV_PREEMPT_V | RISCV_PREEMPT_V_DIRTY);
> -	if (__riscv_v_vstate_check(task_pt_regs(current)->status, DIRTY)) {
> -		uvstate = &current->thread.vstate;
> -		__riscv_v_vstate_save(uvstate, uvstate->datap);
> -	}
> -	riscv_preempt_v_clear_dirty(current);
> -	riscv_v_vstate_set_restore(current, task_pt_regs(current));
> +	get_cpu_vector_context();
> +	__riscv_flush_vector_context();
> +	put_cpu_vector_context();
> +	riscv_v_start(RISCV_PREEMPT_V);

[Severity: High]
Does the gap between put_cpu_vector_context() and riscv_v_start() create
a preemption window that allows guest state to overwrite host user state?

Calling put_cpu_vector_context() enables preemption and bottom halves. If a
context switch occurs exactly here, switch_to() calls __switch_to_vector(),
which observes that riscv_preempt_v_started(prev) is false. 

Since the host's status was not cleared of SR_VS_DIRTY, __switch_to_vector()
might save the physical vector registers (which contain guest state) into
the host QEMU process's prev->thread.vstate.

>  	return 0;
>  }
>  
[ ... ]
> diff --git a/arch/riscv/kvm/vcpu.c b/arch/riscv/kvm/vcpu.c
> index cf6e231e76e25..60288bf6e9be4 100644
> --- a/arch/riscv/kvm/vcpu.c
> +++ b/arch/riscv/kvm/vcpu.c
[ ... ]
> @@ -765,6 +765,12 @@ static void noinstr kvm_riscv_vcpu_enter_exit(struct kvm_vcpu *vcpu,
>  	kvm_riscv_vcpu_swap_in_guest_state(vcpu);
>  	guest_state_enter_irqoff();
>  
> +	if (current->thread.riscv_v_flags & RISCV_V_VCPU_NEED_RESTORE) {
> +		current->thread.riscv_v_flags &= ~RISCV_V_VCPU_NEED_RESTORE;
> +		__kvm_riscv_vector_restore(gcntx);
> +		gcntx->sstatus = (gcntx->sstatus & ~SR_VS) | SR_VS_CLEAN;
> +	}
> +

[Severity: High]
Could this forcefully overwrite the guest's sstatus register based on a false
positive in kvm_riscv_vcpu_flush_vector(), breaking guest OS lazy vector
state?

>  	if (kvm_riscv_nacl_sync_sret_available()) {
>  		nsh = nacl_shmem();
[ ... ]
> diff --git a/arch/riscv/kvm/vcpu_vector.c b/arch/riscv/kvm/vcpu_vector.c
> index 62d2fb77bb9b9..3b94e6e7fec84 100644
> --- a/arch/riscv/kvm/vcpu_vector.c
> +++ b/arch/riscv/kvm/vcpu_vector.c
[ ... ]
> @@ -95,6 +95,24 @@ void kvm_riscv_vcpu_free_vector_context(struct kvm_vcpu *vcpu)
>  	kfree(vcpu->arch.guest_context.vector.datap);
>  	kfree(vcpu->arch.host_context.vector.datap);
>  }
> +
> +bool kvm_riscv_vcpu_flush_vector(void)
> +{
> +	struct kvm_vcpu *vcpu;
> +
> +	vcpu = *this_cpu_ptr(kvm_get_running_vcpus());

[Severity: High]
Can a softirq occurring during vcpu_load() or vcpu_put() mistakenly use this
guest context?

In virt/kvm/kvm_main.c, vcpu_load() assigns the kvm_running_vcpu per-cpu
variable while preemption is disabled via get_cpu(), but hardware interrupts
can still occur. On return from a hardirq, invoke_softirq() could run.

If a softirq executes kernel_vector_begin(), it would find a non-NULL
kvm_running_vcpu here and wrongly assume the physical registers hold guest
state. This could result in saving the host userspace vector state into
guest_context, causing an information leak, and later clobbering the host
context in kvm_arch_vcpu_load().

> +	if (vcpu && riscv_isa_extension_available(vcpu->arch.isa, v)) {
> +		kvm_riscv_vcpu_guest_vector_save(&vcpu->arch.guest_context,
> +						 vcpu->arch.isa);
> +
> +		if (vcpu->arch.guest_context.sstatus != SR_VS_OFF)

[Severity: High]
Does this comparison unconditionally evaluate to true because it is missing a
bitwise AND mask?

Because sstatus usually contains other active flags (like SR_SPP), comparing
the entire 64-bit register against SR_VS_OFF (which is 0) will almost always
evaluate to true. 

Could this use (vcpu->arch.guest_context.sstatus & SR_VS) != SR_VS_OFF
instead?

> +			current->thread.riscv_v_flags |= RISCV_V_VCPU_NEED_RESTORE;
> +
> +		return true;
> +	}
> +	return false;
> +}

-- 
Sashiko AI review · https://sashiko.dev/#/patchset/20260711015835.767259-1-tchiu@tenstorrent.com?part=2