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/arm/mm.c
+++ b/xen/arch/arm/mm.c
@@ -988,7 +988,7 @@ void share_xen_page_with_privileged_gues
 int xenmem_add_to_physmap_one(
     struct domain *d,
     unsigned int space,
-    domid_t foreign_domid,
+    union add_to_physmap_extra extra,
     unsigned long idx,
     xen_pfn_t gpfn)
 {
@@ -1066,7 +1066,7 @@ int xenmem_add_to_physmap_one(
         struct domain *od;
         struct page_info *page;
         p2m_type_t p2mt;
-        od = rcu_lock_domain_by_any_id(foreign_domid);
+        od = rcu_lock_domain_by_any_id(extra.foreign_domid);
         if ( od == NULL )
             return -ESRCH;
 
--- a/xen/arch/x86/mm.c
+++ b/xen/arch/x86/mm.c
@@ -5451,7 +5451,7 @@ static int handle_iomem_range(unsigned l
 int xenmem_add_to_physmap_one(
     struct domain *d,
     unsigned int space,
-    domid_t foreign_domid,
+    union add_to_physmap_extra extra,
     unsigned long idx,
     xen_pfn_t gpfn)
 {
@@ -5561,9 +5561,20 @@ int xenmem_add_to_physmap_one(
         rc = guest_physmap_add_page(d, 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 ( space == XENMAPSPACE_gmfn || space == XENMAPSPACE_gmfn_range )
+    {
         put_gfn(d, gfn);
+        if ( !rc && extra.ppage )
+        {
+            *extra.ppage = page;
+            page = NULL;
+        }
+    }
 
     if ( page )
         put_page(page);
--- a/xen/common/memory.c
+++ b/xen/common/memory.c
@@ -713,9 +713,14 @@ static int xenmem_add_to_physmap(struct
 {
     unsigned int done = 0;
     long rc = 0;
+    union add_to_physmap_extra extra = {};
+    struct page_info *pages[16];
+
+    if ( xatp->space == XENMAPSPACE_gmfn_foreign )
+        extra.foreign_domid = DOMID_INVALID;
 
     if ( xatp->space != XENMAPSPACE_gmfn_range )
-        return xenmem_add_to_physmap_one(d, xatp->space, DOMID_INVALID,
+        return xenmem_add_to_physmap_one(d, xatp->space, extra,
                                          xatp->idx, xatp->gpfn);
 
     if ( xatp->size < start )
@@ -727,12 +732,15 @@ static int xenmem_add_to_physmap(struct
 
 #ifdef HAS_PASSTHROUGH
     if ( need_iommu(d) )
+    {
         this_cpu(iommu_dont_flush_iotlb) = 1;
+        extra.ppage = &pages[0];
+    }
 #endif
 
     while ( xatp->size > done )
     {
-        rc = xenmem_add_to_physmap_one(d, xatp->space, DOMID_INVALID,
+        rc = xenmem_add_to_physmap_one(d, xatp->space, extra,
                                        xatp->idx, xatp->gpfn);
         if ( rc < 0 )
             break;
@@ -740,8 +748,12 @@ static int xenmem_add_to_physmap(struct
         xatp->idx++;
         xatp->gpfn++;
 
+        if ( extra.ppage )
+            ++extra.ppage;
+
         /* Check for continuation if it's not the last iteration. */
-        if ( xatp->size > ++done && hypercall_preempt_check() )
+        if ( (++done > ARRAY_SIZE(pages) && extra.ppage) ||
+             (xatp->size > done && hypercall_preempt_check()) )
         {
             rc = start + done;
             break;
@@ -751,8 +763,18 @@ static int xenmem_add_to_physmap(struct
 #ifdef HAS_PASSTHROUGH
     if ( need_iommu(d) )
     {
+        unsigned int i;
+
         this_cpu(iommu_dont_flush_iotlb) = 0;
         iommu_iotlb_flush(d, xatp->idx - done, done);
+        /*
+         * 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 < done; ++i )
+            put_page(pages[i]);
         iommu_iotlb_flush(d, xatp->gpfn - done, done);
     }
 #endif
@@ -766,6 +788,7 @@ static int xenmem_add_to_physmap_batch(s
 {
     unsigned int done = 0;
     int rc;
+    union add_to_physmap_extra extra = {};
 
     if ( xatpb->size < start )
         return -EILSEQ;
@@ -780,6 +803,9 @@ static int xenmem_add_to_physmap_batch(s
          !guest_handle_okay(xatpb->errs, xatpb->size) )
         return -EFAULT;
 
+    if ( xatpb->space == XENMAPSPACE_gmfn_foreign )
+        extra.foreign_domid = xatpb->foreign_domid;
+
     while ( xatpb->size > done )
     {
         xen_ulong_t idx;
@@ -797,9 +823,7 @@ static int xenmem_add_to_physmap_batch(s
             goto out;
         }
 
-        rc = xenmem_add_to_physmap_one(d, xatpb->space,
-                                       xatpb->foreign_domid,
-                                       idx, gpfn);
+        rc = xenmem_add_to_physmap_one(d, xatpb->space, extra, idx, gpfn);
 
         if ( unlikely(__copy_to_guest_offset(xatpb->errs, 0, &rc, 1)) )
         {
--- a/xen/include/xen/mm.h
+++ b/xen/include/xen/mm.h
@@ -364,8 +364,22 @@ static inline unsigned int get_order_fro
 
 void scrub_one_page(struct page_info *);
 
+union add_to_physmap_extra {
+    /*
+     * XENMAPSPACE_gmfn: When deferring TLB flushes, a page reference needs
+     * to be kept until after the flush, so the page can't get removed from
+     * the domain (and re-used for another purpose) beforehand. By passing
+     * non-NULL, the caller of xenmem_add_to_physmap_one() indicates it wants
+     * to have ownership of such a reference transferred in the success case.
+     */
+    struct page_info **ppage;
+
+    /* XENMAPSPACE_gmfn_foreign */
+    domid_t foreign_domid;
+};
+
 int xenmem_add_to_physmap_one(struct domain *d, unsigned int space,
-                              domid_t foreign_domid,
+                              union add_to_physmap_extra extra,
                               unsigned long idx, xen_pfn_t gpfn);
 
 /* Returns 0 on success, or negative on error. */
