[RFC PATCH bpf-next 0/5] static branches

[RFC PATCH bpf-next 0/5] static branches

[Anton Protopopov]

This series adds support for mapping between xlated and original
instructions offsets, mapping between xlated and jitted instructions
offsets (x86), support for two new BPF instruction JA[SRC=1] and
JA[SRC=3], and a new syscall to configure the jitted values of such
instructions.

This a follow up to the previous attempt to add static keys support
(see [1], [2]) which implements lower-level functionality than what
was proposed before.

It's RFC because 1) to run self-tests it requires a patched llvm
(https://github.com/llvm/llvm-project/pull/75110, thanks a lot
Yonghong!) 2) before spending time writing selftests and proper
bpftool updates [this time] I would prefer to get some initial
feedback first (for examples, see below)

The first patch is a formal fix.
The second patch adds xlated -> original mapping.
The third patch adds xlated -> jitted mapping.
The fourth patch adds support for new instructions.
And the fifth patch adds support for new syscall.

Altogether this provides enough functionality to dynamically patch
programs and support simple static keys. See the third patch which
displays an example of mappings between xlated->orig, xlated->jited.
See the last patch which includes description of how to implement
simple static key functionality.

This would be also interesting to hear what people think about more
high-level api where one "key" controls a set of static branches in
multiple programs and can serialize access to all the branches. What
kind of BPF object can represent such functionality, besides a map,
if any?

  [1] https://lpc.events/event/17/contributions/1608/attachments/1278/2578/bpf-static-keys.pdf
  [2] https://lore.kernel.org/bpf/20231206141030.1478753-1-aspsk@isovalent.com/

Anton Protopopov (5):
  bpf: fix potential error return
  bpf: keep track of and expose xlated insn offsets
  bpf: x86: expose how xlated insns map to jitted insns
  bpf: add support for an extended JA instruction
  bpf: x86: add BPF_STATIC_BRANCH_UPDATE syscall

 arch/x86/net/bpf_jit_comp.c       |  71 ++++++++++++++++++-
 include/linux/bpf.h               |  11 +++
 include/linux/bpf_verifier.h      |   1 -
 include/linux/filter.h            |   1 +
 include/uapi/linux/bpf.h          |  26 +++++++
 kernel/bpf/core.c                 |  69 +++++++++++++++++-
 kernel/bpf/syscall.c              | 112 ++++++++++++++++++++++++++++++
 kernel/bpf/verifier.c             |  58 ++++++++++++----
 tools/bpf/bpftool/prog.c          |  14 ++++
 tools/bpf/bpftool/xlated_dumper.c |   2 +-
 tools/bpf/bpftool/xlated_dumper.h |   2 +
 tools/include/uapi/linux/bpf.h    |  26 +++++++
 12 files changed, 375 insertions(+), 18 deletions(-)

-- 
2.34.1

[RFC PATCH bpf-next 1/5] bpf: fix potential error return

[Anton Protopopov]

The bpf_remove_insns() function returns WARN_ON_ONCE(error), where
error is a result of bpf_adj_branches(), and thus should be always 0
However, if for any reason it is not 0, then it will be converted to
boolean by WARN_ON_ONCE and returned to user space as 1, not an actual
error value. Fix this by returning the original err after the WARN check.

Signed-off-by: Anton Protopopov <aspsk@isovalent.com>
---
 kernel/bpf/core.c | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/kernel/bpf/core.c b/kernel/bpf/core.c
index fbb1d95a9b44..9ba9e0ea9c45 100644
--- a/kernel/bpf/core.c
+++ b/kernel/bpf/core.c
@@ -532,6 +532,8 @@ struct bpf_prog *bpf_patch_insn_single(struct bpf_prog *prog, u32 off,
 
 int bpf_remove_insns(struct bpf_prog *prog, u32 off, u32 cnt)
 {
+	int err;
+
 	/* Branch offsets can't overflow when program is shrinking, no need
 	 * to call bpf_adj_branches(..., true) here
 	 */
@@ -539,7 +541,12 @@ int bpf_remove_insns(struct bpf_prog *prog, u32 off, u32 cnt)
 		sizeof(struct bpf_insn) * (prog->len - off - cnt));
 	prog->len -= cnt;
 
-	return WARN_ON_ONCE(bpf_adj_branches(prog, off, off + cnt, off, false));
+	err = bpf_adj_branches(prog, off, off + cnt, off, false);
+	WARN_ON_ONCE(err);
+	if (err)
+		return err;
+
+	return 0;
 }
 
 static void bpf_prog_kallsyms_del_subprogs(struct bpf_prog *fp)
-- 
2.34.1

Re: [RFC PATCH bpf-next 1/5] bpf: fix potential error return

[Jiri Olsa]

On Mon, Jan 22, 2024 at 04:49:32PM +0000, Anton Protopopov wrote:
> The bpf_remove_insns() function returns WARN_ON_ONCE(error), where
> error is a result of bpf_adj_branches(), and thus should be always 0
> However, if for any reason it is not 0, then it will be converted to
> boolean by WARN_ON_ONCE and returned to user space as 1, not an actual
> error value. Fix this by returning the original err after the WARN check.
> 
> Signed-off-by: Anton Protopopov <aspsk@isovalent.com>

nice catch

Acked-by: Jiri Olsa <jolsa@kernel.org>

> ---
>  kernel/bpf/core.c | 9 ++++++++-
>  1 file changed, 8 insertions(+), 1 deletion(-)
> 
> diff --git a/kernel/bpf/core.c b/kernel/bpf/core.c
> index fbb1d95a9b44..9ba9e0ea9c45 100644
> --- a/kernel/bpf/core.c
> +++ b/kernel/bpf/core.c
> @@ -532,6 +532,8 @@ struct bpf_prog *bpf_patch_insn_single(struct bpf_prog *prog, u32 off,
>  
>  int bpf_remove_insns(struct bpf_prog *prog, u32 off, u32 cnt)
>  {
> +	int err;
> +
>  	/* Branch offsets can't overflow when program is shrinking, no need
>  	 * to call bpf_adj_branches(..., true) here
>  	 */
> @@ -539,7 +541,12 @@ int bpf_remove_insns(struct bpf_prog *prog, u32 off, u32 cnt)
>  		sizeof(struct bpf_insn) * (prog->len - off - cnt));
>  	prog->len -= cnt;
>  
> -	return WARN_ON_ONCE(bpf_adj_branches(prog, off, off + cnt, off, false));
> +	err = bpf_adj_branches(prog, off, off + cnt, off, false);
> +	WARN_ON_ONCE(err);
> +	if (err)
> +		return err;
> +
> +	return 0;

could be just 'return err'

jirka

>  }
>  
>  static void bpf_prog_kallsyms_del_subprogs(struct bpf_prog *fp)
> -- 
> 2.34.1
> 

Re: [RFC PATCH bpf-next 1/5] bpf: fix potential error return

[Anton Protopopov]

On Tue, Jan 30, 2024 at 12:20:06PM +0100, Jiri Olsa wrote:
> On Mon, Jan 22, 2024 at 04:49:32PM +0000, Anton Protopopov wrote:
> > The bpf_remove_insns() function returns WARN_ON_ONCE(error), where
> > error is a result of bpf_adj_branches(), and thus should be always 0
> > However, if for any reason it is not 0, then it will be converted to
> > boolean by WARN_ON_ONCE and returned to user space as 1, not an actual
> > error value. Fix this by returning the original err after the WARN check.
> > 
> > Signed-off-by: Anton Protopopov <aspsk@isovalent.com>
> 
> nice catch
> 
> Acked-by: Jiri Olsa <jolsa@kernel.org>
> 
> > ---
> >  kernel/bpf/core.c | 9 ++++++++-
> >  1 file changed, 8 insertions(+), 1 deletion(-)
> > 
> > diff --git a/kernel/bpf/core.c b/kernel/bpf/core.c
> > index fbb1d95a9b44..9ba9e0ea9c45 100644
> > --- a/kernel/bpf/core.c
> > +++ b/kernel/bpf/core.c
> > @@ -532,6 +532,8 @@ struct bpf_prog *bpf_patch_insn_single(struct bpf_prog *prog, u32 off,
> >  
> >  int bpf_remove_insns(struct bpf_prog *prog, u32 off, u32 cnt)
> >  {
> > +	int err;
> > +
> >  	/* Branch offsets can't overflow when program is shrinking, no need
> >  	 * to call bpf_adj_branches(..., true) here
> >  	 */
> > @@ -539,7 +541,12 @@ int bpf_remove_insns(struct bpf_prog *prog, u32 off, u32 cnt)
> >  		sizeof(struct bpf_insn) * (prog->len - off - cnt));
> >  	prog->len -= cnt;
> >  
> > -	return WARN_ON_ONCE(bpf_adj_branches(prog, off, off + cnt, off, false));
> > +	err = bpf_adj_branches(prog, off, off + cnt, off, false);
> > +	WARN_ON_ONCE(err);
> > +	if (err)
> > +		return err;
> > +
> > +	return 0;
> 
> could be just 'return err'

Thanks.  I am inserting some code in a consequent patch in between,
so left this in this form

> jirka
> 
> >  }
> >  
> >  static void bpf_prog_kallsyms_del_subprogs(struct bpf_prog *fp)
> > -- 
> > 2.34.1
> > 

[RFC PATCH bpf-next 2/5] bpf: keep track of and expose xlated insn offsets

[Anton Protopopov]

On bpf(BPF_PROG_LOAD) syscall a user-supplied program is translated by
the verifier into an "xlated" program. During this process the original
instruction offsets might be adjusted and/or individual instructions
might be replaced by a new set of instructions:

  User-supplied prog:   --->   Xlated prog:

         -- func 0 --          -- func 0 --
             insn                  insn'
             ...                   ...
             insn                  insn'
         -- func 1 --          -- func 1 --
             insn                  insn'
             ...                   ...
             insn                  insn'
         -- func N --          -- func N --
             insn                  insn'
             ...                   ...
             insn                  insn'

We want to provide users (and ourselves) with the off(insn') -> off(insn)
mapping so that when an xlated program is returned to the userspace by the
bpf_prog_get_info_by_fd() function, users can determine the real offsets
of instructions of interest.

Since the 9e4c24e7ee7d ("bpf: verifier: record original instruction
index") commit the verifier saves the original instruction index in
env->insn_aux_data. This information was, however, lost when we patched
instructions. Also, the information about original index was kept in the
verifier env only, so was inaccessible by later stages, like constants
blinding during the jit stage.

To address the above issues save the information about the original
indexes in a separate array inside the prog->aux so that it doesn't
depend on the verifier environment and can be adjusted, and accessed,
during later stages.

To let users access the information after the program was loaded, add new
fields, orig_idx_len and orig_idx to struct bpf_prog_info and patch the
bpf_prog_get_info_by_fd function correspondingly.

Example mapping would be something like this:

    Original prog:                      Xlated prog:

     0:  r1 = 0x0                        0: r1 = 0
     1:  *(u32 *)(r10 - 0x4) = r1        1: *(u32 *)(r10 -4) = r1
     2:  r2 = r10                        2: r2 = r10
     3:  r2 += -0x4                      3: r2 += -4
     4:  r1 = 0x0 ll                     4: r1 = map[id:88]
     6:  call 0x1                        6: r1 += 272
                                         7: r0 = *(u32 *)(r2 +0)
                                         8: if r0 >= 0x1 goto pc+3
                                         9: r0 <<= 3
                                        10: r0 += r1
                                        11: goto pc+1
                                        12 r0 = 0
     7:  r6 = r0                        13 r6 = r0
     8:  if r6 == 0x0 goto +0x2         14: if r6 == 0x0 goto pc+4
     9:  call 0x76                      15: r0 = 0xffffffff8d2079c0
                                        17: r0 = *(u64 *)(r0 +0)
    10:  *(u64 *)(r6 + 0x0) = r0        18: *(u64 *)(r6 +0) = r0
    11:  r0 = 0x0                       19: r0 = 0x0
    12:  exit                           20: exit

Here the orig_idx array has length 21 and is equal to

    (0, 1, 2, 3, 4, 0/*undefined*/, 6, 6, 6, 6, 6, 6, 6, 7, 8, 9, 9, 10, 11, 12)

The item 6 is undefined because the r1=0ll occupies 16 bytes.

Signed-off-by: Anton Protopopov <aspsk@isovalent.com>
---
 include/linux/bpf.h            |  2 ++
 include/linux/bpf_verifier.h   |  1 -
 include/uapi/linux/bpf.h       |  2 ++
 kernel/bpf/core.c              | 30 ++++++++++++++++++++++++++++++
 kernel/bpf/syscall.c           | 30 ++++++++++++++++++++++++++++++
 kernel/bpf/verifier.c          |  6 ++----
 tools/include/uapi/linux/bpf.h |  2 ++
 7 files changed, 68 insertions(+), 5 deletions(-)

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 377857b232c6..dff4c697b674 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -1518,6 +1518,8 @@ struct bpf_prog_aux {
 		struct work_struct work;
 		struct rcu_head	rcu;
 	};
+	/* an array of original indexes for all xlated instructions */
+	u32 *orig_idx;
 };
 
 struct bpf_prog {
diff --git a/include/linux/bpf_verifier.h b/include/linux/bpf_verifier.h
index e11baecbde68..2728c83ea46d 100644
--- a/include/linux/bpf_verifier.h
+++ b/include/linux/bpf_verifier.h
@@ -553,7 +553,6 @@ struct bpf_insn_aux_data {
 	u8 alu_state; /* used in combination with alu_limit */
 
 	/* below fields are initialized once */
-	unsigned int orig_idx; /* original instruction index */
 	bool jmp_point;
 	bool prune_point;
 	/* ensure we check state equivalence and save state checkpoint and
diff --git a/include/uapi/linux/bpf.h b/include/uapi/linux/bpf.h
index a00f8a5623e1..b15e167941fd 100644
--- a/include/uapi/linux/bpf.h
+++ b/include/uapi/linux/bpf.h
@@ -6470,6 +6470,8 @@ struct bpf_prog_info {
 	__u32 verified_insns;
 	__u32 attach_btf_obj_id;
 	__u32 attach_btf_id;
+	__u32 orig_idx_len;
+	__aligned_u64 orig_idx;
 } __attribute__((aligned(8)));
 
 struct bpf_map_info {
diff --git a/kernel/bpf/core.c b/kernel/bpf/core.c
index 9ba9e0ea9c45..11eccc477b83 100644
--- a/kernel/bpf/core.c
+++ b/kernel/bpf/core.c
@@ -469,6 +469,30 @@ static void bpf_adj_linfo(struct bpf_prog *prog, u32 off, u32 delta)
 		linfo[i].insn_off += delta;
 }
 
+static int bpf_prog_realloc_orig_idx(struct bpf_prog *prog, u32 off, u32 patch_len)
+{
+	u32 *old_idx = prog->aux->orig_idx, *new_idx;
+	u32 new_prog_len = prog->len + patch_len - 1;
+	int i;
+
+	if (patch_len <= 1)
+		return 0;
+
+	new_idx = kzalloc(array_size(new_prog_len, sizeof(u32)), GFP_KERNEL);
+	if (!new_idx)
+		return -ENOMEM;
+
+	memcpy(new_idx, old_idx, sizeof(*old_idx) * off);
+	for (i = off; i < off + patch_len; i++)
+		new_idx[i] = old_idx[off];
+	memcpy(new_idx + off + patch_len, old_idx + off + 1,
+			sizeof(*old_idx) * (prog->len - off));
+
+	prog->aux->orig_idx = new_idx;
+	kfree(old_idx);
+	return 0;
+}
+
 struct bpf_prog *bpf_patch_insn_single(struct bpf_prog *prog, u32 off,
 				       const struct bpf_insn *patch, u32 len)
 {
@@ -494,6 +518,10 @@ struct bpf_prog *bpf_patch_insn_single(struct bpf_prog *prog, u32 off,
 	    (err = bpf_adj_branches(prog, off, off + 1, off + len, true)))
 		return ERR_PTR(err);
 
+	err = bpf_prog_realloc_orig_idx(prog, off, len);
+	if (err)
+		return ERR_PTR(err);
+
 	/* Several new instructions need to be inserted. Make room
 	 * for them. Likely, there's no need for a new allocation as
 	 * last page could have large enough tailroom.
@@ -2778,6 +2806,8 @@ static void bpf_prog_free_deferred(struct work_struct *work)
 	} else {
 		bpf_jit_free(aux->prog);
 	}
+	if (aux->orig_idx)
+		kfree(aux->orig_idx);
 }
 
 void bpf_prog_free(struct bpf_prog *fp)
diff --git a/kernel/bpf/syscall.c b/kernel/bpf/syscall.c
index a1f18681721c..e264dbe285b2 100644
--- a/kernel/bpf/syscall.c
+++ b/kernel/bpf/syscall.c
@@ -2589,6 +2589,18 @@ static bool is_perfmon_prog_type(enum bpf_prog_type prog_type)
 	}
 }
 
+static void *bpf_prog_alloc_orig_idx(u32 insn_cnt)
+{
+	u32 *orig_idx;
+	int i;
+
+	orig_idx = kzalloc(sizeof(*orig_idx) * insn_cnt, GFP_KERNEL);
+	if (orig_idx)
+		for (i = 0; i < insn_cnt; i++)
+			orig_idx[i] = i;
+	return orig_idx;
+}
+
 /* last field in 'union bpf_attr' used by this command */
 #define	BPF_PROG_LOAD_LAST_FIELD log_true_size
 
@@ -2690,6 +2702,12 @@ static int bpf_prog_load(union bpf_attr *attr, bpfptr_t uattr, u32 uattr_size)
 		return -ENOMEM;
 	}
 
+	prog->aux->orig_idx = bpf_prog_alloc_orig_idx(attr->insn_cnt);
+	if (!prog->aux->orig_idx) {
+		err = -ENOMEM;
+		goto free_prog;
+	}
+
 	prog->expected_attach_type = attr->expected_attach_type;
 	prog->aux->attach_btf = attach_btf;
 	prog->aux->attach_btf_id = attr->attach_btf_id;
@@ -4460,6 +4478,18 @@ static int bpf_prog_get_info_by_fd(struct file *file,
 			return -EFAULT;
 	}
 
+	ulen = info.orig_idx_len;
+	if (prog->aux->orig_idx)
+		info.orig_idx_len = prog->len * sizeof(*prog->aux->orig_idx);
+	else
+		info.orig_idx_len = 0;
+	if (info.orig_idx_len && ulen) {
+		if (copy_to_user(u64_to_user_ptr(info.orig_idx),
+				 prog->aux->orig_idx,
+				 min_t(u32, info.orig_idx_len, ulen)))
+			return -EFAULT;
+	}
+
 	if (bpf_prog_is_offloaded(prog->aux)) {
 		err = bpf_prog_offload_info_fill(&info, prog);
 		if (err)
diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index 9507800026cf..64c7036b8b56 100644
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c
@@ -18266,7 +18266,7 @@ static struct bpf_prog *bpf_patch_insn_data(struct bpf_verifier_env *env, u32 of
 		if (PTR_ERR(new_prog) == -ERANGE)
 			verbose(env,
 				"insn %d cannot be patched due to 16-bit range\n",
-				env->insn_aux_data[off].orig_idx);
+				env->prog->aux->orig_idx[off]);
 		vfree(new_data);
 		return NULL;
 	}
@@ -20777,7 +20777,7 @@ int bpf_check(struct bpf_prog **prog, union bpf_attr *attr, bpfptr_t uattr, __u3
 {
 	u64 start_time = ktime_get_ns();
 	struct bpf_verifier_env *env;
-	int i, len, ret = -EINVAL, err;
+	int len, ret = -EINVAL, err;
 	u32 log_true_size;
 	bool is_priv;
 
@@ -20800,8 +20800,6 @@ int bpf_check(struct bpf_prog **prog, union bpf_attr *attr, bpfptr_t uattr, __u3
 	ret = -ENOMEM;
 	if (!env->insn_aux_data)
 		goto err_free_env;
-	for (i = 0; i < len; i++)
-		env->insn_aux_data[i].orig_idx = i;
 	env->prog = *prog;
 	env->ops = bpf_verifier_ops[env->prog->type];
 	env->fd_array = make_bpfptr(attr->fd_array, uattr.is_kernel);
diff --git a/tools/include/uapi/linux/bpf.h b/tools/include/uapi/linux/bpf.h
index a00f8a5623e1..b15e167941fd 100644
--- a/tools/include/uapi/linux/bpf.h
+++ b/tools/include/uapi/linux/bpf.h
@@ -6470,6 +6470,8 @@ struct bpf_prog_info {
 	__u32 verified_insns;
 	__u32 attach_btf_obj_id;
 	__u32 attach_btf_id;
+	__u32 orig_idx_len;
+	__aligned_u64 orig_idx;
 } __attribute__((aligned(8)));
 
 struct bpf_map_info {
-- 
2.34.1

[RFC PATCH bpf-next 3/5] bpf: x86: expose how xlated insns map to jitted insns

[Anton Protopopov]

Allow users to get the exact xlated -> jitted instructions mapping.
This is done by using a new field xlated_to_jit in bpf_prog_info
which can return up to prog->len

    struct xlated_to_jit {
            u32 off;
            u32 len;
    };

elements. The xlated_to_jit[insn_off] contains jitted offset within
a function and the length of the resulting jitted instruction.

Example:

       Original:            Xlated:                              Jitted:

                                                                 0:   nopl    (%rax,%rax)
                                                                 5:   nop
                                                                 7:   pushq   %rbp
                                                                 8:   movq    %rsp, %rbp
       0:  call 0x76        0:  r0 = 0xfffffbeef                 b:   movabsq $-1923847220, %rax
                            2:  r0 = *(u64 *)(r0 +0)            15:   movq    (%rax), %rax
       1:  r1 = 0x9 ll      3:  r1 = map[id:666][0]+9           19:   movabsq $-102223334445559, %rdi
       3:  r2 = 0x6         5:  r2 = 6                          23:   movl    $6, %esi
       4:  r3 = r0          6:  r3 = r0                         28:   movq    %rax, %rdx
       5:  call 0x6         7:  call bpf_trace_printk           2b:   callq   0xffffffffcdead4dc
       6:  call pc+2        8:  call pc+2                       30:   callq   0x7c
       7:  r0 = 0x0         9:  r0 = 0                          35:   xorl    %eax, %eax
       8:  exit            10:  exit                            37:   leave
                                                                38:   jmp     0xffffffffcbeeffbc
       ---                 ---                                  ---
                                                                 0:   nopl    (%rax,%rax)
                                                                 5:   nop
                                                                 7:   pushq   %rbp
                                                                 8:   movq    %rsp, %rbp
       9:  goto +0x1       11:  goto pc+1                        b:   jmp     0xf
      10:  goto +0x1       12:  goto pc+1                        d:   jmp     0x11
      11:  goto -0x2       13:  goto pc-2                        f:   jmp     0xd
      12:  r0 = 0x0        14:  r0 = 0                          11:   xorl    %eax, %eax
      13:  exit            15:  exit                            13:   leave
                                                                14:   jmp     0xffffffffcbffbeef

Here the xlated_to_jit array will be of length 16 (11 + 6) and equal to

     0: (0xb, 10)
     1: (0,0) /* undefined, as the previous instruction is 16 bytes */
     2: (0x15, 4)
     3: (0x19, 10)
     4: (0,0) /* undefined, as the previous instruction is 16 bytes */
     5: (0x23, 5)
     6: (0x28, 3)
     7: (0x2b, 5)
     8: (0x30, 5)
     9: (0x35, 2)
    10: (0x37, 6)
    11: (0xb, 2)
    12: (0xd, 2)
    13: (0xf, 2)
    14: (0x11, 2)
    15: (0x13, 6)

The prologues are "unmapped": no mapping exists for xlated -> [0,b)

Signed-off-by: Anton Protopopov <aspsk@isovalent.com>
---
 arch/x86/net/bpf_jit_comp.c    | 13 +++++++++++++
 include/linux/bpf.h            |  7 +++++++
 include/uapi/linux/bpf.h       |  7 +++++++
 kernel/bpf/core.c              | 25 +++++++++++++++++++++++++
 kernel/bpf/syscall.c           | 25 +++++++++++++++++++++++++
 kernel/bpf/verifier.c          |  9 +++++++++
 tools/include/uapi/linux/bpf.h |  7 +++++++
 7 files changed, 93 insertions(+)

diff --git a/arch/x86/net/bpf_jit_comp.c b/arch/x86/net/bpf_jit_comp.c
index e1390d1e331b..736aec2565b8 100644
--- a/arch/x86/net/bpf_jit_comp.c
+++ b/arch/x86/net/bpf_jit_comp.c
@@ -1186,6 +1186,7 @@ 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;
 		u8 b2 = 0, b3 = 0;
 		u8 *start_of_ldx;
 		s64 jmp_offset;
@@ -1290,6 +1291,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 */
@@ -2073,6 +2075,17 @@ st:			if (is_imm8(insn->off))
 				return -EFAULT;
 			}
 			memcpy(rw_image + proglen, temp, ilen);
+
+			if (bpf_prog->aux->xlated_to_jit) {
+				int off;
+
+				off = i - 1 - adjust_off;
+				if (bpf_prog->aux->func_idx)
+					off += bpf_prog->aux->func_info[bpf_prog->aux->func_idx].insn_off;
+
+				bpf_prog->aux->xlated_to_jit[off].off = proglen;
+				bpf_prog->aux->xlated_to_jit[off].len = ilen;
+			}
 		}
 		proglen += ilen;
 		addrs[i] = proglen;
diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index dff4c697b674..660df06cb541 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -1520,6 +1520,13 @@ struct bpf_prog_aux {
 	};
 	/* an array of original indexes for all xlated instructions */
 	u32 *orig_idx;
+	/* for every xlated instruction point to all generated jited
+	 * instructions, if allocated
+	 */
+	struct {
+		u32 off;	/* local offset in the jitted code */
+		u32 len;	/* the total len of generated jit code */
+	} *xlated_to_jit;
 };
 
 struct bpf_prog {
diff --git a/include/uapi/linux/bpf.h b/include/uapi/linux/bpf.h
index b15e167941fd..83dad9ea7a3b 100644
--- a/include/uapi/linux/bpf.h
+++ b/include/uapi/linux/bpf.h
@@ -6430,6 +6430,11 @@ struct sk_reuseport_md {
 
 #define BPF_TAG_SIZE	8
 
+struct xlated_to_jit {
+	__u32 off;
+	__u32 len;
+};
+
 struct bpf_prog_info {
 	__u32 type;
 	__u32 id;
@@ -6472,6 +6477,8 @@ struct bpf_prog_info {
 	__u32 attach_btf_id;
 	__u32 orig_idx_len;
 	__aligned_u64 orig_idx;
+	__u32 xlated_to_jit_len;
+	__aligned_u64 xlated_to_jit;
 } __attribute__((aligned(8)));
 
 struct bpf_map_info {
diff --git a/kernel/bpf/core.c b/kernel/bpf/core.c
index 11eccc477b83..e502485c757a 100644
--- a/kernel/bpf/core.c
+++ b/kernel/bpf/core.c
@@ -493,6 +493,26 @@ static int bpf_prog_realloc_orig_idx(struct bpf_prog *prog, u32 off, u32 patch_l
 	return 0;
 }
 
+static void adjust_func_info(struct bpf_prog *prog, u32 off, u32 insn_delta)
+{
+	int i;
+
+	if (insn_delta == 0)
+		return;
+
+	for (i = 0; i < prog->aux->func_info_cnt; i++) {
+		if (prog->aux->func_info[i].insn_off <= off)
+			continue;
+		prog->aux->func_info[i].insn_off += insn_delta;
+	}
+}
+
+static void bpf_prog_adj_orig_idx_after_remove(struct bpf_prog *prog, u32 off, u32 len)
+{
+	memmove(prog->aux->orig_idx + off, prog->aux->orig_idx + off + len,
+		sizeof(*prog->aux->orig_idx) * (prog->len - off));
+}
+
 struct bpf_prog *bpf_patch_insn_single(struct bpf_prog *prog, u32 off,
 				       const struct bpf_insn *patch, u32 len)
 {
@@ -554,6 +574,7 @@ struct bpf_prog *bpf_patch_insn_single(struct bpf_prog *prog, u32 off,
 	BUG_ON(bpf_adj_branches(prog_adj, off, off + 1, off + len, false));
 
 	bpf_adj_linfo(prog_adj, off, insn_delta);
+	adjust_func_info(prog_adj, off, insn_delta);
 
 	return prog_adj;
 }
@@ -574,6 +595,8 @@ int bpf_remove_insns(struct bpf_prog *prog, u32 off, u32 cnt)
 	if (err)
 		return err;
 
+	bpf_prog_adj_orig_idx_after_remove(prog, off, cnt);
+
 	return 0;
 }
 
@@ -2808,6 +2831,8 @@ static void bpf_prog_free_deferred(struct work_struct *work)
 	}
 	if (aux->orig_idx)
 		kfree(aux->orig_idx);
+	if (aux->xlated_to_jit)
+		kfree(aux->xlated_to_jit);
 }
 
 void bpf_prog_free(struct bpf_prog *fp)
diff --git a/kernel/bpf/syscall.c b/kernel/bpf/syscall.c
index e264dbe285b2..97b0ba6ecf65 100644
--- a/kernel/bpf/syscall.c
+++ b/kernel/bpf/syscall.c
@@ -4490,6 +4490,31 @@ static int bpf_prog_get_info_by_fd(struct file *file,
 			return -EFAULT;
 	}
 
+	ulen = info.xlated_to_jit_len;
+	if (prog->aux->xlated_to_jit)
+		info.xlated_to_jit_len = prog->len * sizeof(struct xlated_to_jit);
+	else
+		info.xlated_to_jit_len = 0;
+	if (info.xlated_to_jit_len && ulen) {
+		struct xlated_to_jit *xlated_to_jit;
+		int i;
+
+		xlated_to_jit = kzalloc(info.xlated_to_jit_len, GFP_KERNEL);
+		if (!xlated_to_jit)
+			return -ENOMEM;
+		for (i = 0; i < prog->len; i++) {
+			xlated_to_jit[i].off = prog->aux->xlated_to_jit[i].off;
+			xlated_to_jit[i].len = prog->aux->xlated_to_jit[i].len;
+		}
+		if (copy_to_user(u64_to_user_ptr(info.xlated_to_jit),
+				 xlated_to_jit,
+				 min_t(u32, info.xlated_to_jit_len, ulen))) {
+			kfree(xlated_to_jit);
+			return -EFAULT;
+		}
+		kfree(xlated_to_jit);
+	}
+
 	if (bpf_prog_is_offloaded(prog->aux)) {
 		err = bpf_prog_offload_info_fill(&info, prog);
 		if (err)
diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index 64c7036b8b56..fad47044ccce 100644
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c
@@ -18951,6 +18951,7 @@ static int jit_subprogs(struct bpf_verifier_env *env)
 		func[i]->aux->exception_cb = env->subprog_info[i].is_exception_cb;
 		if (!i)
 			func[i]->aux->exception_boundary = env->seen_exception;
+		func[i]->aux->xlated_to_jit = prog->aux->xlated_to_jit;
 		func[i] = bpf_int_jit_compile(func[i]);
 		if (!func[i]->jited) {
 			err = -ENOTSUPP;
@@ -20780,6 +20781,7 @@ int bpf_check(struct bpf_prog **prog, union bpf_attr *attr, bpfptr_t uattr, __u3
 	int len, ret = -EINVAL, err;
 	u32 log_true_size;
 	bool is_priv;
+	u32 size;
 
 	/* no program is valid */
 	if (ARRAY_SIZE(bpf_verifier_ops) == 0)
@@ -20930,6 +20932,13 @@ int bpf_check(struct bpf_prog **prog, union bpf_attr *attr, bpfptr_t uattr, __u3
 								     : false;
 	}
 
+	if (ret == 0) {
+		size = array_size(sizeof(*env->prog->aux->xlated_to_jit), env->prog->len);
+		env->prog->aux->xlated_to_jit = kzalloc(size, GFP_KERNEL);
+		if (!env->prog->aux->xlated_to_jit)
+			ret = -ENOMEM;
+	}
+
 	if (ret == 0)
 		ret = fixup_call_args(env);
 
diff --git a/tools/include/uapi/linux/bpf.h b/tools/include/uapi/linux/bpf.h
index b15e167941fd..83dad9ea7a3b 100644
--- a/tools/include/uapi/linux/bpf.h
+++ b/tools/include/uapi/linux/bpf.h
@@ -6430,6 +6430,11 @@ struct sk_reuseport_md {
 
 #define BPF_TAG_SIZE	8
 
+struct xlated_to_jit {
+	__u32 off;
+	__u32 len;
+};
+
 struct bpf_prog_info {
 	__u32 type;
 	__u32 id;
@@ -6472,6 +6477,8 @@ struct bpf_prog_info {
 	__u32 attach_btf_id;
 	__u32 orig_idx_len;
 	__aligned_u64 orig_idx;
+	__u32 xlated_to_jit_len;
+	__aligned_u64 xlated_to_jit;
 } __attribute__((aligned(8)));
 
 struct bpf_map_info {
-- 
2.34.1

[RFC PATCH bpf-next 4/5] bpf: add support for an extended JA instruction

[Anton Protopopov]

Add support for a new version of JA instruction, a static branch JA. Such
instructions may either jump to the specified offset or act as nops. To
distinguish such instructions from normal JA the BPF_STATIC_BRANCH_JA flag
should be set for the SRC register.

By default on program load such instructions are jitted as a normal JA.
However, if the BPF_STATIC_BRANCH_NOP flag is set in the SRC register,
then the instruction is jitted to a NOP.

In order to generate BPF_STATIC_BRANCH_JA instructions using llvm two new
instructions were added:

	asm volatile goto ("nop_or_gotol %l[label]" :::: label);

will generate the BPF_STATIC_BRANCH_JA|BPF_STATIC_BRANCH_NOP instuction and

	asm volatile goto ("gotol_or_nop %l[label]" :::: label);

will generate a BPF_STATIC_BRANCH_JA instruction, without an extra bit set.
The reason for adding two instructions is that both are required to implement
static keys functionality for BPF.

The verifier logic is extended to check both possible paths: jump and nop.

Signed-off-by: Anton Protopopov <aspsk@isovalent.com>
---
 arch/x86/net/bpf_jit_comp.c | 19 ++++++++++++++--
 include/uapi/linux/bpf.h    | 10 +++++++++
 kernel/bpf/verifier.c       | 43 +++++++++++++++++++++++++++++--------
 3 files changed, 61 insertions(+), 11 deletions(-)

diff --git a/arch/x86/net/bpf_jit_comp.c b/arch/x86/net/bpf_jit_comp.c
index 736aec2565b8..52b9de134ab3 100644
--- a/arch/x86/net/bpf_jit_comp.c
+++ b/arch/x86/net/bpf_jit_comp.c
@@ -1131,6 +1131,15 @@ static void emit_shiftx(u8 **pprog, u32 dst_reg, u8 src_reg, bool is64, u8 op)
 	*pprog = prog;
 }
 
+static bool is_static_ja_nop(const struct bpf_insn *insn)
+{
+	u8 code = insn->code;
+
+	return (code == (BPF_JMP | BPF_JA) || code == (BPF_JMP32 | BPF_JA)) &&
+	       (insn->src_reg & BPF_STATIC_BRANCH_JA) &&
+	       (insn->src_reg & BPF_STATIC_BRANCH_NOP);
+}
+
 #define INSN_SZ_DIFF (((addrs[i] - addrs[i - 1]) - (prog - temp)))
 
 /* mov rax, qword ptr [rbp - rounded_stack_depth - 8] */
@@ -2016,9 +2025,15 @@ st:			if (is_imm8(insn->off))
 					}
 					emit_nops(&prog, INSN_SZ_DIFF - 2);
 				}
-				EMIT2(0xEB, jmp_offset);
+				if (is_static_ja_nop(insn))
+					emit_nops(&prog, 2);
+				else
+					EMIT2(0xEB, jmp_offset);
 			} else if (is_simm32(jmp_offset)) {
-				EMIT1_off32(0xE9, jmp_offset);
+				if (is_static_ja_nop(insn))
+					emit_nops(&prog, 5);
+				else
+					EMIT1_off32(0xE9, jmp_offset);
 			} else {
 				pr_err("jmp gen bug %llx\n", jmp_offset);
 				return -EFAULT;
diff --git a/include/uapi/linux/bpf.h b/include/uapi/linux/bpf.h
index 83dad9ea7a3b..43ad332ffbee 100644
--- a/include/uapi/linux/bpf.h
+++ b/include/uapi/linux/bpf.h
@@ -1372,6 +1372,16 @@ struct bpf_stack_build_id {
 	};
 };
 
+/* Flags for JA insn, passed in SRC_REG */
+enum {
+	BPF_STATIC_BRANCH_JA  = 1 << 0,
+	BPF_STATIC_BRANCH_NOP = 1 << 1,
+};
+
+#define BPF_STATIC_BRANCH_MASK (BPF_STATIC_BRANCH_JA | \
+				BPF_STATIC_BRANCH_NOP)
+
+
 #define BPF_OBJ_NAME_LEN 16U
 
 union bpf_attr {
diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index fad47044ccce..50d19755b8fb 100644
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c
@@ -15607,14 +15607,24 @@ static int visit_insn(int t, struct bpf_verifier_env *env)
 		else
 			off = insn->imm;
 
-		/* unconditional jump with single edge */
-		ret = push_insn(t, t + off + 1, FALLTHROUGH, env);
-		if (ret)
-			return ret;
+		if (insn->src_reg & BPF_STATIC_BRANCH_JA) {
+			/* static branch - jump with two edges */
+			mark_prune_point(env, t);
+
+			ret = push_insn(t, t + 1, FALLTHROUGH, env);
+			if (ret)
+				return ret;
 
-		mark_prune_point(env, t + off + 1);
-		mark_jmp_point(env, t + off + 1);
+			ret = push_insn(t, t + off + 1, BRANCH, env);
+		} else {
+			/* unconditional jump with single edge */
+			ret = push_insn(t, t + off + 1, FALLTHROUGH, env);
+			if (ret)
+				return ret;
 
+			mark_prune_point(env, t + off + 1);
+			mark_jmp_point(env, t + off + 1);
+		}
 		return ret;
 
 	default:
@@ -17584,8 +17594,11 @@ static int do_check(struct bpf_verifier_env *env)
 
 				mark_reg_scratched(env, BPF_REG_0);
 			} else if (opcode == BPF_JA) {
+				struct bpf_verifier_state *other_branch;
+				u32 jmp_offset;
+
 				if (BPF_SRC(insn->code) != BPF_K ||
-				    insn->src_reg != BPF_REG_0 ||
+				    (insn->src_reg & ~BPF_STATIC_BRANCH_MASK) ||
 				    insn->dst_reg != BPF_REG_0 ||
 				    (class == BPF_JMP && insn->imm != 0) ||
 				    (class == BPF_JMP32 && insn->off != 0)) {
@@ -17594,9 +17607,21 @@ static int do_check(struct bpf_verifier_env *env)
 				}
 
 				if (class == BPF_JMP)
-					env->insn_idx += insn->off + 1;
+					jmp_offset = insn->off + 1;
 				else
-					env->insn_idx += insn->imm + 1;
+					jmp_offset = insn->imm + 1;
+
+				/* Staic branch can either jump to +off or fall through */
+				if (insn->src_reg & BPF_STATIC_BRANCH_JA) {
+					other_branch = push_stack(env, env->insn_idx + jmp_offset,
+							env->insn_idx, false);
+					if (!other_branch)
+						return -EFAULT;
+
+					jmp_offset = 1;
+				}
+
+				env->insn_idx += jmp_offset;
 				continue;
 
 			} else if (opcode == BPF_EXIT) {
-- 
2.34.1

[RFC PATCH bpf-next 5/5] bpf: x86: add BPF_STATIC_BRANCH_UPDATE syscall

[Anton Protopopov]

Add a new bpf system call, BPF_STATIC_BRANCH_UPDATE, which allows users to
update static branches in BPF. Namely, this system call is executed as

    bpf(BPF_STATIC_BRANCH_UPDATE, attrs={prog_fd, insn_off, on})

where prog_fd points to a BPF program, insn_off is an _xlated_ offset in
this program, on is a boolean value to set this branch on or off.
The instruction at insn_off must be a JA with SRC_REG or'ed with
BPF_STATIC_BRANCH_JA and, optionally, with BPF_STATIC_BRANCH_INVERSE.

To implement this for a particular architecture, re-define the weak
bpf_arch_poke_static_branch() function in the corresponding bpf_jit_comp.c

Example of usage can be found below.  Lets load and compile the
following [dummy] program:

    SEC("kprobe/__x64_sys_getpgid")
    int worker(void *ctx)
    {
            if (bpf_static_branch_unlikely(&key))
                    return 1;
            else
                    return 0;
    }

Here key is some map and bpf_static_branch_unlikely() is defined as
follows:

    static __always_inline int __bpf_static_branch_nop(void *static_key)
    {
            asm goto("1:\n\t"
                    "nop_or_gotol %l[l_yes]\n\t"
                    ".pushsection .jump_table, \"aw\"\n\t"
                    ".balign 8\n\t"
                    ".long 1b - .\n\t"
                    ".long %l[l_yes] - .\n\t"
                    ".quad %c0 - .\n\t"
                    ".popsection\n\t"
                    :: "i" (static_key)
                    :: l_yes);
            return 0;
    l_yes:
            return 1;
    }

    #define bpf_static_branch_unlikely(static_key) \
            unlikely(__bpf_static_branch_nop(static_key))

Here the extra code is needed to automate search for the static
branch location, and the main part is the usage of asm goto + the
nop_or_gotol instruction.

After compilation and load the program will look like this:

    # bpftool prog dump x id 42
    int worker(void * ctx):
       0: (b7) r0 = 1
       1: (06) nop_or_gotol pc+1
       2: (b7) r0 = 0
       3: (95) exit

And the jitted program will have nop_or_gotol (jitted offset 0x10)
translated to a NOP (as the branch is not activated by default):

    # bpftool prog dump j id 42
    int worker(void * ctx):
       0:   nopl    (%rax,%rax)
       5:   nop
       7:   pushq   %rbp
       8:   movq    %rsp, %rbp
       b:   movl    $1, %eax
    ; asm goto("1:\n\t"
      10:   nop
      12:   xorl    %eax, %eax
      14:   leave
      15:   jmp     0xffffffffcbc16ed8

If we issue a

    bpf(BPF_STATIC_BRANCH_UPDATE, {bpf_prog_get_fd_by_id(42), .off=1, .on=1})

syscall (xlated offset = 1, on = 1), then the jitted code will change
to

    # bpftool prog dump j id 42
    int worker(void * ctx):
       0:   nopl    (%rax,%rax)
       5:   nop
       7:   pushq   %rbp
       8:   movq    %rsp, %rbp
       b:   movl    $1, %eax
    ; asm goto("1:\n\t"
      10:   jmp     0x14
      12:   xorl    %eax, %eax
      14:   leave
      15:   jmp     0xffffffffcbc16ed8

as expected.

A "likely" variant can be implemented using the 'gotol_or_nop'
instruction.

Signed-off-by: Anton Protopopov <aspsk@isovalent.com>
---
 arch/x86/net/bpf_jit_comp.c | 39 +++++++++++++++++++++++++
 include/linux/bpf.h         |  2 ++
 include/linux/filter.h      |  1 +
 include/uapi/linux/bpf.h    |  7 +++++
 kernel/bpf/core.c           |  5 ++++
 kernel/bpf/syscall.c        | 57 +++++++++++++++++++++++++++++++++++++
 6 files changed, 111 insertions(+)

diff --git a/arch/x86/net/bpf_jit_comp.c b/arch/x86/net/bpf_jit_comp.c
index 52b9de134ab3..c757e4d997a7 100644
--- a/arch/x86/net/bpf_jit_comp.c
+++ b/arch/x86/net/bpf_jit_comp.c
@@ -2098,8 +2098,16 @@ st:			if (is_imm8(insn->off))
 				if (bpf_prog->aux->func_idx)
 					off += bpf_prog->aux->func_info[bpf_prog->aux->func_idx].insn_off;
 
+				bpf_prog->aux->xlated_to_jit[off].ip = image + proglen;
 				bpf_prog->aux->xlated_to_jit[off].off = proglen;
 				bpf_prog->aux->xlated_to_jit[off].len = ilen;
+
+				/*
+				 * save the offset so that it can later be accessed
+				 * by the bpf(BPF_STATIC_BRANCH_UPDATE) syscall
+				 */
+				if (insn->code == (BPF_JMP | BPF_JA) || insn->code == (BPF_JMP32 | BPF_JA))
+					bpf_prog->aux->xlated_to_jit[off].jmp_offset = jmp_offset;
 			}
 		}
 		proglen += ilen;
@@ -3275,3 +3283,34 @@ bool bpf_jit_supports_ptr_xchg(void)
 {
 	return true;
 }
+
+int bpf_arch_poke_static_branch(struct bpf_prog *prog,
+				u32 insn_off,
+				bool on)
+{
+	int jmp_offset = prog->aux->xlated_to_jit[insn_off].jmp_offset;
+	u32 len = prog->aux->xlated_to_jit[insn_off].len;
+	u8 op[5];
+
+	if (is_imm8(jmp_offset) && len != 2)
+		return -EINVAL;
+
+	if (!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);
+	}
+
+	text_poke_bp(prog->aux->xlated_to_jit[insn_off].ip, op, len, NULL);
+
+	return 0;
+}
diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 660df06cb541..ba77e0c6f390 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -1524,8 +1524,10 @@ struct bpf_prog_aux {
 	 * instructions, if allocated
 	 */
 	struct {
+		void *ip;	/* the address of the jitted insn */
 		u32 off;	/* local offset in the jitted code */
 		u32 len;	/* the total len of generated jit code */
+		u32 jmp_offset;	/* jitted jump offset for BPF_JA insns */
 	} *xlated_to_jit;
 };
 
diff --git a/include/linux/filter.h b/include/linux/filter.h
index 35f067fd3840..ff76a60cf247 100644
--- a/include/linux/filter.h
+++ b/include/linux/filter.h
@@ -957,6 +957,7 @@ bool bpf_jit_supports_far_kfunc_call(void);
 bool bpf_jit_supports_exceptions(void);
 bool bpf_jit_supports_ptr_xchg(void);
 void arch_bpf_stack_walk(bool (*consume_fn)(void *cookie, u64 ip, u64 sp, u64 bp), void *cookie);
+int bpf_arch_poke_static_branch(struct bpf_prog *prog, u32 off, bool on);
 bool bpf_helper_changes_pkt_data(void *func);
 
 static inline bool bpf_dump_raw_ok(const struct cred *cred)
diff --git a/include/uapi/linux/bpf.h b/include/uapi/linux/bpf.h
index 43ad332ffbee..e5d226838a3d 100644
--- a/include/uapi/linux/bpf.h
+++ b/include/uapi/linux/bpf.h
@@ -901,6 +901,7 @@ enum bpf_cmd {
 	BPF_ITER_CREATE,
 	BPF_LINK_DETACH,
 	BPF_PROG_BIND_MAP,
+	BPF_STATIC_BRANCH_UPDATE,
 };
 
 enum bpf_map_type {
@@ -1724,6 +1725,12 @@ union bpf_attr {
 		__u32		flags;		/* extra flags */
 	} prog_bind_map;
 
+	struct { /* struct used by BPF_STATIC_BRANCH_UPDATE command */
+		__u32		prog_fd;
+		__u32		insn_off;
+		__u32		on;
+	} static_branch;
+
 } __attribute__((aligned(8)));
 
 /* The description below is an attempt at providing documentation to eBPF
diff --git a/kernel/bpf/core.c b/kernel/bpf/core.c
index e502485c757a..5272879449d8 100644
--- a/kernel/bpf/core.c
+++ b/kernel/bpf/core.c
@@ -3043,6 +3043,11 @@ static int __init bpf_global_ma_init(void)
 late_initcall(bpf_global_ma_init);
 #endif
 
+int __weak bpf_arch_poke_static_branch(struct bpf_prog *prog, u32 off, bool on)
+{
+	return -EOPNOTSUPP;
+}
+
 DEFINE_STATIC_KEY_FALSE(bpf_stats_enabled_key);
 EXPORT_SYMBOL(bpf_stats_enabled_key);
 
diff --git a/kernel/bpf/syscall.c b/kernel/bpf/syscall.c
index 97b0ba6ecf65..c3509e59f82d 100644
--- a/kernel/bpf/syscall.c
+++ b/kernel/bpf/syscall.c
@@ -1504,6 +1504,60 @@ static int map_lookup_elem(union bpf_attr *attr)
 	return err;
 }
 
+int parse_static_branch_insn(struct bpf_insn *insn, bool *inverse)
+{
+	__u8 code = insn->code;
+
+	if (code != (BPF_JMP | BPF_JA) && code != (BPF_JMP32 | BPF_JA))
+		return -EINVAL;
+
+	if (insn->src_reg & ~BPF_STATIC_BRANCH_MASK)
+		return -EINVAL;
+
+	if (!(insn->src_reg & BPF_STATIC_BRANCH_JA))
+		return -EINVAL;
+
+	if (insn->dst_reg)
+		return -EINVAL;
+
+	*inverse = !(insn->src_reg & BPF_STATIC_BRANCH_NOP);
+
+	return 0;
+}
+
+#define BPF_STATIC_BRANCH_UPDATE_LAST_FIELD static_branch.on
+
+static int bpf_static_branch_update(union bpf_attr *attr)
+{
+	bool on = attr->static_branch.on & 1;
+	struct bpf_prog *prog;
+	u32 insn_off;
+	bool inverse;
+	int ret;
+
+	if (CHECK_ATTR(BPF_STATIC_BRANCH_UPDATE))
+		return -EINVAL;
+
+	prog = bpf_prog_get(attr->static_branch.prog_fd);
+	if (IS_ERR(prog))
+		return PTR_ERR(prog);
+
+	insn_off = attr->static_branch.insn_off;
+	if (insn_off >= prog->len) {
+		ret = -ERANGE;
+		goto put_prog;
+	}
+
+	ret = parse_static_branch_insn(&prog->insnsi[insn_off], &inverse);
+	if (ret)
+		goto put_prog;
+
+	ret = bpf_arch_poke_static_branch(prog, insn_off, on ^ inverse);
+
+put_prog:
+	bpf_prog_put(prog);
+	return ret;
+}
 
 #define BPF_MAP_UPDATE_ELEM_LAST_FIELD flags
 
@@ -5578,6 +5632,9 @@ static int __sys_bpf(int cmd, bpfptr_t uattr, unsigned int size)
 	case BPF_MAP_DELETE_BATCH:
 		err = bpf_map_do_batch(&attr, uattr.user, BPF_MAP_DELETE_BATCH);
 		break;
+	case BPF_STATIC_BRANCH_UPDATE:
+		err = bpf_static_branch_update(&attr);
+		break;
 	case BPF_LINK_CREATE:
 		err = link_create(&attr, uattr);
 		break;
-- 
2.34.1