main.go

package main

import (
	"bsc-shamir/election"
	"bsc-shamir/network"
	"bufio"
	"fmt"
	"log"
	"math/big"
	"os"
	"strconv"
	"sync"
	"time"

	"github.com/jessevdk/go-flags"
)

// See https://godoc.org/github.com/jessevdk/go-flags for available tags
var opts struct {
	Config string  `long:"config" description:"This specifies the configuration file" value-name:"filename" default:"servers"`
	Vote   []int64 `short:"v" long:"vote" description:"This specifies the vote. Must be either 0 or 1" choice:"0" choice:"1"`
	Root   string  `short:"r" long:"root" description:"This specifies the root certificate to use" value-name:"filename" default:"config/cert/root.crt"`
	Cert   string  `short:"c" long:"cert" description:"This specifies the certificate to use" value-name:"filename" default:"config/cert/client/c5.crt"`
	Key    string  `short:"k" long:"key" description:"This specifies the private key to use" value-name:"filename" default:"config/cert/client/c5.key"`
}

var voter struct {
	sync.Mutex

	Config *network.Config

	Client *network.Client

	ConnectionCount int
	RecievedCount   int
	Acks            map[string]bool
	Done            chan bool
}

func init() {
	_, err := flags.Parse(&opts)
	if err != nil {
		panic(err)
	}

	// Load config
	config, err := network.LoadConfig()
	if err != nil {
		log.Fatal("fatal: " + err.Error())
	}
	config.CertPath = opts.Cert
	config.KeyPath = opts.Key

	// Init client
	voter.RecievedCount = 0
	voter.Client = network.NewClient(config)
	voter.Client.OnMessage(election.MsgBallotAck, func(conn *network.Conn, data network.MessageData) *network.Message {
		voter.Lock()
		defer voter.Unlock()

		voter.Acks[conn.RemoteAddr().String()] = true
		conn.Close()

		voter.RecievedCount++
		if voter.RecievedCount == voter.ConnectionCount {
			voter.Done <- true
		}
		return nil
	})

	voter.Mutex = sync.Mutex{}
	voter.Acks = make(map[string]bool)
	voter.Done = make(chan bool, 1)
}

func main() {
	// Parse commandline arguments and exit if help message was printed.

	vote := inputVote()
	if vote != 0 && vote != 1 {
		fmt.Println("error: vote must be 0 or 1")
		os.Exit(0)
	}

	config, err := network.LoadConfig()
	if err != nil {
		log.Fatal("fatal: " + err.Error())
	}
	voter.Config = config

	t := (len(config.Servers) / 2) + 1 // One more than half must be honest!

	xs := make([]*big.Int, 0, len(config.Servers))
	for _, x := range voter.Config.Servers {
		xs = append(xs, x)
	}

	ballots := election.CreateBallots(t, xs, big.NewInt(vote))
	sendBallots(ballots)

	// Wait for acknowledgements or timeout
	timeout := network.TimeoutAfter(10 * time.Second)
	select {
	case <-voter.Done:
		log.Println("Main: all servers acknowledged")
	case <-timeout:
		for address, status := range voter.Acks {
			if !status {
				log.Println("Main: no acknowledgment from ", address)
			}
		}
	}

	err = voter.Client.Close()
	if err != nil {
		log.Println("Main: error: ", err.Error())
	}
}

func inputVote() int64 {
	if len(opts.Vote) != 0 {
		return opts.Vote[0]
	}

	fmt.Print("Input vote: ")
	reader := bufio.NewReader(os.Stdin)
	text, err := reader.ReadString('\n')
	if err != nil {
		log.Fatal("fatal: could not read input")
	}

	vote, err := strconv.Atoi(text[0 : len(text)-1])
	if err != nil {
		fmt.Println("error: vote must be an integer")
		os.Exit(0)
	}

	return int64(vote)
}

func sendBallots(ballots map[election.UniqueID]*election.Ballot) {
	voter.Lock()
	defer voter.Unlock()

	i := 0
	for address := range voter.Config.Servers {
		conn, err := voter.Client.Connect(address)
		if err != nil {
			log.Println("Main: error:", err)
			continue
		}
		ballot := ballots[conn.SerialNumber()]
		message := network.NewMessage(election.MsgBallot, ballot)

		err = conn.WriteMessage(message)
		if err != nil {
			log.Println("Main: error:", err)
			continue
		}
		voter.Acks[conn.RemoteAddr().String()] = false
		voter.ConnectionCount++
		i++
	}

	if voter.ConnectionCount == 0 {
		log.Fatal("fatal: could not connect to servers")
	}
}