diff --git a/src/StdStorage.sol b/src/StdStorage.sol index 3035880e..3b56eb8e 100644 --- a/src/StdStorage.sol +++ b/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; } } diff --git a/test/StdStorage.t.sol b/test/StdStorage.t.sol index ab87da38..34782ac1 100644 --- a/test/StdStorage.t.sol +++ b/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;