pytrap: fix OOB write in iplist.c

pytrap/src/iplist.c

@@ -14,13 +14,19 @@
 static ipps_network_list_t *
 load_networks(PyDictObject *d)
 {
-    uint32_t i = 0;
-    uint32_t struct_count = 50; // Starting v4_count of structs to alloc
+    uint64_t i = 0;
 
     // ************* LOAD NETWORKS ********************** //
 
+    // Get size of passed dictionary
+    Py_ssize_t dict_size = PyDict_Size((PyObject *) d);
+    if (dict_size < 0) {
+        /* PyDict_Size already set an exception. */
+        return NULL;
+    }
+
     // Alloc memory for networks structs, if malloc fails return NULL
-    ipps_network_t *networks = malloc(struct_count * sizeof(ipps_network_t));
+    ipps_network_t *networks = malloc((size_t) dict_size * sizeof(ipps_network_t));
     if (networks == NULL) {
         PyErr_SetString(PyExc_MemoryError, "Failed allocating memory for IP prefix search structures.");
         return NULL;
@@ -42,15 +48,6 @@ load_networks(PyDictObject *d)
             pytrap_unirecipaddrrange *r = (pytrap_unirecipaddrrange *) key;
             network->mask = r->mask;
             memcpy(&network->addr, &r->start->ip, sizeof(ip_addr_t));
-            // If limit is reached alloc new memory
-            if (i >= struct_count) {
-                struct_count += 10;
-                // If realloc fails return NULL
-                if ((networks = realloc(networks, struct_count * sizeof(ipps_network_t))) == NULL) {
-                    PyErr_SetString(PyExc_MemoryError, "Failed in reallocating network structure.");
-                    return NULL;
-                }
-            }
         } else {
             PyErr_SetString(PyExc_TypeError, "Unsupported type.");
             return NULL;

pytrap/test/ipblacklist_unittest.py

@@ -125,3 +125,51 @@ def runTest(self):
         res = iplist.find(pytrap.UnirecIPAddr("1::1"))
         self.assertEqual(res, "abc")
 
+class DictionaryInitTest(unittest.TestCase):
+    def runTest(self):
+        import itertools
+        import ipaddress
+        import pytrap
+
+        # An empty dictionary must be rejected.
+        with self.assertRaises(ValueError):
+            pytrap.UnirecIPList({})
+
+        # Generate ip addresses
+        ipv4_count = 1000
+        ipv6_count = 1000
+        ipv4_networks = itertools.islice(
+            ipaddress.IPv4Network("10.0.0.0/8").subnets(new_prefix=24),
+            ipv4_count,
+        )
+        ipv6_networks = itertools.islice(
+            ipaddress.IPv6Network("2001:db8::/32").subnets(new_prefix=48),
+            ipv6_count,
+        )
+
+        networks = list(ipv4_networks) + list(ipv6_networks)
+
+        ranges = {
+            pytrap.UnirecIPAddrRange(str(network)): f"network-{index}"
+            for index, network in enumerate(networks)
+        }
+
+        self.assertEqual(len(ranges), (ipv4_count+ipv6_count))
+
+        ip_list = pytrap.UnirecIPList(ranges)
+
+        for index, network in enumerate(networks):
+            with self.subTest(network=str(network)):
+                address = pytrap.UnirecIPAddr(
+                    str(network.network_address + 42)
+                )
+
+                self.assertIn(address, ip_list)
+                self.assertEqual(
+                    ip_list.find(address),
+                    f"network-{index}",
+                )
+
+        missing = pytrap.UnirecIPAddr("192.0.2.1")
+        self.assertNotIn(missing, ip_list)
+        self.assertIsNone(ip_list.find(missing))