If the primary plane of a realm wishes to change the access permissions
of memory for the other planes then this causes an exit to the normal
world. KVM then must complete the request using RMI_RTT_SET_S2AP which
may fail if there are missing RTTs. In this case KVM must allocate the
missing RTTs and retry.

Signed-off-by: Steven Price <steven.price@arm.com>
---
 arch/arm64/include/asm/kvm_rme.h |  3 +++
 arch/arm64/kvm/rme-exit.c        | 27 +++++++++++++++++++++++++++
 arch/arm64/kvm/rme.c             | 25 +++++++++++++++++++++----
 3 files changed, 51 insertions(+), 4 deletions(-)

diff --git a/arch/arm64/include/asm/kvm_rme.h b/arch/arm64/include/asm/kvm_rme.h
index e5c0c8274bf8..934b30a8e607 100644
--- a/arch/arm64/include/asm/kvm_rme.h
+++ b/arch/arm64/include/asm/kvm_rme.h
@@ -112,6 +112,9 @@ int kvm_rec_pre_enter(struct kvm_vcpu *vcpu);
 int handle_rec_exit(struct kvm_vcpu *vcpu, int rec_run_status);
 
 int realm_aux_map(struct kvm_vcpu *vcpu, phys_addr_t ipa);
+int kvm_realm_set_s2ap(struct kvm_vcpu *vcpu,
+		       unsigned long start,
+		       unsigned long end);
 
 void kvm_realm_unmap_range(struct kvm *kvm,
 			   unsigned long ipa,
diff --git a/arch/arm64/kvm/rme-exit.c b/arch/arm64/kvm/rme-exit.c
index 04c8af8642af..b7e615f7b3a9 100644
--- a/arch/arm64/kvm/rme-exit.c
+++ b/arch/arm64/kvm/rme-exit.c
@@ -112,6 +112,31 @@ static int rec_exit_ripas_change(struct kvm_vcpu *vcpu)
 	return -EFAULT;
 }
 
+static int rec_exit_s2ap_change(struct kvm_vcpu *vcpu)
+{
+	struct kvm *kvm = vcpu->kvm;
+	struct realm *realm = &kvm->arch.realm;
+	struct realm_rec *rec = &vcpu->arch.rec;
+	unsigned long base = rec->run->exit.s2ap_base;
+	unsigned long top = rec->run->exit.s2ap_top;
+	int ret = -EINVAL;
+
+	if (kvm_realm_is_private_address(realm, base) &&
+	    kvm_realm_is_private_address(realm, top)) {
+		kvm_mmu_topup_memory_cache(&vcpu->arch.mmu_page_cache,
+					   kvm_mmu_cache_min_pages(vcpu->arch.hw_mmu));
+		write_lock(&kvm->mmu_lock);
+		ret = kvm_realm_set_s2ap(vcpu, base, top);
+		write_unlock(&kvm->mmu_lock);
+	}
+
+	WARN_RATELIMIT(ret && ret != -ENOMEM,
+		       "Unable to satisfy SET_S2AP for %#lx - %#lx\n",
+		       base, top);
+
+	return 1;
+}
+
 static int rec_exit_host_call(struct kvm_vcpu *vcpu)
 {
 	int i;
@@ -192,6 +217,8 @@ int handle_rec_exit(struct kvm_vcpu *vcpu, int rec_run_ret)
 		return rec_exit_psci(vcpu);
 	case RMI_EXIT_RIPAS_CHANGE:
 		return rec_exit_ripas_change(vcpu);
+	case RMI_EXIT_S2AP_CHANGE:
+		return rec_exit_s2ap_change(vcpu);
 	case RMI_EXIT_HOST_CALL:
 		return rec_exit_host_call(vcpu);
 	}
diff --git a/arch/arm64/kvm/rme.c b/arch/arm64/kvm/rme.c
index c420546d26f3..fa39a8393d53 100644
--- a/arch/arm64/kvm/rme.c
+++ b/arch/arm64/kvm/rme.c
@@ -1329,6 +1329,7 @@ static int kvm_populate_realm(struct kvm *kvm,
 enum ripas_action {
 	RIPAS_INIT,
 	RIPAS_SET,
+	SET_S2AP,
 };
 
 static int ripas_change(struct kvm *kvm,
@@ -1348,12 +1349,13 @@ static int ripas_change(struct kvm *kvm,
 		rec_phys = virt_to_phys(vcpu->arch.rec.rec_page);
 		memcache = &vcpu->arch.mmu_page_cache;
 
-		WARN_ON(action != RIPAS_SET);
+		WARN_ON(action == RIPAS_INIT);
 	} else {
 		WARN_ON(action != RIPAS_INIT);
 	}
 
 	while (ipa < end) {
+		unsigned long rtt_tree_idx = 0;
 		unsigned long next;
 
 		switch (action) {
@@ -1364,21 +1366,27 @@ static int ripas_change(struct kvm *kvm,
 			ret = rmi_rtt_set_ripas(rd_phys, rec_phys, ipa, end,
 						&next);
 			break;
+		case SET_S2AP:
+			ret = rmi_rtt_set_s2ap(rd_phys, rec_phys, ipa, end,
+					       &next, &rtt_tree_idx);
+			break;
 		}
 
 		switch (RMI_RETURN_STATUS(ret)) {
 		case RMI_SUCCESS:
 			ipa = next;
 			break;
-		case RMI_ERROR_RTT: {
+		case RMI_ERROR_RTT:
+		case RMI_ERROR_RTT_AUX: {
 			int err_level = RMI_RETURN_INDEX(ret);
 			int level = find_map_level(realm, ipa, end);
 
 			if (err_level >= level)
 				return -EINVAL;
 
-			ret = realm_create_rtt_levels(realm, ipa, err_level,
-						      level, memcache);
+			ret = realm_create_rtt_aux_levels(realm, ipa, err_level,
+							  level, rtt_tree_idx,
+							  memcache);
 			if (ret)
 				return ret;
 			/* Retry with the RTT levels in place */
@@ -1396,6 +1404,15 @@ static int ripas_change(struct kvm *kvm,
 	return 0;
 }
 
+int kvm_realm_set_s2ap(struct kvm_vcpu *vcpu,
+		       unsigned long start,
+		       unsigned long end)
+{
+	struct kvm *kvm = vcpu->kvm;
+
+	return ripas_change(kvm, vcpu, start, end, SET_S2AP, NULL);
+}
+
 static int realm_set_ipa_state(struct kvm_vcpu *vcpu,
 			       unsigned long start,
 			       unsigned long end,
-- 
2.43.0