@openzeppelin/upgrades-core 1.42.2 → 1.44.0

package/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)
   );
 }

package/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,
@@ -218,7 +240,7 @@ testAccepts('Recursive_Ok', 'transparent');
 testAccepts('ERC20_Ok', 'uups');
 testRejects('ERC20_Bad', 'uups', {
   contains: [
-    'Missing initializer calls for one or more parent contracts: `ERC20Upgradeable, ERC20PausableUpgradeable, OwnableUpgradeable, ERC20PermitUpgradeable`',
+    'Missing initializer calls for one or more parent contracts: `ERC20Upgradeable, OwnableUpgradeable, ERC20PermitUpgradeable`',
   ],
   count: 1,
 });