x86/vmware: Share steal-time counters in encrypted guests

index 34b73573b108..7b7469f2e7d6 100644
@@ -25,9 +25,12 @@
 #include <linux/init.h>
 #include <linux/export.h>
 #include <linux/clocksource.h>
+#include <linux/cc_platform.h>
 #include <linux/cpu.h>
 #include <linux/efi.h>
+#include <linux/mm.h>
 #include <linux/reboot.h>
+#include <linux/set_memory.h>
 #include <linux/static_call.h>
 #include <linux/sched/cputime.h>
 #include <asm/div64.h>

@@ -145,8 +148,8 @@ static unsigned long vmware_get_tsc_khz(void)
 #ifdef CONFIG_PARAVIRT
 static struct cyc2ns_data vmware_cyc2ns __ro_after_init;
 static bool vmw_sched_clock __initdata = true;
-static DEFINE_PER_CPU_DECRYPTED(struct vmware_steal_time, vmw_steal_time) __aligned(64);
-static bool has_steal_clock;
+static struct vmware_steal_time *vmw_steal_time __ro_after_init;
+static bool has_steal_clock __ro_after_init;
 static bool steal_acc __initdata = true; /* steal time accounting */
 
 static __init int setup_vmw_sched_clock(char *s)

@@ -228,7 +231,7 @@ static bool vmware_is_stealclock_available(void)
  */
 static u64 vmware_steal_clock(int cpu)
 {
-	struct vmware_steal_time *steal = &per_cpu(vmw_steal_time, cpu);
+	struct vmware_steal_time *steal = &vmw_steal_time[cpu];
 	u64 clock;
 
 	if (IS_ENABLED(CONFIG_64BIT))

@@ -253,21 +256,19 @@ static u64 vmware_steal_clock(int cpu)
 			     vmware_cyc2ns.cyc2ns_shift);
 }
 
-static void vmware_register_steal_time(void)
+static bool vmware_register_steal_time(void)
 {
 	int cpu = smp_processor_id();
-	struct vmware_steal_time *st = &per_cpu(vmw_steal_time, cpu);
+	struct vmware_steal_time *st = &vmw_steal_time[cpu];
 
-	if (!has_steal_clock)
-		return;
-
-	if (!stealclock_enable(slow_virt_to_phys(st))) {
-		has_steal_clock = false;
-		return;
+	if (!stealclock_enable(virt_to_phys(st))) {
+		pr_warn("failed to register steal time for CPU %d\n", cpu);
+		return false;
 	}
 
 	pr_info("vmware-stealtime: cpu %d, pa %llx\n",
-		cpu, (unsigned long long) slow_virt_to_phys(st));
+		cpu, (unsigned long long)virt_to_phys(st));
+	return true;
 }
 
 static void vmware_disable_steal_time(void)

@@ -278,12 +279,6 @@ static void vmware_disable_steal_time(void)
 	stealclock_disable();
 }
 
-static void vmware_guest_cpu_init(void)
-{
-	if (has_steal_clock)
-		vmware_register_steal_time();
-}
-
 static void vmware_pv_guest_cpu_reboot(void *unused)
 {
 	vmware_disable_steal_time();

@@ -302,16 +297,11 @@ static struct notifier_block vmware_pv_reboot_nb = {
 };
 
 #ifdef CONFIG_SMP
-static void __init vmware_smp_prepare_boot_cpu(void)
-{
-	vmware_guest_cpu_init();
-	native_smp_prepare_boot_cpu();
-}
-
 static int vmware_cpu_online(unsigned int cpu)
 {
 	local_irq_disable();
-	vmware_guest_cpu_init();
+	if (has_steal_clock)
+		vmware_register_steal_time();
 	local_irq_enable();
 	return 0;
 }

@@ -325,6 +315,76 @@ static int vmware_cpu_down_prepare(unsigned int cpu)
 }
 #endif
 
+static int __init vmware_steal_time_init(void)
+{
+	size_t size = PAGE_ALIGN(nr_cpu_ids * sizeof(*vmw_steal_time));
+	bool encrypted = cc_platform_has(CC_ATTR_GUEST_MEM_ENCRYPT);
+	unsigned long addr, flags;
+	int ret;
+#ifdef CONFIG_SMP
+	int hp_state;
+#endif
+
+	if (!has_steal_clock)
+		return 0;
+
+	vmw_steal_time = alloc_pages_exact(size, GFP_KERNEL | __GFP_ZERO);
+	if (!vmw_steal_time) {
+		ret = -ENOMEM;
+		goto disable;
+	}
+	addr = (unsigned long)vmw_steal_time;
+
+	ret = register_reboot_notifier(&vmware_pv_reboot_nb);
+	if (ret)
+		goto free;
+#ifdef CONFIG_SMP
+	hp_state = cpuhp_setup_state_nocalls(CPUHP_AP_ONLINE_DYN,
+					     "x86/vmware:online",
+					     vmware_cpu_online,
+					     vmware_cpu_down_prepare);
+	if (hp_state < 0) {
+		ret = hp_state;
+		unregister_reboot_notifier(&vmware_pv_reboot_nb);
+		goto free;
+	}
+#endif
+
+	if (encrypted) {
+		ret = set_memory_decrypted(addr, size >> PAGE_SHIFT);
+		if (ret)
+			goto remove;
+		/* Conversion need not preserve the zeroes. */
+		memset(vmw_steal_time, 0, size);
+	}
+
+	local_irq_save(flags);
+	ret = vmware_register_steal_time() ? 0 : -EIO;
+	local_irq_restore(flags);
+	if (!ret) {
+		static_call_update(pv_steal_clock, vmware_steal_clock);
+		return 0;
+	}
+
+	if (!encrypted || !set_memory_encrypted(addr, size >> PAGE_SHIFT))
+		free_pages_exact(vmw_steal_time, size);
+remove:
+	/* Keep pages with an unknown encryption state out of the allocator. */
+#ifdef CONFIG_SMP
+	cpuhp_remove_state_nocalls(hp_state);
+#endif
+	unregister_reboot_notifier(&vmware_pv_reboot_nb);
+	goto disable;
+free:
+	free_pages_exact(vmw_steal_time, size);
+disable:
+	vmw_steal_time = NULL;
+	has_steal_clock = false;
+	pr_warn("steal time disabled: %d\n", ret);
+	return 0;
+}
+early_initcall(vmware_steal_time_init);
+
 static __init int activate_jump_labels(void)
 {
 	if (has_steal_clock) {

@@ -350,25 +410,7 @@ static void __init vmware_paravirt_ops_setup(void)
 	if (vmw_sched_clock)
 		paravirt_set_sched_clock(vmware_sched_clock);
 
-	if (vmware_is_stealclock_available()) {
-		has_steal_clock = true;
-		static_call_update(pv_steal_clock, vmware_steal_clock);
-
-		/* We use reboot notifier only to disable steal clock */
-		register_reboot_notifier(&vmware_pv_reboot_nb);
-
-#ifdef CONFIG_SMP
-		smp_ops.smp_prepare_boot_cpu =
-			vmware_smp_prepare_boot_cpu;
-		if (cpuhp_setup_state_nocalls(CPUHP_AP_ONLINE_DYN,
-					      "x86/vmware:online",
-					      vmware_cpu_online,
-					      vmware_cpu_down_prepare) < 0)
-			pr_err("vmware_guest: Failed to install cpu hotplug callbacks\n");
-#else
-		vmware_guest_cpu_init();
-#endif
-	}
+	has_steal_clock = vmware_is_stealclock_available();
 }
 #else
 #define vmware_paravirt_ops_setup() do {} while (0)