Allow kvm_gmem_get_pfn to return with the folio lock held by adding a
KVM_GMEM_GET_PFN_LOCKED option to `flags`.

When accessing the content of gmem folios, the lock must be held until
kvm_gmem_put_pfn, to avoid concurrent direct map modifications of the
same folio causing use-after-free-like problems. However,
kvm_gmem_get_pfn so far unconditionally drops the folio lock, making it
currently impossible to use the KVM_GMEM_GET_PFN_SHARED flag safely.

Signed-off-by: Patrick Roy <roypat@amazon.co.uk>
---
 include/linux/kvm_host.h | 1 +
 virt/kvm/guest_memfd.c   | 5 +++--
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/include/linux/kvm_host.h b/include/linux/kvm_host.h
index 8a2975674de4b..cd28eb34aaeb1 100644
--- a/include/linux/kvm_host.h
+++ b/include/linux/kvm_host.h
@@ -2433,6 +2433,7 @@ static inline bool kvm_mem_is_private(struct kvm *kvm, gfn_t gfn)
 #endif /* CONFIG_KVM_GENERIC_MEMORY_ATTRIBUTES */
 
 #define KVM_GMEM_GET_PFN_SHARED         BIT(0)
+#define KVM_GMEM_GET_PFN_LOCKED         BIT(1)
 #define KVM_GMEM_GET_PFN_PREPARE        BIT(31)  /* internal */
 
 #ifdef CONFIG_KVM_PRIVATE_MEM
diff --git a/virt/kvm/guest_memfd.c b/virt/kvm/guest_memfd.c
index 492b04f4e5c18..f637abc6045ba 100644
--- a/virt/kvm/guest_memfd.c
+++ b/virt/kvm/guest_memfd.c
@@ -670,7 +670,8 @@ static int __kvm_gmem_get_pfn(struct file *file, struct kvm_memory_slot *slot,
 
 	r = 0;
 
-	folio_unlock(folio);
+	if (!(flags & KVM_GMEM_GET_PFN_LOCKED))
+		folio_unlock(folio);
 
 	return r;
 }
@@ -680,7 +681,7 @@ int kvm_gmem_get_pfn(struct kvm *kvm, struct kvm_memory_slot *slot,
 {
 	struct file *file = kvm_gmem_get_file(slot);
 	int r;
-	int valid_flags = KVM_GMEM_GET_PFN_SHARED;
+	int valid_flags = KVM_GMEM_GET_PFN_SHARED | KVM_GMEM_GET_PFN_LOCKED;
 
 	if ((flags & valid_flags) != flags)
 		return -EINVAL;
-- 
2.46.0