IOMMU: hold page ref until after deferred TLB flush

When moving around a page via XENMAPSPACE_gmfn_range, deferring the TLB
flush for the "from" GFN range requires that the page remains allocated
to the guest until the TLB flush has actually occurred. Otherwise a
parallel hypercall to remove the page would only flush the TLB for the
GFN it has been moved to, but not the one is was mapped at originally.

This is part of XSA-346.

Fixes: cf95b2a9fd5a ("iommu: Introduce per cpu flag (iommu_dont_flush_iotlb) to avoid unnecessary iotlb... ")
Reported-by: Julien Grall <jgrall@amazon.com>
Signed-off-by: Jan Beulich <jbeulich@suse.com>
Acked-by: Julien Grall <jgrall@amazon.com>

--- a/xen/arch/x86/mm.c
+++ b/xen/arch/x86/mm.c
@@ -5626,7 +5626,8 @@ static int handle_iomem_range(unsigned l
 
 static int xenmem_add_to_physmap_once(
     struct domain *d,
-    const struct xen_add_to_physmap *xatp)
+    const struct xen_add_to_physmap *xatp,
+    struct page_info **ppage)
 {
     struct page_info *page = NULL;
     unsigned long gfn = 0; /* gcc ... */
@@ -5701,9 +5702,6 @@ static int xenmem_add_to_physmap_once(
 
     domain_lock(d);
 
-    if ( page )
-        put_page(page);
-
     /* Remove previously mapped page if it was present. */
     prev_mfn = mfn_x(get_gfn(d, xatp->gpfn, &p2mt));
     if ( mfn_valid(prev_mfn) )
@@ -5749,12 +5747,27 @@ static int xenmem_add_to_physmap_once(
         rc = guest_physmap_add_page(d, xatp->gpfn, mfn, PAGE_ORDER_4K);
 
  put_both:
-    /* In the XENMAPSPACE_gmfn, we took a ref of the gfn at the top */
+    /*
+     * In the XENMAPSPACE_gmfn case, we took a ref of the gfn at the top.
+     * We also may need to transfer ownership of the page reference to our
+     * caller.
+     */
     if ( xatp->space == XENMAPSPACE_gmfn ||
          xatp->space == XENMAPSPACE_gmfn_range )
+    {
         put_gfn(d, gfn);
+        if ( !rc && ppage )
+        {
+            *ppage = page;
+            page = NULL;
+        }
+    }
+
     domain_unlock(d);
 
+    if ( page )
+        put_page(page);
+
     return rc;
 }
 
@@ -5766,13 +5779,18 @@ static int xenmem_add_to_physmap(struct 
 
     if ( xatp->space == XENMAPSPACE_gmfn_range )
     {
+        struct page_info **ppage = NULL, *pages[16];
+
         if ( need_iommu(d) )
+        {
             this_cpu(iommu_dont_flush_iotlb) = 1;
+            ppage = &pages[0];
+        }
 
         start_xatp = *xatp;
         while ( xatp->size > 0 )
         {
-            rc = xenmem_add_to_physmap_once(d, xatp);
+            rc = xenmem_add_to_physmap_once(d, xatp, ppage);
             if ( rc < 0 )
                 break;
 
@@ -5780,8 +5798,13 @@ static int xenmem_add_to_physmap(struct 
             xatp->gpfn++;
             xatp->size--;
 
+            if ( ppage )
+                ++ppage;
+
             /* Check for continuation if it's not the last interation */
-            if ( xatp->size > 0 && hypercall_preempt_check() )
+            if ( xatp->size > 0 &&
+                 ((ppage && start_xatp.size - xatp->size >= ARRAY_SIZE(pages)) ||
+                  hypercall_preempt_check()) )
             {
                 rc = -EAGAIN;
                 break;
@@ -5790,15 +5813,25 @@ static int xenmem_add_to_physmap(struct 
 
         if ( need_iommu(d) )
         {
+            unsigned int i;
+
             this_cpu(iommu_dont_flush_iotlb) = 0;
             iommu_iotlb_flush(d, start_xatp.idx, start_xatp.size - xatp->size);
+            /*
+             * Now that the IOMMU TLB flush was done for the original GFN, drop
+             * the page references. The 2nd flush below is fine to make later, as
+             * whoever removes the page again from its new GFN will have to do
+             * another flush anyway.
+             */
+            for ( i = 0; i < start_xatp.size - xatp->size; ++i )
+                put_page(pages[i]);
             iommu_iotlb_flush(d, start_xatp.gpfn, start_xatp.size - xatp->size);
         }
 
         return rc;
     }
 
-    return xenmem_add_to_physmap_once(d, xatp);
+    return xenmem_add_to_physmap_once(d, xatp, NULL);
 }
 
 long arch_memory_op(int op, XEN_GUEST_HANDLE(void) arg)
