Skip to content

Commit a7edda4

Browse files
0xCourtney0xCourtney
committed
diamond utils draft
1 parent f204227 commit a7edda4

File tree

3 files changed

+329
-0
lines changed

3 files changed

+329
-0
lines changed

‎lib/diamonds.ts

+327
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,327 @@
1+
import {
2+
IDiamondReadable,
3+
IDiamondWritable,
4+
} from '@solidstate/typechain-types';
5+
import { Contract, constants } from 'ethers';
6+
7+
const Table = require('cli-table3');
8+
9+
type Diamond = IDiamondReadable & IDiamondWritable;
10+
11+
export interface Facet {
12+
target: string;
13+
selectors: string[];
14+
}
15+
16+
enum FacetCutAction {
17+
Add,
18+
Replace,
19+
Remove,
20+
}
21+
22+
export interface FacetCut extends Facet {
23+
action: FacetCutAction;
24+
}
25+
26+
export function getSignatures(contract: Contract): string[] {
27+
return Object.keys(contract.interface.functions);
28+
}
29+
30+
export function getSelectors(contract: Contract): string[] {
31+
const signatures = getSignatures(contract);
32+
return signatures.reduce((acc: string[], val: string) => {
33+
acc.push(contract.interface.getSighash(val));
34+
return acc;
35+
}, []);
36+
}
37+
38+
export function getFacets(contracts: Contract[]): Facet[] {
39+
return contracts.map((contract) => {
40+
return {
41+
target: contract.address,
42+
selectors: getSelectors(contract),
43+
};
44+
});
45+
}
46+
47+
export function selectorExistsInFacets(
48+
selector: string,
49+
facets: Facet[],
50+
): boolean {
51+
for (const facet of facets) {
52+
if (facet.selectors.includes(selector)) return true;
53+
}
54+
55+
return false;
56+
}
57+
58+
// adds unregistered selectors
59+
export async function addUnregisteredSelectors(
60+
diamond: Diamond,
61+
contracts: Contract[],
62+
exclude: string[] = [],
63+
): Promise<FacetCut[]> {
64+
const diamondFacets: Facet[] = await diamond.facets();
65+
const facets = getFacets(contracts);
66+
let facetCuts: FacetCut[] = [];
67+
68+
// if facet selector is unregistered then it should be added to the diamond.
69+
for (const facet of facets) {
70+
for (const selector of facet.selectors) {
71+
const target = facet.target;
72+
73+
if (
74+
target !== diamond.address &&
75+
selector.length > 0 &&
76+
!selectorExistsInFacets(selector, diamondFacets) &&
77+
!exclude.includes(selector)
78+
) {
79+
facetCuts.push(
80+
printFacetCuts(facet.target, [selector], FacetCutAction.Add),
81+
);
82+
}
83+
}
84+
}
85+
86+
return groupFacetCuts(facetCuts);
87+
}
88+
89+
// replace registered selectors
90+
export async function replaceRegisteredSelectors(
91+
diamond: Diamond,
92+
contracts: Contract[],
93+
exclude: string[] = [],
94+
): Promise<FacetCut[]> {
95+
const diamondFacets: Facet[] = await diamond.facets();
96+
const facets = getFacets(contracts);
97+
let facetCuts: FacetCut[] = [];
98+
99+
// if a facet selector is registered with a different target address, the target will be replaced
100+
for (const facet of facets) {
101+
for (const selector of facet.selectors) {
102+
const target = facet.target;
103+
const oldTarget = await diamond.facetAddress(selector);
104+
105+
if (
106+
target != oldTarget &&
107+
target != constants.AddressZero &&
108+
target != diamond.address &&
109+
selector.length > 0 &&
110+
selectorExistsInFacets(selector, diamondFacets) &&
111+
!exclude.includes(selector)
112+
) {
113+
facetCuts.push(
114+
printFacetCuts(target, [selector], FacetCutAction.Replace),
115+
);
116+
}
117+
}
118+
}
119+
120+
return groupFacetCuts(facetCuts);
121+
}
122+
123+
// removes registered selectors
124+
export async function removeRegisteredSelectors(
125+
diamond: Diamond,
126+
contracts: Contract[],
127+
exclude: string[] = [],
128+
): Promise<FacetCut[]> {
129+
const diamondFacets: Facet[] = await diamond.facets();
130+
const facets = getFacets(contracts);
131+
let facetCuts: FacetCut[] = [];
132+
133+
// if a registered selector is not found in the facets then it should be removed from the diamond
134+
for (const diamondFacet of diamondFacets) {
135+
for (const selector of diamondFacet.selectors) {
136+
const target = diamondFacet.target;
137+
138+
if (
139+
target != constants.AddressZero &&
140+
target != diamond.address &&
141+
selector.length > 0 &&
142+
!selectorExistsInFacets(selector, facets) &&
143+
!exclude.includes(selector)
144+
) {
145+
facetCuts.push(
146+
printFacetCuts(
147+
constants.AddressZero,
148+
[selector],
149+
FacetCutAction.Remove,
150+
),
151+
);
152+
}
153+
}
154+
}
155+
156+
return groupFacetCuts(facetCuts);
157+
}
158+
159+
export async function diamondCut(
160+
diamond: Diamond,
161+
facetCut: FacetCut[],
162+
target: string = constants.AddressZero,
163+
data: string = '0x',
164+
) {
165+
(await diamond.diamondCut(facetCut, target, data)).wait(1);
166+
}
167+
168+
// groups facet cuts by target address and action type
169+
export function groupFacetCuts(facetCuts: FacetCut[]): FacetCut[] {
170+
const cuts = facetCuts.reduce((acc: FacetCut[], facetCut: FacetCut) => {
171+
if (acc.length == 0) acc.push(facetCut);
172+
173+
let exists = false;
174+
175+
acc.forEach((_, i) => {
176+
if (
177+
acc[i].action == facetCut.action &&
178+
acc[i].target == facetCut.target
179+
) {
180+
acc[i].selectors.push(...facetCut.selectors);
181+
// removes duplicates, if there are any
182+
acc[i].selectors = [...new Set(acc[i].selectors)];
183+
exists = true;
184+
}
185+
});
186+
187+
// push facet cut if it does not already exist
188+
if (!exists) acc.push(facetCut);
189+
190+
return acc;
191+
}, []);
192+
193+
let cache: any = {};
194+
195+
// checks if selector is used multiple times, emits warning
196+
cuts.forEach((cut) => {
197+
cut.selectors.forEach((selector: string) => {
198+
if (cache[selector]) {
199+
console.log(
200+
`WARNING: selector: ${selector}, target: ${cut.target} is defined in multiple cuts`,
201+
);
202+
} else {
203+
cache[selector] = true;
204+
}
205+
});
206+
});
207+
208+
return cuts;
209+
}
210+
211+
export function printFacetCuts(
212+
target: string,
213+
selectors: string[],
214+
action: number = 0,
215+
): FacetCut {
216+
return {
217+
target: target,
218+
action: action,
219+
selectors: selectors,
220+
};
221+
}
222+
223+
// generates table of diamond and facet selectors
224+
export async function printDiamond(diamond: Diamond, contracts: Contract[]) {
225+
const padding = 2;
226+
227+
const table = new Table({
228+
style: {
229+
head: [],
230+
border: [],
231+
'padding-left': padding,
232+
'padding-right': padding,
233+
},
234+
chars: {
235+
mid: '·',
236+
'top-mid': '|',
237+
'left-mid': ' ·',
238+
'mid-mid': '|',
239+
'right-mid': '·',
240+
left: ' |',
241+
'top-left': ' ·',
242+
'top-right': '·',
243+
'bottom-left': ' ·',
244+
'bottom-right': '·',
245+
middle: '·',
246+
top: '-',
247+
bottom: '-',
248+
'bottom-mid': '|',
249+
},
250+
});
251+
252+
table.push([
253+
{
254+
hAlign: 'center',
255+
content: `Target`,
256+
},
257+
{
258+
hAlign: 'center',
259+
content: `Signature`,
260+
},
261+
{
262+
hAlign: 'center',
263+
content: `Selector`,
264+
},
265+
{
266+
hAlign: 'center',
267+
content: `Registered`,
268+
},
269+
]);
270+
271+
let diamondTable = [];
272+
const signatures = await getSignatures(diamond);
273+
274+
for (const signature of signatures) {
275+
diamondTable.push({
276+
target: diamond.address,
277+
signature: signature,
278+
selector: diamond.interface.getSighash(signature),
279+
registered: true,
280+
});
281+
}
282+
283+
for (const contract of contracts) {
284+
const signatures = await getSignatures(contract);
285+
for (const signature of signatures) {
286+
diamondTable.push({
287+
target: contract.address,
288+
signature: signature,
289+
selector: contract.interface.getSighash(signature),
290+
registered: false,
291+
});
292+
}
293+
}
294+
295+
const diamondFacets: Facet[] = await diamond.facets();
296+
297+
for (const facet of diamondFacets) {
298+
const target = facet.target;
299+
300+
for (const selector of facet.selectors) {
301+
for (const row of diamondTable) {
302+
if (row.target == target && row.selector == selector) {
303+
row.registered = true;
304+
}
305+
}
306+
}
307+
}
308+
309+
for (const row of diamondTable) {
310+
table.push([
311+
{
312+
content: row.target,
313+
},
314+
{
315+
content: row.signature,
316+
},
317+
{
318+
content: row.selector,
319+
},
320+
{
321+
content: row.registered,
322+
},
323+
]);
324+
}
325+
326+
console.log(table.toString());
327+
}

‎lib/index.ts

+1
Original file line numberDiff line numberDiff line change
@@ -2,3 +2,4 @@ export * from './bn_conversion';
22
export * from './erc20_permit';
33
export * from './mocha_describe_filter';
44
export * from './sign_data';
5+
export * from './diamonds';

‎lib/package.json

+1
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
"tsc-clean": "tsc --build --clean tsconfig.json"
2626
},
2727
"dependencies": {
28+
"cli-table3": "^0.6.3",
2829
"eth-permit": "^0.1.10"
2930
},
3031
"files": [

0 commit comments

Comments
 (0)