diff mbox series

[bpf-next,5/6] bpf: Improve tracing recursion prevention mechanism

Message ID 20230417154737.12740-6-laoar.shao@gmail.com (mailing list archive)
State Rejected
Delegated to: BPF
Headers show
Series bpf: Tracing recursion prevention mechanism improvement | expand

Checks

Context Check Description
bpf/vmtest-bpf-next-PR success PR summary
bpf/vmtest-bpf-next-VM_Test-1 success Logs for ShellCheck
bpf/vmtest-bpf-next-VM_Test-7 success Logs for set-matrix
bpf/vmtest-bpf-next-VM_Test-2 success Logs for build for aarch64 with gcc
bpf/vmtest-bpf-next-VM_Test-3 success Logs for build for aarch64 with llvm-16
bpf/vmtest-bpf-next-VM_Test-5 success Logs for build for x86_64 with gcc
bpf/vmtest-bpf-next-VM_Test-6 success Logs for build for x86_64 with llvm-16
bpf/vmtest-bpf-next-VM_Test-4 success Logs for build for s390x with gcc
bpf/vmtest-bpf-next-VM_Test-8 success Logs for test_maps on aarch64 with gcc
bpf/vmtest-bpf-next-VM_Test-9 success Logs for test_maps on aarch64 with llvm-16
bpf/vmtest-bpf-next-VM_Test-11 success Logs for test_maps on x86_64 with gcc
bpf/vmtest-bpf-next-VM_Test-12 success Logs for test_maps on x86_64 with llvm-16
bpf/vmtest-bpf-next-VM_Test-13 success Logs for test_progs on aarch64 with gcc
bpf/vmtest-bpf-next-VM_Test-14 fail Logs for test_progs on aarch64 with llvm-16
bpf/vmtest-bpf-next-VM_Test-16 fail Logs for test_progs on x86_64 with gcc
bpf/vmtest-bpf-next-VM_Test-17 fail Logs for test_progs on x86_64 with llvm-16
bpf/vmtest-bpf-next-VM_Test-18 success Logs for test_progs_no_alu32 on aarch64 with gcc
bpf/vmtest-bpf-next-VM_Test-19 success Logs for test_progs_no_alu32 on aarch64 with llvm-16
bpf/vmtest-bpf-next-VM_Test-21 fail Logs for test_progs_no_alu32 on x86_64 with gcc
bpf/vmtest-bpf-next-VM_Test-22 fail Logs for test_progs_no_alu32 on x86_64 with llvm-16
bpf/vmtest-bpf-next-VM_Test-23 success Logs for test_progs_no_alu32_parallel on aarch64 with gcc
bpf/vmtest-bpf-next-VM_Test-24 success Logs for test_progs_no_alu32_parallel on aarch64 with llvm-16
bpf/vmtest-bpf-next-VM_Test-25 success Logs for test_progs_no_alu32_parallel on x86_64 with gcc
bpf/vmtest-bpf-next-VM_Test-26 success Logs for test_progs_no_alu32_parallel on x86_64 with llvm-16
bpf/vmtest-bpf-next-VM_Test-27 success Logs for test_progs_parallel on aarch64 with gcc
bpf/vmtest-bpf-next-VM_Test-28 success Logs for test_progs_parallel on aarch64 with llvm-16
bpf/vmtest-bpf-next-VM_Test-29 success Logs for test_progs_parallel on x86_64 with gcc
bpf/vmtest-bpf-next-VM_Test-30 success Logs for test_progs_parallel on x86_64 with llvm-16
bpf/vmtest-bpf-next-VM_Test-31 success Logs for test_verifier on aarch64 with gcc
bpf/vmtest-bpf-next-VM_Test-32 success Logs for test_verifier on aarch64 with llvm-16
bpf/vmtest-bpf-next-VM_Test-33 success Logs for test_verifier on s390x with gcc
bpf/vmtest-bpf-next-VM_Test-34 success Logs for test_verifier on x86_64 with gcc
bpf/vmtest-bpf-next-VM_Test-35 success Logs for test_verifier on x86_64 with llvm-16
bpf/vmtest-bpf-next-VM_Test-20 success Logs for test_progs_no_alu32 on s390x with gcc
bpf/vmtest-bpf-next-VM_Test-15 success Logs for test_progs on s390x with gcc
bpf/vmtest-bpf-next-VM_Test-10 success Logs for test_maps on s390x with gcc
netdev/series_format success Posting correctly formatted
netdev/tree_selection success Clearly marked for bpf-next, async
netdev/fixes_present success Fixes tag not required for -next series
netdev/header_inline success No static functions without inline keyword in header files
netdev/build_32bit success Errors and warnings before: 1465 this patch: 1465
netdev/cc_maintainers warning 2 maintainers not CCed: martin.lau@linux.dev song@kernel.org
netdev/build_clang success Errors and warnings before: 164 this patch: 164
netdev/verify_signedoff success Signed-off-by tag matches author and committer
netdev/deprecated_api success None detected
netdev/check_selftest success No net selftest shell script
netdev/verify_fixes success No Fixes tag
netdev/build_allmodconfig_warn success Errors and warnings before: 1459 this patch: 1459
netdev/checkpatch success total: 0 errors, 0 warnings, 0 checks, 162 lines checked
netdev/kdoc success Errors and warnings before: 0 this patch: 0
netdev/source_inline success Was 0 now: 0

Commit Message

Yafang Shao April 17, 2023, 3:47 p.m. UTC
Currently we use prog->active to prevent tracing recursion, but it has
some downsides,

- It can't identify different contexts
  That said, if a process context is interrupted by a irq context and
  the irq context runs the same code path, it will be considered as
  recursion. For example,
    normal:
      this_cpu_inc_return(*(prog->active)) == 1 <- OK

      irq:
        this_cpu_inc_return(*(prog->active)) == 1 <- FAIL!
        [ Considered as recusion ]

- It has to maintain a percpu area
  A percpu area will be allocated for each prog when the prog is loaded
  and be freed when the prog is destroyed.

Let's replace it with the generic tracing recursion prevention mechanism,
which can work fine with anything. In the above example, the irq context
won't be considered as recursion again,
  normal:
    test_recursion_try_acquire() <- OK

    softirq:
      test_recursion_try_acquire() <- OK

      irq:
        test_recursion_try_acquire() <- OK

Note that, currently one single recursion in process context is allowed
due to the TRACE_CTX_TRANSITION workaround, which can be fixed in the
future. That said, below behavior is expected currently,
  normal:
    test_recursion_try_acquire() <- OK
    [ recursion happens ]        <- one single recursion is allowed
    test_recursion_try_acquire() <- OK
    [ recursion happens ]
    test_recursion_try_acquire() <- RECURSION!

Signed-off-by: Yafang Shao <laoar.shao@gmail.com>
---
 include/linux/bpf.h      |  2 +-
 kernel/bpf/core.c        | 10 ----------
 kernel/bpf/trampoline.c  | 44 +++++++++++++++++++++++++++++++++-----------
 kernel/trace/bpf_trace.c | 12 +++++++-----
 4 files changed, 41 insertions(+), 27 deletions(-)

Comments

Alexei Starovoitov April 17, 2023, 8:14 p.m. UTC | #1
On Mon, Apr 17, 2023 at 03:47:36PM +0000, Yafang Shao wrote:
> diff --git a/kernel/bpf/trampoline.c b/kernel/bpf/trampoline.c
> index f61d513..3df39a5 100644
> --- a/kernel/bpf/trampoline.c
> +++ b/kernel/bpf/trampoline.c
> @@ -842,15 +842,21 @@ static __always_inline u64 notrace bpf_prog_start_time(void)
>  static u64 notrace __bpf_prog_enter_recur(struct bpf_prog *prog, struct bpf_tramp_run_ctx *run_ctx)
>  	__acquires(RCU)
>  {
> -	rcu_read_lock();
> -	migrate_disable();
> -
> -	run_ctx->saved_run_ctx = bpf_set_run_ctx(&run_ctx->run_ctx);
> +	int bit;
>  
> -	if (unlikely(this_cpu_inc_return(*(prog->active)) != 1)) {
> +	rcu_read_lock();
> +	bit = test_recursion_try_acquire(_THIS_IP_, _RET_IP_);

and bpf will prevent ftrace to run and vice versa.
Not a good idea.

One bpf prog will prevent different bpf prog to run since they share current task.
Not a good idea either.
kernel test robot April 17, 2023, 11:29 p.m. UTC | #2
Hi Yafang,

kernel test robot noticed the following build errors:

[auto build test ERROR on bpf-next/master]

url:    https://github.com/intel-lab-lkp/linux/commits/Yafang-Shao/bpf-Add-__rcu_read_-lock-unlock-into-btf-id-deny-list/20230417-235009
base:   https://git.kernel.org/pub/scm/linux/kernel/git/bpf/bpf-next.git master
patch link:    https://lore.kernel.org/r/20230417154737.12740-6-laoar.shao%40gmail.com
patch subject: [PATCH bpf-next 5/6] bpf: Improve tracing recursion prevention mechanism
config: loongarch-defconfig (https://download.01.org/0day-ci/archive/20230418/202304180736.cWjpwhs6-lkp@intel.com/config)
compiler: loongarch64-linux-gcc (GCC) 12.1.0
reproduce (this is a W=1 build):
        wget https://raw.githubusercontent.com/intel/lkp-tests/master/sbin/make.cross -O ~/bin/make.cross
        chmod +x ~/bin/make.cross
        # https://github.com/intel-lab-lkp/linux/commit/ac84d2623c0469a703245030f2b23612ab4505dd
        git remote add linux-review https://github.com/intel-lab-lkp/linux
        git fetch --no-tags linux-review Yafang-Shao/bpf-Add-__rcu_read_-lock-unlock-into-btf-id-deny-list/20230417-235009
        git checkout ac84d2623c0469a703245030f2b23612ab4505dd
        # save the config file
        mkdir build_dir && cp config build_dir/.config
        COMPILER_INSTALL_PATH=$HOME/0day COMPILER=gcc-12.1.0 make.cross W=1 O=build_dir ARCH=loongarch olddefconfig
        COMPILER_INSTALL_PATH=$HOME/0day COMPILER=gcc-12.1.0 make.cross W=1 O=build_dir ARCH=loongarch SHELL=/bin/bash

If you fix the issue, kindly add following tag where applicable
| Reported-by: kernel test robot <lkp@intel.com>
| Link: https://lore.kernel.org/oe-kbuild-all/202304180736.cWjpwhs6-lkp@intel.com/

All errors (new ones prefixed by >>):

   kernel/bpf/trampoline.c: In function '__bpf_prog_enter_recur':
>> kernel/bpf/trampoline.c:848:15: error: implicit declaration of function 'test_recursion_try_acquire' [-Werror=implicit-function-declaration]
     848 |         bit = test_recursion_try_acquire(_THIS_IP_, _RET_IP_);
         |               ^~~~~~~~~~~~~~~~~~~~~~~~~~
   kernel/bpf/trampoline.c: In function '__bpf_prog_exit_recur':
>> kernel/bpf/trampoline.c:896:9: error: implicit declaration of function 'test_recursion_release'; did you mean 'dev_recursion_level'? [-Werror=implicit-function-declaration]
     896 |         test_recursion_release(run_ctx->recursion_bit);
         |         ^~~~~~~~~~~~~~~~~~~~~~
         |         dev_recursion_level
   cc1: some warnings being treated as errors


vim +/test_recursion_try_acquire +848 kernel/bpf/trampoline.c

   828	
   829	/* The logic is similar to bpf_prog_run(), but with an explicit
   830	 * rcu_read_lock() and migrate_disable() which are required
   831	 * for the trampoline. The macro is split into
   832	 * call __bpf_prog_enter
   833	 * call prog->bpf_func
   834	 * call __bpf_prog_exit
   835	 *
   836	 * __bpf_prog_enter returns:
   837	 * 0 - skip execution of the bpf prog
   838	 * 1 - execute bpf prog
   839	 * [2..MAX_U64] - execute bpf prog and record execution time.
   840	 *     This is start time.
   841	 */
   842	static u64 notrace __bpf_prog_enter_recur(struct bpf_prog *prog, struct bpf_tramp_run_ctx *run_ctx)
   843		__acquires(RCU)
   844	{
   845		int bit;
   846	
   847		rcu_read_lock();
 > 848		bit = test_recursion_try_acquire(_THIS_IP_, _RET_IP_);
   849		run_ctx->recursion_bit = bit;
   850		if (bit < 0) {
   851			preempt_disable_notrace();
   852			bpf_prog_inc_misses_counter(prog);
   853			preempt_enable_notrace();
   854			return 0;
   855		}
   856	
   857		migrate_disable();
   858	
   859		run_ctx->saved_run_ctx = bpf_set_run_ctx(&run_ctx->run_ctx);
   860		return bpf_prog_start_time();
   861	}
   862	
   863	static void notrace update_prog_stats(struct bpf_prog *prog,
   864					      u64 start)
   865	{
   866		struct bpf_prog_stats *stats;
   867	
   868		if (static_branch_unlikely(&bpf_stats_enabled_key) &&
   869		    /* static_key could be enabled in __bpf_prog_enter*
   870		     * and disabled in __bpf_prog_exit*.
   871		     * And vice versa.
   872		     * Hence check that 'start' is valid.
   873		     */
   874		    start > NO_START_TIME) {
   875			unsigned long flags;
   876	
   877			stats = this_cpu_ptr(prog->stats);
   878			flags = u64_stats_update_begin_irqsave(&stats->syncp);
   879			u64_stats_inc(&stats->cnt);
   880			u64_stats_add(&stats->nsecs, sched_clock() - start);
   881			u64_stats_update_end_irqrestore(&stats->syncp, flags);
   882		}
   883	}
   884	
   885	static void notrace __bpf_prog_exit_recur(struct bpf_prog *prog, u64 start,
   886						  struct bpf_tramp_run_ctx *run_ctx)
   887		__releases(RCU)
   888	{
   889		if (run_ctx->recursion_bit < 0)
   890			goto out;
   891	
   892		bpf_reset_run_ctx(run_ctx->saved_run_ctx);
   893	
   894		update_prog_stats(prog, start);
   895		migrate_enable();
 > 896		test_recursion_release(run_ctx->recursion_bit);
   897	
   898	out:
   899		rcu_read_unlock();
   900	}
   901
Yafang Shao April 18, 2023, 1:49 a.m. UTC | #3
On Tue, Apr 18, 2023 at 4:15 AM Alexei Starovoitov
<alexei.starovoitov@gmail.com> wrote:
>
> On Mon, Apr 17, 2023 at 03:47:36PM +0000, Yafang Shao wrote:
> > diff --git a/kernel/bpf/trampoline.c b/kernel/bpf/trampoline.c
> > index f61d513..3df39a5 100644
> > --- a/kernel/bpf/trampoline.c
> > +++ b/kernel/bpf/trampoline.c
> > @@ -842,15 +842,21 @@ static __always_inline u64 notrace bpf_prog_start_time(void)
> >  static u64 notrace __bpf_prog_enter_recur(struct bpf_prog *prog, struct bpf_tramp_run_ctx *run_ctx)
> >       __acquires(RCU)
> >  {
> > -     rcu_read_lock();
> > -     migrate_disable();
> > -
> > -     run_ctx->saved_run_ctx = bpf_set_run_ctx(&run_ctx->run_ctx);
> > +     int bit;
> >
> > -     if (unlikely(this_cpu_inc_return(*(prog->active)) != 1)) {
> > +     rcu_read_lock();
> > +     bit = test_recursion_try_acquire(_THIS_IP_, _RET_IP_);
>
> and bpf will prevent ftrace to run and vice versa.
> Not a good idea.
>
> One bpf prog will prevent different bpf prog to run since they share current task.
> Not a good idea either.

That shouldn't happen. test_recursion_try_acquire() uses a
per-task_struct value. One single task_struct can't run in parallel,
right?
Note that the bpf program running in softirq or irq context won't be
prevented by it.
IIUC, the bpf program should run in serial in one single task, right?
That said, one bpf program can only run after another bpf program
finished in the same task?
Alexei Starovoitov April 18, 2023, 3:38 p.m. UTC | #4
On Mon, Apr 17, 2023 at 6:49 PM Yafang Shao <laoar.shao@gmail.com> wrote:
>
> On Tue, Apr 18, 2023 at 4:15 AM Alexei Starovoitov
> <alexei.starovoitov@gmail.com> wrote:
> >
> > On Mon, Apr 17, 2023 at 03:47:36PM +0000, Yafang Shao wrote:
> > > diff --git a/kernel/bpf/trampoline.c b/kernel/bpf/trampoline.c
> > > index f61d513..3df39a5 100644
> > > --- a/kernel/bpf/trampoline.c
> > > +++ b/kernel/bpf/trampoline.c
> > > @@ -842,15 +842,21 @@ static __always_inline u64 notrace bpf_prog_start_time(void)
> > >  static u64 notrace __bpf_prog_enter_recur(struct bpf_prog *prog, struct bpf_tramp_run_ctx *run_ctx)
> > >       __acquires(RCU)
> > >  {
> > > -     rcu_read_lock();
> > > -     migrate_disable();
> > > -
> > > -     run_ctx->saved_run_ctx = bpf_set_run_ctx(&run_ctx->run_ctx);
> > > +     int bit;
> > >
> > > -     if (unlikely(this_cpu_inc_return(*(prog->active)) != 1)) {
> > > +     rcu_read_lock();
> > > +     bit = test_recursion_try_acquire(_THIS_IP_, _RET_IP_);
> >
> > and bpf will prevent ftrace to run and vice versa.
> > Not a good idea.
> >
> > One bpf prog will prevent different bpf prog to run since they share current task.
> > Not a good idea either.
>
> That shouldn't happen. test_recursion_try_acquire() uses a
> per-task_struct value. One single task_struct can't run in parallel,
> right?
> Note that the bpf program running in softirq or irq context won't be
> prevented by it.
> IIUC, the bpf program should run in serial in one single task, right?
> That said, one bpf program can only run after another bpf program
> finished in the same task?

bpf progs can nest in the same task.
Yafang Shao April 19, 2023, 11:46 a.m. UTC | #5
On Tue, Apr 18, 2023 at 11:38 PM Alexei Starovoitov
<alexei.starovoitov@gmail.com> wrote:
>
> On Mon, Apr 17, 2023 at 6:49 PM Yafang Shao <laoar.shao@gmail.com> wrote:
> >
> > On Tue, Apr 18, 2023 at 4:15 AM Alexei Starovoitov
> > <alexei.starovoitov@gmail.com> wrote:
> > >
> > > On Mon, Apr 17, 2023 at 03:47:36PM +0000, Yafang Shao wrote:
> > > > diff --git a/kernel/bpf/trampoline.c b/kernel/bpf/trampoline.c
> > > > index f61d513..3df39a5 100644
> > > > --- a/kernel/bpf/trampoline.c
> > > > +++ b/kernel/bpf/trampoline.c
> > > > @@ -842,15 +842,21 @@ static __always_inline u64 notrace bpf_prog_start_time(void)
> > > >  static u64 notrace __bpf_prog_enter_recur(struct bpf_prog *prog, struct bpf_tramp_run_ctx *run_ctx)
> > > >       __acquires(RCU)
> > > >  {
> > > > -     rcu_read_lock();
> > > > -     migrate_disable();
> > > > -
> > > > -     run_ctx->saved_run_ctx = bpf_set_run_ctx(&run_ctx->run_ctx);
> > > > +     int bit;
> > > >
> > > > -     if (unlikely(this_cpu_inc_return(*(prog->active)) != 1)) {
> > > > +     rcu_read_lock();
> > > > +     bit = test_recursion_try_acquire(_THIS_IP_, _RET_IP_);
> > >
> > > and bpf will prevent ftrace to run and vice versa.
> > > Not a good idea.
> > >
> > > One bpf prog will prevent different bpf prog to run since they share current task.
> > > Not a good idea either.
> >
> > That shouldn't happen. test_recursion_try_acquire() uses a
> > per-task_struct value. One single task_struct can't run in parallel,
> > right?
> > Note that the bpf program running in softirq or irq context won't be
> > prevented by it.
> > IIUC, the bpf program should run in serial in one single task, right?
> > That said, one bpf program can only run after another bpf program
> > finished in the same task?
>
> bpf progs can nest in the same task.

Do you mean the tail_call ?
Steven Rostedt April 24, 2023, 9:40 p.m. UTC | #6
On Wed, 19 Apr 2023 15:46:34 -0700
Alexei Starovoitov <alexei.starovoitov@gmail.com> wrote:

> No. Just one prog at entry into any of the kernel functions
> and another prog at entry of funcs that 1st bpf prog called indirectly.
> Like one prog is tracing networking events while another
> is focusing on mm. They should not conflict.

You mean that you have:

function start:
  __bpf_prog_enter_recur()
    bpf_program1()
      __bpf_prog_enter_recur()
        bpf_program2();
      __bpf_prog_exit_recur()
  __bpf_prog_exit_recur()

  rest of function

That is, a bpf program can be called within another bpf pogram between
the prog_enter and prog_exit(), that is in the same context (normal,
softirq, irq, etc)?

The protection is on the trampoline where the bpf program is called.
Not sure how ftrace can stop BPF or BPF stop ftrace, unless bpf is
tracing a ftrace callback, or ftrace is tracing a bpf function.

-- Steve
Yafang Shao April 27, 2023, 9:57 a.m. UTC | #7
On Tue, Apr 25, 2023 at 5:40 AM Steven Rostedt <rostedt@goodmis.org> wrote:
>
> On Wed, 19 Apr 2023 15:46:34 -0700
> Alexei Starovoitov <alexei.starovoitov@gmail.com> wrote:
>
> > No. Just one prog at entry into any of the kernel functions
> > and another prog at entry of funcs that 1st bpf prog called indirectly.
> > Like one prog is tracing networking events while another
> > is focusing on mm. They should not conflict.
>
> You mean that you have:
>
> function start:
>   __bpf_prog_enter_recur()
>     bpf_program1()
>       __bpf_prog_enter_recur()
>         bpf_program2();
>       __bpf_prog_exit_recur()
>   __bpf_prog_exit_recur()
>
>   rest of function
>
> That is, a bpf program can be called within another bpf pogram between
> the prog_enter and prog_exit(), that is in the same context (normal,
> softirq, irq, etc)?
>

Right, that can happen per my verification. Below is a simple bpf
program to verify it.

struct {
    __uint(type, BPF_MAP_TYPE_LPM_TRIE);
    __type(key, __u64);
    __type(value, __u64);
    __uint(max_entries, 1024);
    __uint(map_flags, BPF_F_NO_PREALLOC);
} write_map SEC(".maps");

__u64 key;

SEC("fentry/kernel_clone")
int program1()
{
    __u64 value = 1;

    bpf_printk("before update");
    // It will call trie_update_elem and thus trigger program2.
    bpf_map_update_elem(&write_map, &key, &value, BPF_ANY);
    __sync_fetch_and_add(&key, 1);
    bpf_printk("after update");
    return 0;
}

SEC("fentry/trie_update_elem")
int program2()
{
    bpf_printk("trie_update_elem");
    return 0;
}

The result as follows,

         kubelet-203203  [018] ....1  9579.862862:
__bpf_prog_enter_recur: __bpf_prog_enter_recur
         kubelet-203203  [018] ...11  9579.862869: bpf_trace_printk:
before update
         kubelet-203203  [018] ....2  9579.862869:
__bpf_prog_enter_recur: __bpf_prog_enter_recur
         kubelet-203203  [018] ...12  9579.862870: bpf_trace_printk:
trie_update_elem
         kubelet-203203  [018] ....2  9579.862870:
__bpf_prog_exit_recur: __bpf_prog_exit_recur
         kubelet-203203  [018] ...11  9579.862870: bpf_trace_printk:
after update
         kubelet-203203  [018] ....1  9579.862871:
__bpf_prog_exit_recur: __bpf_prog_exit_recur

Note that we can't trace __bpf_prog_enter_recur and
__bpf_prog_exit_recur, so we have to modify the kernel to print them.

> The protection is on the trampoline where the bpf program is called.
> Not sure how ftrace can stop BPF or BPF stop ftrace, unless bpf is
> tracing a ftrace callback, or ftrace is tracing a bpf function.
>
> -- Steve
Yafang Shao April 27, 2023, 12:15 p.m. UTC | #8
On Thu, Apr 27, 2023 at 5:57 PM Yafang Shao <laoar.shao@gmail.com> wrote:
>
> On Tue, Apr 25, 2023 at 5:40 AM Steven Rostedt <rostedt@goodmis.org> wrote:
> >
> > On Wed, 19 Apr 2023 15:46:34 -0700
> > Alexei Starovoitov <alexei.starovoitov@gmail.com> wrote:
> >
> > > No. Just one prog at entry into any of the kernel functions
> > > and another prog at entry of funcs that 1st bpf prog called indirectly.
> > > Like one prog is tracing networking events while another
> > > is focusing on mm. They should not conflict.
> >
> > You mean that you have:
> >
> > function start:
> >   __bpf_prog_enter_recur()
> >     bpf_program1()
> >       __bpf_prog_enter_recur()
> >         bpf_program2();
> >       __bpf_prog_exit_recur()
> >   __bpf_prog_exit_recur()
> >
> >   rest of function
> >
> > That is, a bpf program can be called within another bpf pogram between
> > the prog_enter and prog_exit(), that is in the same context (normal,
> > softirq, irq, etc)?
> >
>
> Right, that can happen per my verification. Below is a simple bpf
> program to verify it.
>
> struct {
>     __uint(type, BPF_MAP_TYPE_LPM_TRIE);
>     __type(key, __u64);
>     __type(value, __u64);
>     __uint(max_entries, 1024);
>     __uint(map_flags, BPF_F_NO_PREALLOC);
> } write_map SEC(".maps");
>
> __u64 key;
>
> SEC("fentry/kernel_clone")
> int program1()
> {
>     __u64 value = 1;
>
>     bpf_printk("before update");
>     // It will call trie_update_elem and thus trigger program2.
>     bpf_map_update_elem(&write_map, &key, &value, BPF_ANY);
>     __sync_fetch_and_add(&key, 1);
>     bpf_printk("after update");
>     return 0;
> }
>
> SEC("fentry/trie_update_elem")
> int program2()
> {
>     bpf_printk("trie_update_elem");
>     return 0;
> }
>
> The result as follows,
>
>          kubelet-203203  [018] ....1  9579.862862:
> __bpf_prog_enter_recur: __bpf_prog_enter_recur
>          kubelet-203203  [018] ...11  9579.862869: bpf_trace_printk:
> before update
>          kubelet-203203  [018] ....2  9579.862869:
> __bpf_prog_enter_recur: __bpf_prog_enter_recur
>          kubelet-203203  [018] ...12  9579.862870: bpf_trace_printk:
> trie_update_elem
>          kubelet-203203  [018] ....2  9579.862870:
> __bpf_prog_exit_recur: __bpf_prog_exit_recur
>          kubelet-203203  [018] ...11  9579.862870: bpf_trace_printk:
> after update
>          kubelet-203203  [018] ....1  9579.862871:
> __bpf_prog_exit_recur: __bpf_prog_exit_recur
>
> Note that we can't trace __bpf_prog_enter_recur and
> __bpf_prog_exit_recur, so we have to modify the kernel to print them.
>

... However, surprisingly it still works even after this patchset is
applied, because the hardirq/softirq flag is set when the program2 is
running, see also the flags in the above trace_pipe output. Is that
expected ?!
I need  some time to figure it out, but maybe you have a quick answer...

> > The protection is on the trampoline where the bpf program is called.
> > Not sure how ftrace can stop BPF or BPF stop ftrace, unless bpf is
> > tracing a ftrace callback, or ftrace is tracing a bpf function.
> >
> > -- Steve
>
Yafang Shao April 27, 2023, 12:35 p.m. UTC | #9
On Thu, Apr 27, 2023 at 8:15 PM Yafang Shao <laoar.shao@gmail.com> wrote:
>
> On Thu, Apr 27, 2023 at 5:57 PM Yafang Shao <laoar.shao@gmail.com> wrote:
> >
> > On Tue, Apr 25, 2023 at 5:40 AM Steven Rostedt <rostedt@goodmis.org> wrote:
> > >
> > > On Wed, 19 Apr 2023 15:46:34 -0700
> > > Alexei Starovoitov <alexei.starovoitov@gmail.com> wrote:
> > >
> > > > No. Just one prog at entry into any of the kernel functions
> > > > and another prog at entry of funcs that 1st bpf prog called indirectly.
> > > > Like one prog is tracing networking events while another
> > > > is focusing on mm. They should not conflict.
> > >
> > > You mean that you have:
> > >
> > > function start:
> > >   __bpf_prog_enter_recur()
> > >     bpf_program1()
> > >       __bpf_prog_enter_recur()
> > >         bpf_program2();
> > >       __bpf_prog_exit_recur()
> > >   __bpf_prog_exit_recur()
> > >
> > >   rest of function
> > >
> > > That is, a bpf program can be called within another bpf pogram between
> > > the prog_enter and prog_exit(), that is in the same context (normal,
> > > softirq, irq, etc)?
> > >
> >
> > Right, that can happen per my verification. Below is a simple bpf
> > program to verify it.
> >
> > struct {
> >     __uint(type, BPF_MAP_TYPE_LPM_TRIE);
> >     __type(key, __u64);
> >     __type(value, __u64);
> >     __uint(max_entries, 1024);
> >     __uint(map_flags, BPF_F_NO_PREALLOC);
> > } write_map SEC(".maps");
> >
> > __u64 key;
> >
> > SEC("fentry/kernel_clone")
> > int program1()
> > {
> >     __u64 value = 1;
> >
> >     bpf_printk("before update");
> >     // It will call trie_update_elem and thus trigger program2.
> >     bpf_map_update_elem(&write_map, &key, &value, BPF_ANY);
> >     __sync_fetch_and_add(&key, 1);
> >     bpf_printk("after update");
> >     return 0;
> > }
> >
> > SEC("fentry/trie_update_elem")
> > int program2()
> > {
> >     bpf_printk("trie_update_elem");
> >     return 0;
> > }
> >
> > The result as follows,
> >
> >          kubelet-203203  [018] ....1  9579.862862:
> > __bpf_prog_enter_recur: __bpf_prog_enter_recur
> >          kubelet-203203  [018] ...11  9579.862869: bpf_trace_printk:
> > before update
> >          kubelet-203203  [018] ....2  9579.862869:
> > __bpf_prog_enter_recur: __bpf_prog_enter_recur
> >          kubelet-203203  [018] ...12  9579.862870: bpf_trace_printk:
> > trie_update_elem
> >          kubelet-203203  [018] ....2  9579.862870:
> > __bpf_prog_exit_recur: __bpf_prog_exit_recur
> >          kubelet-203203  [018] ...11  9579.862870: bpf_trace_printk:
> > after update
> >          kubelet-203203  [018] ....1  9579.862871:
> > __bpf_prog_exit_recur: __bpf_prog_exit_recur
> >
> > Note that we can't trace __bpf_prog_enter_recur and
> > __bpf_prog_exit_recur, so we have to modify the kernel to print them.
> >
>
> ... However, surprisingly it still works even after this patchset is
> applied, because the hardirq/softirq flag is set when the program2 is
> running, see also the flags in the above trace_pipe output. Is that
> expected ?!
> I need  some time to figure it out, but maybe you have a quick answer...

Answer it by myself, that is because of the
allowing-one-single-recursion rule. I misread the trace flags before.
Sorry about the noise.
Steven Rostedt April 27, 2023, 1:26 p.m. UTC | #10
On Mon, 17 Apr 2023 15:47:36 +0000
Yafang Shao <laoar.shao@gmail.com> wrote:

> diff --git a/kernel/bpf/trampoline.c b/kernel/bpf/trampoline.c
> index f61d513..3df39a5 100644
> --- a/kernel/bpf/trampoline.c
> +++ b/kernel/bpf/trampoline.c
> @@ -842,15 +842,21 @@ static __always_inline u64 notrace bpf_prog_start_time(void)
>  static u64 notrace __bpf_prog_enter_recur(struct bpf_prog *prog, struct bpf_tramp_run_ctx *run_ctx)
>  	__acquires(RCU)

Because __bpf_prog_enter_recur() and __bpf_prog_exit_recur() can
legitimately nest (as you pointed out later in the thread), I think my
original plan is the way to go.



>  {
> -	rcu_read_lock();
> -	migrate_disable();
> -
> -	run_ctx->saved_run_ctx = bpf_set_run_ctx(&run_ctx->run_ctx);
> +	int bit;
>  
> -	if (unlikely(this_cpu_inc_return(*(prog->active)) != 1)) {
> +	rcu_read_lock();
> +	bit = test_recursion_try_acquire(_THIS_IP_, _RET_IP_);
> +	run_ctx->recursion_bit = bit;
> +	if (bit < 0) {
> +		preempt_disable_notrace();
>  		bpf_prog_inc_misses_counter(prog);
> +		preempt_enable_notrace();
>  		return 0;
>  	}
> +
> +	migrate_disable();

Just encompass the migrate_disable/enable() with the recursion protection.

That is, here add:

	test_recursion_release(recursion_bit);

No need to save it in the run_ctx, as you can use a local variable.

As I mentioned, if it passes when checking migrate_disable() it will also
pass when checking around migrate_enable() so the two will still be paired
properly, even if only the migrate_enable() starts recursing.


  bit = test_recursion_try_acquire() // OK
  if (bit < 0)
	return;
  migrate_disable();
  test_recursion_release(bit);

  [..]

  bit = test_recursion_try_acquire() // OK
  migrate_enable() // traced and recurses...

    bit = test_recursion_try_acquire() // fails
    if (bit < 0)
          return; // returns here
    migrate_disable() // does not get called.

The recursion around migrate_disable/enable() is needed because it's done
before other checks. You can't attach the test_recursion logic to the
__bpf_prog_enter/exit() routines, because those can legitimately recurse.

-- Steve


> +
> +	run_ctx->saved_run_ctx = bpf_set_run_ctx(&run_ctx->run_ctx);
>  	return bpf_prog_start_time();
>  }
Yafang Shao April 27, 2023, 2:22 p.m. UTC | #11
On Thu, Apr 27, 2023 at 9:26 PM Steven Rostedt <rostedt@goodmis.org> wrote:
>
> On Mon, 17 Apr 2023 15:47:36 +0000
> Yafang Shao <laoar.shao@gmail.com> wrote:
>
> > diff --git a/kernel/bpf/trampoline.c b/kernel/bpf/trampoline.c
> > index f61d513..3df39a5 100644
> > --- a/kernel/bpf/trampoline.c
> > +++ b/kernel/bpf/trampoline.c
> > @@ -842,15 +842,21 @@ static __always_inline u64 notrace bpf_prog_start_time(void)
> >  static u64 notrace __bpf_prog_enter_recur(struct bpf_prog *prog, struct bpf_tramp_run_ctx *run_ctx)
> >       __acquires(RCU)
>
> Because __bpf_prog_enter_recur() and __bpf_prog_exit_recur() can
> legitimately nest (as you pointed out later in the thread), I think my
> original plan is the way to go.
>
>
>
> >  {
> > -     rcu_read_lock();
> > -     migrate_disable();
> > -
> > -     run_ctx->saved_run_ctx = bpf_set_run_ctx(&run_ctx->run_ctx);
> > +     int bit;
> >
> > -     if (unlikely(this_cpu_inc_return(*(prog->active)) != 1)) {
> > +     rcu_read_lock();
> > +     bit = test_recursion_try_acquire(_THIS_IP_, _RET_IP_);
> > +     run_ctx->recursion_bit = bit;
> > +     if (bit < 0) {
> > +             preempt_disable_notrace();
> >               bpf_prog_inc_misses_counter(prog);
> > +             preempt_enable_notrace();
> >               return 0;
> >       }
> > +
> > +     migrate_disable();
>
> Just encompass the migrate_disable/enable() with the recursion protection.
>
> That is, here add:
>
>         test_recursion_release(recursion_bit);
>
> No need to save it in the run_ctx, as you can use a local variable.
>
> As I mentioned, if it passes when checking migrate_disable() it will also
> pass when checking around migrate_enable() so the two will still be paired
> properly, even if only the migrate_enable() starts recursing.
>
>
>   bit = test_recursion_try_acquire() // OK
>   if (bit < 0)
>         return;
>   migrate_disable();
>   test_recursion_release(bit);
>
>   [..]
>
>   bit = test_recursion_try_acquire() // OK
>   migrate_enable() // traced and recurses...
>
>     bit = test_recursion_try_acquire() // fails
>     if (bit < 0)
>           return; // returns here
>     migrate_disable() // does not get called.
>
> The recursion around migrate_disable/enable() is needed because it's done
> before other checks. You can't attach the test_recursion logic to the
> __bpf_prog_enter/exit() routines, because those can legitimately recurse.
>

IIUC, the acquire/release pair works as follows,

   test_recursion_try_acquire
     [ protection area ]
   test_recursion_release

After release, there will be no protection, and thus it will fail the
tools/testing/selftests/bpf/progs/recursion.c[1] test case, because
the recursion occurs in the bpf_prog_run() itself,

  __bpf_prog_enter
     test_recursion_try_acquire
     [...]
     test_recursion_release
  // no protection after the release
  bpf_prog_run()
    bpf_prog_run() // the recursion can't be prevented.
        __bpf_prog_enter
            test_recursion_try_acquire
            [...]
            test_recursion_release
       bpf_prog_run()
           bpf_prog_run()
               __bpf_prog_enter
                  test_recursion_try_acquire
                  [...]
                  test_recursion_release
              bpf_prog_run()
              [ And so on ... ]

[1]. https://git.kernel.org/pub/scm/linux/kernel/git/bpf/bpf-next.git/tree/tools/testing/selftests/bpf/progs/recursion.c#n38
Steven Rostedt April 27, 2023, 3:18 p.m. UTC | #12
On Thu, 27 Apr 2023 22:22:22 +0800
Yafang Shao <laoar.shao@gmail.com> wrote:

> IIUC, the acquire/release pair works as follows,
> 
>    test_recursion_try_acquire
>      [ protection area ]
>    test_recursion_release
> 
> After release, there will be no protection, and thus it will fail the
> tools/testing/selftests/bpf/progs/recursion.c[1] test case, because
> the recursion occurs in the bpf_prog_run() itself,

But bpf programs are allowed to recurs. Hence, you need separate logic to
detect that. The test_recursion_*() code is for cases that are not allowed
to recurs.

> 
>   __bpf_prog_enter
>      test_recursion_try_acquire
>      [...]
>      test_recursion_release
>   // no protection after the release
>   bpf_prog_run()
>     bpf_prog_run() // the recursion can't be prevented.

But I thought you can run a bpf_prog from another bpf_prog. So you don't
want to prevent it. You need other logic to detect if it was not suppose to
recurs.

-- Steve


>         __bpf_prog_enter
>             test_recursion_try_acquire
>             [...]
>             test_recursion_release
>        bpf_prog_run()
>            bpf_prog_run()
>                __bpf_prog_enter
>                   test_recursion_try_acquire
>                   [...]
>                   test_recursion_release
>               bpf_prog_run()
>               [ And so on ... ]
> 
> [1]. https://git.kernel.org/pub/scm/linux/kernel/git/bpf/bpf-next.git/tree/tools/testing/selftests/bpf/progs/recursion.c#n38
Yafang Shao April 27, 2023, 3:23 p.m. UTC | #13
On Thu, Apr 27, 2023 at 11:18 PM Steven Rostedt <rostedt@goodmis.org> wrote:
>
> On Thu, 27 Apr 2023 22:22:22 +0800
> Yafang Shao <laoar.shao@gmail.com> wrote:
>
> > IIUC, the acquire/release pair works as follows,
> >
> >    test_recursion_try_acquire
> >      [ protection area ]
> >    test_recursion_release
> >
> > After release, there will be no protection, and thus it will fail the
> > tools/testing/selftests/bpf/progs/recursion.c[1] test case, because
> > the recursion occurs in the bpf_prog_run() itself,
>
> But bpf programs are allowed to recurs. Hence, you need separate logic to
> detect that. The test_recursion_*() code is for cases that are not allowed
> to recurs.
>

Agreed.

> >
> >   __bpf_prog_enter
> >      test_recursion_try_acquire
> >      [...]
> >      test_recursion_release
> >   // no protection after the release
> >   bpf_prog_run()
> >     bpf_prog_run() // the recursion can't be prevented.
>
> But I thought you can run a bpf_prog from another bpf_prog. So you don't
> want to prevent it. You need other logic to detect if it was not suppose to
> recurs.
>

If so, we have to keep the prog->active to prevent it, then I'm not
sure if it is worth adding test_recursion_*().
Steven Rostedt April 27, 2023, 3:36 p.m. UTC | #14
On Thu, 27 Apr 2023 23:23:31 +0800
Yafang Shao <laoar.shao@gmail.com> wrote:

> > But I thought you can run a bpf_prog from another bpf_prog. So you don't
> > want to prevent it. You need other logic to detect if it was not suppose to
> > recurs.
> >  
> 
> If so, we have to keep the prog->active to prevent it, then I'm not
> sure if it is worth adding test_recursion_*().

I thought that the whole point of this exercise was because the
migrate_disable() itself could be traced (or call something that can), and
that's outside of prog->active protection. Which the test_recursion_*()
code was created for.

-- Steve
Alexei Starovoitov April 27, 2023, 3:39 p.m. UTC | #15
On Thu, Apr 27, 2023 at 8:36 AM Steven Rostedt <rostedt@goodmis.org> wrote:
>
> On Thu, 27 Apr 2023 23:23:31 +0800
> Yafang Shao <laoar.shao@gmail.com> wrote:
>
> > > But I thought you can run a bpf_prog from another bpf_prog. So you don't
> > > want to prevent it. You need other logic to detect if it was not suppose to
> > > recurs.
> > >
> >
> > If so, we have to keep the prog->active to prevent it, then I'm not
> > sure if it is worth adding test_recursion_*().
>
> I thought that the whole point of this exercise was because the
> migrate_disable() itself could be traced (or call something that can), and
> that's outside of prog->active protection. Which the test_recursion_*()
> code was created for.

Not sure where did this come from.
migrate_enable/disable were added to deny list back in 2021.
Yafang Shao April 27, 2023, 3:43 p.m. UTC | #16
On Thu, Apr 27, 2023 at 11:39 PM Alexei Starovoitov
<alexei.starovoitov@gmail.com> wrote:
>
> On Thu, Apr 27, 2023 at 8:36 AM Steven Rostedt <rostedt@goodmis.org> wrote:
> >
> > On Thu, 27 Apr 2023 23:23:31 +0800
> > Yafang Shao <laoar.shao@gmail.com> wrote:
> >
> > > > But I thought you can run a bpf_prog from another bpf_prog. So you don't
> > > > want to prevent it. You need other logic to detect if it was not suppose to
> > > > recurs.
> > > >
> > >
> > > If so, we have to keep the prog->active to prevent it, then I'm not
> > > sure if it is worth adding test_recursion_*().
> >
> > I thought that the whole point of this exercise was because the
> > migrate_disable() itself could be traced (or call something that can), and
> > that's outside of prog->active protection. Which the test_recursion_*()
> > code was created for.
>
> Not sure where did this come from.
> migrate_enable/disable were added to deny list back in 2021.

Hi Alexei,

Don't be uneasy.  It is not good to play word games.
What Steven really meant is the preempt_count_{sub, add}.
Anyway thanks Steven for the help with this exercise.
Steven Rostedt April 27, 2023, 3:46 p.m. UTC | #17
On Thu, 27 Apr 2023 23:43:35 +0800
Yafang Shao <laoar.shao@gmail.com> wrote:

> > > I thought that the whole point of this exercise was because the
> > > migrate_disable() itself could be traced (or call something that can), and
> > > that's outside of prog->active protection. Which the test_recursion_*()
> > > code was created for.  
> >
> > Not sure where did this come from.
> > migrate_enable/disable were added to deny list back in 2021.  
> 
> Hi Alexei,
> 
> Don't be uneasy.  It is not good to play word games.
> What Steven really meant is the preempt_count_{sub, add}.
> Anyway thanks Steven for the help with this exercise.

Right, it was the "(or call something that can)" part that this came from.
As Yafang said, migrate_disable() calls preempt_count_add() (on some
configs) which is traced by ftrace, and thus traced by bpf. Or was that
added to the deny list? I think that was one of the solutions as well.

-- Steve
diff mbox series

Patch

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 18b592f..c42ff90 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -1467,7 +1467,6 @@  struct bpf_prog {
 	u32			jited_len;	/* Size of jited insns in bytes */
 	u8			tag[BPF_TAG_SIZE];
 	struct bpf_prog_stats __percpu *stats;
-	int __percpu		*active;
 	unsigned int		(*bpf_func)(const void *ctx,
 					    const struct bpf_insn *insn);
 	struct bpf_prog_aux	*aux;		/* Auxiliary fields */
@@ -1813,6 +1812,7 @@  struct bpf_tramp_run_ctx {
 	struct bpf_run_ctx run_ctx;
 	u64 bpf_cookie;
 	struct bpf_run_ctx *saved_run_ctx;
+	int recursion_bit;
 };
 
 static inline struct bpf_run_ctx *bpf_set_run_ctx(struct bpf_run_ctx *new_ctx)
diff --git a/kernel/bpf/core.c b/kernel/bpf/core.c
index 7421487..0942ab2 100644
--- a/kernel/bpf/core.c
+++ b/kernel/bpf/core.c
@@ -103,12 +103,6 @@  struct bpf_prog *bpf_prog_alloc_no_stats(unsigned int size, gfp_t gfp_extra_flag
 		vfree(fp);
 		return NULL;
 	}
-	fp->active = alloc_percpu_gfp(int, bpf_memcg_flags(GFP_KERNEL | gfp_extra_flags));
-	if (!fp->active) {
-		vfree(fp);
-		kfree(aux);
-		return NULL;
-	}
 
 	fp->pages = size / PAGE_SIZE;
 	fp->aux = aux;
@@ -138,7 +132,6 @@  struct bpf_prog *bpf_prog_alloc(unsigned int size, gfp_t gfp_extra_flags)
 
 	prog->stats = alloc_percpu_gfp(struct bpf_prog_stats, gfp_flags);
 	if (!prog->stats) {
-		free_percpu(prog->active);
 		kfree(prog->aux);
 		vfree(prog);
 		return NULL;
@@ -256,7 +249,6 @@  struct bpf_prog *bpf_prog_realloc(struct bpf_prog *fp_old, unsigned int size,
 		 */
 		fp_old->aux = NULL;
 		fp_old->stats = NULL;
-		fp_old->active = NULL;
 		__bpf_prog_free(fp_old);
 	}
 
@@ -272,7 +264,6 @@  void __bpf_prog_free(struct bpf_prog *fp)
 		kfree(fp->aux);
 	}
 	free_percpu(fp->stats);
-	free_percpu(fp->active);
 	vfree(fp);
 }
 
@@ -1385,7 +1376,6 @@  static void bpf_prog_clone_free(struct bpf_prog *fp)
 	 */
 	fp->aux = NULL;
 	fp->stats = NULL;
-	fp->active = NULL;
 	__bpf_prog_free(fp);
 }
 
diff --git a/kernel/bpf/trampoline.c b/kernel/bpf/trampoline.c
index f61d513..3df39a5 100644
--- a/kernel/bpf/trampoline.c
+++ b/kernel/bpf/trampoline.c
@@ -842,15 +842,21 @@  static __always_inline u64 notrace bpf_prog_start_time(void)
 static u64 notrace __bpf_prog_enter_recur(struct bpf_prog *prog, struct bpf_tramp_run_ctx *run_ctx)
 	__acquires(RCU)
 {
-	rcu_read_lock();
-	migrate_disable();
-
-	run_ctx->saved_run_ctx = bpf_set_run_ctx(&run_ctx->run_ctx);
+	int bit;
 
-	if (unlikely(this_cpu_inc_return(*(prog->active)) != 1)) {
+	rcu_read_lock();
+	bit = test_recursion_try_acquire(_THIS_IP_, _RET_IP_);
+	run_ctx->recursion_bit = bit;
+	if (bit < 0) {
+		preempt_disable_notrace();
 		bpf_prog_inc_misses_counter(prog);
+		preempt_enable_notrace();
 		return 0;
 	}
+
+	migrate_disable();
+
+	run_ctx->saved_run_ctx = bpf_set_run_ctx(&run_ctx->run_ctx);
 	return bpf_prog_start_time();
 }
 
@@ -880,11 +886,16 @@  static void notrace __bpf_prog_exit_recur(struct bpf_prog *prog, u64 start,
 					  struct bpf_tramp_run_ctx *run_ctx)
 	__releases(RCU)
 {
+	if (run_ctx->recursion_bit < 0)
+		goto out;
+
 	bpf_reset_run_ctx(run_ctx->saved_run_ctx);
 
 	update_prog_stats(prog, start);
-	this_cpu_dec(*(prog->active));
 	migrate_enable();
+	test_recursion_release(run_ctx->recursion_bit);
+
+out:
 	rcu_read_unlock();
 }
 
@@ -916,15 +927,21 @@  static void notrace __bpf_prog_exit_lsm_cgroup(struct bpf_prog *prog, u64 start,
 u64 notrace __bpf_prog_enter_sleepable_recur(struct bpf_prog *prog,
 					     struct bpf_tramp_run_ctx *run_ctx)
 {
-	rcu_read_lock_trace();
-	migrate_disable();
-	might_fault();
+	int bit;
 
-	if (unlikely(this_cpu_inc_return(*(prog->active)) != 1)) {
+	rcu_read_lock_trace();
+	bit = test_recursion_try_acquire(_THIS_IP_, _RET_IP_);
+	run_ctx->recursion_bit = bit;
+	if (bit < 0) {
+		preempt_disable_notrace();
 		bpf_prog_inc_misses_counter(prog);
+		preempt_enable_notrace();
 		return 0;
 	}
 
+	migrate_disable();
+	might_fault();
+
 	run_ctx->saved_run_ctx = bpf_set_run_ctx(&run_ctx->run_ctx);
 
 	return bpf_prog_start_time();
@@ -933,11 +950,16 @@  u64 notrace __bpf_prog_enter_sleepable_recur(struct bpf_prog *prog,
 void notrace __bpf_prog_exit_sleepable_recur(struct bpf_prog *prog, u64 start,
 					     struct bpf_tramp_run_ctx *run_ctx)
 {
+	if (run_ctx->recursion_bit < 0)
+		goto out;
+
 	bpf_reset_run_ctx(run_ctx->saved_run_ctx);
 
 	update_prog_stats(prog, start);
-	this_cpu_dec(*(prog->active));
 	migrate_enable();
+	test_recursion_release(run_ctx->recursion_bit);
+
+out:
 	rcu_read_unlock_trace();
 }
 
diff --git a/kernel/trace/bpf_trace.c b/kernel/trace/bpf_trace.c
index bcf91bc..bb9a4c9 100644
--- a/kernel/trace/bpf_trace.c
+++ b/kernel/trace/bpf_trace.c
@@ -2250,16 +2250,18 @@  void bpf_put_raw_tracepoint(struct bpf_raw_event_map *btp)
 static __always_inline
 void __bpf_trace_run(struct bpf_prog *prog, u64 *args)
 {
-	cant_sleep();
-	if (unlikely(this_cpu_inc_return(*(prog->active)) != 1)) {
+	int bit;
+
+	bit = test_recursion_try_acquire(_THIS_IP_, _RET_IP_);
+	if (bit < 0) {
 		bpf_prog_inc_misses_counter(prog);
-		goto out;
+		return;
 	}
+	cant_sleep();
 	rcu_read_lock();
 	(void) bpf_prog_run(prog, args);
 	rcu_read_unlock();
-out:
-	this_cpu_dec(*(prog->active));
+	test_recursion_release(bit);
 }
 
 #define UNPACK(...)			__VA_ARGS__