Skip to content

Commit

Permalink
[#787] Store addresses instead of code hash
Browse files Browse the repository at this point in the history
  • Loading branch information
@akshay-ap
akshay-ap committed Jul 16, 2024
1 parent e5f1af7 commit 6cc9356
Show file tree
Hide file tree
Showing 2 changed files with 40 additions and 200 deletions.
57 changes: 28 additions & 29 deletions contracts/libraries/SafeMigration.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,12 +20,18 @@ contract SafeMigration is SafeStorage {
* @notice Address of this contract
*/
address public immutable MIGRATION_SINGLETON;
// codehash of the Safe singleton implementation
bytes32 public immutable SAFE_SINGLETON_CODEHASH;
// codehash of the Safe singleton (L2) implementation
bytes32 public immutable SAFE_L2_SINGLETON_CODEHASH;
// codehash of the fallbackhandler
bytes32 public immutable SAFE_FALLBACK_HANDLER_CODEHASH;
/**
* @notice Address of the Safe singleton implementation
*/
address public immutable SAFE_SINGLETON;
/**
* @notice Address of the Safe singleton (L2) implementation
*/
address public immutable SAFE_L2_SINGLETON;
/**
* @notice Addresss of the fallbackhandler
*/
address public immutable SAFE_FALLBACK_HANDLER;

/**
* @notice Event indicating a change of master copy address.
Expand All @@ -46,6 +52,7 @@ contract SafeMigration is SafeStorage {
* @notice Constructor
* @param safeSingleton Address of the Safe singleton implementation
* @param safeL2Singleton Address of the SafeL2 singleton implementation
* @param fallbackHandler Address of the fallback handler implmentation
*/
constructor(address safeSingleton, address safeL2Singleton, address fallbackHandler) {
MIGRATION_SINGLETON = address(this);
Expand All @@ -54,52 +61,44 @@ contract SafeMigration is SafeStorage {
require(isContract(safeL2Singleton), "Safe Singleton (L2) is not deployed");
require(isContract(fallbackHandler), "fallback handler is not deployed");

SAFE_SINGLETON_CODEHASH = getCodehash(safeSingleton);
SAFE_L2_SINGLETON_CODEHASH = getCodehash(safeL2Singleton);
SAFE_FALLBACK_HANDLER_CODEHASH = getCodehash(fallbackHandler);
SAFE_SINGLETON = safeSingleton;
SAFE_L2_SINGLETON = safeL2Singleton;
SAFE_FALLBACK_HANDLER = fallbackHandler;
}

/**
* @notice Migrate the Safe contract to a new Safe singleton implementation.
* Checks whether the given target singleton address codehash matches the expected Safe/SafeL2 singleton codehash.
* @param target Address of the new Safe singleton implementation.
*/
function migrateSingleton(address target) public onlyDelegateCall {
require(getCodehash(target) == SAFE_SINGLETON_CODEHASH, "Invalid Safe singleton");
singleton = target;
emit ChangedMasterCopy(target);
function migrateSingleton() public onlyDelegateCall {
singleton = SAFE_SINGLETON;
emit ChangedMasterCopy(SAFE_SINGLETON);
}

/**
* @notice Migrate to Safe Singleton and set the fallback handler. This function is intended to be used when migrating
* a Safe to a version which also requires updating fallback handler.
*/
function migrateWithFallbackHandler(address target, address fallbackHandler) public onlyDelegateCall {
migrateSingleton(target);
require(getCodehash(fallbackHandler) == SAFE_FALLBACK_HANDLER_CODEHASH, "Invalid fallbackhandler");
ISafe(address(this)).setFallbackHandler(fallbackHandler);
function migrateWithFallbackHandler() public onlyDelegateCall {
migrateSingleton();
ISafe(address(this)).setFallbackHandler(SAFE_FALLBACK_HANDLER);
}

/**
* @notice Migrate the Safe contract to a new Safe singleton implementation.
* Checks whether the given target singleton address codehash matches the expected Safe/SafeL2 singleton codehash.
* @param target Address of the new Safe singleton implementation.
*/
function migrateL2Singleton(address target) public onlyDelegateCall {
require(getCodehash(target) == SAFE_L2_SINGLETON_CODEHASH, "Invalid SafeL2 singleton");
singleton = target;
emit ChangedMasterCopy(target);
function migrateL2Singleton() public onlyDelegateCall {
singleton = SAFE_L2_SINGLETON;
emit ChangedMasterCopy(SAFE_L2_SINGLETON);
}

/**
* @notice Migrate to Safe Singleton (L2) and set the fallback handler. This function is intended to be used when migrating
* a Safe to a version which also requires updating fallback handler.
*/
function migrateL2WithFallbackHandler(address target, address fallbackHandler) public onlyDelegateCall {
migrateL2Singleton(target);
require(getCodehash(fallbackHandler) == SAFE_FALLBACK_HANDLER_CODEHASH, "Invalid fallbackhandler");

ISafe(address(this)).setFallbackHandler(fallbackHandler);
function migrateL2WithFallbackHandler() public onlyDelegateCall {
migrateL2Singleton();
ISafe(address(this)).setFallbackHandler(SAFE_FALLBACK_HANDLER);
}

/**
Expand Down
183 changes: 12 additions & 171 deletions test/libraries/SafeMigration.spec.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -111,18 +111,7 @@ describe("SafeMigration Library", () => {
signers: [user1],
} = await setupTests();
await expect(
executeContractCallWithSigners(safe, migration, "migrateSingleton", [SAFE_SINGLETON_ADDRESS], [user1], false),
).to.be.revertedWith("GS013");
});

it("reverts on target singleton codehash mismatch", async () => {
const {
safe,
migration,
signers: [user1],
} = await setupTests();
await expect(
executeContractCallWithSigners(safe, migration, "migrateSingleton", [SAFE_SINGLETON_L2_ADDRESS], [user1], true),
executeContractCallWithSigners(safe, migration, "migrateSingleton", [], [user1], false),
).to.be.revertedWith("GS013");
});

Expand All @@ -136,9 +125,7 @@ describe("SafeMigration Library", () => {
// The emit matcher checks the address, which is the Safe as delegatecall is used
const migrationSafe = migration.attach(safeAddress);

await expect(
executeContractCallWithSigners(safe, migration, "migrateSingleton", [SAFE_SINGLETON_ADDRESS], [user1], true),
)
await expect(executeContractCallWithSigners(safe, migration, "migrateSingleton", [], [user1], true))
.to.emit(migrationSafe, "ChangedMasterCopy")
.withArgs(SAFE_SINGLETON_ADDRESS);

Expand All @@ -160,9 +147,7 @@ describe("SafeMigration Library", () => {
const guardBeforeMigration = await hre.ethers.provider.getStorage(safeAddress, GUARD_STORAGE_SLOT);
const fallbackHandlerBeforeMigration = await hre.ethers.provider.getStorage(safeAddress, FALLBACK_HANDLER_STORAGE_SLOT);

expect(
await executeContractCallWithSigners(safe, migration, "migrateSingleton", [SAFE_SINGLETON_ADDRESS], [user1], true),
);
expect(await executeContractCallWithSigners(safe, migration, "migrateSingleton", [], [user1], true));

expect(await hre.ethers.provider.getStorage(safeAddress, 3)).to.be.eq(ownerCountBeforeMigration);
expect(await hre.ethers.provider.getStorage(safeAddress, 4)).to.be.eq(thresholdBeforeMigration);
Expand All @@ -184,32 +169,7 @@ describe("SafeMigration Library", () => {
signers: [user1],
} = await setupTests();
await expect(
executeContractCallWithSigners(
safe,
migration,
"migrateWithFallbackHandler",
[SAFE_SINGLETON_ADDRESS, COMPATIBILITY_FALLBACK_HANDLER],
[user1],
false,
),
).to.be.revertedWith("GS013");
});

it("reverts on target singleton codehash mismatch", async () => {
const {
safe,
migration,
signers: [user1],
} = await setupTests();
await expect(
executeContractCallWithSigners(
safe,
migration,
"migrateWithFallbackHandler",
[SAFE_SINGLETON_L2_ADDRESS, COMPATIBILITY_FALLBACK_HANDLER],
[user1],
true,
),
executeContractCallWithSigners(safe, migration, "migrateWithFallbackHandler", [], [user1], false),
).to.be.revertedWith("GS013");
});

Expand All @@ -223,16 +183,7 @@ describe("SafeMigration Library", () => {
// The emit matcher checks the address, which is the Safe as delegatecall is used
const migrationSafe = migration.attach(safeAddress);

await expect(
executeContractCallWithSigners(
safe,
migration,
"migrateWithFallbackHandler",
[SAFE_SINGLETON_ADDRESS, COMPATIBILITY_FALLBACK_HANDLER],
[user1],
true,
),
)
await expect(executeContractCallWithSigners(safe, migration, "migrateWithFallbackHandler", [], [user1], true))
.to.emit(migrationSafe, "ChangedMasterCopy")
.withArgs(SAFE_SINGLETON_ADDRESS)
.and.to.emit(safe, "ChangedFallbackHandler")
Expand All @@ -259,16 +210,7 @@ describe("SafeMigration Library", () => {
const nonceBeforeMigration = await hre.ethers.provider.getStorage(safeAddress, 5);
const guardBeforeMigration = await hre.ethers.provider.getStorage(safeAddress, GUARD_STORAGE_SLOT);

expect(
await executeContractCallWithSigners(
safe,
migration,
"migrateWithFallbackHandler",
[SAFE_SINGLETON_ADDRESS, COMPATIBILITY_FALLBACK_HANDLER],
[user1],
true,
),
);
expect(await executeContractCallWithSigners(safe, migration, "migrateWithFallbackHandler", [], [user1], true));

expect(await hre.ethers.provider.getStorage(safeAddress, 3)).to.be.eq(ownerCountBeforeMigration);
expect(await hre.ethers.provider.getStorage(safeAddress, 4)).to.be.eq(thresholdBeforeMigration);
Expand All @@ -287,36 +229,7 @@ describe("SafeMigration Library", () => {
signers: [user1],
} = await setupTests();
await expect(
executeContractCallWithSigners(safe, migration, "migrateL2Singleton", [SAFE_SINGLETON_L2_ADDRESS], [user1], false),
).to.be.revertedWith("GS013");
});

it("reverts on target singleton codehash mismatch", async () => {
const {
safe,
migration,
signers: [user1],
} = await setupTests();
await expect(
executeContractCallWithSigners(safe, migration, "migrateL2Singleton", [SAFE_SINGLETON_ADDRESS], [user1], true),
).to.be.revertedWith("GS013");
});

it("reverts on fallback handler codehash mismatch", async () => {
const {
safe,
migration,
signers: [user1],
} = await setupTests();
await expect(
executeContractCallWithSigners(
safe,
migration,
"migrateWithFallbackHandler",
[SAFE_SINGLETON_ADDRESS, SAFE_SINGLETON_L2_ADDRESS],
[user1],
true,
),
executeContractCallWithSigners(safe, migration, "migrateL2Singleton", [], [user1], false),
).to.be.revertedWith("GS013");
});

Expand All @@ -330,9 +243,7 @@ describe("SafeMigration Library", () => {
// The emit matcher checks the address, which is the Safe as delegatecall is used
const migrationSafe = migration.attach(safeAddress);

await expect(
executeContractCallWithSigners(safeL2, migration, "migrateL2Singleton", [SAFE_SINGLETON_L2_ADDRESS], [user1], true),
)
await expect(executeContractCallWithSigners(safeL2, migration, "migrateL2Singleton", [], [user1], true))
.to.emit(migrationSafe, "ChangedMasterCopy")
.withArgs(SAFE_SINGLETON_L2_ADDRESS);

Expand All @@ -354,16 +265,7 @@ describe("SafeMigration Library", () => {
const guardBeforeMigration = await hre.ethers.provider.getStorage(safeAddress, GUARD_STORAGE_SLOT);
const fallbackHandlerBeforeMigration = await hre.ethers.provider.getStorage(safeAddress, FALLBACK_HANDLER_STORAGE_SLOT);

expect(
await executeContractCallWithSigners(
safeL2,
migration,
"migrateL2Singleton",
[SAFE_SINGLETON_L2_ADDRESS],
[user1],
true,
),
);
expect(await executeContractCallWithSigners(safeL2, migration, "migrateL2Singleton", [], [user1], true));

expect(await hre.ethers.provider.getStorage(safeAddress, 3)).to.be.eq(ownerCountBeforeMigration);
expect(await hre.ethers.provider.getStorage(safeAddress, 4)).to.be.eq(thresholdBeforeMigration);
Expand All @@ -385,50 +287,7 @@ describe("SafeMigration Library", () => {
signers: [user1],
} = await setupTests();
await expect(
executeContractCallWithSigners(
safe,
migration,
"migrateL2WithFallbackHandler",
[SAFE_SINGLETON_L2_ADDRESS, COMPATIBILITY_FALLBACK_HANDLER],
[user1],
false,
),
).to.be.revertedWith("GS013");
});

it("reverts on target singleton codehash mismatch", async () => {
const {
safe,
migration,
signers: [user1],
} = await setupTests();
await expect(
executeContractCallWithSigners(
safe,
migration,
"migrateL2WithFallbackHandler",
[SAFE_SINGLETON_ADDRESS, COMPATIBILITY_FALLBACK_HANDLER],
[user1],
true,
),
).to.be.revertedWith("GS013");
});

it("reverts on fallback handler codehash mismatch", async () => {
const {
safe,
migration,
signers: [user1],
} = await setupTests();
await expect(
executeContractCallWithSigners(
safe,
migration,
"migrateL2WithFallbackHandler",
[SAFE_SINGLETON_L2_ADDRESS, SAFE_SINGLETON_ADDRESS],
[user1],
true,
),
executeContractCallWithSigners(safe, migration, "migrateL2WithFallbackHandler", [], [user1], false),
).to.be.revertedWith("GS013");
});

Expand All @@ -442,16 +301,7 @@ describe("SafeMigration Library", () => {
// The emit matcher checks the address, which is the Safe as delegatecall is used
const migrationSafe = migration.attach(safeAddress);

await expect(
executeContractCallWithSigners(
safeL2,
migration,
"migrateL2WithFallbackHandler",
[SAFE_SINGLETON_L2_ADDRESS, COMPATIBILITY_FALLBACK_HANDLER],
[user1],
true,
),
)
await expect(executeContractCallWithSigners(safeL2, migration, "migrateL2WithFallbackHandler", [], [user1], true))
.to.emit(migrationSafe, "ChangedMasterCopy")
.withArgs(SAFE_SINGLETON_L2_ADDRESS)
.and.to.emit(safeL2, "ChangedFallbackHandler")
Expand All @@ -478,16 +328,7 @@ describe("SafeMigration Library", () => {
const nonceBeforeMigration = await hre.ethers.provider.getStorage(safeAddress, 5);
const guardBeforeMigration = await hre.ethers.provider.getStorage(safeAddress, GUARD_STORAGE_SLOT);

expect(
await executeContractCallWithSigners(
safeL2,
migration,
"migrateL2WithFallbackHandler",
[SAFE_SINGLETON_L2_ADDRESS, COMPATIBILITY_FALLBACK_HANDLER],
[user1],
true,
),
);
expect(await executeContractCallWithSigners(safeL2, migration, "migrateL2WithFallbackHandler", [], [user1], true));

expect(await hre.ethers.provider.getStorage(safeAddress, 3)).to.be.eq(ownerCountBeforeMigration);
expect(await hre.ethers.provider.getStorage(safeAddress, 4)).to.be.eq(thresholdBeforeMigration);
Expand Down

0 comments on commit 6cc9356

Please sign in to comment.