RISC-V: KVM: Fix error code returned for Smstateen ONE_REG

Return -ENOENT for Smstateen ONE_REG when:
1) Smstateen is not enabled for a VCPU
2) ONE_REG id is out of range

This will make Smstateen ONE_REG error codes consistent
with other ONE_REG interfaces of KVM RISC-V.

Fixes: c04913f ("RISCV: KVM: Add sstateen0 to ONE_REG")
Signed-off-by: Anup Patel <anup.patel@oss.qualcomm.com>
Link: https://lore.kernel.org/r/20260120080013.2153519-2-anup.patel@oss.qualcomm.com
Signed-off-by: Anup Patel <anup@brainfault.org>

Author: avpatel

arch/riscv/kvm/vcpu_onereg.c

@@ -565,8 +565,10 @@ static inline int kvm_riscv_vcpu_smstateen_set_csr(struct kvm_vcpu *vcpu,
 	unsigned long regs_max = sizeof(struct kvm_riscv_smstateen_csr) /
 		sizeof(unsigned long);
 
+	if (!riscv_isa_extension_available(vcpu->arch.isa, SMSTATEEN))
+		return -ENOENT;
 	if (reg_num >= regs_max)
-		return -EINVAL;
+		return -ENOENT;
 
 	reg_num = array_index_nospec(reg_num, regs_max);
 
@@ -582,8 +584,10 @@ static int kvm_riscv_vcpu_smstateen_get_csr(struct kvm_vcpu *vcpu,
 	unsigned long regs_max = sizeof(struct kvm_riscv_smstateen_csr) /
 		sizeof(unsigned long);
 
+	if (!riscv_isa_extension_available(vcpu->arch.isa, SMSTATEEN))
+		return -ENOENT;
 	if (reg_num >= regs_max)
-		return -EINVAL;
+		return -ENOENT;
 
 	reg_num = array_index_nospec(reg_num, regs_max);
 
@@ -615,10 +619,7 @@ static int kvm_riscv_vcpu_get_reg_csr(struct kvm_vcpu *vcpu,
 		rc = kvm_riscv_vcpu_aia_get_csr(vcpu, reg_num, &reg_val);
 		break;
 	case KVM_REG_RISCV_CSR_SMSTATEEN:
-		rc = -EINVAL;
-		if (riscv_has_extension_unlikely(RISCV_ISA_EXT_SMSTATEEN))
-			rc = kvm_riscv_vcpu_smstateen_get_csr(vcpu, reg_num,
-							      &reg_val);
+		rc = kvm_riscv_vcpu_smstateen_get_csr(vcpu, reg_num, &reg_val);
 		break;
 	default:
 		rc = -ENOENT;
@@ -660,10 +661,7 @@ static int kvm_riscv_vcpu_set_reg_csr(struct kvm_vcpu *vcpu,
 		rc = kvm_riscv_vcpu_aia_set_csr(vcpu, reg_num, reg_val);
 		break;
 	case KVM_REG_RISCV_CSR_SMSTATEEN:
-		rc = -EINVAL;
-		if (riscv_has_extension_unlikely(RISCV_ISA_EXT_SMSTATEEN))
-			rc = kvm_riscv_vcpu_smstateen_set_csr(vcpu, reg_num,
-							      reg_val);
+		rc = kvm_riscv_vcpu_smstateen_set_csr(vcpu, reg_num, reg_val);
 		break;
 	default:
 		rc = -ENOENT;