[RFC PATCH] mm/thp: order huge zero folio PFN invalidation before removal

[Hengbin Zhang <uqbarz@gmail.com>]

The nonpersistent huge-zero shrinker removes huge_zero_folio with
xchg() and then invalidates huge_zero_pfn. A concurrent fault can
publish a new folio and its PFN between these operations, after which
the old shrinker invalidates the new generation's PFN identity.

A later partial mprotect() can misclassify the live special PMD and
enter the ordinary anonymous THP split path.

Invalidate huge_zero_pfn before making huge_zero_folio NULL. A getter
which observes a zero refcount while the old pointer is still present
cannot publish a new folio: its cmpxchg() fails and it retries. A
getter which succeeds does so after the invalidation and publishes the
new PFN afterwards.

Fixes: 3b77e8c8cde5 ("mm/thp: make is_huge_zero_pmd() safe and quicker")

Signed-off-by: Hengbin Zhang <uqbarz@gmail.com>
---
 mm/huge_memory.c | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/mm/huge_memory.c b/mm/huge_memory.c
index b5d1e9d4463d..fcc492368160 100644
--- a/mm/huge_memory.c
+++ b/mm/huge_memory.c
@@ -298,9 +298,11 @@ static unsigned long shrink_huge_zero_folio_scan(struct shrinker *shrink,
 						 struct shrink_control *sc)
 {
 	if (atomic_cmpxchg(&huge_zero_refcount, 1, 0) == 1) {
-		struct folio *zero_folio = xchg(&huge_zero_folio, NULL);
-		BUG_ON(zero_folio == NULL);
+		struct folio *zero_folio;
+
 		WRITE_ONCE(huge_zero_pfn, ~0UL);
+		zero_folio = xchg(&huge_zero_folio, NULL);
+		BUG_ON(zero_folio == NULL);
 		folio_put(zero_folio);
 		return HPAGE_PMD_NR;
 	}
-- 
2.34.1