diff mbox series

[v1] x86/mm/pat: fix VM_PAT handling when fork() fails in copy_page_range()

Message ID 20241029210331.1339581-1-david@redhat.com (mailing list archive)
State New
Headers show
Series [v1] x86/mm/pat: fix VM_PAT handling when fork() fails in copy_page_range() | expand

Commit Message

David Hildenbrand Oct. 29, 2024, 9:03 p.m. UTC
If track_pfn_copy() fails, we already added the dst VMA to the maple
tree. As fork() fails, we'll cleanup the maple tree, and stumble over
the dst VMA for which we neither performed any reservation nor copied
any page tables.

Consequently untrack_pfn() will see VM_PAT and try obtaining the
PAT information from the page table -- which fails because the page
table was not copied.

The easiest fix would be to simply clear the VM_PAT flag of the dst VMA
if track_pfn_copy() fails. However, the whole thing is about "simply"
clearing the VM_PAT flag is shaky as well: if we passed track_pfn_copy()
and performed a reservation, but copying the page tables fails, we'll
simply clear the VM_PAT flag, not properly undoing the reservation ...
which is also wrong.

So let's fix it properly: set the VM_PAT flag only if the reservation
succeeded (leaving it clear initially), and undo the reservation if
anything goes wrong while copying the page tables: clearing the VM_PAT
flag after undoing the reservation.

Note that any copied page table entries will get zapped when the VMA will
get removed later, after copy_page_range() succeeded; as VM_PAT is not set
then, we won't try cleaning VM_PAT up once more and untrack_pfn() will be
happy. Note that leaving these page tables in place without a reservation
is not a problem, as we are aborting fork(); this process will never run.

A reproducer [1] can trigger this usually at the first try:

[   45.239440] WARNING: CPU: 26 PID: 11650 at arch/x86/mm/pat/memtype.c:983 get_pat_info+0xf6/0x110
[   45.241082] Modules linked in: ...
[   45.249119] CPU: 26 UID: 0 PID: 11650 Comm: repro3 Not tainted 6.12.0-rc5+ #92
[   45.250598] Hardware name: QEMU Standard PC (Q35 + ICH9, 2009), BIOS 1.16.3-2.fc40 04/01/2014
[   45.252181] RIP: 0010:get_pat_info+0xf6/0x110
...
[   45.268513] Call Trace:
[   45.269003]  <TASK>
[   45.269425]  ? __warn.cold+0xb7/0x14d
[   45.270131]  ? get_pat_info+0xf6/0x110
[   45.270846]  ? report_bug+0xff/0x140
[   45.271519]  ? handle_bug+0x58/0x90
[   45.272192]  ? exc_invalid_op+0x17/0x70
[   45.272935]  ? asm_exc_invalid_op+0x1a/0x20
[   45.273717]  ? get_pat_info+0xf6/0x110
[   45.274438]  ? get_pat_info+0x71/0x110
[   45.275165]  untrack_pfn+0x52/0x110
[   45.275835]  unmap_single_vma+0xa6/0xe0
[   45.276549]  unmap_vmas+0x105/0x1f0
[   45.277256]  exit_mmap+0xf6/0x460
[   45.277913]  __mmput+0x4b/0x120
[   45.278512]  copy_process+0x1bf6/0x2aa0
[   45.279264]  kernel_clone+0xab/0x440
[   45.279959]  __do_sys_clone+0x66/0x90
[   45.280650]  do_syscall_64+0x95/0x180

Likely this case was missed in commit d155df53f310
("x86/mm/pat: clear VM_PAT if copy_p4d_range failed"), and instead of
undoing the reservation we simply cleared the VM_PAT flag.

Keep the documentation of these functions in include/linux/pgtable.h,
one place is more than sufficient -- we should clean that up for the other
functions like track_pfn_remap/untrack_pfn separately.

[1] https://gitlab.com/davidhildenbrand/scratchspace/-/raw/main/reproducers/pat_fork.c

Reported-by: xingwei lee <xrivendell7@gmail.com>
Reported-by: yuxin wang <wang1315768607@163.com>
Closes: https://lore.kernel.org/lkml/CABOYnLx_dnqzpCW99G81DmOr+2UzdmZMk=T3uxwNxwz+R1RAwg@mail.gmail.com/
Reported-by: Marius Fleischer <fleischermarius@gmail.com>
Closes: https://lore.kernel.org/lkml/CAJg=8jwijTP5fre8woS4JVJQ8iUA6v+iNcsOgtj9Zfpc3obDOQ@mail.gmail.com/
Fixes: d155df53f310 ("x86/mm/pat: clear VM_PAT if copy_p4d_range failed")
Fixes: 2ab640379a0a ("x86: PAT: hooks in generic vm code to help archs to track pfnmap regions - v3")
Cc: Dave Hansen <dave.hansen@linux.intel.com>
Cc: Andy Lutomirski <luto@kernel.org>
Cc: Peter Zijlstra <peterz@infradead.org>
Cc: Thomas Gleixner <tglx@linutronix.de>
Cc: Ingo Molnar <mingo@redhat.com>
Cc: Borislav Petkov <bp@alien8.de>
Cc: "H. Peter Anvin" <hpa@zytor.com>
Cc: Peter Xu <peterx@redhat.com>
Cc: Andrew Morton <akpm@linux-foundation.org>
Cc: Ma Wupeng <mawupeng1@huawei.com>
Signed-off-by: David Hildenbrand <david@redhat.com>
---
 arch/x86/mm/pat/memtype.c | 66 +++++++++++++++++++++++++--------------
 include/linux/pgtable.h   | 27 ++++++++++++----
 kernel/fork.c             |  4 +++
 mm/memory.c               |  9 ++----
 4 files changed, 70 insertions(+), 36 deletions(-)


base-commit: 0f4cb420b38489c9bab9d091c3815714be8cb69d

Comments

David Hildenbrand Oct. 31, 2024, 9:47 a.m. UTC | #1
On 30.10.24 22:32, Peter Xu wrote:
> On Tue, Oct 29, 2024 at 10:03:31PM +0100, David Hildenbrand wrote:
>> If track_pfn_copy() fails, we already added the dst VMA to the maple
>> tree. As fork() fails, we'll cleanup the maple tree, and stumble over
>> the dst VMA for which we neither performed any reservation nor copied
>> any page tables.
>>
>> Consequently untrack_pfn() will see VM_PAT and try obtaining the
>> PAT information from the page table -- which fails because the page
>> table was not copied.
>>
>> The easiest fix would be to simply clear the VM_PAT flag of the dst VMA
>> if track_pfn_copy() fails. However, the whole thing is about "simply"
>> clearing the VM_PAT flag is shaky as well: if we passed track_pfn_copy()
>> and performed a reservation, but copying the page tables fails, we'll
>> simply clear the VM_PAT flag, not properly undoing the reservation ...
>> which is also wrong.
> 
> David,
> 

Hi Peter,

> Sorry to not have chance yet reply to your other email..
> 
> The only concern I have with the current fix to fork() is.. we started to
> have device drivers providing fault() on PFNMAPs as vfio-pci does, then I
> think it means we could potentially start to hit the same issue even
> without fork(), but as long as the 1st pgtable entry of the PFNMAP range is
> not mapped when the process with VM_PAT vma exit()s, or munmap() the vma.

As these drivers are not using remap_pfn_range, there is no way they 
could currently get VM_PAT set.

So what you describe is independent of the current state we are fixing 
here, and this fix should sort out the issues with current VM_PAT handling.

It indeed is an interesting question how to handle reservations when 
*not* using remap_pfn_range() to cover the whole area.

remap_pfn_range() handles VM_PAT automatically because it can do it: it 
knows that the whole range will map consecutive PFNs with the same 
protection, and we expect not parts of the range suddenly getting 
unmapped (and any driver that does that is buggy).

This behavior is, however, not guaranteed to be the case when 
remap_pfn_range() is *not* called on the whole range.

For that case (i.e., vfio-pci) I still wonder if the driver shouldn't do 
the reservation and leave VM_PAT alone.

In the driver, we'd do the reservation once and not worry about fork() 
etc ... and we'd undo the reservation once the last relevant VM_PFNMAP 
VMA is gone or the driver let's go of the device. I assume there are 
already mechanisms in place to deal with that to some degree, because 
the driver cannot go away while any VMA still has the VM_PFNMAP mapping 
-- otherwise something would be seriously messed up.

Long story short: let's look into not using VM_PAT for that use case.

Looking at the VM_PAT issues we had over time, not making it more 
complicated sounds like a very reasonable thing to me :)
mawupeng Nov. 7, 2024, 8:43 a.m. UTC | #2
On 2024/10/31 17:47, David Hildenbrand wrote:
> On 30.10.24 22:32, Peter Xu wrote:
>> On Tue, Oct 29, 2024 at 10:03:31PM +0100, David Hildenbrand wrote:
>>> If track_pfn_copy() fails, we already added the dst VMA to the maple
>>> tree. As fork() fails, we'll cleanup the maple tree, and stumble over
>>> the dst VMA for which we neither performed any reservation nor copied
>>> any page tables.
>>>
>>> Consequently untrack_pfn() will see VM_PAT and try obtaining the
>>> PAT information from the page table -- which fails because the page
>>> table was not copied.
>>>
>>> The easiest fix would be to simply clear the VM_PAT flag of the dst VMA
>>> if track_pfn_copy() fails. However, the whole thing is about "simply"
>>> clearing the VM_PAT flag is shaky as well: if we passed track_pfn_copy()
>>> and performed a reservation, but copying the page tables fails, we'll
>>> simply clear the VM_PAT flag, not properly undoing the reservation ...
>>> which is also wrong.
>>
>> David,
>>
> 
> Hi Peter,
> 
>> Sorry to not have chance yet reply to your other email..
>>
>> The only concern I have with the current fix to fork() is.. we started to
>> have device drivers providing fault() on PFNMAPs as vfio-pci does, then I
>> think it means we could potentially start to hit the same issue even
>> without fork(), but as long as the 1st pgtable entry of the PFNMAP range is
>> not mapped when the process with VM_PAT vma exit()s, or munmap() the vma.
> 
> As these drivers are not using remap_pfn_range, there is no way they could currently get VM_PAT set.
> 
> So what you describe is independent of the current state we are fixing here, and this fix should sort out the issues with current VM_PAT handling.
> 
> It indeed is an interesting question how to handle reservations when *not* using remap_pfn_range() to cover the whole area.
> 
> remap_pfn_range() handles VM_PAT automatically because it can do it: it knows that the whole range will map consecutive PFNs with the same protection, and we expect not parts of the range suddenly getting unmapped (and any driver that does that is buggy).
> 
> This behavior is, however, not guaranteed to be the case when remap_pfn_range() is *not* called on the whole range.
> 
> For that case (i.e., vfio-pci) I still wonder if the driver shouldn't do the reservation and leave VM_PAT alone.
> 
> In the driver, we'd do the reservation once and not worry about fork() etc ... and we'd undo the reservation once the last relevant VM_PFNMAP VMA is gone or the driver let's go of the device. I assume there are already mechanisms in place to deal with that to some degree, because the driver cannot go away while any VMA still has the VM_PFNMAP mapping -- otherwise something would be seriously messed up.
> 
> Long story short: let's look into not using VM_PAT for that use case.
> 
> Looking at the VM_PAT issues we had over time, not making it more complicated sounds like a very reasonable thing to me :)

Hi David,

The VM_PAT reservation do seems complicated. It can trigger the same warning in get_pat_info if remap_p4d_range fails:

remap_pfn_range
  remap_pfn_range_notrack
    remap_pfn_range_internal
      remap_p4d_range	// page allocation can failed here
    zap_page_range_single
      unmap_single_vma
        untrack_pfn
          get_pat_info
            WARN_ON_ONCE(1);

Any idea on this problem?

>
David Hildenbrand Nov. 7, 2024, 9:08 a.m. UTC | #3
On 07.11.24 09:43, mawupeng wrote:
> 
> 
> On 2024/10/31 17:47, David Hildenbrand wrote:
>> On 30.10.24 22:32, Peter Xu wrote:
>>> On Tue, Oct 29, 2024 at 10:03:31PM +0100, David Hildenbrand wrote:
>>>> If track_pfn_copy() fails, we already added the dst VMA to the maple
>>>> tree. As fork() fails, we'll cleanup the maple tree, and stumble over
>>>> the dst VMA for which we neither performed any reservation nor copied
>>>> any page tables.
>>>>
>>>> Consequently untrack_pfn() will see VM_PAT and try obtaining the
>>>> PAT information from the page table -- which fails because the page
>>>> table was not copied.
>>>>
>>>> The easiest fix would be to simply clear the VM_PAT flag of the dst VMA
>>>> if track_pfn_copy() fails. However, the whole thing is about "simply"
>>>> clearing the VM_PAT flag is shaky as well: if we passed track_pfn_copy()
>>>> and performed a reservation, but copying the page tables fails, we'll
>>>> simply clear the VM_PAT flag, not properly undoing the reservation ...
>>>> which is also wrong.
>>>
>>> David,
>>>
>>
>> Hi Peter,
>>
>>> Sorry to not have chance yet reply to your other email..
>>>
>>> The only concern I have with the current fix to fork() is.. we started to
>>> have device drivers providing fault() on PFNMAPs as vfio-pci does, then I
>>> think it means we could potentially start to hit the same issue even
>>> without fork(), but as long as the 1st pgtable entry of the PFNMAP range is
>>> not mapped when the process with VM_PAT vma exit()s, or munmap() the vma.
>>
>> As these drivers are not using remap_pfn_range, there is no way they could currently get VM_PAT set.
>>
>> So what you describe is independent of the current state we are fixing here, and this fix should sort out the issues with current VM_PAT handling.
>>
>> It indeed is an interesting question how to handle reservations when *not* using remap_pfn_range() to cover the whole area.
>>
>> remap_pfn_range() handles VM_PAT automatically because it can do it: it knows that the whole range will map consecutive PFNs with the same protection, and we expect not parts of the range suddenly getting unmapped (and any driver that does that is buggy).
>>
>> This behavior is, however, not guaranteed to be the case when remap_pfn_range() is *not* called on the whole range.
>>
>> For that case (i.e., vfio-pci) I still wonder if the driver shouldn't do the reservation and leave VM_PAT alone.
>>
>> In the driver, we'd do the reservation once and not worry about fork() etc ... and we'd undo the reservation once the last relevant VM_PFNMAP VMA is gone or the driver let's go of the device. I assume there are already mechanisms in place to deal with that to some degree, because the driver cannot go away while any VMA still has the VM_PFNMAP mapping -- otherwise something would be seriously messed up.
>>
>> Long story short: let's look into not using VM_PAT for that use case.
>>
>> Looking at the VM_PAT issues we had over time, not making it more complicated sounds like a very reasonable thing to me :)
> 
> Hi David,
> 
> The VM_PAT reservation do seems complicated. It can trigger the same warning in get_pat_info if remap_p4d_range fails:
> 
> remap_pfn_range
>    remap_pfn_range_notrack
>      remap_pfn_range_internal
>        remap_p4d_range	// page allocation can failed here
>      zap_page_range_single
>        unmap_single_vma
>          untrack_pfn
>            get_pat_info
>              WARN_ON_ONCE(1);
> 
> Any idea on this problem?

In remap_pfn_range(), if remap_pfn_range_notrack() fails, we call 
untrack_pfn(), to undo the tracking.

The problem is that zap_page_range_single() shouldn't do that 
untrack_pfn() call.

That should be fixed by Peter's patch:

https://lore.kernel.org/all/20240712144244.3090089-1-peterx@redhat.com/T/#u
mawupeng Nov. 7, 2024, 9:30 a.m. UTC | #4
On 2024/11/7 17:08, David Hildenbrand wrote:
> On 07.11.24 09:43, mawupeng wrote:
>>
>>
>> On 2024/10/31 17:47, David Hildenbrand wrote:
>>> On 30.10.24 22:32, Peter Xu wrote:
>>>> On Tue, Oct 29, 2024 at 10:03:31PM +0100, David Hildenbrand wrote:
>>>>> If track_pfn_copy() fails, we already added the dst VMA to the maple
>>>>> tree. As fork() fails, we'll cleanup the maple tree, and stumble over
>>>>> the dst VMA for which we neither performed any reservation nor copied
>>>>> any page tables.
>>>>>
>>>>> Consequently untrack_pfn() will see VM_PAT and try obtaining the
>>>>> PAT information from the page table -- which fails because the page
>>>>> table was not copied.
>>>>>
>>>>> The easiest fix would be to simply clear the VM_PAT flag of the dst VMA
>>>>> if track_pfn_copy() fails. However, the whole thing is about "simply"
>>>>> clearing the VM_PAT flag is shaky as well: if we passed track_pfn_copy()
>>>>> and performed a reservation, but copying the page tables fails, we'll
>>>>> simply clear the VM_PAT flag, not properly undoing the reservation ...
>>>>> which is also wrong.
>>>>
>>>> David,
>>>>
>>>
>>> Hi Peter,
>>>
>>>> Sorry to not have chance yet reply to your other email..
>>>>
>>>> The only concern I have with the current fix to fork() is.. we started to
>>>> have device drivers providing fault() on PFNMAPs as vfio-pci does, then I
>>>> think it means we could potentially start to hit the same issue even
>>>> without fork(), but as long as the 1st pgtable entry of the PFNMAP range is
>>>> not mapped when the process with VM_PAT vma exit()s, or munmap() the vma.
>>>
>>> As these drivers are not using remap_pfn_range, there is no way they could currently get VM_PAT set.
>>>
>>> So what you describe is independent of the current state we are fixing here, and this fix should sort out the issues with current VM_PAT handling.
>>>
>>> It indeed is an interesting question how to handle reservations when *not* using remap_pfn_range() to cover the whole area.
>>>
>>> remap_pfn_range() handles VM_PAT automatically because it can do it: it knows that the whole range will map consecutive PFNs with the same protection, and we expect not parts of the range suddenly getting unmapped (and any driver that does that is buggy).
>>>
>>> This behavior is, however, not guaranteed to be the case when remap_pfn_range() is *not* called on the whole range.
>>>
>>> For that case (i.e., vfio-pci) I still wonder if the driver shouldn't do the reservation and leave VM_PAT alone.
>>>
>>> In the driver, we'd do the reservation once and not worry about fork() etc ... and we'd undo the reservation once the last relevant VM_PFNMAP VMA is gone or the driver let's go of the device. I assume there are already mechanisms in place to deal with that to some degree, because the driver cannot go away while any VMA still has the VM_PFNMAP mapping -- otherwise something would be seriously messed up.
>>>
>>> Long story short: let's look into not using VM_PAT for that use case.
>>>
>>> Looking at the VM_PAT issues we had over time, not making it more complicated sounds like a very reasonable thing to me :)
>>
>> Hi David,
>>
>> The VM_PAT reservation do seems complicated. It can trigger the same warning in get_pat_info if remap_p4d_range fails:
>>
>> remap_pfn_range
>>    remap_pfn_range_notrack
>>      remap_pfn_range_internal
>>        remap_p4d_range    // page allocation can failed here
>>      zap_page_range_single
>>        unmap_single_vma
>>          untrack_pfn
>>            get_pat_info
>>              WARN_ON_ONCE(1);
>>
>> Any idea on this problem?
> 
> In remap_pfn_range(), if remap_pfn_range_notrack() fails, we call untrack_pfn(), to undo the tracking.
> 
> The problem is that zap_page_range_single() shouldn't do that untrack_pfn() call.
> 
> That should be fixed by Peter's patch:
> 
> https://lore.kernel.org/all/20240712144244.3090089-1-peterx@redhat.com/T/#u

Thank you for your prompt reply.

This do fix this issue.

>
diff mbox series

Patch

diff --git a/arch/x86/mm/pat/memtype.c b/arch/x86/mm/pat/memtype.c
index feb8cc6a12bf..3a9e6dd58e2f 100644
--- a/arch/x86/mm/pat/memtype.c
+++ b/arch/x86/mm/pat/memtype.c
@@ -984,27 +984,54 @@  static int get_pat_info(struct vm_area_struct *vma, resource_size_t *paddr,
 	return -EINVAL;
 }
 
-/*
- * track_pfn_copy is called when vma that is covering the pfnmap gets
- * copied through copy_page_range().
- *
- * If the vma has a linear pfn mapping for the entire range, we get the prot
- * from pte and reserve the entire vma range with single reserve_pfn_range call.
- */
-int track_pfn_copy(struct vm_area_struct *vma)
+int track_pfn_copy(struct vm_area_struct *dst_vma,
+		struct vm_area_struct *src_vma)
 {
+	const unsigned long vma_size = src_vma->vm_end - src_vma->vm_start;
 	resource_size_t paddr;
-	unsigned long vma_size = vma->vm_end - vma->vm_start;
 	pgprot_t pgprot;
+	int rc;
 
-	if (vma->vm_flags & VM_PAT) {
-		if (get_pat_info(vma, &paddr, &pgprot))
-			return -EINVAL;
-		/* reserve the whole chunk covered by vma. */
-		return reserve_pfn_range(paddr, vma_size, &pgprot, 1);
+	if (!(src_vma->vm_flags & VM_PAT))
+		return 0;
+
+	/*
+	 * Duplicate the PAT information for the dst VMA based on the src
+	 * VMA.
+	 */
+	if (get_pat_info(src_vma, &paddr, &pgprot))
+		return -EINVAL;
+	rc = reserve_pfn_range(paddr, vma_size, &pgprot, 1);
+	if (!rc)
+		/* Reservation for the destination VMA succeeded. */
+		vm_flags_set(dst_vma, VM_PAT);
+	return rc;
+}
+
+void untrack_pfn_copy(struct vm_area_struct *dst_vma,
+		struct vm_area_struct *src_vma)
+{
+	resource_size_t paddr;
+	unsigned long size;
+
+	if (!(dst_vma->vm_flags & VM_PAT))
+		return;
+
+	/*
+	 * As the page tables might not have been copied yet, the PAT
+	 * information is obtained from the src VMA, just like during
+	 * track_pfn_copy().
+	 */
+	if (get_pat_info(src_vma, &paddr, NULL)) {
+		size = src_vma->vm_end - src_vma->vm_start;
+		free_pfn_range(paddr, size);
 	}
 
-	return 0;
+	/*
+	 * Reservation was freed, any copied page tables will get cleaned
+	 * up later, but without getting PAT involved again.
+	 */
+	vm_flags_clear(dst_vma, VM_PAT);
 }
 
 /*
@@ -1095,15 +1122,6 @@  void untrack_pfn(struct vm_area_struct *vma, unsigned long pfn,
 	}
 }
 
-/*
- * untrack_pfn_clear is called if the following situation fits:
- *
- * 1) while mremapping a pfnmap for a new region,  with the old vma after
- * its pfnmap page table has been removed.  The new vma has a new pfnmap
- * to the same pfn & cache type with VM_PAT set.
- * 2) while duplicating vm area, the new vma fails to copy the pgtable from
- * old vma.
- */
 void untrack_pfn_clear(struct vm_area_struct *vma)
 {
 	vm_flags_clear(vma, VM_PAT);
diff --git a/include/linux/pgtable.h b/include/linux/pgtable.h
index e8b2ac6bd2ae..616707b4ecb8 100644
--- a/include/linux/pgtable.h
+++ b/include/linux/pgtable.h
@@ -1518,14 +1518,24 @@  static inline void track_pfn_insert(struct vm_area_struct *vma, pgprot_t *prot,
 }
 
 /*
- * track_pfn_copy is called when vma that is covering the pfnmap gets
- * copied through copy_page_range().
+ * track_pfn_copy is called when a VM_PFNMAP VMA is about to get the page
+ * tables copied during copy_page_range().
  */
-static inline int track_pfn_copy(struct vm_area_struct *vma)
+static inline int track_pfn_copy(struct vm_area_struct *dst_vma,
+		struct vm_area_struct *src_vma)
 {
 	return 0;
 }
 
+/*
+ * untrack_pfn_copy is called when a VM_PFNMAP VMA failed to copy during
+ * copy_page_range(), but after track_pfn_copy() was already called.
+ */
+static inline void untrack_pfn_copy(struct vm_area_struct *dst_vma,
+		struct vm_area_struct *src_vma)
+{
+}
+
 /*
  * untrack_pfn is called while unmapping a pfnmap for a region.
  * untrack can be called for a specific region indicated by pfn and size or
@@ -1538,8 +1548,10 @@  static inline void untrack_pfn(struct vm_area_struct *vma,
 }
 
 /*
- * untrack_pfn_clear is called while mremapping a pfnmap for a new region
- * or fails to copy pgtable during duplicate vm area.
+ * untrack_pfn_clear is called in the following cases on a VM_PFNMAP VMA:
+ *
+ * 1) During mremap() on the src VMA after the page tables were moved.
+ * 2) During fork() on the dst VMA, immediately after duplicating the src VMA.
  */
 static inline void untrack_pfn_clear(struct vm_area_struct *vma)
 {
@@ -1550,7 +1562,10 @@  extern int track_pfn_remap(struct vm_area_struct *vma, pgprot_t *prot,
 			   unsigned long size);
 extern void track_pfn_insert(struct vm_area_struct *vma, pgprot_t *prot,
 			     pfn_t pfn);
-extern int track_pfn_copy(struct vm_area_struct *vma);
+extern int track_pfn_copy(struct vm_area_struct *dst_vma,
+		struct vm_area_struct *src_vma);
+extern void untrack_pfn_copy(struct vm_area_struct *dst_vma,
+		struct vm_area_struct *src_vma);
 extern void untrack_pfn(struct vm_area_struct *vma, unsigned long pfn,
 			unsigned long size, bool mm_wr_locked);
 extern void untrack_pfn_clear(struct vm_area_struct *vma);
diff --git a/kernel/fork.c b/kernel/fork.c
index 89ceb4a68af2..02a7a8b44107 100644
--- a/kernel/fork.c
+++ b/kernel/fork.c
@@ -504,6 +504,10 @@  struct vm_area_struct *vm_area_dup(struct vm_area_struct *orig)
 	vma_numab_state_init(new);
 	dup_anon_vma_name(orig, new);
 
+	/* track_pfn_copy() will later take care of copying internal state. */
+	if (unlikely(new->vm_flags & VM_PFNMAP))
+		untrack_pfn_clear(new);
+
 	return new;
 }
 
diff --git a/mm/memory.c b/mm/memory.c
index 3ccee51adfbb..f7fbf099e8f9 100644
--- a/mm/memory.c
+++ b/mm/memory.c
@@ -1372,11 +1372,7 @@  copy_page_range(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma)
 		return copy_hugetlb_page_range(dst_mm, src_mm, dst_vma, src_vma);
 
 	if (unlikely(src_vma->vm_flags & VM_PFNMAP)) {
-		/*
-		 * We do not free on error cases below as remove_vma
-		 * gets called on error from higher level routine
-		 */
-		ret = track_pfn_copy(src_vma);
+		ret = track_pfn_copy(dst_vma, src_vma);
 		if (ret)
 			return ret;
 	}
@@ -1413,7 +1409,6 @@  copy_page_range(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma)
 			continue;
 		if (unlikely(copy_p4d_range(dst_vma, src_vma, dst_pgd, src_pgd,
 					    addr, next))) {
-			untrack_pfn_clear(dst_vma);
 			ret = -ENOMEM;
 			break;
 		}
@@ -1423,6 +1418,8 @@  copy_page_range(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma)
 		raw_write_seqcount_end(&src_mm->write_protect_seq);
 		mmu_notifier_invalidate_range_end(&range);
 	}
+	if (ret && unlikely(src_vma->vm_flags & VM_PFNMAP))
+		untrack_pfn_copy(dst_vma, src_vma);
 	return ret;
 }