Implementation of the Stake table contract (register function)

contract-bindings/src/lib.rs

@@ -16,5 +16,6 @@ pub mod i_stake_table;
 pub mod plonk_verifier;
 pub mod polynomial_eval;
 pub mod shared_types;
+pub mod stake_table;
 pub mod transcript;
 pub mod utils;

contract-bindings/src/shared_types.rs

@@ -34,6 +34,23 @@ pub struct G2Point {
     pub y_0: ::ethers::core::types::U256,
     pub y_1: ::ethers::core::types::U256,
 }
+///`EdOnBN254Point(uint256,uint256)`
+#[derive(
+    Clone,
+    ::ethers::contract::EthAbiType,
+    ::ethers::contract::EthAbiCodec,
+    serde::Serialize,
+    serde::Deserialize,
+    Default,
+    Debug,
+    PartialEq,
+    Eq,
+    Hash,
+)]
+pub struct EdOnBN254Point {
+    pub x: ::ethers::core::types::U256,
+    pub y: ::ethers::core::types::U256,
+}
 ///`PlonkProof((uint256,uint256),(uint256,uint256),(uint256,uint256),(uint256,uint256),(uint256,uint256),(uint256,uint256),(uint256,uint256),(uint256,uint256),(uint256,uint256),(uint256,uint256),(uint256,uint256),(uint256,uint256),(uint256,uint256),uint256,uint256,uint256,uint256,uint256,uint256,uint256,uint256,uint256,uint256)`
 #[derive(
     Clone,
@@ -107,3 +124,24 @@ pub struct VerifyingKey {
     pub q_h4: G1Point,
     pub q_ecc: G1Point,
 }
+///`Node(address,uint8,uint64,uint64,uint64,(uint256,uint256))`
+#[derive(
+    Clone,
+    ::ethers::contract::EthAbiType,
+    ::ethers::contract::EthAbiCodec,
+    serde::Serialize,
+    serde::Deserialize,
+    Default,
+    Debug,
+    PartialEq,
+    Eq,
+    Hash,
+)]
+pub struct Node {
+    pub account: ::ethers::core::types::Address,
+    pub stake_type: u8,
+    pub balance: u64,
+    pub register_epoch: u64,
+    pub exit_epoch: u64,
+    pub schnorr_vk: EdOnBN254Point,
+}

contracts/rust/src/bin/diff_test.rs

@@ -26,7 +26,11 @@ use jf_plonk::{
     testing_apis::Verifier,
     transcript::{PlonkTranscript, SolidityTranscript},
 };
+use jf_primitives::constants::CS_ID_BLS_BN254;
+
 use jf_primitives::pcs::prelude::{Commitment, UnivariateUniversalParams};
+use jf_primitives::signatures::bls_over_bn254::KeyPair;
+use jf_primitives::signatures::bls_over_bn254::Signature;
 use jf_relation::{Arithmetization, Circuit, PlonkCircuit};
 use num_bigint::BigUint;
 use num_traits::Num;
@@ -76,6 +80,8 @@ enum Action {
     DummyProof,
     /// Test only logic
     TestOnly,
+    /// Generate a BLS Signature
+    GenBLSSig,
 }
 
 #[allow(clippy::type_complexity)]
@@ -387,6 +393,26 @@ fn main() {
         Action::TestOnly => {
             println!("args: {:?}", cli.args);
         }
+        Action::GenBLSSig => {
+            let mut rng = jf_utils::test_rng();
+
+            let key_pair = KeyPair::generate(&mut rng);
+            let vk = key_pair.ver_key();
+            let vk_g2_proj: G2Affine = vk.to_affine();
+            let msg = Bytes::from(vec![54u8, 12u8]);
+            let pk_x_c0 = field_to_u256::<Fq>(vk_g2_proj.x.c0);
+            let pk_x_c1 = field_to_u256::<Fq>(vk_g2_proj.x.c1);
+            let pk_y_c0 = field_to_u256::<Fq>(vk_g2_proj.y.c0);
+            let pk_y_c1 = field_to_u256::<Fq>(vk_g2_proj.y.c1);
+
+            let sig: Signature = key_pair.sign(&msg, CS_ID_BLS_BN254);
+            let sig_affine_point = sig.sigma.into_affine();
+            let sig_x = field_to_u256::<Fq>(sig_affine_point.x);
+            let sig_y = field_to_u256::<Fq>(sig_affine_point.y);
+
+            let res = (sig_x, sig_y, pk_x_c0, pk_x_c1, pk_y_c0, pk_y_c1, msg);
+            println!("{}", res.encode_hex());
+        }
     };
 }
 

contracts/src/StakeTable.sol

@@ -0,0 +1,112 @@
+pragma solidity ^0.8.0;
+
+import { BN254 } from "bn254/BN254.sol";
+import { BLSSig } from "./libraries/BLSSig.sol";
+
+import "./interfaces/IStakeTable.sol";
+
+contract StakeTable is IStakeTable {
+    mapping(bytes32 keyHash => Node node) private nodesTable;
+    uint256[2] private totalStakeArr;
+    uint32 private totalNumberOfKeys;
+    uint256 private totalVotingStakeVal;
+    uint64 private numRegistrations;
+    uint64 private numPendingExits;
+
+    // TODO check
+    function hashBlsKey(BN254.G2Point calldata blsVK) private pure returns (bytes32) {
+        uint256 x0 = blsVK.x0;
+        uint256 x1 = blsVK.x1;
+        uint256 y0 = blsVK.y0;
+        uint256 y1 = blsVK.y1;
+        bytes32 hash = keccak256(abi.encode(x0, x1, y0, y1));
+        return hash;
+    }
+
+    function totalStake() external view returns (uint256, uint256) {
+        return (totalStakeArr[0], totalStakeArr[1]);
+    }
+
+    function totalKeys() external view returns (uint32) {
+        return totalNumberOfKeys;
+    }
+
+    function totalVotingStake() external view returns (uint256) {
+        return totalVotingStakeVal;
+    }
+
+    function lookupStake(BN254.G2Point calldata blsVK) external view returns (uint64) {
+        Node memory node = this.lookupNode(blsVK);
+        return node.balance;
+    }
+
+    function lookupNode(BN254.G2Point calldata blsVK) external view returns (Node memory) {
+        bytes32 hash = hashBlsKey(blsVK);
+        return nodesTable[hash];
+    }
+
+    function nextRegistrationEpoch() external view returns (uint64) {
+        if (numRegistrations == 0) {
+            return 0;
+        } else {
+            return 1;
+        }
+    }
+
+    function numPendingRegistrations() external view returns (uint64) {
+        return numRegistrations;
+    }
+
+    function nextExitEpoch() external view returns (uint64) {
+        if (numPendingExits == 0) {
+            return 0;
+        } else {
+            return 1;
+        }
+    }
+
+    function numPendingExit() external view returns (uint64) {
+        return numPendingExits;
+    }
+
+    function register(
+        BN254.G2Point calldata blsVK,
+        EdOnBN254.EdOnBN254Point calldata schnorrVK,
+        uint64 amount,
+        StakeType stakeType,
+        BN254.G1Point calldata blsSig,
+        uint64 validUntilEpoch
+    ) external returns (bool) {
+        uint64 thisEpoch = 0; // TODO
+        Node memory node =
+            Node(msg.sender, stakeType, amount, thisEpoch, validUntilEpoch, schnorrVK);
+        // TODO Check blsSig
+        if (blsSig.x == 0) {
+            return false;
+        }
+        bytes32 hash = hashBlsKey(blsVK);
+        nodesTable[hash] = node;
+        return true; // TODO
+    }
+
+    function deposit(BN254.G2Point calldata blsVK, uint64 amount)
+        external
+        returns (uint64, uint64)
+    {
+        bytes32 hash = hashBlsKey(blsVK);
+        nodesTable[hash].balance += amount;
+        return (0, 0);
+    }
+
+    function requestExit(BN254.G2Point calldata blsVK) external returns (bool) {
+        bytes32 hash = hashBlsKey(blsVK);
+        nodesTable[hash].exitEpoch = 0;
+        return true;
+    }
+
+    function withdrawFunds(BN254.G2Point calldata blsVK) external returns (uint64) {
+        bytes32 hash = hashBlsKey(blsVK);
+        nodesTable[hash].balance = 0;
+        return 0;
+    }
+}

contracts/src/interfaces/IStakeTable.sol

@@ -21,8 +21,8 @@ interface IStakeTable {
     // We avoid declaring another struct even if the type info helps with readability,
     // extra layer of struct introduces overhead and more gas cost.
     //
-    /// @dev BLS verification key (used for Consensus voting in HotShot) is `BN254.G1Point`.
-    // `type BlsVerKey is BN254.G1Point;`
+    /// @dev BLS verification key (used for Consensus voting in HotShot) is `BN254.G2Point`.
+    // `type BlsVerKey is BN254.G2Point;`
     //
     /// @dev Verification key of a Schnorr signature is just a `EdOnBN254.EdOnBN254Point`
     // `type SchnorrVerKey is EdOnBN254.EdOnBN254Point;`
@@ -56,9 +56,9 @@ interface IStakeTable {
     function totalVotingStake() external view returns (uint256);
 
     /// @notice Look up the balance of `blsVK`
-    function lookupStake(BN254.G1Point calldata blsVK) external view returns (uint64);
+    function lookupStake(BN254.G2Point calldata blsVK) external view returns (uint64);
     /// @notice Look up the full `Node` state associated with `blsVK`
-    function lookupNode(BN254.G1Point calldata blsVK) external view returns (Node memory);
+    function lookupNode(BN254.G2Point calldata blsVK) external view returns (Node memory);
 
     // === Queuing Stats ===
 
@@ -89,11 +89,11 @@ interface IStakeTable {
     /// the contract only treat it as auxiliary info submitted by `blsVK`.
     /// @dev `blsSig` field is necessary to prevent "rogue public-key attack".
     function register(
-        BN254.G1Point calldata blsVK,
+        BN254.G2Point calldata blsVK,
         EdOnBN254.EdOnBN254Point calldata schnorrVK,
         uint64 amount,
         StakeType stakeType,
-        bytes calldata blsSig,
+        BN254.G1Point calldata blsSig,
         uint64 validUntilEpoch
     ) external returns (bool);
 
@@ -102,20 +102,20 @@ interface IStakeTable {
     /// @param blsVK The BLS verification key
     /// @param amount The amount to deposit
     /// @return (newBalance, effectiveEpoch) the new balance effective at a future epoch
-    function deposit(BN254.G1Point calldata blsVK, uint64 amount)
+    function deposit(BN254.G2Point calldata blsVK, uint64 amount)
         external
         returns (uint64, uint64);
 
     /// @notice Request to exit from the stake table, not immediately withdrawable!
     ///
     /// @param blsVK The BLS verification key to exit
     /// @return success status
-    function requestExit(BN254.G1Point calldata blsVK) external returns (bool);
+    function requestExit(BN254.G2Point calldata blsVK) external returns (bool);
 
     /// @notice Withdraw from the staking pool. Transfers occur! Only successfully exited keys can
     /// withdraw past their `exitEpoch`.
     ///
     /// @param blsVK The BLS verification key to withdraw
     /// @return The total amount withdrawn, equal to `Node.balance` associated with `blsVK`
-    function withdrawFunds(BN254.G1Point calldata blsVK) external returns (uint64);
+    function withdrawFunds(BN254.G2Point calldata blsVK) external returns (uint64);
 }

contracts/test/StakeTable.t.sol

@@ -0,0 +1,47 @@
+// SPDX-License-Identifier: Unlicensed
+
+/* solhint-disable contract-name-camelcase, func-name-mixedcase, one-contract-per-file */
+
+pragma solidity ^0.8.0;
+
+// Libraries
+import "forge-std/Test.sol";
+import { BN254 } from "bn254/BN254.sol";
+import { BLSSig } from "../src/libraries/BLSSig.sol";
+
+// Target contract
+import { StakeTable as S } from "../src/StakeTable.sol";
+
+contract StableTable_keyRegister_Test is Test {
+    S public stakeTable;
+
+    function setUp() public {
+        stakeTable = new S();
+    }
+
+    /// @dev Tests the key registering process
+    function testKeyRegister() external {
+        // Generate a BLS signature in rust and read it in solidity
+        string[] memory cmds = new string[](2);
+        cmds[0] = "diff-test";
+        cmds[1] = "gen-bls-sig";
+
+        bytes memory result = vm.ffi(cmds);
+        (
+            uint256 blsSigX,
+            uint256 blsSigY,
+            uint256 blsVkx0,
+            uint256 blsVkx1,
+            uint256 blsVky0,
+            uint256 blsVky1,
+            bytes memory message
+        ) = abi.decode(result, (uint256, uint256, uint256, uint256, uint256, uint256, bytes));
+
+        BN254.G1Point memory sig = BN254.G1Point(blsSigX, blsSigY);
+
+        // Note: (x,y) coordinates for each field component must be inverted.
+        BN254.G2Point memory vk = BN254.G2Point(blsVkx1, blsVkx0, blsVky1, blsVky0);
+
+        BLSSig.verifyBlsSig(message, sig, vk);
+    }
+}