ANDROID: KVM: arm64: mem_protect to use a walker for guest memory transition

In preparation for handling ranges in guest HVCs, add a new walker to
validate guest stage-2 before a transition.

Other than a mismatch with the checked state, this walker will bail-out
in 3 different cases:

  * No mapping has been found (-EINVAL): We need an existing mapping to
    keep track of the memory ownership

  * The range isn't physically contiguous (-E2BIG): We don't want any
    complex data struct to keep track of physical addresses.

  * More than 512 PTEs have been walked (-E2BIG): We don't want to block
    at EL2 for too long.

A mask (desired_mask) enables what bits of the state are actually
meaningful for the walk. This will later be used for more complex
checks, where several states combinations are allowed.

Bug: 278749606
Bug: 243642516
Change-Id: Ib88725274f1912a9a3bcd5d9c1b5ec2ea4200800
Signed-off-by: Vincent Donnefort <vdonnefort@google.com>

arch/arm64/kvm/hyp/nvhe/mem_protect.c

@@ -1338,64 +1338,96 @@ err_undo_psci:
 	return err;
 }
 
-static int __guest_get_completer_addr(u64 *completer_addr, phys_addr_t phys,
+struct guest_request_walker_data {
-				      const struct pkvm_mem_transition *tx)
+	unsigned long		ipa_start;
+	phys_addr_t		phys_start;
+	u64			size;
+	enum pkvm_page_state	desired_state;
+	enum pkvm_page_state	desired_mask;
+	int			max_ptes;
+};
+
+#define GUEST_WALKER_DATA_INIT(__state)			\
+{						 	\
+	.size		= 0,				\
+	.desired_state	= __state,			\
+	.desired_mask	= ~0,				\
+	/*						\
+	 * Arbitrary limit of walked PTEs to restrict	\
+	 * the time spent at EL2			\
+	 */						\
+	.max_ptes	= 512,				\
+}
+
+static int guest_request_walker(const struct kvm_pgtable_visit_ctx *ctx,
+				enum kvm_pgtable_walk_flags visit)
 {
-	switch (tx->completer.id) {
+	struct guest_request_walker_data *data = (struct guest_request_walker_data *)ctx->arg;
-	case PKVM_ID_HOST:
+	enum pkvm_page_state state;
-		*completer_addr = phys;
+	kvm_pte_t pte = *ctx->ptep;
-		break;
+	u32 level = ctx->level;
-	case PKVM_ID_HYP:
+	phys_addr_t phys;
-		*completer_addr = (u64)__hyp_va(phys);
+
-		break;
+	state = guest_get_page_state(pte, 0);
-	default:
+	if ((data->desired_state & data->desired_mask) != state)
-		return -EINVAL;
+		return (state & PKVM_NOPAGE) ? -EFAULT : -EINVAL;
+
+	if (state & PKVM_NOPAGE) {
+		phys = PHYS_ADDR_MAX;
+	} else {
+		phys = kvm_pte_to_phys(pte);
+		if (!addr_is_allowed_memory(phys))
+			return -EINVAL;
 	}
 
-	return 0;
+	data->max_ptes--;
+
+	if (!data->size) {
+		data->phys_start = phys;
+		data->size = kvm_granule_size(level);
+		data->ipa_start = ctx->addr & ~(kvm_granule_size(level) - 1);
+		goto end;
+	}
+
+	/* Can only describe physically contiguous mappings */
+	if ((data->phys_start != PHYS_ADDR_MAX) &&
+	    (phys != data->phys_start + data->size))
+		return -E2BIG;
+
+	data->size += kvm_granule_size(level);
+end:
+	return data->max_ptes > 0 ? 0 : -E2BIG;
 }
 
 static int __guest_request_page_transition(struct pkvm_checked_mem_transition *checked_tx,
 					   enum pkvm_page_state desired)
 {
+	struct guest_request_walker_data data = GUEST_WALKER_DATA_INIT(desired);
 	const struct pkvm_mem_transition *tx = checked_tx->tx;
 	struct pkvm_hyp_vm *vm = tx->initiator.guest.hyp_vm;
-	enum pkvm_page_state state;
+	struct kvm_pgtable_walker walker = {
-	phys_addr_t phys;
+		.cb     = guest_request_walker,
-	kvm_pte_t pte;
+		.flags  = KVM_PGTABLE_WALK_LEAF,
-	u32 level;
+		.arg    = (void *)&data,
+	};
 	int ret;
 
-	if (tx->nr_pages != 1)
+	ret = kvm_pgtable_walk(&vm->pgt, tx->initiator.addr,
-		return -E2BIG;
+			       tx->nr_pages * PAGE_SIZE, &walker);
-
+	/* Walker reached data.max_ptes or a non physically contiguous block */
-	ret = kvm_pgtable_get_leaf(&vm->pgt, tx->initiator.addr, &pte, &level);
+	if (ret == -E2BIG)
-	if (ret)
+		ret = 0;
+	else if (ret)
 		return ret;
 
-	state = guest_get_page_state(pte, tx->initiator.addr);
+	/* Not aligned with a block mapping */
-	if (state == PKVM_NOPAGE)
+	if (data.ipa_start != tx->initiator.addr)
-		return -EFAULT;
-
-	if (state != desired)
-		return -EPERM;
-
-	/*
-	 * We only deal with page granular mappings in the guest for now as
-	 * the pgtable code relies on being able to recreate page mappings
-	 * lazily after zapping a block mapping, which doesn't work once the
-	 * pages have been donated.
-	 */
-	if (level != KVM_PGTABLE_MAX_LEVELS - 1)
 		return -EINVAL;
 
-	phys = kvm_pte_to_phys(pte);
+	checked_tx->completer_addr = data.phys_start;
-	if (!addr_is_allowed_memory(phys))
+	checked_tx->nr_pages = min_t(u64, data.size >> PAGE_SHIFT, tx->nr_pages);
-		return -EINVAL;
 
-	checked_tx->nr_pages = tx->nr_pages;
+	return 0;
-
-	return __guest_get_completer_addr(&checked_tx->completer_addr, phys, tx);
 }
 
 static int guest_request_share(struct pkvm_checked_mem_transition *checked_tx)
@@ -1446,6 +1478,9 @@ static int check_share(struct pkvm_checked_mem_transition *checked_tx)
 	const struct pkvm_mem_transition *tx = checked_tx->tx;
 	int ret;
 
+	if (!tx->nr_pages)
+		return -EINVAL;
+
 	switch (tx->initiator.id) {
 	case PKVM_ID_HOST:
 		ret = host_request_owned_transition(&checked_tx->completer_addr, tx);
@@ -1564,6 +1599,9 @@ static int check_unshare(struct pkvm_checked_mem_transition *checked_tx)
 	const struct pkvm_mem_transition *tx = checked_tx->tx;
 	int ret;
 
+	if (!tx->nr_pages)
+		return -EINVAL;
+
 	switch (tx->initiator.id) {
 	case PKVM_ID_HOST:
 		ret = host_request_unshare(checked_tx);