From bc3ec6456a0a56a1664d5cf5769a5920bdd80eae Mon Sep 17 00:00:00 2001 From: Alex Halderman Date: Tue, 1 Oct 2013 17:05:47 -0400 Subject: [PATCH] What Eric and Zakir said --- lib/blacklist.c | 15 +++++++++ lib/constraint.c | 81 +++++++++++++++++++++++++++++++++++++++++++----- lib/constraint.h | 6 ++-- 3 files changed, 93 insertions(+), 9 deletions(-) diff --git a/lib/blacklist.c b/lib/blacklist.c index 49eaa32..a486444 100644 --- a/lib/blacklist.c +++ b/lib/blacklist.c @@ -125,5 +125,20 @@ int blacklist_init_from_files(char *whitelist_filename, char *blacklist_filename uint64_t allowed = blacklist_count_allowed(); log_debug("blacklist", "%lu addresses allowed to be scanned (%0.0f%% of address space)", allowed, allowed*100./((long long int)1 << 32)); + + /* + // test + log_debug("blacklist", "testing started"); + uint64_t count = constraint_count_ips(constraint, ADDR_ALLOWED); + for (unsigned int i=0; i < count; i++) { + int ip = constraint_lookup_index(constraint, i, ADDR_ALLOWED); + if ((i & 0xFFFFFF) == 0) + log_info("blacklist", "%x", i & 0xFF000000); + if (constraint_lookup_ip(constraint, ip) != ADDR_ALLOWED) { + log_error("blacklist", "test failed for index %d", i); + } + } + log_debug("blacklist", "testing complete"); + */ return 0; } diff --git a/lib/constraint.c b/lib/constraint.c index aa8d0b9..17a5164 100644 --- a/lib/constraint.c +++ b/lib/constraint.c @@ -50,6 +50,7 @@ typedef struct node { struct node *l; struct node *r; value_t value; + uint64_t count; } node_t; // As an optimization, we precompute lookups for every prefix of this @@ -60,6 +61,8 @@ struct _constraint { node_t *root; // root node of the tree node_t **radix; // array of nodes for every RADIX_LENGTH prefix int optimized; // is radix populated and up-to-date? + int painted; // have we precomputed counts for each node? + value_t paint_value; // value for which we precomputed counts }; // Tree operations respect the invariant that every node that isn't a @@ -176,7 +179,7 @@ static int _lookup_ip(node_t *root, uint32_t address) // Return the value pertaining to an address. // (Note: address must be in host byte order.) -int constraint_lookup_ip(constraint_t *con, uint32_t address) +value_t constraint_lookup_ip(constraint_t *con, uint32_t address) { assert(con); if (con->optimized) { @@ -193,27 +196,91 @@ int constraint_lookup_ip(constraint_t *con, uint32_t address) } } +// Return the nth painted IP address. +static int _lookup_index(node_t *root, uint64_t n) +{ + assert(root); + node_t *node = root; + uint32_t ip = 0; + uint32_t mask = 0x80000000; + for (;;) { + if (IS_LEAF(node)) { + return ip | n; + } + if (n < node->l->count) { + node = node->l; + } else { + n -= node->l->count; + node = node->r; + ip |= mask; + } + mask >>= 1; + } +} + +// For a given value, return the IP address with zero-based index n. +// (i.e., if there are three addresses with value 0xFF, looking up index 1 +// will return the second one). +// Note that the tree must have been previously painted with this value. +uint32_t constraint_lookup_index(constraint_t *con, uint64_t index, value_t value) +{ + assert(con); + if (!con->painted || con->paint_value != value) { + constraint_paint_value(con, value); + } + if (con->optimized) { + // TK TK TK + } + + assert(index < con->root->count); + return _lookup_index(con->root, index); +} + + // Implement count_ips by recursing on halves of the tree. Size represents // the number of addresses in a prefix at the current level of the tree. -static uint64_t _count_ips_recurse(node_t *node, value_t value, uint64_t size) +// If paint is specified, each node will have its count set to the number of +// leaves under it set to value. +static uint64_t _count_ips_recurse(node_t *node, value_t value, uint64_t size, int paint) { assert(node); + uint64_t n; if (IS_LEAF(node)) { if (node->value == value) { - return size; + n = size; } else { - return 0; + n = 0; } + } else { + n = _count_ips_recurse(node->l, value, size >> 1, paint) + + _count_ips_recurse(node->r, value, size >> 1, paint); } - return _count_ips_recurse(node->l, value, size >> 1) + - _count_ips_recurse(node->r, value, size >> 1); + if (paint) { + node->count = n; + } + return n; +} + +// For each node, precompute the count of leaves beneath it set to value. +// Note that the tree can be painted for only one value at a time. +void constraint_paint_value(constraint_t *con, value_t value) +{ + assert(con); + log_info("constraint", "Painting value %lu", value); + _count_ips_recurse(con->root, value, (uint64_t)1 << 32, 1); + con->painted = 1; + con->paint_value = value; } // Return the number of addresses that have a given value. uint64_t constraint_count_ips(constraint_t *con, value_t value) { assert(con); - return _count_ips_recurse(con->root, value, (uint64_t)1 << 32); + if (con->painted && con->paint_value == value) { + return con->root->count; + } else { + return _count_ips_recurse(con->root, value, (uint64_t)1 << 32, 0); + } } // Initialize the tree. diff --git a/lib/constraint.h b/lib/constraint.h index 3e914a3..2e4f28b 100644 --- a/lib/constraint.h +++ b/lib/constraint.h @@ -2,13 +2,15 @@ #define _CONSTRAINT_H typedef struct _constraint constraint_t; -typedef int value_t; +typedef unsigned int value_t; constraint_t* constraint_init(value_t value); void constraint_free(constraint_t *con); void constraint_set(constraint_t *con, uint32_t prefix, int len, value_t value); void constraint_optimize(constraint_t *con); -int constraint_lookup_ip(constraint_t *con, uint32_t address); +value_t constraint_lookup_ip(constraint_t *con, uint32_t address); uint64_t constraint_count_ips(constraint_t *con, value_t value); +uint32_t constraint_lookup_index(constraint_t *con, uint64_t index, value_t value); +void constraint_paint_value(constraint_t *con, value_t value); #endif //_CONSTRAINT_H