Add BLS Public Keys field in staking contract (#26)

* Add validator BLS PublicKey in contract * Fix type of _addressToBLSPublicKey * Fix script * Add test for registerBLSPublicKey and validatorBLSPublicKeys
0xPolygon · Sep 5, 2022 · 436f108 · 436f108
1 parent c5ccf1a
commit 436f108
Show file tree

Hide file tree

Showing 7 changed files with 149 additions and 29 deletions.
diff --git a/contracts/Staking.sol b/contracts/Staking.sol
@@ -10,18 +10,23 @@ contract Staking {
 
     // Properties
     address[] public _validators;
+
     mapping(address => bool) public _addressToIsValidator;
     mapping(address => uint256) public _addressToStakedAmount;
     mapping(address => uint256) public _addressToValidatorIndex;
     uint256 public _stakedAmount;
     uint256 public _minimumNumValidators;
     uint256 public _maximumNumValidators;
 
+    mapping(address => bytes) public _addressToBLSPublicKey;
+
     // Events
     event Staked(address indexed account, uint256 amount);
 
     event Unstaked(address indexed account, uint256 amount);
 
+    event BLSPublicKeyRegistered(address indexed accout, bytes key);
+
     // Modifiers
     modifier onlyEOA() {
         require(!msg.sender.isContract(), "Only EOA can call function");
@@ -36,6 +41,11 @@ contract Staking {
         _;
     }
 
+    modifier onlyValidator() {
+        require(_isValidator(msg.sender), "Only validator can call function");
+        _;
+    }
+
     constructor(uint256 minNumValidators, uint256 maxNumValidators) {
         require(
             minNumValidators <= maxNumValidators,
@@ -54,6 +64,16 @@ contract Staking {
         return _validators;
     }
 
+    function validatorBLSPublicKeys() public view returns (bytes[] memory) {
+        bytes[] memory keys = new bytes[](_validators.length);
+
+        for (uint256 i = 0; i < _validators.length; i++) {
+            keys[i] = _addressToBLSPublicKey[_validators[i]];
+        }
+
+        return keys;
+    }
+
     function isValidator(address addr) public view returns (bool) {
         return _addressToIsValidator[addr];
     }
@@ -83,6 +103,12 @@ contract Staking {
         _unstake();
     }
 
+    function registerBLSPublicKey(bytes memory blsPubKey) public {
+        _addressToBLSPublicKey[msg.sender] = blsPubKey;
+
+        emit BLSPublicKeyRegistered(msg.sender, blsPubKey);
+    }
+
     // Private functions
     function _stake() private {
         _stakedAmount += msg.value;

diff --git a/package.json b/package.json
@@ -10,6 +10,7 @@
     "stake": "hardhat run scripts/stake.ts --network polygonedge",
     "transfer": "hardhat run scripts/transfer.ts --network polygonedge",
     "unstake": "hardhat run scripts/unstake.ts --network polygonedge",
+    "register-blskey": "hardhat run scripts/register-bls-key.ts --network polygonedge",
     "info": "hardhat run scripts/info.ts --network polygonedge",
     "test": "npx hardhat test",
     "clean": "npx hardhat clean",

diff --git a/scripts/info.ts b/scripts/info.ts
@@ -1,26 +1,35 @@
-import {ethers} from "hardhat";
-import {Staking} from "../types/Staking";
+import { ethers } from "hardhat";
+import { Staking } from "../types/Staking";
 
-const STAKING_CONTRACT_ADDRESS = process.env.STAKING_CONTRACT_ADDRESS ?? '';
+const STAKING_CONTRACT_ADDRESS = process.env.STAKING_CONTRACT_ADDRESS ?? "";
 
 async function main() {
   console.log("Check current contract information");
 
-  const StakingContractFactory = await ethers.getContractFactory("Staking");
-  const stakingContract = await StakingContractFactory.attach(STAKING_CONTRACT_ADDRESS) as Staking;
+  const stakingContract = (await ethers.getContractAt(
+    "Staking",
+    STAKING_CONTRACT_ADDRESS
+  )) as Staking;
 
-  const [stakedAmount, validators, minimumNumValidators, maximumNumValidators] = await Promise.all([
+  const [
+    stakedAmount,
+    validators,
+    blsPublicKeys,
+    minimumNumValidators,
+    maximumNumValidators,
+  ] = await Promise.all([
     stakingContract.stakedAmount(),
     stakingContract.validators(),
+    stakingContract.validatorBLSPublicKeys(),
     stakingContract.minimumNumValidators(),
     stakingContract.maximumNumValidators(),
-  ])
-
-  console.log(`Total staked amount: ${stakedAmount.toString()}`)
-  console.log('Minimum number of validators', minimumNumValidators.toNumber());
-  console.log('Maximum number of validators', maximumNumValidators.toNumber());
-  console.log('Current validators list', validators);
+  ]);
 
+  console.log(`Total staked amount: ${stakedAmount.toString()}`);
+  console.log("Minimum number of validators", minimumNumValidators.toNumber());
+  console.log("Maximum number of validators", maximumNumValidators.toNumber());
+  console.log("Current validators list", validators);
+  console.log("BLS Public Keys", blsPublicKeys);
 }
 
 main()

diff --git a/scripts/register-bls-key.ts b/scripts/register-bls-key.ts
@@ -0,0 +1,32 @@
+import { ethers } from "hardhat";
+import { Staking } from "../types/Staking";
+
+const STAKING_CONTRACT_ADDRESS = process.env.STAKING_CONTRACT_ADDRESS ?? "";
+const BLS_PUBLIC_KEY = process.env.BLS_PUBLIC_KEY ?? "";
+
+async function main() {
+  const [account] = await ethers.getSigners();
+
+  console.log(
+    `Register BLS Public Key: address=${STAKING_CONTRACT_ADDRESS}, account=${account.address}, key=${BLS_PUBLIC_KEY}`
+  );
+  console.log(`Account balance: ${(await account.getBalance()).toString()}`);
+
+  const stakingContract = (await ethers.getContractAt(
+    "Staking",
+    STAKING_CONTRACT_ADDRESS,
+    account
+  )) as Staking;
+
+  const tx = await stakingContract.registerBLSPublicKey(BLS_PUBLIC_KEY);
+  const receipt = await tx.wait();
+
+  console.log("Registered", tx.hash, receipt);
+}
+
+main()
+  .then(() => process.exit(0))
+  .catch((error) => {
+    console.error(error);
+    process.exit(1);
+  });
diff --git a/scripts/stake.ts b/scripts/stake.ts
@@ -1,20 +1,24 @@
-import {ethers} from "hardhat";
-import {Staking} from "../types/Staking";
+import { ethers } from "hardhat";
+import { Staking } from "../types/Staking";
 
-const STAKING_CONTRACT_ADDRESS = process.env.STAKING_CONTRACT_ADDRESS ?? '';
-const STAKE_AMOUNT = ethers.utils.parseEther("1")
+const STAKING_CONTRACT_ADDRESS = process.env.STAKING_CONTRACT_ADDRESS ?? "";
+const STAKE_AMOUNT = ethers.utils.parseEther("1");
 
 async function main() {
   const [account] = await ethers.getSigners();
 
-  console.log(`Stake: address=${STAKING_CONTRACT_ADDRESS}, account=${account.address}`);
+  console.log(
+    `Stake: address=${STAKING_CONTRACT_ADDRESS}, account=${account.address}`
+  );
   console.log(`Account balance: ${(await account.getBalance()).toString()}`);
 
-  const StakingContractFactory = await ethers.getContractFactory("Staking");
-  let stakingContract = await StakingContractFactory.attach(STAKING_CONTRACT_ADDRESS) as Staking;
-  stakingContract = stakingContract.connect(account);
+  const stakingContract = (await ethers.getContractAt(
+    "Staking",
+    STAKING_CONTRACT_ADDRESS,
+    account
+  )) as Staking;
 
-  const tx = await stakingContract.stake({ value: STAKE_AMOUNT })
+  const tx = await stakingContract.stake({ value: STAKE_AMOUNT });
   const receipt = await tx.wait();
 
   console.log("Staked", tx.hash, receipt);

diff --git a/scripts/unstake.ts b/scripts/unstake.ts
@@ -1,19 +1,23 @@
-import {ethers} from "hardhat";
-import {Staking} from "../types/Staking";
+import { ethers } from "hardhat";
+import { Staking } from "../types/Staking";
 
-const STAKING_CONTRACT_ADDRESS = process.env.STAKING_CONTRACT_ADDRESS ?? '';
+const STAKING_CONTRACT_ADDRESS = process.env.STAKING_CONTRACT_ADDRESS ?? "";
 
 async function main() {
   const [account] = await ethers.getSigners();
 
-  console.log(`Unstake: contract=${STAKING_CONTRACT_ADDRESS}, account=${account.address}`);
+  console.log(
+    `Unstake: contract=${STAKING_CONTRACT_ADDRESS}, account=${account.address}`
+  );
   console.log("Account balance", (await account.getBalance()).toString());
 
-  const StakingContractFactory = await ethers.getContractFactory("Staking");
-  let stakingContract = await StakingContractFactory.attach(STAKING_CONTRACT_ADDRESS) as Staking;
-  stakingContract = stakingContract.connect(account);
+  const stakingContract = (await ethers.getContractAt(
+    "Staking",
+    STAKING_CONTRACT_ADDRESS,
+    account
+  )) as Staking;
 
-  const tx = await stakingContract.unstake()
+  const tx = await stakingContract.unstake();
   const receipt = await tx.wait();
 
   console.log("Unstaked", tx.hash, receipt);

diff --git a/test/Staking.ts b/test/Staking.ts
@@ -301,4 +301,48 @@ describe("Staking and Unstaking", function () {
       );
     });
   });
+
+  describe("Register BLS Public Key", () => {
+    it("should succeed and register BLS Public Key", async () => {
+      const data = "0x12345678"
+      const tx = contract.connect(accounts[0]).registerBLSPublicKey(
+        data,
+      )
+
+      await expect(tx)
+      .to.emit(contract, "BLSPublicKeyRegistered")
+      .withArgs(accounts[0].address, data);
+    })
+  });
+
+  describe("Get BLS Public Keys", () => {
+    const numValidators = 5;
+    const numRegisteredAccounts = 10;
+    const stakedAmount = ethers.utils.parseEther("1");
+    const blsPublicKeys = new Array(numRegisteredAccounts).fill(null).map((_, idx) => (
+      BigNumber.from(idx).toHexString()
+    ))
+
+    beforeEach(async () => {
+      // set required number of validators
+      await Promise.all(
+        accounts
+          .slice(0, numValidators)
+          .map((account) => stake(account, stakedAmount))
+      );
+
+      await Promise.all(
+        accounts.splice(0, numRegisteredAccounts)
+        .map((account, i) => contract.connect(account).registerBLSPublicKey(
+          blsPublicKeys[i]
+        ))
+      )
+    });
+
+    it("should return only the BLS Public Keys of Validators", async () => {
+      expect(await contract.validatorBLSPublicKeys())
+        .to.have.length(numValidators)
+        .and.to.deep.equal(blsPublicKeys.slice(0, numValidators));
+    })
+  });
 });