3 files changed, 39 insertions, 16 deletions

diff --git a/arch/s390/boot/compressed/decompressor.h b/arch/s390/boot/compressed/decompressor.h
index 011cbb6e0e08..90d382d501d7 100644
--- a/arch/s390/boot/compressed/decompressor.h
+++ b/arch/s390/boot/compressed/decompressor.h
@@ -7,11 +7,13 @@ static inline void *decompress_kernel(void) {}
 #else
 void *decompress_kernel(void);
 #endif
+unsigned long mem_safe_offset(void);
 
 struct vmlinux_info {
 	unsigned long default_lma;
 	void (*entry)(void);
 	unsigned long image_size;	/* does not include .bss */
+	unsigned long bss_size;		/* uncompressed image .bss size */
 };
 
 extern char _vmlinux_info[];
diff --git a/arch/s390/boot/compressed/misc.c b/arch/s390/boot/compressed/misc.c
index 5dcf34e31f8d..b773f81f5bff 100644
--- a/arch/s390/boot/compressed/misc.c
+++ b/arch/s390/boot/compressed/misc.c
@@ -83,25 +83,22 @@ static void error(char *x)
 	asm volatile("lpsw %0" : : "Q" (psw));
 }
 
-void *decompress_kernel(void)
-{
-	void *output, *kernel_end;
-
-	output = (void *) ALIGN((unsigned long) _end + HEAP_SIZE, PAGE_SIZE);
-	kernel_end = output + vmlinux.image_size;
+#define decompress_offset ALIGN((unsigned long)_end + HEAP_SIZE, PAGE_SIZE)
 
-#ifdef CONFIG_BLK_DEV_INITRD
+unsigned long mem_safe_offset(void)
+{
 	/*
-	 * Move the initrd right behind the end of the decompressed
-	 * kernel image. This also prevents initrd corruption caused by
-	 * bss clearing since kernel_end will always be located behind the
-	 * current bss section..
+	 * due to 4MB HEAD_SIZE for bzip2
+	 * 'decompress_offset + vmlinux.image_size' could be larger than
+	 * kernel at final position + its .bss, so take the larger of two
 	 */
-	if (INITRD_START && INITRD_SIZE && kernel_end > (void *) INITRD_START) {
-		memmove(kernel_end, (void *) INITRD_START, INITRD_SIZE);
-		INITRD_START = (unsigned long) kernel_end;
-	}
-#endif
+	return max(decompress_offset + vmlinux.image_size,
+		   vmlinux.default_lma + vmlinux.image_size + vmlinux.bss_size);
+}
+
+void *decompress_kernel(void)
+{
+	void *output = (void *)decompress_offset;
 
 	__decompress(_compressed_start, _compressed_end - _compressed_start,
 		     NULL, NULL, output, 0, NULL, error);
diff --git a/arch/s390/boot/startup.c b/arch/s390/boot/startup.c
index 474dee84d8a8..5aeac7564e67 100644
--- a/arch/s390/boot/startup.c
+++ b/arch/s390/boot/startup.c
@@ -1,12 +1,36 @@
 // SPDX-License-Identifier: GPL-2.0
 #include <linux/string.h>
+#include <asm/setup.h>
 #include "compressed/decompressor.h"
 #include "boot.h"
 
+#ifdef CONFIG_KERNEL_UNCOMPRESSED
+unsigned long mem_safe_offset(void)
+{
+	return vmlinux.default_lma + vmlinux.image_size + vmlinux.bss_size;
+}
+#endif
+
+static void rescue_initrd(void)
+{
+	unsigned long min_initrd_addr;
+
+	if (!IS_ENABLED(CONFIG_BLK_DEV_INITRD))
+		return;
+	if (!INITRD_START || !INITRD_SIZE)
+		return;
+	min_initrd_addr = mem_safe_offset();
+	if (min_initrd_addr <= INITRD_START)
+		return;
+	memmove((void *)min_initrd_addr, (void *)INITRD_START, INITRD_SIZE);
+	INITRD_START = min_initrd_addr;
+}
+
 void startup_kernel(void)
 {
 	void *img;
 
+	rescue_initrd();
 	if (!IS_ENABLED(CONFIG_KERNEL_UNCOMPRESSED)) {
 		img = decompress_kernel();
 		memmove((void *)vmlinux.default_lma, img, vmlinux.image_size);