[PATCH] MIPS: relocatable: Provide kaslr_offset() to get the kernel offset

[Jinyang He <hejinyang@loongson.cn>]

Use kimage_vaddr to indicate kernel start address. Provide kaslr_offset()
to get the kernel offset when KASLR is enabled. Error may occur before
update_kimage_vaddr(), so put it at the end of the offset branch.

Fixes: a307a4ce9ecd ("MIPS: Loongson64: Add KASLR support")
Reported-by: kernel test robot <lkp@intel.com>
Signed-off-by: Jinyang He <hejinyang@loongson.cn>
---
 arch/mips/include/asm/page.h |  6 ++++++
 arch/mips/kernel/relocate.c  | 12 ++++++++++++
 arch/mips/kernel/setup.c     |  3 +++
 3 files changed, 21 insertions(+)

diff --git a/arch/mips/include/asm/page.h b/arch/mips/include/asm/page.h
index 6a77bc4..9429520 100644
--- a/arch/mips/include/asm/page.h
+++ b/arch/mips/include/asm/page.h
@@ -255,6 +255,12 @@ extern bool __virt_addr_valid(const volatile void *kaddr);
 
 #define VM_DATA_DEFAULT_FLAGS	VM_DATA_FLAGS_TSK_EXEC
 
+extern unsigned long kimage_vaddr;
+static inline unsigned long kaslr_offset(void)
+{
+	return kimage_vaddr - VMLINUX_LOAD_ADDRESS;
+}
+
 #include <asm-generic/memory_model.h>
 #include <asm-generic/getorder.h>
 
diff --git a/arch/mips/kernel/relocate.c b/arch/mips/kernel/relocate.c
index c643c81..2cbc3b1 100644
--- a/arch/mips/kernel/relocate.c
+++ b/arch/mips/kernel/relocate.c
@@ -300,6 +300,15 @@ static inline int __init relocation_addr_valid(void *loc_new)
 	return 1;
 }
 
+static inline void __init update_kimage_vaddr(unsigned long value,
+					      unsigned long *addr, long offset)
+{
+	unsigned long new_value = value + offset;
+	unsigned long *new_addr = (void *)addr + offset;
+
+	*new_addr = new_value;
+}
+
 #if defined(CONFIG_USE_OF)
 void __weak *plat_get_fdt(void)
 {
@@ -410,6 +419,9 @@ void *__init relocate_kernel(void)
 
 		/* Return the new kernel's entry point */
 		kernel_entry = RELOCATED(start_kernel);
+
+		/* Error may occur before, so keep it at last */
+		update_kimage_vaddr(kimage_vaddr, &kimage_vaddr, offset);
 	}
 out:
 	return kernel_entry;
diff --git a/arch/mips/kernel/setup.c b/arch/mips/kernel/setup.c
index 7e1f8e2..ad37b24 100644
--- a/arch/mips/kernel/setup.c
+++ b/arch/mips/kernel/setup.c
@@ -84,6 +84,9 @@ static struct resource code_resource = { .name = "Kernel code", };
 static struct resource data_resource = { .name = "Kernel data", };
 static struct resource bss_resource = { .name = "Kernel bss", };
 
+unsigned long kimage_vaddr __initdata = VMLINUX_LOAD_ADDRESS;
+EXPORT_SYMBOL(kimage_vaddr);
+
 static void *detect_magic __initdata = detect_memory_region;
 
 #ifdef CONFIG_MIPS_AUTO_PFN_OFFSET
-- 
2.1.0