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>

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[];

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

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

arch/s390/kernel/early_nobss.c

@@ -13,7 +13,6 @@
 #include <linux/string.h>
 #include <asm/sections.h>
 #include <asm/lowcore.h>
-#include <asm/setup.h>
 #include <asm/timex.h>
 #include "entry.h"
 
@@ -32,26 +31,6 @@ static void __init reset_tod_clock(void)
 	S390_lowcore.last_update_clock = TOD_UNIX_EPOCH;
 }
 
-static void __init rescue_initrd(void)
-{
-	unsigned long min_initrd_addr = (unsigned long) _end + (4UL << 20);
-
-	/*
-	 * Just like in case of IPL from VM reader we make sure there is a
-	 * gap of 4MB between end of kernel and start of initrd.
-	 * That way we can also be sure that saving an NSS will succeed,
-	 * which however only requires different segments.
-	 */
-	if (!IS_ENABLED(CONFIG_BLK_DEV_INITRD))
-		return;
-	if (!INITRD_START || !INITRD_SIZE)
-		return;
-	if (INITRD_START >= min_initrd_addr)
-		return;
-	memmove((void *) min_initrd_addr, (void *) INITRD_START, INITRD_SIZE);
-	INITRD_START = min_initrd_addr;
-}
-
 static void __init clear_bss_section(void)
 {
 	memset(__bss_start, 0, __bss_stop - __bss_start);
@@ -60,6 +39,5 @@ static void __init clear_bss_section(void)
 void __init startup_init_nobss(void)
 {
 	reset_tod_clock();
-	rescue_initrd();
 	clear_bss_section();
 }

arch/s390/kernel/vmlinux.lds.S

@@ -154,6 +154,7 @@ SECTIONS
 		QUAD(_stext)			/* default_lma */
 		QUAD(startup_continue)		/* entry */
 		QUAD(__bss_start - _stext)	/* image_size */
+		QUAD(__bss_stop - __bss_start)	/* bss_size */
 	}
 
 	/* Debugging sections.	*/