Now that the driver tells us whether TRBE was used or not we can use
that. Except in pKVM where the host isn't trusted we keep the existing
feature + sysreg check.

Now in the normal nVHE case, TRBE save and restore are gated by flag
checks on kvm_host_data.

Instead of using a magic value of host_debug_state.trfcr_el1 to
determine whether to restore, add a flag. This will also simplify the
logic in the next commit where restoration but no disabling is required.

Signed-off-by: James Clark <james.clark@linaro.org>
---
 arch/arm64/include/asm/kvm_host.h  |  2 ++
 arch/arm64/kvm/hyp/nvhe/debug-sr.c | 51 +++++++++++++++++++++++-------
 2 files changed, 41 insertions(+), 12 deletions(-)

diff --git a/arch/arm64/include/asm/kvm_host.h b/arch/arm64/include/asm/kvm_host.h
index b1dccac996a6..a8846689512b 100644
--- a/arch/arm64/include/asm/kvm_host.h
+++ b/arch/arm64/include/asm/kvm_host.h
@@ -947,6 +947,8 @@ struct kvm_vcpu_arch {
 #define HOST_STATE_SPE_EN	__kvm_single_flag(state, BIT(0))
 /* TRBLIMITR_EL1_E is set (TRBE trace buffer enabled) */
 #define HOST_STATE_TRBE_EN	__kvm_single_flag(state, BIT(1))
+/* Hyp modified TRFCR */
+#define HOST_STATE_RESTORE_TRFCR __kvm_single_flag(state, BIT(2))
 
 /* Pointer to the vcpu's SVE FFR for sve_{save,load}_state() */
 #define vcpu_sve_pffr(vcpu) (kern_hyp_va((vcpu)->arch.sve_state) +	\
diff --git a/arch/arm64/kvm/hyp/nvhe/debug-sr.c b/arch/arm64/kvm/hyp/nvhe/debug-sr.c
index 578c549af3c6..17c23e52f5f4 100644
--- a/arch/arm64/kvm/hyp/nvhe/debug-sr.c
+++ b/arch/arm64/kvm/hyp/nvhe/debug-sr.c
@@ -63,32 +63,55 @@ static void __debug_restore_spe(void)
 	*host_data_ptr(host_debug_state.pmscr_el1) = 0;
 }
 
-static void __debug_save_trace(u64 *trfcr_el1)
+static bool __debug_should_save_trace(void)
 {
-	*trfcr_el1 = 0;
+	/* pKVM reads the state for itself rather than trusting the host */
+	if (unlikely(is_protected_kvm_enabled())) {
+		/* Always disable any trace regardless of TRBE */
+		if (read_sysreg_el1(SYS_TRFCR) &
+		    (TRFCR_ELx_E0TRE | TRFCR_ELx_ExTRE))
+			return true;
+
+		/*
+		 * Trace could already be disabled but TRBE buffer
+		 * might still need to be drained if it was in use.
+		 */
+		if (host_data_get_flag(HOST_FEAT_HAS_TRBE))
+			return read_sysreg_s(SYS_TRBLIMITR_EL1) &
+			       TRBLIMITR_EL1_E;
+	}
+
+	return host_data_get_flag(HOST_STATE_TRBE_EN);
+}
 
-	/* Check if the TRBE is enabled */
-	if (!(read_sysreg_s(SYS_TRBLIMITR_EL1) & TRBLIMITR_EL1_E))
-		return;
+static void __debug_save_trace(void)
+{
 	/*
 	 * Prohibit trace generation while we are in guest.
 	 * Since access to TRFCR_EL1 is trapped, the guest can't
 	 * modify the filtering set by the host.
 	 */
-	*trfcr_el1 = read_sysreg_el1(SYS_TRFCR);
+	*host_data_ptr(host_debug_state.trfcr_el1) = read_sysreg_el1(SYS_TRFCR);
 	write_sysreg_el1(0, SYS_TRFCR);
 	isb();
 	/* Drain the trace buffer to memory */
 	tsb_csync();
+
+	host_data_set_flag(HOST_STATE_RESTORE_TRFCR);
 }
 
-static void __debug_restore_trace(u64 trfcr_el1)
+static void __debug_restore_trace(void)
 {
-	if (!trfcr_el1)
+	u64 trfcr_el1;
+
+	if (!host_data_get_flag(HOST_STATE_RESTORE_TRFCR))
 		return;
 
 	/* Restore trace filter controls */
+	trfcr_el1 = *host_data_ptr(host_debug_state.trfcr_el1);
+	*host_data_ptr(host_debug_state.trfcr_el1) = read_sysreg_el1(SYS_TRFCR);
 	write_sysreg_el1(trfcr_el1, SYS_TRFCR);
+	host_data_clear_flag(HOST_STATE_RESTORE_TRFCR);
 }
 
 void __debug_save_host_buffers_nvhe(void)
@@ -97,9 +120,14 @@ void __debug_save_host_buffers_nvhe(void)
 	if (__debug_spe_enabled())
 		__debug_save_spe();
 
+	/* Any trace filtering requires TRFCR register */
+	if (!host_data_get_flag(HOST_FEAT_HAS_TRF))
+		return;
+
 	/* Disable and flush Self-Hosted Trace generation */
-	if (host_data_get_flag(HOST_FEAT_HAS_TRBE))
-		__debug_save_trace(host_data_ptr(host_debug_state.trfcr_el1));
+	if (__debug_should_save_trace())
+		__debug_save_trace();
+
 }
 
 void __debug_switch_to_guest(struct kvm_vcpu *vcpu)
@@ -110,8 +138,7 @@ void __debug_switch_to_guest(struct kvm_vcpu *vcpu)
 void __debug_restore_host_buffers_nvhe(struct kvm_vcpu *vcpu)
 {
 	__debug_restore_spe();
-	if (host_data_get_flag(HOST_FEAT_HAS_TRBE))
-		__debug_restore_trace(*host_data_ptr(host_debug_state.trfcr_el1));
+	__debug_restore_trace();
 }
 
 void __debug_switch_to_host(struct kvm_vcpu *vcpu)
-- 
2.34.1