1 files changed, 56 insertions, 19 deletions

diff --git a/net/sched/cls_cgroup.c b/net/sched/cls_cgroup.c
index e4877ca6727c..78ef2c5e130b 100644
--- a/net/sched/cls_cgroup.c
+++ b/net/sched/cls_cgroup.c
@@ -10,20 +10,37 @@
  */
 
 #include <linux/module.h>
+#include <linux/slab.h>
 #include <linux/types.h>
 #include <linux/string.h>
 #include <linux/errno.h>
 #include <linux/skbuff.h>
 #include <linux/cgroup.h>
+#include <linux/rcupdate.h>
 #include <net/rtnetlink.h>
 #include <net/pkt_cls.h>
+#include <net/sock.h>
+#include <net/cls_cgroup.h>
 
-struct cgroup_cls_state
-{
-        struct cgroup_subsys_state css;
-        u32 classid;
+static struct cgroup_subsys_state *cgrp_create(struct cgroup_subsys *ss,
+                                               struct cgroup *cgrp);
+static void cgrp_destroy(struct cgroup_subsys *ss, struct cgroup *cgrp);
+static int cgrp_populate(struct cgroup_subsys *ss, struct cgroup *cgrp);
+
+struct cgroup_subsys net_cls_subsys = {
+        .name           = "net_cls",
+        .create         = cgrp_create,
+        .destroy        = cgrp_destroy,
+        .populate       = cgrp_populate,
+#ifdef CONFIG_NET_CLS_CGROUP
+        .subsys_id      = net_cls_subsys_id,
+#else
+#define net_cls_subsys_id net_cls_subsys.subsys_id
+#endif
+        .module         = THIS_MODULE,
 };
 
+
 static inline struct cgroup_cls_state *cgrp_cls_state(struct cgroup *cgrp)
 {
         return container_of(cgroup_subsys_state(cgrp, net_cls_subsys_id),
@@ -79,14 +96,6 @@ static int cgrp_populate(struct cgroup_subsys *ss, struct cgroup *cgrp)
         return cgroup_add_files(cgrp, ss, ss_files, ARRAY_SIZE(ss_files));
 }
 
-struct cgroup_subsys net_cls_subsys = {
-        .name           = "net_cls",
-        .create         = cgrp_create,
-        .destroy        = cgrp_destroy,
-        .populate       = cgrp_populate,
-        .subsys_id      = net_cls_subsys_id,
-};
-
 struct cls_cgroup_head
 {
         u32                     handle;
@@ -100,6 +109,10 @@ static int cls_cgroup_classify(struct sk_buff *skb, struct tcf_proto *tp,
         struct cls_cgroup_head *head = tp->root;
         u32 classid;
 
+        rcu_read_lock();
+        classid = task_cls_state(current)->classid;
+        rcu_read_unlock();
+
         /*
          * Due to the nature of the classifier it is required to ignore all
          * packets originating from softirq context as accessing `current'
@@ -110,12 +123,12 @@ static int cls_cgroup_classify(struct sk_buff *skb, struct tcf_proto *tp,
          * calls by looking at the number of nested bh disable calls because
          * softirqs always disables bh.
          */
-        if (softirq_count() != SOFTIRQ_OFFSET)
-                return -1;
-
-        rcu_read_lock();
-        classid = task_cls_state(current)->classid;
-        rcu_read_unlock();
+        if (softirq_count() != SOFTIRQ_OFFSET) {
+                /* If there is an sk_classid we'll use that. */
+                if (!skb->sk)
+                        return -1;
+                classid = skb->sk->sk_classid;
+        }
 
         if (!classid)
                 return -1;
@@ -277,12 +290,36 @@ static struct tcf_proto_ops cls_cgroup_ops __read_mostly = {
 
 static int __init init_cgroup_cls(void)
 {
-        return register_tcf_proto_ops(&cls_cgroup_ops);
+        int ret;
+
+        ret = cgroup_load_subsys(&net_cls_subsys);
+        if (ret)
+                goto out;
+
+#ifndef CONFIG_NET_CLS_CGROUP
+        /* We can't use rcu_assign_pointer because this is an int. */
+        smp_wmb();
+        net_cls_subsys_id = net_cls_subsys.subsys_id;
+#endif
+
+        ret = register_tcf_proto_ops(&cls_cgroup_ops);
+        if (ret)
+                cgroup_unload_subsys(&net_cls_subsys);
+
+out:
+        return ret;
 }
 
 static void __exit exit_cgroup_cls(void)
 {
         unregister_tcf_proto_ops(&cls_cgroup_ops);
+
+#ifndef CONFIG_NET_CLS_CGROUP
+        net_cls_subsys_id = -1;
+        synchronize_rcu();
+#endif
+
+        cgroup_unload_subsys(&net_cls_subsys);
 }
 
 module_init(init_cgroup_cls);