1 files changed, 9 insertions, 24 deletions

diff --git a/arch/s390/boot/kaslr.c b/arch/s390/boot/kaslr.c
index d844a5ef9089..0dd48fbdbaa4 100644
--- a/arch/s390/boot/kaslr.c
+++ b/arch/s390/boot/kaslr.c
@@ -7,6 +7,7 @@
 #include <asm/cpacf.h>
 #include <asm/timex.h>
 #include <asm/sclp.h>
+#include <asm/kasan.h>
 #include "compressed/decompressor.h"
 #include "boot.h"
 
@@ -176,8 +177,14 @@ unsigned long get_random_base(unsigned long safe_addr)
 	unsigned long kasan_needs;
 	int i;
 
-	if (memory_end_set)
-		memory_limit = min(memory_limit, memory_end);
+	memory_limit = min(memory_limit, ident_map_size);
+
+	/*
+	 * Avoid putting kernel in the end of physical memory
+	 * which kasan will use for shadow memory and early pgtable
+	 * mapping allocations.
+	 */
+	memory_limit -= kasan_estimate_memory_needs(memory_limit);
 
 	if (IS_ENABLED(CONFIG_BLK_DEV_INITRD) && INITRD_START && INITRD_SIZE) {
 		if (safe_addr < INITRD_START + INITRD_SIZE)
@@ -185,28 +192,6 @@ unsigned long get_random_base(unsigned long safe_addr)
 	}
 	safe_addr = ALIGN(safe_addr, THREAD_SIZE);
 
-	if ((IS_ENABLED(CONFIG_KASAN))) {
-		/*
-		 * Estimate kasan memory requirements, which it will reserve
-		 * at the very end of available physical memory. To estimate
-		 * that, we take into account that kasan would require
-		 * 1/8 of available physical memory (for shadow memory) +
-		 * creating page tables for the whole memory + shadow memory
-		 * region (1 + 1/8). To keep page tables estimates simple take
-		 * the double of combined ptes size.
-		 */
-		memory_limit = get_mem_detect_end();
-		if (memory_end_set && memory_limit > memory_end)
-			memory_limit = memory_end;
-
-		/* for shadow memory */
-		kasan_needs = memory_limit / 8;
-		/* for paging structures */
-		kasan_needs += (memory_limit + kasan_needs) / PAGE_SIZE /
-			       _PAGE_ENTRIES * _PAGE_TABLE_SIZE * 2;
-		memory_limit -= kasan_needs;
-	}
-
 	kernel_size = vmlinux.image_size + vmlinux.bss_size;
 	if (safe_addr + kernel_size > memory_limit)
 		return 0;