struct bpf_local_storage_map;
 struct kobject;
 struct mem_cgroup;
+struct module;
 
 extern struct idr btf_idr;
 extern spinlock_t btf_idr_lock;
        /* Executable image of trampoline */
        struct bpf_tramp_image *cur_image;
        u64 selector;
+       struct module *mod;
 };
 
 struct bpf_attach_target_info {
 
 #include <linux/btf.h>
 #include <linux/rcupdate_trace.h>
 #include <linux/rcupdate_wait.h>
+#include <linux/module.h>
 
 /* dummy _ops. The verifier will operate on target program's ops. */
 const struct bpf_verifier_ops bpf_extension_verifier_ops = {
        return tr;
 }
 
+static int bpf_trampoline_module_get(struct bpf_trampoline *tr)
+{
+       struct module *mod;
+       int err = 0;
+
+       preempt_disable();
+       mod = __module_text_address((unsigned long) tr->func.addr);
+       if (mod && !try_module_get(mod))
+               err = -ENOENT;
+       preempt_enable();
+       tr->mod = mod;
+       return err;
+}
+
+static void bpf_trampoline_module_put(struct bpf_trampoline *tr)
+{
+       module_put(tr->mod);
+       tr->mod = NULL;
+}
+
 static int is_ftrace_location(void *ip)
 {
        long 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)
+               bpf_trampoline_module_put(tr);
        return ret;
 }
 
                return ret;
        tr->func.ftrace_managed = ret;
 
+       if (bpf_trampoline_module_get(tr))
+               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)
+               bpf_trampoline_module_put(tr);
        return ret;
 }