@@ -7,6 +7,20 @@
 #ifndef _ASM_POWERPC_LIVEPATCH_H
 #define _ASM_POWERPC_LIVEPATCH_H
 
+#ifdef CONFIG_LIVEPATCH_64
+#define LIVEPATCH_STACK_MAGIC_OFFSET	8
+#define LIVEPATCH_STACK_LR_OFFSET	16
+#define LIVEPATCH_STACK_TOC_OFFSET	24
+
+#if defined(CONFIG_PPC_FTRACE_OUT_OF_LINE) && defined(CONFIG_DYNAMIC_FTRACE_WITH_DIRECT_CALLS)
+#define LIVEPATCH_STACK_FRAME_SIZE	32	/* Allocate 4 x 8 bytes (to save new NIP as well) */
+#define LIVEPATCH_STACK_NIP_OFFSET	32
+#else
+#define LIVEPATCH_STACK_FRAME_SIZE	24	/* Allocate 3 x 8 bytes */
+#endif
+#endif
+
+#ifndef __ASSEMBLY__
 #include <linux/sched.h>
 #include <linux/sched/task_stack.h>
 

@@ -20,4 +34,5 @@ static inline void klp_init_thread_info(struct task_struct *p)
 static inline void klp_init_thread_info(struct task_struct *p) { }
 #endif
 
+#endif /* !__ASSEMBLY__ */
 #endif /* _ASM_POWERPC_LIVEPATCH_H */

@@ -8,6 +8,7 @@
 #include <asm/ppc_asm.h>
 #include <asm/asm-offsets.h>
 #include <asm/ftrace.h>
+#include <asm/livepatch.h>
 #include <asm/ppc-opcode.h>
 #include <asm/thread_info.h>
 #include <asm/bug.h>

@@ -244,6 +245,8 @@
 	/* jump after _mcount site */
 #ifdef CONFIG_PPC_FTRACE_OUT_OF_LINE
 #ifdef CONFIG_DYNAMIC_FTRACE_WITH_DIRECT_CALLS
+	/* For direct_call, set cr0.eq bit only if livepatch is active */
+	crclr	4*cr0+eq
 	bnectr	cr1
 #endif
 	/*

@@ -306,10 +309,14 @@ ftrace_no_trace:
 	mtctr	r12
 	REST_GPRS(11, 12, r1)
 	addi	r1, r1, SWITCH_FRAME_SIZE+STACK_FRAME_MIN_SIZE
+	/* For direct_call, set cr0.eq bit only if livepatch is active */
+	crclr	4*cr0+eq
 	bctr
 .Lftrace_direct_call:
 	mtctr	r12
 	addi	r1, r1, SWITCH_FRAME_SIZE+STACK_FRAME_MIN_SIZE
+	/* For direct_call, set cr0.eq bit only if livepatch is active */
+	crclr	4*cr0+eq
 	bctr
 SYM_FUNC_START(ftrace_stub_direct_tramp)
 	blr

@@ -340,25 +347,72 @@ SYM_FUNC_END(ftrace_stub_direct_tramp)
 livepatch_handler:
 	ld	r12, PACA_THREAD_INFO(r13)
 
-	/* Allocate 3 x 8 bytes */
 	ld	r11, TI_livepatch_sp(r12)
-	addi	r11, r11, 24
+	/* Allocate stack to save LR, TOC & optionally NIP (in case of direct_call) */
+	addi	r11, r11, LIVEPATCH_STACK_FRAME_SIZE
 	std	r11, TI_livepatch_sp(r12)
 
 	/* Store stack end marker */
 	lis     r12, STACK_END_MAGIC@h
 	ori     r12, r12, STACK_END_MAGIC@l
-	std	r12, -8(r11)
+	std	r12, -LIVEPATCH_STACK_MAGIC_OFFSET(r11)
 
 	/* Save toc & real LR on livepatch stack */
-	std	r2,  -24(r11)
+	std	r2,  -LIVEPATCH_STACK_TOC_OFFSET(r11)
 #ifndef CONFIG_PPC_FTRACE_OUT_OF_LINE
 	mflr	r12
-	std	r12, -16(r11)
+	std	r12, -LIVEPATCH_STACK_LR_OFFSET(r11)
 	mfctr	r12
 #else
-	std	r0, -16(r11)
+	std	r0, -LIVEPATCH_STACK_LR_OFFSET(r11)
 	mflr	r12
+
+	/* Also, save new NIP on livepatch stack before the direct_call */
+#ifdef CONFIG_DYNAMIC_FTRACE_WITH_DIRECT_CALLS
+	std	r12, -LIVEPATCH_STACK_NIP_OFFSET(r11)
+
+	/* For direct_call, set cr0.eq bit to indicate livepatch is active */
+	crset	4*cr0+eq
+	/* Jump to the direct_call */
+	bnectrl	cr1
+
+	/*
+	 * The address to jump after direct call is deduced based on ftrace OOL stub sequence.
+	 * The seemingly insignificant couple of instructions below is to mimic that here to
+	 * jump back to the livepatch handler code below.
+	 */
+	nop
+	b	1f
+
+	/*
+	 * Restore the state for livepatching from the livepatch stack.
+	 * Before that, check if livepatch stack is intact. Use r0 for it.
+	 */
+1:	mtctr	r0
+	ld	r12, PACA_THREAD_INFO(r13)
+	ld	r11, TI_livepatch_sp(r12)
+	lis     r0,  STACK_END_MAGIC@h
+	ori     r0,  r0, STACK_END_MAGIC@l
+	ld	r12, -LIVEPATCH_STACK_MAGIC_OFFSET(r11)
+1:	tdne	r12, r0
+	EMIT_BUG_ENTRY 1b, __FILE__, __LINE__ - 1, 0
+	mfctr	r0
+
+	/*
+	 * A change in r0 implies the direct_call is not done yet. The direct_call
+	 * will take care of calling the original LR. Update r0 in livepatch stack
+	 * with the new LR in the direct_call.
+	 */
+	ld	r12, -LIVEPATCH_STACK_LR_OFFSET(r11)
+	cmpd	r12, r0
+	beq	1f
+	mflr	r0
+	std	r0, -LIVEPATCH_STACK_LR_OFFSET(r11)
+
+	/* Put new NIP back in r12 to proceed with livepatch handling */
+1:	ld	r12, -LIVEPATCH_STACK_NIP_OFFSET(r11)
+#endif /* CONFIG_DYNAMIC_FTRACE_WITH_DIRECT_CALLS */
+
 	/* Put ctr in r12 for global entry and branch there */
 	mtctr	r12
 #endif

@@ -377,18 +431,18 @@ livepatch_handler:
 	/* Check stack marker hasn't been trashed */
 	lis     r2,  STACK_END_MAGIC@h
 	ori     r2,  r2, STACK_END_MAGIC@l
-	ld	r12, -8(r11)
+	ld	r12, -LIVEPATCH_STACK_MAGIC_OFFSET(r11)
 1:	tdne	r12, r2
 	EMIT_BUG_ENTRY 1b, __FILE__, __LINE__ - 1, 0
 
 	/* Restore LR & toc from livepatch stack */
-	ld	r12, -16(r11)
+	ld	r12, -LIVEPATCH_STACK_LR_OFFSET(r11)
 	mtlr	r12
-	ld	r2,  -24(r11)
+	ld	r2,  -LIVEPATCH_STACK_TOC_OFFSET(r11)
 
 	/* Pop livepatch stack frame */
 	ld	r12, PACA_THREAD_INFO(r13)
-	subi	r11, r11, 24
+	subi	r11, r11, LIVEPATCH_STACK_FRAME_SIZE
 	std	r11, TI_livepatch_sp(r12)
 
 	/* Return to original caller of live patched function */

@@ -19,6 +19,7 @@
 
 #include <asm/kprobes.h>
 #include <asm/text-patching.h>
+#include <asm/livepatch.h>
 
 #include "bpf_jit.h"
 

@@ -684,6 +685,8 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
 	struct bpf_tramp_links *fentry = &tlinks[BPF_TRAMP_FENTRY];
 	struct bpf_tramp_links *fexit = &tlinks[BPF_TRAMP_FEXIT];
 	struct codegen_context codegen_ctx, *ctx;
+	int __maybe_unused livepatch_sp_off = 0;
+	bool __maybe_unused handle_lp = false;
 	u32 *image = (u32 *)rw_image;
 	ppc_inst_t branch_insn;
 	u32 *branches = NULL;

@@ -716,6 +719,8 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
 	 * dummy frame for unwind       [ back chain 1      ] --
 	 *                              [ padding           ] align stack frame
 	 *       r4_off                 [ r4 (tailcallcnt)  ] optional - 32-bit powerpc
+	 *                              [ *current.TI.lp_sp ]
+	 *    livepatch_sp_off          [ current.TI.lp_sp  ] optional - livepatch stack info
 	 *       alt_lr_off             [ real lr (ool stub)] optional - actual lr
 	 *                              [ r26               ]
 	 *       nvr_off                [ r25               ] nvr save area

@@ -780,11 +785,23 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
 	nvr_off = bpf_frame_size;
 	bpf_frame_size += 2 * SZL;
 
+
 	/* Optional save area for actual LR in case of ool ftrace */
 	if (IS_ENABLED(CONFIG_PPC_FTRACE_OUT_OF_LINE)) {
 		alt_lr_off = bpf_frame_size;
 		bpf_frame_size += SZL;
+		if (IS_ENABLED(CONFIG_DYNAMIC_FTRACE_WITH_DIRECT_CALLS)) {
+			handle_lp = (func_ptr_is_kernel_text(func_addr) && fmod_ret->nr_links &&
+				     (flags & BPF_TRAMP_F_CALL_ORIG));
+		}
+	}
+
+#ifdef CONFIG_LIVEPATCH_64
+	if (handle_lp) {
+		livepatch_sp_off = bpf_frame_size;
+		bpf_frame_size += 2 * SZL;
 	}
+#endif
 
 	if (IS_ENABLED(CONFIG_PPC32)) {
 		if (nr_regs < 2) {

@@ -822,6 +839,32 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
 	if (IS_ENABLED(CONFIG_PPC32) && nr_regs < 2)
 		EMIT(PPC_RAW_STL(_R4, _R1, r4_off));
 
+#ifdef CONFIG_LIVEPATCH_64
+	/* Save expected livepatch stack state on the trampoline stack */
+	if (handle_lp) {
+		/*
+		 * The caller is expected to set cr0.eq bit, if livepatch was active on it.
+		 *
+		 * If livepatch is active, save address & the expected value of
+		 * livepatch stack pointer on the trampoline stack.
+		 * Else, set both of them to 0.
+		 */
+		PPC_BCC_SHORT(COND_EQ, (ctx->idx + 5) * 4);
+		EMIT(PPC_RAW_LI(_R12, 0));
+		EMIT(PPC_RAW_STL(_R12, _R1, livepatch_sp_off));
+		EMIT(PPC_RAW_STL(_R12, _R1, livepatch_sp_off + SZL));
+		PPC_JMP((ctx->idx + 7) * 4);
+
+		EMIT(PPC_RAW_LL(_R12, _R13, offsetof(struct paca_struct, __current) +
+					    offsetof(struct task_struct, thread_info)));
+		EMIT(PPC_RAW_ADDI(_R12, _R12, offsetof(struct thread_info, livepatch_sp)));
+		EMIT(PPC_RAW_STL(_R12, _R1, livepatch_sp_off));
+		EMIT(PPC_RAW_LL(_R12, _R12, 0));
+		EMIT(PPC_RAW_ADDI(_R12, _R12, -LIVEPATCH_STACK_FRAME_SIZE));
+		EMIT(PPC_RAW_STL(_R12, _R1, livepatch_sp_off + SZL));
+	}
+#endif
+
 	bpf_trampoline_save_args(image, ctx, func_frame_offset, nr_regs, regs_off);
 
 	/* Save our return address */

@@ -932,6 +975,38 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
 		image[branches[i]] = ppc_inst_val(branch_insn);
 	}
 
+#ifdef CONFIG_LIVEPATCH_64
+	/*
+	 * Restore livepatch stack state if livepatch was active & an attached
+	 * fmod_ret program failed.
+	 */
+	if (handle_lp) {
+		EMIT(PPC_RAW_LL(_R12, _R1, livepatch_sp_off + SZL));
+		EMIT(PPC_RAW_CMPLI(_R12, 0));
+
+		/*
+		 * If expected value (_R12) of livepatch stack pointer saved on the
+		 * trampoline stack is 0, livepatch was not active. Skip the rest.
+		 */
+		PPC_BCC_SHORT(COND_EQ, (ctx->idx + 7) * 4);
+
+		EMIT(PPC_RAW_LL(_R25, _R1, livepatch_sp_off));
+		EMIT(PPC_RAW_LL(_R25, _R25, 0));
+
+		/*
+		 * If the expected value (_R12) of livepatch stack pointer saved on the
+		 * trampoline stack is not the same as actual value (_R25), it implies
+		 * fmod_ret program failed and skipped calling the traced/livepatch'ed
+		 * function. The livepatch'ed function did not get a chance to tear down
+		 * the livepatch stack it setup. Take care of that here in do_fexit.
+		 */
+		EMIT(PPC_RAW_CMPD(_R12, _R25));
+		PPC_BCC_SHORT(COND_EQ, (ctx->idx + 3) * 4);
+		EMIT(PPC_RAW_LL(_R25, _R1, livepatch_sp_off));
+		EMIT(PPC_RAW_STL(_R12, _R25, 0));
+	}
+#endif
+
 	for (i = 0; i < fexit->nr_links; i++)
 		if (invoke_bpf_prog(image, ro_image, ctx, fexit->links[i], regs_off, retval_off,
 				    run_ctx_off, false)) {