// SPDX-License-Identifier: MIT pragma solidity >=0.6.2 <0.9.0; import {Vm} from "./Vm.sol"; struct FindData { uint256 slot; uint256 offsetLeft; uint256 offsetRight; bool found; } struct StdStorage { mapping(address => mapping(bytes4 => mapping(bytes32 => FindData))) finds; bytes32[] _keys; bytes4 _sig; uint256 _depth; address _target; bytes32 _set; bool _enable_packed_slots; bytes _calldata; } library stdStorageSafe { event SlotFound(address who, bytes4 fsig, bytes32 keysHash, uint256 slot); event WARNING_UninitedSlot(address who, uint256 slot); Vm private constant vm = Vm(address(uint160(uint256(keccak256("hevm cheat code"))))); uint256 constant UINT256_MAX = 115792089237316195423570985008687907853269984665640564039457584007913129639935; function sigs(string memory sigStr) internal pure returns (bytes4) { return bytes4(keccak256(bytes(sigStr))); } function getCallParams(StdStorage storage self) internal view returns (bytes memory) { if (self._calldata.length == 0) { return flatten(self._keys); } else { return self._calldata; } } // Calls target contract with configured parameters function callTarget(StdStorage storage self) internal view returns (bool, bytes32) { bytes memory cald = abi.encodePacked(self._sig, getCallParams(self)); (bool success, bytes memory rdat) = self._target.staticcall(cald); bytes32 result = bytesToBytes32(rdat, 32 * self._depth); return (success, result); } // Tries mutating slot value to determine if the targeted value is stored in it. // If current value is 0, then we are setting slot value to type(uint256).max // Otherwise, we set it to 0. That way, return value should always be affected. function checkSlotMutatesCall(StdStorage storage self, bytes32 slot) internal returns (bool) { bytes32 prevSlotValue = vm.load(self._target, slot); (bool success, bytes32 prevReturnValue) = callTarget(self); bytes32 testVal = prevReturnValue == bytes32(0) ? bytes32(UINT256_MAX) : bytes32(0); vm.store(self._target, slot, testVal); (, bytes32 newReturnValue) = callTarget(self); vm.store(self._target, slot, prevSlotValue); return (success && (prevReturnValue != newReturnValue)); } // Tries setting one of the bits in slot to 1 until return value changes. // Index of resulted bit is an offset packed slot has from left/right side function findOffset(StdStorage storage self, bytes32 slot, bool left) internal returns (bool, uint256) { for (uint256 offset = 0; offset < 256; offset++) { uint256 valueToPut = left ? (1 << (255 - offset)) : (1 << offset); vm.store(self._target, slot, bytes32(valueToPut)); (bool success, bytes32 data) = callTarget(self); if (success && (uint256(data) > 0)) { return (true, offset); } } return (false, 0); } function findOffsets(StdStorage storage self, bytes32 slot) internal returns (bool, uint256, uint256) { bytes32 prevSlotValue = vm.load(self._target, slot); (bool foundLeft, uint256 offsetLeft) = findOffset(self, slot, true); (bool foundRight, uint256 offsetRight) = findOffset(self, slot, false); // `findOffset` may mutate slot value, so we are setting it to initial value vm.store(self._target, slot, prevSlotValue); return (foundLeft && foundRight, offsetLeft, offsetRight); } function find(StdStorage storage self) internal returns (FindData storage) { return find(self, true); } /// @notice find an arbitrary storage slot given a function sig, input data, address of the contract and a value to check against // slot complexity: // if flat, will be bytes32(uint256(uint)); // if map, will be keccak256(abi.encode(key, uint(slot))); // if deep map, will be keccak256(abi.encode(key1, keccak256(abi.encode(key0, uint(slot))))); // if map struct, will be bytes32(uint256(keccak256(abi.encode(key1, keccak256(abi.encode(key0, uint(slot)))))) + structFieldDepth); function find(StdStorage storage self, bool _clear) internal returns (FindData storage) { address who = self._target; bytes4 fsig = self._sig; uint256 field_depth = self._depth; bytes memory params = getCallParams(self); // calldata to test against if (self.finds[who][fsig][keccak256(abi.encodePacked(params, field_depth))].found) { if (_clear) { clear(self); } return self.finds[who][fsig][keccak256(abi.encodePacked(params, field_depth))]; } vm.record(); (, bytes32 callResult) = callTarget(self); (bytes32[] memory reads,) = vm.accesses(address(who)); if (reads.length == 0) { revert("stdStorage find(StdStorage): No storage use detected for target."); } else { for (uint256 i = reads.length; --i >= 0;) { bytes32 prev = vm.load(who, reads[i]); if (prev == bytes32(0)) { emit WARNING_UninitedSlot(who, uint256(reads[i])); } if (!checkSlotMutatesCall(self, reads[i])) { continue; } (uint256 offsetLeft, uint256 offsetRight) = (0, 0); if (self._enable_packed_slots) { bool found; (found, offsetLeft, offsetRight) = findOffsets(self, reads[i]); if (!found) { continue; } } // 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) { continue; } emit SlotFound(who, fsig, keccak256(abi.encodePacked(params, field_depth)), uint256(reads[i])); self.finds[who][fsig][keccak256(abi.encodePacked(params, field_depth))] = FindData(uint256(reads[i]), offsetLeft, offsetRight, true); break; } } require( self.finds[who][fsig][keccak256(abi.encodePacked(params, field_depth))].found, "stdStorage find(StdStorage): Slot(s) not found." ); if (_clear) { clear(self); } return self.finds[who][fsig][keccak256(abi.encodePacked(params, field_depth))]; } function target(StdStorage storage self, address _target) internal returns (StdStorage storage) { self._target = _target; return self; } function sig(StdStorage storage self, bytes4 _sig) internal returns (StdStorage storage) { self._sig = _sig; return self; } function sig(StdStorage storage self, string memory _sig) internal returns (StdStorage storage) { self._sig = sigs(_sig); return self; } function with_calldata(StdStorage storage self, bytes memory _calldata) internal returns (StdStorage storage) { self._calldata = _calldata; return self; } function with_key(StdStorage storage self, address who) internal returns (StdStorage storage) { self._keys.push(bytes32(uint256(uint160(who)))); return self; } function with_key(StdStorage storage self, uint256 amt) internal returns (StdStorage storage) { self._keys.push(bytes32(amt)); return self; } function with_key(StdStorage storage self, bytes32 key) internal returns (StdStorage storage) { self._keys.push(key); return self; } function enable_packed_slots(StdStorage storage self) internal returns (StdStorage storage) { self._enable_packed_slots = true; return self; } function depth(StdStorage storage self, uint256 _depth) internal returns (StdStorage storage) { self._depth = _depth; return self; } function read(StdStorage storage self) private returns (bytes memory) { FindData storage data = find(self, false); uint256 mask = getMaskByOffsets(data.offsetLeft, data.offsetRight); uint256 value = (uint256(vm.load(self._target, bytes32(data.slot))) & mask) >> data.offsetRight; clear(self); return abi.encode(value); } function read_bytes32(StdStorage storage self) internal returns (bytes32) { return abi.decode(read(self), (bytes32)); } function read_bool(StdStorage storage self) internal returns (bool) { int256 v = read_int(self); if (v == 0) return false; if (v == 1) return true; revert("stdStorage read_bool(StdStorage): Cannot decode. Make sure you are reading a bool."); } function read_address(StdStorage storage self) internal returns (address) { return abi.decode(read(self), (address)); } function read_uint(StdStorage storage self) internal returns (uint256) { return abi.decode(read(self), (uint256)); } function read_int(StdStorage storage self) internal returns (int256) { return abi.decode(read(self), (int256)); } function parent(StdStorage storage self) internal returns (uint256, bytes32) { address who = self._target; uint256 field_depth = self._depth; vm.startMappingRecording(); uint256 child = find(self, true).slot - field_depth; (bool found, bytes32 key, bytes32 parent_slot) = vm.getMappingKeyAndParentOf(who, bytes32(child)); if (!found) { revert( "stdStorage read_bool(StdStorage): Cannot find parent. Make sure you give a slot and startMappingRecording() has been called." ); } return (uint256(parent_slot), key); } function root(StdStorage storage self) internal returns (uint256) { address who = self._target; uint256 field_depth = self._depth; vm.startMappingRecording(); uint256 child = find(self, true).slot - field_depth; bool found; bytes32 root_slot; bytes32 parent_slot; (found,, parent_slot) = vm.getMappingKeyAndParentOf(who, bytes32(child)); if (!found) { revert( "stdStorage read_bool(StdStorage): Cannot find parent. Make sure you give a slot and startMappingRecording() has been called." ); } while (found) { root_slot = parent_slot; (found,, parent_slot) = vm.getMappingKeyAndParentOf(who, bytes32(root_slot)); } return uint256(root_slot); } function bytesToBytes32(bytes memory b, uint256 offset) private pure returns (bytes32) { bytes32 out; uint256 max = b.length > 32 ? 32 : b.length; for (uint256 i = 0; i < max; i++) { out |= bytes32(b[offset + i] & 0xFF) >> (i * 8); } return out; } function flatten(bytes32[] memory b) private pure returns (bytes memory) { bytes memory result = new bytes(b.length * 32); for (uint256 i = 0; i < b.length; i++) { bytes32 k = b[i]; /// @solidity memory-safe-assembly assembly { mstore(add(result, add(32, mul(32, i))), k) } } return result; } function clear(StdStorage storage self) internal { delete self._target; delete self._sig; delete self._keys; delete self._depth; delete self._enable_packed_slots; delete self._calldata; } // Returns mask which contains non-zero bits for values between `offsetLeft` and `offsetRight` // (slotValue & mask) >> offsetRight will be the value of the given packed variable function getMaskByOffsets(uint256 offsetLeft, uint256 offsetRight) internal pure returns (uint256 mask) { // mask = ((1 << (256 - (offsetRight + offsetLeft))) - 1) << offsetRight; // using assembly because (1 << 256) causes overflow assembly { mask := shl(offsetRight, sub(shl(sub(256, add(offsetRight, offsetLeft)), 1), 1)) } } // Returns slot value with updated packed variable. function getUpdatedSlotValue(bytes32 curValue, uint256 varValue, uint256 offsetLeft, uint256 offsetRight) internal pure returns (bytes32 newValue) { return bytes32((uint256(curValue) & ~getMaskByOffsets(offsetLeft, offsetRight)) | (varValue << offsetRight)); } } library stdStorage { Vm private constant vm = Vm(address(uint160(uint256(keccak256("hevm cheat code"))))); function sigs(string memory sigStr) internal pure returns (bytes4) { return stdStorageSafe.sigs(sigStr); } function find(StdStorage storage self) internal returns (uint256) { return find(self, true); } function find(StdStorage storage self, bool _clear) internal returns (uint256) { return stdStorageSafe.find(self, _clear).slot; } function target(StdStorage storage self, address _target) internal returns (StdStorage storage) { return stdStorageSafe.target(self, _target); } function sig(StdStorage storage self, bytes4 _sig) internal returns (StdStorage storage) { return stdStorageSafe.sig(self, _sig); } function sig(StdStorage storage self, string memory _sig) internal returns (StdStorage storage) { return stdStorageSafe.sig(self, _sig); } function with_key(StdStorage storage self, address who) internal returns (StdStorage storage) { return stdStorageSafe.with_key(self, who); } function with_key(StdStorage storage self, uint256 amt) internal returns (StdStorage storage) { return stdStorageSafe.with_key(self, amt); } function with_key(StdStorage storage self, bytes32 key) internal returns (StdStorage storage) { return stdStorageSafe.with_key(self, key); } function with_calldata(StdStorage storage self, bytes memory _calldata) internal returns (StdStorage storage) { return stdStorageSafe.with_calldata(self, _calldata); } function enable_packed_slots(StdStorage storage self) internal returns (StdStorage storage) { return stdStorageSafe.enable_packed_slots(self); } function depth(StdStorage storage self, uint256 _depth) internal returns (StdStorage storage) { return stdStorageSafe.depth(self, _depth); } function clear(StdStorage storage self) internal { stdStorageSafe.clear(self); } function checked_write(StdStorage storage self, address who) internal { checked_write(self, bytes32(uint256(uint160(who)))); } function checked_write(StdStorage storage self, uint256 amt) internal { checked_write(self, bytes32(amt)); } function checked_write_int(StdStorage storage self, int256 val) internal { checked_write(self, bytes32(uint256(val))); } function checked_write(StdStorage storage self, bool write) internal { bytes32 t; /// @solidity memory-safe-assembly assembly { t := write } checked_write(self, t); } function checked_write(StdStorage storage self, bytes32 set) internal { address who = self._target; bytes4 fsig = self._sig; uint256 field_depth = self._depth; bytes memory params = stdStorageSafe.getCallParams(self); if (!self.finds[who][fsig][keccak256(abi.encodePacked(params, field_depth))].found) { find(self, false); } FindData storage data = self.finds[who][fsig][keccak256(abi.encodePacked(params, field_depth))]; if ((data.offsetLeft + data.offsetRight) > 0) { uint256 maxVal = 2 ** (256 - (data.offsetLeft + data.offsetRight)); require( uint256(set) < maxVal, string( abi.encodePacked( "stdStorage find(StdStorage): Packed slot. We can't fit value greater than ", vm.toString(maxVal) ) ) ); } bytes32 curVal = vm.load(who, bytes32(data.slot)); bytes32 valToSet = stdStorageSafe.getUpdatedSlotValue(curVal, uint256(set), data.offsetLeft, data.offsetRight); vm.store(who, bytes32(data.slot), valToSet); (bool success, bytes32 callResult) = stdStorageSafe.callTarget(self); if (!success || callResult != set) { vm.store(who, bytes32(data.slot), curVal); revert("stdStorage find(StdStorage): Failed to write value."); } clear(self); } function read_bytes32(StdStorage storage self) internal returns (bytes32) { return stdStorageSafe.read_bytes32(self); } function read_bool(StdStorage storage self) internal returns (bool) { return stdStorageSafe.read_bool(self); } function read_address(StdStorage storage self) internal returns (address) { return stdStorageSafe.read_address(self); } function read_uint(StdStorage storage self) internal returns (uint256) { return stdStorageSafe.read_uint(self); } function read_int(StdStorage storage self) internal returns (int256) { return stdStorageSafe.read_int(self); } function parent(StdStorage storage self) internal returns (uint256, bytes32) { return stdStorageSafe.parent(self); } function root(StdStorage storage self) internal returns (uint256) { return stdStorageSafe.root(self); } }