Make NodeDriver upgradable

contracts/sfc/NetworkInitializer.sol

@@ -21,7 +21,7 @@ contract NetworkInitializer {
         address _evmWriter,
         address _owner
     ) external {
-        NodeDriver(_driver).initialize(_auth, _evmWriter);
+        NodeDriver(_driver).initialize(_auth, _evmWriter, _owner);
         NodeDriverAuth(_auth).initialize(_sfc, _driver, _owner);
 
         ConstantsManager consts = new ConstantsManager();

contracts/sfc/NodeDriver.sol

@@ -1,7 +1,9 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity 0.8.27;
 
+import {UUPSUpgradeable} from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
 import {Initializable} from "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
+import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
 import {NodeDriverAuth} from "./NodeDriverAuth.sol";
 import {IEVMWriter} from "../interfaces/IEVMWriter.sol";
 import {INodeDriver} from "../interfaces/INodeDriver.sol";
@@ -12,7 +14,7 @@ import {INodeDriver} from "../interfaces/INodeDriver.sol";
  * @dev Methods with onlyNode modifier are called by Sonic internal txs during epoch sealing.
  * @custom:security-contact [email protected]
  */
-contract NodeDriver is Initializable, INodeDriver {
+contract NodeDriver is Initializable, OwnableUpgradeable, UUPSUpgradeable, INodeDriver {
     NodeDriverAuth internal backend;
     IEVMWriter internal evmWriter;
 
@@ -44,12 +46,18 @@ contract NodeDriver is Initializable, INodeDriver {
 
     /// Initialization is called only once, after the contract deployment.
     /// Because the contract code is written directly into genesis, constructor cannot be used.
-    function initialize(address _backend, address _evmWriterAddress) external initializer {
+    function initialize(address _backend, address _evmWriterAddress, address _owner) external initializer {
+        __Ownable_init(_owner);
+        __UUPSUpgradeable_init();
         backend = NodeDriverAuth(_backend);
         emit UpdatedBackend(_backend);
         evmWriter = IEVMWriter(_evmWriterAddress);
     }
 
+    /// Override the upgrade authorization check to allow upgrades only from the owner.
+    // solhint-disable-next-line no-empty-blocks
+    function _authorizeUpgrade(address) internal override onlyOwner {}
+
     function setBalance(address acc, uint256 value) external onlyBackend {
         evmWriter.setBalance(acc, value);
     }

contracts/test/UnitTestSFC.sol

@@ -58,7 +58,7 @@ contract UnitTestNetworkInitializer {
         address _evmWriter,
         address _owner
     ) external {
-        NodeDriver(_driver).initialize(_auth, _evmWriter);
+        NodeDriver(_driver).initialize(_auth, _evmWriter, _owner);
         NodeDriverAuth(_auth).initialize(_sfc, _driver, _owner);
 
         UnitTestConstantsManager consts = new UnitTestConstantsManager();

test/NodeDriver.ts

@@ -1,7 +1,7 @@
 import { ethers, upgrades } from 'hardhat';
 import { expect } from 'chai';
 import { loadFixture } from '@nomicfoundation/hardhat-network-helpers';
-import { IEVMWriter, NetworkInitializer, NodeDriver, NodeDriverAuth } from '../typechain-types';
+import { IEVMWriter, NetworkInitializer, NodeDriverAuth } from '../typechain-types';
 
 describe('NodeDriver', () => {
   const fixture = async () => {
@@ -10,7 +10,10 @@ describe('NodeDriver', () => {
       kind: 'uups',
       initializer: false,
     });
-    const nodeDriver: NodeDriver = await ethers.deployContract('NodeDriver');
+    const nodeDriver = await upgrades.deployProxy(await ethers.getContractFactory('NodeDriver'), {
+      kind: 'uups',
+      initializer: false,
+    });
     const evmWriter: IEVMWriter = await ethers.deployContract('StubEvmWriter');
     const nodeDriverAuth: NodeDriverAuth = await ethers.deployContract('NodeDriverAuth');
     const initializer: NetworkInitializer = await ethers.deployContract('NetworkInitializer');

test/Proxy.ts

@@ -84,3 +84,77 @@ describe('SFC', () => {
     });
   });
 });
+
+describe('NodeDriver', () => {
+  const fixture = async () => {
+    const [user, owner] = await ethers.getSigners();
+    const nodeDriver = await upgrades.deployProxy(await ethers.getContractFactory('NodeDriver'), {
+      kind: 'uups',
+      initializer: false,
+    });
+
+    const backend = ethers.Wallet.createRandom();
+    const evmWriter = ethers.Wallet.createRandom();
+
+    await nodeDriver.initialize(backend, evmWriter, owner);
+
+    return {
+      owner,
+      user,
+      nodeDriver,
+      backend,
+      evmWriter,
+    };
+  };
+
+  beforeEach(async function () {
+    Object.assign(this, await loadFixture(fixture));
+  });
+
+  describe('Initialization', () => {
+    it('Should succeed and initialize', async function () {
+      expect(await this.nodeDriver.owner()).to.equal(this.owner);
+    });
+
+    it('Should revert when already initialized', async function () {
+      await expect(this.nodeDriver.initialize(this.backend, this.evmWriter, this.owner)).to.be.revertedWithCustomError(
+        this.nodeDriver,
+        'InvalidInitialization',
+      );
+    });
+
+    describe('Upgrade', () => {
+      it('Should revert when not owner', async function () {
+        await expect(
+          upgrades.upgradeProxy(this.nodeDriver, (await ethers.getContractFactory('NodeDriver')).connect(this.user)),
+        ).to.be.revertedWithCustomError(this.nodeDriver, 'OwnableUnauthorizedAccount');
+      });
+
+      it('Should succeed and upgrade', async function () {
+        // get the implementation address
+        // the address is stored at slot keccak-256 hash of "eip1967.proxy.implementation" subtracted by 1
+        const implementationAddr = await ethers.provider.getStorage(
+          this.nodeDriver,
+          '0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc',
+        );
+
+        // deploy new implementation contract
+        const newImpl = await ethers.deployContract('NodeDriver');
+
+        // upgrade proxy with new implementation
+        await this.nodeDriver.connect(this.owner).upgradeToAndCall(newImpl, '0x');
+
+        // check that the implementation address changed
+        const newImplementationAddress = await ethers.provider.getStorage(
+          this.nodeDriver,
+          '0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc',
+        );
+
+        expect(newImplementationAddress).to.not.equal(implementationAddr);
+        expect(ethers.AbiCoder.defaultAbiCoder().decode(['address'], newImplementationAddress)[0]).to.equal(
+          await newImpl.getAddress(),
+        );
+      });
+    });
+  });
+});

test/SFC.ts

@@ -1,13 +1,7 @@
 import { ethers, upgrades } from 'hardhat';
 import { expect } from 'chai';
 import { loadFixture } from '@nomicfoundation/hardhat-network-helpers';
-import {
-  IEVMWriter,
-  NodeDriver,
-  NodeDriverAuth,
-  UnitTestConstantsManager,
-  UnitTestNetworkInitializer,
-} from '../typechain-types';
+import { IEVMWriter, NodeDriverAuth, UnitTestConstantsManager, UnitTestNetworkInitializer } from '../typechain-types';
 import { beforeEach, Context } from 'mocha';
 import { BlockchainNode, ValidatorMetrics } from './helpers/BlockchainNode';
 
@@ -18,7 +12,10 @@ describe('SFC', () => {
       kind: 'uups',
       initializer: false,
     });
-    const nodeDriver: NodeDriver = await ethers.deployContract('NodeDriver');
+    const nodeDriver = await upgrades.deployProxy(await ethers.getContractFactory('NodeDriver'), {
+      kind: 'uups',
+      initializer: false,
+    });
     const evmWriter: IEVMWriter = await ethers.deployContract('StubEvmWriter');
     const nodeDriverAuth: NodeDriverAuth = await ethers.deployContract('NodeDriverAuth');
     const initializer: UnitTestNetworkInitializer = await ethers.deployContract('UnitTestNetworkInitializer');