[HVM] Save/restore: merge xc_linux_restore and xc_hvm_restore

into one function (and one file) since they share a lot of code
Signed-off-by: Tim Deegan <Tim.Deegan@øensource.com>

--HG--
rename : tools/libxc/xc_linux_restore.c => tools/libxc/xc_domain_restore.c

7 files changed, 277 insertions, 498 deletions

diff --git a/tools/libxc/Makefile b/tools/libxc/Makefile
index c924ea509a..5e23bef968 100644
--- a/tools/libxc/Makefile
+++ b/tools/libxc/Makefile
@@ -26,8 +26,8 @@ CTRL_SRCS-$(CONFIG_X86_Linux) += xc_ptrace.c xc_ptrace_core.c
 
 GUEST_SRCS-y :=
 GUEST_SRCS-y += xg_private.c
-GUEST_SRCS-$(CONFIG_MIGRATE) += xc_linux_restore.c xc_linux_save.c
-GUEST_SRCS-$(CONFIG_HVM) += xc_hvm_build.c xc_hvm_restore.c xc_hvm_save.c
+GUEST_SRCS-$(CONFIG_MIGRATE) += xc_domain_restore.c xc_linux_save.c
+GUEST_SRCS-$(CONFIG_HVM) += xc_hvm_build.c xc_hvm_save.c
 
 # symlink libelf from xen/common/libelf/
 LIBELF_SRCS := libelf-tools.c libelf-loader.c
diff --git a/tools/libxc/xc_linux_restore.c b/tools/libxc/xc_domain_restore.c
index 02342b4509..624c675b9e 100644
--- a/tools/libxc/xc_linux_restore.c
+++ b/tools/libxc/xc_domain_restore.c
@@ -1,9 +1,25 @@
 /******************************************************************************
- * xc_linux_restore.c
+ * xc_domain_restore.c
  *
- * Restore the state of a Linux session.
+ * Restore the state of a guest session.
  *
  * Copyright (c) 2003, K A Fraser.
+ * Copyright (c) 2006, Intel Corporation
+ * Copyright (c) 2007, XenSource Inc.
+ *
+ * This program is free software; you can redistribute it and/or modify it
+ * under the terms and conditions of the GNU General Public License,
+ * version 2, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
+ * more details.
+ *
+ * You should have received a copy of the GNU General Public License along with
+ * this program; if not, write to the Free Software Foundation, Inc., 59 Temple
+ * Place - Suite 330, Boston, MA 02111-1307 USA.
+ *
  */
 
 #include <stdlib.h>
@@ -13,6 +29,9 @@
 #include "xg_save_restore.h"
 #include "xc_dom.h"
 
+#include <xen/hvm/ioreq.h>
+#include <xen/hvm/params.h>
+
 /* max mfn of the current host machine */
 static unsigned long max_mfn;
 
@@ -102,7 +121,7 @@ static int uncanonicalize_pagetable(int xc_handle, uint32_t dom,
     }
     
     
-    /* Alllocate the requistite number of mfns */
+    /* Allocate the requistite number of mfns */
     if (nr_mfns && xc_domain_memory_populate_physmap(
             xc_handle, dom, nr_mfns, 0, 0, p2m_batch) != 0) { 
         ERROR("Failed to allocate memory for batch.!\n"); 
@@ -141,9 +160,93 @@ static int uncanonicalize_pagetable(int xc_handle, uint32_t dom,
 }
 
 
-int xc_linux_restore(int xc_handle, int io_fd, uint32_t dom,
-                     unsigned int store_evtchn, unsigned long *store_mfn,
-                     unsigned int console_evtchn, unsigned long *console_mfn)
+/* Load the p2m frame list, plus potential extended info chunk */
+static xen_pfn_t * load_p2m_frame_list(int io_fd, int *pae_extended_cr3)
+{
+    xen_pfn_t *p2m_frame_list;
+    vcpu_guest_context_t ctxt;
+
+    if (!(p2m_frame_list = malloc(P2M_FL_SIZE))) {
+        ERROR("Couldn't allocate p2m_frame_list array");
+        return NULL;
+    }
+    
+    /* Read first entry of P2M list, or extended-info signature (~0UL). */
+    if (!read_exact(io_fd, p2m_frame_list, sizeof(long))) {
+            ERROR("read extended-info signature failed");
+            return NULL;
+        }
+    
+    if (p2m_frame_list[0] == ~0UL) {
+        uint32_t tot_bytes;
+        
+        /* Next 4 bytes: total size of following extended info. */
+        if (!read_exact(io_fd, &tot_bytes, sizeof(tot_bytes))) {
+            ERROR("read extended-info size failed");
+            return NULL;
+        }
+        
+        while (tot_bytes) {
+            uint32_t chunk_bytes;
+            char     chunk_sig[4];
+            
+            /* 4-character chunk signature + 4-byte remaining chunk size. */
+            if (!read_exact(io_fd, chunk_sig, sizeof(chunk_sig)) ||
+                !read_exact(io_fd, &chunk_bytes, sizeof(chunk_bytes))) {
+                ERROR("read extended-info chunk signature failed");
+                return NULL;
+            }
+            tot_bytes -= 8;
+            
+            /* VCPU context structure? */
+            if (!strncmp(chunk_sig, "vcpu", 4)) {
+                if (!read_exact(io_fd, &ctxt, sizeof(ctxt))) {
+                    ERROR("read extended-info vcpu context failed");
+                    return NULL;
+                }
+                tot_bytes   -= sizeof(struct vcpu_guest_context);
+                chunk_bytes -= sizeof(struct vcpu_guest_context);
+                
+                if (ctxt.vm_assist & (1UL << VMASST_TYPE_pae_extended_cr3))
+                    *pae_extended_cr3 = 1;
+            }
+            
+            /* Any remaining bytes of this chunk: read and discard. */
+            while (chunk_bytes) {
+                unsigned long sz = chunk_bytes;
+                if ( sz > P2M_FL_SIZE )
+                    sz = P2M_FL_SIZE;
+                if (!read_exact(io_fd, p2m_frame_list, sz)) {
+                    ERROR("read-and-discard extended-info chunk bytes failed");
+                    return NULL;
+                }
+                chunk_bytes -= sz;
+                tot_bytes   -= sz;
+            }
+        }
+        
+        /* Now read the real first entry of P2M list. */
+        if (!read_exact(io_fd, p2m_frame_list, sizeof(long))) {
+            ERROR("read first entry of p2m_frame_list failed");
+            return NULL;
+        }
+    }
+    
+    /* First entry is already read into the p2m array. */
+    if (!read_exact(io_fd, &p2m_frame_list[1], P2M_FL_SIZE - sizeof(long))) {
+            ERROR("read p2m_frame_list failed");
+            return NULL;
+    }
+    
+    return p2m_frame_list;
+}
+
+
+
+int xc_domain_restore(int xc_handle, int io_fd, uint32_t dom,
+                      unsigned int store_evtchn, unsigned long *store_mfn,
+                      unsigned int console_evtchn, unsigned long *console_mfn,
+                      unsigned int hvm, unsigned int pae)
 {
     DECLARE_DOMCTL;
     int rc = 1, i, j, n, m, pae_extended_cr3 = 0;
@@ -174,7 +277,7 @@ int xc_linux_restore(int xc_handle, int io_fd, uint32_t dom,
 
     /* A copy of the pfn-to-mfn table frame list. */
     xen_pfn_t *p2m_frame_list = NULL;
-
+    
     /* A temporary mapping of the guest's start_info page. */
     start_info_t *start_info;
 
@@ -193,6 +296,12 @@ int xc_linux_restore(int xc_handle, int io_fd, uint32_t dom,
     unsigned int max_vcpu_id = 0;
     int new_ctxt_format = 0;
 
+    /* Magic frames in HVM guests: ioreqs and xenstore comms. */
+    uint64_t magic_pfns[3]; /* ioreq_pfn, bufioreq_pfn, store_pfn */
+
+    /* Buffer for holding HVM context */
+    uint8_t *hvm_buf = NULL;
+
     /* For info only */
     nr_pfns = 0;
 
@@ -201,20 +310,24 @@ int xc_linux_restore(int xc_handle, int io_fd, uint32_t dom,
         ERROR("read: p2m_size");
         goto out;
     }
-    DPRINTF("xc_linux_restore start: p2m_size = %lx\n", p2m_size);
+    DPRINTF("xc_domain_restore start: p2m_size = %lx\n", p2m_size);
 
-    /*
-     * XXX For now, 32bit dom0's can only save/restore 32bit domUs
-     * on 64bit hypervisors.
-     */
-    memset(&domctl, 0, sizeof(domctl));
-    domctl.domain = dom;
-    domctl.cmd    = XEN_DOMCTL_set_address_size;
-    domctl.u.address_size.size = sizeof(unsigned long) * 8;
-    rc = do_domctl(xc_handle, &domctl);
-    if ( rc != 0 ) {
-	ERROR("Unable to set guest address size.");
-	goto out;
+    if ( !hvm )
+    {
+        /*
+         * XXX For now, 32bit dom0's can only save/restore 32bit domUs
+         * on 64bit hypervisors.
+         */
+        memset(&domctl, 0, sizeof(domctl));
+        domctl.domain = dom;
+        domctl.cmd    = XEN_DOMCTL_set_address_size;
+        domctl.u.address_size.size = sizeof(unsigned long) * 8;
+        rc = do_domctl(xc_handle, &domctl);
+        if ( rc != 0 ) {
+            ERROR("Unable to set guest address size.");
+            goto out;
+        }
+        rc = 1;
     }
 
     if(!get_platform_info(xc_handle, dom,
@@ -229,76 +342,12 @@ int xc_linux_restore(int xc_handle, int io_fd, uint32_t dom,
         return 1;
     }
 
-    if (!(p2m_frame_list = malloc(P2M_FL_SIZE))) {
-        ERROR("Couldn't allocate p2m_frame_list array");
-        goto out;
-    }
-
-    /* Read first entry of P2M list, or extended-info signature (~0UL). */
-    if (!read_exact(io_fd, p2m_frame_list, sizeof(long))) {
-        ERROR("read extended-info signature failed");
-        goto out;
-    }
-
-    if (p2m_frame_list[0] == ~0UL) {
-        uint32_t tot_bytes;
-
-        /* Next 4 bytes: total size of following extended info. */
-        if (!read_exact(io_fd, &tot_bytes, sizeof(tot_bytes))) {
-            ERROR("read extended-info size failed");
-            goto out;
-        }
-
-        while (tot_bytes) {
-            uint32_t chunk_bytes;
-            char     chunk_sig[4];
-
-            /* 4-character chunk signature + 4-byte remaining chunk size. */
-            if (!read_exact(io_fd, chunk_sig, sizeof(chunk_sig)) ||
-                !read_exact(io_fd, &chunk_bytes, sizeof(chunk_bytes))) {
-                ERROR("read extended-info chunk signature failed");
-                goto out;
-            }
-            tot_bytes -= 8;
-
-            /* VCPU context structure? */
-            if (!strncmp(chunk_sig, "vcpu", 4)) {
-                if (!read_exact(io_fd, &ctxt, sizeof(ctxt))) {
-                    ERROR("read extended-info vcpu context failed");
-                    goto out;
-                }
-                tot_bytes   -= sizeof(struct vcpu_guest_context);
-                chunk_bytes -= sizeof(struct vcpu_guest_context);
-
-                if (ctxt.vm_assist & (1UL << VMASST_TYPE_pae_extended_cr3))
-                    pae_extended_cr3 = 1;
-            }
-
-            /* Any remaining bytes of this chunk: read and discard. */
-            while (chunk_bytes) {
-                unsigned long sz = chunk_bytes;
-                if ( sz > P2M_FL_SIZE )
-                    sz = P2M_FL_SIZE;
-                if (!read_exact(io_fd, p2m_frame_list, sz)) {
-                    ERROR("read-and-discard extended-info chunk bytes failed");
-                    goto out;
-                }
-                chunk_bytes -= sz;
-                tot_bytes   -= sz;
-            }
-        }
-
-        /* Now read the real first entry of P2M list. */
-        if (!read_exact(io_fd, p2m_frame_list, sizeof(long))) {
-            ERROR("read first entry of p2m_frame_list failed");
+    /* Load the p2m frame list, plus potential extended info chunk */
+    if ( !hvm ) 
+    {
+        p2m_frame_list = load_p2m_frame_list(io_fd, &pae_extended_cr3);
+        if ( !p2m_frame_list )
             goto out;
-        }
-    }
-
-    /* First entry is already read into the p2m array. */
-    if (!read_exact(io_fd, &p2m_frame_list[1], P2M_FL_SIZE - sizeof(long))) {
-        ERROR("read p2m_frame_list failed");
-        goto out;
     }
 
     /* We want zeroed memory so use calloc rather than malloc. */
@@ -442,8 +491,9 @@ int xc_linux_restore(int xc_handle, int io_fd, uint32_t dom,
                     nr_pfns++; 
                 }
 
-                /* setup region_mfn[] for batch map */
-                region_mfn[i] = p2m[pfn]; 
+                /* setup region_mfn[] for batch map.
+                 * For HVM guests, this interface takes PFNs, not MFNs */
+                region_mfn[i] = hvm ? pfn : p2m[pfn]; 
             }
         } 
 
@@ -551,9 +601,10 @@ int xc_linux_restore(int xc_handle, int io_fd, uint32_t dom,
                 }
             }
 
-            if (xc_add_mmu_update(xc_handle, mmu,
-                                  (((unsigned long long)mfn) << PAGE_SHIFT)
-                                  | MMU_MACHPHYS_UPDATE, pfn)) {
+            if (!hvm 
+                && xc_add_mmu_update(xc_handle, mmu,
+                                     (((unsigned long long)mfn) << PAGE_SHIFT)
+                                     | MMU_MACHPHYS_UPDATE, pfn)) {
                 ERROR("failed machpys update mfn=%lx pfn=%lx", mfn, pfn);
                 goto out;
             }
@@ -578,13 +629,90 @@ int xc_linux_restore(int xc_handle, int io_fd, uint32_t dom,
      * Ensure we flush all machphys updates before potential PAE-specific
      * reallocations below.
      */
-    if (xc_finish_mmu_updates(xc_handle, mmu)) {
+    if (!hvm && xc_finish_mmu_updates(xc_handle, mmu)) {
         ERROR("Error doing finish_mmu_updates()");
         goto out;
     }
 
     DPRINTF("Received all pages (%d races)\n", nraces);
 
+    if ( hvm ) 
+    {
+        uint32_t rec_len;
+
+        /* Set HVM-specific parameters */
+        if ( !read_exact(io_fd, magic_pfns, sizeof(magic_pfns)) )
+        {
+            ERROR("error reading magic page addresses");
+            goto out;
+        }
+        
+        /* These comms pages need to be zeroed at the start of day */
+        if ( xc_clear_domain_page(xc_handle, dom, magic_pfns[0]) ||
+             xc_clear_domain_page(xc_handle, dom, magic_pfns[1]) ||
+             xc_clear_domain_page(xc_handle, dom, magic_pfns[2]) )
+        {
+            ERROR("error zeroing magic pages");
+            goto out;
+        }
+        
+        xc_set_hvm_param(xc_handle, dom, HVM_PARAM_IOREQ_PFN, magic_pfns[0]);
+        xc_set_hvm_param(xc_handle, dom, HVM_PARAM_BUFIOREQ_PFN, magic_pfns[1]);
+        xc_set_hvm_param(xc_handle, dom, HVM_PARAM_STORE_PFN, magic_pfns[2]);
+        xc_set_hvm_param(xc_handle, dom, HVM_PARAM_PAE_ENABLED, pae);
+        xc_set_hvm_param(xc_handle, dom, HVM_PARAM_STORE_EVTCHN, store_evtchn);
+        *store_mfn = magic_pfns[2];
+
+        /* Read vcpu contexts */
+        for (i = 0; i <= max_vcpu_id; i++) 
+        {
+            if (!(vcpumap & (1ULL << i)))
+                continue;
+
+            if ( !read_exact(io_fd, &(ctxt), sizeof(ctxt)) )
+            {
+                ERROR("error read vcpu context.\n");
+                goto out;
+            }
+            
+            if ( (rc = xc_vcpu_setcontext(xc_handle, dom, i, &ctxt)) )
+            {
+                ERROR("Could not set vcpu context, rc=%d", rc);
+                goto out;
+            }
+            rc = 1;
+        }
+
+        /* Read HVM context */
+        if ( !read_exact(io_fd, &rec_len, sizeof(uint32_t)) )
+        {
+            ERROR("error read hvm context size!\n");
+            goto out;
+        }
+        
+        hvm_buf = malloc(rec_len);
+        if ( hvm_buf == NULL )
+        {
+            ERROR("memory alloc for hvm context buffer failed");
+            errno = ENOMEM;
+            goto out;
+        }
+        
+        if ( !read_exact(io_fd, hvm_buf, rec_len) )
+        {
+            ERROR("error loading the HVM context");
+            goto out;
+        }
+        
+        rc = xc_domain_hvm_setcontext(xc_handle, dom, hvm_buf, rec_len);
+        if ( rc ) 
+            ERROR("error setting the HVM context");
+       
+        goto out;
+    }
+
+    /* Non-HVM guests only from here on */
+
     if ((pt_levels == 3) && !pae_extended_cr3) {
 
         /*
@@ -897,6 +1025,7 @@ int xc_linux_restore(int xc_handle, int io_fd, uint32_t dom,
             ERROR("Couldn't build vcpu%d", i);
             goto out;
         }
+        rc = 1;
     }
 
     if (!read_exact(io_fd, shared_info_page, PAGE_SIZE)) {
@@ -938,6 +1067,7 @@ int xc_linux_restore(int xc_handle, int io_fd, uint32_t dom,
     munmap(live_p2m, ROUNDUP(p2m_size * sizeof(xen_pfn_t), PAGE_SHIFT));
 
     DPRINTF("Domain ready to be built.\n");
+    rc = 0;
 
  out:
     if ( (rc != 0) && (dom != 0) )
@@ -945,6 +1075,7 @@ int xc_linux_restore(int xc_handle, int io_fd, uint32_t dom,
     free(mmu);
     free(p2m);
     free(pfn_type);
+    free(hvm_buf);
 
     /* discard cache for save file  */
     discard_file_cache(io_fd, 1 /*flush*/);
diff --git a/tools/libxc/xc_hvm_restore.c b/tools/libxc/xc_hvm_restore.c
deleted file mode 100644
index d9f7435951..0000000000
--- a/tools/libxc/xc_hvm_restore.c
+++ /dev/null
@@ -1,351 +0,0 @@
-/******************************************************************************
- * xc_hvm_restore.c
- *
- * Restore the state of a HVM guest.
- *
- * Copyright (c) 2003, K A Fraser.
- * Copyright (c) 2006 Intel Corperation
- * rewriten for hvm guest by Zhai Edwin <edwin.zhai@intel.com>
- *
- * This program is free software; you can redistribute it and/or modify it
- * under the terms and conditions of the GNU General Public License,
- * version 2, as published by the Free Software Foundation.
- *
- * This program is distributed in the hope it will be useful, but WITHOUT
- * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
- * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
- * more details.
- *
- * You should have received a copy of the GNU General Public License along with
- * this program; if not, write to the Free Software Foundation, Inc., 59 Temple
- * Place - Suite 330, Boston, MA 02111-1307 USA.
- *
- */
-
-#include <stdlib.h>
-#include <unistd.h>
-
-#include "xg_private.h"
-#include "xg_save_restore.h"
-
-#include <xen/hvm/ioreq.h>
-#include <xen/hvm/params.h>
-#include <xen/hvm/e820.h>
-
-static ssize_t
-read_exact(int fd, void *buf, size_t count)
-{
-    int r = 0, s;
-    unsigned char *b = buf;
-
-    while ( r < count )
-    {
-        s = read(fd, &b[r], count - r);
-        if ( (s == -1) && (errno == EINTR) )
-            continue;
-        if ( s <= 0 )
-            break;
-        r += s;
-    }
-
-    return (r == count) ? 1 : 0;
-}
-
-#define BPL (sizeof(long)*8)
-#define test_bit(bit, map) !!((map)[(bit)/BPL] & (1UL << ((bit) % BPL)))
-#define set_bit(bit, map)  ((map)[(bit)/BPL] |= (1UL << ((bit) % BPL)))
-static int test_and_set_bit(unsigned long nr, unsigned long *map)
-{
-    int rc = test_bit(nr, map);
-    if ( !rc )
-        set_bit(nr, map);
-    return rc;
-}
-
-int xc_hvm_restore(int xc_handle, int io_fd, uint32_t dom,
-                   unsigned int store_evtchn, unsigned long *store_mfn,
-                   unsigned int pae, unsigned int apic)
-{
-    DECLARE_DOMCTL;
-
-    /* A copy of the CPU context of the guest. */
-    vcpu_guest_context_t ctxt;
-
-    char *region_base;
-
-    unsigned long buf[PAGE_SIZE/sizeof(unsigned long)];
-
-    xc_dominfo_t info;
-    unsigned int rc = 1, n, i;
-    uint32_t rec_len, nr_vcpus;
-    uint8_t *hvm_buf = NULL;
-
-    /* Magic frames: ioreqs and xenstore comms. */
-    uint64_t magic_pfns[3]; /* ioreq_pfn, bufioreq_pfn, store_pfn */
-
-    unsigned long pfn;
-    int verify = 0;
-
-    /* Types of the pfns in the current region */
-    unsigned long region_pfn_type[MAX_BATCH_SIZE];
-    xen_pfn_t pfn_alloc_batch[MAX_BATCH_SIZE];
-    unsigned int pfn_alloc_batch_size;
-
-    /* The size of an array big enough to contain all guest pfns */
-    unsigned long max_pfn = 0xfffffUL; /* initial memory map guess: 4GB */
-    unsigned long *pfn_bitmap = NULL, *new_pfn_bitmap;
-
-    DPRINTF("xc_hvm_restore:dom=%d, store_evtchn=%d, "
-            "pae=%u, apic=%u.\n", dom, store_evtchn, pae, apic);
-
-    DPRINTF("xc_hvm_restore start: max_pfn = %lx\n", max_pfn);
-
-    if ( mlock(&ctxt, sizeof(ctxt)) )
-    {
-        /* needed for build dom0 op, but might as well do early */
-        ERROR("Unable to mlock ctxt");
-        return 1;
-    }
-
-    if ( xc_domain_getinfo(xc_handle, dom, 1, &info) != 1 )
-    {
-        ERROR("Could not get domain info");
-        return 1;
-    }
-
-    domctl.cmd = XEN_DOMCTL_getdomaininfo;
-    domctl.domain = (domid_t)dom;
-    if ( xc_domctl(xc_handle, &domctl) < 0 )
-    {
-        ERROR("Could not get information on new domain");
-        goto out;
-    }
-
-    pfn_bitmap = calloc((max_pfn+1)/8, 1);
-    if ( pfn_bitmap == NULL )
-    {
-        ERROR("Could not allocate pfn bitmap");
-        goto out;
-    }
-
-    n = 0;
-    for ( ; ; )
-    {
-        int j;
-
-        if ( !read_exact(io_fd, &j, sizeof(int)) )
-        {
-            ERROR("HVM restore Error when reading batch size");
-            goto out;
-        }
-
-        PPRINTF("batch %d\n",j);
-
-        if ( j == -1 )
-        {
-            verify = 1;
-            DPRINTF("Entering page verify mode\n");
-            continue;
-        }
-
-        if ( j == 0 )
-            break;  /* our work here is done */
-
-        if ( j > MAX_BATCH_SIZE )
-        {
-            ERROR("Max batch size exceeded. Giving up.");
-            goto out;
-        }
-
-        if ( !read_exact(io_fd, region_pfn_type, j*sizeof(unsigned long)) )
-        {
-            ERROR("Error when reading region pfn types");
-            goto out;
-        }
-
-        pfn_alloc_batch_size = 0;
-        for ( i = 0; i < j; i++ )
-        {
-            pfn = region_pfn_type[i];
-            if ( pfn & XEN_DOMCTL_PFINFO_LTAB_MASK )
-                continue;
-
-            while ( pfn > max_pfn )
-            {
-                if ( max_pfn >= 0xfffffff )
-                {
-                    ERROR("Maximum PFN beyond reason (1TB) %lx\n", pfn);
-                    goto out;
-                }
-                max_pfn = 2*max_pfn + 1;
-                new_pfn_bitmap = realloc(pfn_bitmap, (max_pfn+1)/8);
-                if ( new_pfn_bitmap == NULL )
-                {
-                    ERROR("Could not realloc pfn bitmap for max_pfn=%lx\n",
-                          max_pfn);
-                    goto out;
-                }
-                pfn_bitmap = new_pfn_bitmap;
-                memset(&pfn_bitmap[(max_pfn+1)/(2*BPL)], 0, (max_pfn+1)/(2*8));
-            }
-
-            if ( !test_and_set_bit(pfn, pfn_bitmap) )
-                pfn_alloc_batch[pfn_alloc_batch_size++] = pfn;
-        }
-
-        if ( pfn_alloc_batch_size != 0 )
-        {
-             rc = xc_domain_memory_populate_physmap(
-                 xc_handle, dom, pfn_alloc_batch_size, 0, 0, pfn_alloc_batch);
-             if ( rc != 0 )
-             {
-                 PERROR("Could not allocate %u pages for HVM guest.\n",
-                        pfn_alloc_batch_size);
-                 goto out;
-             }
-        }
-
-        region_base = xc_map_foreign_batch(
-            xc_handle, dom, PROT_WRITE, region_pfn_type, j);
-
-        for ( i = 0; i < j; i++ )
-        {
-            void *page;
-
-            pfn = region_pfn_type[i];
-            if ( pfn & XEN_DOMCTL_PFINFO_LTAB_MASK )
-                continue;
-
-            /* In verify mode, we use a copy; otherwise we work in place */
-            page = verify ? (void *)buf : (region_base + i*PAGE_SIZE);
-
-            if ( !read_exact(io_fd, page, PAGE_SIZE) )
-            {
-                ERROR("Error when reading page (%x)", i);
-                goto out;
-            }
-
-            if ( verify )
-            {
-                int res = memcmp(buf, (region_base + i*PAGE_SIZE), PAGE_SIZE);
-                if ( res )
-                {
-                    int v;
-
-                    DPRINTF("************** pfn=%lx gotcs=%08lx "
-                            "actualcs=%08lx\n", pfn, 
-                            csum_page(region_base + i*PAGE_SIZE),
-                            csum_page(buf));
-
-                    for ( v = 0; v < 4; v++ )
-                    {
-                        unsigned long *p = (unsigned long *)
-                            (region_base + i*PAGE_SIZE);
-                        if (buf[v] != p[v])
-                            DPRINTF("    %d: %08lx %08lx\n", v, buf[v], p[v]);
-                    }
-                }
-            }
-
-        } /* end of 'batch' for loop */
-
-        munmap(region_base, j*PAGE_SIZE);
-        n += j; /* crude stats */
-    }
-    
-    xc_set_hvm_param(xc_handle, dom, HVM_PARAM_PAE_ENABLED, pae);
-    xc_set_hvm_param(xc_handle, dom, HVM_PARAM_STORE_EVTCHN, store_evtchn);
-
-    if ( !read_exact(io_fd, magic_pfns, sizeof(magic_pfns)) )
-    {
-        ERROR("error reading magic page addresses\n");
-        goto out;
-    }
-
-    if ( xc_clear_domain_page(xc_handle, dom, magic_pfns[0]) ||
-         xc_clear_domain_page(xc_handle, dom, magic_pfns[1]) ||
-         xc_clear_domain_page(xc_handle, dom, magic_pfns[2]) )
-    {
-        rc = -1;
-        goto out;
-    }
-
-    xc_set_hvm_param(xc_handle, dom, HVM_PARAM_IOREQ_PFN, magic_pfns[0]);
-    xc_set_hvm_param(xc_handle, dom, HVM_PARAM_BUFIOREQ_PFN, magic_pfns[1]);
-    xc_set_hvm_param(xc_handle, dom, HVM_PARAM_STORE_PFN, magic_pfns[2]);
-    *store_mfn = magic_pfns[2];
-    DPRINTF("hvm restore: calculate new store_mfn=0x%lx.\n", *store_mfn);
-
-    if ( !read_exact(io_fd, &nr_vcpus, sizeof(uint32_t)) )
-    {
-        ERROR("error read nr vcpu !\n");
-        goto out;
-    }
-    DPRINTF("hvm restore:get nr_vcpus=%d.\n", nr_vcpus);
-
-    for ( i = 0; i < nr_vcpus; i++ )
-    {
-        if ( !read_exact(io_fd, &rec_len, sizeof(uint32_t)) )
-        {
-            ERROR("error read vcpu context size!\n");
-            goto out;
-        }
-        if ( rec_len != sizeof(ctxt) )
-        {
-            ERROR("vcpu context size dismatch!\n");
-            goto out;
-        }
-
-        if ( !read_exact(io_fd, &(ctxt), sizeof(ctxt)) )
-        {
-            ERROR("error read vcpu context.\n");
-            goto out;
-        }
-
-        if ( (rc = xc_vcpu_setcontext(xc_handle, dom, i, &ctxt)) )
-        {
-            ERROR("Could not set vcpu context, rc=%d", rc);
-            goto out;
-        }
-    }
-
-    /* restore hvm context including pic/pit/shpage */
-    if ( !read_exact(io_fd, &rec_len, sizeof(uint32_t)) )
-    {
-        ERROR("error read hvm context size!\n");
-        goto out;
-    }
-
-    hvm_buf = malloc(rec_len);
-    if ( hvm_buf == NULL )
-    {
-        ERROR("memory alloc for hvm context buffer failed");
-        errno = ENOMEM;
-        goto out;
-    }
-
-    if ( !read_exact(io_fd, hvm_buf, rec_len) )
-    {
-        ERROR("error read hvm buffer!\n");
-        goto out;
-    }
-
-    if ( (rc = xc_domain_hvm_setcontext(xc_handle, dom, hvm_buf, rec_len)) )
-    {
-        ERROR("error set hvm buffer!\n");
-        goto out;
-    }
-
-    rc = 0;
-    goto out;
-
- out:
-    if ( (rc != 0) && (dom != 0) )
-        xc_domain_destroy(xc_handle, dom);
-    free(hvm_buf);
-    free(pfn_bitmap);
-
-    DPRINTF("Restore exit with rc=%d\n", rc);
-
-    return rc;
-}
diff --git a/tools/libxc/xc_hvm_save.c b/tools/libxc/xc_hvm_save.c
index 1b3ad5d356..8d34bf6643 100644
--- a/tools/libxc/xc_hvm_save.c
+++ b/tools/libxc/xc_hvm_save.c
@@ -305,6 +305,8 @@ int xc_hvm_save(int xc_handle, int io_fd, uint32_t dom, uint32_t max_iters,
 
     unsigned long total_sent = 0;
 
+    uint64_t vcpumap = 1ULL;
+
     DPRINTF("xc_hvm_save: dom=%d, max_iters=%d, max_factor=%d, flags=0x%x, "
             "live=%d, debug=%d.\n", dom, max_iters, max_factor, flags,
             live, debug);
@@ -371,6 +373,12 @@ int xc_hvm_save(int xc_handle, int io_fd, uint32_t dom, uint32_t max_iters,
 
     /* Size of any array that covers 0 ... max_pfn */
     pfn_array_size = max_pfn + 1;
+    if ( !write_exact(io_fd, &pfn_array_size, sizeof(unsigned long)) )
+    {
+        ERROR("Error when writing to state file (1)");
+        goto out;
+    }
+    
 
     /* pretend we sent all the pages last iteration */
     sent_last_iter = pfn_array_size;
@@ -644,6 +652,32 @@ int xc_hvm_save(int xc_handle, int io_fd, uint32_t dom, uint32_t max_iters,
 
     DPRINTF("All HVM memory is saved\n");
 
+    {
+        struct {
+            int minustwo;
+            int max_vcpu_id;
+            uint64_t vcpumap;
+        } chunk = { -2, info.max_vcpu_id };
+
+        if (info.max_vcpu_id >= 64) {
+            ERROR("Too many VCPUS in guest!");
+            goto out;
+        }
+
+        for (i = 1; i <= info.max_vcpu_id; i++) {
+            xc_vcpuinfo_t vinfo;
+            if ((xc_vcpu_getinfo(xc_handle, dom, i, &vinfo) == 0) &&
+                vinfo.online)
+                vcpumap |= 1ULL << i;
+        }
+
+        chunk.vcpumap = vcpumap;
+        if(!write_exact(io_fd, &chunk, sizeof(chunk))) {
+            ERROR("Error when writing to state file (errno %d)", errno);
+            goto out;
+        }
+    }
+
     /* Zero terminate */
     i = 0;
     if ( !write_exact(io_fd, &i, sizeof(int)) )
@@ -666,33 +700,22 @@ int xc_hvm_save(int xc_handle, int io_fd, uint32_t dom, uint32_t max_iters,
         goto out;
     }
 
-    /* save vcpu/vmcs context */
-    if ( !write_exact(io_fd, &nr_vcpus, sizeof(uint32_t)) )
-    {
-        ERROR("error write nr vcpus");
-        goto out;
-    }
-
-    /*XXX: need a online map to exclude down cpu */
+    /* save vcpu/vmcs contexts */
     for ( i = 0; i < nr_vcpus; i++ )
     {
+        if (!(vcpumap & (1ULL << i)))
+            continue;
+
         if ( xc_vcpu_getcontext(xc_handle, dom, i, &ctxt) )
         {
             ERROR("HVM:Could not get vcpu context");
             goto out;
         }
 
-        rec_size = sizeof(ctxt);
-        DPRINTF("write %d vcpucontext of total %d.\n", i, nr_vcpus); 
-        if ( !write_exact(io_fd, &rec_size, sizeof(uint32_t)) )
-        {
-            ERROR("error write vcpu ctxt size");
-            goto out;
-        }
-
+        DPRINTF("write vcpu %d context.\n", i); 
         if ( !write_exact(io_fd, &(ctxt), sizeof(ctxt)) )
         {
-            ERROR("write vmcs failed!\n");
+            ERROR("write vcpu context failed!\n");
             goto out;
         }
     }
diff --git a/tools/libxc/xenguest.h b/tools/libxc/xenguest.h
index 988c85feb6..8b636abb81 100644
--- a/tools/libxc/xenguest.h
+++ b/tools/libxc/xenguest.h
@@ -38,29 +38,21 @@ int xc_hvm_save(int xc_handle, int io_fd, uint32_t dom, uint32_t max_iters,
                 void (*qemu_flip_buffer)(int, int));
 
 /**
- * This function will restore a saved domain running Linux.
+ * This function will restore a saved domain.
  *
  * @parm xc_handle a handle to an open hypervisor interface
  * @parm fd the file descriptor to restore a domain from
  * @parm dom the id of the domain
  * @parm store_evtchn the store event channel for this domain to use
  * @parm store_mfn returned with the mfn of the store page
+ * @parm hvm non-zero if this is a HVM restore
+ * @parm pae non-zero if this HVM domain has PAE support enabled
  * @return 0 on success, -1 on failure
  */
-int xc_linux_restore(int xc_handle, int io_fd, uint32_t dom,
-                     unsigned int store_evtchn, unsigned long *store_mfn,
-                     unsigned int console_evtchn, unsigned long *console_mfn);
-
-/**
- * This function will restore a saved hvm domain running unmodified guest.
- *
- * @parm store_mfn pass mem size & returned with the mfn of the store page
- * @return 0 on success, -1 on failure
- */
-int xc_hvm_restore(int xc_handle, int io_fd, uint32_t dom,
-                   unsigned int store_evtchn,
-                   unsigned long *store_mfn, 
-                   unsigned int pae, unsigned int apic);
+int xc_domain_restore(int xc_handle, int io_fd, uint32_t dom,
+                      unsigned int store_evtchn, unsigned long *store_mfn,
+                      unsigned int console_evtchn, unsigned long *console_mfn,
+                      unsigned int hvm, unsigned int pae);
 
 /**
  * This function will create a domain for a paravirtualized Linux
diff --git a/tools/libxc/xg_private.c b/tools/libxc/xg_private.c
index 9110518d69..5ca3302479 100644
--- a/tools/libxc/xg_private.c
+++ b/tools/libxc/xg_private.c
@@ -209,16 +209,6 @@ __attribute__((weak))
     return -1;
 }
 
-__attribute__((weak)) 
-    int xc_hvm_restore(int xc_handle, int io_fd, uint32_t dom,
-                       unsigned int store_evtchn,
-                       unsigned long *store_mfn,
-                       unsigned int pae, unsigned int apic)
-{
-    errno = ENOSYS;
-    return -1;
-}
-
 __attribute__((weak)) int xc_get_hvm_param(
     int handle, domid_t dom, int param, unsigned long *value)
 {
diff --git a/tools/xcutils/xc_restore.c b/tools/xcutils/xc_restore.c
index fc1a59073e..e769c0be31 100644
--- a/tools/xcutils/xc_restore.c
+++ b/tools/xcutils/xc_restore.c
@@ -39,14 +39,8 @@ main(int argc, char **argv)
     pae  = atoi(argv[6]);
     apic = atoi(argv[7]);
 
-    if ( hvm )
-        ret = xc_hvm_restore(xc_fd, io_fd, domid,
-                             store_evtchn, &store_mfn,
-                             pae, apic);
-    else
-        ret = xc_linux_restore(xc_fd, io_fd, domid,
-                               store_evtchn, &store_mfn,
-                               console_evtchn, &console_mfn);
+    ret = xc_domain_restore(xc_fd, io_fd, domid, store_evtchn, &store_mfn,
+                            console_evtchn, &console_mfn, hvm, pae);
 
     if ( ret == 0 )
     {