Support annotation to validate arbitrary functions as initializers (#1148)

Author: ericglau

docs/modules/ROOT/pages/writing-upgradeable.adoc

@@ -191,12 +191,19 @@ contract MyContract {
 
 Do not leave an implementation contract uninitialized. An uninitialized implementation contract can be taken over by an attacker, which may impact the proxy. To prevent the implementation contract from being used, you should invoke the `_disableInitializers` function in the constructor to automatically lock it when it is deployed:
 
-```
+[source,solidity]
+----
 /// @custom:oz-upgrades-unsafe-allow constructor
 constructor() {
     _disableInitializers();
 }
-```
+----
+
+=== Validating Initializers
+
+The OpenZeppelin Upgrades plugins will automatically detect specific issues with initializers in implementation contracts. These include checking if your implementation contract is missing an initializer when there are parent initializers that need to be called, if a parent initializer is not called, or if a parent initializer is called more than once. The plugins will also warn if parent initializers are not called in linearized order.
+
+Reinitializers are not included in validations by default, because the Upgrades plugins cannot determine whether they are intended to be used for new deployments. If you want to validate a reinitializer function as an initializer that can be used for new deployments, annotate it with `@custom:oz-upgrades-validate-as-initializer`. Note that functions which cannot possibly be initializers are always ignored, such as private functions which cannot be called externally or by child contracts.
 
 [[creating-new-instances-from-your-contract-code]]
 == Creating New Instances From Your Contract Code

packages/core/CHANGELOG.md

@@ -1,8 +1,9 @@
 # Changelog
 
-## Unreleased
+## 1.43.0 (2025-04-11)
 
 - Add Sonic network to manifest file names. ([#1146](https://github.com/OpenZeppelin/openzeppelin-upgrades/pull/1146))
+- Support annotation to validate functions as initializers. ([#1148](https://github.com/OpenZeppelin/openzeppelin-upgrades/pull/1148))
 
 ## 1.42.2 (2025-03-19)
 

packages/core/contracts/test/ValidationsInitializer.sol

@@ -24,6 +24,14 @@ contract Parent_NoInitializer {
   }
 }
 
+contract Parent_ValidateAsInitializer {
+  uint8 x;
+  /// @custom:oz-upgrades-validate-as-initializer
+  function parentAssumeInit() public {
+    x = 1;
+  }
+}
+
 contract Parent_InitializerModifier is Initializable {
   uint8 x;
   function parentInit() initializer internal {
@@ -117,6 +125,36 @@ contract MissingInitializer_UnsafeAllow_Contract is Parent_InitializerModifier {
   }
 }
 
+contract ValidateAsInitializer_Ok is Parent_InitializerModifier {
+  uint y;
+  /// @custom:oz-upgrades-validate-as-initializer
+  function regularFn() public {
+    parentInit();
+    y = 2;
+  }
+}
+
+contract Reinitializer_NotDetected is Parent_InitializerModifier {
+  uint y;
+  function initializeV2() public reinitializer(2) {
+    parentInit();
+    y = 2;
+  }
+}
+
+contract Reinitializer_ValidateAsInitializer_Ok is Parent_InitializerModifier {
+  uint y;
+  /**
+   * Text before
+   * @custom:oz-upgrades-validate-as-initializer
+   * Text after
+   */
+  function initializeV2() public reinitializer(2) {
+    parentInit();
+    y = 2;
+  }
+}
+
 contract A is Initializable {
   uint a;
   function __A_init() onlyInitializing internal {
@@ -263,6 +301,58 @@ contract InitializationOrder_UnsafeAllowDuplicate_But_WrongOrder is A, B, C, Par
   }
 }
 
+contract InitializationOrder_ValidateAsInitializer_Ok is A, B, C, Parent_ValidateAsInitializer {
+  function initialize() public {
+    __A_init();
+    __B_init();
+    __C_init();
+    parentAssumeInit();
+  }
+}
+
+contract InitializationOrder_ValidateAsInitializer_WrongOrder is A, B, C, Parent_ValidateAsInitializer {
+  function initialize() public {
+    __A_init();
+    __B_init();
+    parentAssumeInit();
+    __C_init();
+  }
+}
+
+contract InitializationOrder_ValidateAsInitializer_MissingCall is A, B, C, Parent_ValidateAsInitializer {
+  function initialize() public {
+    __A_init();
+    __B_init();
+    __C_init();
+  }
+}
+
+contract InitializationOrder_ValidateAsInitializer_DuplicateCall is A, B, C, Parent_ValidateAsInitializer {
+  function initialize() public {
+    __A_init();
+    __B_init();
+    __C_init();
+    parentAssumeInit();
+    parentAssumeInit();
+  }
+}
+
+contract Parent_ValidateAsInitializer_External_Ok {
+  uint8 x;
+  /// @custom:oz-upgrades-validate-as-initializer
+  function parentAssumeInit() external {
+    // this is a valid initializer, but it does not need to be called by the child (and it is not possible to do so), because it is external
+    x = 1;
+  }
+}
+
+contract Child_Of_ValidateAsInitializer_External_Ok is Parent_ValidateAsInitializer_External_Ok {
+  uint y;
+  function initialize() public {
+    y = 2;
+  }
+}
+
 contract WithRequire_Ok is Parent__OnlyInitializingModifier {
   uint y;
   function initialize(bool foo) initializer public {

packages/core/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@openzeppelin/upgrades-core",
-  "version": "1.42.2",
+  "version": "1.43.0",
   "description": "",
   "repository": "https://github.com/OpenZeppelin/openzeppelin-upgrades/tree/master/packages/core",
   "license": "MIT",

packages/core/src/validate-initializers.test.ts

@@ -42,7 +42,7 @@ interface ExpectedErrors {
   count: number;
 }
 
-function testRejects(name: string, kind: ValidationOptions['kind'], expectedErrors?: ExpectedErrors) {
+function testRejects(name: string, kind: ValidationOptions['kind'], expectedErrors: ExpectedErrors) {
   testOverride(name, kind, {}, expectedErrors);
 }
 
@@ -98,10 +98,18 @@ testRejects('MissingInitializer_Bad', 'transparent', {
 testAccepts('MissingInitializer_UnsafeAllow_Contract', 'transparent');
 testOverride('MissingInitializer_Bad', 'transparent', { unsafeAllow: ['missing-initializer'] });
 
+testAccepts('ValidateAsInitializer_Ok', 'transparent');
+
+testRejects('Reinitializer_NotDetected', 'transparent', {
+  contains: ['Missing initializer'],
+  count: 1,
+});
+testAccepts('Reinitializer_ValidateAsInitializer_Ok', 'transparent');
+
 testAccepts('InitializationOrder_Ok', 'transparent');
 testAccepts('InitializationOrder_Ok_2', 'transparent');
 
-testAccepts('InitializationOrder_WrongOrder_Bad', 'transparent'); // warn 'Expected initializers to be called for parent contracts in the following order: A, B, C'
+testAccepts('InitializationOrder_WrongOrder_Bad', 'transparent'); // warn 'Expected: A, B, C'
 testAccepts('InitializationOrder_WrongOrder_UnsafeAllow_Contract', 'transparent');
 testAccepts('InitializationOrder_WrongOrder_UnsafeAllow_Function', 'transparent');
 testOverride('InitializationOrder_WrongOrder_Bad', 'transparent', { unsafeAllow: ['incorrect-initializer-order'] }); // skips the warning
@@ -122,7 +130,21 @@ testAccepts('InitializationOrder_Duplicate_UnsafeAllow_Contract', 'transparent')
 testAccepts('InitializationOrder_Duplicate_UnsafeAllow_Function', 'transparent');
 testAccepts('InitializationOrder_Duplicate_UnsafeAllow_Call', 'transparent');
 testOverride('InitializationOrder_Duplicate_Bad', 'transparent', { unsafeAllow: ['duplicate-initializer-call'] });
-testAccepts('InitializationOrder_UnsafeAllowDuplicate_But_WrongOrder', 'transparent'); // warn 'Expected initializers to be called for parent contracts in the following order: A, B, C'
+testAccepts('InitializationOrder_UnsafeAllowDuplicate_But_WrongOrder', 'transparent'); // warn 'Expected: A, B, C'
+
+testAccepts('InitializationOrder_ValidateAsInitializer_Ok', 'transparent');
+testAccepts('InitializationOrder_ValidateAsInitializer_WrongOrder', 'transparent'); // warn 'Expected: A, B, C, Parent_ValidateAsInitializer'
+testRejects('InitializationOrder_ValidateAsInitializer_MissingCall', 'transparent', {
+  contains: ['Missing initializer calls for one or more parent contracts: `Parent_ValidateAsInitializer`'],
+  count: 1,
+});
+testRejects('InitializationOrder_ValidateAsInitializer_DuplicateCall', 'transparent', {
+  contains: ['Duplicate calls found to initializer `parentAssumeInit` for contract `Parent_ValidateAsInitializer`'],
+  count: 1,
+});
+
+testAccepts('Parent_ValidateAsInitializer_External_Ok', 'transparent');
+testAccepts('Child_Of_ValidateAsInitializer_External_Ok', 'transparent');
 
 testAccepts('WithRequire_Ok', 'transparent');
 testRejects('WithRequire_Bad', 'transparent', {
@@ -178,11 +200,11 @@ testRejects('TransitiveChild_Bad_Parent', 'transparent', {
     'Missing initializer calls for one or more parent contracts: `TransitiveGrandparent2`', // occurs twice: missing initializer for child, missing initializer for parent
   ],
   count: 2,
-}); // warn 'Expected initializers to be called for parent contracts in the following order: TransitiveGrandparent2, TransitiveParent_Bad'
+}); // warn 'Expected: TransitiveGrandparent2, TransitiveParent_Bad'
 testRejects('TransitiveChild_Bad_Order', 'transparent', {
   contains: ['Missing initializer calls for one or more parent contracts: `TransitiveGrandparent2`'],
   count: 1,
-}); // warn 'Expected initializers to be called for parent contracts in the following order: TransitiveGrandparent2, TransitiveParent_Bad'
+}); // warn 'Expected: TransitiveGrandparent2, TransitiveParent_Bad'
 testRejects('TransitiveChild_Good_Order_Bad_Parent', 'transparent', {
   contains: ['Missing initializer calls for one or more parent contracts: `TransitiveGrandparent2`'],
   count: 1,

packages/core/src/validate/run/initializer.ts

@@ -1,7 +1,10 @@
 import type { ContractDefinition, FunctionDefinition } from 'solidity-ast';
+import type { Node } from 'solidity-ast/node';
 import { ASTDereferencer, findAll } from 'solidity-ast/utils';
 import { SrcDecoder } from '../../src-decoder';
 import { ValidationExceptionInitializer, skipCheck, tryDerefFunction } from '../run';
+import { getAnnotationArgs, getDocumentation, hasAnnotationTag } from '../../utils/annotations';
+import { logNote } from '../../utils/log';
 
 /**
  * Reports if this contract is non-abstract and any of the following are true:
@@ -20,12 +23,12 @@ export function* getInitializerExceptions(
   }
 
   const linearizedParentContracts = getLinearizedParentContracts(contractDef, deref);
-  const parentNameToInitializersMap = getParentNameToInitializersMap(linearizedParentContracts);
+  const parentNameToInitializersMap = getParentNameToInitializersMap(linearizedParentContracts, decodeSrc);
 
   const remainingParents = getParentsNotInitializedByOtherParents(parentNameToInitializersMap);
 
   if (remainingParents.length > 0) {
-    const contractInitializers = getPossibleInitializers(contractDef, false);
+    const contractInitializers = getPossibleInitializers(contractDef, false, decodeSrc);
 
     // Report if there is no initializer but parents need initialization
     if (
@@ -71,10 +74,10 @@ function getLinearizedParentContracts(contractDef: ContractDefinition, deref: AS
  * Gets a map of parent contract names to their possible initializers.
  * If a parent contract has no initializers, it is not included in the map.
  */
-function getParentNameToInitializersMap(linearizedParentContracts: ContractDefinition[]) {
+function getParentNameToInitializersMap(linearizedParentContracts: ContractDefinition[], decodeSrc: SrcDecoder) {
   const map = new Map();
   for (const parent of linearizedParentContracts) {
-    const initializers = getPossibleInitializers(parent, true);
+    const initializers = getPossibleInitializers(parent, true, decodeSrc);
     if (initializers.length > 0) {
       map.set(parent.name, initializers);
     }
@@ -318,22 +321,65 @@ function getRecursiveFunctionIds(
 }
 
 /**
- * Get all functions that could be initializers. Does not include private functions.
- * For parent contracts, only internal and public functions which contain statements are included.
+ * Get all functions that are annotated as initializers or are inferred to be initializers.
+ * Logs a note if any reinitializer is found.
  */
-function getPossibleInitializers(contractDef: ContractDefinition, isParentContract: boolean) {
+function getPossibleInitializers(
+  contractDef: ContractDefinition,
+  isParentContract: boolean,
+  decodeSrc: SrcDecoder,
+): FunctionDefinition[] {
   const fns = [...findAll('FunctionDefinition', contractDef)];
-  return fns.filter(
-    fnDef =>
-      (fnDef.modifiers.some(modifier => ['initializer', 'onlyInitializing'].includes(modifier.modifierName.name)) ||
-        ['initialize', 'initializer'].includes(fnDef.name)) &&
-      // Skip virtual functions without a body, since that indicates an abstract function and is not itself an initializer
-      !(fnDef.virtual && !fnDef.body) &&
-      // Ignore private functions, since they cannot be called outside the contract
-      fnDef.visibility !== 'private' &&
-      // For parent contracts, only internal and public functions which contain statements need to be called
-      (isParentContract
-        ? fnDef.body?.statements?.length && (fnDef.visibility === 'internal' || fnDef.visibility === 'public')
-        : true),
+  return fns.filter((fnDef: FunctionDefinition) => {
+    const validateAsInitializer = hasValidateAsInitializerAnnotation(fnDef, decodeSrc);
+    if (!validateAsInitializer && fnDef.modifiers.some(modifier => 'reinitializer' === modifier.modifierName.name)) {
+      logNote(`Reinitializers are not included in validations by default`, [
+        `${decodeSrc(fnDef)}: If you want to validate this function as an initializer, annotate it with '@custom:oz-upgrades-validate-as-initializer'`,
+      ]);
+    }
+
+    return inferPossibleInitializer(fnDef, validateAsInitializer, isParentContract);
+  });
+}
+
+function hasValidateAsInitializerAnnotation(node: Node, decodeSrc: SrcDecoder): boolean {
+  const doc = getDocumentation(node);
+  const tag = 'oz-upgrades-validate-as-initializer';
+  const validateAsInitializer = hasAnnotationTag(doc, tag);
+  if (validateAsInitializer) {
+    const annotationArgs = getAnnotationArgs(doc, tag);
+    if (annotationArgs.length !== 0) {
+      throw new Error(`${decodeSrc(node)}: @custom:${tag} annotation must not have any arguments`);
+    }
+  }
+  return validateAsInitializer;
+}
+
+function hasInitializerNameOrModifier(fnDef: FunctionDefinition) {
+  return (
+    fnDef.modifiers.some(modifier => ['initializer', 'onlyInitializing'].includes(modifier.modifierName.name)) ||
+    ['initialize', 'initializer'].includes(fnDef.name)
+  );
+}
+
+/**
+ * Infers whether a function could be an initializer. Does not include private functions.
+ * For parent contracts, only internal and public functions which contain statements are included.
+ */
+function inferPossibleInitializer(
+  fnDef: FunctionDefinition,
+  validateAsInitializer: boolean,
+  isParentContract: boolean,
+): boolean {
+  return (
+    (validateAsInitializer || hasInitializerNameOrModifier(fnDef)) &&
+    // Skip virtual functions without a body, since that indicates an abstract function and is not itself an initializer
+    !(fnDef.virtual && !fnDef.body) &&
+    // Ignore private functions, since they cannot be called outside the contract
+    fnDef.visibility !== 'private' &&
+    // For parent contracts, only internal and public functions which contain statements need to be called
+    (isParentContract
+      ? Boolean(fnDef.body?.statements?.length) && (fnDef.visibility === 'internal' || fnDef.visibility === 'public')
+      : true)
   );
 }