Fix TLB flush on grant unmap

diff -r 87cc3035108f xen/common/grant_table.c
--- a/xen/common/grant_table.c	Fri Jul 13 14:32:11 2007 +0100
+++ b/xen/common/grant_table.c	Fri Jul 13 16:36:19 2007 +0100
@@ -505,16 +505,49 @@ __gnttab_unmap_common(
         }
     }
 
-    if ( (map->flags & (GNTMAP_device_map|GNTMAP_host_map)) == 0 )
-    {
-        map->flags = 0;
-        put_maptrack_handle(ld->grant_table, op->handle);
-    }
-
     /* If just unmapped a writable mapping, mark as dirtied */
     if ( !(flags & GNTMAP_readonly) )
          gnttab_mark_dirty(rd, frame);
 
+ unmap_out:
+    op->status = rc;
+    spin_unlock(&rd->grant_table->lock);
+    rcu_unlock_domain(rd);
+}
+
+static void
+__gnttab_unmap_common_complete(struct gnttab_unmap_common *op)
+{
+    struct domain   *ld, *rd;
+    domid_t          dom;
+    struct active_grant_entry *act;
+    grant_entry_t   *sha;
+    struct grant_mapping *map;
+    u16 flags;
+
+    ld = current->domain;
+    map = &maptrack_entry(ld->grant_table, op->handle);
+    dom = map->domid;
+    flags = map->flags;
+
+    if ( unlikely((rd = rcu_lock_domain_by_id(dom)) == NULL) ) {
+        /* This shouldn't happen - __gnttab_unmap_common should have
+           already crashed the domain */
+        gdprintk(XENLOG_INFO, "Could not find domain %d\n", dom);
+        return;
+    }
+       
+    spin_lock(&rd->grant_table->lock);
+
+    act = &active_entry(rd->grant_table, map->ref);
+    sha = &shared_entry(rd->grant_table, map->ref);
+
+    if ( (map->flags & (GNTMAP_device_map|GNTMAP_host_map)) == 0 )
+    {
+        map->flags = 0;
+        put_maptrack_handle(ld->grant_table, op->handle);
+    }
+
     if ( ((act->pin & (GNTPIN_devw_mask|GNTPIN_hstw_mask)) == 0) &&
          !(flags & GNTMAP_readonly) )
         gnttab_clear_flag(_GTF_writing, &sha->flags);
@@ -522,8 +555,6 @@ __gnttab_unmap_common(
     if ( act->pin == 0 )
         gnttab_clear_flag(_GTF_reading, &sha->flags);
 
- unmap_out:
-    op->status = rc;
     spin_unlock(&rd->grant_table->lock);
     rcu_unlock_domain(rd);
 }
@@ -542,28 +573,51 @@ __gnttab_unmap_grant_ref(
     op->status = common.status;
 }
 
+static void
+__gnttab_unmap_grant_ref_complete(struct gnttab_unmap_grant_ref *op)
+{
+    struct gnttab_unmap_common common = {
+	.host_addr = op->host_addr,
+	.dev_bus_addr = op->dev_bus_addr,
+	.handle = op->handle,
+    };
+
+    __gnttab_unmap_common_complete(&common);
+}
+
 static long
 gnttab_unmap_grant_ref(
     XEN_GUEST_HANDLE(gnttab_unmap_grant_ref_t) uop, unsigned int count)
 {
-    int i;
+    int i, done = 0;
+    long rc = -EFAULT;
     struct gnttab_unmap_grant_ref op;
 
     for ( i = 0; i < count; i++ )
     {
         if ( unlikely(__copy_from_guest_offset(&op, uop, i, 1)) )
-            goto fault;
+            goto out;
         __gnttab_unmap_grant_ref(&op);
+        ++done;
         if ( unlikely(__copy_to_guest_offset(uop, i, &op, 1)) )
-            goto fault;
-    }
-
+            goto out;
+    }
+
+    rc = 0;
+
+out:
     flush_tlb_mask(current->domain->domain_dirty_cpumask);
-    return 0;
-
-fault:
-    flush_tlb_mask(current->domain->domain_dirty_cpumask);
-    return -EFAULT;    
+
+    for ( i = 0; i < done; i++ )
+    {
+        if ( unlikely(__copy_from_guest_offset(&op, uop, i, 1)) )
+            /* This really shouldn't happen as it worked earlier in
+               the function */
+            continue;
+        __gnttab_unmap_grant_ref_complete(&op);
+
+    }
+    return rc;
 }
 
 static void
@@ -580,28 +634,51 @@ __gnttab_unmap_and_replace(
     op->status = common.status;
 }
 
+static void
+__gnttab_unmap_and_replace_complete(
+    struct gnttab_unmap_and_replace *op)
+{
+    struct gnttab_unmap_common common = {
+	.host_addr = op->host_addr,
+	.new_addr = op->new_addr,
+	.handle = op->handle,
+    };
+
+    __gnttab_unmap_common_complete(&common);
+}
+
 static long
 gnttab_unmap_and_replace(
     XEN_GUEST_HANDLE(gnttab_unmap_and_replace_t) uop, unsigned int count)
 {
-    int i;
+    int i, done = 0;
+    long rc = -EFAULT;
     struct gnttab_unmap_and_replace op;
 
     for ( i = 0; i < count; i++ )
     {
         if ( unlikely(__copy_from_guest_offset(&op, uop, i, 1)) )
-            goto fault;
+            goto out;
         __gnttab_unmap_and_replace(&op);
+        ++done;
         if ( unlikely(__copy_to_guest_offset(uop, i, &op, 1)) )
-            goto fault;
-    }
-
+            goto out;
+    }
+
+    rc = 0;
+
+out:
     flush_tlb_mask(current->domain->domain_dirty_cpumask);
-    return 0;
-
-fault:
-    flush_tlb_mask(current->domain->domain_dirty_cpumask);
-    return -EFAULT;    
+
+    for ( i = 0; i < done; i++ )
+    {
+        if ( unlikely(__copy_from_guest_offset(&op, uop, i, 1)) )
+            /* This really shouldn't happen as it worked earlier in
+               the function */
+            continue;
+        __gnttab_unmap_and_replace_complete(&op);
+    }
+    return rc;    
 }
 
 int