Implementation for the Notary to Vote on Collations (#131)

sharding: notary functionality for voting on collations

Former-commit-id: 3ac6bb4a7a269966dd8d526c13cce4d0f1aea680 [formerly 2d527347f708a462bfbc206a581bd9b79963ad4c]
Former-commit-id: 68ac7bb8cffc3211a78e15a733101ef035ca50d4

sharding/contracts/sharding_manager.sol

@@ -145,7 +145,7 @@ contract SMC {
         uint deregisteredPeriod = block.number / PERIOD_LENGTH;
         notaryRegistry[notaryAddress].deregisteredPeriod = deregisteredPeriod;
 
-        stackPush(index);
+        stackPush(index); 
         delete notaryPool[index];
         --notaryPoolLength;
         emit NotaryDeregistered(notaryAddress, index, deregisteredPeriod);

sharding/contracts/types.go

@@ -0,0 +1,13 @@
+package contracts
+
+import (
+	"math/big"
+)
+
+// Registry describes the Notary Registry in the SMC.
+type Registry struct {
+	DeregisteredPeriod *big.Int
+	PoolIndex          *big.Int
+	Balance            *big.Int
+	Deposited          bool
+}

sharding/mainchain/interfaces.go

@@ -3,8 +3,10 @@ package mainchain
 import (
 	"context"
 	"math/big"
+	"time"
 
 	ethereum "github.com/ethereum/go-ethereum"
+	"github.com/ethereum/go-ethereum/accounts"
 	"github.com/ethereum/go-ethereum/accounts/abi/bind"
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/types"
@@ -38,6 +40,15 @@ type ContractTransactor interface {
 	CreateTXOpts(value *big.Int) (*bind.TransactOpts, error)
 }
 
+// EthClient defines the methods that will be used to perform rpc calls
+// to the main geth node, and be responsible for other user-specific data
+type EthClient interface {
+	Account() *accounts.Account
+	WaitForTransaction(ctx context.Context, hash common.Hash, durationInSeconds time.Duration) error
+	TransactionReceipt(hash common.Hash) (*types.Receipt, error)
+	DepositFlag() bool
+}
+
 // Reader defines an interface for a struct that can read mainchain information
 // such as blocks, transactions, receipts, and more. Useful for testing.
 type Reader interface {

sharding/notary/notary.go

@@ -1,15 +1,21 @@
 package notary
 
 import (
+	"bytes"
 	"context"
+	"errors"
 	"fmt"
 	"math/big"
+	"time"
 
 	"github.com/ethereum/go-ethereum/accounts"
 	"github.com/ethereum/go-ethereum/accounts/abi/bind"
+	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/params"
+	"github.com/ethereum/go-ethereum/sharding"
+	"github.com/ethereum/go-ethereum/sharding/contracts"
 	"github.com/ethereum/go-ethereum/sharding/mainchain"
 	shardparams "github.com/ethereum/go-ethereum/sharding/params"
 )
@@ -67,76 +73,207 @@ func checkSMCForNotary(caller mainchain.ContractCaller, account *accounts.Accoun
 			return err
 		}
 
-		// If output is non-empty and the addr == coinbase.
 		if addr == account.Address {
 			log.Info(fmt.Sprintf("Selected as notary on shard: %d", s))
-			err := submitCollation(s)
-			if err != nil {
-				return err
-			}
-
-			// If the account is selected as notary, submit collation.
-			if addr == account.Address {
-				log.Info(fmt.Sprintf("Selected as notary on shard: %d", s))
-				err := submitCollation(s)
-				if err != nil {
-					return fmt.Errorf("could not add collation. %v", err)
-				}
-			}
 		}
+
 	}
 
 	return nil
 }
 
+// getNotaryRegistry retrieves the registry of the registered account.
+func getNotaryRegistry(caller mainchain.ContractCaller, account *accounts.Account) (*contracts.Registry, error) {
+
+	var nreg contracts.Registry
+	nreg, err := caller.SMCCaller().NotaryRegistry(&bind.CallOpts{}, account.Address)
+	if err != nil {
+		return nil, fmt.Errorf("unable to retrieve notary registry: %v", err)
+	}
+
+	return &nreg, nil
+}
+
 // isAccountInNotaryPool checks if the user is in the notary pool because
 // we can't guarantee our tx for deposit will be in the next block header we receive.
 // The function calls IsNotaryDeposited from the SMC and returns true if
 // the user is in the notary pool.
 func isAccountInNotaryPool(caller mainchain.ContractCaller, account *accounts.Account) (bool, error) {
-	// Checks if our deposit has gone through according to the SMC.
-	nreg, err := caller.SMCCaller().NotaryRegistry(&bind.CallOpts{}, account.Address)
-	if !nreg.Deposited && err != nil {
-		log.Warn(fmt.Sprintf("Account %s not in notary pool.", account.Address.String()))
+
+	nreg, err := getNotaryRegistry(caller, account)
+	if err != nil {
+		return false, err
 	}
-	return nreg.Deposited, err
+
+	if !nreg.Deposited {
+		log.Warn(fmt.Sprintf("Account %s not in notary pool.", account.Address.Hex()))
+	}
+
+	return nreg.Deposited, nil
 }
 
-// submitCollation interacts with the SMC directly to add a collation header.
-func submitCollation(shardID int64) error {
-	// TODO: Adds a collation header to the SMC with the following fields:
-	// [
-	//  shard_id: uint256,
-	//  expected_period_number: uint256,
-	//  period_start_prevhash: bytes32,
-	//  parent_hash: bytes32,
-	//  transactions_root: bytes32,
-	//  coinbase: address,
-	//  state_root: bytes32,
-	//  receipts_root: bytes32,
-	//  number: uint256,
-	//  sig: bytes
-	// ]
-	//
-	// Before calling this, we would need to have access to the state of
-	// the period_start_prevhash. Refer to the comments in:
-	// https://github.com/ethereum/py-evm/issues/258#issuecomment-359879350
-	//
-	// This function will call FetchCandidateHead() of the SMC to obtain
-	// more necessary information.
-	//
-	// This functions will fetch the transactions in the proposer tx pool and and apply
-	// them to finish up the collation. It will then need to broadcast the
-	// collation to the main chain using JSON-RPC.
-	log.Info("Submit collation function called")
+// hasAccountBeenDeregistered checks if the account has been deregistered from the notary pool.
+func hasAccountBeenDeregistered(caller mainchain.ContractCaller, account *accounts.Account) (bool, error) {
+	nreg, err := getNotaryRegistry(caller, account)
+
+	if err != nil {
+		return false, err
+	}
+
+	return nreg.DeregisteredPeriod.Cmp(big.NewInt(0)) == 1, err
+}
+
+// isLockUpOver checks if the lock up period is over
+// which will allow the notary to call the releaseNotary function
+// in the SMC and get their deposit back.
+func isLockUpOver(caller mainchain.ContractCaller, reader mainchain.Reader, account *accounts.Account) (bool, error) {
+
+	//TODO: When chainreader for tests is implemented, get block using the method
+	//get BlockByNumber instead of passing as an argument to this function.
+	nreg, err := getNotaryRegistry(caller, account)
+	if err != nil {
+		return false, err
+	}
+	block, err := reader.BlockByNumber(context.Background(), nil)
+	if err != nil {
+		return false, err
+	}
+
+	return (block.Number().Int64() / shardparams.DefaultConfig.PeriodLength) > nreg.DeregisteredPeriod.Int64()+shardparams.DefaultConfig.NotaryLockupLength, nil
+
+}
+
+func transactionWaiting(client mainchain.EthClient, tx *types.Transaction, duration time.Duration) error {
+
+	err := client.WaitForTransaction(context.Background(), tx.Hash(), duration)
+	if err != nil {
+		return err
+	}
+
+	receipt, err := client.TransactionReceipt(tx.Hash())
+	if err != nil {
+		return err
+	}
+
+	if receipt.Status == types.ReceiptStatusFailed {
+		return errors.New("transaction was not successful, unable to release Notary")
+	}
 	return nil
+
+}
+
+func settingCanonicalShardChain(shard sharding.Shard, manager mainchain.ContractManager, period *big.Int, headerHash *common.Hash) error {
+
+	shardID := shard.ShardID()
+	collationRecords, err := manager.SMCCaller().CollationRecords(&bind.CallOpts{}, shardID, period)
+
+	if err != nil {
+		return fmt.Errorf("unable to get collation record: %v", err)
+	}
+
+	// Logs if quorum has been reached and collation is added to the canonical shard chain
+	if collationRecords.IsElected {
+		log.Info(fmt.Sprintf(
+			"Shard %v in period %v has chosen the collation with its header hash %v to be added to the canonical shard chain",
+			shardID, period, headerHash))
+
+		// Setting collation header as canonical in the shard chain
+		header, err := shard.HeaderByHash(headerHash)
+		if err != nil {
+			return fmt.Errorf("unable to set Header from hash: %v", err)
+		}
+
+		err = shard.SetCanonical(header)
+		if err != nil {
+			return fmt.Errorf("unable to add collation to canonical shard chain: %v", err)
+		}
+	}
+
+	return nil
+
+}
+
+func getCurrentNetworkState(manager mainchain.ContractManager, shard sharding.Shard, reader mainchain.Reader) (int64, *big.Int, *types.Block, error) {
+
+	shardcount, err := manager.GetShardCount()
+	if err != nil {
+		return 0, nil, nil, fmt.Errorf("could not get shard count: %v", err)
+	}
+
+	shardID := shard.ShardID()
+	// checks if the shardID is valid
+	if shardID.Int64() <= int64(0) || shardID.Int64() > shardcount {
+		return 0, nil, nil, fmt.Errorf("shardId is invalid, it has to be between %d and %d, instead it is %v", 0, shardcount, shardID)
+	}
+
+	block, err := reader.BlockByNumber(context.Background(), nil)
+	if err != nil {
+		return 0, nil, nil, fmt.Errorf("unable to retrieve block: %v", err)
+	}
+
+	return shardcount, shardID, block, nil
+
+}
+
+func checkCollationPeriod(manager mainchain.ContractManager, block *types.Block, shardID *big.Int) (*big.Int, *big.Int, error) {
+
+	period := big.NewInt(0).Div(block.Number(), big.NewInt(shardparams.DefaultConfig.PeriodLength))
+	collPeriod, err := manager.SMCCaller().LastSubmittedCollation(&bind.CallOpts{}, shardID)
+
+	if err != nil {
+		return nil, nil, fmt.Errorf("unable to get period from last submitted collation: %v", err)
+	}
+
+	// Checks if the current period is valid in order to vote for the collation on the shard
+	if period.Int64() != collPeriod.Int64() {
+		return nil, nil, fmt.Errorf("period in collation is not equal to current period : %d , %d", collPeriod, period)
+	}
+	return period, collPeriod, nil
+
+}
+
+func hasNotaryVoted(manager mainchain.ContractManager, shardID *big.Int, poolIndex *big.Int) (bool, error) {
+	hasVoted, err := manager.SMCCaller().HasVoted(&bind.CallOpts{}, shardID, poolIndex)
+
+	if err != nil {
+		return false, fmt.Errorf("unable to know if notary voted: %v", err)
+	}
+
+	return hasVoted, nil
+}
+
+func verifyNotary(manager mainchain.ContractManager, client mainchain.EthClient) (*contracts.Registry, error) {
+	nreg, err := getNotaryRegistry(manager, client.Account())
+
+	if err != nil {
+		return nil, err
+	}
+
+	if !nreg.Deposited {
+		return nil, fmt.Errorf("notary has not deposited to the SMC")
+	}
+
+	// Checking if the pool index is valid
+	if nreg.PoolIndex.Int64() >= shardparams.DefaultConfig.NotaryCommitteeSize {
+		return nil, fmt.Errorf("invalid pool index %d as it is more than the committee size of %d", nreg.PoolIndex, shardparams.DefaultConfig.NotaryCommitteeSize)
+	}
+
+	return nreg, nil
 }
 
 // joinNotaryPool checks if the deposit flag is true and the account is a
 // notary in the SMC. If the account is not in the set, it will deposit ETH
 // into contract.
-func joinNotaryPool(manager mainchain.ContractManager, account *accounts.Account, config *shardparams.Config) error {
-	if b, err := isAccountInNotaryPool(manager, account); b || err != nil {
+func joinNotaryPool(manager mainchain.ContractManager, client mainchain.EthClient, config *shardparams.Config) error {
+	if !client.DepositFlag() {
+		return errors.New("joinNotaryPool called when deposit flag was not set")
+	}
+
+	if b, err := isAccountInNotaryPool(manager, client.Account()); b || err != nil {
+		if b {
+			log.Info(fmt.Sprint("Already joined notary pool"))
+			return nil
+		}
 		return err
 	}
 
@@ -150,7 +287,210 @@ func joinNotaryPool(manager mainchain.ContractManager, account *accounts.Account
 	if err != nil {
 		return fmt.Errorf("unable to deposit eth and become a notary: %v", err)
 	}
-	log.Info(fmt.Sprintf("Deposited %dETH into contract with transaction hash: %s", new(big.Int).Div(config.NotaryDeposit, big.NewInt(params.Ether)), tx.Hash().String()))
+
+	err = client.WaitForTransaction(context.Background(), tx.Hash(), 400)
+	if err != nil {
+		return err
+	}
+
+	receipt, err := client.TransactionReceipt(tx.Hash())
+	if err != nil {
+		return err
+	}
+	if receipt.Status == types.ReceiptStatusFailed {
+		return errors.New("transaction was not successful, unable to deposit ETH and become a notary")
+	}
+
+	if inPool, err := isAccountInNotaryPool(manager, client.Account()); !inPool || err != nil {
+		if err != nil {
+			return err
+		}
+		return errors.New("account has not been able to be deposited in notary pool")
+	}
+
+	log.Info(fmt.Sprintf("Deposited %dETH into contract with transaction hash: %s", new(big.Int).Div(shardparams.DefaultConfig.NotaryDeposit, big.NewInt(params.Ether)), tx.Hash().Hex()))
 
 	return nil
 }
+
+// leaveNotaryPool causes the notary to deregister and leave the notary pool
+// by calling the DeregisterNotary function in the SMC.
+func leaveNotaryPool(manager mainchain.ContractManager, client mainchain.EthClient) error {
+
+	if inPool, err := isAccountInNotaryPool(manager, client.Account()); !inPool || err != nil {
+		if err != nil {
+			return err
+		}
+		return errors.New("account has not been able to be deposited in notary pool")
+	}
+
+	txOps, err := manager.CreateTXOpts(nil)
+	if err != nil {
+		return fmt.Errorf("unable to create txOpts: %v", err)
+	}
+
+	tx, err := manager.SMCTransactor().DeregisterNotary(txOps)
+	if err != nil {
+		return fmt.Errorf("unable to deregister notary: %v", err)
+	}
+
+	err = client.WaitForTransaction(context.Background(), tx.Hash(), 400)
+	if err != nil {
+		return err
+	}
+
+	receipt, err := client.TransactionReceipt(tx.Hash())
+	if err != nil {
+		return err
+	}
+
+	if receipt.Status == types.ReceiptStatusFailed {
+		return errors.New("transaction was not successful, unable to deregister notary")
+	}
+
+	if dreg, err := hasAccountBeenDeregistered(manager, client.Account()); !dreg || err != nil {
+		if err != nil {
+			return err
+		}
+		return errors.New("notary unable to be deregistered successfully from pool")
+	}
+
+	log.Info(fmt.Sprintf("Notary deregistered from the pool with hash: %s", tx.Hash().Hex()))
+	return nil
+
+}
+
+// releaseNotary releases the Notary from the pool by deleting the notary from
+// the registry and transferring back the deposit
+func releaseNotary(manager mainchain.ContractManager, client mainchain.EthClient, reader mainchain.Reader) error {
+
+	if dreg, err := hasAccountBeenDeregistered(manager, client.Account()); !dreg || err != nil {
+		if err != nil {
+			return err
+		}
+		return errors.New("notary has not been deregistered from the pool")
+	}
+
+	if lockup, err := isLockUpOver(manager, reader, client.Account()); !lockup || err != nil {
+		if err != nil {
+			return err
+		}
+		return errors.New("lockup period is not over")
+	}
+
+	txOps, err := manager.CreateTXOpts(nil)
+	if err != nil {
+		return fmt.Errorf("unable to create txOpts: %v", err)
+	}
+
+	tx, err := manager.SMCTransactor().ReleaseNotary(txOps)
+	if err != nil {
+		return fmt.Errorf("unable to Release Notary: %v", err)
+	}
+
+	err = transactionWaiting(client, tx, 400)
+	if err != nil {
+		return err
+	}
+
+	nreg, err := getNotaryRegistry(manager, client.Account())
+	if err != nil {
+		return err
+	}
+
+	if nreg.Deposited || nreg.Balance.Cmp(big.NewInt(0)) != 0 || nreg.DeregisteredPeriod.Cmp(big.NewInt(0)) != 0 || nreg.PoolIndex.Cmp(big.NewInt(0)) != 0 {
+		return errors.New("notary unable to be released from the pool")
+	}
+
+	log.Info(fmt.Sprintf("notary with address: %s released from pool", client.Account().Address.Hex()))
+
+	return nil
+
+}
+
+// submitVote votes for a collation on the shard
+// by taking in the shard and the hash of the collation header
+func submitVote(shard sharding.Shard, manager mainchain.ContractManager, client mainchain.EthClient, reader mainchain.Reader, headerHash *common.Hash) error {
+
+	_, shardID, block, err := getCurrentNetworkState(manager, shard, reader)
+	if err != nil {
+		return err
+	}
+
+	period, _, err := checkCollationPeriod(manager, block, shardID)
+	if err != nil {
+		return err
+	}
+
+	nreg, err := verifyNotary(manager, client)
+
+	if err != nil {
+		return err
+	}
+
+	collationRecords, err := manager.SMCCaller().CollationRecords(&bind.CallOpts{}, shardID, period)
+	if err != nil {
+		return fmt.Errorf("unable to get collation record: %v", err)
+	}
+
+	chunkroot, err := shard.ChunkRootfromHeaderHash(headerHash)
+	if err != nil {
+		return fmt.Errorf("unable to get chunk root: %v", err)
+	}
+
+	// Checking if the chunkroot is valid
+	if !bytes.Equal(collationRecords.ChunkRoot[:], chunkroot.Bytes()) {
+		return fmt.Errorf("submmitted collation header has a different chunkroot to the one saved in the SMC")
+
+	}
+
+	hasVoted, err := hasNotaryVoted(manager, shardID, nreg.PoolIndex)
+	if err != nil {
+		return err
+	}
+	if hasVoted {
+		return errors.New("notary has already voted")
+	}
+
+	inCommitee, err := manager.SMCCaller().GetNotaryInCommittee(&bind.CallOpts{}, shardID)
+	if err != nil {
+		return fmt.Errorf("unable to know if notary is in committee: %v", err)
+	}
+
+	if inCommitee != client.Account().Address {
+		return errors.New("notary is not eligible to vote in this shard at the current period")
+	}
+
+	txOps, err := manager.CreateTXOpts(nil)
+	if err != nil {
+		return fmt.Errorf("unable to create txOpts: %v", err)
+	}
+
+	tx, err := manager.SMCTransactor().SubmitVote(txOps, shardID, period, nreg.PoolIndex, collationRecords.ChunkRoot)
+	if err != nil {
+
+		return fmt.Errorf("unable to submit Vote: %v", err)
+	}
+
+	err = transactionWaiting(client, tx, 400)
+
+	if err != nil {
+		return err
+	}
+
+	hasVoted, err = hasNotaryVoted(manager, shardID, nreg.PoolIndex)
+
+	if err != nil {
+		return fmt.Errorf("unable to know if notary voted: %v", err)
+	}
+
+	if !hasVoted {
+		return errors.New("notary has not voted")
+	}
+
+	log.Info(fmt.Sprintf("Notary has voted for shard: %v in the %v period", shardID, period))
+
+	err = settingCanonicalShardChain(shard, manager, period, headerHash)
+
+	return err
+}

sharding/notary/service.go

@@ -46,7 +46,7 @@ func (n *Notary) notarizeCollations() {
 	// TODO: handle this better through goroutines. Right now, these methods
 	// are blocking.
 	if n.smcClient.DepositFlag() {
-		if err := joinNotaryPool(n.smcClient, n.smcClient.Account(), n.config); err != nil {
+		if err := joinNotaryPool(n.smcClient, n.smcClient, n.config); err != nil {
 			log.Error(fmt.Sprintf("Could not fetch current block number: %v", err))
 			return
 		}

sharding/notary/service_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"math/big"
 	"testing"
+	"time"
 
 	ethereum "github.com/ethereum/go-ethereum"
 	"github.com/ethereum/go-ethereum/accounts"
@@ -15,7 +16,7 @@ import (
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/sharding"
 	"github.com/ethereum/go-ethereum/sharding/contracts"
-	"github.com/ethereum/go-ethereum/sharding/params"
+	shardparams "github.com/ethereum/go-ethereum/sharding/params"
 )
 
 var (
@@ -32,6 +33,12 @@ type smcClient struct {
 	smc         *contracts.SMC
 	depositFlag bool
 	t           *testing.T
+	backend     *backends.SimulatedBackend
+	blocknumber int64
+}
+
+func (s *smcClient) SubscribeNewHead(ctx context.Context, ch chan<- *types.Header) (ethereum.Subscription, error) {
+	return nil, nil
 }
 
 func (s *smcClient) Account() *accounts.Account {
@@ -54,12 +61,14 @@ func (s *smcClient) SMCFilterer() *contracts.SMCFilterer {
 	return &s.smc.SMCFilterer
 }
 
-func (s *smcClient) WaitForTransaction(ctx context.Context, hash common.Hash, durationInSeconds int64) error {
+func (s *smcClient) WaitForTransaction(ctx context.Context, hash common.Hash, durationInSeconds time.Duration) error {
+	s.CommitWithBlock()
+	s.fastForward(1)
 	return nil
 }
 
 func (s *smcClient) TransactionReceipt(hash common.Hash) (*types.Receipt, error) {
-	return nil, nil
+	return s.backend.TransactionReceipt(context.Background(), hash)
 }
 
 func (s *smcClient) CreateTXOpts(value *big.Int) (*bind.TransactOpts, error) {
@@ -68,10 +77,46 @@ func (s *smcClient) CreateTXOpts(value *big.Int) (*bind.TransactOpts, error) {
 	return txOpts, nil
 }
 
+func (s *smcClient) DepositFlag() bool {
+	return s.depositFlag
+}
+
+func (s *smcClient) SetDepositFlag(deposit bool) {
+	s.depositFlag = deposit
+}
+
+func (s *smcClient) Sign(hash common.Hash) ([]byte, error) {
+	return nil, nil
+}
+
+// Unused mockClient methods.
+func (s *smcClient) Start() error {
+	s.t.Fatal("Start called")
+	return nil
+}
+
+func (s *smcClient) Close() {
+	s.t.Fatal("Close called")
+}
+
+func (s *smcClient) DataDirPath() string {
+	return "/tmp/datadir"
+}
+
+func (s *smcClient) CommitWithBlock() {
+	s.backend.Commit()
+	s.blocknumber = s.blocknumber + 1
+
+}
+
 func (s *smcClient) GetShardCount() (int64, error) {
 	return 100, nil
 }
 
+func (s *smcClient) BlockByNumber(ctx context.Context, number *big.Int) (*types.Block, error) {
+	return types.NewBlockWithHeader(&types.Header{Number: big.NewInt(s.blocknumber)}), nil
+}
+
 // Helper/setup methods.
 // TODO: consider moving these to common sharding testing package as the notary and smc tests
 // use them.
@@ -79,6 +124,14 @@ func transactOpts() *bind.TransactOpts {
 	return bind.NewKeyedTransactor(key)
 }
 
+// fastForward is a helper function to skip through n period.
+func (s *smcClient) fastForward(p int) {
+
+	for i := 0; i < p*int(shardparams.DefaultConfig.PeriodLength); i++ {
+		s.CommitWithBlock()
+	}
+}
+
 func setup() (*backends.SimulatedBackend, *contracts.SMC) {
 	backend := backends.NewSimulatedBackend(core.GenesisAlloc{addr: {Balance: accountBalance1001Eth}})
 	_, _, smc, _ := contracts.DeploySMC(transactOpts(), backend)
@@ -86,9 +139,72 @@ func setup() (*backends.SimulatedBackend, *contracts.SMC) {
 	return backend, smc
 }
 
+func TestHasAccountBeenDeregistered(t *testing.T) {
+	backend, smc := setup()
+	client := &smcClient{smc: smc, t: t, backend: backend, blocknumber: 1}
+
+	client.SetDepositFlag(true)
+	err := joinNotaryPool(client, client, nil)
+	if err != nil {
+		t.Error(err)
+	}
+
+	err = leaveNotaryPool(client, client)
+
+	if err != nil {
+		t.Error(err)
+	}
+
+	dreg, err := hasAccountBeenDeregistered(client, client.Account())
+
+	if err != nil {
+		t.Error(err)
+	}
+
+	if !dreg {
+		t.Error("account unable to be deregistered from notary pool")
+	}
+
+}
+
+func TestIsLockupOver(t *testing.T) {
+	backend, smc := setup()
+	client := &smcClient{smc: smc, t: t, backend: backend}
+
+	client.SetDepositFlag(true)
+	err := joinNotaryPool(client, client, shardparams.DefaultConfig)
+	if err != nil {
+		t.Error(err)
+	}
+
+	err = leaveNotaryPool(client, client)
+
+	if err != nil {
+		t.Error(err)
+	}
+
+	client.fastForward(int(shardparams.DefaultConfig.NotaryLockupLength + 100))
+
+	err = releaseNotary(client, client, client)
+	if err != nil {
+		t.Error(err)
+	}
+
+	lockup, err := isLockUpOver(client, client, client.Account())
+
+	if err != nil {
+		t.Error(err)
+	}
+
+	if !lockup {
+		t.Error("lockup period is not over despite account being relased from registry")
+	}
+
+}
+
 func TestIsAccountInNotaryPool(t *testing.T) {
 	backend, smc := setup()
-	client := &smcClient{smc: smc, t: t}
+	client := &smcClient{smc: smc, t: t, backend: backend}
 
 	// address should not be in pool initially.
 	b, err := isAccountInNotaryPool(client, client.Account())
@@ -96,65 +212,178 @@ func TestIsAccountInNotaryPool(t *testing.T) {
 		t.Fatal(err)
 	}
 	if b {
-		t.Fatal("Account unexpectedly in notary pool")
+		t.Fatal("account unexpectedly in notary pool")
 	}
 
 	txOpts := transactOpts()
 	// deposit in notary pool, then it should return true.
-	txOpts.Value = params.DefaultConfig.NotaryDeposit
+	txOpts.Value = shardparams.DefaultConfig.NotaryDeposit
 	if _, err := smc.RegisterNotary(txOpts); err != nil {
 		t.Fatalf("Failed to deposit: %v", err)
 	}
-	backend.Commit()
+	client.CommitWithBlock()
 	b, err = isAccountInNotaryPool(client, client.Account())
 	if err != nil {
-		t.Fatal(err)
+		t.Error(err)
 	}
 	if !b {
-		t.Fatal("Account not in notary pool when expected to be")
+		t.Error("account not in notary pool when expected to be")
 	}
 }
 
 func TestJoinNotaryPool(t *testing.T) {
 	backend, smc := setup()
-	client := &smcClient{smc: smc, t: t}
+	client := &smcClient{smc: smc, depositFlag: false, t: t, backend: backend}
 	// There should be no notary initially.
 	numNotaries, err := smc.NotaryPoolLength(&bind.CallOpts{})
 	if err != nil {
-		t.Fatal(err)
+		t.Error(err)
 	}
 	if big.NewInt(0).Cmp(numNotaries) != 0 {
-		t.Fatalf("Unexpected number of notaries. Got %d, wanted 0.", numNotaries)
+		t.Errorf("unexpected number of notaries. Got %d, wanted 0.", numNotaries)
 	}
 
-	err = joinNotaryPool(client, client.Account(), params.DefaultConfig)
-	if err != nil {
-		t.Fatal(err)
+	client.SetDepositFlag(false)
+	err = joinNotaryPool(client, client, shardparams.DefaultConfig)
+	if err == nil {
+		t.Error("joined notary pool while --deposit was not present")
+	}
+
+	client.SetDepositFlag(true)
+	err = joinNotaryPool(client, client, shardparams.DefaultConfig)
+	if err != nil {
+		t.Error(err)
 	}
-	backend.Commit()
 
 	// Now there should be one notary.
 	numNotaries, err = smc.NotaryPoolLength(&bind.CallOpts{})
 	if err != nil {
-		t.Fatal(err)
+		t.Error(err)
 	}
 	if big.NewInt(1).Cmp(numNotaries) != 0 {
-		t.Fatalf("Unexpected number of notaries. Got %d, wanted 1.", numNotaries)
+		t.Errorf("unexpected number of notaries. Got %d, wanted 1", numNotaries)
 	}
 
 	// Trying to join while deposited should do nothing
-	err = joinNotaryPool(client, client.Account(), params.DefaultConfig)
+	err = joinNotaryPool(client, client, shardparams.DefaultConfig)
 	if err != nil {
 		t.Error(err)
 	}
-	backend.Commit()
 
 	numNotaries, err = smc.NotaryPoolLength(&bind.CallOpts{})
 	if err != nil {
-		t.Fatal(err)
+		t.Error(err)
 	}
 	if big.NewInt(1).Cmp(numNotaries) != 0 {
-		t.Fatalf("Unexpected number of notaries. Got %d, wanted 1.", numNotaries)
+		t.Errorf("unexpected number of notaries. Got %d, wanted 1", numNotaries)
+	}
+
+}
+
+func TestLeaveNotaryPool(t *testing.T) {
+	backend, smc := setup()
+	client := &smcClient{smc: smc, t: t, depositFlag: true, backend: backend}
+
+	// Test Leaving Notary Pool Before Joining it
+
+	err := leaveNotaryPool(client, client)
+	if err == nil {
+		t.Error("able to leave notary pool despite having not joined it")
+	}
+
+	// Roundtrip Test , Join and leave pool
+
+	err = joinNotaryPool(client, client, shardparams.DefaultConfig)
+	if err != nil {
+		t.Error(err)
+	}
+	client.CommitWithBlock()
+
+	// Now there should be one notary.
+	numNotaries, err := smc.NotaryPoolLength(&bind.CallOpts{})
+	if err != nil {
+		t.Error(err)
+	}
+	if big.NewInt(1).Cmp(numNotaries) != 0 {
+		t.Errorf("unexpected number of notaries. Got %d, wanted 1", numNotaries)
+	}
+
+	err = leaveNotaryPool(client, client)
+	if err != nil {
+		t.Error(err)
+	}
+	client.CommitWithBlock()
+
+	numNotaries, err = smc.NotaryPoolLength(&bind.CallOpts{})
+	if err != nil {
+		t.Error(err)
+	}
+	if big.NewInt(0).Cmp(numNotaries) != 0 {
+		t.Errorf("unexpected number of notaries. Got %d, wanted 0", numNotaries)
+	}
+
+}
+
+func TestReleaseNotary(t *testing.T) {
+	backend, smc := setup()
+	client := &smcClient{smc: smc, t: t, depositFlag: true, backend: backend}
+
+	// Test Release Notary Before Joining it
+
+	err := releaseNotary(client, client, client)
+	if err == nil {
+		t.Error("released From notary despite never joining pool")
+	}
+
+	// Roundtrip Test , Join and leave pool and release Notary
+
+	err = joinNotaryPool(client, client, shardparams.DefaultConfig)
+	if err != nil {
+		t.Error(err)
+	}
+
+	err = leaveNotaryPool(client, client)
+	if err != nil {
+		t.Error(err)
+	}
+
+	balance, err := backend.BalanceAt(context.Background(), addr, nil)
+	if err != nil {
+		t.Error("unable to retrieve balance")
+	}
+	client.fastForward(int(shardparams.DefaultConfig.NotaryLockupLength + 10))
+
+	err = releaseNotary(client, client, client)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	nreg, err := smc.NotaryRegistry(&bind.CallOpts{}, addr)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if nreg.Deposited {
+		t.Error("Unable to release Notary and deposit money back")
+	}
+
+	newbalance, err := client.backend.BalanceAt(context.Background(), addr, nil)
+	if err != nil {
+		t.Error("unable to retrieve balance")
+	}
+
+	if balance.Cmp(newbalance) != -1 {
+		t.Errorf("Deposit was not returned, balance is currently: %v", newbalance)
+	}
+
+}
+
+func TestSubmitVote(t *testing.T) {
+	backend, smc := setup()
+	client := &smcClient{smc: smc, t: t, depositFlag: true, backend: backend}
+
+	err := joinNotaryPool(client, client, shardparams.DefaultConfig)
+	if err != nil {
+		t.Error(err)
 	}
 
 }

sharding/proposer/service_test.go

@@ -5,6 +5,7 @@ import (
 	"crypto/rand"
 	"math/big"
 	"testing"
+	"time"
 
 	"github.com/ethereum/go-ethereum"
 	"github.com/ethereum/go-ethereum/accounts"
@@ -48,7 +49,11 @@ func (m *mockNode) SMCTransactor() *contracts.SMCTransactor {
 	return &m.smc.SMCTransactor
 }
 
-func (m *mockNode) WaitForTransaction(ctx context.Context, hash common.Hash, durationInSeconds int64) error {
+func (m *mockNode) SMCFilterer() *contracts.SMCFilterer {
+	return &m.smc.SMCFilterer
+}
+
+func (m *mockNode) WaitForTransaction(ctx context.Context, hash common.Hash, durationInSeconds time.Duration) error {
 	return nil
 }
 

sharding/shard.go

@@ -81,6 +81,15 @@ func (s *Shard) CollationByHeaderHash(headerHash *common.Hash) (*Collation, erro
 	return NewCollation(header, body, *txs), nil
 }
 
+// ChunkRootfromHeaderHash gets the chunk root of a collation body from the hash of its header.
+func (s *Shard) ChunkRootfromHeaderHash(headerHash *common.Hash) (*common.Hash, error) {
+	collation, err := s.CollationByHeaderHash(headerHash)
+	if err != nil {
+		return nil, err
+	}
+	return collation.Header().ChunkRoot(), nil
+}
+
 // CanonicalHeaderHash gets a collation header hash that has been set as
 // canonical for shardID/period pair.
 func (s *Shard) CanonicalHeaderHash(shardID *big.Int, period *big.Int) (*common.Hash, error) {

sharding/shard_test.go

@@ -162,6 +162,41 @@ func TestShard_CollationByHeaderHash(t *testing.T) {
 	}
 }
 
+func TestShard_ChunkRootfromHeaderHash(t *testing.T) {
+	shardID := big.NewInt(1)
+	period := big.NewInt(1)
+	proposerAddress := common.BytesToAddress([]byte{})
+	proposerSignature := []byte{}
+	header := NewCollationHeader(shardID, nil, period, &proposerAddress, proposerSignature)
+
+	collation := NewCollation(header, []byte{1, 2, 3}, nil)
+	collation.CalculateChunkRoot()
+	shardDB := database.NewShardKV()
+	shard := NewShard(shardID, shardDB)
+
+	if err := shard.SaveCollation(collation); err != nil {
+		t.Fatalf("failed to save collation to shardDB: %v", err)
+	}
+
+	if err := shard.SetCanonical(header); err != nil {
+		t.Fatalf("failed to set header as canonical: %v", err)
+	}
+
+	chunkroot := collation.header.ChunkRoot()
+	headerHash := collation.header.Hash()
+
+	newchunkroot, err := shard.ChunkRootfromHeaderHash(&headerHash)
+
+	if err != nil {
+		t.Errorf("Unable to retrieve chunk root from collation header hash: %v", err)
+	}
+
+	if !bytes.Equal(newchunkroot.Bytes(), chunkroot.Bytes()) {
+		t.Errorf("Calculated chunk root, %v, and chunk root, %v, retrieved from headerHash is different", chunkroot, newchunkroot)
+	}
+
+}
+
 func TestShard_CanonicalHeaderHash(t *testing.T) {
 	shardID := big.NewInt(1)
 	period := big.NewInt(1)