2 files changed, 45 insertions, 3 deletions

diff --git a/net/Kconfig b/net/Kconfig
index 4c06c7c513a4..da1282359303 100644
--- a/net/Kconfig
+++ b/net/Kconfig
@@ -296,7 +296,7 @@ config NET_TCPPROBE
         module will be called tcp_probe.
 
 config NET_DROP_MONITOR
-        boolean "Network packet drop alerting service"
+        tristate "Network packet drop alerting service"
         depends on INET && EXPERIMENTAL && TRACEPOINTS
         ---help---
         This feature provides an alerting service to userspace in the
diff --git a/net/core/drop_monitor.c b/net/core/drop_monitor.c
index eca00a96bcf3..3252e7e0a005 100644
--- a/net/core/drop_monitor.c
+++ b/net/core/drop_monitor.c
@@ -24,6 +24,7 @@
 #include <linux/timer.h>
 #include <linux/bitops.h>
 #include <linux/slab.h>
+#include <linux/module.h>
 #include <net/genetlink.h>
 #include <net/netevent.h>
 
@@ -263,9 +264,15 @@ static int set_all_monitor_traces(int state)
 
         switch (state) {
         case TRACE_ON:
+                if (!try_module_get(THIS_MODULE)) {
+                        rc = -ENODEV;
+                        break;
+                }
+
                 rc |= register_trace_kfree_skb(trace_kfree_skb_hit, NULL);
                 rc |= register_trace_napi_poll(trace_napi_poll_hit, NULL);
                 break;
+
         case TRACE_OFF:
                 rc |= unregister_trace_kfree_skb(trace_kfree_skb_hit, NULL);
                 rc |= unregister_trace_napi_poll(trace_napi_poll_hit, NULL);
@@ -281,6 +288,9 @@ static int set_all_monitor_traces(int state)
                                 kfree_rcu(new_stat, rcu);
                         }
                 }
+
+                module_put(THIS_MODULE);
+
                 break;
         default:
                 rc = 1;
@@ -406,7 +416,7 @@ static int __init init_net_drop_monitor(void)
 
         rc = 0;
 
-        for_each_present_cpu(cpu) {
+        for_each_possible_cpu(cpu) {
                 data = &per_cpu(dm_cpu_data, cpu);
                 data->cpu = cpu;
                 INIT_WORK(&data->dm_alert_work, send_dm_alert);
@@ -425,4 +435,36 @@ out:
         return rc;
 }
 
-late_initcall(init_net_drop_monitor);
+static void exit_net_drop_monitor(void)
+{
+        struct per_cpu_dm_data *data;
+        int cpu;
+
+        BUG_ON(unregister_netdevice_notifier(&dropmon_net_notifier));
+
+        /*
+         * Because of the module_get/put we do in the trace state change path
+         * we are guarnateed not to have any current users when we get here
+         * all we need to do is make sure that we don't have any running timers
+         * or pending schedule calls
+         */
+
+        for_each_possible_cpu(cpu) {
+                data = &per_cpu(dm_cpu_data, cpu);
+                del_timer_sync(&data->send_timer);
+                cancel_work_sync(&data->dm_alert_work);
+                /*
+                 * At this point, we should have exclusive access
+                 * to this struct and can free the skb inside it
+                 */
+                kfree_skb(data->skb);
+        }
+
+        BUG_ON(genl_unregister_family(&net_drop_monitor_family));
+}
+
+module_init(init_net_drop_monitor);
+module_exit(exit_net_drop_monitor);
+
+MODULE_LICENSE("GPL v2");
+MODULE_AUTHOR("Neil Horman <nhorman@tuxdriver.com>");