fix(StdStorage): find short string and bytes slots

src/StdStorage.sol

@@ -22,6 +22,12 @@ struct StdStorage {
 }
 
 library stdStorageSafe {
+    struct FindCallData {
+        bytes32 result;
+        bytes32 shortBytesStorageValue;
+        bytes32[] reads;
+    }
+
     event SlotFound(address who, bytes4 fsig, bytes32 keysHash, uint256 slot);
     event WARNING_UninitedSlot(address who, uint256 slot);
 
@@ -42,15 +48,59 @@ library stdStorageSafe {
         }
     }
 
-    /// @notice Calls the target contract with the configured parameters and returns the success flag and return value.
-    function callTarget(StdStorage storage self) internal view returns (bool, bytes32) {
+    /// @notice Calls the target contract with the configured parameters and returns its raw return data.
+    function callTargetRaw(StdStorage storage self) private view returns (bool, bytes memory) {
         bytes memory cd = abi.encodePacked(self._sig, getCallParams(self));
         (bool success, bytes memory rdat) = self._target.staticcall(cd);
+
+        return (success, rdat);
+    }
+
+    /// @notice Calls the target contract with the configured parameters and returns the success flag and return value.
+    function callTarget(StdStorage storage self) internal view returns (bool, bytes32) {
+        (bool success, bytes memory rdat) = callTargetRaw(self);
         bytes32 result = _bytesToBytes32(rdat, 32 * self._depth);
 
         return (success, result);
     }
 
+    /// @notice Returns the storage encoding when `rdat` is one non-empty short `bytes` or `string` value.
+    function _parseShortBytesReturn(bytes memory rdat) private pure returns (bytes32) {
+        if (rdat.length != 96 || uint256(_bytesToBytes32(rdat, 0)) != 32) {
+            return bytes32(0);
+        }
+
+        uint256 length = uint256(_bytesToBytes32(rdat, 32));
+        if (length == 0 || length > 31) {
+            return bytes32(0);
+        }
+
+        bytes32 value = _bytesToBytes32(rdat, 64);
+        if (uint256(value) << (length * 8) != 0) {
+            return bytes32(0);
+        }
+
+        return value | bytes32(length * 2);
+    }
+
+    /// @notice Returns whether the payload and length marker match a short `bytes` or `string` storage value.
+    function _matchesShortBytes(bytes32 slotValue, bytes32 expected) private pure returns (bool) {
+        uint256 length = uint8(uint256(expected)) / 2;
+        uint256 mask = (type(uint256).max << ((32 - length) * 8)) | 0xFF;
+        return uint256(slotValue) & mask == uint256(expected);
+    }
+
+    /// @notice Returns whether clearing `slot` makes the configured target return an empty dynamic byte array.
+    function _checkShortBytesSlot(StdStorage storage self, bytes32 slot) private returns (bool) {
+        bytes32 prevSlotValue = vm.load(self._target, slot);
+        vm.store(self._target, slot, bytes32(0));
+        (bool success, bytes memory rdat) = callTargetRaw(self);
+        vm.store(self._target, slot, prevSlotValue);
+
+        return success && rdat.length == 64 && uint256(_bytesToBytes32(rdat, 0)) == 32
+            && uint256(_bytesToBytes32(rdat, 32)) == 0;
+    }
+
     /// @notice Returns whether mutating `slot` changes the return value of the configured target call.
     /// @dev Temporarily writes either `type(uint256).max` or `0` to the slot to detect sensitivity.
     function checkSlotMutatesCall(StdStorage storage self, bytes32 slot) internal returns (bool) {
@@ -118,29 +168,40 @@ library stdStorageSafe {
             }
             return self.finds[who][fsig][keccak256(abi.encodePacked(params, field_depth))];
         }
-        vm.record();
-        (, bytes32 callResult) = callTarget(self);
-        (bytes32[] memory reads,) = vm.accesses(address(who));
+        FindCallData memory callData;
+        {
+            vm.record();
+            (bool callSuccess, bytes memory rdat) = callTargetRaw(self);
+            callData.result = _bytesToBytes32(rdat, 32 * field_depth);
+            if (callSuccess && field_depth == 0) {
+                callData.shortBytesStorageValue = _parseShortBytesReturn(rdat);
+            }
+            (callData.reads,) = vm.accesses(address(who));
+        }
 
-        if (reads.length == 0) {
+        if (callData.reads.length == 0) {
             revert("stdStorage find(StdStorage): No storage use detected for target.");
         } else {
-            for (uint256 i = reads.length; i > 0;) {
+            for (uint256 i = callData.reads.length; i > 0;) {
                 --i;
-                bytes32 prev = vm.load(who, reads[i]);
+                bytes32 slot = callData.reads[i];
+                bytes32 prev = vm.load(who, slot);
                 if (prev == bytes32(0)) {
-                    emit WARNING_UninitedSlot(who, uint256(reads[i]));
+                    emit WARNING_UninitedSlot(who, uint256(slot));
                 }
 
-                if (!checkSlotMutatesCall(self, reads[i])) {
+                bool shortBytesFound = callData.shortBytesStorageValue != bytes32(0)
+                    && _matchesShortBytes(prev, callData.shortBytesStorageValue) && _checkShortBytesSlot(self, slot);
+
+                if (!shortBytesFound && !checkSlotMutatesCall(self, slot)) {
                     continue;
                 }
 
                 (uint256 offsetLeft, uint256 offsetRight) = (0, 0);
 
-                if (self._enable_packed_slots) {
+                if (!shortBytesFound && self._enable_packed_slots) {
                     bool found;
-                    (found, offsetLeft, offsetRight) = findOffsets(self, reads[i]);
+                    (found, offsetLeft, offsetRight) = findOffsets(self, slot);
                     if (!found) {
                         continue;
                     }
@@ -149,13 +210,13 @@ library stdStorageSafe {
                 // Check that value between found offsets is equal to the current call result
                 uint256 curVal = (uint256(prev) & getMaskByOffsets(offsetLeft, offsetRight)) >> offsetRight;
 
-                if (uint256(callResult) != curVal) {
+                if (!shortBytesFound && uint256(callData.result) != curVal) {
                     continue;
                 }
 
-                emit SlotFound(who, fsig, keccak256(abi.encodePacked(params, field_depth)), uint256(reads[i]));
+                emit SlotFound(who, fsig, keccak256(abi.encodePacked(params, field_depth)), uint256(slot));
                 self.finds[who][fsig][keccak256(abi.encodePacked(params, field_depth))] =
-                    FindData(uint256(reads[i]), offsetLeft, offsetRight, true);
+                    FindData(uint256(slot), offsetLeft, offsetRight, true);
                 break;
             }
         }

test/StdStorage.t.sol

@@ -351,6 +351,39 @@ contract StdStorageTest is Test {
         assertEq(test.edgeCaseArray(0), 1);
     }
 
+    // Regression tests for https://github.com/foundry-rs/forge-std/issues/345
+    function test_StorageFindShortString() public {
+        ShortBytesStorage target = new ShortBytesStorage();
+        assertEq(stdstore.target(address(target)).sig(target.exists.selector).find(), 0);
+    }
+
+    function test_StorageFindShortBytesWithDirtyPadding() public {
+        ShortBytesStorage target = new ShortBytesStorage();
+        bytes32 slot = bytes32(uint256(1));
+        bytes32 dirtyValue = vm.load(address(target), slot) | bytes32(uint256(0xDEADBE) << 8);
+        vm.store(address(target), slot, dirtyValue);
+        assertEq(target.data(), hex"616263");
+        assertEq(stdstore.target(address(target)).sig(target.data.selector).find(), 1);
+        assertEq(vm.load(address(target), slot), dirtyValue);
+    }
+
+    function test_StorageFindCanonicalLookingStaticReturn() public {
+        CanonicalLookingStaticReturn target = new CanonicalLookingStaticReturn();
+        assertEq(stdstore.target(address(target)).sig(target.values.selector).depth(0).find(), 0);
+        assertEq(stdstore.target(address(target)).sig(target.values.selector).depth(1).find(), 1);
+    }
+
+    function test_RevertStorageFindRestoresFailedShortBytesProbe() public {
+        RevertingEmptyString target = new RevertingEmptyString();
+        vm.expectRevert(bytes("stdStorage find(StdStorage): Slot(s) not found."));
+        this.findSlot(address(target), target.value.selector);
+        assertEq(target.value(), "A");
+    }
+
+    function findSlot(address target, bytes4 sig) external {
+        stdstore.target(target).sig(sig).find();
+    }
+
     // Regression test for https://github.com/foundry-rs/forge-std/issues/740
     // `find()` used to infinite-loop on tokens whose `balanceOf` reads multiple
     // storage slots and returns a derived value (reflection tokens).
@@ -377,6 +410,30 @@ contract StorageTestTarget {
     }
 }
 
+contract ShortBytesStorage {
+    string public exists = "thequickbrownfoxjumpsoverthelaz";
+    bytes public data = hex"616263";
+}
+
+contract CanonicalLookingStaticReturn {
+    uint256 private first = 32;
+    uint256 private second = 1;
+    bytes32 private decoy = bytes32("A") | bytes32(uint256(2));
+
+    function values() public view returns (uint256, uint256, bytes32) {
+        return (first, second, decoy & ~bytes32(uint256(0xFF)));
+    }
+}
+
+contract RevertingEmptyString {
+    string private storedValue = "A";
+
+    function value() public view returns (string memory) {
+        require(bytes(storedValue).length != 0);
+        return storedValue;
+    }
+}
+
 contract ReflectionTokenTarget {
     using stdStorage for StdStorage;