Fix storage layout comparison for function types, disallow internal functions in storage

docs/modules/ROOT/pages/api-core.adoc

@@ -119,7 +119,7 @@ If any errors are found, the command will exit with a non-zero exit code and pri
 * `--contract <CONTRACT>`  The name or fully qualified name of the contract to validate. If not specified, all upgradeable contracts in the build info directory will be validated.
 * `--reference <REFERENCE_CONTRACT>`  Can only be used when the `--contract` option is also provided. The name or fully qualified name of the reference contract to use for storage layout comparisons. If not specified, uses the `@custom:oz-upgrades-from` annotation if it is defined in the contract that is being validated.
 * `--requireReference`  Can only be used when the `--contract` option is also provided. Not compatible with `--unsafeSkipStorageCheck`. If specified, requires either the `--reference` option to be provided or the contract to have a `@custom:oz-upgrades-from` annotation.
-* `--unsafeAllow "<VALIDATION_ERRORS>"` - Selectively disable one or more validation errors. Comma-separated list with one or more of the following: `state-variable-assignment, state-variable-immutable, external-library-linking, struct-definition, enum-definition, constructor, delegatecall, selfdestruct, missing-public-upgradeto`
+* `--unsafeAllow "<VALIDATION_ERRORS>"` - Selectively disable one or more validation errors. Comma-separated list with one or more of the following: `state-variable-assignment, state-variable-immutable, external-library-linking, struct-definition, enum-definition, constructor, delegatecall, selfdestruct, missing-public-upgradeto, internal-function-storage`
 * `--unsafeAllowRenames` - Configure storage layout check to allow variable renaming.
 * `--unsafeSkipStorageCheck` - Skips checking for storage layout compatibility errors. This is a dangerous option meant to be used as a last resort.
 

docs/modules/ROOT/pages/api-hardhat-upgrades.adoc

@@ -16,6 +16,7 @@ The following options are common to some functions.
 ** `"constructor"`: Allows defining a constructor. See `constructorArgs`.
 ** `"delegatecall"`, `"selfdestruct"`: Allow the use of these operations. Incorrect use of this option can put funds at risk of permanent loss. See xref:faq.adoc#delegatecall-selfdestruct[Can I safely use `delegatecall` and `selfdestruct`?]
 ** `"missing-public-upgradeto"`: Allow UUPS implementations that do not contain a public `upgradeTo` or `upgradeToAndCall` function. Enabling this option is likely to cause a revert due to the built-in UUPS safety mechanism.
+** `"internal-function-storage"`: Allow internal functions in storage variables. Internal functions are code pointers which will no longer be valid after an upgrade, so they must be reassigned during upgrades. See xref:faq.adoc#internal-function-storage[How can I use internal functions in storage variables?]
 * `unsafeAllowRenames`: (`boolean`) Configure storage layout check to allow variable renaming.
 * `unsafeSkipStorageCheck`: (`boolean`) upgrades the proxy or beacon without first checking for storage layout compatibility errors. This is a dangerous option meant to be used as a last resort.
 * `constructorArgs`: (`unknown[]`) Provide arguments for the constructor of the implementation contract. Note that these are different from initializer arguments, and will be used in the deployment of the implementation contract itself. Can be used to initialize immutable variables.

docs/modules/ROOT/pages/faq.adoc

@@ -22,9 +22,7 @@ As a replacement for the constructor, it is common to set up an `initialize` fun
 
 [source,solidity]
 ----
-import "@openzeppelin/contracts/proxy/utils/Initializable.sol";
-// Alternatively, if you are using @openzeppelin/contracts-upgradeable:
-// import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
+import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
 
 contract MyContract is Initializable {
   uint256 value;
@@ -50,6 +48,7 @@ Deployment and upgrade related functions come with an optional `opts` object, wh
   * `delegatecall`
   * `selfdestruct`
   * `missing-public-upgradeto`
+  * `internal-function-storage`
 
 This function is a generalized version of the original `unsafeAllowCustomTypes` and `unsafeAllowLinkedLibraries` allowing any check to be manually disabled.
 
@@ -193,13 +192,14 @@ When using UUPS proxies (through the `kind: 'uups'` option), the implementation
 
 The recommended way to include one or both of these functions is by inheriting the `UUPSUpgradeable` contract provided in OpenZeppelin Contracts, as shown below. This contract adds the required function(s), but also contains a built-in mechanism that will check on-chain, at the time of an upgrade, that the new implementation proposed also inherits `UUPSUpgradeable` or implements the same interface. In this way, when using the Upgrades Plugins there are two layers of mitigations to prevent accidentally disabling upgradeability: an off-chain check by the plugins, and an on-chain fallback in the contract itself.
 
-```solidity
+[source,solidity]
+----
 import "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
 
 contract MyContract is Initializable, ..., UUPSUpgradeable {
     ...
 }
-```
+----
 
 Read more about the differences with the Transparent Proxy Pattern in xref:contracts:api:proxy.adoc#transparent-vs-uups[Transparent vs UUPS].
 
@@ -217,26 +217,90 @@ Renaming a variable is disallowed by default because there is a chance that a re
 
 It is possible to disable this check by passing the option `unsafeAllowRenames: true`. A more granular approach is to use a docstring comment `/// @custom:oz-renamed-from <previous name>` right above the variable that is being renamed, for example:
 
-```
+[source,solidity]
+----
 contract V1 {
     uint x;
 }
 contract V2 {
     /// @custom:oz-renamed-from x
     uint y;
 }
-```
+----
 
 Changing the type of a variable is not allowed either, even in cases where the types have the same size and alignment, for the similar reason explained above. As long as we can guarantee that the rest of the layout is not affected by this type change, it is also possible to override this check by placing a docstring comment `/// @custom:oz-retyped-from <previous type>`.
 
-```
+[source,solidity]
+----
 contract V1 {
     bool x;
 }
 contract V2 {
     /// @custom:oz-retyped-from bool
     uint8 x;
 }
-```
+----
 
 Docstring comments don't yet work for struct members, due to a current Solidity limitation.
+
+[[internal-function-storage]]
+== How can I use internal functions in storage variables?
+
+Internal functions in storage variables are code pointers which will no longer be valid after an upgrade, because the code will move around and the pointer would change. To avoid this issue, you can either declare those functions as external, or reassign internal functions during upgrades.
+
+For example, the `messageFunction` variable in the following contract is not upgrade safe. Attempting to call `showMessage()` after an upgrade would likely result in a revert.
+[source,solidity]
+----
+import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
+
+contract V1 is Initializable {
+  function() internal pure returns (string memory) messageFunction;
+
+  function initialize() initializer public {
+    messageFunction = hello;
+  }
+
+  function hello() internal pure returns (string memory) {
+    return "Hello, World!";
+  }
+
+  function showMessage() public view returns (string memory) {
+    return messageFunction();
+  }
+  ...
+}
+----
+
+To allow the above contract to be deployed by the Upgrades Plugins, you can disable the `internal-function-storage` check according to <<how-can-i-disable-checks>>, but ensure you follow the steps below to reassign the internal function during upgrades.
+
+In new versions of this contract, assign the internal function in the storage variable again, for example by using a https://docs.openzeppelin.com/contracts/api/proxy#Initializable-reinitializer-uint64-[reinitializer]:
+[source,solidity]
+----
+import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
+import "./V1.sol";
+
+contract V2 is V1 {
+  function initializeV2() reinitializer(2) public {
+    messageFunction = hello;
+  }
+  ...
+}
+----
+
+Then when upgrading, call the reinitializer function as part of the upgrade process, for example in Hardhat:
+[source,ts]
+----
+await upgrades.upgradeProxy(PROXY_ADDRESS, ContractFactoryV2, {
+  call: 'initializeV2',
+  unsafeAllow: ['internal-function-storage']
+});
+----
+or in Foundry:
+[source,solidity]
+----
+Upgrades.upgradeProxy(
+    PROXY_ADDRESS,
+    "V2.sol",
+    abi.encodeCall(V2.initializeV2, ())
+);
+----

packages/core/CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## Unreleased
+
+- Fix storage layout comparison for function types, disallow internal functions in storage. ([#1032](https://github.com/OpenZeppelin/openzeppelin-upgrades/pull/1032))
+
 ## 1.33.1 (2024-04-25)
 
 - Fix Hardhat compile error when variable has whitespace before semicolon. ([#1020](https://github.com/OpenZeppelin/openzeppelin-upgrades/pull/1020))

packages/core/contracts/test/Storage.sol

@@ -899,4 +899,31 @@ contract V2Storage {
 }
 
 contract StorageUpgrade_ConsumeAndAddGap_Storage_V2 is V2Storage, Parent1, Parent2 {
+}
+
+contract StorageUpgrade_FunctionPointer_V1 {
+    struct S {
+        uint256 a;
+        function(bool) external b;
+    }
+    S private s;
+    function(bool) external c;
+}
+
+contract StorageUpgrade_FunctionPointer_V2_Ok {
+    struct S {
+        uint256 a;
+        function(bool) external b;
+    }
+    S private s;
+    function(bool) external c;
+}
+
+contract StorageUpgrade_FunctionPointer_V2_Bad {
+    struct S {
+        uint256 a;
+        function(bool) b;
+    }
+    S private s;
+    function(bool) c;
 }

packages/core/contracts/test/Validations.sol

@@ -206,3 +206,51 @@ contract TransitiveLibraryIsUnsafe {
         DirectLibrary.f2();
     }
 }
+
+contract StructExternalFunctionPointer {
+    struct S {
+        function(bool) external foo;
+    }
+}
+
+contract StructInternalFunctionPointer {
+    struct S {
+        function(bool) internal foo;
+    }
+}
+
+contract StructImpliedInternalFunctionPointer {
+    struct S {
+        function(bool) foo;
+    }
+}
+
+struct StandaloneStructInternalFn {
+    function(bool) internal foo;
+}
+
+contract UsesStandaloneStructInternalFn {
+    StandaloneStructInternalFn s;
+}
+
+contract StructUsesStandaloneStructInternalFn {
+    struct S2 {
+        StandaloneStructInternalFn s;
+    }
+}
+
+contract RecursiveStructUsesStandaloneStructInternalFn {
+    StructUsesStandaloneStructInternalFn.S2 s2;
+}
+
+contract ExternalFunctionPointer {
+    function(bool) external foo;
+}
+
+contract InternalFunctionPointer {
+    function(bool) internal foo;
+}
+
+contract ImpliedInternalFunctionPointer {
+    function(bool) foo;
+}

packages/core/src/cli/cli.test.ts.md

@@ -19,7 +19,7 @@ Generated by [AVA](https://avajs.dev).
       --contract <CONTRACT>  The name or fully qualified name of the contract to validate. If not specified, all upgradeable contracts in the build info directory will be validated.␊
       --reference <REFERENCE_CONTRACT>  Can only be used when the --contract option is also provided. The name or fully qualified name of the reference contract to use for storage layout comparisons. If not specified, uses the @custom:oz-upgrades-from annotation if it is defined in the contract that is being validated.␊
       --requireReference  Can only be used when the --contract option is also provided. Not compatible with --unsafeSkipStorageCheck. If specified, requires either the --reference option to be provided or the contract to have a @custom:oz-upgrades-from annotation.␊
-      --unsafeAllow "<VALIDATION_ERRORS>"  Selectively disable one or more validation errors. Comma-separated list with one or more of the following: state-variable-assignment, state-variable-immutable, external-library-linking, struct-definition, enum-definition, constructor, delegatecall, selfdestruct, missing-public-upgradeto␊
+      --unsafeAllow "<VALIDATION_ERRORS>"  Selectively disable one or more validation errors. Comma-separated list with one or more of the following: state-variable-assignment, state-variable-immutable, external-library-linking, struct-definition, enum-definition, constructor, delegatecall, selfdestruct, missing-public-upgradeto, internal-function-storage␊
       --unsafeAllowRenames  Configure storage layout check to allow variable renaming.␊
       --unsafeSkipStorageCheck  Skips checking for storage layout compatibility errors. This is a dangerous option meant to be used as a last resort.␊
     `
@@ -39,7 +39,7 @@ Generated by [AVA](https://avajs.dev).
       --contract <CONTRACT>  The name or fully qualified name of the contract to validate. If not specified, all upgradeable contracts in the build info directory will be validated.␊
       --reference <REFERENCE_CONTRACT>  Can only be used when the --contract option is also provided. The name or fully qualified name of the reference contract to use for storage layout comparisons. If not specified, uses the @custom:oz-upgrades-from annotation if it is defined in the contract that is being validated.␊
       --requireReference  Can only be used when the --contract option is also provided. Not compatible with --unsafeSkipStorageCheck. If specified, requires either the --reference option to be provided or the contract to have a @custom:oz-upgrades-from annotation.␊
-      --unsafeAllow "<VALIDATION_ERRORS>"  Selectively disable one or more validation errors. Comma-separated list with one or more of the following: state-variable-assignment, state-variable-immutable, external-library-linking, struct-definition, enum-definition, constructor, delegatecall, selfdestruct, missing-public-upgradeto␊
+      --unsafeAllow "<VALIDATION_ERRORS>"  Selectively disable one or more validation errors. Comma-separated list with one or more of the following: state-variable-assignment, state-variable-immutable, external-library-linking, struct-definition, enum-definition, constructor, delegatecall, selfdestruct, missing-public-upgradeto, internal-function-storage␊
       --unsafeAllowRenames  Configure storage layout check to allow variable renaming.␊
       --unsafeSkipStorageCheck  Skips checking for storage layout compatibility errors. This is a dangerous option meant to be used as a last resort.␊
     `

packages/core/src/storage.test.ts

@@ -937,3 +937,40 @@ test('storage upgrade with struct gap', t => {
     },
   });
 });
+
+test('storage upgrade with function pointers', t => {
+  const v1 = t.context.extractStorageLayout('StorageUpgrade_FunctionPointer_V1');
+  const v2_Ok = t.context.extractStorageLayout('StorageUpgrade_FunctionPointer_V2_Ok');
+  const v2_Bad = t.context.extractStorageLayout('StorageUpgrade_FunctionPointer_V2_Bad');
+
+  t.deepEqual(getStorageUpgradeErrors(v1, v2_Ok), []);
+
+  t.like(getStorageUpgradeErrors(v1, v2_Bad), {
+    length: 2,
+    0: {
+      kind: 'typechange',
+      change: {
+        kind: 'struct members',
+        ops: {
+          length: 1,
+          0: {
+            kind: 'typechange',
+            change: {
+              kind: 'visibility change',
+            },
+          },
+        },
+      },
+      original: { label: 's' },
+      updated: { label: 's' },
+    },
+    1: {
+      kind: 'typechange',
+      change: {
+        kind: 'visibility change',
+      },
+      original: { label: 'c' },
+      updated: { label: 'c' },
+    },
+  });
+});

packages/core/src/storage/compare.ts

@@ -283,7 +283,7 @@ export class StorageLayoutComparator {
     updated: ParsedTypeDetailed,
     { allowAppend }: { allowAppend: boolean },
   ): TypeChange | undefined {
-    if (updated.head.startsWith('t_function')) {
+    if (original.head.startsWith('t_function') && updated.head.startsWith('t_function')) {
       return this.getVisibilityChange(original, updated);
     }
 

packages/core/src/validate.test.ts

@@ -140,6 +140,25 @@ testValid('TransitiveLibraryIsUnsafe', 'transparent', false);
 testValid('contracts/test/ValidationsSameNameSafe.sol:SameName', 'transparent', true);
 testValid('contracts/test/ValidationsSameNameUnsafe.sol:SameName', 'transparent', false);
 
+testValid('StructExternalFunctionPointer', 'transparent', true);
+testValid('StructInternalFunctionPointer', 'transparent', false);
+testValid('StructImpliedInternalFunctionPointer', 'transparent', false);
+testOverride(
+  'StructImpliedInternalFunctionPointer',
+  'transparent',
+  { unsafeAllow: ['internal-function-storage'] },
+  true,
+);
+
+testValid('UsesStandaloneStructInternalFn', 'transparent', false);
+testValid('StructUsesStandaloneStructInternalFn', 'transparent', false);
+testValid('RecursiveStructUsesStandaloneStructInternalFn', 'transparent', false);
+
+testValid('ExternalFunctionPointer', 'transparent', true);
+testValid('InternalFunctionPointer', 'transparent', false);
+testValid('ImpliedInternalFunctionPointer', 'transparent', false);
+testOverride('ImpliedInternalFunctionPointer', 'transparent', { unsafeAllow: ['internal-function-storage'] }, true);
+
 test('ambiguous name', t => {
   const error = t.throws(() => getContractVersion(t.context.validation, 'SameName'));
   t.is(

packages/core/src/validate/overrides.ts

@@ -76,6 +76,11 @@ export const ValidationErrorUnsafeMessages: Record<ValidationError['kind'], stri
     `Not having a public upgradeTo or upgradeToAndCall function in your implementation can break upgradeability.`,
     `Some implementation might check that onchain, and cause the upgrade to revert.`,
   ],
+  'internal-function-storage': [
+    `You are using the \`unsafeAllow.internal-function-storage\` flag.`,
+    `Internal functions are code pointers which will no longer be valid after an upgrade.`,
+    `Make sure you reassign internal functions in storage variables during upgrades.`,
+  ],
 };
 
 export function withValidationDefaults(opts: ValidationOptions): Required<ValidationOptions> {

packages/core/src/validate/report.ts

@@ -64,6 +64,14 @@ const errorInfo: ErrorDescriptions<ValidationError> = {
       `    @openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol`,
     link: 'https://zpl.in/upgrades/error-008',
   },
+  'internal-function-storage': {
+    msg: e => `Variable \`${e.name}\` is an internal function`,
+    hint: () =>
+      `Use external functions in storage variables, or skip this check with the\n` +
+      `     \`unsafeAllow.internal-function-storage\` flag and ensure you always\n` +
+      `    reassign internal functions in storage during upgrades`,
+    link: 'https://zpl.in/upgrades/error-009',
+  },
 };
 
 function describeError(e: ValidationError, color = true): string {