s390/kasan: add option for 4-level paging support

By default 3-level paging is used when the kernel is compiled with
kasan support. Add 4-level paging option to support systems with more
then 3TB of physical memory and to cover 4-level paging specific code
with kasan as well.
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/Kconfig

@@ -59,6 +59,7 @@ config ARCH_SUPPORTS_UPROBES
 config KASAN_SHADOW_OFFSET
 	hex
 	depends on KASAN
+	default 0x18000000000000 if KASAN_S390_4_LEVEL_PAGING
 	default 0x30000000000
 
 config S390

arch/s390/include/asm/kasan.h

@@ -7,8 +7,13 @@
 #ifdef CONFIG_KASAN
 
 #define KASAN_SHADOW_SCALE_SHIFT 3
+#ifdef CONFIG_KASAN_S390_4_LEVEL_PAGING
+#define KASAN_SHADOW_SIZE						       \
+	(_AC(1, UL) << (_REGION1_SHIFT - KASAN_SHADOW_SCALE_SHIFT))
+#else
 #define KASAN_SHADOW_SIZE						       \
 	(_AC(1, UL) << (_REGION2_SHIFT - KASAN_SHADOW_SCALE_SHIFT))
+#endif
 #define KASAN_SHADOW_OFFSET	_AC(CONFIG_KASAN_SHADOW_OFFSET, UL)
 #define KASAN_SHADOW_START	KASAN_SHADOW_OFFSET
 #define KASAN_SHADOW_END	(KASAN_SHADOW_START + KASAN_SHADOW_SIZE)

arch/s390/kernel/setup.c

@@ -535,7 +535,9 @@ static void __init setup_memory_end(void)
 	/* Choose kernel address space layout: 3 or 4 levels. */
 	vmalloc_size = VMALLOC_END ?: (128UL << 30) - MODULES_LEN;
 	if (IS_ENABLED(CONFIG_KASAN)) {
-		vmax = _REGION2_SIZE; /* 3-level kernel page table */
+		vmax = IS_ENABLED(CONFIG_KASAN_S390_4_LEVEL_PAGING)
+			   ? _REGION1_SIZE
+			   : _REGION2_SIZE;
 	} else {
 		tmp = (memory_end ?: max_physmem_end) / PAGE_SIZE;
 		tmp = tmp * (sizeof(struct page) + PAGE_SIZE);

arch/s390/mm/kasan_init.c

@@ -252,12 +252,23 @@ void __init kasan_early_init(void)
 		pgt_prot &= ~_PAGE_NOEXEC;
 	pte_z = __pte(__pa(kasan_zero_page) | pgt_prot);
 
+	if (IS_ENABLED(CONFIG_KASAN_S390_4_LEVEL_PAGING)) {
+		/* 4 level paging */
+		BUILD_BUG_ON(!IS_ALIGNED(KASAN_SHADOW_START, P4D_SIZE));
+		BUILD_BUG_ON(!IS_ALIGNED(KASAN_SHADOW_END, P4D_SIZE));
+		crst_table_init((unsigned long *)early_pg_dir,
+				_REGION2_ENTRY_EMPTY);
+		untracked_mem_end = vmax = _REGION1_SIZE;
+		asce_type = _ASCE_TYPE_REGION2;
+	} else {
 		/* 3 level paging */
 		BUILD_BUG_ON(!IS_ALIGNED(KASAN_SHADOW_START, PUD_SIZE));
 		BUILD_BUG_ON(!IS_ALIGNED(KASAN_SHADOW_END, PUD_SIZE));
-	crst_table_init((unsigned long *)early_pg_dir, _REGION3_ENTRY_EMPTY);
+		crst_table_init((unsigned long *)early_pg_dir,
+				_REGION3_ENTRY_EMPTY);
 		untracked_mem_end = vmax = _REGION2_SIZE;
 		asce_type = _ASCE_TYPE_REGION3;
+	}
 
 	/* init kasan zero shadow */
 	crst_table_init((unsigned long *)kasan_zero_p4d, p4d_val(p4d_z));

lib/Kconfig.kasan

@@ -57,6 +57,15 @@ config KASAN_INLINE
 
 endchoice
 
+config KASAN_S390_4_LEVEL_PAGING
+	bool "KASan: use 4-level paging"
+	depends on KASAN && S390
+	help
+	  Compiling the kernel with KASan disables automatic 3-level vs
+	  4-level paging selection. 3-level paging is used by default (up
+	  to 3TB of RAM with KASan enabled). This options allows to force
+	  4-level paging instead.
+
 config TEST_KASAN
 	tristate "Module for testing kasan for bug detection"
 	depends on m && KASAN