prevent set overflow

src/lib/LibMemoryKV.sol

@@ -16,6 +16,11 @@ type MemoryKVVal is bytes32;
 
 /// @title LibMemoryKV
 library LibMemoryKV {
+    /// Thrown when the memory allocation for a new key/value pair would exceed
+    /// the maximum pointer value of `0xFFFF` which would cause corruption of
+    /// the linked list and potentially overwriting of unrelated memory.
+    error MemoryKVOverflow(uint256 pointer);
+
     /// Gets the value associated with a given key.
     /// The value returned will be `0` if the key exists and was set to zero OR
     /// the key DOES NOT exist, i.e. was never set.
@@ -59,6 +64,7 @@ library LibMemoryKV {
     /// @return The final value of `kv` as it MAY be modified if the upsert
     /// resulted in an insert operation.
     function set(MemoryKV kv, MemoryKVKey key, MemoryKVVal value) internal pure returns (MemoryKV) {
+        uint256 pointer;
         assembly ("memory-safe") {
             // Hash to spread inserts across internal lists.
             // This MUST remain in sync with `get` logic.
@@ -70,7 +76,7 @@ library LibMemoryKV {
             let startPointer := and(shr(bitOffset, kv), 0xFFFF)
 
             // Find a key match then break so that we populate a nonzero pointer.
-            let pointer := startPointer
+            pointer := startPointer
             for {} iszero(iszero(pointer)) { pointer := mload(add(pointer, 0x40)) } {
                 if eq(key, mload(pointer)) { break }
             }
@@ -106,6 +112,9 @@ library LibMemoryKV {
                 )
             }
         }
+        if (pointer > 0xFFFF) {
+            revert MemoryKVOverflow(pointer);
+        }
         return kv;
     }
 

test/src/lib/LibMemoryKV.getset.t.sol

@@ -8,6 +8,23 @@ import {LibPointer, Pointer} from "rain.solmem/lib/LibPointer.sol";
 import {LibMemoryKV, MemoryKVKey, MemoryKVVal, MemoryKV} from "src/lib/LibMemoryKV.sol";
 
 contract LibMemoryKVGetSetTest is Test {
+    function setOverflowExternal(MemoryKV kv, MemoryKVKey key, MemoryKVVal value) external pure returns (MemoryKV) {
+        assembly ("memory-safe") {
+            // Set the pointer past 0xFFFF to cause an overflow on the next
+            // insert.
+            mstore(0x40, 0x10000)
+        }
+
+        return LibMemoryKV.set(kv, key, value);
+    }
+
+    function testSetOverflow(MemoryKVKey key, MemoryKVVal value) external {
+        MemoryKV kv = MemoryKV.wrap(0);
+        // The next set should revert with a MemoryKVOverflow error.
+        vm.expectRevert(abi.encodeWithSelector(LibMemoryKV.MemoryKVOverflow.selector, 0x10000));
+        this.setOverflowExternal(kv, key, value);
+    }
+
     function testSetGet0(MemoryKVKey key, MemoryKVVal value) public pure {
         MemoryKV kv = MemoryKV.wrap(0);