4 files changed, 72 insertions, 9 deletions

diff --git a/include/linux/bpf-cgroup.h b/include/linux/bpf-cgroup.h
index cb3c6b3b89c8..9f100fc422c3 100644
--- a/include/linux/bpf-cgroup.h
+++ b/include/linux/bpf-cgroup.h
@@ -6,6 +6,7 @@
 #include <linux/errno.h>
 #include <linux/jump_label.h>
 #include <linux/percpu.h>
+#include <linux/percpu-refcount.h>
 #include <linux/rbtree.h>
 #include <uapi/linux/bpf.h>
 
@@ -72,10 +73,16 @@ struct cgroup_bpf {
 
         /* temp storage for effective prog array used by prog_attach/detach */
         struct bpf_prog_array __rcu *inactive;
+
+        /* reference counter used to detach bpf programs after cgroup removal */
+        struct percpu_ref refcnt;
+
+        /* cgroup_bpf is released using a work queue */
+        struct work_struct release_work;
 };
 
-void cgroup_bpf_put(struct cgroup *cgrp);
 int cgroup_bpf_inherit(struct cgroup *cgrp);
+void cgroup_bpf_offline(struct cgroup *cgrp);
 
 int __cgroup_bpf_attach(struct cgroup *cgrp, struct bpf_prog *prog,
                         enum bpf_attach_type type, u32 flags);
@@ -283,8 +290,8 @@ int cgroup_bpf_prog_query(const union bpf_attr *attr,
 
 struct bpf_prog;
 struct cgroup_bpf {};
-static inline void cgroup_bpf_put(struct cgroup *cgrp) {}
 static inline int cgroup_bpf_inherit(struct cgroup *cgrp) { return 0; }
+static inline void cgroup_bpf_offline(struct cgroup *cgrp) {}
 
 static inline int cgroup_bpf_prog_attach(const union bpf_attr *attr,
                                          enum bpf_prog_type ptype,
diff --git a/include/linux/cgroup.h b/include/linux/cgroup.h
index c0077adeea83..49e8facf7c4a 100644
--- a/include/linux/cgroup.h
+++ b/include/linux/cgroup.h
@@ -924,4 +924,22 @@ static inline bool cgroup_task_frozen(struct task_struct *task)
 
 #endif /* !CONFIG_CGROUPS */
 
+#ifdef CONFIG_CGROUP_BPF
+static inline void cgroup_bpf_get(struct cgroup *cgrp)
+{
+        percpu_ref_get(&cgrp->bpf.refcnt);
+}
+
+static inline void cgroup_bpf_put(struct cgroup *cgrp)
+{
+        percpu_ref_put(&cgrp->bpf.refcnt);
+}
+
+#else /* CONFIG_CGROUP_BPF */
+
+static inline void cgroup_bpf_get(struct cgroup *cgrp) {}
+static inline void cgroup_bpf_put(struct cgroup *cgrp) {}
+
+#endif /* CONFIG_CGROUP_BPF */
+
 #endif /* _LINUX_CGROUP_H */
diff --git a/kernel/bpf/cgroup.c b/kernel/bpf/cgroup.c
index fcde0f7b2585..d995edbe816d 100644
--- a/kernel/bpf/cgroup.c
+++ b/kernel/bpf/cgroup.c
@@ -22,12 +22,21 @@
 DEFINE_STATIC_KEY_FALSE(cgroup_bpf_enabled_key);
 EXPORT_SYMBOL(cgroup_bpf_enabled_key);
 
+void cgroup_bpf_offline(struct cgroup *cgrp)
+{
+        cgroup_get(cgrp);
+        percpu_ref_kill(&cgrp->bpf.refcnt);
+}
+
 /**
- * cgroup_bpf_put() - put references of all bpf programs
- * @cgrp: the cgroup to modify
+ * cgroup_bpf_release() - put references of all bpf programs and
+ *                        release all cgroup bpf data
+ * @work: work structure embedded into the cgroup to modify
  */
-void cgroup_bpf_put(struct cgroup *cgrp)
+static void cgroup_bpf_release(struct work_struct *work)
 {
+        struct cgroup *cgrp = container_of(work, struct cgroup,
+                                           bpf.release_work);
         enum bpf_cgroup_storage_type stype;
         unsigned int type;
 
@@ -47,6 +56,22 @@ void cgroup_bpf_put(struct cgroup *cgrp)
                 }
                 bpf_prog_array_free(cgrp->bpf.effective[type]);
         }
+
+        percpu_ref_exit(&cgrp->bpf.refcnt);
+        cgroup_put(cgrp);
+}
+
+/**
+ * cgroup_bpf_release_fn() - callback used to schedule releasing
+ *                           of bpf cgroup data
+ * @ref: percpu ref counter structure
+ */
+static void cgroup_bpf_release_fn(struct percpu_ref *ref)
+{
+        struct cgroup *cgrp = container_of(ref, struct cgroup, bpf.refcnt);
+
+        INIT_WORK(&cgrp->bpf.release_work, cgroup_bpf_release);
+        queue_work(system_wq, &cgrp->bpf.release_work);
 }
 
 /* count number of elements in the list.
@@ -167,7 +192,12 @@ int cgroup_bpf_inherit(struct cgroup *cgrp)
  */
 #define NR ARRAY_SIZE(cgrp->bpf.effective)
         struct bpf_prog_array __rcu *arrays[NR] = {};
-        int i;
+        int ret, i;
+
+        ret = percpu_ref_init(&cgrp->bpf.refcnt, cgroup_bpf_release_fn, 0,
+                              GFP_KERNEL);
+        if (ret)
+                return ret;
 
         for (i = 0; i < NR; i++)
                 INIT_LIST_HEAD(&cgrp->bpf.progs[i]);
@@ -183,6 +213,9 @@ int cgroup_bpf_inherit(struct cgroup *cgrp)
 cleanup:
         for (i = 0; i < NR; i++)
                 bpf_prog_array_free(arrays[i]);
+
+        percpu_ref_exit(&cgrp->bpf.refcnt);
+
         return -ENOMEM;
 }
 
diff --git a/kernel/cgroup/cgroup.c b/kernel/cgroup/cgroup.c
index 217cec4e22c6..ef9cfbfc82a9 100644
--- a/kernel/cgroup/cgroup.c
+++ b/kernel/cgroup/cgroup.c
@@ -4955,8 +4955,6 @@ static void css_release_work_fn(struct work_struct *work)
                 if (cgrp->kn)
                         RCU_INIT_POINTER(*(void __rcu __force **)&cgrp->kn->priv,
                                          NULL);
-
-                cgroup_bpf_put(cgrp);
         }
 
         mutex_unlock(&cgroup_mutex);
@@ -5482,6 +5480,8 @@ static int cgroup_destroy_locked(struct cgroup *cgrp)
 
         cgroup1_check_for_release(parent);
 
+        cgroup_bpf_offline(cgrp);
+
         /* put the base reference */
         percpu_ref_kill(&cgrp->self.refcnt);
 
@@ -6221,6 +6221,7 @@ void cgroup_sk_alloc(struct sock_cgroup_data *skcd)
                  * Don't use cgroup_get_live().
                  */
                 cgroup_get(sock_cgroup_ptr(skcd));
+                cgroup_bpf_get(sock_cgroup_ptr(skcd));
                 return;
         }
 
@@ -6232,6 +6233,7 @@ void cgroup_sk_alloc(struct sock_cgroup_data *skcd)
                 cset = task_css_set(current);
                 if (likely(cgroup_tryget(cset->dfl_cgrp))) {
                         skcd->val = (unsigned long)cset->dfl_cgrp;
+                        cgroup_bpf_get(cset->dfl_cgrp);
                         break;
                 }
                 cpu_relax();
@@ -6242,7 +6244,10 @@ void cgroup_sk_alloc(struct sock_cgroup_data *skcd)
 
 void cgroup_sk_free(struct sock_cgroup_data *skcd)
 {
-        cgroup_put(sock_cgroup_ptr(skcd));
+        struct cgroup *cgrp = sock_cgroup_ptr(skcd);
+
+        cgroup_bpf_put(cgrp);
+        cgroup_put(cgrp);
 }
 
 #endif  /* CONFIG_SOCK_CGROUP_DATA */