From 458e81ecf742382f585af598262cce81ab882c83 Mon Sep 17 00:00:00 2001
From: Minchan Kim <minchan@google.com>
Date: Mon, 3 May 2021 10:52:58 -0700
Subject: [PATCH] ANDROID: mm: spf: fix task fault accounting

SPF has missed per-process fault account logic.
Fix it.

Bug: 187072626
Signed-off-by: Minchan Kim <minchan@google.com>
Change-Id: Ibd854bf350721917ef0be1af65055691d16b52f0
---
 arch/arm64/mm/fault.c |  2 +-
 arch/x86/mm/fault.c   |  2 +-
 include/linux/mm.h    | 11 +++++++----
 mm/memory.c           |  4 +++-
 4 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/arch/arm64/mm/fault.c b/arch/arm64/mm/fault.c
index b50f8b062b8b..3e8dfa52b2c4 100644
--- a/arch/arm64/mm/fault.c
+++ b/arch/arm64/mm/fault.c
@@ -576,7 +576,7 @@ static int __kprobes do_page_fault(unsigned long far, unsigned int esr,
 	 * let's try a speculative page fault without grabbing the
 	 * mmap_sem.
 	 */
-	fault = handle_speculative_fault(mm, addr, mm_flags, &vma);
+	fault = handle_speculative_fault(mm, addr, mm_flags, &vma, regs);
 	if (fault != VM_FAULT_RETRY)
 		goto done;
 
diff --git a/arch/x86/mm/fault.c b/arch/x86/mm/fault.c
index 50a2d03c398a..e9afbf85e600 100644
--- a/arch/x86/mm/fault.c
+++ b/arch/x86/mm/fault.c
@@ -1316,7 +1316,7 @@ void do_user_addr_fault(struct pt_regs *regs,
 	 * protection keys since it can't be resolved.
 	 */
 	if (!(hw_error_code & X86_PF_PK)) {
-		fault = handle_speculative_fault(mm, address, flags, &vma);
+		fault = handle_speculative_fault(mm, address, flags, &vma, regs);
 		if (fault != VM_FAULT_RETRY)
 			goto done;
 	}
diff --git a/include/linux/mm.h b/include/linux/mm.h
index 1883e5922114..93c56a30cd6e 100644
--- a/include/linux/mm.h
+++ b/include/linux/mm.h
@@ -1792,11 +1792,13 @@ extern int fixup_user_fault(struct mm_struct *mm,
 extern vm_fault_t __handle_speculative_fault(struct mm_struct *mm,
 					     unsigned long address,
 					     unsigned int flags,
-					     struct vm_area_struct **vma);
+					     struct vm_area_struct **vma,
+					     struct pt_regs *regs);
 static inline vm_fault_t handle_speculative_fault(struct mm_struct *mm,
 						  unsigned long address,
 						  unsigned int flags,
-						  struct vm_area_struct **vma)
+						  struct vm_area_struct **vma,
+						  struct pt_regs *regs)
 {
 	/*
 	 * Try speculative page fault for multithreaded user space task only.
@@ -1805,7 +1807,7 @@ static inline vm_fault_t handle_speculative_fault(struct mm_struct *mm,
 		*vma = NULL;
 		return VM_FAULT_RETRY;
 	}
-	return __handle_speculative_fault(mm, address, flags, vma);
+	return __handle_speculative_fault(mm, address, flags, vma, regs);
 }
 extern bool can_reuse_spf_vma(struct vm_area_struct *vma,
 			      unsigned long address);
@@ -1813,7 +1815,8 @@ extern bool can_reuse_spf_vma(struct vm_area_struct *vma,
 static inline vm_fault_t handle_speculative_fault(struct mm_struct *mm,
 						  unsigned long address,
 						  unsigned int flags,
-						  struct vm_area_struct **vma)
+						  struct vm_area_struct **vma,
+						  struct pt_regs *regs)
 {
 	return VM_FAULT_RETRY;
 }
diff --git a/mm/memory.c b/mm/memory.c
index b14457490bfd..3e76655b279b 100644
--- a/mm/memory.c
+++ b/mm/memory.c
@@ -5063,7 +5063,8 @@ static vm_fault_t ___handle_speculative_fault(struct mm_struct *mm,
 
 vm_fault_t __handle_speculative_fault(struct mm_struct *mm,
 				unsigned long address, unsigned int flags,
-				struct vm_area_struct **vma)
+				struct vm_area_struct **vma,
+				struct pt_regs *regs)
 {
 	vm_fault_t ret;
 
@@ -5081,6 +5082,7 @@ vm_fault_t __handle_speculative_fault(struct mm_struct *mm,
 	if (ret != VM_FAULT_RETRY) {
 		put_vma(*vma);
 		*vma = NULL;
+		mm_account_fault(regs, address, flags, ret);
 	}
 
 	return ret;