diff --git a/.gitignore b/.gitignore index 69480f4..34676a2 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,6 @@ update.sh MANIFEST build dist +.vscode/ +venv/ +*.swp diff --git a/README.md b/README.md index 559cfe0..96d19c7 100644 --- a/README.md +++ b/README.md @@ -8,11 +8,16 @@ to recommend it over related modules (including py-radix and SubnetTree): 2. it works in Python 3, and 3. there are a few nicer library features for manipulating the structure. -Copyright (c) 2012-2017 Joel Sommers. All rights reserved. +Copyright (c) 2012-2025 Joel Sommers. All rights reserved. Pytricia is released under terms of the GNU Lesser General Public License, version 3.0 and greater. +## Support further development of Pytricia + +I originally wrote this code with funding from the US National Science Foundation. Development since 2016 has been on an "as I have time and motivation" basis. If you or your organization gets benefit from this software and you'd like to see further development, [please consider donating](https://www.buymeacoffee.com/joelsommers). + + # Building Building pytricia is done in the standard pythonic way: @@ -200,6 +205,18 @@ Although it is possible to store IPv4 and IPv6 subnets in the same trie, this is IPv4 address `32.0.0.1` matches `2000::/8` prefix due to the first octet being the same in both. In order to avoid this, separate tries should be used for IPv4 and IPv6 prefixes. Alternatively, [IPv4 addresses can be mapped to IPv6 addresses](https://en.wikipedia.org/wiki/IPv6#IPv4-mapped_IPv6_addresses). +``PyTricia`` objects can be be pickled, but you must first ``freeze()`` to reconfigure them to a more efficient representation suitable for serialization. Note that while in this more compact representation you can not modify the object. To restore the ability to modify you can use ``thaw()``. + + >>> import pytricia + >>> import pickle + >>> pyt = pytricia.PyTricia() + >>> pyt.freeze() + >>> s = pickle.dumps(pyt) + >>> pyt = pickle.loads(s) + >>> pyt.thaw() + + + # Performance For API usage, the usual Python advice applies: using indexing is the fastest method for insertion, lookup, and removal. See the ``apiperf.py`` script in the repo for some comparative numbers. For Python 3, using ``ipaddress``-module objects is the slowest. There's a price to pay for the convenience, unfortunately. diff --git a/apiperf.py b/apiperf.py index c99712b..558c348 100644 --- a/apiperf.py +++ b/apiperf.py @@ -65,8 +65,8 @@ def main(): for fn in expts: t = Timer("{}()".format(fn), "from __main__ import {}".format(fn)) v = t.timeit(number=iterations) - v = v / float(iterations) - print ("{:.08f}s: average execution time for {}".format(v, fn)) + v_iter = v / float(iterations) + print ("{:.08f}s: (average execution time for {} total_sec={:.1f})".format(v_iter, fn, v)) if __name__ == '__main__': main() diff --git a/patricia.c b/patricia.c index 1278675..b27ab99 100644 --- a/patricia.c +++ b/patricia.c @@ -274,8 +274,8 @@ prefix_toa (prefix_t * prefix) return (prefix_toa2 (prefix, (char *) NULL)); } -prefix_t * -New_Prefix2 (int family, void *dest, int bitlen, prefix_t *prefix) +int +New_Prefix (int family, void *dest, int bitlen, prefix_t *prefix) { int dynamic_allocated = 0; int default_bitlen = 32; @@ -283,11 +283,11 @@ New_Prefix2 (int family, void *dest, int bitlen, prefix_t *prefix) #ifdef HAVE_IPV6 if (family == AF_INET6) { default_bitlen = 128; - if (prefix == NULL) { + if (prefix == NULL) { prefix = calloc(1, sizeof (prefix6_t)); - dynamic_allocated++; - } - memcpy (&prefix->add.sin6, dest, 16); + dynamic_allocated++; + } + memcpy (&prefix->add.sin6, dest, 16); } else #endif /* HAVE_IPV6 */ @@ -306,7 +306,7 @@ New_Prefix2 (int family, void *dest, int bitlen, prefix_t *prefix) memcpy (&prefix->add.sin, dest, 4); } else { - return (NULL); + return 0; } prefix->bitlen = (bitlen >= 0)? bitlen: default_bitlen; @@ -316,112 +316,7 @@ New_Prefix2 (int family, void *dest, int bitlen, prefix_t *prefix) prefix->ref_count++; } /* fprintf(stderr, "[C %s, %d]\n", prefix_toa (prefix), prefix->ref_count); */ - return (prefix); -} - -prefix_t * -New_Prefix (int family, void *dest, int bitlen) -{ - return (New_Prefix2 (family, dest, bitlen, NULL)); -} - -/* ascii2prefix - */ -prefix_t * -ascii2prefix (int family, char *string) -{ - u_long bitlen = 0, maxbitlen = 0; - char *cp = NULL; - struct in_addr sin; -#ifdef HAVE_IPV6 - struct in6_addr sin6; -#endif /* HAVE_IPV6 */ - int result = 0; - char save[MAXLINE]; - - if (string == NULL) - return (NULL); - - /* easy way to handle both families */ - if (family == 0) { - family = AF_INET; -#ifdef HAVE_IPV6 - if (strchr (string, ':')) family = AF_INET6; -#endif /* HAVE_IPV6 */ - } - - if (family == AF_INET) { - maxbitlen = 32; - } -#ifdef HAVE_IPV6 - else if (family == AF_INET6) { - maxbitlen = 128; - } -#endif /* HAVE_IPV6 */ - - if ((cp = strchr (string, '/')) != NULL) { - bitlen = atol (cp + 1); - /* *cp = '\0'; */ - /* copy the string to save. Avoid destroying the string */ - assert (cp - string < MAXLINE); - memcpy (save, string, cp - string); - save[cp - string] = '\0'; - string = save; - if (bitlen > maxbitlen) - bitlen = maxbitlen; - } - - if (family == AF_INET) { - if ((result = my_inet_pton (AF_INET, string, &sin)) <= 0) - return (NULL); - return (New_Prefix (AF_INET, &sin, (int)bitlen)); - } - -#ifdef HAVE_IPV6 - else if (family == AF_INET6) { -// Get rid of this with next IPv6 upgrade -#if defined(NT) && !defined(HAVE_INET_NTOP) - inet6_addr(string, &sin6); - return (New_Prefix (AF_INET6, &sin6, (int)bitlen)); -#else - if ((result = local_inet_pton (AF_INET6, string, &sin6)) <= 0) - return (NULL); -#endif /* NT */ - return (New_Prefix (AF_INET6, &sin6, (int)bitlen)); - } -#endif /* HAVE_IPV6 */ - else - return (NULL); -} - -prefix_t * -Ref_Prefix (prefix_t * prefix) -{ - if (prefix == NULL) - return (NULL); - if (prefix->ref_count == 0) { - /* make a copy in case of a static prefix */ - return (New_Prefix2 (prefix->family, &prefix->add, prefix->bitlen, NULL)); - } - prefix->ref_count++; -/* fprintf(stderr, "[A %s, %d]\n", prefix_toa (prefix), prefix->ref_count); */ - return (prefix); -} - -void -Deref_Prefix (prefix_t * prefix) -{ - if (prefix == NULL) - return; - /* for secure programming, raise an assert. no static prefix can call this */ - assert (prefix->ref_count > 0); - - prefix->ref_count--; - assert (prefix->ref_count >= 0); - if (prefix->ref_count <= 0) { - Delete (prefix); - return; - } + return 1; } /* } */ @@ -440,6 +335,7 @@ New_Patricia (int maxbits) patricia->maxbits = maxbits; patricia->head = NULL; patricia->num_active_node = 0; + patricia->frozen = 0; assert (maxbits <= PATRICIA_MAXBITS); /* XXX */ num_active_patricia++; return (patricia); @@ -456,7 +352,6 @@ Clear_Patricia (patricia_tree_t *patricia, void_fn1_t func) { assert (patricia); if (patricia->head) { - patricia_node_t *Xstack[PATRICIA_MAXBITS+1]; patricia_node_t **Xsp = Xstack; patricia_node_t *Xrn = patricia->head; @@ -465,16 +360,17 @@ Clear_Patricia (patricia_tree_t *patricia, void_fn1_t func) patricia_node_t *l = Xrn->l; patricia_node_t *r = Xrn->r; - if (Xrn->prefix) { - Deref_Prefix (Xrn->prefix); - if (Xrn->data && func) - func (Xrn->data); + if (Xrn->data) { + if (func) + func (Xrn->data); } else { - assert (Xrn->data == NULL); + assert (Xrn->data == NULL); } - Delete (Xrn); - patricia->num_active_node--; + if(!patricia->frozen) { + Delete (Xrn); + } + patricia->num_active_node--; if (l) { if (r) { @@ -491,7 +387,10 @@ Clear_Patricia (patricia_tree_t *patricia, void_fn1_t func) } } assert (patricia->num_active_node == 0); - /* Delete (patricia); */ + if(patricia->frozen && patricia->head) { + Delete (patricia->head); + } + /* Delete (patricia); */ } @@ -515,7 +414,7 @@ patricia_process (patricia_tree_t *patricia, void_fn2_t func) assert (func); PATRICIA_WALK (patricia->head, node) { - func (node->prefix, node->data); + func (&node->prefix, node->data); } PATRICIA_WALK_END; } @@ -542,9 +441,9 @@ patricia_search_exact (patricia_tree_t *patricia, prefix_t *prefix) if (BIT_TEST (addr[node->bit >> 3], 0x80 >> (node->bit & 0x07))) { #ifdef PATRICIA_DEBUG - if (node->prefix) + if (node->data) fprintf (stderr, "patricia_search_exact: take right %s/%d\n", - prefix_toa (node->prefix), node->prefix->bitlen); + prefix_toa (&node->prefix), node->prefix.bitlen); else fprintf (stderr, "patricia_search_exact: take right at %d\n", node->bit); @@ -553,9 +452,9 @@ patricia_search_exact (patricia_tree_t *patricia, prefix_t *prefix) } else { #ifdef PATRICIA_DEBUG - if (node->prefix) - fprintf (stderr, "patricia_search_exact: take left %s/%d\n", - prefix_toa (node->prefix), node->prefix->bitlen); + if (node->data) + fprintf (stderr, "patricia_search_exact: take left %s/%d\n", + prefix_toa (&node->prefix), node->prefix.bitlen); else fprintf (stderr, "patricia_search_exact: take left at %d\n", node->bit); @@ -568,23 +467,23 @@ patricia_search_exact (patricia_tree_t *patricia, prefix_t *prefix) } #ifdef PATRICIA_DEBUG - if (node->prefix) + if (&node->prefix) fprintf (stderr, "patricia_search_exact: stop at %s/%d\n", - prefix_toa (node->prefix), node->prefix->bitlen); + prefix_toa (&node->prefix), node->prefix.bitlen); else fprintf (stderr, "patricia_search_exact: stop at %d\n", node->bit); #endif /* PATRICIA_DEBUG */ - if (node->bit > bitlen || node->prefix == NULL) - return (NULL); + if (node->bit > bitlen || node->data == NULL) + return (NULL); assert (node->bit == bitlen); - assert (node->bit == node->prefix->bitlen); - if (comp_with_mask (prefix_tochar (node->prefix), prefix_tochar (prefix), + assert (node->bit == node->prefix.bitlen); + if (comp_with_mask (prefix_tochar (&node->prefix), prefix_tochar (prefix), bitlen)) { #ifdef PATRICIA_DEBUG fprintf (stderr, "patricia_search_exact: found %s/%d\n", - prefix_toa (node->prefix), node->prefix->bitlen); + prefix_toa (&node->prefix), node->prefix.bitlen); #endif /* PATRICIA_DEBUG */ - return (node); + return (node); } return (NULL); } @@ -613,72 +512,72 @@ patricia_search_best2 (patricia_tree_t *patricia, prefix_t *prefix, int inclusiv while (node->bit < bitlen) { - if (node->prefix) { + if (node->data) { #ifdef PATRICIA_DEBUG fprintf (stderr, "patricia_search_best: push %s/%d\n", - prefix_toa (node->prefix), node->prefix->bitlen); + prefix_toa (&node->prefix), node->prefix.bitlen); #endif /* PATRICIA_DEBUG */ - stack[cnt++] = node; - } + stack[cnt++] = node; + } - if (BIT_TEST (addr[node->bit >> 3], 0x80 >> (node->bit & 0x07))) { + if (BIT_TEST (addr[node->bit >> 3], 0x80 >> (node->bit & 0x07))) { #ifdef PATRICIA_DEBUG - if (node->prefix) + if (&node->prefix) fprintf (stderr, "patricia_search_best: take right %s/%d\n", - prefix_toa (node->prefix), node->prefix->bitlen); - else + prefix_toa (&node->prefix), node->prefix.bitlen); + else fprintf (stderr, "patricia_search_best: take right at %d\n", - node->bit); + node->bit); #endif /* PATRICIA_DEBUG */ - node = node->r; - } - else { + node = node->r; + } + else { #ifdef PATRICIA_DEBUG - if (node->prefix) + if (&node->prefix) fprintf (stderr, "patricia_search_best: take left %s/%d\n", - prefix_toa (node->prefix), node->prefix->bitlen); - else + prefix_toa (&node->prefix), node->prefix.bitlen); + else fprintf (stderr, "patricia_search_best: take left at %d\n", - node->bit); + node->bit); #endif /* PATRICIA_DEBUG */ - node = node->l; - } + node = node->l; + } - if (node == NULL) - break; + if (node == NULL) + break; } - if (inclusive && node && node->prefix && node->bit <= bitlen) - stack[cnt++] = node; + if (inclusive && node && node->data && node->bit <= bitlen) + stack[cnt++] = node; #ifdef PATRICIA_DEBUG if (node == NULL) fprintf (stderr, "patricia_search_best: stop at null\n"); - else if (node->prefix) + else if (&node->prefix) fprintf (stderr, "patricia_search_best: stop at %s/%d\n", - prefix_toa (node->prefix), node->prefix->bitlen); + prefix_toa (&node->prefix), node->prefix.bitlen); else fprintf (stderr, "patricia_search_best: stop at %d\n", node->bit); #endif /* PATRICIA_DEBUG */ if (cnt <= 0) - return (NULL); + return (NULL); while (--cnt >= 0) { - node = stack[cnt]; + node = stack[cnt]; #ifdef PATRICIA_DEBUG fprintf (stderr, "patricia_search_best: pop %s/%d\n", - prefix_toa (node->prefix), node->prefix->bitlen); + prefix_toa (&node->prefix), node->prefix.bitlen); #endif /* PATRICIA_DEBUG */ - if (comp_with_mask (prefix_tochar (node->prefix), + if (comp_with_mask (prefix_tochar (&node->prefix), prefix_tochar (prefix), - node->prefix->bitlen)) { + node->prefix.bitlen)) { #ifdef PATRICIA_DEBUG fprintf (stderr, "patricia_search_best: found %s/%d\n", - prefix_toa (node->prefix), node->prefix->bitlen); + prefix_toa (&node->prefix), node->prefix.bitlen); #endif /* PATRICIA_DEBUG */ - return (node); - } + return (node); + } } return (NULL); } @@ -705,200 +604,200 @@ patricia_lookup (patricia_tree_t *patricia, prefix_t *prefix) assert (prefix->bitlen <= patricia->maxbits); if (patricia->head == NULL) { - node = calloc(1, sizeof *node); - node->bit = prefix->bitlen; - node->prefix = Ref_Prefix (prefix); - node->parent = NULL; - node->l = node->r = NULL; - node->data = NULL; - patricia->head = node; + node = calloc(1, sizeof *node); + node->bit = prefix->bitlen; + node->prefix = *prefix; + node->parent = NULL; + node->l = node->r = NULL; + node->data = NULL; + patricia->head = node; #ifdef PATRICIA_DEBUG - fprintf (stderr, "patricia_lookup: new_node #0 %s/%d (head)\n", - prefix_toa (prefix), prefix->bitlen); + fprintf (stderr, "patricia_lookup: new_node #0 %s/%d (head)\n", + prefix_toa (prefix), prefix->bitlen); #endif /* PATRICIA_DEBUG */ - patricia->num_active_node++; - return (node); + patricia->num_active_node++; + return (node); } addr = prefix_touchar (prefix); bitlen = prefix->bitlen; node = patricia->head; - while (node->bit < bitlen || node->prefix == NULL) { + while (node->bit < bitlen || node->data == NULL) { - if (node->bit < patricia->maxbits && - BIT_TEST (addr[node->bit >> 3], 0x80 >> (node->bit & 0x07))) { - if (node->r == NULL) - break; + if (node->bit < patricia->maxbits && + BIT_TEST (addr[node->bit >> 3], 0x80 >> (node->bit & 0x07))) { + if (node->r == NULL) + break; #ifdef PATRICIA_DEBUG - if (node->prefix) + if (&node->prefix) fprintf (stderr, "patricia_lookup: take right %s/%d\n", - prefix_toa (node->prefix), node->prefix->bitlen); - else + prefix_toa (&node->prefix), node->prefix.bitlen); + else fprintf (stderr, "patricia_lookup: take right at %d\n", node->bit); #endif /* PATRICIA_DEBUG */ - node = node->r; - } - else { - if (node->l == NULL) - break; + node = node->r; + } + else { + if (node->l == NULL) + break; #ifdef PATRICIA_DEBUG - if (node->prefix) + if (&node->prefix) fprintf (stderr, "patricia_lookup: take left %s/%d\n", - prefix_toa (node->prefix), node->prefix->bitlen); - else + prefix_toa (&node->prefix), node->prefix.bitlen); + else fprintf (stderr, "patricia_lookup: take left at %d\n", node->bit); #endif /* PATRICIA_DEBUG */ - node = node->l; - } + node = node->l; + } - assert (node); + assert (node); } - assert (node->prefix); + assert (node->data); #ifdef PATRICIA_DEBUG fprintf (stderr, "patricia_lookup: stop at %s/%d\n", - prefix_toa (node->prefix), node->prefix->bitlen); + prefix_toa (&node->prefix), node->prefix.bitlen); #endif /* PATRICIA_DEBUG */ - test_addr = prefix_touchar (node->prefix); + test_addr = prefix_touchar (&(node->prefix)); /* find the first bit different */ check_bit = (node->bit < bitlen)? node->bit: bitlen; differ_bit = 0; for (i = 0; i*8 < check_bit; i++) { - if ((r = (addr[i] ^ test_addr[i])) == 0) { - differ_bit = (i + 1) * 8; - continue; - } - /* I know the better way, but for now */ - for (j = 0; j < 8; j++) { - if (BIT_TEST (r, (0x80 >> j))) + if ((r = (addr[i] ^ test_addr[i])) == 0) { + differ_bit = (i + 1) * 8; + continue; + } + /* I know the better way, but for now */ + for (j = 0; j < 8; j++) { + if (BIT_TEST (r, (0x80 >> j))) + break; + } + /* must be found */ + assert (j < 8); + differ_bit = i * 8 + j; break; } - /* must be found */ - assert (j < 8); - differ_bit = i * 8 + j; - break; - } if (differ_bit > check_bit) - differ_bit = check_bit; + differ_bit = check_bit; #ifdef PATRICIA_DEBUG fprintf (stderr, "patricia_lookup: differ_bit %d\n", differ_bit); #endif /* PATRICIA_DEBUG */ parent = node->parent; while (parent && parent->bit >= differ_bit) { - node = parent; - parent = node->parent; + node = parent; + parent = node->parent; #ifdef PATRICIA_DEBUG - if (node->prefix) + if (&node->prefix) fprintf (stderr, "patricia_lookup: up to %s/%d\n", - prefix_toa (node->prefix), node->prefix->bitlen); - else + prefix_toa (&node->prefix), node->prefix.bitlen); + else fprintf (stderr, "patricia_lookup: up to %d\n", node->bit); #endif /* PATRICIA_DEBUG */ } if (differ_bit == bitlen && node->bit == bitlen) { - if (node->prefix) { + if (node->data) { #ifdef PATRICIA_DEBUG fprintf (stderr, "patricia_lookup: found %s/%d\n", - prefix_toa (node->prefix), node->prefix->bitlen); + prefix_toa (&node->prefix), node->prefix.bitlen); #endif /* PATRICIA_DEBUG */ - return (node); - } - node->prefix = Ref_Prefix (prefix); + return (node); + } + node->prefix = *prefix; #ifdef PATRICIA_DEBUG - fprintf (stderr, "patricia_lookup: new node #1 %s/%d (glue mod)\n", - prefix_toa (prefix), prefix->bitlen); + fprintf (stderr, "patricia_lookup: new node #1 %s/%d (glue mod)\n", + prefix_toa (prefix), prefix->bitlen); #endif /* PATRICIA_DEBUG */ - assert (node->data == NULL); - return (node); + assert (node->data == NULL); + return (node); } new_node = calloc(1, sizeof *new_node); new_node->bit = prefix->bitlen; - new_node->prefix = Ref_Prefix (prefix); + new_node->prefix = *prefix; new_node->parent = NULL; new_node->l = new_node->r = NULL; new_node->data = NULL; patricia->num_active_node++; if (node->bit == differ_bit) { - new_node->parent = node; - if (node->bit < patricia->maxbits && - BIT_TEST (addr[node->bit >> 3], 0x80 >> (node->bit & 0x07))) { - assert (node->r == NULL); - node->r = new_node; - } - else { - assert (node->l == NULL); - node->l = new_node; - } + new_node->parent = node; + if (node->bit < patricia->maxbits && + BIT_TEST (addr[node->bit >> 3], 0x80 >> (node->bit & 0x07))) { + assert (node->r == NULL); + node->r = new_node; + } + else { + assert (node->l == NULL); + node->l = new_node; + } #ifdef PATRICIA_DEBUG - fprintf (stderr, "patricia_lookup: new_node #2 %s/%d (child)\n", - prefix_toa (prefix), prefix->bitlen); + fprintf (stderr, "patricia_lookup: new_node #2 %s/%d (child)\n", + prefix_toa (prefix), prefix->bitlen); #endif /* PATRICIA_DEBUG */ - return (new_node); + return (new_node); } if (bitlen == differ_bit) { - if (bitlen < patricia->maxbits && - BIT_TEST (test_addr[bitlen >> 3], 0x80 >> (bitlen & 0x07))) { - new_node->r = node; - } - else { - new_node->l = node; - } - new_node->parent = node->parent; - if (node->parent == NULL) { - assert (patricia->head == node); - patricia->head = new_node; - } - else if (node->parent->r == node) { - node->parent->r = new_node; - } - else { - node->parent->l = new_node; - } - node->parent = new_node; + if (bitlen < patricia->maxbits && + BIT_TEST (test_addr[bitlen >> 3], 0x80 >> (bitlen & 0x07))) { + new_node->r = node; + } + else { + new_node->l = node; + } + new_node->parent = node->parent; + if (node->parent == NULL) { + assert (patricia->head == node); + patricia->head = new_node; + } + else if (node->parent->r == node) { + node->parent->r = new_node; + } + else { + node->parent->l = new_node; + } + node->parent = new_node; #ifdef PATRICIA_DEBUG - fprintf (stderr, "patricia_lookup: new_node #3 %s/%d (parent)\n", - prefix_toa (prefix), prefix->bitlen); + fprintf (stderr, "patricia_lookup: new_node #3 %s/%d (parent)\n", + prefix_toa (prefix), prefix->bitlen); #endif /* PATRICIA_DEBUG */ } else { glue = calloc(1, sizeof *glue); glue->bit = differ_bit; - glue->prefix = NULL; + memset(&glue->prefix, 0, sizeof(prefix_t)); glue->parent = node->parent; glue->data = NULL; patricia->num_active_node++; - if (differ_bit < patricia->maxbits && - BIT_TEST (addr[differ_bit >> 3], 0x80 >> (differ_bit & 0x07))) { - glue->r = new_node; - glue->l = node; - } - else { - glue->r = node; - glue->l = new_node; - } - new_node->parent = glue; + if (differ_bit < patricia->maxbits && + BIT_TEST (addr[differ_bit >> 3], 0x80 >> (differ_bit & 0x07))) { + glue->r = new_node; + glue->l = node; + } + else { + glue->r = node; + glue->l = new_node; + } + new_node->parent = glue; - if (node->parent == NULL) { - assert (patricia->head == node); - patricia->head = glue; - } - else if (node->parent->r == node) { - node->parent->r = glue; - } - else { - node->parent->l = glue; - } - node->parent = glue; + if (node->parent == NULL) { + assert (patricia->head == node); + patricia->head = glue; + } + else if (node->parent->r == node) { + node->parent->r = glue; + } + else { + node->parent->l = glue; + } + node->parent = glue; #ifdef PATRICIA_DEBUG - fprintf (stderr, "patricia_lookup: new_node #4 %s/%d (glue+node)\n", - prefix_toa (prefix), prefix->bitlen); + fprintf (stderr, "patricia_lookup: new_node #4 %s/%d (glue+node)\n", + prefix_toa (prefix), prefix->bitlen); #endif /* PATRICIA_DEBUG */ } return (new_node); @@ -915,159 +814,90 @@ patricia_remove (patricia_tree_t *patricia, patricia_node_t *node) if (node->r && node->l) { #ifdef PATRICIA_DEBUG - fprintf (stderr, "patricia_remove: #0 %s/%d (r & l)\n", - prefix_toa (node->prefix), node->prefix->bitlen); + fprintf (stderr, "patricia_remove: #0 %s/%d (r & l)\n", + prefix_toa (&node->prefix), node->prefix.bitlen); #endif /* PATRICIA_DEBUG */ - /* this might be a placeholder node -- have to check and make sure - * there is a prefix aossciated with it ! */ - if (node->prefix != NULL) - Deref_Prefix (node->prefix); - node->prefix = NULL; - /* Also I needed to clear data pointer -- masaki */ - node->data = NULL; - return; + /* Also I needed to clear data pointer -- masaki */ + node->data = NULL; + return; } if (node->r == NULL && node->l == NULL) { #ifdef PATRICIA_DEBUG - fprintf (stderr, "patricia_remove: #1 %s/%d (!r & !l)\n", - prefix_toa (node->prefix), node->prefix->bitlen); + fprintf (stderr, "patricia_remove: #1 %s/%d (!r & !l)\n", + prefix_toa (&node->prefix), node->prefix.bitlen); #endif /* PATRICIA_DEBUG */ - parent = node->parent; - Deref_Prefix (node->prefix); - Delete (node); + parent = node->parent; + Delete (node); patricia->num_active_node--; - if (parent == NULL) { - assert (patricia->head == node); - patricia->head = NULL; - return; - } + if (parent == NULL) { + assert (patricia->head == node); + patricia->head = NULL; + return; + } - if (parent->r == node) { - parent->r = NULL; - child = parent->l; - } - else { - assert (parent->l == node); - parent->l = NULL; - child = parent->r; - } + if (parent->r == node) { + parent->r = NULL; + child = parent->l; + } + else { + assert (parent->l == node); + parent->l = NULL; + child = parent->r; + } - if (parent->prefix) - return; + if (parent->data) + return; - /* we need to remove parent too */ + /* we need to remove parent too */ - if (parent->parent == NULL) { - assert (patricia->head == parent); - patricia->head = child; - } - else if (parent->parent->r == parent) { - parent->parent->r = child; - } - else { - assert (parent->parent->l == parent); - parent->parent->l = child; - } - child->parent = parent->parent; - Delete (parent); + if (parent->parent == NULL) { + assert (patricia->head == parent); + patricia->head = child; + } + else if (parent->parent->r == parent) { + parent->parent->r = child; + } + else { + assert (parent->parent->l == parent); + parent->parent->l = child; + } + child->parent = parent->parent; + Delete (parent); patricia->num_active_node--; - return; + return; } #ifdef PATRICIA_DEBUG fprintf (stderr, "patricia_remove: #2 %s/%d (r ^ l)\n", - prefix_toa (node->prefix), node->prefix->bitlen); + prefix_toa (&node->prefix), node->prefix.bitlen); #endif /* PATRICIA_DEBUG */ if (node->r) { - child = node->r; + child = node->r; } else { - assert (node->l); - child = node->l; + assert (node->l); + child = node->l; } parent = node->parent; child->parent = parent; - Deref_Prefix (node->prefix); Delete (node); patricia->num_active_node--; if (parent == NULL) { - assert (patricia->head == node); - patricia->head = child; - return; + assert (patricia->head == node); + patricia->head = child; + return; } if (parent->r == node) { - parent->r = child; + parent->r = child; } else { assert (parent->l == node); - parent->l = child; + parent->l = child; } } - -/* { from demo.c */ - -patricia_node_t * -make_and_lookup (patricia_tree_t *tree, char *string) -{ - prefix_t *prefix; - patricia_node_t *node; - - prefix = ascii2prefix (AF_INET, string); - // printf ("make_and_lookup: %s/%d\n", prefix_toa (prefix), prefix->bitlen); - node = patricia_lookup (tree, prefix); - Deref_Prefix (prefix); - return (node); -} - -patricia_node_t * -try_search_exact (patricia_tree_t *tree, char *string) -{ - prefix_t *prefix; - patricia_node_t *node; - - prefix = ascii2prefix (AF_INET, string); - printf ("try_search_exact: %s/%d\n", prefix_toa (prefix), prefix->bitlen); - if ((node = patricia_search_exact (tree, prefix)) == NULL) { - printf ("try_search_exact: not found\n"); - } - else { - printf ("try_search_exact: %s/%d found\n", - prefix_toa (node->prefix), node->prefix->bitlen); - } - Deref_Prefix (prefix); - return (node); -} - -void -lookup_then_remove (patricia_tree_t *tree, char *string) -{ - patricia_node_t *node; - - if ((node = try_search_exact (tree, string))) - patricia_remove (tree, node); -} - -patricia_node_t * -try_search_best (patricia_tree_t *tree, char *string) -{ - prefix_t *prefix; - patricia_node_t *node; - - prefix = ascii2prefix (AF_INET, string); - printf ("try_search_best: %s/%d\n", prefix_toa (prefix), prefix->bitlen); - if ((node = patricia_search_best (tree, prefix)) == NULL) - printf ("try_search_best: not found\n"); - else - printf ("try_search_best: %s/%d found\n", - prefix_toa (node->prefix), node->prefix->bitlen); - Deref_Prefix (prefix); - return 0; // [RS] What is supposed to be returned here? - } - -/* } */ diff --git a/patricia.h b/patricia.h index 7597248..3e03f83 100644 --- a/patricia.h +++ b/patricia.h @@ -70,6 +70,8 @@ #define _PATRICIA_H #include +#include + #define HAVE_IPV6 1 // JS: force use of ip6 @@ -80,8 +82,6 @@ #define BIT_TEST(f, b) ((f) & (b)) /* } */ -#define addroute make_and_lookup - #if defined(_WIN32) || defined(_WIN64) #include #include @@ -126,8 +126,8 @@ typedef void (*void_fn1_t)(void *); typedef void (*void_fn2_t)(struct _prefix_t *, void *); typedef struct _patricia_node_t { - u_int bit; - prefix_t *prefix; + u_int bit; + prefix_t prefix; struct _patricia_node_t *l, *r; struct _patricia_node_t *parent; void *data; @@ -138,6 +138,7 @@ typedef struct _patricia_tree_t { patricia_node_t *head; u_int maxbits; int num_active_node; + u_short frozen; } patricia_tree_t; @@ -152,20 +153,13 @@ void Clear_Patricia (patricia_tree_t *patricia, void_fn1_t func); void Destroy_Patricia (patricia_tree_t *patricia, void_fn1_t func); void patricia_process (patricia_tree_t *patricia, void_fn2_t func); -void Deref_Prefix (prefix_t * prefix); -prefix_t * New_Prefix(int, void *, int); +int New_Prefix(int, void *, int, prefix_t*); /* { from demo.c */ -prefix_t * -ascii2prefix (int family, char *string); - char * prefix_toa2x(prefix_t *prefix, char *buff, int with_len); -patricia_node_t * -make_and_lookup (patricia_tree_t *tree, char *string); - /* } */ #define PATRICIA_MAXBITS 128 @@ -181,7 +175,7 @@ make_and_lookup (patricia_tree_t *tree, char *string); patricia_node_t **Xsp = Xstack; \ patricia_node_t *Xrn = (Xhead); \ while ((Xnode = Xrn)) { \ - if (Xnode->prefix) + if (Xnode->data) #define PATRICIA_WALK_ALL(Xhead, Xnode) \ do { \ diff --git a/pytricia.c b/pytricia.c index 59ac182..a86bd3a 100644 --- a/pytricia.c +++ b/pytricia.c @@ -74,14 +74,14 @@ static void _set_ipaddr_refs(void) { } #endif -static prefix_t * -_prefix_convert(int family, const char *addr) { +static int +_prefix_convert(int family, const char *addr, prefix_t* prefix) { int prefixlen = -1; char addrcopy[128]; - if (strlen(addr) < 4) { - return NULL; - } + if (strlen(addr) < 2) { + return 0; + } strncpy(addrcopy, addr, 128); char *slash = strchr(addrcopy, '/'); @@ -108,9 +108,9 @@ _prefix_convert(int family, const char *addr) { } struct in_addr sin; if (inet_pton(AF_INET, addrcopy, &sin) != 1) { - return NULL; + return 0; } - return New_Prefix(AF_INET, &sin, prefixlen); + return New_Prefix(AF_INET, &sin, prefixlen, prefix); } else if (family == AF_INET6) { if (prefixlen == -1 || prefixlen < 0 || prefixlen > 128) { prefixlen = 128; @@ -118,43 +118,43 @@ _prefix_convert(int family, const char *addr) { struct in6_addr sin6; if (inet_pton(AF_INET6, addrcopy, &sin6) != 1) { - return NULL; + return 0; } - return New_Prefix(AF_INET6, &sin6, prefixlen); + return New_Prefix(AF_INET6, &sin6, prefixlen, prefix); } else { - return NULL; + return 0; } } -static prefix_t * -_packed_addr_to_prefix(char *addrbuf, long len) { - prefix_t *pfx_rv = NULL; +static int +_packed_addr_to_prefix(char *addrbuf, long len, prefix_t* pfx_rv) { + int ret_ok = 0; if (len == 4) { - pfx_rv = New_Prefix(AF_INET, addrbuf, 32); + ret_ok = New_Prefix(AF_INET, addrbuf, 32, pfx_rv); } else if (len == 16) { - pfx_rv = New_Prefix(AF_INET6, addrbuf, 128); + ret_ok = New_Prefix(AF_INET6, addrbuf, 128, pfx_rv); } else { PyErr_SetString(PyExc_ValueError, "Address bytes must be of length 4 or 16"); } - return pfx_rv; + return ret_ok; } #if PY_MAJOR_VERSION == 3 -static prefix_t * -_bytes_to_prefix(PyObject *key) { +static int +_bytes_to_prefix(PyObject *key, prefix_t* pfx_rv) { char *addrbuf = NULL; Py_ssize_t len = 0; if (PyBytes_AsStringAndSize(key, &addrbuf, &len) < 0) { PyErr_SetString(PyExc_ValueError, "Error decoding bytes"); - return NULL; + return 0; } - return _packed_addr_to_prefix(addrbuf, len); + return _packed_addr_to_prefix(addrbuf, len, pfx_rv); } #endif -static prefix_t * -_key_object_to_prefix(PyObject *key) { - prefix_t *pfx_rv = NULL; +static int +_key_object_to_prefix(PyObject *key, prefix_t* pfx_rv) { + int ret_ok = 0; #if PY_MAJOR_VERSION == 3 #if PY_MINOR_VERSION >= 4 if (!_ipaddr_isset) { @@ -166,41 +166,41 @@ _key_object_to_prefix(PyObject *key) { int rv = PyUnicode_READY(key); if (rv < 0) { PyErr_SetString(PyExc_ValueError, "Error parsing string prefix"); - return NULL; + return 0; } - char* temp = PyUnicode_AsUTF8(key); - if (temp == NULL) { + const char* temp = PyUnicode_AsUTF8(key); + if (temp == 0) { PyErr_SetString(PyExc_ValueError, "Error parsing string prefix"); - return NULL; + return 0; } if (strchr(temp, '.') || strchr(temp, ':')) { - pfx_rv = _prefix_convert(0, temp); + ret_ok = _prefix_convert(0, temp, pfx_rv); } else { PyErr_SetString(PyExc_ValueError, "Invalid key type"); - return NULL; + return 0; } } else if (PyLong_Check(key)) { unsigned long packed_addr = htonl(PyLong_AsUnsignedLong(key)); - pfx_rv = New_Prefix(AF_INET, &packed_addr, 32); + ret_ok = New_Prefix(AF_INET, &packed_addr, 32, pfx_rv); } else if (PyBytes_Check(key)) { - pfx_rv = _bytes_to_prefix(key); + ret_ok = _bytes_to_prefix(key, pfx_rv); } else if (PyTuple_Check(key)) { PyObject* value = PyTuple_GetItem(key, 0); PyObject* size = PyTuple_GetItem(key, 1); if (!PyBytes_Check(value)) { PyErr_SetString(PyExc_ValueError, "Invalid key tuple value type"); - return NULL; + return 0; } Py_ssize_t slen = PyBytes_Size(value); if (slen != 4 && slen != 16) { PyErr_SetString(PyExc_ValueError, "Invalid key tuple value"); - return NULL; + return 0; } - pfx_rv = _bytes_to_prefix(value); - if (pfx_rv) { + ret_ok = _bytes_to_prefix(value, pfx_rv); + if (ret_ok) { if (!PyLong_Check(size)) { PyErr_SetString(PyExc_ValueError, "Invalid key tuple size type"); - return NULL; + return 0; } unsigned long bitlen = PyLong_AsUnsignedLong(size); if (bitlen < pfx_rv->bitlen) @@ -215,7 +215,7 @@ _key_object_to_prefix(PyObject *key) { if (netaddr) { PyObject *packed = PyObject_GetAttrString(netaddr, "packed"); if (packed && PyBytes_Check(packed)) { - pfx_rv = _bytes_to_prefix(packed); + ret_ok = _bytes_to_prefix(packed, pfx_rv); PyObject *prefixlen = PyObject_GetAttrString(key, "prefixlen"); if (prefixlen && PyLong_Check(prefixlen)) { long bitlen = PyLong_AsLong(prefixlen); @@ -233,7 +233,7 @@ _key_object_to_prefix(PyObject *key) { } else if (ipaddr_base && PyObject_IsInstance(key, ipaddr_base)) { PyObject *packed = PyObject_GetAttrString(key, "packed"); if (packed && PyBytes_Check(packed)) { - pfx_rv = _bytes_to_prefix(packed); + ret_ok = _bytes_to_prefix(packed, pfx_rv); Py_DECREF(packed); } else { PyErr_SetString(PyExc_ValueError, "Error getting raw representation of IPAddress"); @@ -243,38 +243,38 @@ _key_object_to_prefix(PyObject *key) { else { PyErr_SetString(PyExc_ValueError, "Invalid key type"); } -#else // python2 +#else //python2 if (PyString_Check(key)) { char* temp = PyString_AsString(key); Py_ssize_t slen = PyString_Size(key); if (strchr(temp, '.') || strchr(temp, ':')) { - pfx_rv = _prefix_convert(0, temp); + ret_ok = _prefix_convert(0, temp, pfx_rv); } else if (slen == 4 || slen == 16) { - pfx_rv = _packed_addr_to_prefix(temp, slen); + ret_ok = _packed_addr_to_prefix(temp, slen, pfx_rv); } else { PyErr_SetString(PyExc_ValueError, "Invalid key type"); } } else if (PyLong_Check(key) || PyInt_Check(key)) { unsigned long packed_addr = htonl(PyInt_AsUnsignedLongMask(key)); - pfx_rv = New_Prefix(AF_INET, &packed_addr, 32); + ret_ok = New_Prefix(AF_INET, &packed_addr, 32, pfx_rv); } else if (PyTuple_Check(key)) { PyObject* value = PyTuple_GetItem(key, 0); PyObject* size = PyTuple_GetItem(key, 1); if (!PyString_Check(value)) { PyErr_SetString(PyExc_ValueError, "Invalid key tuple value type"); - return NULL; + return 0; } Py_ssize_t slen = PyString_Size(value); if (slen != 4 && slen != 16) { PyErr_SetString(PyExc_ValueError, "Invalid key tuple value"); - return NULL; + return 0; } char* temp = PyString_AsString(value); - pfx_rv = _packed_addr_to_prefix(temp, slen); + ret_ok = _packed_addr_to_prefix(temp, slen, pfx_rv); if (pfx_rv) { if (!PyLong_Check(size) && !PyInt_Check(size)) { PyErr_SetString(PyExc_ValueError, "Invalid key tuple size type"); - return NULL; + return 0; } unsigned long bitlen = PyInt_AsLong(size); if (bitlen < pfx_rv->bitlen) @@ -284,7 +284,7 @@ _key_object_to_prefix(PyObject *key) { PyErr_SetString(PyExc_ValueError, "Invalid key type"); } #endif - return pfx_rv; + return ret_ok; } static PyObject * @@ -383,13 +383,13 @@ pytricia_length(PyTricia *self) static PyObject* pytricia_subscript(PyTricia *self, PyObject *key) { - prefix_t *subnet = _key_object_to_prefix(key); - if (subnet == NULL) { + prefix_t prefix; memset(&prefix, 0, sizeof(prefix)); + int ret_ok = _key_object_to_prefix(key, &prefix); + if (!ret_ok) { PyErr_SetString(PyExc_ValueError, "Invalid prefix."); return NULL; } - patricia_node_t* node = patricia_search_best(self->m_tree, subnet); - Deref_Prefix(subnet); + patricia_node_t* node = patricia_search_best(self->m_tree, &prefix); if (!node) { PyErr_SetString(PyExc_KeyError, "Prefix not found."); @@ -404,13 +404,18 @@ pytricia_subscript(PyTricia *self, PyObject *key) { static int pytricia_internal_delete(PyTricia *self, PyObject *key) { - prefix_t *prefix = _key_object_to_prefix(key); - if (prefix == NULL) { + if (self->m_tree->frozen) { + PyErr_SetString(PyExc_ValueError, "can not modify a frozen pytricia! Thaw?"); + return -1; + } + prefix_t prefix; memset(&prefix, 0, sizeof(prefix)); + + int ret_ok = _key_object_to_prefix(key, &prefix); + if (!ret_ok) { PyErr_SetString(PyExc_ValueError, "Invalid prefix."); return -1; } - patricia_node_t* node = patricia_search_exact(self->m_tree, prefix); - Deref_Prefix(prefix); + patricia_node_t* node = patricia_search_exact(self->m_tree, &prefix); if (!node) { PyErr_SetString(PyExc_KeyError, "Prefix doesn't exist."); @@ -430,18 +435,23 @@ _pytricia_assign_subscript_internal(PyTricia *self, PyObject *key, PyObject *val if (!value) { return pytricia_internal_delete(self, key); } - - prefix_t *prefix = _key_object_to_prefix(key); - if (!prefix) { + + if (self->m_tree->frozen) { + PyErr_SetString(PyExc_ValueError, "can not modify a frozen pytricia! Thaw?"); + return -1; + } + + prefix_t prefix; memset(&prefix, 0, sizeof(prefix)); + int ret_ok = _key_object_to_prefix(key, &prefix); + if (!ret_ok) { return -1; } // if prefixlen > -1, it should override (possibly) parsed prefix len in key if (prefixlen != -1) { - prefix->bitlen = prefixlen; + prefix.bitlen = prefixlen; } - patricia_node_t *node = patricia_lookup(self->m_tree, prefix); - Deref_Prefix(prefix); + patricia_node_t *node = patricia_lookup(self->m_tree, &prefix); if (!node) { PyErr_SetString(PyExc_ValueError, "Error inserting into patricia tree"); @@ -528,13 +538,13 @@ pytricia_get(register PyTricia *obj, PyObject *args) { if (!PyArg_ParseTuple(args, "O|O:get", &key, &defvalue)) { return NULL; } - prefix_t *prefix = _key_object_to_prefix(key); - if (!prefix) { + prefix_t prefix; memset(&prefix, 0, sizeof(prefix)); + int ret_ok = _key_object_to_prefix(key, &prefix); + if (!ret_ok) { PyErr_SetString(PyExc_ValueError, "Invalid prefix."); return NULL; } - patricia_node_t* node = patricia_search_best(obj->m_tree, prefix); - Deref_Prefix(prefix); + patricia_node_t* node = patricia_search_best(obj->m_tree, &prefix); if (!node) { if (defvalue) { @@ -557,29 +567,29 @@ pytricia_get_key(register PyTricia *obj, PyObject *args) { return NULL; } - prefix_t *prefix = _key_object_to_prefix(key); - if (!prefix) { + prefix_t prefix; memset(&prefix, 0, sizeof(prefix)); + int ret_ok = _key_object_to_prefix(key, &prefix); + if (!ret_ok) { PyErr_SetString(PyExc_ValueError, "Invalid prefix."); return NULL; } - patricia_node_t* node = patricia_search_best(obj->m_tree, prefix); - Deref_Prefix(prefix); + patricia_node_t* node = patricia_search_best(obj->m_tree, &prefix); if (!node) { Py_RETURN_NONE; } - return _prefix_to_key_object(node->prefix, obj->m_raw_output); + return _prefix_to_key_object(&node->prefix, obj->m_raw_output); } static int pytricia_contains(PyTricia *self, PyObject *key) { - prefix_t *prefix = _key_object_to_prefix(key); - if (!prefix) { + prefix_t prefix; memset(&prefix, 0, sizeof(prefix)); + int ret_ok = _key_object_to_prefix(key, &prefix); + if (!ret_ok) { return -1; } - patricia_node_t* node = patricia_search_best(self->m_tree, prefix); - Deref_Prefix(prefix); + patricia_node_t* node = patricia_search_best(self->m_tree, &prefix); if (node) { return 1; } @@ -591,14 +601,13 @@ pytricia_has_key(PyTricia *self, PyObject *args) { PyObject *key = NULL; if (!PyArg_ParseTuple(args, "O", &key)) return NULL; - - prefix_t *prefix = _key_object_to_prefix(key); - if (!prefix) { + prefix_t prefix; memset(&prefix, 0, sizeof(prefix)); + int ret_ok = _key_object_to_prefix(key, &prefix); + if (!ret_ok) { PyErr_SetString(PyExc_ValueError, "Invalid prefix."); return NULL; } - patricia_node_t* node = patricia_search_exact(self->m_tree, prefix); - Deref_Prefix(prefix); + patricia_node_t* node = patricia_search_exact(self->m_tree, &prefix); if (node) { Py_RETURN_TRUE; } @@ -616,7 +625,7 @@ pytricia_keys(register PyTricia *self, PyObject *unused) { int err = 0; PATRICIA_WALK (self->m_tree->head, node) { - PyObject *item = _prefix_to_key_object(node->prefix, self->m_raw_output); + PyObject *item = _prefix_to_key_object(&node->prefix, self->m_raw_output); if (!item) { Py_DECREF(rvlist); return NULL; @@ -639,8 +648,9 @@ pytricia_children(register PyTricia *self, PyObject *args) { return NULL; } - prefix_t *prefix = _key_object_to_prefix(key); - if (!prefix) { + prefix_t prefix; memset(&prefix, 0, sizeof(prefix)); + int ret_ok = _key_object_to_prefix(key, &prefix); + if (!ret_ok) { PyErr_SetString(PyExc_ValueError, "Invalid prefix."); return NULL; } @@ -650,8 +660,7 @@ pytricia_children(register PyTricia *self, PyObject *args) { return NULL; } - patricia_node_t* base_node = patricia_search_exact(self->m_tree, prefix); - Deref_Prefix(prefix); + patricia_node_t* base_node = patricia_search_exact(self->m_tree, &prefix); if (!base_node) { PyErr_SetString(PyExc_KeyError, "Prefix doesn't exist."); Py_DECREF(rvlist); @@ -663,7 +672,7 @@ pytricia_children(register PyTricia *self, PyObject *args) { PATRICIA_WALK (base_node, node) { /* Discard first prefix (we want strict children) */ if (node != base_node) { - PyObject *item = _prefix_to_key_object(node->prefix, self->m_raw_output); + PyObject *item = _prefix_to_key_object(&node->prefix, self->m_raw_output); if (!item) { Py_DECREF(rvlist); return NULL; @@ -687,24 +696,288 @@ pytricia_parent(register PyTricia *self, PyObject *args) { return NULL; } - prefix_t *prefix = _key_object_to_prefix(key); - if (!prefix) { + prefix_t prefix; memset(&prefix, 0, sizeof(prefix)); + int ret_ok = _key_object_to_prefix(key, &prefix); + if (!ret_ok) { PyErr_SetString(PyExc_ValueError, "Invalid prefix."); return NULL; } - patricia_node_t* node = patricia_search_exact(self->m_tree, prefix); - Deref_Prefix(prefix); + patricia_node_t* node = patricia_search_exact(self->m_tree, &prefix); if (!node) { PyErr_SetString(PyExc_KeyError, "Prefix doesn't exist."); return NULL; } - patricia_node_t* parent_node = patricia_search_best2(self->m_tree, node->prefix, 0); + patricia_node_t* parent_node = patricia_search_best2(self->m_tree, &node->prefix, 0); if (!parent_node) { Py_RETURN_NONE; } - return _prefix_to_key_object(parent_node->prefix, self->m_raw_output); + return _prefix_to_key_object(&parent_node->prefix, self->m_raw_output); +} + +static PyObject* +pytricia_freeze(register PyTricia *self, PyObject *unused) { + if (self->m_tree->frozen) { + Py_RETURN_NONE; + } + + patricia_node_t *node = NULL; + + // Get count of all nodes + size_t count = 0; + PATRICIA_WALK_ALL (self->m_tree->head, node) { + count += 1; + } PATRICIA_WALK_END; + + // allocate contiguous space for all nodes + patricia_node_t* new_node = calloc(count, sizeof(patricia_node_t)); + + // copy all nodes to new array and discard originals + size_t idx = 0; + patricia_node_t** free_list = calloc(count, sizeof(patricia_node_t*)); + PATRICIA_WALK_ALL (self->m_tree->head, node) { + new_node[idx] = *node; + if(node->l) + node->l->parent = &(new_node[idx]); + if(node->r) + node->r->parent = &(new_node[idx]); + if(node->parent == NULL) { + assert (self->m_tree->head == node); + self->m_tree->head = new_node; + } + else if (node->parent->r == node) { + node->parent->r = &(new_node[idx]); + } + else { + node->parent->l = &(new_node[idx]); + } + free_list[idx] = node; + idx++; + } PATRICIA_WALK_END; + + // clearing via free_list because otherise PATRICIA_WALK_END + // will illegally reference items from node to determine next steps + for(size_t i=0; im_tree->frozen = 1; + + Py_RETURN_NONE; +} + +static PyObject* +pytricia_thaw(register PyTricia *self, PyObject *unused) { + if (!self->m_tree->frozen) { + Py_RETURN_NONE; // already thaw'd + } + if (self->m_tree->head == NULL) { + self->m_tree->frozen = 0; + Py_RETURN_NONE; + } + patricia_node_t* original_head = self->m_tree->head; + + // walk all nodes, allocating heap space for individual + // nodes and re-linking + patricia_node_t *node = NULL; + PATRICIA_WALK_ALL (self->m_tree->head, node) { + patricia_node_t* new_node = calloc(1, sizeof(patricia_node_t)); + *new_node = *node; + if(node->l) + node->l->parent = new_node; + if(node->r) + node->r->parent = new_node; + if(node->parent == NULL) { + assert (self->m_tree->head == node); + self->m_tree->head = new_node; + } + else if (node->parent->r == node) { + node->parent->r = new_node; + } + else { + node->parent->l = new_node; + } + } PATRICIA_WALK_END; + + free(original_head); + + // mark as NOT frozen + self->m_tree->frozen = 0; + + Py_RETURN_NONE; +} + +// forward declaration +static PyTypeObject PyTriciaType; + +static PyObject* pytricia_reduce(PyTricia *self, PyObject *Py_UNUSED(ignored)) { + if (!self->m_tree->frozen) { + PyErr_SetString(PyExc_RuntimeError, "pytri must be frozen before attempting to pickle!"); + return NULL; + } + + PyObject *args = PyTuple_New(1); + if (!args) { + return NULL; + } + PyTuple_SET_ITEM(args, 0, PyLong_FromLong(self->m_tree->maxbits)); + + // setup the extra dictionary info for setstate + PyObject* dict = PyDict_New(); + PyObject* bytes = PyBytes_FromStringAndSize((const char*)self->m_tree, sizeof(patricia_tree_t)); + + if (!dict || !bytes) { + PyErr_SetString(PyExc_MemoryError, "unable to allocate space for tri"); + Py_XDECREF(dict); + Py_XDECREF(bytes); + return NULL; + } + + int ret_err = PyDict_SetItemString(dict, "tree", bytes); + Py_DECREF(bytes); // dictionary now owns the reference + if(ret_err) { + PyErr_SetString(PyExc_MemoryError, "error writing tree_t to dictionary"); + Py_DECREF(dict); + return NULL; + } + + // Now set nodes internal block + patricia_node_t *node = NULL; + size_t count = 0; + PATRICIA_WALK_ALL (self->m_tree->head, node) { + count += 1; + } PATRICIA_WALK_END; + + bytes = PyBytes_FromStringAndSize((const char*)self->m_tree->head, sizeof(patricia_node_t)*count); + ret_err = PyDict_SetItemString(dict, "nodes", bytes); + Py_DECREF(bytes); // dictionary now owns the reference + if(ret_err) { + PyErr_SetString(PyExc_MemoryError, "error writing nodes to dictionary"); + Py_DECREF(dict); + return NULL; + } + + PyObject* list = PyList_New(count); + if (!list) { + PyErr_SetString(PyExc_MemoryError, "error allocating data list"); + return NULL; + } + + count = 0; + if (self->m_tree->head) { + PATRICIA_WALK_ALL (self->m_tree->head, node) { + // some nodes may be glued in middle and will have prefix, but not data + // In case of no data, we'll leave the list entry as default (None) + if(node->data) { + // SET_ITEM steal reference so increment in advance to keep original + Py_INCREF(node->data); + PyList_SET_ITEM(list, count, node->data); + } + else { + Py_INCREF(Py_None); + PyList_SET_ITEM(list, count, Py_None); + } + count += 1; + } PATRICIA_WALK_END; + } + + ret_err = PyDict_SetItemString(dict, "data", list); + Py_DECREF(list); // dictionary now owns the reference + if(ret_err) { + PyErr_SetString(PyExc_MemoryError, "error writing data list to dictionary"); + Py_DECREF(dict); + return NULL; + } + + PyObject* out_tuple = PyTuple_Pack(3, (PyObject*)Py_TYPE(self), args, dict); + Py_DECREF(args); + Py_DECREF(dict); + return out_tuple; +} + +static PyObject* pytricia_setstate(PyTricia *self, PyObject *args) { + PyObject *state; + if (!PyArg_ParseTuple(args, "O", &state)) { + return NULL; + } + if (!PyDict_Check(state)) { + PyErr_SetString(PyExc_TypeError, "__setstate__ argument must be a dictionary"); + return NULL; + } + + PyObject* bytes = PyDict_GetItemString(state, "tree"); + if (!bytes || !PyBytes_Check(bytes) || PyBytes_Size(bytes) != sizeof(patricia_tree_t)) { + PyErr_SetString(PyExc_TypeError, "__setstate__ failed tree type checking"); + return NULL; + } + memcpy(self->m_tree, PyBytes_AsString(bytes), sizeof(patricia_tree_t)); + + // restore head/node data + PyObject* nodebytes = PyDict_GetItemString(state, "nodes"); + if (!PyBytes_Check(nodebytes)) { + PyErr_SetString(PyExc_TypeError, "__setstate__ failed nodes type checking"); + return NULL; + } + if (PyBytes_Size(nodebytes)) { + patricia_node_t* original_head = self->m_tree->head; + self->m_tree->head = calloc(1, PyBytes_Size(nodebytes)); + ssize_t offset_bytes = (char*)self->m_tree->head - (char*)original_head; + size_t num_nodes = PyBytes_Size(nodebytes) / sizeof(patricia_node_t); + if(self->m_tree->head == NULL) { + PyErr_SetString(PyExc_MemoryError, "__setstate__ error allocating space for nodes"); + return NULL; + } + memcpy(self->m_tree->head, PyBytes_AsString(nodebytes), PyBytes_Size(nodebytes)); + + // Now re-write the links relative to the start of the contiguous memory block + patricia_node_t *node = self->m_tree->head; + for(size_t i=0; iparent) { + node->parent = (patricia_node_t*)((char*)node->parent + offset_bytes); + } + if(node->l) { + node->l = (patricia_node_t*)((char*)node->l + offset_bytes); + } + if(node->r) { + node->r = (patricia_node_t*)((char*)node->r + offset_bytes); + } + node++; + } + } + + // Restore node data items + PyObject* list = PyDict_GetItemString(state, "data"); + if (!PyList_Check(list)) { + PyErr_SetString(PyExc_TypeError, "__setstate__ data is not list as expected!"); + return NULL; + } + else { + Py_ssize_t list_len = PyList_Size(list); + if (list_len * (Py_ssize_t)sizeof(patricia_node_t) != PyBytes_Size(nodebytes)) { + PyErr_SetString(PyExc_TypeError, "__setstate__ node and data list sizes inconsistent!"); + return NULL; + } + } + patricia_node_t *node = NULL; + size_t count = 0; + if (self->m_tree->head) { + PATRICIA_WALK_ALL (self->m_tree->head, node) { + node->data = PyList_GET_ITEM(list, count); + Py_INCREF(node->data); // make our own reference + // in special case of glue node with data value of None + // we'll want to remove python None and instead use NULL + if (node->data == Py_None) { + Py_DECREF(node->data); + node->data = NULL; + } + count += 1; + } PATRICIA_WALK_END; + } + + Py_RETURN_NONE; } static PyMappingMethods pytricia_as_mapping = { @@ -740,6 +1013,10 @@ static PyMethodDef pytricia_methods[] = { {"insert", (PyCFunction)pytricia_insert, METH_VARARGS, "insert(prefix, data) -> data\nCreate mapping between prefix and data in tree."}, {"children", (PyCFunction)pytricia_children, METH_VARARGS, "children(prefix) -> list\nReturn a list of all prefixes that are more specific than the given prefix (the prefix must be present as an exact match)."}, {"parent", (PyCFunction)pytricia_parent, METH_VARARGS, "parent(prefix) -> prefix\nReturn the immediate parent of the given prefix (the prefix must be present as an exact match)."}, + {"freeze", (PyCFunction)pytricia_freeze, METH_NOARGS, "freeze() -> \nCompacts pytricia object for efficient access, but disallows updates"}, + {"thaw", (PyCFunction)pytricia_thaw, METH_NOARGS, "thaw() -> \nreverses a frozen pytricia object to allow updates"}, + {"__reduce__", (PyCFunction)pytricia_reduce, METH_NOARGS, "Return state information for pickling"}, + {"__setstate__", (PyCFunction)pytricia_setstate, METH_VARARGS, "Set state information for unpickling"}, {NULL, NULL} /* sentinel */ }; @@ -770,8 +1047,8 @@ pytriciaiter_next(PyTriciaIter *iter) iter->m_Xrn = (patricia_node_t *) 0; } - if (iter->m_Xnode->prefix) { - return _prefix_to_key_object(iter->m_Xnode->prefix, iter->m_parent->m_raw_output); + if (iter->m_Xnode->data) { + return _prefix_to_key_object(&iter->m_Xnode->prefix, iter->m_parent->m_raw_output); } } else { PyErr_SetNone(PyExc_StopIteration); diff --git a/setup.py b/setup.py index 59eae58..557bf9a 100644 --- a/setup.py +++ b/setup.py @@ -16,27 +16,30 @@ # along with Pytricia. If not, see . # -from setuptools import setup, Extension +from setuptools import setup, Extension, find_packages setup(name="pytricia", - version="1.0.2", + version="1.3.0", description="An efficient IP address storage and lookup module for Python.", author="Joel Sommers", author_email="jsommers@acm.org", url="https://github.com/jsommers/pytricia", - # download_url="http://cs.colgate.edu/~jsommers/downloads/pytricia-0.1.tar.gz", + include_package_data=True, + packages=find_packages(), keywords=['patricia tree','IP addresses'], + license_files=["COPYING*"], classifiers=[ "Programming Language :: Python :: 2", "Programming Language :: Python :: 3", "Intended Audience :: Developers", - "License :: OSI Approved :: GNU Lesser General Public License v3 or later (LGPLv3+)", "Operating System :: OS Independent", "Topic :: Software Development :: Libraries :: Python Modules", "Topic :: Internet :: Log Analysis", "Topic :: Scientific/Engineering", ], ext_modules=[ - Extension("pytricia", ["pytricia.c","patricia.c"]), + Extension("pytricia", ["pytricia.c","patricia.c"], + # extra_compile_args = ["-g", "-O0"] # Enable debug info, disable optimization + ), ], long_description=''' Pytricia is a Python module to store IP prefixes in a diff --git a/test.py b/test.py index 6c52299..71ae395 100644 --- a/test.py +++ b/test.py @@ -23,6 +23,10 @@ import socket import struct import sys +import pickle +import multiprocessing +from multiprocessing import Process, Queue + def dumppyt(t): print ("\nDumping Pytricia") @@ -279,6 +283,24 @@ def testInsert4(self): with self.assertRaises(ValueError) as cm: val = pyt.insert("fe80::1") # should raise an exception + def testInsertValidShortV6(self): + pyt = pytricia.PyTricia(128) + val = pyt.insert("::2", "a") # A valid short IPv6 address should succeed + self.assertIs(val, None) + self.assertEqual(len(pyt), 1) + self.assertEqual(pyt["::2"], "a") + self.assertIn("::2", pyt) + + def testInsertInvalidShortV6(self): + pyt = pytricia.PyTricia(128) + with self.assertRaisesRegex(ValueError, "Invalid key.") as cm: + val = pyt.insert(":2:", "a") # should raise an exception + + def testInsertInvalidShortV4(self): + pyt = pytricia.PyTricia() + with self.assertRaisesRegex(ValueError, "Invalid key.") as cm: + val = pyt.insert("192.", "a") # should raise an exception + def testGet(self): pyt = pytricia.PyTricia() pyt.insert("10.0.0.0/8", "a") @@ -449,6 +471,199 @@ def testRawIP6(self): self.assertListEqual(list(pyt.children((b'\xAA\xBB\xCC\xDD\xAA\xBB\xCC\xDD\xAA\xBB\xCC\xDD\x01\x02\x03\x00', 96+24))), [(b'\xAA\xBB\xCC\xDD\xAA\xBB\xCC\xDD\xAA\xBB\xCC\xDD\x01\x02\x03\x04', 96+32)]) self.assertListEqual(sorted(list(pyt)), sorted(prefixes)) + def testPickleMustFreeze(self): + pyt = pytricia.PyTricia() + pyt["10.0.0.0/8"] = 'a' + with self.assertRaises(RuntimeError): + s = pickle.dumps(pyt) + pyt.freeze() + s = pickle.dumps(pyt) + + # no modification allowed while frozen + with self.assertRaises(ValueError): + pyt["10.1.0.0/16"] = 'b' + with self.assertRaises(ValueError): + del pyt["10.1.0.0/16"] + with self.assertRaises(ValueError): + pyt.delete("10.1.0.0/16") + with self.assertRaises(ValueError): + pyt.insert("10.1.0.0/16") + + # still frozen ofter pickle/unpickle + pyt = pickle.loads(s) + with self.assertRaises(ValueError): + pyt["10.1.0.0/16"] = 'b' + with self.assertRaises(ValueError): + del pyt["10.1.0.0/16"] + with self.assertRaises(ValueError): + pyt.delete("10.1.0.0/16") + with self.assertRaises(ValueError): + pyt.insert("10.1.0.0/16") + + + def testPickleEmpty(self): + """Make sure things function when pytri empty""" + pyt = pytricia.PyTricia() + pyt.freeze() + s = pickle.dumps(pyt) + pyt = pickle.loads(s) + + def testPickleBasic(self): + pyt = pytricia.PyTricia() + pyt["10.0.0.0/8"] = 'a' + pyt["10.1.0.0/16"] = 'b' + pyt.freeze() + s = pickle.dumps(pyt) + pyt = pickle.loads(s) + + self.assertEqual(pyt["10.0.0.0/8"], 'a') + self.assertEqual(pyt["10.1.0.0/16"], 'b') + self.assertEqual(pyt["10.0.0.0"], 'a') + self.assertEqual(pyt["10.1.0.0"], 'b') + self.assertEqual(pyt["10.1.0.1"], 'b') + self.assertEqual(pyt["10.0.0.1"], 'a') + + self.assertTrue('10.0.0.0' in pyt) + self.assertTrue('10.1.0.0' in pyt) + self.assertTrue('10.0.0.1' in pyt) + self.assertFalse('9.0.0.0' in pyt) + self.assertFalse('0.0.0.0' in pyt) + + self.assertTrue(pyt.has_key('10.0.0.0/8')) + self.assertTrue(pyt.has_key('10.1.0.0/16')) + self.assertFalse(pyt.has_key('10.2.0.0/16')) + self.assertFalse(pyt.has_key('9.0.0.0/8')) + self.assertFalse(pyt.has_key('10.0.0.1')) + + self.assertTrue(pyt.has_key('10.0.0.0/8')) + self.assertTrue(pyt.has_key('10.1.0.0/16')) + self.assertFalse(pyt.has_key('10.2.0.0/16')) + self.assertFalse(pyt.has_key('9.0.0.0/8')) + self.assertFalse(pyt.has_key('10.0.0.0')) + + self.assertListEqual(sorted(['10.0.0.0/8','10.1.0.0/16']), sorted(pyt.keys())) + + def testPickleBasicSubprocess(self): + #NOTE: This is a bit odd for unittest but I've done this + # specifically to ensure shared memory not possible + # between the pickle/unpickle side. Without this isolation + # it is possible for references to be put together incorrectly + # but still *appear* to work due to old memory locations still + # being accessible. + pyt = pytricia.PyTricia() + pyt["10.0.0.0/8"] = 'a' + pyt["10.1.0.0/16"] = 'b' + pyt.freeze() + pickled_data = pickle.dumps(pyt) + + q = Queue() + p = Process(target=basicAsserts, args=(q, pickled_data)) + p.start() + p.join() + + if not q.empty(): + status, msg = q.get() + if status == "success": + pass + else: + self.fail(f"Assertion failed in subprocess: {msg}") + else: + self.fail("No result received from subprocess") + + def testPickleMoreComplex(self): + pyt = pytricia.PyTricia() + pyt["10.0.0.0/8"] = 'a' + pyt["10.1.0.0/16"] = 'b' + pyt["10.0.1.0/24"] = 'c' + pyt["0.0.0.0/0"] = 'default route' + pyt.freeze() + s = pickle.dumps(pyt) + pyt = pickle.loads(s) + + self.assertEqual(pyt['10.0.0.1/32'], 'a') + self.assertEqual(pyt['10.0.0.1'], 'a') + + self.assertFalse(pyt.has_key('1.0.0.0/8')) + # with 0.0.0.0/0, everything should be 'in' + for i in range(256): + self.assertTrue('{}.2.3.4'.format(i) in pyt) + # default for all but 10.0.0.0/8 + if i != 10: + self.assertEqual(pyt['{}.2.3.4'.format(i)], 'default route') + + def testPickleFreezeThaw(self): + pyt = pytricia.PyTricia() + pyt["10.0.0.0/8"] = 'a' + pyt.freeze() + s = pickle.dumps(pyt) + pyt = pickle.loads(s) + pyt.thaw() + pyt["10.1.0.0/16"] = 'b' + + self.assertEqual(pyt["10.0.0.0/8"], 'a') + self.assertEqual(pyt["10.1.0.0/16"], 'b') + self.assertEqual(pyt["10.0.0.0"], 'a') + self.assertEqual(pyt["10.1.0.0"], 'b') + self.assertEqual(pyt["10.1.0.1"], 'b') + self.assertEqual(pyt["10.0.0.1"], 'a') + + self.assertTrue('10.0.0.0' in pyt) + self.assertTrue('10.1.0.0' in pyt) + self.assertTrue('10.0.0.1' in pyt) + self.assertFalse('9.0.0.0' in pyt) + self.assertFalse('0.0.0.0' in pyt) + + self.assertTrue(pyt.has_key('10.0.0.0/8')) + self.assertTrue(pyt.has_key('10.1.0.0/16')) + self.assertFalse(pyt.has_key('10.2.0.0/16')) + self.assertFalse(pyt.has_key('9.0.0.0/8')) + self.assertFalse(pyt.has_key('10.0.0.1')) + + self.assertTrue(pyt.has_key('10.0.0.0/8')) + self.assertTrue(pyt.has_key('10.1.0.0/16')) + self.assertFalse(pyt.has_key('10.2.0.0/16')) + self.assertFalse(pyt.has_key('9.0.0.0/8')) + self.assertFalse(pyt.has_key('10.0.0.0')) + + self.assertListEqual(sorted(['10.0.0.0/8','10.1.0.0/16']), sorted(pyt.keys())) + + +def basicAsserts(q, pickled_data): + #NOTE: done this way specifically to ensure shared memory not possible + try: + pyt = pickle.loads(pickled_data) + assert(pyt["10.0.0.0/8"] == 'a') + assert(pyt["10.1.0.0/16"] == 'b') + assert(pyt["10.0.0.0"] == 'a') + assert(pyt["10.1.0.0"] == 'b') + assert(pyt["10.1.0.1"] == 'b') + assert(pyt["10.0.0.1"] == 'a') + + assert('10.0.0.0' in pyt) + assert('10.1.0.0' in pyt) + assert('10.0.0.1' in pyt) + assert(not '9.0.0.0' in pyt) + assert(not '0.0.0.0' in pyt) + + assert(pyt.has_key('10.0.0.0/8')) + assert(pyt.has_key('10.1.0.0/16')) + assert(not pyt.has_key('10.2.0.0/16')) + assert(not pyt.has_key('9.0.0.0/8')) + assert(not pyt.has_key('10.0.0.1')) + + assert(pyt.has_key('10.0.0.0/8')) + assert(pyt.has_key('10.1.0.0/16')) + assert(not pyt.has_key('10.2.0.0/16')) + assert(not pyt.has_key('9.0.0.0/8')) + assert(not pyt.has_key('10.0.0.0')) + except Exception as e: + q.put(('error', str(e))) + return + q.put(('success',None)) + + if __name__ == '__main__': + # Set this way specifically for the multiprocess test above + multiprocessing.set_start_method('spawn') unittest.main()