u32 *btf_kfunc_id_set_contains(const struct btf *btf,
                               enum bpf_prog_type prog_type,
                               u32 kfunc_btf_id);
+u32 *btf_kfunc_is_modify_return(const struct btf *btf, u32 kfunc_btf_id);
 int register_btf_kfunc_id_set(enum bpf_prog_type prog_type,
                              const struct btf_kfunc_id_set *s);
+int register_btf_fmodret_id_set(const struct btf_kfunc_id_set *kset);
 s32 btf_find_dtor_kfunc(struct btf *btf, u32 btf_id);
 int register_btf_id_dtor_kfuncs(const struct btf_id_dtor_kfunc *dtors, u32 add_cnt,
                                struct module *owner);
 
        BTF_KFUNC_HOOK_STRUCT_OPS,
        BTF_KFUNC_HOOK_TRACING,
        BTF_KFUNC_HOOK_SYSCALL,
+       BTF_KFUNC_HOOK_FMODRET,
        BTF_KFUNC_HOOK_MAX,
 };
 
        return __btf_kfunc_id_set_contains(btf, hook, kfunc_btf_id);
 }
 
-/* This function must be invoked only from initcalls/module init functions */
-int register_btf_kfunc_id_set(enum bpf_prog_type prog_type,
-                             const struct btf_kfunc_id_set *kset)
+u32 *btf_kfunc_is_modify_return(const struct btf *btf, u32 kfunc_btf_id)
+{
+       return __btf_kfunc_id_set_contains(btf, BTF_KFUNC_HOOK_FMODRET, kfunc_btf_id);
+}
+
+static int __register_btf_kfunc_id_set(enum btf_kfunc_hook hook,
+                                      const struct btf_kfunc_id_set *kset)
 {
-       enum btf_kfunc_hook hook;
        struct btf *btf;
        int ret;
 
        if (IS_ERR(btf))
                return PTR_ERR(btf);
 
-       hook = bpf_prog_type_to_kfunc_hook(prog_type);
        ret = btf_populate_kfunc_set(btf, hook, kset->set);
        btf_put(btf);
        return ret;
 }
+
+/* This function must be invoked only from initcalls/module init functions */
+int register_btf_kfunc_id_set(enum bpf_prog_type prog_type,
+                             const struct btf_kfunc_id_set *kset)
+{
+       enum btf_kfunc_hook hook;
+
+       hook = bpf_prog_type_to_kfunc_hook(prog_type);
+       return __register_btf_kfunc_id_set(hook, kset);
+}
 EXPORT_SYMBOL_GPL(register_btf_kfunc_id_set);
 
+/* This function must be invoked only from initcalls/module init functions */
+int register_btf_fmodret_id_set(const struct btf_kfunc_id_set *kset)
+{
+       return __register_btf_kfunc_id_set(BTF_KFUNC_HOOK_FMODRET, kset);
+}
+EXPORT_SYMBOL_GPL(register_btf_fmodret_id_set);
+
 s32 btf_find_dtor_kfunc(struct btf *btf, u32 btf_id)
 {
        struct btf_id_dtor_kfunc_tab *tab = btf->dtor_kfunc_tab;
 
                        ret = -EINVAL;
                        switch (prog->type) {
                        case BPF_PROG_TYPE_TRACING:
-                               /* fentry/fexit/fmod_ret progs can be sleepable only if they are
+
+                               /* fentry/fexit/fmod_ret progs can be sleepable if they are
                                 * attached to ALLOW_ERROR_INJECTION and are not in denylist.
                                 */
                                if (!check_non_sleepable_error_inject(btf_id) &&
                                    within_error_injection_list(addr))
                                        ret = 0;
+                               /* fentry/fexit/fmod_ret progs can also be sleepable if they are
+                                * in the fmodret id set with the KF_SLEEPABLE flag.
+                                */
+                               else {
+                                       u32 *flags = btf_kfunc_is_modify_return(btf, btf_id);
+
+                                       if (flags && (*flags & KF_SLEEPABLE))
+                                               ret = 0;
+                               }
                                break;
                        case BPF_PROG_TYPE_LSM:
                                /* LSM progs check that they are attached to bpf_lsm_*() funcs.
                                bpf_log(log, "can't modify return codes of BPF programs\n");
                                return -EINVAL;
                        }
-                       ret = check_attach_modify_return(addr, tname);
+                       ret = -EINVAL;
+                       if (btf_kfunc_is_modify_return(btf, btf_id) ||
+                           !check_attach_modify_return(addr, tname))
+                               ret = 0;
                        if (ret) {
                                bpf_log(log, "%s() is not modifiable\n", tname);
                                return ret;
 
        return a + 1;
 }
 EXPORT_SYMBOL_GPL(bpf_fentry_test1);
-ALLOW_ERROR_INJECTION(bpf_fentry_test1, ERRNO);
 
 int noinline bpf_fentry_test2(int a, u64 b)
 {
 
 __diag_pop();
 
-ALLOW_ERROR_INJECTION(bpf_modify_return_test, ERRNO);
+BTF_SET8_START(bpf_test_modify_return_ids)
+BTF_ID_FLAGS(func, bpf_modify_return_test)
+BTF_ID_FLAGS(func, bpf_fentry_test1, KF_SLEEPABLE)
+BTF_SET8_END(bpf_test_modify_return_ids)
+
+static const struct btf_kfunc_id_set bpf_test_modify_return_set = {
+       .owner = THIS_MODULE,
+       .set   = &bpf_test_modify_return_ids,
+};
 
 BTF_SET8_START(test_sk_check_kfunc_ids)
 BTF_ID_FLAGS(func, bpf_kfunc_call_test1)
        };
        int ret;
 
-       ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_SCHED_CLS, &bpf_prog_test_kfunc_set);
+       ret = register_btf_fmodret_id_set(&bpf_test_modify_return_set);
+       ret = ret ?: register_btf_kfunc_id_set(BPF_PROG_TYPE_SCHED_CLS, &bpf_prog_test_kfunc_set);
        ret = ret ?: register_btf_kfunc_id_set(BPF_PROG_TYPE_TRACING, &bpf_prog_test_kfunc_set);
        ret = ret ?: register_btf_kfunc_id_set(BPF_PROG_TYPE_SYSCALL, &bpf_prog_test_kfunc_set);
        return ret ?: register_btf_id_dtor_kfuncs(bpf_prog_test_dtor_kfunc,