diff --git a/packages/core/contracts/test/StructAppend.sol b/packages/core/contracts/test/StructAppend.sol new file mode 100644 index 000000000..c0a03e386 --- /dev/null +++ b/packages/core/contracts/test/StructAppend.sol @@ -0,0 +1,42 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.0; + +contract StructAppendV1 { + struct MyStruct { + uint256 x; + uint256 y; + } + + MyStruct public myStruct; +} + +contract StructAppendV2_Ok { + struct MyStruct { + uint256 x; + uint256 y; + uint256 z; // appended field + } + + MyStruct public myStruct; +} + +contract StructAppendBeforeStorageEndV1 { + struct MyStruct { + uint256 x; + uint256 y; + } + + MyStruct public myStruct; + uint256 public zz; +} + +contract StructAppendBeforeStorageEndV2_Bad { + struct MyStruct { + uint256 x; + uint256 y; + uint256 z; // appended field + } + + MyStruct public myStruct; + uint256 public zz; +} diff --git a/packages/core/src/storage/compare.ts b/packages/core/src/storage/compare.ts index 5ce260458..9177a5154 100644 --- a/packages/core/src/storage/compare.ts +++ b/packages/core/src/storage/compare.ts @@ -176,7 +176,32 @@ export class StorageLayoutComparator { // Gap shrink or gap replacement that finishes on the same slot is safe return false; } else if (o.kind === 'append' && allowAppend) { + // Allow append operation if the allowAppend flag is set + // This applies to both direct storage appends and appends within structs return false; + } else if (o.kind === 'typechange' && o.change?.kind === 'struct members' && o.change.allowAppend) { + const structOps = o.change.ops; + const hasNonAppendOps = structOps.some(op => op.kind !== 'append'); + if (hasNonAppendOps) { + return true; + } + + const structEnd = storageFieldEnd(o.original); + if (structEnd === undefined) { + return true; + } + + const hasFieldsAfter = original + .filter(field => !isGap(field)) + .some(field => { + const fieldStart = storageFieldBegin(field); + if (field === o.original || fieldStart === undefined) { + return false; + } + return fieldStart > structEnd; + }); + + return hasFieldsAfter; } return true; }); @@ -324,6 +349,7 @@ export class StorageLayoutComparator { } } assert(isStructMembers(originalMembers) && isStructMembers(updatedMembers)); + const ops = this.layoutLevenshtein(originalMembers, updatedMembers, { allowAppend }); if (ops.length > 0) { return { kind: 'struct members', ops, original, updated, allowAppend }; diff --git a/packages/core/src/storage/struct-append.test.ts b/packages/core/src/storage/struct-append.test.ts new file mode 100644 index 000000000..cc8592e33 --- /dev/null +++ b/packages/core/src/storage/struct-append.test.ts @@ -0,0 +1,73 @@ +import _test, { TestFn } from 'ava'; +import { ContractDefinition } from 'solidity-ast'; +import { findAll, astDereferencer } from 'solidity-ast/utils'; +import { artifacts } from 'hardhat'; + +import { SolcOutput } from '../solc-api'; +import { getStorageUpgradeErrors } from '../storage'; +import { StorageLayout } from './layout'; +import { extractStorageLayout } from './extract'; + +interface Context { + extractStorageLayout: (contract: string) => StorageLayout; +} + +const test = _test as TestFn; + +test.before(async t => { + const buildInfo = await artifacts.getBuildInfo('contracts/test/StructAppend.sol:StructAppendV1'); + if (buildInfo === undefined) { + throw new Error('Build info not found'); + } + const solcOutput: SolcOutput = buildInfo.output; + const contracts: Record = {}; + const storageLayouts: Record = {}; + for (const def of findAll('ContractDefinition', solcOutput.sources['contracts/test/StructAppend.sol'].ast)) { + contracts[def.name] = def; + // eslint-disable-next-line @typescript-eslint/no-non-null-assertion + storageLayouts[def.name] = solcOutput.contracts['contracts/test/StructAppend.sol'][def.name].storageLayout!; + } + const deref = astDereferencer(solcOutput); + const dummyDecodeSrc = () => 'file.sol:1'; + t.context.extractStorageLayout = name => + extractStorageLayout(contracts[name], dummyDecodeSrc, deref, storageLayouts[name]); +}); + +test('struct append at end of storage - ok', t => { + const v1 = t.context.extractStorageLayout('StructAppendV1'); + const v2 = t.context.extractStorageLayout('StructAppendV2_Ok'); + const comparison = getStorageUpgradeErrors(v1, v2); + t.deepEqual(comparison, [], JSON.stringify(comparison, null, 2)); +}); + +test('struct append not at end of storage - bad', t => { + const v1 = t.context.extractStorageLayout('StructAppendBeforeStorageEndV1'); + const v2 = t.context.extractStorageLayout('StructAppendBeforeStorageEndV2_Bad'); + const comparison = getStorageUpgradeErrors(v1, v2); + t.like(comparison, { + length: 2, + 0: { + kind: 'typechange', + change: { + kind: 'struct members', + ops: { + length: 1, + 0: { kind: 'append' }, + }, + }, + original: { label: 'myStruct' }, + updated: { label: 'myStruct' }, + }, + 1: { + kind: 'layoutchange', + change: { + slot: { + from: '2', + to: '3', + }, + }, + original: { label: 'zz' }, + updated: { label: 'zz' }, + }, + }); +}); \ No newline at end of file