2 files changed, 21 insertions, 0 deletions

diff --git a/kernel/power/Kconfig b/kernel/power/Kconfig
index 3744c594b19b..e01e6899592c 100644
--- a/kernel/power/Kconfig
+++ b/kernel/power/Kconfig
@@ -65,6 +65,9 @@ config HIBERNATION
 
           For more information take a look at <file:Documentation/power/swsusp.txt>.
 
+config ARCH_SAVE_PAGE_KEYS
+        bool
+
 config PM_STD_PARTITION
         string "Default resume partition"
         depends on HIBERNATION
diff --git a/kernel/power/snapshot.c b/kernel/power/snapshot.c
index 06efa54f93d6..cbe2c1441392 100644
--- a/kernel/power/snapshot.c
+++ b/kernel/power/snapshot.c
@@ -1339,6 +1339,9 @@ int hibernate_preallocate_memory(void)
         count += highmem;
         count -= totalreserve_pages;
 
+        /* Add number of pages required for page keys (s390 only). */
+        size += page_key_additional_pages(saveable);
+
         /* Compute the maximum number of saveable pages to leave in memory. */
         max_size = (count - (size + PAGES_FOR_IO)) / 2
                         - 2 * DIV_ROUND_UP(reserved_size, PAGE_SIZE);
@@ -1662,6 +1665,8 @@ pack_pfns(unsigned long *buf, struct memory_bitmap *bm)
                 buf[j] = memory_bm_next_pfn(bm);
                 if (unlikely(buf[j] == BM_END_OF_MAP))
                         break;
+                /* Save page key for data page (s390 only). */
+                page_key_read(buf + j);
         }
 }
 
@@ -1821,6 +1826,9 @@ static int unpack_orig_pfns(unsigned long *buf, struct memory_bitmap *bm)
                 if (unlikely(buf[j] == BM_END_OF_MAP))
                         break;
 
+                /* Extract and buffer page key for data page (s390 only). */
+                page_key_memorize(buf + j);
+
                 if (memory_bm_pfn_present(bm, buf[j]))
                         memory_bm_set_bit(bm, buf[j]);
                 else
@@ -2223,6 +2231,11 @@ int snapshot_write_next(struct snapshot_handle *handle)
                 if (error)
                         return error;
 
+                /* Allocate buffer for page keys. */
+                error = page_key_alloc(nr_copy_pages);
+                if (error)
+                        return error;
+
         } else if (handle->cur <= nr_meta_pages + 1) {
                 error = unpack_orig_pfns(buffer, &copy_bm);
                 if (error)
@@ -2243,6 +2256,8 @@ int snapshot_write_next(struct snapshot_handle *handle)
                 }
         } else {
                 copy_last_highmem_page();
+                /* Restore page key for data page (s390 only). */
+                page_key_write(handle->buffer);
                 handle->buffer = get_buffer(&orig_bm, &ca);
                 if (IS_ERR(handle->buffer))
                         return PTR_ERR(handle->buffer);
@@ -2264,6 +2279,9 @@ int snapshot_write_next(struct snapshot_handle *handle)
 void snapshot_write_finalize(struct snapshot_handle *handle)
 {
         copy_last_highmem_page();
+        /* Restore page key for data page (s390 only). */
+        page_key_write(handle->buffer);
+        page_key_free();
         /* Free only if we have loaded the image entirely */
         if (handle->cur > 1 && handle->cur > nr_meta_pages + nr_copy_pages) {
                 memory_bm_free(&orig_bm, PG_UNSAFE_CLEAR);