From c86dcbc5c0d88d40929c4052a3e866f4a24e1fce Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Thu, 3 Oct 2024 10:52:14 -0400 Subject: trie: Add some extra bounds checking --- src/trie.c | 37 ++++++++++++++++++++++++------------- 1 file changed, 24 insertions(+), 13 deletions(-) (limited to 'src') diff --git a/src/trie.c b/src/trie.c index a92950b..ede8811 100644 --- a/src/trie.c +++ b/src/trie.c @@ -171,9 +171,10 @@ void trie_init(struct trie *trie) { } /** Extract the nibble at a certain offset from a byte sequence. */ -static unsigned char trie_key_nibble(const void *key, size_t offset) { +static unsigned char trie_key_nibble(const void *key, size_t length, size_t offset) { const unsigned char *bytes = key; size_t byte = offset >> 1; + bfs_assert(byte < length); // A branchless version of // if (offset & 1) { @@ -185,6 +186,11 @@ static unsigned char trie_key_nibble(const void *key, size_t offset) { return (bytes[byte] >> shift) & 0xF; } +/** Extract the nibble at a certain offset from a leaf. */ +static unsigned char trie_leaf_nibble(const struct trie_leaf *leaf, size_t offset) { + return trie_key_nibble(leaf->key, leaf->length, offset); +} + /** * Finds a leaf in the trie that matches the key at every branch. If the key * exists in the trie, the representative will match the searched key. But @@ -206,7 +212,7 @@ static struct trie_leaf *trie_representative(const struct trie *trie, const void unsigned int index = 0; if ((offset >> 1) < length) { - unsigned char nibble = trie_key_nibble(key, offset); + unsigned char nibble = trie_key_nibble(key, length, offset); unsigned int bit = 1U << nibble; // bits = bitmap & bit ? bitmap & (bit - 1) : 0 unsigned int mask = -!!(node->bitmap & bit); @@ -307,7 +313,7 @@ static struct trie_leaf *trie_find_prefix_impl(const struct trie *trie, const ch skip = offset >> 1; } - unsigned char nibble = trie_key_nibble(key, offset); + unsigned char nibble = trie_key_nibble(key, length, offset); unsigned int bit = 1U << nibble; if (node->bitmap & bit) { unsigned int index = count_ones(node->bitmap & (bit - 1)); @@ -494,10 +500,10 @@ static struct trie_leaf *trie_node_insert(struct trie *trie, uintptr_t *ptr, str * | Y * +--->key */ -static uintptr_t *trie_jump(struct trie *trie, uintptr_t *ptr, const char *key, size_t *offset) { +static uintptr_t *trie_jump(struct trie *trie, uintptr_t *ptr, size_t *offset) { // We only ever need to jump to leaf nodes, since internal nodes are // guaranteed to be within OFFSET_MAX anyway - bfs_assert(trie_is_leaf(*ptr)); + struct trie_leaf *leaf = trie_decode_leaf(*ptr); struct trie_node *node = trie_node_alloc(trie, 1); if (!node) { @@ -507,7 +513,7 @@ static uintptr_t *trie_jump(struct trie *trie, uintptr_t *ptr, const char *key, *offset += OFFSET_MAX; node->offset = OFFSET_MAX; - unsigned char nibble = trie_key_nibble(key, *offset); + unsigned char nibble = trie_leaf_nibble(leaf, *offset); node->bitmap = 1 << nibble; node->children[0] = *ptr; @@ -533,8 +539,8 @@ static uintptr_t *trie_jump(struct trie *trie, uintptr_t *ptr, const char *key, * +--->leaf */ static struct trie_leaf *trie_split(struct trie *trie, uintptr_t *ptr, struct trie_leaf *leaf, struct trie_leaf *rep, size_t offset, size_t mismatch) { - unsigned char key_nibble = trie_key_nibble(leaf->key, mismatch); - unsigned char rep_nibble = trie_key_nibble(rep->key, mismatch); + unsigned char key_nibble = trie_leaf_nibble(leaf, mismatch); + unsigned char rep_nibble = trie_leaf_nibble(rep, mismatch); bfs_assert(key_nibble != rep_nibble); struct trie_node *node = trie_node_alloc(trie, 2); @@ -567,8 +573,14 @@ _trie_clones static struct trie_leaf *trie_insert_mem_impl(struct trie *trie, const void *key, size_t length) { struct trie_leaf *rep = trie_representative(trie, key, length); size_t mismatch = trie_mismatch(rep, key, length); - if (mismatch >= (length << 1)) { + size_t misbyte = mismatch / 2; + if (misbyte >= length) { + bfs_assert(misbyte == length); return rep; + } else if (rep && misbyte >= rep->length) { + bfs_bug("trie keys must be prefix-free"); + errno = EINVAL; + return NULL; } struct trie_leaf *leaf = trie_leaf_alloc(trie, key, length); @@ -590,7 +602,7 @@ static struct trie_leaf *trie_insert_mem_impl(struct trie *trie, const void *key } offset += node->offset; - unsigned char nibble = trie_key_nibble(key, offset); + unsigned char nibble = trie_leaf_nibble(leaf, offset); unsigned int bit = 1U << nibble; if (node->bitmap & bit) { bfs_assert(offset < mismatch); @@ -603,7 +615,7 @@ static struct trie_leaf *trie_insert_mem_impl(struct trie *trie, const void *key } while (mismatch - offset > OFFSET_MAX) { - ptr = trie_jump(trie, ptr, key, &offset); + ptr = trie_jump(trie, ptr, &offset); if (!ptr) { trie_leaf_free(trie, leaf); return NULL; @@ -674,9 +686,8 @@ static void trie_remove_impl(struct trie *trie, struct trie_leaf *leaf) { while (!trie_is_leaf(*child)) { struct trie_node *node = trie_decode_node(*child); offset += node->offset; - bfs_assert((offset >> 1) < leaf->length); - unsigned char nibble = trie_key_nibble(leaf->key, offset); + unsigned char nibble = trie_leaf_nibble(leaf, offset); unsigned int bit = 1U << nibble; unsigned int bitmap = node->bitmap; bfs_assert(bitmap & bit); -- cgit v1.2.3