From 981ca639c8a03ee6ce44dc2a3076b8ed267ae280 Mon Sep 17 00:00:00 2001
From: Robsdedude <rouven.bauer@neo4j.com>
Date: Wed, 13 Aug 2025 16:50:37 +0200
Subject: [PATCH] Fix decoding of map keys of certain sizes.

---
 changelog.d/59.fix.md                   |  1 +
 src/v1/unpack.rs                        |  4 +--
 tests/v1/from_driver/test_packstream.py | 40 +++++++++++++++++++++++++
 3 files changed, 43 insertions(+), 2 deletions(-)
 create mode 100644 changelog.d/59.fix.md

diff --git a/changelog.d/59.fix.md b/changelog.d/59.fix.md
new file mode 100644
index 0000000..1d97248
--- /dev/null
+++ b/changelog.d/59.fix.md
@@ -0,0 +1 @@
+Fix decoding of map keys of certain sizes<ISSUES_LIST>.
diff --git a/src/v1/unpack.rs b/src/v1/unpack.rs
index a258ed9..4cbf3d6 100644
--- a/src/v1/unpack.rs
+++ b/src/v1/unpack.rs
@@ -236,8 +236,8 @@ impl<'a> PackStreamDecoder<'a> {
     fn read_string_length(&mut self) -> PyResult<usize> {
         let marker = self.read_byte()?;
         let high_nibble = marker & 0xF0;
-        match high_nibble {
-            TINY_STRING => Ok((marker & 0x0F) as usize),
+        match marker {
+            _ if high_nibble == TINY_STRING => Ok((marker & 0x0F).into()),
             STRING_8 => self.read_u8(),
             STRING_16 => self.read_u16(),
             STRING_32 => self.read_u32(),
diff --git a/tests/v1/from_driver/test_packstream.py b/tests/v1/from_driver/test_packstream.py
index ad5c169..324b2c0 100644
--- a/tests/v1/from_driver/test_packstream.py
+++ b/tests/v1/from_driver/test_packstream.py
@@ -564,11 +564,27 @@ def test_tiny_maps(self, assert_packable, size):
             data_out += bytearray([0x81, 64 + el, el])
         assert_packable(data_in, bytes(data_out))
 
+    @pytest.mark.parametrize("size", range(0x10))
+    def test_tiny_maps_padded_key(self, assert_packable, size):
+        data_in = {}
+        data_out = bytearray([0xA0 + size])
+        padding = b"1234567890abcdefghijklmnopqrstuvwxyz"
+        for el in range(1, size + 1):
+            data_in[padding.decode("ascii") + chr(64 + el)] = el
+            data_out += bytearray([0xD0, 37, *padding, 64 + el, el])
+        assert_packable(data_in, bytes(data_out))
+
     def test_map_8(self, pack, assert_packable):
         d = {f"A{i}": 1 for i in range(40)}
         b = b"".join(pack(f"A{i}", 1) for i in range(40))
         assert_packable(d, b"\xd8\x28" + b)
 
+    def test_map_8_padded_key(self, pack, assert_packable):
+        padding = "1234567890abcdefghijklmnopqrstuvwxyz"
+        d = {f"{padding}-{i}": 1 for i in range(40)}
+        b = b"".join(pack(f"{padding}-{i}", 1) for i in range(40))
+        assert_packable(d, b"\xd8\x28" + b)
+
     def test_map_16(self, pack, assert_packable):
         d = {f"A{i}": 1 for i in range(40000)}
         b = b"".join(pack(f"A{i}", 1) for i in range(40000))
@@ -579,6 +595,30 @@ def test_map_32(self, pack, assert_packable):
         b = b"".join(pack(f"A{i}", 1) for i in range(80000))
         assert_packable(d, b"\xda\x00\x01\x38\x80" + b)
 
+    def test_map_key_tiny_string(self, assert_packable):
+        key = "A"
+        d = {key: 1}
+        data_out = b"\xa1\x81" + key.encode("utf-8") + b"\x01"
+        assert_packable(d, bytes(data_out))
+
+    def test_map_key_string_8(self, assert_packable):
+        key = "A" * 40
+        d = {key: 1}
+        data_out = b"\xa1\xd0\x28" + key.encode("utf-8") + b"\x01"
+        assert_packable(d, data_out)
+
+    def test_map_key_string_16(self, assert_packable):
+        key = "A" * 40000
+        d = {key: 1}
+        data_out = b"\xa1\xd1\x9c\x40" + key.encode("utf-8") + b"\x01"
+        assert_packable(d, data_out)
+
+    def test_map_key_string_32(self, assert_packable):
+        key = "A" * 80000
+        d = {key: 1}
+        data_out = b"\xa1\xd2\x00\x01\x38\x80" + key.encode("utf-8") + b"\x01"
+        assert_packable(d, data_out)
+
     def test_empty_dataframe_maps(self, assert_packable):
         df = pd.DataFrame()
         assert_packable(df, b"\xa0", {})