To effectively operate with memory cgroups in bpf there is a need
to convert css pointers to memcg pointers. A simple container_of
cast which is used in the kernel code can't be used in bpf because
from the verifier's point of view that's a out-of-bounds memory access.

Introduce helper get/put kfuncs which can be used to get
a refcounted memcg pointer from the css pointer:
  - bpf_get_mem_cgroup,
  - bpf_put_mem_cgroup.

bpf_get_mem_cgroup() can take both memcg's css and the corresponding
cgroup's "self" css. It allows it to be used with the existing cgroup
iterator which iterates over cgroup tree, not memcg tree.

Signed-off-by: Roman Gushchin <roman.gushchin@linux.dev>
---
 include/linux/memcontrol.h |   2 +
 mm/Makefile                |   1 +
 mm/bpf_memcontrol.c        | 151 +++++++++++++++++++++++++++++++++++++
 3 files changed, 154 insertions(+)
 create mode 100644 mm/bpf_memcontrol.c

diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
index 87b6688f124a..785a064000cd 100644
--- a/include/linux/memcontrol.h
+++ b/include/linux/memcontrol.h
@@ -932,6 +932,8 @@ static inline void mod_memcg_page_state(struct page *page,
 	rcu_read_unlock();
 }
 
+unsigned long memcg_events(struct mem_cgroup *memcg, int event);
+unsigned long mem_cgroup_usage(struct mem_cgroup *memcg, bool swap);
 unsigned long memcg_page_state(struct mem_cgroup *memcg, int idx);
 unsigned long lruvec_page_state(struct lruvec *lruvec, enum node_stat_item idx);
 unsigned long lruvec_page_state_local(struct lruvec *lruvec,
diff --git a/mm/Makefile b/mm/Makefile
index a714aba03759..c397af904a87 100644
--- a/mm/Makefile
+++ b/mm/Makefile
@@ -107,6 +107,7 @@ obj-$(CONFIG_MEMCG) += swap_cgroup.o
 endif
 ifdef CONFIG_BPF_SYSCALL
 obj-y += bpf_oom.o
+obj-$(CONFIG_MEMCG) += bpf_memcontrol.o
 endif
 obj-$(CONFIG_CGROUP_HUGETLB) += hugetlb_cgroup.o
 obj-$(CONFIG_GUP_TEST) += gup_test.o
diff --git a/mm/bpf_memcontrol.c b/mm/bpf_memcontrol.c
new file mode 100644
index 000000000000..66f2a359af7e
--- /dev/null
+++ b/mm/bpf_memcontrol.c
@@ -0,0 +1,151 @@
+// SPDX-License-Identifier: GPL-2.0-or-later
+/*
+ * Memory Controller-related BPF kfuncs and auxiliary code
+ *
+ * Author: Roman Gushchin <roman.gushchin@linux.dev>
+ */
+
+#include <linux/memcontrol.h>
+#include <linux/bpf.h>
+
+__bpf_kfunc_start_defs();
+
+/**
+ * bpf_get_mem_cgroup - Get a reference to a memory cgroup
+ * @css: pointer to the css structure
+ *
+ * Returns a pointer to a mem_cgroup structure after bumping
+ * the corresponding css's reference counter.
+ *
+ * It's fine to pass a css which belongs to any cgroup controller,
+ * e.g. unified hierarchy's main css.
+ *
+ * Implements KF_ACQUIRE semantics.
+ */
+__bpf_kfunc struct mem_cgroup *
+bpf_get_mem_cgroup(struct cgroup_subsys_state *css)
+{
+	struct mem_cgroup *memcg = NULL;
+	bool rcu_unlock = false;
+
+	if (!root_mem_cgroup)
+		return NULL;
+
+	if (root_mem_cgroup->css.ss != css->ss) {
+		struct cgroup *cgroup = css->cgroup;
+		int ssid = root_mem_cgroup->css.ss->id;
+
+		rcu_read_lock();
+		rcu_unlock = true;
+		css = rcu_dereference_raw(cgroup->subsys[ssid]);
+	}
+
+	if (css && css_tryget(css))
+		memcg = container_of(css, struct mem_cgroup, css);
+
+	if (rcu_unlock)
+		rcu_read_unlock();
+
+	return memcg;
+}
+
+/**
+ * bpf_put_mem_cgroup - Put a reference to a memory cgroup
+ * @memcg: memory cgroup to release
+ *
+ * Releases a previously acquired memcg reference.
+ * Implements KF_RELEASE semantics.
+ */
+__bpf_kfunc void bpf_put_mem_cgroup(struct mem_cgroup *memcg)
+{
+	css_put(&memcg->css);
+}
+
+/**
+ * bpf_mem_cgroup_events - Read memory cgroup's event counter
+ * @memcg: memory cgroup
+ * @event: event idx
+ *
+ * Allows to read memory cgroup event counters.
+ */
+__bpf_kfunc unsigned long bpf_mem_cgroup_events(struct mem_cgroup *memcg, int event)
+{
+
+	if (event < 0 || event >= NR_VM_EVENT_ITEMS)
+		return (unsigned long)-1;
+
+	return memcg_events(memcg, event);
+}
+
+/**
+ * bpf_mem_cgroup_usage - Read memory cgroup's usage
+ * @memcg: memory cgroup
+ *
+ * Returns current memory cgroup size in bytes.
+ */
+__bpf_kfunc unsigned long bpf_mem_cgroup_usage(struct mem_cgroup *memcg)
+{
+	return page_counter_read(&memcg->memory);
+}
+
+/**
+ * bpf_mem_cgroup_events - Read memory cgroup's page state counter
+ * @memcg: memory cgroup
+ * @event: event idx
+ *
+ * Allows to read memory cgroup statistics.
+ */
+__bpf_kfunc unsigned long bpf_mem_cgroup_page_state(struct mem_cgroup *memcg, int idx)
+{
+	if (idx < 0 || idx >= MEMCG_NR_STAT)
+		return (unsigned long)-1;
+
+	return memcg_page_state(memcg, idx);
+}
+
+/**
+ * bpf_mem_cgroup_flush_stats - Flush memory cgroup's statistics
+ * @memcg: memory cgroup
+ *
+ * Propagate memory cgroup's statistics up the cgroup tree.
+ *
+ * Note, that this function uses the rate-limited version of
+ * mem_cgroup_flush_stats() to avoid hurting the system-wide
+ * performance. So bpf_mem_cgroup_flush_stats() guarantees only
+ * that statistics is not stale beyond 2*FLUSH_TIME.
+ */
+__bpf_kfunc void bpf_mem_cgroup_flush_stats(struct mem_cgroup *memcg)
+{
+	mem_cgroup_flush_stats_ratelimited(memcg);
+}
+
+__bpf_kfunc_end_defs();
+
+BTF_KFUNCS_START(bpf_memcontrol_kfuncs)
+BTF_ID_FLAGS(func, bpf_get_mem_cgroup, KF_ACQUIRE | KF_RET_NULL)
+BTF_ID_FLAGS(func, bpf_put_mem_cgroup, KF_RELEASE)
+
+BTF_ID_FLAGS(func, bpf_mem_cgroup_events, KF_TRUSTED_ARGS)
+BTF_ID_FLAGS(func, bpf_mem_cgroup_usage, KF_TRUSTED_ARGS)
+BTF_ID_FLAGS(func, bpf_mem_cgroup_page_state, KF_TRUSTED_ARGS)
+BTF_ID_FLAGS(func, bpf_mem_cgroup_flush_stats, KF_TRUSTED_ARGS)
+
+BTF_KFUNCS_END(bpf_memcontrol_kfuncs)
+
+static const struct btf_kfunc_id_set bpf_memcontrol_kfunc_set = {
+	.owner          = THIS_MODULE,
+	.set            = &bpf_memcontrol_kfuncs,
+};
+
+static int __init bpf_memcontrol_init(void)
+{
+	int err;
+
+	err = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
+					&bpf_memcontrol_kfunc_set);
+	if (err)
+		pr_warn("error while registering bpf memcontrol kfuncs: %d", err);
+
+	return err;
+}
+late_initcall(bpf_memcontrol_init);
-- 
2.50.1

On Mon, 18 Aug 2025 at 19:02, Roman Gushchin <roman.gushchin@linux.dev> wrote:
>
> To effectively operate with memory cgroups in bpf there is a need
> to convert css pointers to memcg pointers. A simple container_of
> cast which is used in the kernel code can't be used in bpf because
> from the verifier's point of view that's a out-of-bounds memory access.
>
> Introduce helper get/put kfuncs which can be used to get
> a refcounted memcg pointer from the css pointer:
>   - bpf_get_mem_cgroup,
>   - bpf_put_mem_cgroup.
>
> bpf_get_mem_cgroup() can take both memcg's css and the corresponding
> cgroup's "self" css. It allows it to be used with the existing cgroup
> iterator which iterates over cgroup tree, not memcg tree.
>
> Signed-off-by: Roman Gushchin <roman.gushchin@linux.dev>
> ---
>  include/linux/memcontrol.h |   2 +
>  mm/Makefile                |   1 +
>  mm/bpf_memcontrol.c        | 151 +++++++++++++++++++++++++++++++++++++
>  3 files changed, 154 insertions(+)
>  create mode 100644 mm/bpf_memcontrol.c
>
> diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
> index 87b6688f124a..785a064000cd 100644
> --- a/include/linux/memcontrol.h
> +++ b/include/linux/memcontrol.h
> @@ -932,6 +932,8 @@ static inline void mod_memcg_page_state(struct page *page,
>         rcu_read_unlock();
>  }
>
> +unsigned long memcg_events(struct mem_cgroup *memcg, int event);
> +unsigned long mem_cgroup_usage(struct mem_cgroup *memcg, bool swap);
>  unsigned long memcg_page_state(struct mem_cgroup *memcg, int idx);
>  unsigned long lruvec_page_state(struct lruvec *lruvec, enum node_stat_item idx);
>  unsigned long lruvec_page_state_local(struct lruvec *lruvec,
> diff --git a/mm/Makefile b/mm/Makefile
> index a714aba03759..c397af904a87 100644
> --- a/mm/Makefile
> +++ b/mm/Makefile
> @@ -107,6 +107,7 @@ obj-$(CONFIG_MEMCG) += swap_cgroup.o
>  endif
>  ifdef CONFIG_BPF_SYSCALL
>  obj-y += bpf_oom.o
> +obj-$(CONFIG_MEMCG) += bpf_memcontrol.o
>  endif
>  obj-$(CONFIG_CGROUP_HUGETLB) += hugetlb_cgroup.o
>  obj-$(CONFIG_GUP_TEST) += gup_test.o
> diff --git a/mm/bpf_memcontrol.c b/mm/bpf_memcontrol.c
> new file mode 100644
> index 000000000000..66f2a359af7e
> --- /dev/null
> +++ b/mm/bpf_memcontrol.c
> @@ -0,0 +1,151 @@
> +// SPDX-License-Identifier: GPL-2.0-or-later
> +/*
> + * Memory Controller-related BPF kfuncs and auxiliary code
> + *
> + * Author: Roman Gushchin <roman.gushchin@linux.dev>
> + */
> +
> +#include <linux/memcontrol.h>
> +#include <linux/bpf.h>
> +
> +__bpf_kfunc_start_defs();
> +
> +/**
> + * bpf_get_mem_cgroup - Get a reference to a memory cgroup
> + * @css: pointer to the css structure
> + *
> + * Returns a pointer to a mem_cgroup structure after bumping
> + * the corresponding css's reference counter.
> + *
> + * It's fine to pass a css which belongs to any cgroup controller,
> + * e.g. unified hierarchy's main css.
> + *
> + * Implements KF_ACQUIRE semantics.
> + */
> +__bpf_kfunc struct mem_cgroup *
> +bpf_get_mem_cgroup(struct cgroup_subsys_state *css)
> +{
> +       struct mem_cgroup *memcg = NULL;
> +       bool rcu_unlock = false;
> +
> +       if (!root_mem_cgroup)
> +               return NULL;
> +
> +       if (root_mem_cgroup->css.ss != css->ss) {
> +               struct cgroup *cgroup = css->cgroup;
> +               int ssid = root_mem_cgroup->css.ss->id;
> +
> +               rcu_read_lock();
> +               rcu_unlock = true;
> +               css = rcu_dereference_raw(cgroup->subsys[ssid]);
> +       }
> +
> +       if (css && css_tryget(css))
> +               memcg = container_of(css, struct mem_cgroup, css);
> +
> +       if (rcu_unlock)
> +               rcu_read_unlock();
> +
> +       return memcg;
> +}
> +
> +/**
> + * bpf_put_mem_cgroup - Put a reference to a memory cgroup
> + * @memcg: memory cgroup to release
> + *
> + * Releases a previously acquired memcg reference.
> + * Implements KF_RELEASE semantics.
> + */
> +__bpf_kfunc void bpf_put_mem_cgroup(struct mem_cgroup *memcg)
> +{
> +       css_put(&memcg->css);
> +}
> +
> +/**
> + * bpf_mem_cgroup_events - Read memory cgroup's event counter
> + * @memcg: memory cgroup
> + * @event: event idx
> + *
> + * Allows to read memory cgroup event counters.
> + */
> +__bpf_kfunc unsigned long bpf_mem_cgroup_events(struct mem_cgroup *memcg, int event)
> +{
> +
> +       if (event < 0 || event >= NR_VM_EVENT_ITEMS)
> +               return (unsigned long)-1;
> +
> +       return memcg_events(memcg, event);
> +}
> +
> +/**
> + * bpf_mem_cgroup_usage - Read memory cgroup's usage
> + * @memcg: memory cgroup
> + *
> + * Returns current memory cgroup size in bytes.
> + */
> +__bpf_kfunc unsigned long bpf_mem_cgroup_usage(struct mem_cgroup *memcg)
> +{
> +       return page_counter_read(&memcg->memory);
> +}
> +
> +/**
> + * bpf_mem_cgroup_events - Read memory cgroup's page state counter
> + * @memcg: memory cgroup
> + * @event: event idx
> + *
> + * Allows to read memory cgroup statistics.
> + */
> +__bpf_kfunc unsigned long bpf_mem_cgroup_page_state(struct mem_cgroup *memcg, int idx)
> +{
> +       if (idx < 0 || idx >= MEMCG_NR_STAT)
> +               return (unsigned long)-1;
> +
> +       return memcg_page_state(memcg, idx);
> +}
> +
> +/**
> + * bpf_mem_cgroup_flush_stats - Flush memory cgroup's statistics
> + * @memcg: memory cgroup
> + *
> + * Propagate memory cgroup's statistics up the cgroup tree.
> + *
> + * Note, that this function uses the rate-limited version of
> + * mem_cgroup_flush_stats() to avoid hurting the system-wide
> + * performance. So bpf_mem_cgroup_flush_stats() guarantees only
> + * that statistics is not stale beyond 2*FLUSH_TIME.
> + */
> +__bpf_kfunc void bpf_mem_cgroup_flush_stats(struct mem_cgroup *memcg)
> +{
> +       mem_cgroup_flush_stats_ratelimited(memcg);
> +}
> +
> +__bpf_kfunc_end_defs();
> +
> +BTF_KFUNCS_START(bpf_memcontrol_kfuncs)
> +BTF_ID_FLAGS(func, bpf_get_mem_cgroup, KF_ACQUIRE | KF_RET_NULL)

I think you could set KF_TRUSTED_ARGS for this as well.

> +BTF_ID_FLAGS(func, bpf_put_mem_cgroup, KF_RELEASE)
> +
> +BTF_ID_FLAGS(func, bpf_mem_cgroup_events, KF_TRUSTED_ARGS)
> +BTF_ID_FLAGS(func, bpf_mem_cgroup_usage, KF_TRUSTED_ARGS)
> +BTF_ID_FLAGS(func, bpf_mem_cgroup_page_state, KF_TRUSTED_ARGS)
> +BTF_ID_FLAGS(func, bpf_mem_cgroup_flush_stats, KF_TRUSTED_ARGS)
> +
> +BTF_KFUNCS_END(bpf_memcontrol_kfuncs)
> +
> +static const struct btf_kfunc_id_set bpf_memcontrol_kfunc_set = {
> +       .owner          = THIS_MODULE,
> +       .set            = &bpf_memcontrol_kfuncs,
> +};
> +
> +static int __init bpf_memcontrol_init(void)
> +{
> +       int err;
> +
> +       err = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
> +                                       &bpf_memcontrol_kfunc_set);
> +       if (err)
> +               pr_warn("error while registering bpf memcontrol kfuncs: %d", err);
> +
> +       return err;
> +}
> +late_initcall(bpf_memcontrol_init);
> --
> 2.50.1
>

Kumar Kartikeya Dwivedi <memxor@gmail.com> writes:

> On Mon, 18 Aug 2025 at 19:02, Roman Gushchin <roman.gushchin@linux.dev> wrote:
>>
>> To effectively operate with memory cgroups in bpf there is a need
>> to convert css pointers to memcg pointers. A simple container_of
>> cast which is used in the kernel code can't be used in bpf because
>> from the verifier's point of view that's a out-of-bounds memory access.
>>
>> Introduce helper get/put kfuncs which can be used to get
>> a refcounted memcg pointer from the css pointer:
>>   - bpf_get_mem_cgroup,
>>   - bpf_put_mem_cgroup.
>>
>> bpf_get_mem_cgroup() can take both memcg's css and the corresponding
>> cgroup's "self" css. It allows it to be used with the existing cgroup
>> iterator which iterates over cgroup tree, not memcg tree.
>>
>> Signed-off-by: Roman Gushchin <roman.gushchin@linux.dev>
>> ---
>>  include/linux/memcontrol.h |   2 +
>>  mm/Makefile                |   1 +
>>  mm/bpf_memcontrol.c        | 151 +++++++++++++++++++++++++++++++++++++
>>  3 files changed, 154 insertions(+)
>>  create mode 100644 mm/bpf_memcontrol.c
>>
>> diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
>> index 87b6688f124a..785a064000cd 100644
>> --- a/include/linux/memcontrol.h
>> +++ b/include/linux/memcontrol.h
>> @@ -932,6 +932,8 @@ static inline void mod_memcg_page_state(struct page *page,
>>         rcu_read_unlock();
>>  }
>>
>> +unsigned long memcg_events(struct mem_cgroup *memcg, int event);
>> +unsigned long mem_cgroup_usage(struct mem_cgroup *memcg, bool swap);
>>  unsigned long memcg_page_state(struct mem_cgroup *memcg, int idx);
>>  unsigned long lruvec_page_state(struct lruvec *lruvec, enum node_stat_item idx);
>>  unsigned long lruvec_page_state_local(struct lruvec *lruvec,
>> diff --git a/mm/Makefile b/mm/Makefile
>> index a714aba03759..c397af904a87 100644
>> --- a/mm/Makefile
>> +++ b/mm/Makefile
>> @@ -107,6 +107,7 @@ obj-$(CONFIG_MEMCG) += swap_cgroup.o
>>  endif
>>  ifdef CONFIG_BPF_SYSCALL
>>  obj-y += bpf_oom.o
>> +obj-$(CONFIG_MEMCG) += bpf_memcontrol.o
>>  endif
>>  obj-$(CONFIG_CGROUP_HUGETLB) += hugetlb_cgroup.o
>>  obj-$(CONFIG_GUP_TEST) += gup_test.o
>> diff --git a/mm/bpf_memcontrol.c b/mm/bpf_memcontrol.c
>> new file mode 100644
>> index 000000000000..66f2a359af7e
>> --- /dev/null
>> +++ b/mm/bpf_memcontrol.c
>> @@ -0,0 +1,151 @@
>> +// SPDX-License-Identifier: GPL-2.0-or-later
>> +/*
>> + * Memory Controller-related BPF kfuncs and auxiliary code
>> + *
>> + * Author: Roman Gushchin <roman.gushchin@linux.dev>
>> + */
>> +
>> +#include <linux/memcontrol.h>
>> +#include <linux/bpf.h>
>> +
>> +__bpf_kfunc_start_defs();
>> +
>> +/**
>> + * bpf_get_mem_cgroup - Get a reference to a memory cgroup
>> + * @css: pointer to the css structure
>> + *
>> + * Returns a pointer to a mem_cgroup structure after bumping
>> + * the corresponding css's reference counter.
>> + *
>> + * It's fine to pass a css which belongs to any cgroup controller,
>> + * e.g. unified hierarchy's main css.
>> + *
>> + * Implements KF_ACQUIRE semantics.
>> + */
>> +__bpf_kfunc struct mem_cgroup *
>> +bpf_get_mem_cgroup(struct cgroup_subsys_state *css)
>> +{
>> +       struct mem_cgroup *memcg = NULL;
>> +       bool rcu_unlock = false;
>> +
>> +       if (!root_mem_cgroup)
>> +               return NULL;
>> +
>> +       if (root_mem_cgroup->css.ss != css->ss) {
>> +               struct cgroup *cgroup = css->cgroup;
>> +               int ssid = root_mem_cgroup->css.ss->id;
>> +
>> +               rcu_read_lock();
>> +               rcu_unlock = true;
>> +               css = rcu_dereference_raw(cgroup->subsys[ssid]);
>> +       }
>> +
>> +       if (css && css_tryget(css))
>> +               memcg = container_of(css, struct mem_cgroup, css);
>> +
>> +       if (rcu_unlock)
>> +               rcu_read_unlock();
>> +
>> +       return memcg;
>> +}
>> +
>> +/**
>> + * bpf_put_mem_cgroup - Put a reference to a memory cgroup
>> + * @memcg: memory cgroup to release
>> + *
>> + * Releases a previously acquired memcg reference.
>> + * Implements KF_RELEASE semantics.
>> + */
>> +__bpf_kfunc void bpf_put_mem_cgroup(struct mem_cgroup *memcg)
>> +{
>> +       css_put(&memcg->css);
>> +}
>> +
>> +/**
>> + * bpf_mem_cgroup_events - Read memory cgroup's event counter
>> + * @memcg: memory cgroup
>> + * @event: event idx
>> + *
>> + * Allows to read memory cgroup event counters.
>> + */
>> +__bpf_kfunc unsigned long bpf_mem_cgroup_events(struct mem_cgroup *memcg, int event)
>> +{
>> +
>> +       if (event < 0 || event >= NR_VM_EVENT_ITEMS)
>> +               return (unsigned long)-1;
>> +
>> +       return memcg_events(memcg, event);
>> +}
>> +
>> +/**
>> + * bpf_mem_cgroup_usage - Read memory cgroup's usage
>> + * @memcg: memory cgroup
>> + *
>> + * Returns current memory cgroup size in bytes.
>> + */
>> +__bpf_kfunc unsigned long bpf_mem_cgroup_usage(struct mem_cgroup *memcg)
>> +{
>> +       return page_counter_read(&memcg->memory);
>> +}
>> +
>> +/**
>> + * bpf_mem_cgroup_events - Read memory cgroup's page state counter
>> + * @memcg: memory cgroup
>> + * @event: event idx
>> + *
>> + * Allows to read memory cgroup statistics.
>> + */
>> +__bpf_kfunc unsigned long bpf_mem_cgroup_page_state(struct mem_cgroup *memcg, int idx)
>> +{
>> +       if (idx < 0 || idx >= MEMCG_NR_STAT)
>> +               return (unsigned long)-1;
>> +
>> +       return memcg_page_state(memcg, idx);
>> +}
>> +
>> +/**
>> + * bpf_mem_cgroup_flush_stats - Flush memory cgroup's statistics
>> + * @memcg: memory cgroup
>> + *
>> + * Propagate memory cgroup's statistics up the cgroup tree.
>> + *
>> + * Note, that this function uses the rate-limited version of
>> + * mem_cgroup_flush_stats() to avoid hurting the system-wide
>> + * performance. So bpf_mem_cgroup_flush_stats() guarantees only
>> + * that statistics is not stale beyond 2*FLUSH_TIME.
>> + */
>> +__bpf_kfunc void bpf_mem_cgroup_flush_stats(struct mem_cgroup *memcg)
>> +{
>> +       mem_cgroup_flush_stats_ratelimited(memcg);
>> +}
>> +
>> +__bpf_kfunc_end_defs();
>> +
>> +BTF_KFUNCS_START(bpf_memcontrol_kfuncs)
>> +BTF_ID_FLAGS(func, bpf_get_mem_cgroup, KF_ACQUIRE | KF_RET_NULL)
>
> I think you could set KF_TRUSTED_ARGS for this as well.

Not really. The intended use case is to iterate over the cgroup tree,
which gives non-trusted css pointers:
	bpf_for_each(css, css_pos, &root_memcg->css, BPF_CGROUP_ITER_DESCENDANTS_POST) {
		memcg = bpf_get_mem_cgroup(css_pos);
	}

Thanks

On Thu, 21 Aug 2025 at 00:43, Roman Gushchin <roman.gushchin@linux.dev> wrote:
>
> Kumar Kartikeya Dwivedi <memxor@gmail.com> writes:
>
> > On Mon, 18 Aug 2025 at 19:02, Roman Gushchin <roman.gushchin@linux.dev> wrote:
> >>
> >> To effectively operate with memory cgroups in bpf there is a need
> >> to convert css pointers to memcg pointers. A simple container_of
> >> cast which is used in the kernel code can't be used in bpf because
> >> from the verifier's point of view that's a out-of-bounds memory access.
> >>
> >> Introduce helper get/put kfuncs which can be used to get
> >> a refcounted memcg pointer from the css pointer:
> >>   - bpf_get_mem_cgroup,
> >>   - bpf_put_mem_cgroup.
> >>
> >> bpf_get_mem_cgroup() can take both memcg's css and the corresponding
> >> cgroup's "self" css. It allows it to be used with the existing cgroup
> >> iterator which iterates over cgroup tree, not memcg tree.
> >>
> >> Signed-off-by: Roman Gushchin <roman.gushchin@linux.dev>
> >> ---
> >>  include/linux/memcontrol.h |   2 +
> >>  mm/Makefile                |   1 +
> >>  mm/bpf_memcontrol.c        | 151 +++++++++++++++++++++++++++++++++++++
> >>  3 files changed, 154 insertions(+)
> >>  create mode 100644 mm/bpf_memcontrol.c
> >>
> >> diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
> >> index 87b6688f124a..785a064000cd 100644
> >> --- a/include/linux/memcontrol.h
> >> +++ b/include/linux/memcontrol.h
> >> @@ -932,6 +932,8 @@ static inline void mod_memcg_page_state(struct page *page,
> >>         rcu_read_unlock();
> >>  }
> >>
> >> +unsigned long memcg_events(struct mem_cgroup *memcg, int event);
> >> +unsigned long mem_cgroup_usage(struct mem_cgroup *memcg, bool swap);
> >>  unsigned long memcg_page_state(struct mem_cgroup *memcg, int idx);
> >>  unsigned long lruvec_page_state(struct lruvec *lruvec, enum node_stat_item idx);
> >>  unsigned long lruvec_page_state_local(struct lruvec *lruvec,
> >> diff --git a/mm/Makefile b/mm/Makefile
> >> index a714aba03759..c397af904a87 100644
> >> --- a/mm/Makefile
> >> +++ b/mm/Makefile
> >> @@ -107,6 +107,7 @@ obj-$(CONFIG_MEMCG) += swap_cgroup.o
> >>  endif
> >>  ifdef CONFIG_BPF_SYSCALL
> >>  obj-y += bpf_oom.o
> >> +obj-$(CONFIG_MEMCG) += bpf_memcontrol.o
> >>  endif
> >>  obj-$(CONFIG_CGROUP_HUGETLB) += hugetlb_cgroup.o
> >>  obj-$(CONFIG_GUP_TEST) += gup_test.o
> >> diff --git a/mm/bpf_memcontrol.c b/mm/bpf_memcontrol.c
> >> new file mode 100644
> >> index 000000000000..66f2a359af7e
> >> --- /dev/null
> >> +++ b/mm/bpf_memcontrol.c
> >> @@ -0,0 +1,151 @@
> >> +// SPDX-License-Identifier: GPL-2.0-or-later
> >> +/*
> >> + * Memory Controller-related BPF kfuncs and auxiliary code
> >> + *
> >> + * Author: Roman Gushchin <roman.gushchin@linux.dev>
> >> + */
> >> +
> >> +#include <linux/memcontrol.h>
> >> +#include <linux/bpf.h>
> >> +
> >> +__bpf_kfunc_start_defs();
> >> +
> >> +/**
> >> + * bpf_get_mem_cgroup - Get a reference to a memory cgroup
> >> + * @css: pointer to the css structure
> >> + *
> >> + * Returns a pointer to a mem_cgroup structure after bumping
> >> + * the corresponding css's reference counter.
> >> + *
> >> + * It's fine to pass a css which belongs to any cgroup controller,
> >> + * e.g. unified hierarchy's main css.
> >> + *
> >> + * Implements KF_ACQUIRE semantics.
> >> + */
> >> +__bpf_kfunc struct mem_cgroup *
> >> +bpf_get_mem_cgroup(struct cgroup_subsys_state *css)
> >> +{
> >> +       struct mem_cgroup *memcg = NULL;
> >> +       bool rcu_unlock = false;
> >> +
> >> +       if (!root_mem_cgroup)
> >> +               return NULL;
> >> +
> >> +       if (root_mem_cgroup->css.ss != css->ss) {
> >> +               struct cgroup *cgroup = css->cgroup;
> >> +               int ssid = root_mem_cgroup->css.ss->id;
> >> +
> >> +               rcu_read_lock();
> >> +               rcu_unlock = true;
> >> +               css = rcu_dereference_raw(cgroup->subsys[ssid]);
> >> +       }
> >> +
> >> +       if (css && css_tryget(css))
> >> +               memcg = container_of(css, struct mem_cgroup, css);
> >> +
> >> +       if (rcu_unlock)
> >> +               rcu_read_unlock();
> >> +
> >> +       return memcg;
> >> +}
> >> +
> >> +/**
> >> + * bpf_put_mem_cgroup - Put a reference to a memory cgroup
> >> + * @memcg: memory cgroup to release
> >> + *
> >> + * Releases a previously acquired memcg reference.
> >> + * Implements KF_RELEASE semantics.
> >> + */
> >> +__bpf_kfunc void bpf_put_mem_cgroup(struct mem_cgroup *memcg)
> >> +{
> >> +       css_put(&memcg->css);
> >> +}
> >> +
> >> +/**
> >> + * bpf_mem_cgroup_events - Read memory cgroup's event counter
> >> + * @memcg: memory cgroup
> >> + * @event: event idx
> >> + *
> >> + * Allows to read memory cgroup event counters.
> >> + */
> >> +__bpf_kfunc unsigned long bpf_mem_cgroup_events(struct mem_cgroup *memcg, int event)
> >> +{
> >> +
> >> +       if (event < 0 || event >= NR_VM_EVENT_ITEMS)
> >> +               return (unsigned long)-1;
> >> +
> >> +       return memcg_events(memcg, event);
> >> +}
> >> +
> >> +/**
> >> + * bpf_mem_cgroup_usage - Read memory cgroup's usage
> >> + * @memcg: memory cgroup
> >> + *
> >> + * Returns current memory cgroup size in bytes.
> >> + */
> >> +__bpf_kfunc unsigned long bpf_mem_cgroup_usage(struct mem_cgroup *memcg)
> >> +{
> >> +       return page_counter_read(&memcg->memory);
> >> +}
> >> +
> >> +/**
> >> + * bpf_mem_cgroup_events - Read memory cgroup's page state counter
> >> + * @memcg: memory cgroup
> >> + * @event: event idx
> >> + *
> >> + * Allows to read memory cgroup statistics.
> >> + */
> >> +__bpf_kfunc unsigned long bpf_mem_cgroup_page_state(struct mem_cgroup *memcg, int idx)
> >> +{
> >> +       if (idx < 0 || idx >= MEMCG_NR_STAT)
> >> +               return (unsigned long)-1;
> >> +
> >> +       return memcg_page_state(memcg, idx);
> >> +}
> >> +
> >> +/**
> >> + * bpf_mem_cgroup_flush_stats - Flush memory cgroup's statistics
> >> + * @memcg: memory cgroup
> >> + *
> >> + * Propagate memory cgroup's statistics up the cgroup tree.
> >> + *
> >> + * Note, that this function uses the rate-limited version of
> >> + * mem_cgroup_flush_stats() to avoid hurting the system-wide
> >> + * performance. So bpf_mem_cgroup_flush_stats() guarantees only
> >> + * that statistics is not stale beyond 2*FLUSH_TIME.
> >> + */
> >> +__bpf_kfunc void bpf_mem_cgroup_flush_stats(struct mem_cgroup *memcg)
> >> +{
> >> +       mem_cgroup_flush_stats_ratelimited(memcg);
> >> +}
> >> +
> >> +__bpf_kfunc_end_defs();
> >> +
> >> +BTF_KFUNCS_START(bpf_memcontrol_kfuncs)
> >> +BTF_ID_FLAGS(func, bpf_get_mem_cgroup, KF_ACQUIRE | KF_RET_NULL)
> >
> > I think you could set KF_TRUSTED_ARGS for this as well.
>
> Not really. The intended use case is to iterate over the cgroup tree,
> which gives non-trusted css pointers:
>         bpf_for_each(css, css_pos, &root_memcg->css, BPF_CGROUP_ITER_DESCENDANTS_POST) {
>                 memcg = bpf_get_mem_cgroup(css_pos);
>         }

Then I assume they're at least RCU protected? You could relax it from
trusted to KF_RCU (since I see css_tryget internally).
Otherwise the default behavior is unconstrained (any ptr matching that
type obtained from random walks --- which is something to fix, but
until then we have to actively mark for taking safe arguments).

>
> Thanks