[RFC 5/6] KVM: arm64: Conditionally expose ID_AA64MMFR4_EL1.TLBID to guest

[Jinqian Yang]

The TLBID feature bit in ID_AA64MMFR4_EL1 is only advertised to the
guest after userspace configures vDomains via KVM_ARM_VCPU_SET_VDOMAIN,
which sets the guest_tlbid_enabled flag. Until then the bit is masked
to prevent the guest from enabling TLBID without a valid vDomain
affinity.

Exposing the feature would otherwise let the guest read the physical
TLBIDIDR_EL1 directly, so keep that register trapped in the guest
FGT context and served by the emulation added earlier. The host's
own register is never affected.

Expose KVM_REG_ARM_ID_AA64MMFR4_EL1 in the arm64 UAPI header so
userspace can read the sanitised value via KVM_GET_ONE_REG.

Signed-off-by: Jinqian Yang <yangjinqian1 at huawei.com>
Signed-off-by: Tian Zheng <zhengtian10 at huawei.com>
---
 arch/arm64/include/uapi/asm/kvm.h |  2 ++
 arch/arm64/kvm/arm.c              |  2 ++
 arch/arm64/kvm/config.c           | 11 +++++++
 arch/arm64/kvm/sys_regs.c         | 49 +++++++++++++++++++++++++++++--
 arch/arm64/kvm/sys_regs.h         |  6 ++++
 5 files changed, 68 insertions(+), 2 deletions(-)

diff --git a/arch/arm64/include/uapi/asm/kvm.h b/arch/arm64/include/uapi/asm/kvm.h
index 019e5e3d892e..e54129ed431d 100644
--- a/arch/arm64/include/uapi/asm/kvm.h
+++ b/arch/arm64/include/uapi/asm/kvm.h
@@ -255,6 +255,8 @@ struct kvm_arm_counter_offset {
 
 #define ARM64_SYS_REG(...) (__ARM64_SYS_REG(__VA_ARGS__) | KVM_REG_SIZE_U64)
 
+#define KVM_REG_ARM_ID_AA64MMFR4_EL1 ARM64_SYS_REG(3, 0, 0, 7, 4)
+
 /* Physical Timer EL0 Registers */
 #define KVM_REG_ARM_PTIMER_CTL		ARM64_SYS_REG(3, 3, 14, 2, 1)
 #define KVM_REG_ARM_PTIMER_CVAL		ARM64_SYS_REG(3, 3, 14, 2, 2)
diff --git a/arch/arm64/kvm/arm.c b/arch/arm64/kvm/arm.c
index ddf7c25f9ed9..b5ab89db1c07 100644
--- a/arch/arm64/kvm/arm.c
+++ b/arch/arm64/kvm/arm.c
@@ -2313,6 +2313,8 @@ static int kvm_arm_vcpu_set_vdomain(struct kvm_vcpu *vcpu,
 
 	kvm_info("vCPU%d: vDomain bitmap: 0x%x\n", vcpu->vcpu_idx, vdomain_bitmap);
 
+	kvm_update_aa64mmfr4_tlbid(vcpu->kvm);
+
 	return 0;
 }
 #else
diff --git a/arch/arm64/kvm/config.c b/arch/arm64/kvm/config.c
index 6c116aa1864a..d0402151a752 100644
--- a/arch/arm64/kvm/config.c
+++ b/arch/arm64/kvm/config.c
@@ -1536,6 +1536,17 @@ void compute_fgu(struct kvm *kvm, enum fgt_group_id fgt)
 	case HFGRTR2_GROUP:
 		val |= compute_fgu_bits(kvm, &hfgrtr2_desc);
 		val |= compute_fgu_bits(kvm, &hfgwtr2_desc);
+#ifdef CONFIG_KVM_ARM_VTLBID
+		/*
+		 * TLBIDIDR_EL1 must stay trapped even when TLBID is
+		 * exposed to the guest, so it can be emulated with the
+		 * synthesized virtual domain counts. This only affects
+		 * the guest FGT context; the host's own HFGRTR2_EL2 is
+		 * never modified.
+		 */
+		if (kvm->arch.tlbidomain.kvm_tlbid_enabled)
+			val |= HFGRTR2_EL2_nTLBIDIDR_EL1;
+#endif
 		break;
 	case HFGITR2_GROUP:
 		val |= compute_fgu_bits(kvm, &hfgitr2_desc);
diff --git a/arch/arm64/kvm/sys_regs.c b/arch/arm64/kvm/sys_regs.c
index cf9804e54ff1..a1ca7309406e 100644
--- a/arch/arm64/kvm/sys_regs.c
+++ b/arch/arm64/kvm/sys_regs.c
@@ -2012,9 +2012,17 @@ static u64 __kvm_read_sanitised_id_reg(const struct kvm_vcpu *vcpu,
 		if (!system_supports_poe())
 			val &= ~ID_AA64MMFR3_EL1_S1POE;
 		break;
-	case SYS_ID_AA64MMFR4_EL1:
+	case SYS_ID_AA64MMFR4_EL1: {
+#ifdef CONFIG_KVM_ARM_VTLBID
+		struct tlbidomain *vdomain = &vcpu->kvm->arch.tlbidomain;
+
+		if (!vdomain->guest_tlbid_enabled)
+			val &= ~ID_AA64MMFR4_EL1_TLBID;
+#else
 		val &= ~ID_AA64MMFR4_EL1_TLBID;
+#endif
 		break;
+	}
 	case SYS_ID_MMFR4_EL1:
 		val &= ~ID_MMFR4_EL1_CCIDX;
 		break;
@@ -3494,7 +3502,8 @@ static const struct sys_reg_desc sys_reg_descs[] = {
 				       ID_AA64MMFR3_EL1_SCTLRX	|
 				       ID_AA64MMFR3_EL1_S1PIE   |
 				       ID_AA64MMFR3_EL1_S1POE)),
-	ID_WRITABLE(ID_AA64MMFR4_EL1, ID_AA64MMFR4_EL1_NV_frac),
+	ID_WRITABLE(ID_AA64MMFR4_EL1, ID_AA64MMFR4_EL1_NV_frac |
+				    ID_AA64MMFR4_EL1_TLBID_MASK),
 	ID_UNALLOCATED(7,5),
 	ID_UNALLOCATED(7,6),
 	ID_UNALLOCATED(7,7),
@@ -5495,6 +5504,42 @@ static bool index_to_params(u64 id, struct sys_reg_params *params)
 	}
 }
 
+/*
+ * Update the VM's stored value for an ID register.
+ */
+#ifdef CONFIG_KVM_ARM_VTLBID
+void kvm_update_aa64mmfr4_tlbid(struct kvm *kvm)
+{
+	struct tlbidomain *vdomain = &kvm->arch.tlbidomain;
+	u64 val, tlbid_val;
+
+	mutex_lock(&kvm->arch.config_lock);
+
+	if (kvm_vm_has_ran_once(kvm)) {
+		/*
+		 * ID_AA64MMFR4 is a VM-level register already set during
+		 * the first vCPU init. Skip the write to avoid triggering
+		 * KVM_BUG_ON in kvm_set_vm_id_reg(). The per-vCPU
+		 * vdomain_bitmap is still updated by the caller.
+		 */
+		mutex_unlock(&kvm->arch.config_lock);
+		return;
+	}
+
+	val = kvm_read_vm_id_reg(kvm, SYS_ID_AA64MMFR4_EL1);
+	tlbid_val = read_sanitised_ftr_reg(SYS_ID_AA64MMFR4_EL1);
+	tlbid_val &= ID_AA64MMFR4_EL1_TLBID_MASK;
+
+	val &= ~ID_AA64MMFR4_EL1_TLBID_MASK;
+	if (vdomain->guest_tlbid_enabled)
+		val |= tlbid_val;
+
+	kvm_set_vm_id_reg(kvm, SYS_ID_AA64MMFR4_EL1, val);
+
+	mutex_unlock(&kvm->arch.config_lock);
+}
+#endif
+
 const struct sys_reg_desc *get_reg_by_id(u64 id,
 					 const struct sys_reg_desc table[],
 					 unsigned int num)
diff --git a/arch/arm64/kvm/sys_regs.h b/arch/arm64/kvm/sys_regs.h
index bd56a45abbf9..0b5487947790 100644
--- a/arch/arm64/kvm/sys_regs.h
+++ b/arch/arm64/kvm/sys_regs.h
@@ -252,6 +252,12 @@ int kvm_sys_reg_get_user(struct kvm_vcpu *vcpu, const struct kvm_one_reg *reg,
 int kvm_sys_reg_set_user(struct kvm_vcpu *vcpu, const struct kvm_one_reg *reg,
 			 const struct sys_reg_desc table[], unsigned int num);
 
+#ifdef CONFIG_KVM_ARM_VTLBID
+void kvm_update_aa64mmfr4_tlbid(struct kvm *kvm);
+#else
+static inline void kvm_update_aa64mmfr4_tlbid(struct kvm *kvm) { }
+#endif
+
 bool triage_sysreg_trap(struct kvm_vcpu *vcpu, int *sr_index);
 
 int kvm_finalize_sys_regs(struct kvm_vcpu *vcpu);
-- 
2.33.0