increased readabilty of memory management, fixed memory leak on unload

[ipt_ACCOUNT] / linux / net / ipv4 / netfilter / ipt_ACCOUNT.c

diff --git a/linux/net/ipv4/netfilter/ipt_ACCOUNT.c b/linux/net/ipv4/netfilter/ipt_ACCOUNT.c
index c380d97..40e207f 100644 (file)
--- a/linux/net/ipv4/netfilter/ipt_ACCOUNT.c
+++ b/linux/net/ipv4/netfilter/ipt_ACCOUNT.c
@@ -22,8 +22,13 @@ struct in_device;
 
 struct ipt_account_table *ipt_account_tables = NULL;
 struct ipt_account_handle *ipt_account_handles = NULL;
+void *ipt_account_tmpbuf = NULL;
 
+// Spinlock used for manipulating the current accounting tables/data
 static spinlock_t ipt_account_lock = SPIN_LOCK_UNLOCKED;
+// Spinlock used for manipulating userspace handles/snapshot data
+static spinlock_t ipt_account_userspace_lock = SPIN_LOCK_UNLOCKED;
+
 
 /* Recursive free of all data structures */
 void ipt_account_data_free(void *data, unsigned char depth)
@@ -150,7 +155,7 @@ int ipt_account_table_insert(char *name, unsigned int ip, unsigned int netmask)
             printk("ACCOUNT: calculated netsize: %u -> ipt_account_table depth %u\n", netsize, ipt_account_tables[i].depth);
                         
             ipt_account_tables[i].refcount++;
-            if (!(ipt_account_tables[i].data = (void *)get_zeroed_page(GFP_KERNEL)))
+            if ((ipt_account_tables[i].data = (void *)get_zeroed_page(GFP_KERNEL)) == NULL)
             {
                 printk("ACCOUNT: out of memory for data of table: %s\n", name);
                 memset(&ipt_account_tables[i], 0, sizeof(struct ipt_account_table));
@@ -323,7 +328,7 @@ void ipt_account_depth1_insert(struct ipt_account_mask_16 *mask_16, unsigned int
         DEBUGP("ACCOUNT: Calculated SRC 16 bit network slot: %d\n", slot);
         
         // Do we need to create a new mask_24 bucket?
-        if (!mask_16->mask_24[slot] && !(mask_16->mask_24[slot] = (void *)get_zeroed_page(GFP_KERNEL)))
+        if (!mask_16->mask_24[slot] && (mask_16->mask_24[slot] = (void *)get_zeroed_page(GFP_KERNEL)) == NULL)
         {
             printk("ACCOUNT: Can't process packet because out of memory!\n");
             return;
@@ -340,7 +345,7 @@ void ipt_account_depth1_insert(struct ipt_account_mask_16 *mask_16, unsigned int
         DEBUGP("ACCOUNT: Calculated DST 16 bit network slot: %d\n", slot);
         
         // Do we need to create a new mask_24 bucket?
-        if (!mask_16->mask_24[slot] && !(mask_16->mask_24[slot] = (void *)get_zeroed_page(GFP_KERNEL)))
+        if (!mask_16->mask_24[slot] && (mask_16->mask_24[slot] = (void *)get_zeroed_page(GFP_KERNEL)) == NULL)
         {
             printk("ACCOUT: Can't process packet because out of memory!\n");
             return;
@@ -361,7 +366,7 @@ void ipt_account_depth2_insert(struct ipt_account_mask_8 *mask_8, unsigned int n
         DEBUGP("ACCOUNT: Calculated SRC 24 bit network slot: %d\n", slot);
                 
         // Do we need to create a new mask_24 bucket?
-        if (!mask_8->mask_16[slot] && !(mask_8->mask_16[slot] = (void *)get_zeroed_page(GFP_KERNEL)))
+        if (!mask_8->mask_16[slot] && (mask_8->mask_16[slot] = (void *)get_zeroed_page(GFP_KERNEL)) == NULL)
         {
             printk("ACCOUNT: Can't process packet because out of memory!\n");
             return;
@@ -378,7 +383,7 @@ void ipt_account_depth2_insert(struct ipt_account_mask_8 *mask_8, unsigned int n
         DEBUGP("ACCOUNT: Calculated DST 24 bit network slot: %d\n", slot);
         
         // Do we need to create a new mask_24 bucket?
-        if (!mask_8->mask_16[slot] && !(mask_8->mask_16[slot] = (void *)get_zeroed_page(GFP_KERNEL)))
+        if (!mask_8->mask_16[slot] && (mask_8->mask_16[slot] = (void *)get_zeroed_page(GFP_KERNEL)) == NULL)
         {
             printk("ACCOUNT: Can't process packet because out of memory!\n");
             return;
@@ -528,7 +533,7 @@ int ipt_account_handle_prepare_read(char *tablename, unsigned int *count)
     ipt_account_handles[handle].itemcount = ipt_account_tables[table_nr].itemcount;
     
     // allocate "root" table
-    if (!(ipt_account_handles[handle].data = (void*)get_zeroed_page(GFP_KERNEL)))
+    if ((ipt_account_handles[handle].data = (void*)get_zeroed_page(GFP_KERNEL)) == NULL)
     {
         printk("ACCOUNT: out of memory for root table in ipt_account_handle_prepare_read()\n");
         memset (&ipt_account_handles[handle], 0, sizeof(struct ipt_account_handle));
@@ -549,7 +554,7 @@ int ipt_account_handle_prepare_read(char *tablename, unsigned int *count)
         {
             if (src_16->mask_24[b])
             {
-                if (!(network_16->mask_24[b] = (void*)get_zeroed_page(GFP_KERNEL)))
+                if ((network_16->mask_24[b] = (void*)get_zeroed_page(GFP_KERNEL)) == NULL)
                 {
                     printk("ACCOUNT: out of memory during copy of 16 bit network in ipt_account_handle_prepare_read()\n");
                     ipt_account_data_free(ipt_account_handles[handle].data, depth);
@@ -569,7 +574,7 @@ int ipt_account_handle_prepare_read(char *tablename, unsigned int *count)
         {
             if (src_8->mask_16[a])
             {
-                if (!(network_8->mask_16[a] = (void*)get_zeroed_page(GFP_KERNEL)))
+                if ((network_8->mask_16[a] = (void*)get_zeroed_page(GFP_KERNEL)) == NULL)
                 {
                     printk("ACCOUNT: out of memory during copy of 24 bit network in ipt_account_handle_prepare_read()\n");
                     ipt_account_data_free(ipt_account_handles[handle].data, depth);
@@ -587,7 +592,7 @@ int ipt_account_handle_prepare_read(char *tablename, unsigned int *count)
                 {
                     if (src_16->mask_24[b])
                     {
-                        if (!(network_16->mask_24[b] = (void*)get_zeroed_page(GFP_KERNEL)))
+                        if ((network_16->mask_24[b] = (void*)get_zeroed_page(GFP_KERNEL)) == NULL)
                         {
                             printk("ACCOUNT: out of memory during copy of 16 bit network in ipt_account_handle_prepare_read()\n");
                             ipt_account_data_free(ipt_account_handles[handle].data, depth);
@@ -650,7 +655,7 @@ int ipt_account_handle_get_data(unsigned int handle, void *buffer)
 {
     struct ipt_account_handle_ip handle_ip;    
     unsigned int handle_ip_size = sizeof (struct ipt_account_handle_ip);
-    unsigned int i;
+    unsigned int i, tmpbuf_pos=0;
 
     if (handle >= ACCOUNT_MAX_HANDLES)
     {
@@ -664,7 +669,6 @@ int ipt_account_handle_get_data(unsigned int handle, void *buffer)
         return -1;
     }
     
-    // Check if at least packet src/dst matches ip/netmask
     unsigned int net_ip = ipt_account_handles[handle].ip;
     unsigned int depth = ipt_account_handles[handle].depth;
     
@@ -682,10 +686,20 @@ int ipt_account_handle_get_data(unsigned int handle, void *buffer)
                 handle_ip.dst_packets = network->ip[i].dst_packets;
                 handle_ip.dst_bytes = network->ip[i].dst_bytes;
 
-                copy_to_user(buffer, &handle_ip, handle_ip_size);
-                buffer += handle_ip_size;
+                // Temporary buffer full? Flush to userspace
+                if (tmpbuf_pos+handle_ip_size >= PAGE_SIZE)
+                {
+                    copy_to_user(buffer, ipt_account_tmpbuf, tmpbuf_pos);
+                    tmpbuf_pos = 0;
+                }
+                memcpy(ipt_account_tmpbuf+tmpbuf_pos, &handle_ip, handle_ip_size);
+                tmpbuf_pos += handle_ip_size;
             }
         }
+
+        // Flush remaining data to userspace
+        if (tmpbuf_pos)
+            copy_to_user(buffer, ipt_account_tmpbuf, tmpbuf_pos);
         
         return 0;
     }
@@ -710,13 +724,23 @@ int ipt_account_handle_get_data(unsigned int handle, void *buffer)
                         handle_ip.dst_packets = network->ip[i].dst_packets;
                         handle_ip.dst_bytes = network->ip[i].dst_bytes;
         
-                        copy_to_user(buffer, &handle_ip, handle_ip_size);
-                        buffer += handle_ip_size;
+                        // Temporary buffer full? Flush to userspace
+                        if (tmpbuf_pos+handle_ip_size >= PAGE_SIZE)
+                        {
+                            copy_to_user(buffer, ipt_account_tmpbuf, tmpbuf_pos);
+                            tmpbuf_pos = 0;
+                        }
+                        memcpy(ipt_account_tmpbuf+tmpbuf_pos, &handle_ip, handle_ip_size);
+                        tmpbuf_pos += handle_ip_size;
                     }
                 }
             }
         }
         
+        // Flush remaining data to userspace
+        if (tmpbuf_pos)
+            copy_to_user(buffer, ipt_account_tmpbuf, tmpbuf_pos);
+        
         return 0;
     }
     
@@ -745,8 +769,14 @@ int ipt_account_handle_get_data(unsigned int handle, void *buffer)
                                 handle_ip.dst_packets = network->ip[i].dst_packets;
                                 handle_ip.dst_bytes = network->ip[i].dst_bytes;
                 
-                                copy_to_user(buffer, &handle_ip, handle_ip_size);
-                                buffer += handle_ip_size;
+                                // Temporary buffer full? Flush to userspace
+                                if (tmpbuf_pos+handle_ip_size >= PAGE_SIZE)
+                                {
+                                    copy_to_user(buffer, ipt_account_tmpbuf, tmpbuf_pos);
+                                    tmpbuf_pos = 0;
+                                }
+                                memcpy(ipt_account_tmpbuf+tmpbuf_pos, &handle_ip, handle_ip_size);
+                                tmpbuf_pos += handle_ip_size;
                             }
                         }
                     }
@@ -754,9 +784,13 @@ int ipt_account_handle_get_data(unsigned int handle, void *buffer)
             }
         }
         
+        // Flush remaining data to userspace
+        if (tmpbuf_pos)
+            copy_to_user(buffer, ipt_account_tmpbuf, tmpbuf_pos);
+        
         return 0;
     }
-        
+    
     return -1;
 }
 
@@ -783,9 +817,9 @@ static int ipt_account_set_ctl(struct sock *sk, int cmd, void *user, unsigned in
                 break;
             }
             
-            spin_lock_bh(&ipt_account_lock);
+            spin_lock_bh(&ipt_account_userspace_lock);
             ret = ipt_account_handle_free(handle.handle_nr);
-            spin_unlock_bh(&ipt_account_lock);
+            spin_unlock_bh(&ipt_account_userspace_lock);
             break;
         default:
             printk("ACCOUNT: ipt_account_set_ctl: unknown request %i\n", cmd);
@@ -820,10 +854,12 @@ static int ipt_account_get_ctl(struct sock *sk, int cmd, void *user, int *len)
             }
             
             spin_lock_bh(&ipt_account_lock);
+            spin_lock_bh(&ipt_account_userspace_lock);
             if (cmd == IPT_SO_GET_ACCOUNT_PREPARE_READ_FLUSH)
                 handle.handle_nr = ipt_account_handle_prepare_read_flush(handle.name, &handle.itemcount);
             else
                 handle.handle_nr = ipt_account_handle_prepare_read(handle.name, &handle.itemcount);
+            spin_unlock_bh(&ipt_account_userspace_lock);
             spin_unlock_bh(&ipt_account_lock);
                 
             if (handle.handle_nr == -1)
@@ -866,9 +902,9 @@ static int ipt_account_get_ctl(struct sock *sk, int cmd, void *user, int *len)
                 break;
             }   
             
-            spin_lock_bh(&ipt_account_lock);
+            spin_lock_bh(&ipt_account_userspace_lock);
             ret = ipt_account_handle_get_data(handle.handle_nr, user);
-            spin_unlock_bh(&ipt_account_lock);
+            spin_unlock_bh(&ipt_account_userspace_lock);
             if (ret)
             {
                 printk("ACCOUNT: ipt_account_get_ctl: ipt_account_handle_get_data failed for handle %u\n", handle.handle_nr);
@@ -901,14 +937,14 @@ static int __init init(void)
         return -EINVAL;
     }
 
-    if (!(ipt_account_tables = kmalloc(ACCOUNT_MAX_TABLES*sizeof(struct ipt_account_table), GFP_KERNEL)))
+    if ((ipt_account_tables = kmalloc(ACCOUNT_MAX_TABLES*sizeof(struct ipt_account_table), GFP_KERNEL)) == NULL)
     {
         printk("ACCOUNT: Out of memory allocating account_tables structure");
         return -EINVAL;
     }
     memset(ipt_account_tables, 0, ACCOUNT_MAX_TABLES*sizeof(struct ipt_account_table));
                             
-    if (!(ipt_account_handles = kmalloc(ACCOUNT_MAX_HANDLES*sizeof(struct ipt_account_handle), GFP_KERNEL)))
+    if ((ipt_account_handles = kmalloc(ACCOUNT_MAX_HANDLES*sizeof(struct ipt_account_handle), GFP_KERNEL)) == NULL)
     {
         printk("ACCOUNT: Out of memory allocating account_handles structure");
         kfree (ipt_account_tables);
@@ -917,15 +953,28 @@ static int __init init(void)
     }
     memset(ipt_account_handles, 0, ACCOUNT_MAX_HANDLES*sizeof(struct ipt_account_handle));
 
+    // Allocate one page as temporary storage
+    if ((ipt_account_tmpbuf = (void*)__get_free_page(GFP_KERNEL)) == NULL)
+    {
+        printk("ACCOUNT: Out of memory for temporary buffer page\n");
+        kfree(ipt_account_tables);
+        kfree(ipt_account_handles);
+        ipt_account_tables = NULL;
+        ipt_account_handles = NULL;
+        return -EINVAL;
+    }
+    
     /* Register setsockopt */
     if (nf_register_sockopt(&ipt_account_sockopts) < 0)
     {
         printk("ACCOUNT: Can't register sockopts. Aborting\n");
         
-        kfree (ipt_account_tables);
-        ipt_account_tables = NULL;
+        kfree(ipt_account_tables);
         kfree(ipt_account_handles);
+        kfree(ipt_account_tmpbuf);
+        ipt_account_tables = NULL;
         ipt_account_handles = NULL;
+        ipt_account_tmpbuf = NULL;
         
         return -EINVAL;
     }
@@ -943,10 +992,12 @@ static void __exit fini(void)
     nf_unregister_sockopt(&ipt_account_sockopts);
     
     kfree(ipt_account_tables);
-    ipt_account_tables = NULL;
-        
     kfree(ipt_account_handles);
+    kfree(ipt_account_tmpbuf);
+    
+    ipt_account_tables = NULL;
     ipt_account_handles = NULL;
+    ipt_account_tmpbuf = NULL;
 }
 
 module_init(init);