3 files changed, 61 insertions, 9 deletions

diff --git a/arch/x86/xen/p2m.c b/arch/x86/xen/p2m.c
index a8691cb08420..f18fd1d411f6 100644
--- a/arch/x86/xen/p2m.c
+++ b/arch/x86/xen/p2m.c
@@ -681,12 +681,8 @@ int set_foreign_p2m_mapping(struct gnttab_map_grant_ref *map_ops,
                 }
                 pfn = page_to_pfn(pages[i]);
 
-                WARN_ON(PagePrivate(pages[i]));
                 WARN(pfn_to_mfn(pfn) != INVALID_P2M_ENTRY, "page must be ballooned");
 
-                SetPagePrivate(pages[i]);
-                set_page_private(pages[i], mfn);
-
                 if (unlikely(!set_phys_to_machine(pfn, FOREIGN_FRAME(mfn)))) {
                         ret = -ENOMEM;
                         goto out;
@@ -716,9 +712,6 @@ int clear_foreign_p2m_mapping(struct gnttab_unmap_grant_ref *unmap_ops,
                         goto out;
                 }
 
-                set_page_private(pages[i], INVALID_P2M_ENTRY);
-                WARN_ON(!PagePrivate(pages[i]));
-                ClearPagePrivate(pages[i]);
                 set_phys_to_machine(pfn, INVALID_P2M_ENTRY);
         }
         if (kunmap_ops)
diff --git a/drivers/xen/grant-table.c b/drivers/xen/grant-table.c
index b4f93c490f83..89dcca448bb6 100644
--- a/drivers/xen/grant-table.c
+++ b/drivers/xen/grant-table.c
@@ -679,12 +679,27 @@ EXPORT_SYMBOL_GPL(gnttab_free_auto_xlat_frames);
  */
 int gnttab_alloc_pages(int nr_pages, struct page **pages)
 {
+        int i;
         int ret;
 
         ret = alloc_xenballooned_pages(nr_pages, pages, false);
         if (ret < 0)
                 return ret;
 
+        for (i = 0; i < nr_pages; i++) {
+#if BITS_PER_LONG < 64
+                struct xen_page_foreign *foreign;
+
+                foreign = kzalloc(sizeof(*foreign), GFP_KERNEL);
+                if (!foreign) {
+                        gnttab_free_pages(nr_pages, pages);
+                        return -ENOMEM;
+                }
+                set_page_private(pages[i], (unsigned long)foreign);
+#endif
+                SetPagePrivate(pages[i]);
+        }
+
         return 0;
 }
 EXPORT_SYMBOL(gnttab_alloc_pages);
@@ -696,6 +711,16 @@ EXPORT_SYMBOL(gnttab_alloc_pages);
  */
 void gnttab_free_pages(int nr_pages, struct page **pages)
 {
+        int i;
+
+        for (i = 0; i < nr_pages; i++) {
+                if (PagePrivate(pages[i])) {
+#if BITS_PER_LONG < 64
+                        kfree((void *)page_private(pages[i]));
+#endif
+                        ClearPagePrivate(pages[i]);
+                }
+        }
         free_xenballooned_pages(nr_pages, pages);
 }
 EXPORT_SYMBOL(gnttab_free_pages);
@@ -756,12 +781,22 @@ int gnttab_map_refs(struct gnttab_map_grant_ref *map_ops,
         if (ret)
                 return ret;
 
-        /* Retry eagain maps */
-        for (i = 0; i < count; i++)
+        for (i = 0; i < count; i++) {
+                /* Retry eagain maps */
                 if (map_ops[i].status == GNTST_eagain)
                         gnttab_retry_eagain_gop(GNTTABOP_map_grant_ref, map_ops + i,
                                                 &map_ops[i].status, __func__);
 
+                if (map_ops[i].status == GNTST_okay) {
+                        struct xen_page_foreign *foreign;
+
+                        SetPageForeign(pages[i]);
+                        foreign = xen_page_foreign(pages[i]);
+                        foreign->domid = map_ops[i].dom;
+                        foreign->gref = map_ops[i].ref;
+                }
+        }
+
         return set_foreign_p2m_mapping(map_ops, kmap_ops, pages, count);
 }
 EXPORT_SYMBOL_GPL(gnttab_map_refs);
@@ -770,12 +805,16 @@ int gnttab_unmap_refs(struct gnttab_unmap_grant_ref *unmap_ops,
                       struct gnttab_unmap_grant_ref *kunmap_ops,
                       struct page **pages, unsigned int count)
 {
+        unsigned int i;
         int ret;
 
         ret = HYPERVISOR_grant_table_op(GNTTABOP_unmap_grant_ref, unmap_ops, count);
         if (ret)
                 return ret;
 
+        for (i = 0; i < count; i++)
+                ClearPageForeign(pages[i]);
+
         return clear_foreign_p2m_mapping(unmap_ops, kunmap_ops, pages, count);
 }
 EXPORT_SYMBOL_GPL(gnttab_unmap_refs);
diff --git a/include/xen/grant_table.h b/include/xen/grant_table.h
index 949803e20872..d3bef563e8da 100644
--- a/include/xen/grant_table.h
+++ b/include/xen/grant_table.h
@@ -45,6 +45,8 @@
 #include <asm/xen/hypervisor.h>
 
 #include <xen/features.h>
+#include <linux/mm_types.h>
+#include <linux/page-flags.h>
 
 #define GNTTAB_RESERVED_XENSTORE 1
 
@@ -185,4 +187,22 @@ int gnttab_unmap_refs(struct gnttab_unmap_grant_ref *unmap_ops,
 void gnttab_batch_map(struct gnttab_map_grant_ref *batch, unsigned count);
 void gnttab_batch_copy(struct gnttab_copy *batch, unsigned count);
 
+
+struct xen_page_foreign {
+        domid_t domid;
+        grant_ref_t gref;
+};
+
+static inline struct xen_page_foreign *xen_page_foreign(struct page *page)
+{
+        if (!PageForeign(page))
+                return NULL;
+#if BITS_PER_LONG < 64
+        return (struct xen_page_foreign *)page->private;
+#else
+        BUILD_BUG_ON(sizeof(struct xen_page_foreign) > BITS_PER_LONG);
+        return (struct xen_page_foreign *)&page->private;
+#endif
+}
+
 #endif /* __ASM_GNTTAB_H__ */