*
  * Authors: Madhav Suresh <madhav@u.northwestern.edu>
  *          Mark Cartwright <mcartwright@gmail.com> (live migration)
+ *          Peter Dinda <pdinda@northwestern.edu> (store interface changes)
  *
  * This is free software.  You are permitted to use,
  * redistribute, and modify it as specified in the file "V3VEE_LICENSE".
 #include <palacios/vmm.h>
 
 
+/*
+  This code implements both checkpointing and live migration.  The
+  main difference between these two is how the memory of the VM is
+  transfered.
+
+  A checkpoint is written to/read from a checkpoint store.  A
+  checkpoint store is conceptually an unseekable stream.  This allows
+  us to reuse the same model for things like:
+
+  - checkpoint to memory
+  - checkpoint to directory of files
+  - checkpoint to gem5
+  - checkpoint to network
+  - live migrate over network
+  
+  To save data to a checkpoint store, you first open a context on the
+  checkpoint store, then you save data to the context. The data
+  consists of blobs tagged with strings.
+
+  Only a single context can be open at a time, and a context cannot be
+  reopened once it has been closed.  These invariants are also essential
+  to maximizing compatability with different implementations. 
+
+  The result of these constraints is that the stream that the
+  checkpoint store records will look like:
+
+  [contextname1][tag1len][tag1]][datalen][data][tag2len][tag2][datalen][data]....
+  [contextname2][tag1len][tag1]][datalen][data][tag2len][tag2][datalen][data]....
+  ...
+
+  The code implemented here assures that
+  [tag?len][tag?][datalen][data] is written as shown above.  The
+  checkpoint store handles [contextname?] internally. For example, it
+  might use the context names as files for a checkpoint, or it might
+  communicate the context names over a network stream.
+
+  To add checkpointing to your code, the primary thing you need to do
+  is implement save and load functions.   For a device, you then add them
+  to your dev_ops structure.   The save function takes an context.  
+  You then write on this context using one of the following functions/macros:
+
+  1. v3_chkpt_save(context, tag, datalen, dataptr);
+  2. V3_CHKPT_SAVE(context, tag, data, faillabel);
+  3. V3_CHKPT_SAVE_AUTOTAG(context, data, faillabel)
+  
+  Here (2) is a macro that computes the data length from sizeof(data)
+  while (3) does the same, and uses as the tag the name of the variable data.
+  We strongly recommend the use of (2). 
+  
+  These functions and macros will return -1 if the save is unsuccessful.
+  The faillabel argumnent is optional, if supplied and an error occurs,
+  the macro will goto the label. 
+
+  Some classes of devices, for example IDE HDs and CD-ROMs, and PCI
+  devices have a class-specific component and a device-specific
+  component.  For such devices, the class implementation (e.g., IDE, PCI) 
+  will provide a class-specific save function that should be called first
+  in the device's save function.  For example:
+
+  #include <devices/pci.h>
+  
+  static int my_device_save(struct v3_chkpt_ctx *ctx, void *priv) 
+  { 
+     struct my_device *dev = priv;
+
+     if (v3_pci_save(ctx, dev->pci)<0) { 
+        goto fail;
+     }
+
+     V3_CHKPT_SAVE(ctx, "myfoo", dev->foo, failout),;
+     V3_CHKPT_SAVE(ctx, "mybar", dev->bar, failout);
+
+     // Success
+     return 0;
+
+failout:
+
+     PrintError("Failed to save device\n");
+     return -1;
+
+  }     
+
+  The load side is symmetric. 
+
+*/
+
 struct v3_chkpt;
 
 
 struct v3_chkpt_ctx {
-    struct v3_chkpt * chkpt;
-    struct v3_chkpt_ctx * parent;
-    void * store_ctx;
+  struct v3_chkpt * chkpt;
+  void *store_ctx;
 };
 
-/* Temporary */
-#define  V3_CHKPT_STD_SAVE(ctx,x) v3_chkpt_save(ctx,#x,sizeof(x),&(x))
-#define  V3_CHKPT_STD_LOAD(ctx,x) v3_chkpt_load(ctx,#x,sizeof(x),&(x))
-
-
-
-int v3_chkpt_save(struct v3_chkpt_ctx * ctx, char * tag, uint64_t len, void * buf);
-int v3_chkpt_load(struct v3_chkpt_ctx * ctx, char * tag, uint64_t len, void * buf);
-
-static inline int v3_chkpt_save_64(struct v3_chkpt_ctx * ctx, char * tag, void * buf) {
-    return v3_chkpt_save(ctx, tag, 8, buf);
-}
-static inline int v3_chkpt_save_32(struct v3_chkpt_ctx * ctx, char * tag, void * buf) {
-    return v3_chkpt_save(ctx, tag, 4, buf);
-}
-static inline int v3_chkpt_save_16(struct v3_chkpt_ctx * ctx, char * tag, void * buf) {
-    return v3_chkpt_save(ctx, tag, 2, buf);
-}
-static inline int v3_chkpt_save_8(struct v3_chkpt_ctx * ctx, char * tag, void * buf) {
-    return v3_chkpt_save(ctx, tag, 1, buf);
-}
-
-static inline int v3_chkpt_load_64(struct v3_chkpt_ctx * ctx, char * tag, void * buf) {
-    return v3_chkpt_load(ctx, tag, 8, buf);
-}
-static inline int v3_chkpt_load_32(struct v3_chkpt_ctx * ctx, char * tag, void * buf) {
-    return v3_chkpt_load(ctx, tag, 4, buf);
-}
-static inline int v3_chkpt_load_16(struct v3_chkpt_ctx * ctx, char * tag, void * buf) {
-    return v3_chkpt_load(ctx, tag, 2, buf);
-}
-static inline int v3_chkpt_load_8(struct v3_chkpt_ctx * ctx, char * tag, void * buf) {
-    return v3_chkpt_load(ctx, tag, 1, buf);
-}
-
-
-
-int v3_chkpt_close_ctx(struct v3_chkpt_ctx * ctx);
-struct v3_chkpt_ctx * v3_chkpt_open_ctx(struct v3_chkpt * chkpt, struct v3_chkpt_ctx * parent, char * name);
+
+
+/*
+ * You do not need to look behind this curtain
+ */
+#define SELECT(x,A,FUNC, ...) FUNC
+#define V3_CHKPT_SAVE_BASE(context, tag, data)  v3_chkpt_save(context,tag,sizeof(data),&(data))
+#define V3_CHKPT_SAVE_LABELED(context, tag, data, faillabel)   (({if (V3_CHKPT_SAVE_BASE(context,tag,data)<0) { goto faillabel; } 0; }), 0)
+#define V3_CHKPT_LOAD_BASE(context, tag, data)  v3_chkpt_load(context,tag,sizeof(data),&(data))
+#define V3_CHKPT_LOAD_LABELED(context, tag, data, faillabel)   (({if (V3_CHKPT_LOAD_BASE(context,tag,data)<0) { goto faillabel; } 0; }), 0)
+/*
+ * Safe to open your eyes again
+ */
+
+
+//
+// Functions and macros to save to a context
+//
+//
+int     v3_chkpt_save(struct v3_chkpt_ctx * ctx, char * tag, uint64_t len, void * buf);
+#define V3_CHKPT_SAVE(context, tag, data, ...) SELECT(,##__VA_ARGS__,V3_CHKPT_SAVE_LABELED(context,tag,data,__VA_ARGS__),V3_CHKPT_SAVE_BASE(context,tag,data))
+#define V3_CHKPT_SAVE_AUTOTAG(context, data, ...) V3_CHKPT_SAVE(context, #data ,data,__VA_ARGS__)
+
+//
+// Funtions and macros to load from a context
+// 
+//
+int     v3_chkpt_load(struct v3_chkpt_ctx * ctx, char * tag, uint64_t len, void * buf);
+#define V3_CHKPT_LOAD(context, tag, data, ...) SELECT(,##__VA_ARGS__,V3_CHKPT_LOAD_LABELED(context,tag,data,__VA_ARGS__),V3_CHKPT_LOAD_BASE(context,tag,data))
+#define V3_CHKPT_LOAD_AUTOTAG(context, data, ...) V3_CHKPT_LOAD(context, #data ,data,__VA_ARGS__)
+
+
+struct v3_chkpt_ctx * v3_chkpt_open_ctx(struct v3_chkpt * chkpt, char * name);
+int                   v3_chkpt_close_ctx(struct v3_chkpt_ctx * ctx);
 
 int v3_chkpt_save_vm(struct v3_vm_info * vm, char * store, char * url);
 int v3_chkpt_load_vm(struct v3_vm_info * vm, char * store, char * url);
 
  *
  * Author: Madhav Suresh <madhav@u.northwestern.edu>
  *        Arefin Huq <fig@arefin.net>
+ *         Peter Dinda <pdinda@northwestern.edu> (store interface changes)
  *
  *
  * This is free software.  You are permitted to use,
 typedef enum {SAVE, LOAD} chkpt_mode_t;
 
 struct chkpt_interface {
-    char name[128];
-    void * (*open_chkpt)(char * url, chkpt_mode_t mode);
-    int (*close_chkpt)(void * store_data);
-    
-    void * (*open_ctx)(void * store_data, void * parent_ctx, char * name);
-    int (*close_ctx)(void * store_data, void * ctx);
-    
-    int (*save)(void * store_data, void * ctx, char * tag, uint64_t len, void * buf);
-    int (*load)(void * store_data, void * ctx, char * tag, uint64_t len, void * buf);
+  char name[128];
+  // Opening a checkpoint should return a pointer to the internal representation
+  // of the checkpoint in the store.  This will be passed back
+  // as "store_data".  Return NULL if the context cannot be opened
+  void * (*open_chkpt)(char * url, chkpt_mode_t mode);
+  // Closing the checkpoint should return -1 on failure, 0 on success
+  int    (*close_chkpt)(void * store_data);
+  
+  // Opening a context on the checkpoint with a given name should return
+  // a pointer to an internal representation of the context.  This pointer
+  // is then passed back as "ctx". 
+  // We will open only a single context at a time.  
+  void * (*open_ctx)(void * store_data, char *name);
+  // Closing the context should return -1 on failure, 0 on success
+  int    (*close_ctx)(void * store_data, void * ctx);
+  
+  // Save and load include a tagged data buffer.  These are 
+  // "all or nothing" writes and reads.  
+  // return -1 on failure, and 0 on success
+  // 
+  int (*save)(void * store_data, void * ctx, char * tag, uint64_t len, void * buf);
+  int (*load)(void * store_data, void * ctx, char * tag, uint64_t len, void * buf);
 };
 
 
 struct v3_chkpt {
-    struct v3_vm_info * vm;
-
-    struct chkpt_interface * interface;
-
-    void * store_data;
+  struct v3_vm_info * vm;
+  
+  struct v3_chkpt_ctx *current_ctx;
+  
+  struct chkpt_interface * interface;
+  
+  void * store_data;
 };
 
 
 static char vmx_chkpt_header[] = "v3vee palacios checkpoint version: x.x, VMX x.x";
 
 static int chkpt_close(struct v3_chkpt * chkpt) {
-    chkpt->interface->close_chkpt(chkpt->store_data);
+  if (chkpt) { 
+    int rc;
+
+    rc = chkpt->interface->close_chkpt(chkpt->store_data);
 
     V3_Free(chkpt);
 
-    return 0;
+    if (rc!=0) { 
+      PrintError("Internal store failed to close valid checkpoint\n");
+      return -1;
+    } else {
+      return 0;
+    }
+  } else {
+    PrintError("Attempt to close null checkpoint\n");
+    return -1;
+  }
 }
 
 
 
 
     chkpt = V3_Malloc(sizeof(struct v3_chkpt));
-
+    
     if (!chkpt) {
-       PrintError("Could not allocate checkpoint state\n");
+       PrintError("Could not allocate checkpoint state, closing checkpoint\n");
+       iface->close_chkpt(store_data);
        return NULL;
     }
 
+    memset(chkpt,0,sizeof(struct v3_chkpt));
+
     chkpt->interface = iface;
     chkpt->vm = vm;
     chkpt->store_data = store_data;
+    chkpt->current_ctx = NULL;
     
     return chkpt;
 }
 
-struct v3_chkpt_ctx * v3_chkpt_open_ctx(struct v3_chkpt * chkpt, struct v3_chkpt_ctx * parent, char * name) {
-    struct v3_chkpt_ctx * ctx = V3_Malloc(sizeof(struct v3_chkpt_ctx));
-    void * parent_store_ctx = NULL;
-
-
-    if (!ctx) { 
-       PrintError("Unable to allocate context\n");
-       return 0;
-    }
+struct v3_chkpt_ctx * v3_chkpt_open_ctx(struct v3_chkpt * chkpt, char * name) {
+  struct v3_chkpt_ctx * ctx;
 
-    memset(ctx, 0, sizeof(struct v3_chkpt_ctx));
+  if (chkpt->current_ctx) { 
+    PrintError("Attempt to open context %s before old context has been closed\n", name);
+    return NULL;
+  }
 
-    ctx->chkpt = chkpt;
-    ctx->parent = parent;
+  ctx = V3_Malloc(sizeof(struct v3_chkpt_ctx));
 
-    if (parent) {
-       parent_store_ctx = parent->store_ctx;
-    }
-
-    ctx->store_ctx = chkpt->interface->open_ctx(chkpt->store_data, parent_store_ctx, name);
+  if (!ctx) { 
+    PrintError("Unable to allocate context\n");
+    return 0;
+  }
+  
+  memset(ctx, 0, sizeof(struct v3_chkpt_ctx));
+  
+  ctx->chkpt = chkpt;
+  ctx->store_ctx = chkpt->interface->open_ctx(chkpt->store_data, name);
+
+  if (!(ctx->store_ctx)) {
+    PrintError("Underlying store failed to open context %s\n",name);
+    V3_Free(ctx);
+    return NULL;
+  }
 
-    if (!(ctx->store_ctx)) {
-       PrintError("Warning: opening underlying representation returned null\n");
-    }
+  chkpt->current_ctx = ctx;
 
-    return ctx;
+  return ctx;
 }
 
 int v3_chkpt_close_ctx(struct v3_chkpt_ctx * ctx) {
     struct v3_chkpt * chkpt = ctx->chkpt;
     int ret = 0;
 
+    if (chkpt->current_ctx != ctx) { 
+      PrintError("Attempt to close a context that is not the current context on the store\n");
+      return -1;
+    }
+
     ret = chkpt->interface->close_ctx(chkpt->store_data, ctx->store_ctx);
 
+    if (ret) { 
+      PrintError("Failed to close context on store, closing device-independent context anyway - bad\n");
+      ret = -1;
+    }
+
+    chkpt->current_ctx=NULL;
+
     V3_Free(ctx);
 
     return ret;
 
 int v3_chkpt_save(struct v3_chkpt_ctx * ctx, char * tag, uint64_t len, void * buf) {
     struct v3_chkpt * chkpt = ctx->chkpt;    
-    
-    return  chkpt->interface->save(chkpt->store_data, ctx->store_ctx, tag, len, buf);
+    int rc;
+
+    if (!ctx) { 
+      PrintError("Attempt to save tag %s on null context\n",tag);
+      return -1;
+    }
+
+    if (chkpt->current_ctx != ctx) { 
+      PrintError("Attempt to save on context that is not the current context for the store\n");
+      return -1;
+    }
+
+    rc = chkpt->interface->save(chkpt->store_data, ctx->store_ctx, tag , len, buf);
 
+    if (rc) { 
+      PrintError("Underlying store failed to save tag %s on valid context\n",tag);
+      return -1;
+    } else {
+      return 0;
+    }
 }
 
 
 int v3_chkpt_load(struct v3_chkpt_ctx * ctx, char * tag, uint64_t len, void * buf) {
     struct v3_chkpt * chkpt = ctx->chkpt;    
+    int rc;
+
+    if (!ctx) { 
+      PrintError("Attempt to load tag %s from null context\n",tag);
+      return -1;
+    }
     
-    return chkpt->interface->load(chkpt->store_data, ctx->store_ctx, tag, len, buf);
+    if (chkpt->current_ctx != ctx) { 
+      PrintError("Attempt to load from context that is not the current context for the store\n");
+      return -1;
+    }
+
+    rc = chkpt->interface->load(chkpt->store_data, ctx->store_ctx, tag, len, buf);
+
+    if (rc) { 
+      PrintError("Underlying store failed to load tag %s from valid context\n",tag);
+      return -1;
+    } else {
+      return 0;
+    }
 }
 
 
 
     guest_mem_base = V3_VAddr((void *)vm->mem_map.base_region.host_addr);
 
-    ctx = v3_chkpt_open_ctx(chkpt, NULL, "memory_img");
+    ctx = v3_chkpt_open_ctx(chkpt, "memory_img");
     
     if (!ctx) { 
        PrintError("Unable to open context for memory load\n");
        return -1;
     }
                     
-    if (v3_chkpt_load(ctx, "memory_img", vm->mem_size, guest_mem_base) == -1) {
+    if (v3_chkpt_load(ctx, "memory_img", vm->mem_size, guest_mem_base)) {
        PrintError("Unable to load all of memory (requested=%llu bytes, result=%llu bytes\n",(uint64_t)(vm->mem_size),ret);
        v3_chkpt_close_ctx(ctx);
        return -1;
 
     guest_mem_base = V3_VAddr((void *)vm->mem_map.base_region.host_addr);
 
-    ctx = v3_chkpt_open_ctx(chkpt, NULL,"memory_img");
+    ctx = v3_chkpt_open_ctx(chkpt, "memory_img");
 
     if (!ctx) { 
        PrintError("Unable to open context to save memory\n");
        return -1;
     }
 
-    if (v3_chkpt_save(ctx, "memory_img", vm->mem_size, guest_mem_base) == -1) {
-       PrintError("Unable to load all of memory (requested=%llu, received=%llu)\n",(uint64_t)(vm->mem_size),ret);
+    if (v3_chkpt_save(ctx, "memory_img", vm->mem_size, guest_mem_base)) {
+       PrintError("Unable to save all of memory (requested=%llu, received=%llu)\n",(uint64_t)(vm->mem_size),ret);
        v3_chkpt_close_ctx(ctx);  
        return -1;
     }
     
     PrintDebug("Saving incremental memory.\n");
 
-    ctx = v3_chkpt_open_ctx(chkpt, NULL,"memory_bitmap_bits");
+    ctx = v3_chkpt_open_ctx(chkpt,"memory_bitmap_bits");
 
     if (!ctx) { 
        PrintError("Cannot open context for dirty memory bitmap\n");
     if (v3_chkpt_save(ctx,
                      "memory_bitmap_bits",
                      bitmap_num_bytes,
-                     mod_pgs_to_send->bits) == -1) {
+                     mod_pgs_to_send->bits)) {
        PrintError("Unable to write all of the dirty memory bitmap\n");
        v3_chkpt_close_ctx(ctx);
        return -1;
     for (i = 0; i < mod_pgs_to_send->num_bits; i++) {
         if (v3_bitmap_check(mod_pgs_to_send, i)) {
            // PrintDebug("Sending memory page %d.\n",i);
-            ctx = v3_chkpt_open_ctx(chkpt, NULL,"memory_page");
+            ctx = v3_chkpt_open_ctx(chkpt, "memory_page");
            if (!ctx) { 
                PrintError("Unable to open context to send memory page\n");
                return -1;
             if (v3_chkpt_save(ctx, 
                              "memory_page", 
                              page_size_bytes,
-                             guest_mem_base + (page_size_bytes * i)) == -1) {
+                             guest_mem_base + (page_size_bytes * i))) {
                PrintError("Unable to send a memory page\n");
                v3_chkpt_close_ctx(ctx);
                return -1;
 
     guest_mem_base = V3_VAddr((void *)vm->mem_map.base_region.host_addr);
 
-    ctx = v3_chkpt_open_ctx(chkpt, NULL,"memory_bitmap_bits");
+    ctx = v3_chkpt_open_ctx(chkpt, "memory_bitmap_bits");
 
     if (!ctx) { 
        PrintError("Cannot open context to receive memory bitmap\n");
     if (v3_chkpt_load(ctx,
                      "memory_bitmap_bits",
                      bitmap_num_bytes,
-                     mod_pgs->bits) == -1) {
+                     mod_pgs->bits)) {
        PrintError("Did not receive all of memory bitmap\n");
        v3_chkpt_close_ctx(ctx);
        return -1;
         if (v3_bitmap_check(mod_pgs, i)) {
             PrintDebug("Loading page %d\n", i);
             empty_bitmap = false;
-            ctx = v3_chkpt_open_ctx(chkpt, NULL,"memory_page");
+            ctx = v3_chkpt_open_ctx(chkpt, "memory_page");
            if (!ctx) { 
                PrintError("Cannot open context to receive memory page\n");
                return -1;
             if (v3_chkpt_load(ctx, 
                              "memory_page", 
                              page_size_bytes,
-                             guest_mem_base + (page_size_bytes * i)) == -1) {
+                             guest_mem_base + (page_size_bytes * i))) {
                PrintError("Did not receive all of memory page\n");
                v3_chkpt_close_ctx(ctx);
                return -1;
     extern v3_cpu_arch_t v3_mach_type;
     void * ctx = NULL;
     
-    ctx = v3_chkpt_open_ctx(chkpt, NULL, "header");
+    ctx = v3_chkpt_open_ctx(chkpt, "header");
     if (!ctx) { 
        PrintError("Cannot open context to save header\n");
        return -1;
     switch (v3_mach_type) {
        case V3_SVM_CPU:
        case V3_SVM_REV3_CPU: {
-           if (v3_chkpt_save(ctx, "header", strlen(svm_chkpt_header), svm_chkpt_header) == -1) { 
+           if (v3_chkpt_save(ctx, "header", strlen(svm_chkpt_header), svm_chkpt_header)) { 
                PrintError("Could not save all of SVM header\n");
                v3_chkpt_close_ctx(ctx);
                return -1;
        case V3_VMX_CPU:
        case V3_VMX_EPT_CPU:
        case V3_VMX_EPT_UG_CPU: {
-           if (v3_chkpt_save(ctx, "header", strlen(vmx_chkpt_header), vmx_chkpt_header) == -1) { 
+           if (v3_chkpt_save(ctx, "header", strlen(vmx_chkpt_header), vmx_chkpt_header)) { 
                PrintError("Could not save all of VMX header\n");
                v3_chkpt_close_ctx(ctx);
                return -1;
     extern v3_cpu_arch_t v3_mach_type;
     void * ctx = NULL;
     
-    ctx = v3_chkpt_open_ctx(chkpt, NULL, "header");
+    ctx = v3_chkpt_open_ctx(chkpt, "header");
 
     switch (v3_mach_type) {
        case V3_SVM_CPU:
        case V3_SVM_REV3_CPU: {
            char header[strlen(svm_chkpt_header) + 1];
         
-           if (v3_chkpt_load(ctx, "header", strlen(svm_chkpt_header), header) == -1) {
+           if (v3_chkpt_load(ctx, "header", strlen(svm_chkpt_header), header)) {
                PrintError("Could not load all of SVM header\n");
                v3_chkpt_close_ctx(ctx);
                return -1;
        case V3_VMX_EPT_UG_CPU: {
            char header[strlen(vmx_chkpt_header) + 1];
            
-           if (v3_chkpt_load(ctx, "header", strlen(vmx_chkpt_header), header) == -1) {
+           if (v3_chkpt_load(ctx, "header", strlen(vmx_chkpt_header), header)) {
                PrintError("Could not load all of VMX header\n");
                v3_chkpt_close_ctx(ctx);
                return -1;
 
     snprintf(key_name, 16, "guest_info%d", info->vcpu_id);
 
-    ctx = v3_chkpt_open_ctx(chkpt, NULL, key_name);
+    ctx = v3_chkpt_open_ctx(chkpt, key_name);
 
     if (!ctx) { 
        PrintError("Could not open context to load core\n");
        return -1;
     }
 
-    // These really need to have error checking
-
-    v3_chkpt_load_64(ctx, "RIP", &(info->rip));
+    V3_CHKPT_LOAD(ctx, "RIP", info->rip, loadfailout);
+    V3_CHKPT_LOAD(ctx, "GPRS", info->vm_regs, loadfailout);
 
-    V3_CHKPT_STD_LOAD(ctx, info->vm_regs);
+    V3_CHKPT_LOAD(ctx, "CR0", info->ctrl_regs.cr0, loadfailout);
+    V3_CHKPT_LOAD(ctx, "CR2", info->ctrl_regs.cr2, loadfailout);
+    V3_CHKPT_LOAD(ctx, "CR4", info->ctrl_regs.cr4, loadfailout);
+    V3_CHKPT_LOAD(ctx, "APIC_TPR", info->ctrl_regs.apic_tpr, loadfailout);
+    V3_CHKPT_LOAD(ctx, "RFLAGS", info->ctrl_regs.rflags, loadfailout);
+    V3_CHKPT_LOAD(ctx, "EFER", info->ctrl_regs.efer, loadfailout);
 
-    V3_CHKPT_STD_LOAD(ctx, info->ctrl_regs.cr0);
-    V3_CHKPT_STD_LOAD(ctx, info->ctrl_regs.cr2);
-    V3_CHKPT_STD_LOAD(ctx, info->ctrl_regs.cr4);
-    V3_CHKPT_STD_LOAD(ctx, info->ctrl_regs.apic_tpr);
-    V3_CHKPT_STD_LOAD(ctx, info->ctrl_regs.rflags);
-    V3_CHKPT_STD_LOAD(ctx, info->ctrl_regs.efer);
-
-    V3_CHKPT_STD_LOAD(ctx, info->dbg_regs);
-    V3_CHKPT_STD_LOAD(ctx, info->segments);
-    V3_CHKPT_STD_LOAD(ctx, info->shdw_pg_state.guest_cr3);
-    V3_CHKPT_STD_LOAD(ctx, info->shdw_pg_state.guest_cr0);
-    V3_CHKPT_STD_LOAD(ctx, info->shdw_pg_state.guest_efer);
+    V3_CHKPT_LOAD(ctx, "DBRS", info->dbg_regs, loadfailout);
+    V3_CHKPT_LOAD(ctx, "SEGS", info->segments, loadfailout);
+    V3_CHKPT_LOAD(ctx, "GUEST_CR3", info->shdw_pg_state.guest_cr3, loadfailout);
+    V3_CHKPT_LOAD(ctx, "GUEST_CRO", info->shdw_pg_state.guest_cr0, loadfailout);
+    V3_CHKPT_LOAD(ctx, "GUEST_EFER", info->shdw_pg_state.guest_efer, loadfailout);
 
     v3_chkpt_close_ctx(ctx);
 
            char key_name[16];
 
            snprintf(key_name, 16, "vmcb_data%d", info->vcpu_id);
-           ctx = v3_chkpt_open_ctx(chkpt, NULL, key_name);
+           ctx = v3_chkpt_open_ctx(chkpt, key_name);
 
            if (!ctx) { 
                PrintError("Could not open context to load SVM core\n");
                return -1;
            }
            
-           if (v3_svm_load_core(info, ctx) == -1) {
+           if (v3_svm_load_core(info, ctx) < 0 ) {
                PrintError("Failed to patch core %d\n", info->vcpu_id);
-               v3_chkpt_close_ctx(ctx);
-               return -1;
+               goto loadfailout;
            }
 
            v3_chkpt_close_ctx(ctx);
 
            snprintf(key_name, 16, "vmcs_data%d", info->vcpu_id);
 
-           ctx = v3_chkpt_open_ctx(chkpt, NULL, key_name);
+           ctx = v3_chkpt_open_ctx(chkpt, key_name);
 
            if (!ctx) { 
                PrintError("Could not open context to load VMX core\n");
            
            if (v3_vmx_load_core(info, ctx) < 0) {
                PrintError("VMX checkpoint failed\n");
-               v3_chkpt_close_ctx(ctx);
-               return -1;
+               goto loadfailout;
            }
 
            v3_chkpt_close_ctx(ctx);
            return -1;
     }
 
+    PrintDebug("Load of core succeeded\n");
+
     v3_print_guest_state(info);
 
     return 0;
+
+ loadfailout:
+    PrintError("Failed to load core due to bad context load\n");
+    v3_chkpt_close_ctx(ctx);
+    return -1;
+
 }
 
 
     void * ctx = NULL;
     char key_name[16];
 
-    memset(key_name, 0, 16);
+    PrintDebug("Saving core\n");
 
     v3_print_guest_state(info);
 
+    memset(key_name, 0, 16);
+
 
     snprintf(key_name, 16, "guest_info%d", info->vcpu_id);
 
-    ctx = v3_chkpt_open_ctx(chkpt, NULL, key_name);
+    ctx = v3_chkpt_open_ctx(chkpt, key_name);
     
     if (!ctx) { 
        PrintError("Unable to open context to save core\n");
     }
 
 
-    // Error checking of all this needs to happen
-    v3_chkpt_save_64(ctx, "RIP", &(info->rip));
+    V3_CHKPT_SAVE(ctx, "RIP", info->rip, savefailout);
+    V3_CHKPT_SAVE(ctx, "GPRS", info->vm_regs, savefailout);
 
-    V3_CHKPT_STD_SAVE(ctx, info->vm_regs);
+    V3_CHKPT_SAVE(ctx, "CR0", info->ctrl_regs.cr0, savefailout);
+    V3_CHKPT_SAVE(ctx, "CR2", info->ctrl_regs.cr2, savefailout);
+    V3_CHKPT_SAVE(ctx, "CR4", info->ctrl_regs.cr4, savefailout);
+    V3_CHKPT_SAVE(ctx, "APIC_TPR", info->ctrl_regs.apic_tpr, savefailout);
+    V3_CHKPT_SAVE(ctx, "RFLAGS", info->ctrl_regs.rflags, savefailout);
+    V3_CHKPT_SAVE(ctx, "EFER", info->ctrl_regs.efer, savefailout);
 
-    V3_CHKPT_STD_SAVE(ctx, info->ctrl_regs.cr0);
-    V3_CHKPT_STD_SAVE(ctx, info->ctrl_regs.cr2);
-    V3_CHKPT_STD_SAVE(ctx, info->ctrl_regs.cr4);
-    V3_CHKPT_STD_SAVE(ctx, info->ctrl_regs.apic_tpr);
-    V3_CHKPT_STD_SAVE(ctx, info->ctrl_regs.rflags);
-    V3_CHKPT_STD_SAVE(ctx, info->ctrl_regs.efer);
-
-    V3_CHKPT_STD_SAVE(ctx, info->dbg_regs);
-    V3_CHKPT_STD_SAVE(ctx, info->segments);
-    V3_CHKPT_STD_SAVE(ctx, info->shdw_pg_state.guest_cr3);
-    V3_CHKPT_STD_SAVE(ctx, info->shdw_pg_state.guest_cr0);
-    V3_CHKPT_STD_SAVE(ctx, info->shdw_pg_state.guest_efer);
+    V3_CHKPT_SAVE(ctx, "DBRS", info->dbg_regs, savefailout);
+    V3_CHKPT_SAVE(ctx, "SEGS", info->segments, savefailout);
+    V3_CHKPT_SAVE(ctx, "GUEST_CR3", info->shdw_pg_state.guest_cr3, savefailout);
+    V3_CHKPT_SAVE(ctx, "GUEST_CRO", info->shdw_pg_state.guest_cr0, savefailout);
+    V3_CHKPT_SAVE(ctx, "GUEST_EFER", info->shdw_pg_state.guest_efer, savefailout);
 
     v3_chkpt_close_ctx(ctx);
 
        case V3_SVM_CPU:
        case V3_SVM_REV3_CPU: {
            char key_name[16];
-           void * ctx = NULL;
            
            snprintf(key_name, 16, "vmcb_data%d", info->vcpu_id);
            
-           ctx = v3_chkpt_open_ctx(chkpt, NULL, key_name);
+           ctx = v3_chkpt_open_ctx(chkpt, key_name);
 
            if (!ctx) { 
                PrintError("Could not open context to store SVM core\n");
                return -1;
            }
            
-           if (v3_svm_save_core(info, ctx) == -1) {
+           if (v3_svm_save_core(info, ctx) < 0) {
                PrintError("VMCB Unable to be written\n");
-               v3_chkpt_close_ctx(ctx);
-               return -1;
+               goto savefailout;
            }
            
            v3_chkpt_close_ctx(ctx);
        case V3_VMX_EPT_CPU:
        case V3_VMX_EPT_UG_CPU: {
            char key_name[16];
-           void * ctx = NULL;
 
            snprintf(key_name, 16, "vmcs_data%d", info->vcpu_id);
            
-           ctx = v3_chkpt_open_ctx(chkpt, NULL, key_name);
+           ctx = v3_chkpt_open_ctx(chkpt, key_name);
            
            if (!ctx) { 
                PrintError("Could not open context to store VMX core\n");
 
            if (v3_vmx_save_core(info, ctx) == -1) {
                PrintError("VMX checkpoint failed\n");
-               v3_chkpt_close_ctx(ctx);
-               return -1;
+               goto savefailout;
            }
 
            v3_chkpt_close_ctx(ctx);
        default:
            PrintError("Invalid CPU Type (%d)\n", v3_mach_type);
            return -1;
+           
     }
     
     return 0;
+
+ savefailout:
+    PrintError("Failed to save core due to bad context save\n");
+    v3_chkpt_close_ctx(ctx);
+    return -1;
+
 }
 
 
     }
 
     if ((ret = load_memory(vm, chkpt)) == -1) {
-       PrintError("Unable to save memory\n");
+       PrintError("Unable to load memory\n");
        goto out;
     }