[RFC PATCH bpf-next 08/14] bpf, x86: implement static key support

[Anton Protopopov <aspsk@xxxxxxxxxxxxx>]

Implement bpf_arch_poke_static_branch() for x86. Namely, during each
JIT loop, save IP values and sizes of jump instructions pointed by
static keys. Then use the text_poke_bp() to toggle jumps/nops.

Signed-off-by: Anton Protopopov <aspsk@xxxxxxxxxxxxx>
---
 arch/x86/net/bpf_jit_comp.c | 46 +++++++++++++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/arch/x86/net/bpf_jit_comp.c b/arch/x86/net/bpf_jit_comp.c
index 5856ac1aab80..31cadded820b 100644
--- a/arch/x86/net/bpf_jit_comp.c
+++ b/arch/x86/net/bpf_jit_comp.c
@@ -1569,6 +1569,8 @@ static int do_jit(struct bpf_prog *bpf_prog, int *addrs, u8 *image, u8 *rw_image
 		const s32 imm32 = insn->imm;
 		u32 dst_reg = insn->dst_reg;
 		u32 src_reg = insn->src_reg;
+		int adjust_off = 0;
+		int abs_xlated_off;
 		u8 b2 = 0, b3 = 0;
 		u8 *start_of_ldx;
 		s64 jmp_offset;
@@ -1724,6 +1726,7 @@ static int do_jit(struct bpf_prog *bpf_prog, int *addrs, u8 *image, u8 *rw_image
 			emit_mov_imm64(&prog, dst_reg, insn[1].imm, insn[0].imm);
 			insn++;
 			i++;
+			adjust_off = 1;
 			break;
 
 			/* dst %= src, dst /= src, dst %= imm32, dst /= imm32 */
@@ -2595,6 +2598,14 @@ st:			if (is_imm8(insn->off))
 				return -EFAULT;
 			}
 			memcpy(rw_image + proglen, temp, ilen);
+
+			/*
+			 * Static keys need to know how the xlated code
+			 * of static ja instructions maps to jited code
+			 */
+			abs_xlated_off = bpf_prog->aux->subprog_start + i - 1 - adjust_off;
+			bpf_prog_update_insn_ptr(bpf_prog, abs_xlated_off, proglen, ilen,
+						 jmp_offset, image + proglen);
 		}
 		proglen += ilen;
 		addrs[i] = proglen;
@@ -3880,3 +3891,38 @@ bool bpf_jit_supports_timed_may_goto(void)
 {
 	return true;
 }
+
+int bpf_arch_poke_static_branch(struct bpf_insn_ptr *ptr, bool on)
+{
+	int jmp_offset = ptr->jitted_jump_offset;
+	void *ip = ptr->jitted_ip;
+	u32 len = ptr->jitted_len;
+	u8 op[5];
+
+	if (WARN_ON_ONCE(!ip))
+		return -EINVAL;
+
+	if (WARN_ON_ONCE(is_imm8(jmp_offset) && len != 2))
+		return -EINVAL;
+
+	if (WARN_ON_ONCE(!is_imm8(jmp_offset) && len != 5))
+		return -EINVAL;
+
+	if (on) {
+		if (len == 2) {
+			op[0] = 0xEB;
+			op[1] = jmp_offset;
+		} else {
+			op[0] = 0xE9;
+			memcpy(&op[1], &jmp_offset, 4);
+		}
+	} else {
+		memcpy(op, x86_nops[len], len);
+	}
+
+	mutex_lock(&text_mutex);
+	text_poke_bp(ip, op, len, NULL);
+	mutex_unlock(&text_mutex);
+
+	return 0;
+}
-- 
2.34.1