diff mbox series

[bpf-next,v6,03/11] bpf: Populate kfunc BTF ID sets in struct btf

Message ID 20220102162115.1506833-4-memxor@gmail.com
State Awaiting Upstream
Delegated to: Pablo Neira
Headers show
Series Introduce unstable CT lookup helpers | expand

Commit Message

Kumar Kartikeya Dwivedi Jan. 2, 2022, 4:21 p.m. UTC 
This patch prepares the kernel to support putting all kinds of kfunc BTF
ID sets in the struct btf itself. The various kernel subsystems will
make register_btf_kfunc_id_set call in the initcalls (for built-in code
and modules).

The 'hook' is one of the many program types, e.g. XDP and TC/SCHED_CLS,
STRUCT_OPS, and 'types' are check (allowed or not), acquire, release,
and ret_null (with PTR_TO_BTF_ID_OR_NULL return type).

A maximum of BTF_KFUNC_SET_MAX_CNT (32) kfunc BTF IDs are permitted in a
set of certain hook and type. They are also allocated on demand, and
otherwise set as NULL.

A new btf_kfunc_id_set_contains function is exposed for use in verifier,
this new method is faster than the existing list searching method, and
is also automatic. It also lets other code not care whether the set is
unallocated or not.

Note that module code can only do single register_btf_kfunc_id_set call
per hook. This is why sort_set is only true for in-kernel vmlinux sets,
because there might be multiple sets for the same hook and type that
must be concatenated, hence sorting them is required to ensure bsearch
in btf_id_set_contains continues to work correctly.

Next commit will update the kernel users to make use of this
infrastructure.

Finally, add __maybe_unused annotation for BTF ID macros for the
!CONFIG_DEBUG_INFO_BTF case, so that they don't produce warnings during
build time.

The previous patch is also needed to provide synchronization against
initialization for module BTF's kfunc_set_tab introduced here, as
described below:

  The kfunc_set_tab pointer in struct btf is write-once (if we consider
  the registration phase (comprised of multiple register_btf_kfunc_id_set
  calls) as a single operation). In this sense, once it has been fully
  prepared, it isn't modified, only used for lookup (from the verifier
  context).

  For btf_vmlinux, it is initialized fully during the do_initcalls phase,
  which happens fairly early in the boot process, before any processes are
  present. This also eliminates the possibility of bpf_check being called
  at that point, thus relieving us of ensuring any synchronization between
  the registration and lookup function (btf_kfunc_id_set_contains).

  However, the case for module BTF is a bit tricky. The BTF is parsed,
  prepared, and published from the MODULE_STATE_COMING notifier callback.
  After this, the module initcalls are invoked, where our registration
  function will be called to populate the kfunc_set_tab for module BTF.

  At this point, BTF may be available to userspace while its corresponding
  module is still intializing. A BTF fd can then be passed to verifier
  using bpf syscall (e.g. for kfunc call insn).

  Hence, there is a race window where verifier may concurrently try to
  lookup the kfunc_set_tab. To prevent this race, we must ensure the
  operations are serialized, or waiting for the __init functions to
  complete.

  In the earlier registration API, this race was alleviated as verifier
  bpf_check_mod_kfunc_call didn't find the kfunc BTF ID until it was added
  by the registration function (called usually at the end of module __init
  function after all module resources have been initialized). If the
  verifier made the check_kfunc_call before kfunc BTF ID was added to the
  list, it would fail verification (saying call isn't allowed). The
  access to list was protected using a mutex.

  Now, it would still fail verification, but for a different reason
  (returning ENXIO due to the failed btf_try_get_module call in
  add_kfunc_call), because if the __init call is in progress the module
  will be in the middle of MODULE_STATE_COMING -> MODULE_STATE_LIVE
  transition, and the try_module_get_live call inside btf_try_get_module
  will fail.

Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 include/linux/btf.h     |  46 ++++++++
 include/linux/btf_ids.h |  13 +--
 kernel/bpf/btf.c        | 228 +++++++++++++++++++++++++++++++++++++++-
 kernel/bpf/verifier.c   |   1 +
 4 files changed, 280 insertions(+), 8 deletions(-)

Comments

Alexei Starovoitov Jan. 5, 2022, 6:19 a.m. UTC | #1 
On Sun, Jan 02, 2022 at 09:51:07PM +0530, Kumar Kartikeya Dwivedi wrote:
>  
> +enum btf_kfunc_hook {
> +	BTF_KFUNC_HOOK_XDP,
> +	BTF_KFUNC_HOOK_TC,
> +	BTF_KFUNC_HOOK_STRUCT_OPS,
> +	_BTF_KFUNC_HOOK_MAX,

Why prefix with _ ?

> +enum {
> +	BTF_KFUNC_SET_MAX_CNT = 32,
> +};
...
> +	if (set_cnt + add_set->cnt > BTF_KFUNC_SET_MAX_CNT) {
> +		ret = -E2BIG;
> +		goto end;
> +	}

This artificial limit wouldn't be needed if you didn't insist on sorting.
The later patches don't take advantage of this sorting feature and
I don't see a test for sorting either.

> +
> +	/* Grow set */
> +	set = krealloc(tab->sets[hook][type], offsetof(struct btf_id_set, ids[set_cnt + add_set->cnt]),
> +		       GFP_KERNEL | __GFP_NOWARN);
> +	if (!set) {
> +		ret = -ENOMEM;
> +		goto end;
> +	}
> +
> +	/* For newly allocated set, initialize set->cnt to 0 */
> +	if (!tab->sets[hook][type])
> +		set->cnt = 0;
> +	tab->sets[hook][type] = set;
> +
> +	/* Concatenate the two sets */
> +	memcpy(set->ids + set->cnt, add_set->ids, add_set->cnt * sizeof(set->ids[0]));
> +	set->cnt += add_set->cnt;

Without sorting this function would just assign the pointer.
No need for krealloc and memcpy.

> +
> +	if (sort_set)
> +		sort(set->ids, set->cnt, sizeof(set->ids[0]), btf_id_cmp_func, NULL);

All that looks like extra code for a dubious feature.

> +bool btf_kfunc_id_set_contains(const struct btf *btf,
> +			       enum bpf_prog_type prog_type,
> +			       enum btf_kfunc_type type, u32 kfunc_btf_id)
> +{
> +	enum btf_kfunc_hook hook;
> +
> +	switch (prog_type) {
> +	case BPF_PROG_TYPE_XDP:
> +		hook = BTF_KFUNC_HOOK_XDP;
> +		break;
> +	case BPF_PROG_TYPE_SCHED_CLS:
> +		hook = BTF_KFUNC_HOOK_TC;
> +		break;
> +	case BPF_PROG_TYPE_STRUCT_OPS:
> +		hook = BTF_KFUNC_HOOK_STRUCT_OPS;
> +		break;
> +	default:
> +		return false;
> +	}

So this switch() is necessary only to compress prog_types into smaller hooks
to save memory in the struct btf_kfunc_set_tab, right ?
If so both kfunc_id_set_contains() and register_btf_kfunc() should
probably use prog_type as an argument for symmetry.
Kumar Kartikeya Dwivedi Jan. 6, 2022, 8:59 a.m. UTC | #2 
On Wed, Jan 05, 2022 at 11:49:11AM IST, Alexei Starovoitov wrote:
> On Sun, Jan 02, 2022 at 09:51:07PM +0530, Kumar Kartikeya Dwivedi wrote:
> >
> > +enum btf_kfunc_hook {
> > +	BTF_KFUNC_HOOK_XDP,
> > +	BTF_KFUNC_HOOK_TC,
> > +	BTF_KFUNC_HOOK_STRUCT_OPS,
> > +	_BTF_KFUNC_HOOK_MAX,
>
> Why prefix with _ ?
>

Will fix.

> > +enum {
> > +	BTF_KFUNC_SET_MAX_CNT = 32,
> > +};
> ...
> > +	if (set_cnt + add_set->cnt > BTF_KFUNC_SET_MAX_CNT) {
> > +		ret = -E2BIG;
> > +		goto end;
> > +	}
>
> This artificial limit wouldn't be needed if you didn't insist on sorting.
> The later patches don't take advantage of this sorting feature and
> I don't see a test for sorting either.
>

I'm not insisting, but for vmlinux we will have multiple
register_btf_kfunc_id_set calls for same hook, so we have to concat multiple
sets into one, which may result in an unsorted set. It's ok to not sort for
modules where only one register call per hook is allowed.

Unless we switch to linear search for now (which is ok by me), we have to
re-sort for vmlinux BTF, to make btf_id_set_contains (in
btf_kfunc_id_set_contains) work.

> > +
> > +	/* Grow set */
> > +	set = krealloc(tab->sets[hook][type], offsetof(struct btf_id_set, ids[set_cnt + add_set->cnt]),
> > +		       GFP_KERNEL | __GFP_NOWARN);
> > +	if (!set) {
> > +		ret = -ENOMEM;
> > +		goto end;
> > +	}
> > +
> > +	/* For newly allocated set, initialize set->cnt to 0 */
> > +	if (!tab->sets[hook][type])
> > +		set->cnt = 0;
> > +	tab->sets[hook][type] = set;
> > +
> > +	/* Concatenate the two sets */
> > +	memcpy(set->ids + set->cnt, add_set->ids, add_set->cnt * sizeof(set->ids[0]));
> > +	set->cnt += add_set->cnt;
>
> Without sorting this function would just assign the pointer.
> No need for krealloc and memcpy.
>

Even if we didn't sort, we'd need to concat multiple sets for vmlinux case, so
krealloc and memcpy would still be needed for the vmlinux BTF case, right? For
modules I could certainly do a direct assignment, even if we keep sorting,
because only one set per hook is permitted.

> > +
> > +	if (sort_set)
> > +		sort(set->ids, set->cnt, sizeof(set->ids[0]), btf_id_cmp_func, NULL);
>
> All that looks like extra code for a dubious feature.
>

It's needed for the vmlinux case. I use WARN_ON_ONCE when modules try to
register more than one set for a certain hook.

> > +bool btf_kfunc_id_set_contains(const struct btf *btf,
> > +			       enum bpf_prog_type prog_type,
> > +			       enum btf_kfunc_type type, u32 kfunc_btf_id)
> > +{
> > +	enum btf_kfunc_hook hook;
> > +
> > +	switch (prog_type) {
> > +	case BPF_PROG_TYPE_XDP:
> > +		hook = BTF_KFUNC_HOOK_XDP;
> > +		break;
> > +	case BPF_PROG_TYPE_SCHED_CLS:
> > +		hook = BTF_KFUNC_HOOK_TC;
> > +		break;
> > +	case BPF_PROG_TYPE_STRUCT_OPS:
> > +		hook = BTF_KFUNC_HOOK_STRUCT_OPS;
> > +		break;
> > +	default:
> > +		return false;
> > +	}
>
> So this switch() is necessary only to compress prog_types into smaller hooks
> to save memory in the struct btf_kfunc_set_tab, right ?
> If so both kfunc_id_set_contains() and register_btf_kfunc() should
> probably use prog_type as an argument for symmetry.

Ok, will fix.

--
Kartikeya
Alexei Starovoitov Jan. 6, 2022, 11:40 p.m. UTC | #3 
On Thu, Jan 6, 2022 at 12:59 AM Kumar Kartikeya Dwivedi
<memxor@gmail.com> wrote:
>
> I'm not insisting, but for vmlinux we will have multiple
> register_btf_kfunc_id_set calls for same hook, so we have to concat multiple
> sets into one, which may result in an unsorted set. It's ok to not sort for
> modules where only one register call per hook is allowed.
>
> Unless we switch to linear search for now (which is ok by me), we have to
> re-sort for vmlinux BTF, to make btf_id_set_contains (in
> btf_kfunc_id_set_contains) work.

Could you give an example when it happens in vmlinux?
Is it when modules are built-in (like tcp_bbr.c, tcp_cubic.c) ?
But then they should go into struct btf of the module?
Or THIS_MODULE becomes vmlinux and everything goes there?
If that's the case then yeah, sorting is needed.
Please make it clear in the comment and the commit log.
It should be clear to reviewers without having to grill the author
with questions.
Would certainly save time for both author and reviewer. Agree?
Kumar Kartikeya Dwivedi Jan. 7, 2022, 6:26 a.m. UTC | #4 
On Fri, Jan 07, 2022 at 05:10:56AM IST, Alexei Starovoitov wrote:
> On Thu, Jan 6, 2022 at 12:59 AM Kumar Kartikeya Dwivedi
> <memxor@gmail.com> wrote:
> >
> > I'm not insisting, but for vmlinux we will have multiple
> > register_btf_kfunc_id_set calls for same hook, so we have to concat multiple
> > sets into one, which may result in an unsorted set. It's ok to not sort for
> > modules where only one register call per hook is allowed.
> >
> > Unless we switch to linear search for now (which is ok by me), we have to
> > re-sort for vmlinux BTF, to make btf_id_set_contains (in
> > btf_kfunc_id_set_contains) work.
>
> Could you give an example when it happens in vmlinux?
> Is it when modules are built-in (like tcp_bbr.c, tcp_cubic.c) ?
> But then they should go into struct btf of the module?
> Or THIS_MODULE becomes vmlinux and everything goes there?
> If that's the case then yeah, sorting is needed.

Yep, THIS_MODULE would be NULL, and it would pick vmlinux BTF for storing the
set.

Your suggestion to do direct assignment for module case is also good, since we
always ensure module refcount is held when looking into set, access to it should
be safe.

> Please make it clear in the comment and the commit log.
> It should be clear to reviewers without having to grill the author
> with questions.
> Would certainly save time for both author and reviewer. Agree?

Agreed, I'll add the comments and respin.

Thanks for your review.

--
Kartikeya
diff mbox series

Patch

diff --git a/include/linux/btf.h b/include/linux/btf.h
index 0c74348cbc9d..01c8f8ad30c6 100644
--- a/include/linux/btf.h
+++ b/include/linux/btf.h
@@ -12,11 +12,40 @@ 
 #define BTF_TYPE_EMIT(type) ((void)(type *)0)
 #define BTF_TYPE_EMIT_ENUM(enum_val) ((void)enum_val)
 
+enum btf_kfunc_hook {
+	BTF_KFUNC_HOOK_XDP,
+	BTF_KFUNC_HOOK_TC,
+	BTF_KFUNC_HOOK_STRUCT_OPS,
+	_BTF_KFUNC_HOOK_MAX,
+};
+
+enum btf_kfunc_type {
+	BTF_KFUNC_TYPE_CHECK,
+	BTF_KFUNC_TYPE_ACQUIRE,
+	BTF_KFUNC_TYPE_RELEASE,
+	BTF_KFUNC_TYPE_RET_NULL,
+	_BTF_KFUNC_TYPE_MAX,
+};
+
 struct btf;
 struct btf_member;
 struct btf_type;
 union bpf_attr;
 struct btf_show;
+struct btf_id_set;
+
+struct btf_kfunc_id_set {
+	struct module *owner;
+	union {
+		struct {
+			struct btf_id_set *check_set;
+			struct btf_id_set *acquire_set;
+			struct btf_id_set *release_set;
+			struct btf_id_set *ret_null_set;
+		};
+		struct btf_id_set *sets[_BTF_KFUNC_TYPE_MAX];
+	};
+};
 
 extern const struct file_operations btf_fops;
 
@@ -307,6 +336,11 @@  const struct btf_type *btf_type_by_id(const struct btf *btf, u32 type_id);
 const char *btf_name_by_offset(const struct btf *btf, u32 offset);
 struct btf *btf_parse_vmlinux(void);
 struct btf *bpf_prog_get_target_btf(const struct bpf_prog *prog);
+bool btf_kfunc_id_set_contains(const struct btf *btf,
+			       enum bpf_prog_type prog_type,
+			       enum btf_kfunc_type type, u32 kfunc_btf_id);
+int register_btf_kfunc_id_set(enum btf_kfunc_hook hook,
+			      const struct btf_kfunc_id_set *s);
 #else
 static inline const struct btf_type *btf_type_by_id(const struct btf *btf,
 						    u32 type_id)
@@ -318,6 +352,18 @@  static inline const char *btf_name_by_offset(const struct btf *btf,
 {
 	return NULL;
 }
+static inline bool btf_kfunc_id_set_contains(const struct btf *btf,
+					     enum bpf_prog_type prog_type,
+					     enum btf_kfunc_type type,
+					     u32 kfunc_btf_id)
+{
+	return false;
+}
+static inline int register_btf_kfunc_id_set(enum btf_kfunc_hook hook,
+					    const struct btf_kfunc_id_set *s)
+{
+	return 0;
+}
 #endif
 
 struct kfunc_btf_id_set {
diff --git a/include/linux/btf_ids.h b/include/linux/btf_ids.h
index 919c0fde1c51..bc5d9cc34e4c 100644
--- a/include/linux/btf_ids.h
+++ b/include/linux/btf_ids.h
@@ -11,6 +11,7 @@  struct btf_id_set {
 #ifdef CONFIG_DEBUG_INFO_BTF
 
 #include <linux/compiler.h> /* for __PASTE */
+#include <linux/compiler_attributes.h> /* for __maybe_unused */
 
 /*
  * Following macros help to define lists of BTF IDs placed
@@ -146,14 +147,14 @@  extern struct btf_id_set name;
 
 #else
 
-#define BTF_ID_LIST(name) static u32 name[5];
+#define BTF_ID_LIST(name) static u32 __maybe_unused name[5];
 #define BTF_ID(prefix, name)
 #define BTF_ID_UNUSED
-#define BTF_ID_LIST_GLOBAL(name, n) u32 name[n];
-#define BTF_ID_LIST_SINGLE(name, prefix, typename) static u32 name[1];
-#define BTF_ID_LIST_GLOBAL_SINGLE(name, prefix, typename) u32 name[1];
-#define BTF_SET_START(name) static struct btf_id_set name = { 0 };
-#define BTF_SET_START_GLOBAL(name) static struct btf_id_set name = { 0 };
+#define BTF_ID_LIST_GLOBAL(name, n) u32 __maybe_unused name[n];
+#define BTF_ID_LIST_SINGLE(name, prefix, typename) static u32 __maybe_unused name[1];
+#define BTF_ID_LIST_GLOBAL_SINGLE(name, prefix, typename) u32 __maybe_unused name[1];
+#define BTF_SET_START(name) static struct btf_id_set __maybe_unused name = { 0 };
+#define BTF_SET_START_GLOBAL(name) static struct btf_id_set __maybe_unused name = { 0 };
 #define BTF_SET_END(name)
 
 #endif /* CONFIG_DEBUG_INFO_BTF */
diff --git a/kernel/bpf/btf.c b/kernel/bpf/btf.c
index b5b423de53ab..02c33b2ad47e 100644
--- a/kernel/bpf/btf.c
+++ b/kernel/bpf/btf.c
@@ -198,6 +198,14 @@ 
 DEFINE_IDR(btf_idr);
 DEFINE_SPINLOCK(btf_idr_lock);
 
+enum {
+	BTF_KFUNC_SET_MAX_CNT = 32,
+};
+
+struct btf_kfunc_set_tab {
+	struct btf_id_set *sets[_BTF_KFUNC_HOOK_MAX][_BTF_KFUNC_TYPE_MAX];
+};
+
 struct btf {
 	void *data;
 	struct btf_type **types;
@@ -212,6 +220,7 @@  struct btf {
 	refcount_t refcnt;
 	u32 id;
 	struct rcu_head rcu;
+	struct btf_kfunc_set_tab *kfunc_set_tab;
 
 	/* split BTF support */
 	struct btf *base_btf;
@@ -1531,8 +1540,23 @@  static void btf_free_id(struct btf *btf)
 	spin_unlock_irqrestore(&btf_idr_lock, flags);
 }
 
+static void btf_free_kfunc_set_tab(struct btf *btf)
+{
+	struct btf_kfunc_set_tab *tab = btf->kfunc_set_tab;
+	int hook, type;
+
+	if (!tab)
+		return;
+	for (hook = 0; hook < ARRAY_SIZE(tab->sets); hook++) {
+		for (type = 0; type < ARRAY_SIZE(tab->sets[0]); type++)
+			kfree(tab->sets[hook][type]);
+	}
+	btf->kfunc_set_tab = NULL;
+}
+
 static void btf_free(struct btf *btf)
 {
+	btf_free_kfunc_set_tab(btf);
 	kvfree(btf->types);
 	kvfree(btf->resolved_sizes);
 	kvfree(btf->resolved_ids);
@@ -6341,8 +6365,11 @@  struct module *btf_try_get_module(const struct btf *btf)
 		/* We must only consider module whose __init routine has
 		 * finished, hence use try_module_get_live.
 		 */
-		if (try_module_get_live(btf_mod->module))
+		if (try_module_get_live(btf_mod->module)) {
+			/* pairs with smp_wmb in register_btf_kfunc_id_set */
+			smp_rmb();
 			res = btf_mod->module;
+		}
 
 		break;
 	}
@@ -6352,6 +6379,36 @@  struct module *btf_try_get_module(const struct btf *btf)
 	return res;
 }
 
+/* Returns struct btf corresponding to the struct module
+ *
+ * This function can return NULL or ERR_PTR. Note that caller must
+ * release reference for struct btf iff btf_is_module is true.
+ */
+static struct btf *btf_get_module_btf(const struct module *module)
+{
+	struct btf *btf = NULL;
+#ifdef CONFIG_DEBUG_INFO_BTF_MODULES
+	struct btf_module *btf_mod, *tmp;
+#endif
+
+	if (!module)
+		return bpf_get_btf_vmlinux();
+#ifdef CONFIG_DEBUG_INFO_BTF_MODULES
+	mutex_lock(&btf_module_mutex);
+	list_for_each_entry_safe(btf_mod, tmp, &btf_modules, list) {
+		if (btf_mod->module != module)
+			continue;
+
+		btf_get(btf_mod->btf);
+		btf = btf_mod->btf;
+		break;
+	}
+	mutex_unlock(&btf_module_mutex);
+#endif
+
+	return btf;
+}
+
 BPF_CALL_4(bpf_btf_find_by_name_kind, char *, name, int, name_sz, u32, kind, int, flags)
 {
 	struct btf *btf;
@@ -6419,7 +6476,174 @@  BTF_ID_LIST_GLOBAL(btf_tracing_ids, MAX_BTF_TRACING_TYPE)
 BTF_TRACING_TYPE_xxx
 #undef BTF_TRACING_TYPE
 
-/* BTF ID set registration API for modules */
+/* Kernel Function (kfunc) BTF ID set registration API */
+
+static int __btf_populate_kfunc_set(struct btf *btf, enum btf_kfunc_hook hook,
+				    enum btf_kfunc_type type,
+				    struct btf_id_set *add_set, bool sort_set)
+{
+	struct btf_kfunc_set_tab *tab;
+	struct btf_id_set *set;
+	u32 set_cnt;
+	int ret;
+
+	if (WARN_ON_ONCE(hook >= _BTF_KFUNC_HOOK_MAX || type >= _BTF_KFUNC_TYPE_MAX)) {
+		ret = -EINVAL;
+		goto end;
+	}
+
+	if (WARN_ON_ONCE(btf_is_module(btf) && sort_set)) {
+		ret = -EINVAL;
+		goto end;
+	}
+
+	if (!add_set->cnt)
+		return 0;
+
+	tab = btf->kfunc_set_tab;
+	if (!tab) {
+		tab = kzalloc(sizeof(*tab), GFP_KERNEL | __GFP_NOWARN);
+		if (!tab)
+			return -ENOMEM;
+		btf->kfunc_set_tab = tab;
+	}
+
+	set = tab->sets[hook][type];
+	if (WARN_ON_ONCE(set && !sort_set)) {
+		ret = -EINVAL;
+		goto end;
+	}
+
+	set_cnt = set ? set->cnt : 0;
+
+	if (set_cnt > U32_MAX - add_set->cnt) {
+		ret = -EOVERFLOW;
+		goto end;
+	}
+
+	if (set_cnt + add_set->cnt > BTF_KFUNC_SET_MAX_CNT) {
+		ret = -E2BIG;
+		goto end;
+	}
+
+	/* Grow set */
+	set = krealloc(tab->sets[hook][type], offsetof(struct btf_id_set, ids[set_cnt + add_set->cnt]),
+		       GFP_KERNEL | __GFP_NOWARN);
+	if (!set) {
+		ret = -ENOMEM;
+		goto end;
+	}
+
+	/* For newly allocated set, initialize set->cnt to 0 */
+	if (!tab->sets[hook][type])
+		set->cnt = 0;
+	tab->sets[hook][type] = set;
+
+	/* Concatenate the two sets */
+	memcpy(set->ids + set->cnt, add_set->ids, add_set->cnt * sizeof(set->ids[0]));
+	set->cnt += add_set->cnt;
+
+	if (sort_set)
+		sort(set->ids, set->cnt, sizeof(set->ids[0]), btf_id_cmp_func, NULL);
+
+	return 0;
+end:
+	btf_free_kfunc_set_tab(btf);
+	return ret;
+}
+
+static int btf_populate_kfunc_set(struct btf *btf, enum btf_kfunc_hook hook,
+				  const struct btf_kfunc_id_set *kset)
+{
+	bool sort_set = !btf_is_module(btf);
+	int type, ret;
+
+	for (type = 0; type < ARRAY_SIZE(kset->sets); type++) {
+		if (!kset->sets[type])
+			continue;
+
+		ret = __btf_populate_kfunc_set(btf, hook, type, kset->sets[type], sort_set);
+		if (ret)
+			break;
+	}
+	return ret;
+}
+
+static bool __btf_kfunc_id_set_contains(const struct btf *btf,
+					enum btf_kfunc_hook hook,
+					enum btf_kfunc_type type,
+					u32 kfunc_btf_id)
+{
+	struct btf_id_set *set;
+
+	if (WARN_ON_ONCE(hook >= _BTF_KFUNC_HOOK_MAX || type >= _BTF_KFUNC_TYPE_MAX))
+		return false;
+	if (!btf->kfunc_set_tab)
+		return false;
+	set = btf->kfunc_set_tab->sets[hook][type];
+	if (!set)
+		return false;
+	return btf_id_set_contains(set, kfunc_btf_id);
+}
+
+/* Caution:
+ * Reference to the module (obtained using btf_try_get_module) corresponding to
+ * the struct btf *MUST* be held when calling this function from verifier
+ * context. This is usually true as we stash references in prog's kfunc_btf_tab;
+ * keeping the reference for the duration of the call provides the necessary
+ * protection for looking up a well-formed btf->kfunc_set_tab.
+ */
+bool btf_kfunc_id_set_contains(const struct btf *btf,
+			       enum bpf_prog_type prog_type,
+			       enum btf_kfunc_type type, u32 kfunc_btf_id)
+{
+	enum btf_kfunc_hook hook;
+
+	switch (prog_type) {
+	case BPF_PROG_TYPE_XDP:
+		hook = BTF_KFUNC_HOOK_XDP;
+		break;
+	case BPF_PROG_TYPE_SCHED_CLS:
+		hook = BTF_KFUNC_HOOK_TC;
+		break;
+	case BPF_PROG_TYPE_STRUCT_OPS:
+		hook = BTF_KFUNC_HOOK_STRUCT_OPS;
+		break;
+	default:
+		return false;
+	}
+	return __btf_kfunc_id_set_contains(btf, hook, type, kfunc_btf_id);
+}
+
+/* This function must be invoked only from initcalls/module init functions */
+int register_btf_kfunc_id_set(enum btf_kfunc_hook hook,
+			      const struct btf_kfunc_id_set *kset)
+{
+	struct btf *btf;
+	int ret;
+
+	btf = btf_get_module_btf(kset->owner);
+	if (IS_ERR_OR_NULL(btf))
+		return btf ? PTR_ERR(btf) : -ENOENT;
+
+	ret = btf_populate_kfunc_set(btf, hook, kset);
+	/* Make sure all updates are visible before we go to MODULE_STATE_LIVE,
+	 * pairs with smp_rmb in btf_try_get_module (for success case).
+	 *
+	 * btf_populate_kfunc_set(...)
+	 * smp_wmb()	<-----------.
+	 * mod->state = LIVE	    |		if (mod->state == LIVE)
+	 *			    |		  atomic_inc_nz(mod)
+	 *			    `--------->	  smp_rmb()
+	 *					  btf_kfunc_id_set_contains(...)
+	 */
+	smp_wmb();
+	/* reference is only taken for module BTF */
+	if (btf_is_module(btf))
+		btf_put(btf);
+	return ret;
+}
+EXPORT_SYMBOL_GPL(register_btf_kfunc_id_set);
 
 #ifdef CONFIG_DEBUG_INFO_BTF_MODULES
 
diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index 133599dfe2a2..da9948fe8468 100644
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c
@@ -1777,6 +1777,7 @@  static struct btf *__find_kfunc_desc_btf(struct bpf_verifier_env *env,
 
 		mod = btf_try_get_module(btf);
 		if (!mod) {
+			verbose(env, "failed to get reference for BTF's module\n");
 			btf_put(btf);
 			return ERR_PTR(-ENXIO);
 		}