KVM: arm64: pvtime: Don't lose stolen time on failed updates

index 4ceabaa4c30b..9b4f0645d92b 100644
@@ -17,6 +17,7 @@ void kvm_update_stolen_time(struct kvm_vcpu *vcpu)
 	u64 last_steal = vcpu->arch.steal.last_steal;
 	u64 offset = offsetof(struct pvclock_vcpu_stolen_time, stolen_time);
 	u64 steal = 0;
+	u64 run_delay;
 	int idx;
 
 	if (base == INVALID_GPA)

@@ -25,9 +26,10 @@ void kvm_update_stolen_time(struct kvm_vcpu *vcpu)
 	idx = srcu_read_lock(&kvm->srcu);
 	if (!kvm_get_guest(kvm, base + offset, steal)) {
 		steal = le64_to_cpu(steal);
-		vcpu->arch.steal.last_steal = READ_ONCE(current->sched_info.run_delay);
-		steal += vcpu->arch.steal.last_steal - last_steal;
-		kvm_put_guest(kvm, base + offset, cpu_to_le64(steal));
+		run_delay = READ_ONCE(current->sched_info.run_delay);
+		steal += run_delay - last_steal;
+		if (!kvm_put_guest(kvm, base + offset, cpu_to_le64(steal)))
+			vcpu->arch.steal.last_steal = run_delay;
 	}
 	srcu_read_unlock(&kvm->srcu, idx);
 }

index bc3c62b72c58..fbdd55fa5ab6 100644
@@ -8,6 +8,8 @@
 #include <time.h>
 #include <sched.h>
 #include <pthread.h>
+#include <sys/mman.h>
+#include <unistd.h>
 #include <linux/kernel.h>
 #include <asm/kvm.h>
 #ifdef __riscv

@@ -162,6 +164,15 @@ static void guest_code(int cpu)
 
 	check_status(st);
 	WRITE_ONCE(guest_stolen_time[cpu], st->st_time);
+	if (!cpu) {
+		/*
+		 * Let userspace trigger a failed stolen-time update, then retry
+		 * the update after restoring write access.
+		 */
+		GUEST_SYNC(2);
+		GUEST_SYNC(3);
+		WRITE_ONCE(guest_stolen_time[cpu], st->st_time);
+	}
 	GUEST_DONE();
 }
 

@@ -469,9 +480,10 @@ static void check_steal_time_uapi(void)
 static void *do_steal_time(void *arg)
 {
 	struct timespec ts, stop;
+	unsigned long duration = arg ? *(unsigned long *)arg : MIN_RUN_DELAY_NS;
 
 	clock_gettime(CLOCK_MONOTONIC, &ts);
-	stop = timespec_add_ns(ts, MIN_RUN_DELAY_NS);
+	stop = timespec_add_ns(ts, duration);
 
 	while (1) {
 		clock_gettime(CLOCK_MONOTONIC, &ts);

@@ -500,6 +512,45 @@ static void run_vcpu(struct kvm_vcpu *vcpu)
 	}
 }
 
+#ifdef __aarch64__
+static void test_failed_stolen_time_update(struct kvm_vcpu *vcpu)
+{
+	struct kvm_vm *vm = vcpu->vm;
+	struct st_time *st = addr_gva2hva(vm, ST_GPA_BASE);
+	unsigned long duration = 100 * 1000 * 1000UL;
+	u64 stolen_time = st->st_time;
+	long run_delay;
+	pthread_t thread;
+
+	/*
+	 * Keep the page readable so that kvm_get_guest() succeeds, but make
+	 * kvm_put_guest() fail when KVM attempts to update stolen time.
+	 */
+	TEST_ASSERT(!mprotect(st, getpagesize(), PROT_READ),
+		    "Failed to make stolen time page read-only");
+
+	run_delay = get_run_delay();
+	kvm_pthread_create(&thread, NULL, do_steal_time, &duration);
+	do
+		sched_yield();
+	while (get_run_delay() - run_delay < duration / 5);
+	kvm_pthread_join(thread, NULL);
+	run_delay = get_run_delay() - run_delay;
+
+	/* The update fails, but the guest must still be able to run. */
+	run_vcpu(vcpu);
+	TEST_ASSERT(st->st_time == stolen_time,
+		    "Stolen time changed on a read-only page");
+
+	TEST_ASSERT(!mprotect(st, getpagesize(), PROT_READ | PROT_WRITE),
+		    "Failed to restore stolen time page write access");
+	run_vcpu(vcpu);
+	TEST_ASSERT(st->st_time - stolen_time >= run_delay,
+		    "Lost stolen time after a failed update: expected >= %ld, got %lu",
+		    run_delay, (unsigned long)(st->st_time - stolen_time));
+}
+#endif
+
 int main(int ac, char **av)
 {
 	struct kvm_vcpu *vcpus[NR_VCPUS];

@@ -569,6 +620,11 @@ int main(int ac, char **av)
 			    "Expected stolen time >= %ld, got %ld",
 			    run_delay, stolen_time);
 
+#ifdef __aarch64__
+		if (!i)
+			test_failed_stolen_time_update(vcpus[i]);
+#endif
+
 		if (verbose) {
 			ksft_print_msg("VCPU%d: total-stolen-time=%ld test-stolen-time=%ld%s\n",
 				       i, guest_stolen_time[i], stolen_time,