enum kvm_pgtable_prot prot);
 int __pkvm_host_unshare_guest(u64 gfn, u64 nr_pages, struct pkvm_hyp_vm *hyp_vm);
 int __pkvm_host_relax_perms_guest(u64 gfn, struct pkvm_hyp_vcpu *vcpu, enum kvm_pgtable_prot prot);
-int __pkvm_host_wrprotect_guest(u64 gfn, struct pkvm_hyp_vm *hyp_vm);
 int __pkvm_host_test_clear_young_guest(u64 gfn, bool mkold, struct pkvm_hyp_vm *vm);
+int __pkvm_host_wrprotect_guest(u64 gfn, u64 nr_pages, struct pkvm_hyp_vm *hyp_vm);
 int __pkvm_host_mkyoung_guest(u64 gfn, struct pkvm_hyp_vcpu *vcpu);
 
 bool addr_is_memory(phys_addr_t phys);
 
 {
        DECLARE_REG(pkvm_handle_t, handle, host_ctxt, 1);
        DECLARE_REG(u64, gfn, host_ctxt, 2);
+       DECLARE_REG(u64, nr_pages, host_ctxt, 3);
        struct pkvm_hyp_vm *hyp_vm;
        int ret = -EINVAL;
 
        if (!hyp_vm)
                goto out;
 
-       ret = __pkvm_host_wrprotect_guest(gfn, hyp_vm);
+       ret = __pkvm_host_wrprotect_guest(gfn, nr_pages, hyp_vm);
        put_pkvm_hyp_vm(hyp_vm);
 out:
        cpu_reg(host_ctxt, 1) = ret;
 
        return ret;
 }
 
-static void assert_host_shared_guest(struct pkvm_hyp_vm *vm, u64 ipa)
+static void assert_host_shared_guest(struct pkvm_hyp_vm *vm, u64 ipa, u64 size)
 {
        u64 phys;
        int ret;
        host_lock_component();
        guest_lock_component(vm);
 
-       ret = __check_host_shared_guest(vm, &phys, ipa, PAGE_SIZE);
+       ret = __check_host_shared_guest(vm, &phys, ipa, size);
 
        guest_unlock_component(vm);
        host_unlock_component();
        if (prot & ~KVM_PGTABLE_PROT_RWX)
                return -EINVAL;
 
-       assert_host_shared_guest(vm, ipa);
+       assert_host_shared_guest(vm, ipa, PAGE_SIZE);
        guest_lock_component(vm);
        ret = kvm_pgtable_stage2_relax_perms(&vm->pgt, ipa, prot, 0);
        guest_unlock_component(vm);
        return ret;
 }
 
-int __pkvm_host_wrprotect_guest(u64 gfn, struct pkvm_hyp_vm *vm)
+int __pkvm_host_wrprotect_guest(u64 gfn, u64 nr_pages, struct pkvm_hyp_vm *vm)
 {
-       u64 ipa = hyp_pfn_to_phys(gfn);
+       u64 size, ipa = hyp_pfn_to_phys(gfn);
        int ret;
 
        if (pkvm_hyp_vm_is_protected(vm))
                return -EPERM;
 
-       assert_host_shared_guest(vm, ipa);
+       ret = __guest_check_transition_size(0, ipa, nr_pages, &size);
+       if (ret)
+               return ret;
+
+       assert_host_shared_guest(vm, ipa, size);
        guest_lock_component(vm);
-       ret = kvm_pgtable_stage2_wrprotect(&vm->pgt, ipa, PAGE_SIZE);
+       ret = kvm_pgtable_stage2_wrprotect(&vm->pgt, ipa, size);
        guest_unlock_component(vm);
 
        return ret;
        if (pkvm_hyp_vm_is_protected(vm))
                return -EPERM;
 
-       assert_host_shared_guest(vm, ipa);
+       assert_host_shared_guest(vm, ipa, PAGE_SIZE);
        guest_lock_component(vm);
        ret = kvm_pgtable_stage2_test_clear_young(&vm->pgt, ipa, PAGE_SIZE, mkold);
        guest_unlock_component(vm);
        if (pkvm_hyp_vm_is_protected(vm))
                return -EPERM;
 
-       assert_host_shared_guest(vm, ipa);
+       assert_host_shared_guest(vm, ipa, PAGE_SIZE);
        guest_lock_component(vm);
        kvm_pgtable_stage2_mkyoung(&vm->pgt, ipa, 0);
        guest_unlock_component(vm);