/* Defined in arch/x86_64/kernel/suspend_asm.S */
 extern int restore_image(void);
 
+/*
+ * Address to jump to in the last phase of restore in order to get to the image
+ * kernel's text (this value is passed in the image header).
+ */
+unsigned long restore_jump_address;
+
 pgd_t *temp_level4_pgt;
 
+void *relocated_restore_code;
+
 static int res_phys_pud_init(pud_t *pud, unsigned long address, unsigned long end)
 {
        long i, j;
 
                        if (paddr >= end)
                                break;
-                       pe = _PAGE_NX | _PAGE_PSE | _KERNPG_TABLE | paddr;
+                       pe = __PAGE_KERNEL_LARGE_EXEC | paddr;
                        pe &= __supported_pte_mask;
                        set_pmd(pmd, __pmd(pe));
                }
        /* We have got enough memory and from now on we cannot recover */
        if ((error = set_up_temporary_mappings()))
                return error;
+
+       relocated_restore_code = (void *)get_safe_page(GFP_ATOMIC);
+       if (!relocated_restore_code)
+               return -ENOMEM;
+       memcpy(relocated_restore_code, &core_restore_code,
+              &restore_registers - &core_restore_code);
+
        restore_image();
        return 0;
 }
        unsigned long nosave_end_pfn = PAGE_ALIGN(__pa_symbol(&__nosave_end)) >> PAGE_SHIFT;
        return (pfn >= nosave_begin_pfn) && (pfn < nosave_end_pfn);
 }
+
+struct restore_data_record {
+       unsigned long jump_address;
+       unsigned long control;
+};
+
+#define RESTORE_MAGIC  0x0123456789ABCDEFUL
+
+/**
+ *     arch_hibernation_header_save - populate the architecture specific part
+ *             of a hibernation image header
+ *     @addr: address to save the data at
+ */
+int arch_hibernation_header_save(void *addr, unsigned int max_size)
+{
+       struct restore_data_record *rdr = addr;
+
+       if (max_size < sizeof(struct restore_data_record))
+               return -EOVERFLOW;
+       rdr->jump_address = restore_jump_address;
+       rdr->control = (restore_jump_address ^ RESTORE_MAGIC);
+       return 0;
+}
+
+/**
+ *     arch_hibernation_header_restore - read the architecture specific data
+ *             from the hibernation image header
+ *     @addr: address to read the data from
+ */
+int arch_hibernation_header_restore(void *addr)
+{
+       struct restore_data_record *rdr = addr;
+
+       restore_jump_address = rdr->jump_address;
+       return (rdr->control == (restore_jump_address ^ RESTORE_MAGIC)) ?
+                       0 : -EINVAL;
+}
 #endif /* CONFIG_HIBERNATION */
 
  *
  * Distribute under GPLv2.
  *
- * swsusp_arch_resume may not use any stack, nor any variable that is
- * not "NoSave" during copying pages:
+ * swsusp_arch_resume must not use any stack or any nonlocal variables while
+ * copying pages:
  *
  * Its rewriting one kernel image with another. What is stack in "old"
  * image could very well be data page in "new" image, and overwriting
        movq %r15, saved_context_r15(%rip)
        pushfq ; popq saved_context_eflags(%rip)
 
+       /* save the address of restore_registers */
+       movq    $restore_registers, %rax
+       movq    %rax, restore_jump_address(%rip)
+
        call swsusp_save
        ret
 
        movq    %rcx, %cr3;
        movq    %rax, %cr4;  # turn PGE back on
 
+       /* prepare to jump to the image kernel */
+       movq    restore_jump_address(%rip), %rax
+
+       /* prepare to copy image data to their original locations */
        movq    restore_pblist(%rip), %rdx
+       movq    relocated_restore_code(%rip), %rcx
+       jmpq    *%rcx
+
+       /* code below has been relocated to a safe page */
+ENTRY(core_restore_code)
 loop:
        testq   %rdx, %rdx
        jz      done
        /* get addresses from the pbe and copy the page */
        movq    pbe_address(%rdx), %rsi
        movq    pbe_orig_address(%rdx), %rdi
-       movq    $512, %rcx
+       movq    $(PAGE_SIZE >> 3), %rcx
        rep
        movsq
 
        movq    pbe_next(%rdx), %rdx
        jmp     loop
 done:
+       /* jump to the restore_registers address from the image header */
+       jmpq    *%rax
+       /*
+        * NOTE: This assumes that the boot kernel's text mapping covers the
+        * image kernel's page containing restore_registers and the address of
+        * this page is the same as in the image kernel's text mapping (it
+        * should always be true, because the text mapping is linear, starting
+        * from 0, and is supposed to cover the entire kernel text for every
+        * kernel).
+        *
+        * code below belongs to the image kernel
+        */
+
+ENTRY(restore_registers)
        /* go back to the original page tables */
        movq    $(init_level4_pgt - __START_KERNEL_map), %rax
        addq    phys_base(%rip), %rax
        movq    %rcx, %cr3
        movq    %rax, %cr4;  # turn PGE back on
 
-       movl    $24, %eax
-       movl    %eax, %ds
-
        movq saved_context_esp(%rip), %rsp
        movq saved_context_ebp(%rip), %rbp
-       /* Don't restore %rax, it must be 0 anyway */
+       /* restore GPRs (we don't restore %rax, it must be 0 anyway) */
        movq saved_context_ebx(%rip), %rbx
        movq saved_context_ecx(%rip), %rcx
        movq saved_context_edx(%rip), %rdx
 
        xorq    %rax, %rax
 
+       /* tell the hibernation core that we've just restored the memory */
+       movq    %rax, in_suspend(%rip)
+
        ret