[PATCH bpf] bpf: Take module reference for ip in module code

[Jiri Olsa]

Currently module can be unloaded even if there's a trampoline
register in it. It's easily reproduced by running in parallel:

  # while :; do ./test_progs -t module_attach; done
  # while :; do ./test_progs -t fentry_test; done

Taking the module reference in case the trampoline's ip is
within the module code. Releasing it when the trampoline's
ip is unregistered.

Signed-off-by: Jiri Olsa <jo...@kernel.org>
---
 kernel/bpf/trampoline.c | 32 ++++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/kernel/bpf/trampoline.c b/kernel/bpf/trampoline.c
index 1f3a4be4b175..f6cb179842b2 100644
--- a/kernel/bpf/trampoline.c
+++ b/kernel/bpf/trampoline.c
@@ -87,6 +87,27 @@ static struct bpf_trampoline *bpf_trampoline_lookup(u64 key)
        return tr;
 }
 
+static struct module *ip_module_get(unsigned long ip)
+{
+       struct module *mod;
+       int err = 0;
+
+       preempt_disable();
+       mod = __module_text_address(ip);
+       if (mod && !try_module_get(mod))
+               err = -ENOENT;
+       preempt_enable();
+       return err ? ERR_PTR(err) : mod;
+}
+
+static void ip_module_put(unsigned long ip)
+{
+       struct module *mod = __module_text_address(ip);
+
+       if (mod)
+               module_put(mod);
+}
+
 static int is_ftrace_location(void *ip)
 {
        long addr;
@@ -108,6 +129,9 @@ static int unregister_fentry(struct bpf_trampoline *tr, 
void *old_addr)
                ret = unregister_ftrace_direct((long)ip, (long)old_addr);
        else
                ret = bpf_arch_text_poke(ip, BPF_MOD_CALL, old_addr, NULL);
+
+       if (!ret)
+               ip_module_put((unsigned long) ip);
        return ret;
 }
 
@@ -126,6 +150,7 @@ static int modify_fentry(struct bpf_trampoline *tr, void 
*old_addr, void *new_ad
 /* first time registering */
 static int register_fentry(struct bpf_trampoline *tr, void *new_addr)
 {
+       struct module *mod;
        void *ip = tr->func.addr;
        int ret;
 
@@ -134,10 +159,17 @@ static int register_fentry(struct bpf_trampoline *tr, 
void *new_addr)
                return ret;
        tr->func.ftrace_managed = ret;
 
+       mod = ip_module_get((unsigned long) ip);
+       if (IS_ERR(mod))
+               return -ENOENT;
+
        if (tr->func.ftrace_managed)
                ret = register_ftrace_direct((long)ip, (long)new_addr);
        else
                ret = bpf_arch_text_poke(ip, BPF_MOD_CALL, NULL, new_addr);
+
+       if (ret)
+               module_put(mod);
        return ret;
 }
 
-- 
2.30.2