s390: rescue initrd as early as possible

To avoid multi-stage initrd rescue operation and to simplify
assumptions during early memory allocations move initrd at some final
safe destination as early as possible. This would also allow us to
drop .bss usage restrictions for some files.

Reviewed-by: Martin Schwidefsky <schwidefsky@de.ibm.com>
Signed-off-by: Vasily Gorbik <gor@linux.ibm.com>
Signed-off-by: Martin Schwidefsky <schwidefsky@de.ibm.com>

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);