@openzeppelin/upgrades-core 1.27.2 → 1.28.0

package/src/cli/cli.test.ts

@@ -4,11 +4,10 @@ import { exec } from 'child_process';
 import { promises as fs } from 'fs';
 import path from 'path';
 import os from 'os';
-import rimrafAsync from 'rimraf';
+import { rimraf } from 'rimraf';
 import { artifacts } from 'hardhat';
 
 const execAsync = promisify(exec);
-const rimraf = promisify(rimrafAsync);
 
 const CLI = 'node dist/cli/cli.js';
 
@@ -38,6 +37,136 @@ test('validate - errors', async t => {
   t.snapshot(expectation.join('\n'));
 });
 
+test('validate - single contract', async t => {
+  // This should check even though the contract is not detected as upgradeable, since the --contract option was used.
+  const temp = await getTempDir(t);
+  const buildInfo = await artifacts.getBuildInfo(`contracts/test/cli/Validate.sol:NonUpgradeable`);
+  await fs.writeFile(path.join(temp, 'validate.json'), JSON.stringify(buildInfo));
+
+  const error = await t.throwsAsync(execAsync(`${CLI} validate ${temp} --contract NonUpgradeable`));
+  const expectation: string[] = [`Stdout: ${(error as any).stdout}`, `Stderr: ${(error as any).stderr}`];
+  t.snapshot(expectation.join('\n'));
+});
+
+test('validate - single contract, has upgrades-from', async t => {
+  const temp = await getTempDir(t);
+  const buildInfo = await artifacts.getBuildInfo(`contracts/test/cli/Validate.sol:UnsafeAndStorageLayoutErrors`);
+  await fs.writeFile(path.join(temp, 'validate.json'), JSON.stringify(buildInfo));
+
+  const error = await t.throwsAsync(execAsync(`${CLI} validate ${temp} --contract UnsafeAndStorageLayoutErrors`));
+  const expectation: string[] = [`Stdout: ${(error as any).stdout}`, `Stderr: ${(error as any).stderr}`];
+  t.snapshot(expectation.join('\n'));
+});
+
+test('validate - single contract, reference overrides upgrades-from', async t => {
+  const temp = await getTempDir(t);
+  const buildInfo = await artifacts.getBuildInfo(`contracts/test/cli/Validate.sol:UnsafeAndStorageLayoutErrors`);
+  await fs.writeFile(path.join(temp, 'validate.json'), JSON.stringify(buildInfo));
+
+  const error = await t.throwsAsync(
+    execAsync(`${CLI} validate ${temp} --contract UnsafeAndStorageLayoutErrors --reference Safe`),
+  );
+  const expectation: string[] = [`Stdout: ${(error as any).stdout}`, `Stderr: ${(error as any).stderr}`];
+  t.snapshot(expectation.join('\n'));
+});
+
+test('validate - single contract, reference is uups, overrides upgrades-from', async t => {
+  const temp = await getTempDir(t);
+  const buildInfo = await artifacts.getBuildInfo(`contracts/test/cli/Validate.sol:UnsafeAndStorageLayoutErrors`);
+  await fs.writeFile(path.join(temp, 'validate.json'), JSON.stringify(buildInfo));
+
+  const error = await t.throwsAsync(
+    execAsync(`${CLI} validate ${temp} --contract BecomesBadLayout --reference HasUpgradeTo`),
+  );
+  const expectation: string[] = [`Stdout: ${(error as any).stdout}`, `Stderr: ${(error as any).stderr}`];
+  t.snapshot(expectation.join('\n'));
+});
+
+test('validate - single contract, reference', async t => {
+  const temp = await getTempDir(t);
+  const buildInfo = await artifacts.getBuildInfo(`contracts/test/cli/Validate.sol:BecomesBadLayout`);
+  await fs.writeFile(path.join(temp, 'validate.json'), JSON.stringify(buildInfo));
+
+  const error = await t.throwsAsync(
+    execAsync(`${CLI} validate ${temp} --contract BecomesBadLayout --reference StorageV1`),
+  );
+  const expectation: string[] = [`Stdout: ${(error as any).stdout}`, `Stderr: ${(error as any).stderr}`];
+  t.snapshot(expectation.join('\n'));
+});
+
+test('validate - single contract, reference, fully qualified names', async t => {
+  const temp = await getTempDir(t);
+  const buildInfo = await artifacts.getBuildInfo(`contracts/test/cli/Validate.sol:BecomesBadLayout`);
+  await fs.writeFile(path.join(temp, 'validate.json'), JSON.stringify(buildInfo));
+
+  const error = await t.throwsAsync(
+    execAsync(
+      `${CLI} validate ${temp} --contract contracts/test/cli/Validate.sol:BecomesBadLayout --reference contracts/test/cli/Validate.sol:StorageV1`,
+    ),
+  );
+  const expectation: string[] = [`Stdout: ${(error as any).stdout}`, `Stderr: ${(error as any).stderr}`];
+  t.snapshot(expectation.join('\n'));
+});
+
+test('validate - reference without contract option', async t => {
+  const temp = await getTempDir(t);
+  const buildInfo = await artifacts.getBuildInfo(`contracts/test/cli/Validate.sol:StorageV1`);
+  await fs.writeFile(path.join(temp, 'validate.json'), JSON.stringify(buildInfo));
+
+  const error = await t.throwsAsync(execAsync(`${CLI} validate ${temp} --reference StorageV1`));
+  t.true(
+    error?.message.includes('The --reference option can only be used along with the --contract option.'),
+    error?.message,
+  );
+});
+
+test('validate - emtpy contract string', async t => {
+  const temp = await getTempDir(t);
+  const buildInfo = await artifacts.getBuildInfo(`contracts/test/cli/Validate.sol:StorageV1`);
+  await fs.writeFile(path.join(temp, 'validate.json'), JSON.stringify(buildInfo));
+
+  const error = await t.throwsAsync(execAsync(`${CLI} validate ${temp} --contract --reference StorageV1`));
+  t.true(error?.message.includes('Invalid option: --contract cannot be empty'), error?.message);
+});
+
+test('validate - blank contract string', async t => {
+  const temp = await getTempDir(t);
+  const buildInfo = await artifacts.getBuildInfo(`contracts/test/cli/Validate.sol:StorageV1`);
+  await fs.writeFile(path.join(temp, 'validate.json'), JSON.stringify(buildInfo));
+
+  const error = await t.throwsAsync(execAsync(`${CLI} validate ${temp} --contract '    ' --reference StorageV1`));
+  t.true(error?.message.includes('Invalid option: --contract cannot be empty'), error?.message);
+});
+
+test('validate - empty reference string', async t => {
+  const temp = await getTempDir(t);
+  const buildInfo = await artifacts.getBuildInfo(`contracts/test/cli/Validate.sol:StorageV1`);
+  await fs.writeFile(path.join(temp, 'validate.json'), JSON.stringify(buildInfo));
+
+  const error = await t.throwsAsync(execAsync(`${CLI} validate ${temp} --contract StorageV1 --reference`));
+  t.true(error?.message.includes('Invalid option: --reference cannot be empty'), error?.message);
+});
+
+test('validate - single contract not found', async t => {
+  const temp = await getTempDir(t);
+  const buildInfo = await artifacts.getBuildInfo(`contracts/test/cli/Validate.sol:BecomesBadLayout`);
+  await fs.writeFile(path.join(temp, 'validate.json'), JSON.stringify(buildInfo));
+
+  const error = await t.throwsAsync(execAsync(`${CLI} validate ${temp} --contract NonExistent`));
+  t.true(error?.message.includes('Could not find contract NonExistent.'), error?.message);
+});
+
+test('validate - reference not found', async t => {
+  const temp = await getTempDir(t);
+  const buildInfo = await artifacts.getBuildInfo(`contracts/test/cli/Validate.sol:BecomesBadLayout`);
+  await fs.writeFile(path.join(temp, 'validate.json'), JSON.stringify(buildInfo));
+
+  const error = await t.throwsAsync(
+    execAsync(`${CLI} validate ${temp} --contract BecomesBadLayout --reference NonExistent`),
+  );
+  t.true(error?.message.includes('Could not find contract NonExistent.'), error?.message);
+});
+
 test('validate - no upgradeable', async t => {
   const temp = await getTempDir(t);
   const buildInfo = await artifacts.getBuildInfo(`contracts/test/cli/Storage088.sol:Storage088`);

package/src/cli/cli.test.ts.md

@@ -16,6 +16,8 @@ Generated by [AVA](https://avajs.dev).
       <BUILD_INFO_DIR>  Optional path to the build info directory which contains JSON files with Solidity compiler input and output. Defaults to 'artifacts/build-info' for Hardhat projects or 'out/build-info' for Foundry projects. If your project uses a custom output directory, you must specify its build info directory here.␊
     ␊
     Options:␊
+      --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>  The name or fully qualified name of the reference contract to use for storage layout comparisons. Can only be used along with the --contract option. If not specified, uses the @custom:oz-upgrades-from annotation in the contract that is being validated.␊
       --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␊
       --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.␊
@@ -33,6 +35,8 @@ Generated by [AVA](https://avajs.dev).
       <BUILD_INFO_DIR>  Optional path to the build info directory which contains JSON files with Solidity compiler input and output. Defaults to 'artifacts/build-info' for Hardhat projects or 'out/build-info' for Foundry projects. If your project uses a custom output directory, you must specify its build info directory here.␊
     ␊
     Options:␊
+      --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>  The name or fully qualified name of the reference contract to use for storage layout comparisons. Can only be used along with the --contract option. If not specified, uses the @custom:oz-upgrades-from annotation in the contract that is being validated.␊
       --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␊
       --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.␊
@@ -121,6 +125,96 @@ Generated by [AVA](https://avajs.dev).
     ␊
     Stderr: `
 
+## validate - single contract
+
+> Snapshot 1
+
+    `Stdout:  ✘  contracts/test/cli/Validate.sol:NonUpgradeable␊
+    ␊
+          contracts/test/cli/Validate.sol:14: Use of selfdestruct is not allowed␊
+              https://zpl.in/upgrades/error-003␊
+    ␊
+    FAILED (1 upgradeable contract detected, 0 passed, 1 failed)␊
+    ␊
+    Stderr: `
+
+## validate - single contract, has upgrades-from
+
+> Snapshot 1
+
+    `Stdout:  ✘  contracts/test/cli/Validate.sol:UnsafeAndStorageLayoutErrors (upgrades from contracts/test/cli/Validate.sol:StorageV1)␊
+    ␊
+          contracts/test/cli/Validate.sol:97: Use of selfdestruct is not allowed␊
+              https://zpl.in/upgrades/error-003␊
+    ␊
+          contracts/test/cli/Validate.sol:94: Layout changed for \`__gap\` (uint256[49] -> uint256[49])␊
+            - Slot changed from 1 to 2␊
+            > Set __gap array to size 48␊
+    ␊
+    FAILED (1 upgradeable contract detected, 0 passed, 1 failed)␊
+    ␊
+    Stderr: `
+
+## validate - single contract, reference overrides upgrades-from
+
+> Snapshot 1
+
+    `Stdout:  ✘  contracts/test/cli/Validate.sol:UnsafeAndStorageLayoutErrors (upgrades from contracts/test/cli/Validate.sol:Safe)␊
+    ␊
+          contracts/test/cli/Validate.sol:97: Use of selfdestruct is not allowed␊
+              https://zpl.in/upgrades/error-003␊
+    ␊
+    FAILED (1 upgradeable contract detected, 0 passed, 1 failed)␊
+    ␊
+    Stderr: `
+
+## validate - single contract, reference is uups, overrides upgrades-from
+
+> Snapshot 1
+
+    `Stdout:  ✘  contracts/test/cli/Validate.sol:BecomesBadLayout (upgrades from contracts/test/cli/Validate.sol:HasUpgradeTo)␊
+    ␊
+          contracts/test/cli/Validate.sol:111: Implementation is missing a public \`upgradeTo(address)\` function␊
+              Inherit UUPSUpgradeable to include this function in your contract␊
+              @openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol␊
+              https://zpl.in/upgrades/error-008␊
+    ␊
+    FAILED (1 upgradeable contract detected, 0 passed, 1 failed)␊
+    ␊
+    Stderr: `
+
+## validate - single contract, reference
+
+> Snapshot 1
+
+    `Stdout:  ✘  contracts/test/cli/Validate.sol:BecomesBadLayout (upgrades from contracts/test/cli/Validate.sol:StorageV1)␊
+    ␊
+          StorageV1: Deleted \`x\`␊
+            > Keep the variable even if unused␊
+          ␊
+          StorageV1: Deleted \`__gap\`␊
+            > Keep the variable even if unused␊
+    ␊
+    FAILED (1 upgradeable contract detected, 0 passed, 1 failed)␊
+    ␊
+    Stderr: `
+
+## validate - single contract, reference, fully qualified names
+
+> Snapshot 1
+
+    `Stdout:  ✘  contracts/test/cli/Validate.sol:BecomesBadLayout (upgrades from contracts/test/cli/Validate.sol:StorageV1)␊
+    ␊
+          StorageV1: Deleted \`x\`␊
+            > Keep the variable even if unused␊
+          ␊
+          StorageV1: Deleted \`__gap\`␊
+            > Keep the variable even if unused␊
+    ␊
+    FAILED (1 upgradeable contract detected, 0 passed, 1 failed)␊
+    ␊
+    Stderr: `
+
 ## validate - no upgradeable
 
 > Snapshot 1

package/src/cli/index.ts

@@ -1,3 +1,3 @@
 export { ValidateUpgradeSafetyOptions, validateUpgradeSafety } from './validate/validate-upgrade-safety';
 export { ProjectReport } from './validate/project-report';
-export { ReferenceContractNotFound } from './validate/upgradeability-assessment';
+export { ReferenceContractNotFound } from './validate/find-contract';

package/src/cli/validate/build-info-file.test.ts

@@ -1,14 +1,11 @@
 import test, { ExecutionContext } from 'ava';
 
 import { promises as fs } from 'fs';
-import rimrafAsync from 'rimraf';
-import util from 'util';
+import { rimraf } from 'rimraf';
 import path from 'path';
 import os from 'os';
 import { BuildInfoFile, getBuildInfoFiles } from './build-info-file';
 
-const rimraf = util.promisify(rimrafAsync);
-
 test.beforeEach(async t => {
   process.chdir(await fs.mkdtemp(path.join(os.tmpdir(), `upgrades-core-test-${t.title.replace(/\s/g, '-')}-`)));
 });

package/src/cli/validate/contract-report.ts

@@ -17,6 +17,7 @@ import { getUpgradeabilityAssessment } from './upgradeability-assessment';
 import { SourceContract } from './validations';
 import { LayoutCompatibilityReport } from '../../storage/report';
 import { indent } from '../../utils/indent';
+import { SpecifiedContracts } from './validate-upgrade-safety';
 
 /**
  * Report for an upgradeable contract.
@@ -63,17 +64,28 @@ export class UpgradeableContractReport implements Report {
  *
  * @param sourceContracts The source contracts to check, which must include all contracts that are referenced by the given contracts. Can also include non-upgradeable contracts, which will be ignored.
  * @param opts The validation options.
+ * @param specifiedContracts If provided, only the specified contract (upgrading from its reference contract) will be reported.
  * @returns The upgradeable contract reports.
  */
-export function getContractReports(sourceContracts: SourceContract[], opts: ValidateUpgradeSafetyOptions) {
+export function getContractReports(
+  sourceContracts: SourceContract[],
+  opts: ValidateUpgradeSafetyOptions,
+  specifiedContracts?: SpecifiedContracts,
+) {
   const upgradeableContractReports: UpgradeableContractReport[] = [];
-  for (const sourceContract of sourceContracts) {
-    const upgradeabilityAssessment = getUpgradeabilityAssessment(sourceContract, sourceContracts);
-    if (upgradeabilityAssessment.upgradeable) {
-      const reference = upgradeabilityAssessment.referenceContract;
-      const uups = upgradeabilityAssessment.uups;
-      const kind = uups ? 'uups' : 'transparent';
 
+  const contractsToReport: SourceContract[] =
+    specifiedContracts !== undefined ? [specifiedContracts.contract] : sourceContracts;
+
+  for (const sourceContract of contractsToReport) {
+    const upgradeabilityAssessment = getUpgradeabilityAssessment(
+      sourceContract,
+      sourceContracts,
+      specifiedContracts?.reference,
+    );
+    if (specifiedContracts !== undefined || upgradeabilityAssessment.upgradeable) {
+      const reference = upgradeabilityAssessment.referenceContract;
+      const kind = upgradeabilityAssessment.uups ? 'uups' : 'transparent';
       const report = getUpgradeableContractReport(sourceContract, reference, { ...opts, kind: kind });
       if (report !== undefined) {
         upgradeableContractReports.push(report);

package/src/cli/validate/find-contract.ts

@@ -0,0 +1,44 @@
+import { ValidateCommandError } from './error';
+import { SourceContract } from './validations';
+
+export class ReferenceContractNotFound extends Error {
+  /**
+   * The contract reference that could not be found.
+   */
+  readonly reference: string;
+
+  /**
+   * The fully qualified name of the contract that referenced the missing contract.
+   */
+  readonly origin?: string;
+
+  constructor(reference: string, origin?: string) {
+    const msg =
+      origin !== undefined
+        ? `Could not find contract ${reference} referenced in ${origin}.`
+        : `Could not find contract ${reference}.`;
+    super(msg);
+    this.reference = reference;
+    this.origin = origin;
+  }
+}
+
+export function findContract(contractName: string, origin: SourceContract | undefined, allContracts: SourceContract[]) {
+  const foundContracts = allContracts.filter(c => c.fullyQualifiedName === contractName || c.name === contractName);
+
+  if (foundContracts.length > 1) {
+    const msg =
+      origin !== undefined
+        ? `Found multiple contracts with name ${contractName} referenced in ${origin.fullyQualifiedName}.`
+        : `Found multiple contracts with name ${contractName}.`;
+    throw new ValidateCommandError(
+      msg,
+      () =>
+        `This may be caused by old copies of build info files. Clean and recompile your project, then run the command again with the updated files.`,
+    );
+  } else if (foundContracts.length === 1) {
+    return foundContracts[0];
+  } else {
+    throw new ReferenceContractNotFound(contractName, origin?.fullyQualifiedName);
+  }
+}

package/src/cli/validate/project-report.test.ts

@@ -3,11 +3,10 @@ import { UpgradeableContractErrorReport } from '../../validate';
 
 import _test, { TestFn } from 'ava';
 import { ContractDefinition } from 'solidity-ast';
-import { findAll } from 'solidity-ast/utils';
+import { findAll, astDereferencer } from 'solidity-ast/utils';
 import { artifacts } from 'hardhat';
 
 import { SolcOutput } from '../../solc-api';
-import { astDereferencer } from '../../ast-dereferencer';
 import { extractStorageLayout } from '../../storage/extract';
 import { StorageLayoutComparator } from '../../storage/compare';
 import { StorageLayout, getDetailedLayout } from '../../storage/layout';

package/src/cli/validate/upgradeability-assessment.ts

@@ -1,26 +1,9 @@
 import { getAnnotationArgs, getDocumentation } from '../../utils/annotations';
 import { inferInitializable, inferUUPS } from '../../validate/query';
 import { ValidateCommandError } from './error';
+import { findContract } from './find-contract';
 import { SourceContract } from './validations';
 
-export class ReferenceContractNotFound extends Error {
-  /**
-   * The contract reference that could not be found.
-   */
-  readonly reference: string;
-
-  /**
-   * The fully qualified name of the contract that referenced the missing contract.
-   */
-  readonly origin: string;
-
-  constructor(reference: string, origin: string) {
-    super(`Could not find contract ${reference} referenced in ${origin}.`);
-    this.reference = reference;
-    this.origin = origin;
-  }
-}
-
 interface AnnotationAssessment {
   upgradeable: boolean;
   referenceName?: string;
@@ -35,6 +18,7 @@ export interface UpgradeabilityAssessment {
 export function getUpgradeabilityAssessment(
   contract: SourceContract,
   allContracts: SourceContract[],
+  overrideReferenceContract?: SourceContract,
 ): UpgradeabilityAssessment {
   const fullContractName = contract.fullyQualifiedName;
   const contractValidation = contract.validationData[fullContractName];
@@ -42,42 +26,26 @@ export function getUpgradeabilityAssessment(
   const isUUPS = inferUUPS(contract.validationData, fullContractName);
 
   const annotationAssessment = getAnnotationAssessment(contract);
-  if (annotationAssessment.upgradeable) {
-    let referenceContract = undefined;
-    let isReferenceUUPS = false;
-    if (annotationAssessment.referenceName !== undefined) {
-      referenceContract = getReferenceContract(annotationAssessment.referenceName, contract, allContracts);
-      isReferenceUUPS = inferUUPS(referenceContract.validationData, referenceContract.fullyQualifiedName);
-    }
 
-    return {
-      upgradeable: true,
-      referenceContract: referenceContract,
-      uups: isReferenceUUPS || isUUPS,
-    };
-  } else {
-    const initializable = inferInitializable(contractValidation);
-    return {
-      upgradeable: initializable || isUUPS,
-      uups: isUUPS,
-    };
+  let referenceContract = overrideReferenceContract;
+  if (referenceContract === undefined && annotationAssessment.referenceName !== undefined) {
+    referenceContract = findContract(annotationAssessment.referenceName, contract, allContracts);
   }
-}
-
-function getReferenceContract(reference: string, origin: SourceContract, allContracts: SourceContract[]) {
-  const referenceContracts = allContracts.filter(c => c.fullyQualifiedName === reference || c.name === reference);
 
-  if (referenceContracts.length > 1) {
-    throw new ValidateCommandError(
-      `Found multiple contracts with name ${reference} referenced in ${origin.fullyQualifiedName}.`,
-      () =>
-        `This may be caused by old copies of build info files. Clean and recompile your project, then run the command again with the updated files.`,
-    );
-  } else if (referenceContracts.length === 1) {
-    return referenceContracts[0];
-  } else {
-    throw new ReferenceContractNotFound(reference, origin.fullyQualifiedName);
+  let isReferenceUUPS = false;
+  if (referenceContract !== undefined) {
+    isReferenceUUPS = inferUUPS(referenceContract.validationData, referenceContract.fullyQualifiedName);
   }
+
+  return {
+    upgradeable:
+      referenceContract !== undefined ||
+      annotationAssessment.upgradeable ||
+      inferInitializable(contractValidation) ||
+      isUUPS,
+    referenceContract: referenceContract,
+    uups: isReferenceUUPS || isUUPS,
+  };
 }
 
 function getAnnotationAssessment(contract: SourceContract): AnnotationAssessment {

package/src/cli/validate/validate-upgrade-safety.test.ts

@@ -1,16 +1,13 @@
 import test from 'ava';
 
 import { promises as fs } from 'fs';
-import rimrafAsync from 'rimraf';
-import util from 'util';
+import { rimraf } from 'rimraf';
 import path from 'path';
 import os from 'os';
 
 import { artifacts } from 'hardhat';
 import { validateUpgradeSafety } from './validate-upgrade-safety';
-import { ReferenceContractNotFound } from './upgradeability-assessment';
-
-const rimraf = util.promisify(rimrafAsync);
+import { ReferenceContractNotFound } from './find-contract';
 
 test.before(async () => {
   process.chdir(await fs.mkdtemp(path.join(os.tmpdir(), 'upgrades-core-test-')));

package/src/cli/validate/validate-upgrade-safety.ts

@@ -2,28 +2,58 @@ import { ValidationOptions } from '../..';
 
 import { getBuildInfoFiles } from './build-info-file';
 import { getContractReports } from './contract-report';
+import { findContract } from './find-contract';
 import { ProjectReport, getProjectReport } from './project-report';
-import { validateBuildInfoContracts } from './validations';
+import { SourceContract, validateBuildInfoContracts } from './validations';
 
 /**
  * Validation options for upgrade safety checks.
  */
 export type ValidateUpgradeSafetyOptions = Omit<ValidationOptions, 'kind'>;
 
+export type SpecifiedContracts = {
+  contract: SourceContract;
+  reference?: SourceContract;
+};
+
 /**
  * Validates the upgrade safety of all contracts in the build info dir's build info files.
  * Only contracts that are detected as upgradeable will be validated.
  *
  * @param buildInfoDir Absolute path of build info directory, or undefined to use the default Hardhat or Foundry build-info dir.
+ * @param 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.
+ * @param reference The name or fully qualified name of the reference contract to use for storage layout comparisons. Can only be used along with `contract`. If not specified, uses the `@custom:oz-upgrades-from` annotation in the contract that is being validated.
  * @param opts Validation options, or undefined to use the default validation options.
  * @returns The project report.
  */
 export async function validateUpgradeSafety(
   buildInfoDir?: string,
+  contract?: string,
+  reference?: string,
   opts: ValidateUpgradeSafetyOptions = {},
 ): Promise<ProjectReport> {
   const buildInfoFiles = await getBuildInfoFiles(buildInfoDir);
   const sourceContracts = validateBuildInfoContracts(buildInfoFiles);
-  const contractReports = getContractReports(sourceContracts, opts);
+
+  const specifiedContracts = findSpecifiedContracts(sourceContracts, contract, reference);
+
+  const contractReports = getContractReports(sourceContracts, opts, specifiedContracts);
   return getProjectReport(contractReports);
 }
+
+function findSpecifiedContracts(
+  sourceContracts: SourceContract[],
+  contractName?: string,
+  referenceName?: string,
+): SpecifiedContracts | undefined {
+  if (contractName !== undefined) {
+    return {
+      contract: findContract(contractName, undefined, sourceContracts),
+      reference: referenceName !== undefined ? findContract(referenceName, undefined, sourceContracts) : undefined,
+    };
+  } else if (referenceName !== undefined) {
+    throw new Error(`The reference option can only be specified when the contract option is also specified.`);
+  } else {
+    return undefined;
+  }
+}

package/src/cli/validate.test.ts

@@ -99,7 +99,7 @@ test('withDefaults - all', t => {
   t.is(opts.unsafeSkipStorageCheck, true);
   t.is(opts.unsafeAllowCustomTypes, true);
   t.is(opts.unsafeAllowLinkedLibraries, true);
-  t.true(opts.unsafeAllow.every((kind: string) => errorKinds.includes(kind as typeof errorKinds[number])));
+  t.true(opts.unsafeAllow.every((kind: string) => errorKinds.includes(kind as (typeof errorKinds)[number])));
 });
 
 test('withDefaults - invalid unsafeAllow', t => {

package/src/cli/validate.ts

@@ -13,6 +13,8 @@ Arguments:
   <BUILD_INFO_DIR>  Optional path to the build info directory which contains JSON files with Solidity compiler input and output. Defaults to 'artifacts/build-info' for Hardhat projects or 'out/build-info' for Foundry projects. If your project uses a custom output directory, you must specify its build info directory here.
 
 Options:
+  --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>  The name or fully qualified name of the reference contract to use for storage layout comparisons. Can only be used along with the --contract option. If not specified, uses the @custom:oz-upgrades-from annotation in the contract that is being validated.
   --unsafeAllow "<VALIDATION_ERRORS>"  Selectively disable one or more validation errors. Comma-separated list with one or more of the following: ${errorKinds.join(
     ', ',
   )}
@@ -24,7 +26,12 @@ export async function main(args: string[]): Promise<void> {
 
   if (!help(parsedArgs, extraArgs)) {
     const functionArgs = getFunctionArgs(parsedArgs, extraArgs);
-    const result = await validateUpgradeSafety(functionArgs.buildInfoDir, functionArgs.opts);
+    const result = await validateUpgradeSafety(
+      functionArgs.buildInfoDir,
+      functionArgs.contract,
+      functionArgs.reference,
+      functionArgs.opts,
+    );
     console.log(result.explain());
     process.exitCode = result.ok ? 0 : 1;
   }
@@ -39,7 +46,7 @@ function parseArgs(args: string[]) {
       'unsafeAllowCustomTypes',
       'unsafeAllowLinkedLibraries',
     ],
-    string: ['unsafeAllow'],
+    string: ['unsafeAllow', 'contract', 'reference'],
     alias: { h: 'help' },
   });
   const extraArgs = parsedArgs._;
@@ -59,6 +66,8 @@ function help(parsedArgs: minimist.ParsedArgs, extraArgs: string[]): boolean {
 
 interface FunctionArgs {
   buildInfoDir?: string;
+  contract?: string;
+  reference?: string;
   opts: Required<ValidateUpgradeSafetyOptions>;
 }
 
@@ -76,9 +85,22 @@ export function getFunctionArgs(parsedArgs: minimist.ParsedArgs, extraArgs: stri
     throw new Error('The validate command takes only one argument: the build info directory.');
   } else {
     const buildInfoDir = extraArgs.length === 1 ? undefined : extraArgs[1];
+    const contract = getAndValidateString(parsedArgs, 'contract');
+    const reference = getAndValidateString(parsedArgs, 'reference');
+    if (reference !== undefined && contract === undefined) {
+      throw new Error('The --reference option can only be used along with the --contract option.');
+    }
     const opts = withDefaults(parsedArgs);
-    return { buildInfoDir, opts };
+    return { buildInfoDir, contract, reference, opts };
+  }
+}
+
+function getAndValidateString(parsedArgs: minimist.ParsedArgs, option: string): string | undefined {
+  const value = parsedArgs[option];
+  if (value !== undefined && value.trim().length === 0) {
+    throw new Error(`Invalid option: --${option} cannot be empty`);
   }
+  return value;
 }
 
 function validateOptions(parsedArgs: minimist.ParsedArgs) {
@@ -93,6 +115,8 @@ function validateOptions(parsedArgs: minimist.ParsedArgs) {
         'unsafeAllowCustomTypes',
         'unsafeAllowLinkedLibraries',
         'unsafeAllow',
+        'contract',
+        'reference',
       ].includes(key),
   );
   if (invalidArgs.length > 0) {
@@ -101,7 +125,7 @@ function validateOptions(parsedArgs: minimist.ParsedArgs) {
 }
 
 function getUnsafeAllowKinds(unsafeAllow: string | undefined): ValidationError['kind'][] {
-  type errorKindsType = typeof errorKinds[number];
+  type errorKindsType = (typeof errorKinds)[number];
 
   if (unsafeAllow === undefined) {
     return [];