🐞 Fix LibBytes.uint8At stale byte (#1550)

* 🐞 Fix LibBytes.uint8At stale byte

* Added Test Case

Author: atarpara

src/utils/LibBytes.sol

@@ -129,21 +129,19 @@ library LibBytes {
         assembly {
             for { let packed := sload($.slot) } 1 {} {
                 if iszero(eq(or(packed, 0xff), packed)) {
+                    if iszero(lt(i, and(packed, 0xff))) { break }
                     if iszero(gt(i, 0x1e)) {
                         result := byte(i, packed)
                         break
                     }
-                    if iszero(gt(i, and(0xff, packed))) {
-                        mstore(0x00, $.slot)
-                        let j := sub(i, 0x1f)
-                        result := byte(and(j, 0x1f), sload(add(keccak256(0x00, 0x20), shr(5, j))))
-                    }
-                    break
-                }
-                if iszero(gt(i, shr(8, packed))) {
                     mstore(0x00, $.slot)
-                    result := byte(and(i, 0x1f), sload(add(keccak256(0x00, 0x20), shr(5, i))))
+                    let j := sub(i, 0x1f)
+                    result := byte(and(j, 0x1f), sload(add(keccak256(0x00, 0x20), shr(5, j))))
+                    break
                 }
+                if iszero(gt(shr(8, packed), i)) { break }
+                mstore(0x00, $.slot)
+                result := byte(and(i, 0x1f), sload(add(keccak256(0x00, 0x20), shr(5, i))))
                 break
             }
         }

src/utils/g/LibBytes.sol

@@ -133,21 +133,19 @@ library LibBytes {
         assembly {
             for { let packed := sload($.slot) } 1 {} {
                 if iszero(eq(or(packed, 0xff), packed)) {
+                    if iszero(lt(i, and(packed, 0xff))) { break }
                     if iszero(gt(i, 0x1e)) {
                         result := byte(i, packed)
                         break
                     }
-                    if iszero(gt(i, and(0xff, packed))) {
-                        mstore(0x00, $.slot)
-                        let j := sub(i, 0x1f)
-                        result := byte(and(j, 0x1f), sload(add(keccak256(0x00, 0x20), shr(5, j))))
-                    }
-                    break
-                }
-                if iszero(gt(i, shr(8, packed))) {
                     mstore(0x00, $.slot)
-                    result := byte(and(i, 0x1f), sload(add(keccak256(0x00, 0x20), shr(5, i))))
+                    let j := sub(i, 0x1f)
+                    result := byte(and(j, 0x1f), sload(add(keccak256(0x00, 0x20), shr(5, j))))
+                    break
                 }
+                if iszero(gt(shr(8, packed), i)) { break }
+                mstore(0x00, $.slot)
+                result := byte(and(i, 0x1f), sload(add(keccak256(0x00, 0x20), shr(5, i))))
                 break
             }
         }

test/LibBytes.t.sol

@@ -4,6 +4,22 @@ pragma solidity ^0.8.4;
 import "./utils/SoladyTest.sol";
 import {LibBytes} from "../src/utils/LibBytes.sol";
 
+contract BytesStore {
+    LibBytes.BytesStorage internal value;
+
+    function set(bytes calldata newValue) external {
+        LibBytes.setCalldata(value, newValue);
+    }
+
+    function length() external view returns (uint256) {
+        return LibBytes.length(value);
+    }
+
+    function at_(uint256 index) external view returns (uint8) {
+        return LibBytes.uint8At(value, index);
+    }
+}
+
 contract LibBytesTest is SoladyTest {
     function testLoad(bytes memory a) public {
         if (a.length < 32) a = abi.encodePacked(a, new bytes(32));
@@ -414,4 +430,55 @@ contract LibBytesTest is SoladyTest {
             require(keccak256(expectedChildren[i]) == keccak256(children[i]));
         }
     }
+
+    function testUint8AtStaleByteAfterShrink() public {
+        BytesStore store = new BytesStore();
+        bytes memory previous = new bytes(32);
+        previous[31] = 0xbb;
+
+        store.set(previous);
+        store.set(new bytes(31));
+
+        assertEq(store.length(), 31);
+        assertEq(store.at_(31), 0);
+    }
+
+    function testUint8AtPaddingAfterShortValue() public {
+        BytesStore store = _storeWithDirtyPadding(5);
+        assertEq(store.length(), 5);
+        assertEq(store.at_(4), 0x11);
+        assertEq(store.at_(5), 0);
+        assertEq(store.at_(30), 0);
+    }
+
+    function testUint8AtPaddingAfterSpilledValue() public {
+        BytesStore store = _storeWithDirtyPadding(40);
+        assertEq(store.length(), 40);
+        assertEq(store.at_(39), 0x11);
+        assertEq(store.at_(40), 0);
+    }
+
+    function testUint8AtPaddingAfterTaggedLongValue() public {
+        BytesStore store = _storeWithDirtyPadding(255);
+        assertEq(store.length(), 255);
+        assertEq(store.at_(254), 0x11);
+        assertEq(store.at_(255), 0);
+        assertEq(store.at_(256), 0);
+    }
+
+    function _storeWithDirtyPadding(uint256 n) internal returns (BytesStore store) {
+        store = new BytesStore();
+
+        bytes memory value = new bytes(n);
+        for (uint256 i; i < n; ++i) {
+            value[i] = 0x11;
+        }
+
+        bytes memory data = abi.encodeWithSelector(BytesStore.set.selector, value);
+        for (uint256 i = 4 + 32 + 32 + n; i < data.length; ++i) {
+            data[i] = 0xaa;
+        }
+        (bool success,) = address(store).call(data);
+        require(success);
+    }
 }