6 files changed, 107 insertions(+), 120 deletions(-)

M config.go
M election.go
M main.go
M merkle.go
M ui.go
M voter.go

@@ 40,7 40,7 @@ type RawElectionInfo struct{
 
 // Saves an election's info into a file that can later be easily retrieved, even
 // on separate invocations of tallyard.  Returns an error if there was one.
-func saveElectionInfo(electionKey string, privKey crypto.PrivKey, pubKey crypto.PubKey) error {
+func (e *Election) saveInfo() error {
 	err := os.MkdirAll(tallyardConfigPath, 0700)
 	if err != nil {
 		return err


@@ 55,9 55,10 @@ func saveElectionInfo(electionKey string, privKey crypto.PrivKey, pubKey crypto.
 		return fmt.Errorf("unable to open election info file: %s", err)
 	}
 	defer file.Close()
-	infos[electionKey] = &ElectionInfo{
-		PrivKey: privKey,
-		PubKey:  pubKey,
+	localVoter := e.localVoter
+	infos[e.electionKey] = &ElectionInfo{
+		PrivKey: localVoter.Peerstore().PrivKey(localVoter.ID()),
+		PubKey:  localVoter.Peerstore().PubKey(localVoter.ID()),
 	}
 	raw, err := marshalElectionInfos(infos)
 	if err != nil {

@@ 1,9 1,12 @@
 package main
 
 import (
+	"fmt"
 	"strings"
 	"sync"
 
+	"github.com/cbergoon/merkletree"
+	"github.com/libp2p/go-libp2p"
 	"github.com/libp2p/go-libp2p-core/peer"
 	"github.com/mr-tron/base58/base58"
 )


@@ 20,48 23,76 @@ type Election struct {
 	close           chan<- int
 	// used by handleCmd to prevent closing election more than once
 	closed          bool
-	ElectionKey     string
+	electionKey     string
+	localVoter      *LocalVoter
 	masterID        peer.ID
 	merkleRoot      []byte
 	remoteVoters    map[peer.ID]*Voter
 	rendezvousNonce Nonce
 }
 
-func NewElection() *Election {
-	return &Election{
+func NewElectionWithCandidates(candidates []Candidate) *Election {
+	localVoter := NewLocalVoter(libp2p.RandomIdentity)
+	e := &Election{
+		Candidates:   candidates,
+		localVoter:   localVoter,
+		masterID:     localVoter.ID(),
 		remoteVoters: make(map[peer.ID]*Voter),
 	}
-}
-
-func (e *Election) MasterID() peer.ID {
-	if e.masterID != "" {
-		return e.masterID
+	e.rendezvousNonce = NewNonce()
+	content := []merkletree.Content{e.rendezvousNonce}
+	for _, cand := range e.Candidates {
+		content = append(content, cand)
 	}
-	zeroi := strings.IndexByte(e.ElectionKey, '0')
-	if zeroi == -1 {
-		panic("invalid election key")
-	}
-	var err error
-	e.masterID, err = peer.Decode(e.ElectionKey[zeroi+1:])
+	optionsMerkle, err := merkletree.NewTree(content)
 	if err != nil {
 		panic(err)
 	}
-	logger.Info("masted ID:", e.masterID)
-	return e.masterID
+	e.merkleRoot = optionsMerkle.MerkleRoot()
+	e.electionKey = fmt.Sprintf("%s0%s",
+		base58.Encode(e.merkleRoot),
+		localVoter.ID())
+	return e
 }
 
-func (e *Election) MerkleRoot() []byte {
-	if len(e.merkleRoot) > 0 {
-		return e.merkleRoot
+func NewElectionWithElectionKey(electionKey string) *Election {
+	var localVoter *LocalVoter
+	// try to recover info from elections file to see if we've
+	// joined/created this election before
+	electionInfo, err := getElectionInfo(electionKey)
+	if err != nil {
+		panic(err)
+	}
+	if electionInfo != nil { // we've joined this election before
+		// ensure public key from config is same as one derived
+		// from private key
+		if !electionInfo.PrivKey.GetPublic().Equals(electionInfo.PubKey) {
+			// TODO: handle properly
+			panic("loaded public key does not match derived one from private key")
+		}
+		localVoter = NewLocalVoter(libp2p.Identity(electionInfo.PrivKey))
+	} else { // we've never joined this election before
+		localVoter = NewLocalVoter(libp2p.RandomIdentity)
+	}
+
+	e := &Election{
+		electionKey:  electionKey,
+		localVoter:   localVoter,
+		remoteVoters: make(map[peer.ID]*Voter),
 	}
-	zeroi := strings.IndexByte(e.ElectionKey, '0')
+	zeroi := strings.IndexByte(electionKey, '0')
 	if zeroi == -1 {
 		panic("invalid election key")
 	}
-	logger.Info("merkle root:", e.ElectionKey[:zeroi])
-	merkleRoot, err := base58.Decode(e.ElectionKey[:zeroi])
+	logger.Info("merkle root:", electionKey[:zeroi])
+	e.merkleRoot, err = base58.Decode(electionKey[:zeroi])
+	if err != nil {
+		panic(err)
+	}
+	e.masterID, err = peer.Decode(electionKey[zeroi+1:])
 	if err != nil {
 		panic(err)
 	}
-	return merkleRoot
+	logger.Info("master ID:", e.masterID)
+	return e
 }

@@ 17,7 17,6 @@ import (
 	"github.com/libp2p/go-libp2p-core/protocol"
 	"github.com/libp2p/go-libp2p-core/routing"
 	dht "github.com/libp2p/go-libp2p-kad-dht"
-	"github.com/mr-tron/base58/base58"
 )
 
 const protocolID = protocol.ID("/tallyard/0.0.0")


@@ 53,7 52,9 @@ func NewLocalVoter(hostOpts ...libp2p.Option) *LocalVoter {
 	return localVoter
 }
 
-func (localVoter *LocalVoter) Bootstrap(election *Election) {
+func bootstrap(election *Election) {
+	localVoter := election.localVoter
+
 	var wg sync.WaitGroup
 	for _, peerAddr := range dht.DefaultBootstrapPeers {
 		peerInfo, _ := peer.AddrInfoFromP2pAddr(peerAddr)


@@ 69,35 70,25 @@ func (localVoter *LocalVoter) Bootstrap(election *Election) {
 	}
 	wg.Wait()
 
-	if election.ElectionKey == "" { // we are the master
+	if election.masterID == localVoter.ID() { // we are the master
 		fmt.Println("share this with peers:")
-		election.ElectionKey = fmt.Sprintf("%s0%s", base58.Encode(election.MerkleRoot()), localVoter.ID())
-		fmt.Printf("%s\n", election.ElectionKey)
-
-		err := saveElectionInfo(
-			election.ElectionKey,
-			localVoter.Peerstore().PrivKey(localVoter.ID()),
-			localVoter.Peerstore().PubKey(localVoter.ID()),
-		)
-		if err != nil {
-			panic(err)
-		}
+		fmt.Printf("%s\n", election.electionKey)
 
 		logger.Info("waiting for incoming streams and finding voters...")
 
 		election.Lock()
 		ch := make(chan int, 1)
 		election.close = ch
-		go findPeers(ch, election, localVoter)
+		go findPeers(ch, election)
 		election.Unlock()
 
 		localVoter.SetStreamHandler(protocolID, func(stream network.Stream) {
-			streamHandler(stream, election, localVoter)
+			streamHandler(stream, election)
 		})
 
 		fmt.Println("press ENTER to solidify group of voters and start voting")
 		stdReader := bufio.NewReader(os.Stdin)
-		_, err = stdReader.ReadString('\n')
+		_, err := stdReader.ReadString('\n')
 		if err != nil {
 			panic(err)
 		}


@@ 121,17 112,8 @@ func (localVoter *LocalVoter) Bootstrap(election *Election) {
 		}
 		election.RUnlock()
 	} else { // we are a slave
-		err := saveElectionInfo(
-			election.ElectionKey,
-			localVoter.Peerstore().PrivKey(localVoter.ID()),
-			localVoter.Peerstore().PubKey(localVoter.ID()),
-		)
-		if err != nil {
-			panic(err)
-		}
-
 		logger.Info("attempting to open stream with master peer...")
-		stream, err := localVoter.NewStream(localVoter.ctx, election.MasterID(), protocolID)
+		stream, err := localVoter.NewStream(localVoter.ctx, election.masterID, protocolID)
 		rw := bufio.NewReadWriter(bufio.NewReader(stream), bufio.NewWriter(stream))
 		if err != nil {
 			panic(err)


@@ 170,16 152,16 @@ func (localVoter *LocalVoter) Bootstrap(election *Election) {
 		logger.Info("done fetching election info")
 
 		logger.Info("checking authenticity of election info...")
-		verifyElectionInfo(election, election.MerkleRoot())
+		verifyElectionInfo(election, election.merkleRoot)
 
 		// channel used to signify when election is closed
 		ch := make(chan int, 1)
 		election.close = ch
 		// now that we have election info, begin handling streams
 		localVoter.SetStreamHandler(protocolID, func(stream network.Stream) {
-			streamHandler(stream, election, localVoter)
+			streamHandler(stream, election)
 		})
-		findPeers(ch, election, localVoter)
+		findPeers(ch, election)
 	}
 }
 


@@ 198,28 180,23 @@ func main() {
 		log.SetLogLevel("tallyard", "critical")
 	}
 
-	var (
-		election    *Election
-		electionKey string
-		identityOpt libp2p.Option = libp2p.RandomIdentity
-	)
+	var election *Election
 
 	if *master {
 		// we are the master and the candidates are the positional
 		// arguments
-		election = NewElection()
+		candidates := []Candidate{}
 		for _, candidate := range flag.Args() {
-			election.Candidates = append(
-				election.Candidates,
-				Candidate(candidate),
-			)
+			candidates = append(
+				candidates,
+				Candidate(candidate))
 		}
-		election.GenMerkleFromCandidates()
-	} else if electionKey = flag.Arg(0); electionKey != "" {
+		// TODO: ensure at least 2 candidates
+		election = NewElectionWithCandidates(candidates)
+	} else if electionKey := flag.Arg(0); electionKey != "" {
 		// we are a slave slave the and election key was given via CLI
 		// args
-		election = NewElection()
-		election.ElectionKey = electionKey
+		election = NewElectionWithElectionKey(electionKey)
 	} else {
 		// create/join election via TUI
 		election = tui()


@@ 229,26 206,11 @@ func main() {
 		}
 	}
 
-	if electionKey != "" {
-		// try to recover election info if we've joined before
-		electionInfo, err := getElectionInfo(electionKey)
-		if err != nil {
-			panic(err)
-		}
-		if electionInfo != nil {
-			// ensure public key from config is same as one derived
-			// from private key
-			if !electionInfo.PrivKey.GetPublic().Equals(electionInfo.PubKey) {
-				// TODO: handle properly
-				panic("loaded public key does not match derived one from private key")
-			}
-			identityOpt = libp2p.Identity(electionInfo.PrivKey)
-		}
-	}
-
-	localVoter := NewLocalVoter(identityOpt)
+	bootstrap(election)
 
-	localVoter.Bootstrap(election)
+	if err := election.saveInfo(); err != nil {
+		panic(err)
+	}
 
-	startVoting(election, localVoter)
+	startVoting(election)
 }

@@ 69,16 69,3 @@ func verifyElectionInfo(election *Election, merkleRoot []byte) {
 		os.Exit(1)
 	}
 }
-
-func (e *Election) GenMerkleFromCandidates() {
-	e.rendezvousNonce = NewNonce()
-	content := []merkletree.Content{e.rendezvousNonce}
-	for _, cand := range e.Candidates {
-		content = append(content, cand)
-	}
-	optionsMerkle, err := merkletree.NewTree(content)
-	if err != nil {
-		panic(err)
-	}
-	e.merkleRoot = optionsMerkle.MerkleRoot()
-}

@@ 14,22 14,22 @@ func tui() (election *Election) {
 	app := tview.NewApplication()
 	done := func(buttonIndex int, buttonLabel string) {
 		app.Stop()
-		election = NewElection()
 		switch buttonLabel {
 		case "Create Election":
-			election.Candidates = getCandidatesTUI()
+			candidates := getCandidatesTUI()
 			// TODO: slaves should check that len candidates >= 2
-			if election.Candidates == nil || len(election.Candidates) == 0 {
+			if len(candidates) == 0 {
 				fmt.Printf("no candidates entered; exiting\n")
 				os.Exit(0)
 			}
-			election.GenMerkleFromCandidates()
+			election = NewElectionWithCandidates(candidates)
 		case "Join Election":
-			election.ElectionKey = getElectionKeyTUI()
-			if election.ElectionKey == "" {
+			electionKey := getElectionKeyTUI()
+			if electionKey == "" {
 				fmt.Printf("no election key given; exiting\n")
 				os.Exit(0)
 			}
+			election = NewElectionWithElectionKey(electionKey)
 		}
 	}
 	modal := tview.NewModal().

@@ 54,7 54,8 @@ func stripNewline(str string) string {
 	return str
 }
 
-func handleCmd(cmd string, rw *bufio.ReadWriter, stream network.Stream, election *Election, localVoter *LocalVoter) {
+func handleCmd(cmd string, rw *bufio.ReadWriter, stream network.Stream, election *Election) {
+	localVoter := election.localVoter
 	switch cmd {
 	case "info":
 		rw.WriteString(fmt.Sprintf("%s\n", election.rendezvousNonce))


@@ 152,7 153,7 @@ func handleCmd(cmd string, rw *bufio.ReadWriter, stream network.Stream, election
 	}
 }
 
-func streamHandler(stream network.Stream, election *Election, localVoter *LocalVoter) {
+func streamHandler(stream network.Stream, election *Election) {
 	logger.Info("got a new stream:", stream)
 	logger.Info("remote peer:", stream.Conn().RemotePeer())
 	rw := bufio.NewReadWriter(bufio.NewReader(stream), bufio.NewWriter(stream))


@@ 164,11 165,12 @@ func streamHandler(stream network.Stream, election *Election, localVoter *LocalV
 	cmd = stripNewline(cmd)
 
 	logger.Info("cmd:", cmd)
-	handleCmd(cmd, rw, stream, election, localVoter)
+	handleCmd(cmd, rw, stream, election)
 	stream.Close()
 }
 
-func findPeers(closeElection <-chan int, election *Election, localVoter *LocalVoter) {
+func findPeers(closeElection <-chan int, election *Election) {
+	localVoter := election.localVoter
 	routingDiscovery := discovery.NewRoutingDiscovery(localVoter.kdht)
 	logger.Info("announcing ourselves")
 	discovery.Advertise(localVoter.ctx, routingDiscovery, string(election.rendezvousNonce))


@@ 217,13 219,14 @@ func findPeers(closeElection <-chan int, election *Election, localVoter *LocalVo
 	logger.Info("done finding peers")
 }
 
-func (voter *Voter) fetchNumber(election *Election, localVoter *LocalVoter, cmd string, args ...string) *big.Int {
+func (voter *Voter) fetchNumber(election *Election, cmd string, args ...string) *big.Int {
 	printErr := func(err error, msg string) {
 		logger.Errorf("%s: %s while fetcing `%s'; retrying in 2 seconds",
 			voter.addrInfo.ID, msg, cmd)
 		time.Sleep(time.Second * 2)
 	}
 retry:
+	localVoter := election.localVoter
 	stream, err := localVoter.NewStream(localVoter.ctx, voter.addrInfo.ID, protocolID)
 	if err != nil {
 		printErr(err, "couldn't open stream")


@@ 267,7 270,9 @@ retry:
 	return new(big.Int).SetBytes(retBytes)
 }
 
-func startVoting(election *Election, localVoter *LocalVoter) {
+func startVoting(election *Election) {
+	localVoter := election.localVoter
+
 	var err error
 	localVoter.inputMu.Lock()
 	localVoter.input, err = RandomBigInt(128, false)


@@ 291,7 296,7 @@ func startVoting(election *Election, localVoter *LocalVoter) {
 	for _, voter := range election.remoteVoters {
 		wg.Add(1)
 		go func(voter *Voter) {
-			voter.output = voter.fetchNumber(election, localVoter, "eval", base58.Encode(localVoter.input.Bytes()))
+			voter.output = voter.fetchNumber(election, "eval", base58.Encode(localVoter.input.Bytes()))
 			logger.Infof("voter %s output: %s", voter.addrInfo.ID, voter.output)
 			wg.Done()
 		}(voter)


@@ 311,7 316,7 @@ func startVoting(election *Election, localVoter *LocalVoter) {
 	for _, voter := range election.remoteVoters {
 		wg.Add(1)
 		go func(voter *Voter) {
-			voter.sum = voter.fetchNumber(election, localVoter, "sum")
+			voter.sum = voter.fetchNumber(election, "sum")
 			logger.Infof("voter %s sum: %s",
 				voter.addrInfo.ID, voter.sum)
 			wg.Done()


@@ 319,7 324,7 @@ func startVoting(election *Election, localVoter *LocalVoter) {
 	}
 	wg.Wait()
 
-	mat := constructPolyMatrix(election, localVoter)
+	mat := constructPolyMatrix(election)
 	mat.RREF()
 
 	constant := mat[0][len(mat[0])-1]


@@ 339,7 344,7 @@ func startVoting(election *Election, localVoter *LocalVoter) {
 	select {}
 }
 
-func constructPolyMatrix(election *Election, localVoter *LocalVoter) Matrix {
+func constructPolyMatrix(election *Election) Matrix {
 	mat := make(Matrix, len(election.remoteVoters) + 1) // row for everyone (including ourselves)
 
 	i := 0


@@ 359,6 364,7 @@ func constructPolyMatrix(election *Election, localVoter *LocalVoter) Matrix {
 	mat[i] = make([]big.Rat, len(mat) + 1)
 	row := mat[i]
 	row[0].SetInt64(1)
+	localVoter := election.localVoter
 	var j int64
 	for j = 1; j <= int64(len(election.remoteVoters)); j++ {
 		row[j].SetInt(new(big.Int).Exp(localVoter.input, big.NewInt(j), nil))