[PATCH] assoc_array: preserve shortcut prefixes when splitting

[Kyle Zeng <kylebot@openai.com>]

assoc_array_insert_mid_shortcut() advances past a parent node's key
segment when deciding whether a pre-shortcut is needed.  A root shortcut
has no such parent.  If the first differing nibble is at bit 4, the code
therefore installs a node indexed at level 4 directly at the root, where
assoc_array_walk() starts at level 0.  A successfully inserted object can
then be unreachable by lookup.

For keyrings, a user can turn these lookup failures into an oracle for
the pointer-derived index hash and recover the KASLR slide of a known
kernel image.

Only advance past the parent segment when the old shortcut actually has
a parent.  This retains the prefix before a root split at bit 4 while
leaving the non-root and bit-0 cases unchanged.

Also preserve the last complete prefix word when the split is
word-aligned.  The unconditional zero-bit shift currently clears that
word and can make both sides of the split unsearchable.  Use the same
partial-word guard as the terminal-node shortcut constructor.

Fixes: 3cb989501c26 ("Add a generic associative array implementation.")
Assisted-by: Codex:gpt-6-astra
Signed-off-by: Kyle Zeng <kylebot@openai.com>
---
 lib/assoc_array.c | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/lib/assoc_array.c b/lib/assoc_array.c
index b6c9723e12ce..ee4f3994f494 100644
--- a/lib/assoc_array.c
+++ b/lib/assoc_array.c
@@ -841,9 +841,10 @@ static bool assoc_array_insert_mid_shortcut(struct assoc_array_edit *edit,
 
 	/* Insert a new shortcut before the new node if this segment isn't of
 	 * zero length - otherwise we just connect the new node directly to the
-	 * parent.
+	 * parent.  A root shortcut has no parent consuming the first level.
 	 */
-	level += ASSOC_ARRAY_LEVEL_STEP;
+	if (shortcut->back_pointer)
+		level += ASSOC_ARRAY_LEVEL_STEP;
 	if (diff > level) {
 		pr_devel("pre-shortcut %d...%d\n", level, diff);
 		keylen = round_up(diff, ASSOC_ARRAY_KEY_CHUNK_SIZE);
@@ -865,9 +866,11 @@ static bool assoc_array_insert_mid_shortcut(struct assoc_array_edit *edit,
 		memcpy(new_s0->index_key, shortcut->index_key,
 		       flex_array_size(new_s0, index_key, keylen));
 
-		blank = ULONG_MAX << (diff & ASSOC_ARRAY_KEY_CHUNK_MASK);
-		pr_devel("blank off [%zu] %d: %lx\n", keylen - 1, diff, blank);
-		new_s0->index_key[keylen - 1] &= ~blank;
+		if (diff & ASSOC_ARRAY_KEY_CHUNK_MASK) {
+			blank = ULONG_MAX << (diff & ASSOC_ARRAY_KEY_CHUNK_MASK);
+			pr_devel("blank off [%zu] %d: %lx\n", keylen - 1, diff, blank);
+			new_s0->index_key[keylen - 1] &= ~blank;
+		}
 	} else {
 		pr_devel("no pre-shortcut\n");
 		edit->set[0].to = assoc_array_node_to_ptr(new_n0);