Similar to alternatives, save the original bytes when patching retpolines
the first time.

Signed-off-by: David Kaplan <david.kaplan@amd.com>
---
 arch/x86/include/asm/alternative.h |  5 ++++
 arch/x86/include/asm/module.h      |  2 ++
 arch/x86/kernel/alternative.c      | 37 +++++++++++++++++++++++++++++-
 3 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/arch/x86/include/asm/alternative.h b/arch/x86/include/asm/alternative.h
index 3ee781d61927..24a4afbf163b 100644
--- a/arch/x86/include/asm/alternative.h
+++ b/arch/x86/include/asm/alternative.h
@@ -102,6 +102,11 @@ struct alt_site {
 	u8 len;
 };
 
+struct retpoline_site {
+	u8 bytes[6];
+	u8 len;
+} __packed;
+
 extern void alternative_instructions(void);
 extern void apply_alternatives(struct alt_instr *start, struct alt_instr *end,
 		struct module *mod);
diff --git a/arch/x86/include/asm/module.h b/arch/x86/include/asm/module.h
index 2bb602f99154..d0c39b921408 100644
--- a/arch/x86/include/asm/module.h
+++ b/arch/x86/include/asm/module.h
@@ -21,6 +21,8 @@ struct mod_arch_specific {
 	struct its_array its_pages;
 #ifdef CONFIG_DYNAMIC_MITIGATIONS
 	struct alt_site *alt_sites;
+	struct retpoline_site *retpoline_sites;
+	int num_retpoline_sites;
 #endif
 };
 
diff --git a/arch/x86/kernel/alternative.c b/arch/x86/kernel/alternative.c
index 8037076e9301..a02dc6bfb696 100644
--- a/arch/x86/kernel/alternative.c
+++ b/arch/x86/kernel/alternative.c
@@ -302,6 +302,8 @@ static bool __maybe_unused repatch_in_progress;
 
 #ifdef CONFIG_DYNAMIC_MITIGATIONS
 static struct alt_site *alt_sites;
+static struct retpoline_site *retpoline_sites;
+static int num_retpoline_sites;
 
 /* Do not patch __init text addresses when repatching */
 static bool should_patch(void *addr, struct module *mod)
@@ -1036,8 +1038,36 @@ static int patch_retpoline(void *addr, struct insn *insn, u8 *bytes)
 void __init_or_module noinline apply_retpolines(s32 *start, s32 *end, struct module *mod)
 {
 	s32 *s;
+	u32 idx = 0;
+	struct retpoline_site *save_site = NULL;
 
-	for (s = start; s < end; s++) {
+#ifdef CONFIG_DYNAMIC_MITIGATIONS
+	u32 size = ((u64)end - (u64)start)/4;
+
+	/* ITS code needs the save_site pointer even on re-patch. */
+	if (!mod) {
+		if (!retpoline_sites) {
+			retpoline_sites = kcalloc(size,	sizeof(struct retpoline_site), GFP_KERNEL);
+			if (WARN_ON(!retpoline_sites))
+				return;
+		}
+
+		save_site = retpoline_sites;
+		num_retpoline_sites = size;
+	} else {
+		if (!mod->arch.retpoline_sites) {
+			mod->arch.retpoline_sites = kcalloc(size, sizeof(struct retpoline_site),
+							    GFP_KERNEL);
+			if (WARN_ON(!mod->arch.retpoline_sites))
+				return;
+		}
+
+		save_site = mod->arch.retpoline_sites;
+		mod->arch.num_retpoline_sites = size;
+	}
+#endif
+
+	for (s = start; s < end; s++, idx++) {
 		void *addr = (void *)s + *s;
 		struct insn insn;
 		int len, ret;
@@ -1085,6 +1115,11 @@ void __init_or_module noinline apply_retpolines(s32 *start, s32 *end, struct mod
 			addr, addr, insn.length,
 			addr + insn.length + insn.immediate.value);
 
+		if (IS_ENABLED(CONFIG_DYNAMIC_MITIGATIONS) && save_site) {
+			save_site[idx].len = insn.length;
+			memcpy(save_site[idx].bytes, addr, insn.length);
+		}
+
 		len = patch_retpoline(addr, &insn, bytes);
 		if (len == insn.length) {
 			optimize_nops(addr, bytes, len);
-- 
2.34.1