Add string and bytes support to the StorageSlots library (#4008)

Co-authored-by: Francisco <[email protected]>
OpenZeppelin · Feb 1, 2023 · 91e8d0b · 91e8d0b
1 parent a70ee4e
commit 91e8d0b
Show file tree

Hide file tree

Showing 7 changed files with 286 additions and 2 deletions.
diff --git a/.changeset/modern-games-exist.md b/.changeset/modern-games-exist.md
@@ -0,0 +1,5 @@
+---
+'openzeppelin-solidity': minor
+---
+
+`StorageSlot`: Add support for `string` and `bytes`.
diff --git a/contracts/mocks/StorageSlotMock.sol b/contracts/mocks/StorageSlotMock.sol
@@ -5,7 +5,7 @@ pragma solidity ^0.8.0;
 import "../utils/StorageSlot.sol";
 
 contract StorageSlotMock {
-    using StorageSlot for bytes32;
+    using StorageSlot for *;
 
     function setBoolean(bytes32 slot, bool value) public {
         slot.getBooleanSlot().value = value;
@@ -38,4 +38,40 @@ contract StorageSlotMock {
     function getUint256(bytes32 slot) public view returns (uint256) {
         return slot.getUint256Slot().value;
     }
+
+    mapping(uint256 => string) public stringMap;
+
+    function setString(bytes32 slot, string calldata value) public {
+        slot.getStringSlot().value = value;
+    }
+
+    function setStringStorage(uint256 key, string calldata value) public {
+        stringMap[key].getStringSlot().value = value;
+    }
+
+    function getString(bytes32 slot) public view returns (string memory) {
+        return slot.getStringSlot().value;
+    }
+
+    function getStringStorage(uint256 key) public view returns (string memory) {
+        return stringMap[key].getStringSlot().value;
+    }
+
+    mapping(uint256 => bytes) public bytesMap;
+
+    function setBytes(bytes32 slot, bytes calldata value) public {
+        slot.getBytesSlot().value = value;
+    }
+
+    function setBytesStorage(uint256 key, bytes calldata value) public {
+        bytesMap[key].getBytesSlot().value = value;
+    }
+
+    function getBytes(bytes32 slot) public view returns (bytes memory) {
+        return slot.getBytesSlot().value;
+    }
+
+    function getBytesStorage(uint256 key) public view returns (bytes memory) {
+        return bytesMap[key].getBytesSlot().value;
+    }
 }
diff --git a/contracts/utils/StorageSlot.sol b/contracts/utils/StorageSlot.sol
@@ -1,5 +1,6 @@
 // SPDX-License-Identifier: MIT
 // OpenZeppelin Contracts (last updated v4.7.0) (utils/StorageSlot.sol)
+// This file was procedurally generated from scripts/generate/templates/StorageSlot.js.
 
 pragma solidity ^0.8.0;
 
@@ -27,7 +28,8 @@ pragma solidity ^0.8.0;
  * }
  * ```
  *
- * _Available since v4.1 for `address`, `bool`, `bytes32`, and `uint256`._
+ * _Available since v4.1 for `address`, `bool`, `bytes32`, `uint256`._
+ * _Available since v4.9 for `string`, `bytes`._
  */
 library StorageSlot {
     struct AddressSlot {
@@ -46,6 +48,14 @@ library StorageSlot {
         uint256 value;
     }
 
+    struct StringSlot {
+        string value;
+    }
+
+    struct BytesSlot {
+        bytes value;
+    }
+
     /**
      * @dev Returns an `AddressSlot` with member `value` located at `slot`.
      */
@@ -85,4 +95,44 @@ library StorageSlot {
             r.slot := slot
         }
     }
+
+    /**
+     * @dev Returns an `StringSlot` with member `value` located at `slot`.
+     */
+    function getStringSlot(bytes32 slot) internal pure returns (StringSlot storage r) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            r.slot := slot
+        }
+    }
+
+    /**
+     * @dev Returns an `StringSlot` representation of the string storage pointer `store`.
+     */
+    function getStringSlot(string storage store) internal pure returns (StringSlot storage r) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            r.slot := store.slot
+        }
+    }
+
+    /**
+     * @dev Returns an `BytesSlot` with member `value` located at `slot`.
+     */
+    function getBytesSlot(bytes32 slot) internal pure returns (BytesSlot storage r) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            r.slot := slot
+        }
+    }
+
+    /**
+     * @dev Returns an `BytesSlot` representation of the bytes storage pointer `store`.
+     */
+    function getBytesSlot(bytes storage store) internal pure returns (BytesSlot storage r) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            r.slot := store.slot
+        }
+    }
 }
diff --git a/scripts/generate/run.js b/scripts/generate/run.js
@@ -18,6 +18,7 @@ for (const [file, template] of Object.entries({
   'utils/structs/EnumerableSet.sol': './templates/EnumerableSet.js',
   'utils/structs/EnumerableMap.sol': './templates/EnumerableMap.js',
   'utils/Checkpoints.sol': './templates/Checkpoints.js',
+  'utils/StorageSlot.sol': './templates/StorageSlot.js',
 })) {
   const script = path.relative(path.join(__dirname, '../..'), __filename);
   const input = path.join(path.dirname(script), template);

diff --git a/scripts/generate/templates/StorageSlot.js b/scripts/generate/templates/StorageSlot.js
@@ -0,0 +1,87 @@
+const format = require('../format-lines');
+const { capitalize, unique } = require('../../helpers');
+
+const TYPES = [
+  { type: 'address', isValueType: true, version: '4.1' },
+  { type: 'bool', isValueType: true, name: 'Boolean', version: '4.1' },
+  { type: 'bytes32', isValueType: true, version: '4.1' },
+  { type: 'uint256', isValueType: true, version: '4.1' },
+  { type: 'string', isValueType: false, version: '4.9' },
+  { type: 'bytes', isValueType: false, version: '4.9' },
+].map(type => Object.assign(type, { struct: (type.name ?? capitalize(type.type)) + 'Slot' }));
+
+const VERSIONS = unique(TYPES.map(t => t.version)).map(
+  version =>
+    `_Available since v${version} for ${TYPES.filter(t => t.version == version)
+      .map(t => `\`${t.type}\``)
+      .join(', ')}._`,
+);
+
+const header = `\
+pragma solidity ^0.8.0;
+
+/**
+ * @dev Library for reading and writing primitive types to specific storage slots.
+ *
+ * Storage slots are often used to avoid storage conflict when dealing with upgradeable contracts.
+ * This library helps with reading and writing to such slots without the need for inline assembly.
+ *
+ * The functions in this library return Slot structs that contain a \`value\` member that can be used to read or write.
+ *
+ * Example usage to set ERC1967 implementation slot:
+ * \`\`\`solidity
+ * contract ERC1967 {
+ *     bytes32 internal constant _IMPLEMENTATION_SLOT = 0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc;
+ *
+ *     function _getImplementation() internal view returns (address) {
+ *         return StorageSlot.getAddressSlot(_IMPLEMENTATION_SLOT).value;
+ *     }
+ *
+ *     function _setImplementation(address newImplementation) internal {
+ *         require(Address.isContract(newImplementation), "ERC1967: new implementation is not a contract");
+ *         StorageSlot.getAddressSlot(_IMPLEMENTATION_SLOT).value = newImplementation;
+ *     }
+ * }
+ * \`\`\`
+ *
+${VERSIONS.map(s => ` * ${s}`).join('\n')}
+ */
+`;
+
+const struct = type => `\
+struct ${type.struct} {
+  ${type.type} value;
+}
+`;
+
+const get = type => `\
+/**
+ * @dev Returns an \`${type.struct}\` with member \`value\` located at \`slot\`.
+ */
+function get${type.struct}(bytes32 slot) internal pure returns (${type.struct} storage r) {
+  /// @solidity memory-safe-assembly
+  assembly {
+      r.slot := slot
+  }
+}
+`;
+
+const getStorage = type => `\
+/**
+ * @dev Returns an \`${type.struct}\` representation of the ${type.type} storage pointer \`store\`.
+ */
+function get${type.struct}(${type.type} storage store) internal pure returns (${type.struct} storage r) {
+  /// @solidity memory-safe-assembly
+  assembly {
+      r.slot := store.slot
+  }
+}
+`;
+
+// GENERATE
+module.exports = format(
+  header.trimEnd(),
+  'library StorageSlot {',
+  [...TYPES.map(struct), ...TYPES.flatMap(type => [get(type), type.isValueType ? '' : getStorage(type)])],
+  '}',
+);
diff --git a/scripts/helpers.js b/scripts/helpers.js
@@ -24,9 +24,14 @@ function zip(...args) {
     .map((_, i) => args.map(arg => arg[i]));
 }
 
+function capitalize(str) {
+  return str.charAt(0).toUpperCase() + str.slice(1);
+}
+
 module.exports = {
   chunk,
   range,
   unique,
   zip,
+  capitalize,
 };
diff --git a/test/utils/StorageSlot.test.js b/test/utils/StorageSlot.test.js
@@ -107,4 +107,104 @@ contract('StorageSlot', function (accounts) {
       });
     });
   });
+
+  describe('string storage slot', function () {
+    beforeEach(async function () {
+      this.value = 'lorem ipsum';
+    });
+
+    it('set', async function () {
+      await this.store.setString(slot, this.value);
+    });
+
+    describe('get', function () {
+      beforeEach(async function () {
+        await this.store.setString(slot, this.value);
+      });
+
+      it('from right slot', async function () {
+        expect(await this.store.getString(slot)).to.be.equal(this.value);
+      });
+
+      it('from other slot', async function () {
+        expect(await this.store.getString(otherSlot)).to.be.equal('');
+      });
+    });
+  });
+
+  describe('string storage pointer', function () {
+    beforeEach(async function () {
+      this.value = 'lorem ipsum';
+    });
+
+    it('set', async function () {
+      await this.store.setStringStorage(slot, this.value);
+    });
+
+    describe('get', function () {
+      beforeEach(async function () {
+        await this.store.setStringStorage(slot, this.value);
+      });
+
+      it('from right slot', async function () {
+        expect(await this.store.stringMap(slot)).to.be.equal(this.value);
+        expect(await this.store.getStringStorage(slot)).to.be.equal(this.value);
+      });
+
+      it('from other slot', async function () {
+        expect(await this.store.stringMap(otherSlot)).to.be.equal('');
+        expect(await this.store.getStringStorage(otherSlot)).to.be.equal('');
+      });
+    });
+  });
+
+  describe('bytes storage slot', function () {
+    beforeEach(async function () {
+      this.value = web3.utils.randomHex(128);
+    });
+
+    it('set', async function () {
+      await this.store.setBytes(slot, this.value);
+    });
+
+    describe('get', function () {
+      beforeEach(async function () {
+        await this.store.setBytes(slot, this.value);
+      });
+
+      it('from right slot', async function () {
+        expect(await this.store.getBytes(slot)).to.be.equal(this.value);
+      });
+
+      it('from other slot', async function () {
+        expect(await this.store.getBytes(otherSlot)).to.be.equal(null);
+      });
+    });
+  });
+
+  describe('bytes storage pointer', function () {
+    beforeEach(async function () {
+      this.value = web3.utils.randomHex(128);
+    });
+
+    it('set', async function () {
+      await this.store.setBytesStorage(slot, this.value);
+    });
+
+    describe('get', function () {
+      beforeEach(async function () {
+        await this.store.setBytesStorage(slot, this.value);
+      });
+
+      it('from right slot', async function () {
+        expect(await this.store.bytesMap(slot)).to.be.equal(this.value);
+        expect(await this.store.getBytesStorage(slot)).to.be.equal(this.value);
+      });
+
+      it('from other slot', async function () {
+        expect(await this.store.bytesMap(otherSlot)).to.be.equal(null);
+        expect(await this.store.getBytesStorage(otherSlot)).to.be.equal(null);
+      });
+    });
+  });
 });