Skip to content

Commit c295620

Browse files
sakulstrasakulstra
authored
feat: diff improvements (#199)
1 parent 09eaecc commit c295620

File tree

4 files changed

+116
-30
lines changed

4 files changed

+116
-30
lines changed

‎src/reports/code-diff.spec.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,15 @@ import {describe, expect, it} from 'vitest';
22
import pre32 from './mocks/pre3-2.json';
33
import post32 from './mocks/post3-2.json';
44
import {diffCode, downloadContract} from './code-diff';
5+
import {diffSlot} from './raw-storage-diff';
56

67
describe.skip('code diffs', () => {
8+
it('should diff slots', () => {
9+
diffSlot(1, '0x0', {
10+
previousValue: '0x4816b2C2895f97fB918f1aE7Da403750a0eE372e',
11+
newValue: '0xE5e48Ad1F9D1A894188b483DcF91f4FaD6AbA43b',
12+
});
13+
});
714
it(
815
'should download contract',
916
() => {

‎src/reports/diff-reports.ts‎

Lines changed: 2 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ import {renderReserve, renderReserveDiff} from './reserve';
66
import {AaveV3Reserve, type AaveV3Snapshot} from './snapshot-types';
77
import {renderStrategy, renderStrategyDiff} from './strategy';
88
import {diffCode, downloadContract} from './code-diff';
9-
import {bytes32ToAddress} from '../utils/storageSlots';
9+
import {diffRawStorage} from './raw-storage-diff';
1010

1111
function hasDiff(input: Record<string, any>): boolean {
1212
if (!input) return false;
@@ -127,21 +127,7 @@ export async function diffReports<A extends AaveV3Snapshot, B extends AaveV3Snap
127127
}
128128

129129
if (raw) {
130-
// ERC1967 slot https://github.com/OpenZeppelin/openzeppelin-contracts/blob/master/contracts/proxy/ERC1967/ERC1967Utils.sol#L21C53-L21C119
131-
const erc1967Slot = '0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc';
132-
Object.keys(raw).map((contract) => {
133-
if (raw[contract].stateDiff[erc1967Slot] && raw[contract].stateDiff[erc1967Slot].previousValue != '0x0000000000000000000000000000000000000000000000000000000000000000') {
134-
const fromAddress = bytes32ToAddress(raw[contract].stateDiff[erc1967Slot].previousValue);
135-
const toAddress = bytes32ToAddress(raw[contract].stateDiff[erc1967Slot].newValue);
136-
const from = downloadContract(pre.chainId, fromAddress);
137-
const to = downloadContract(pre.chainId, toAddress);
138-
const result = diffCode(from, to);
139-
writeFileSync(
140-
`./diffs/${pre.chainId}_${contract}_${fromAddress}_${toAddress}.diff`,
141-
result,
142-
);
143-
}
144-
});
130+
diffRawStorage(pre.chainId, raw);
145131
}
146132

147133
try {

‎src/reports/raw-storage-diff.ts‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
import {writeFileSync, mkdirSync, readFileSync} from 'fs';
2+
import {bytes32ToAddress} from '../utils/storageSlots';
3+
import {diffCode, downloadContract} from './code-diff';
4+
import {RawStorage, SlotDiff} from './snapshot-types';
5+
import {isKnownAddress} from '../govv3/utils/checkAddress';
6+
import {Address, getContract, zeroHash} from 'viem';
7+
import {getClient} from '@bgd-labs/rpc-env';
8+
import {IPool_ABI} from '@bgd-labs/aave-address-book/abis';
9+
10+
export function diffSlot(chainId: number, address: Address, slot: SlotDiff) {
11+
// pure new deployments cannot be diffed, we just download the code in that case
12+
if (slot.previousValue == zeroHash) {
13+
const toAddress = bytes32ToAddress(slot.newValue);
14+
const to = downloadContract(chainId, toAddress);
15+
mkdirSync('./diffs', {recursive: true});
16+
writeFileSync(`./diffs/${chainId}_${address}_${toAddress}.diff`, readFileSync(to), {});
17+
} else {
18+
const fromAddress = bytes32ToAddress(slot.previousValue);
19+
const toAddress = bytes32ToAddress(slot.newValue);
20+
const from = downloadContract(chainId, fromAddress);
21+
const to = downloadContract(chainId, toAddress);
22+
const result = diffCode(from, to);
23+
mkdirSync('./diffs', {recursive: true});
24+
writeFileSync(`./diffs/${chainId}_${address}_${fromAddress}_${toAddress}.diff`, result, {});
25+
}
26+
}
27+
28+
export async function diffRawStorage(chainId: number, raw: RawStorage) {
29+
// ERC1967 slot https://github.com/OpenZeppelin/openzeppelin-contracts/blob/master/contracts/proxy/ERC1967/ERC1967Utils.sol#L21C53-L21C119
30+
const erc1967ImplSlot = '0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc';
31+
const erc1967AdminSlot = '0xb53127684a568b3173ae13b9f8a6016e243e63b6e8ee1178d6a717850b5d6103';
32+
await Promise.all(
33+
(Object.keys(raw) as (keyof typeof raw)[]).map(async (contract) => {
34+
if (raw[contract]) {
35+
const contractName = isKnownAddress(contract, chainId);
36+
// label known contracts if no label was set on foundry
37+
if (!raw[contract].label) {
38+
if (contractName) raw[contract].label = contractName.join(', ');
39+
}
40+
41+
// create contract diff if storage changed
42+
if (raw[contract].stateDiff[erc1967ImplSlot]) {
43+
raw[contract].stateDiff[erc1967ImplSlot].label = 'Implementation slot';
44+
diffSlot(chainId, contract, raw[contract].stateDiff[erc1967ImplSlot]);
45+
}
46+
47+
// admin slot
48+
if (raw[contract].stateDiff[erc1967AdminSlot]) {
49+
raw[contract].stateDiff[erc1967ImplSlot].label = 'Admin slot';
50+
// ... we might want to fetch the owner in that case
51+
}
52+
53+
// Diff code logic libraries
54+
if (contractName && contractName[contractName.length - 1] === 'POOL') {
55+
const oldPool = getContract({
56+
client: getClient(chainId, {}),
57+
abi: IPool_ABI,
58+
address: bytes32ToAddress(raw[contract].stateDiff[erc1967ImplSlot].previousValue),
59+
});
60+
const newPool = getContract({
61+
client: getClient(chainId, {}),
62+
abi: IPool_ABI,
63+
address: bytes32ToAddress(raw[contract].stateDiff[erc1967ImplSlot].previousValue),
64+
});
65+
const addresses = await Promise.all([
66+
oldPool.read.getSupplyLogic(),
67+
newPool.read.getSupplyLogic(),
68+
oldPool.read.getBorrowLogic(),
69+
newPool.read.getBorrowLogic(),
70+
oldPool.read.getLiquidationLogic(),
71+
newPool.read.getLiquidationLogic(),
72+
oldPool.read.getPoolLogic(),
73+
newPool.read.getPoolLogic(),
74+
oldPool.read.getFlashLoanLogic(),
75+
newPool.read.getFlashLoanLogic(),
76+
oldPool.read.getEModeLogic(),
77+
newPool.read.getEModeLogic(),
78+
]);
79+
for (let i = 0; i < addresses.length; i = i + 2) {
80+
diffSlot(chainId, contract, {previousValue: addresses[i], newValue: addresses[i + 1]});
81+
}
82+
}
83+
}
84+
}),
85+
);
86+
}

‎src/reports/snapshot-types.ts‎

Lines changed: 21 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import {Address} from 'viem';
1+
import {Address, Hex} from 'viem';
22
import {zksync} from 'viem/chains';
33
import {z} from 'zod';
44

@@ -103,25 +103,32 @@ const zodChainId = z.nativeEnum(CHAIN_ID);
103103

104104
export type CHAIN_ID = z.infer<typeof zodChainId>;
105105

106+
export const slotDiff = z.object({
107+
previousValue: z.string() as z.ZodType<Hex>,
108+
newValue: z.string() as z.ZodType<Hex>,
109+
label: z.string().optional(), // does not initially exist, but we might want to inject information here
110+
});
111+
112+
export const rawStorageSchema = z.record(
113+
z.string() as z.ZodType<Address>,
114+
z.object({
115+
label: z.string().nullable(),
116+
balanceDiff: z.string().nullable(),
117+
stateDiff: z.record(z.string(), slotDiff),
118+
}),
119+
);
120+
121+
export type RawStorage = z.infer<typeof rawStorageSchema>;
122+
123+
export type SlotDiff = z.infer<typeof slotDiff>;
124+
106125
export const aaveV3SnapshotSchema = z.object({
107126
reserves: z.record(aaveV3ReserveSchema),
108127
strategies: z.record(aaveV3StrategySchema),
109128
eModes: z.record(aaveV3EmodeSchema),
110129
poolConfig: aaveV3ConfigSchema,
111130
chainId: zodChainId,
112-
raw: z
113-
.record(
114-
z.string(),
115-
z.object({
116-
label: z.string().nullable(),
117-
balanceDiff: z.string().nullable(),
118-
stateDiff: z.record(
119-
z.string(),
120-
z.object({previousValue: z.string(), newValue: z.string()}),
121-
),
122-
}),
123-
)
124-
.optional(),
131+
raw: rawStorageSchema.optional(),
125132
});
126133

127134
export const aDIReceiverConfigSchema = z.object({

0 commit comments

Comments
 (0)