4 files changed, 146 insertions(+), 117 deletions(-)

M main.go
M merkle.go
M ui.go
M voter.go

@@ 9,36 9,26 @@ import (
 	"os"
 	"sync"
 
-	"github.com/cbergoon/merkletree"
 	"github.com/ipfs/go-log"
 	"github.com/libp2p/go-libp2p"
 	"github.com/libp2p/go-libp2p-core/host"
+	"github.com/libp2p/go-libp2p-core/network"
 	"github.com/libp2p/go-libp2p-core/peer"
 	"github.com/libp2p/go-libp2p-core/protocol"
 	dht "github.com/libp2p/go-libp2p-kad-dht"
 	routing "github.com/libp2p/go-libp2p-routing"
 	"github.com/mr-tron/base58/base58"
-	"github.com/rivo/tview"
 	"github.com/whyrusleeping/go-logging"
 )
 
-var (
-	logger          = log.Logger("tallyard")
-	protocolID      = protocol.ID("/tallyard/0.0.0")
-	candidates      []Candidate
-	optionsMerkle   *merkletree.MerkleTree
-	rendezvousNonce Nonce
-	merkleRoot      []byte
-	me              Me
-	election        Election
-)
-
-func bootstrap() {
-	var err error
+const protocolID = protocol.ID("/tallyard/0.0.0")
+var logger = log.Logger("tallyard")
 
+func bootstrap(election *Election, me *Me, merkleRoot []byte) {
 	me.ctx = context.Background()
 	election.remoteVoters = make(map[peer.ID]*Voter)
 
+	var err error
 	me.Host, err = libp2p.New(me.ctx,
 		libp2p.Routing(func(h host.Host) (routing.PeerRouting, error) {
 			var err error


@@ 77,19 67,19 @@ func bootstrap() {
 
 	if election.masterID == "" { // we are the master
 		fmt.Println("share this with peers:")
-		fmt.Printf("%s0%s\n",
-			base58.Encode(optionsMerkle.MerkleRoot()),
-			me.ID())
+		fmt.Printf("%s0%s\n", base58.Encode(merkleRoot), me.ID())
 
 		logger.Info("waiting for incoming streams and finding voters...")
 
 		election.Lock()
 		ch := make(chan int, 1)
 		election.close = ch
-		go findPeers(ch)
+		go findPeers(ch, election, me)
 		election.Unlock()
 
-		me.SetStreamHandler(protocolID, streamHandler)
+		me.SetStreamHandler(protocolID, func(stream network.Stream) {
+			streamHandler(stream, election, me)
+		})
 
 		fmt.Println("press ENTER to solidify group of voters and start voting")
 		stdReader := bufio.NewReader(os.Stdin)


@@ 139,7 129,7 @@ func bootstrap() {
 			panic(err)
 		}
 		str = stripNewline(str)
-		rendezvousNonce = Nonce(str)
+		election.rendezvousNonce = Nonce(str)
 		// remaining fields are the candidates
 		for {
 			str, err := rw.ReadString('\n')


@@ 148,7 138,7 @@ func bootstrap() {
 			}
 			str = stripNewline(str)
 			if str != "" {
-				candidates = append(candidates, Candidate(str))
+				election.candidates = append(election.candidates, Candidate(str))
 			}
 			if err == io.EOF {
 				break


@@ 157,14 147,16 @@ func bootstrap() {
 		logger.Info("done fetching election info")
 
 		logger.Info("checking authenticity of election info...")
-		verifyElectionInfo()
+		verifyElectionInfo(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
-		me.SetStreamHandler(protocolID, streamHandler)
-		findPeers(ch)
+		me.SetStreamHandler(protocolID, func(stream network.Stream) {
+			streamHandler(stream, election, me)
+		})
+		findPeers(ch, election, me)
 	}
 }
 


@@ 181,24 173,25 @@ func main() {
 		log.SetLogLevel("tallyard", "critical")
 	}
 
-	app := tview.NewApplication()
-	modal := tview.NewModal().
-		SetText("Welcome to tallyard!").
-		AddButtons([]string{"Create Election", "Join Election"}).
-		SetDoneFunc(func(buttonIndex int, buttonLabel string) {
-			app.Stop()
-			switch buttonLabel {
-			case "Create Election":
-				createElection()
-			case "Join Election":
-				joinElection()
-			}
-		})
-	if err := app.SetRoot(modal, false).EnableMouse(true).Run(); err != nil {
-		panic(err)
+	var (
+		election   *Election
+		me         = new(Me)
+		merkleRoot []byte
+	)
+
+	// check if election key was given via cli
+	if electionKey := flag.Arg(0); electionKey != "" {
+		election = new(Election)
+		merkleRoot, election.masterID = splitElectionKey(electionKey)
+	} else {
+		election, merkleRoot = tui()
+		// tui button was never pressed
+		if election == nil {
+			os.Exit(0)
+		}
 	}
 
-	bootstrap()
+	bootstrap(election, me, merkleRoot)
 
-	startVoting()
+	startVoting(election, me)
 }

@@ 6,8 6,10 @@ import (
 	"crypto/sha256"
 	"fmt"
 	"os"
+	"strings"
 
 	"github.com/cbergoon/merkletree"
+	"github.com/libp2p/go-libp2p-core/peer"
 	"github.com/mr-tron/base58/base58"
 )
 


@@ 52,13 54,13 @@ func (n Nonce) Equals(other merkletree.Content) (bool, error) {
 	return n == other.(Nonce), nil
 }
 
-func verifyElectionInfo() {
-	content := []merkletree.Content{rendezvousNonce}
+func verifyElectionInfo(election *Election, merkleRoot []byte) {
+	content := []merkletree.Content{election.rendezvousNonce}
 	var err error
-	for _, eo := range candidates {
+	for _, eo := range election.candidates {
 		content = append(content, eo)
 	}
-	optionsMerkle, err = merkletree.NewTree(content)
+	optionsMerkle, err := merkletree.NewTree(content)
 	if err != nil {
 		panic(err)
 	}


@@ 69,3 71,22 @@ func verifyElectionInfo() {
 		os.Exit(1)
 	}
 }
+
+func splitElectionKey(electionKey string) (merkleRoot []byte, masterID peer.ID) {
+	var err error
+	zeroi := strings.IndexByte(electionKey, '0')
+	if zeroi == -1 {
+		panic("invalid election key")
+	}
+	logger.Info("merkle root:", electionKey[:zeroi])
+	merkleRoot, err = base58.Decode(electionKey[:zeroi])
+	if err != nil {
+		panic(err)
+	}
+	masterID, err = peer.Decode(electionKey[zeroi+1:])
+	if err != nil {
+		panic(err)
+	}
+	logger.Info("master ID:", masterID)
+	return merkleRoot, masterID
+}

@@ 4,81 4,93 @@ import (
 	"fmt"
 	"sort"
 	"strconv"
-	"strings"
 	"unicode"
 
 	"github.com/cbergoon/merkletree"
 	"github.com/libp2p/go-libp2p-core/peer"
-	"github.com/mr-tron/base58/base58"
 	"github.com/rivo/tview"
 )
 
-func createElection() {
+func tui() (election *Election, merkleRoot []byte) {
+	app := tview.NewApplication()
+	done := func(buttonIndex int, buttonLabel string) {
+		app.Stop()
+		election = new(Election)
+		switch buttonLabel {
+		case "Create Election":
+			merkleRoot, election.rendezvousNonce, election.candidates = createElection()
+		case "Join Election":
+			merkleRoot, election.masterID = joinElection()
+		}
+	}
+	modal := tview.NewModal().
+		SetText("Welcome to tallyard!").
+		AddButtons([]string{"Create Election", "Join Election"}).
+		SetDoneFunc(done)
+	if err := app.SetRoot(modal, false).EnableMouse(true).Run(); err != nil {
+		panic(err)
+	}
+	return election, merkleRoot
+}
+
+func createElection() (merkleRoot []byte, rendezvousNonce Nonce, candidates []Candidate) {
 	var form *tview.Form
 	n := 3
 	app := tview.NewApplication()
+	plus := func() {
+		form.AddInputField(fmt.Sprintf("%d.", n), "", 50, nil, nil)
+		n++
+	}
+	minus := func() {
+		// TODO: ensure from joiner that there are at least two
+		// candidates
+		if n > 3 {
+			form.RemoveFormItem(n - 2)
+			n--
+		}
+	}
+	done := func() {
+		// TODO: ensure none of the candidates are empty
+		app.Stop()
+		rendezvousNonce = NewNonce()
+		content := []merkletree.Content{rendezvousNonce}
+		for i := 0; i < n-1; i++ {
+			eo := Candidate(form.GetFormItem(i).(*tview.InputField).GetText())
+			candidates = append(candidates, eo)
+			content = append(content, eo)
+		}
+		optionsMerkle, err := merkletree.NewTree(content)
+		if err != nil {
+			panic(err)
+		}
+		merkleRoot = optionsMerkle.MerkleRoot()
+	}
 	form = tview.NewForm().
 		AddInputField("1.", "", 50, nil, nil).
 		AddInputField("2.", "", 50, nil, nil).
-		AddButton("+", func() {
-			form.AddInputField(fmt.Sprintf("%d.", n), "", 50, nil, nil)
-			n++
-		}).
-		AddButton("-", func() {
-			// TODO: ensure from joiner that there are at least two
-			// candidates
-			if n > 3 {
-				form.RemoveFormItem(n - 2)
-				n--
-			}
-		}).
-		AddButton("Done", func() {
-			// TODO: ensure none of the candidates are empty
-			app.Stop()
-			rendezvousNonce = NewNonce()
-			content := []merkletree.Content{rendezvousNonce}
-			for i := 0; i < n-1; i++ {
-				eo := Candidate(form.GetFormItem(i).(*tview.InputField).GetText())
-				candidates = append(candidates, eo)
-				content = append(content, eo)
-			}
-			var err error
-			optionsMerkle, err = merkletree.NewTree(content)
-			if err != nil {
-				panic(err)
-			}
-		})
+		AddButton("+", plus).
+		AddButton("-", minus).
+		AddButton("Done", done)
 	if err := app.SetRoot(form, true).EnableMouse(true).Run(); err != nil {
 		panic(err)
 	}
+	return merkleRoot, rendezvousNonce, candidates
 }
 
-func joinElection() {
+func joinElection() (merkleRoot []byte, masterID peer.ID) {
 	app := tview.NewApplication()
 	var form *tview.Form
+	done := func() {
+		app.Stop()
+		merkleRoot, masterID = splitElectionKey(form.GetFormItem(0).(*tview.InputField).GetText())
+	}
 	form = tview.NewForm().
 		AddInputField("Election key:", "", 100, nil, nil).
-		AddButton("Continue", func() {
-			app.Stop()
-			electionKey := form.GetFormItem(0).(*tview.InputField).GetText()
-
-			zeroi := strings.IndexByte(electionKey, '0')
-			var err error
-			logger.Info("merkle root:", electionKey[:zeroi])
-			merkleRoot, err = base58.Decode(electionKey[:zeroi])
-			if err != nil {
-				panic(err)
-			}
-
-			election.masterID, err = peer.Decode(electionKey[zeroi+1:])
-			if err != nil {
-				panic(err)
-			}
-			logger.Info("master ID:", election.masterID)
-		})
+		AddButton("Continue", done)
 	if err := app.SetRoot(form, true).EnableMouse(true).Run(); err != nil {
 		panic(err)
 	}
+	return merkleRoot, masterID
 }
 
 // displays a voting UI to the user and returns the encoded ballot

@@ 52,10 52,13 @@ type Election struct {
 	// for master: signifies when user hits ENTER to close the election
 	//
 	// the number of peers know by master is passed through it
-	close        chan<- int
-	closed       bool // used by handleCmd to prevent closing election more than once
-	masterID     peer.ID
-	remoteVoters map[peer.ID]*Voter
+	close           chan<- int
+	closed          bool // used by handleCmd to prevent closing election more than once
+
+	candidates      []Candidate
+	masterID        peer.ID
+	remoteVoters    map[peer.ID]*Voter
+	rendezvousNonce Nonce
 }
 
 func stripNewline(str string) string {


@@ 68,11 71,11 @@ func stripNewline(str string) string {
 	return str
 }
 
-func handleCmd(cmd string, rw *bufio.ReadWriter, stream network.Stream) {
+func handleCmd(cmd string, rw *bufio.ReadWriter, stream network.Stream, election *Election, me *Me) {
 	switch cmd {
 	case "info":
-		rw.WriteString(fmt.Sprintf("%s\n", rendezvousNonce))
-		for _, option := range candidates {
+		rw.WriteString(fmt.Sprintf("%s\n", election.rendezvousNonce))
+		for _, option := range election.candidates {
 			rw.WriteString(fmt.Sprintf("%s\n", option))
 		}
 		rw.Flush()


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


@@ 178,18 181,18 @@ func streamHandler(stream network.Stream) {
 	cmd = stripNewline(cmd)
 
 	logger.Info("cmd:", cmd)
-	handleCmd(cmd, rw, stream)
+	handleCmd(cmd, rw, stream, election, me)
 	stream.Close()
 }
 
-func findPeers(closeElection <-chan int) {
+func findPeers(closeElection <-chan int, election *Election, me *Me) {
 	routingDiscovery := discovery.NewRoutingDiscovery(me.kdht)
 	logger.Info("announcing ourselves")
-	discovery.Advertise(me.ctx, routingDiscovery, string(rendezvousNonce))
+	discovery.Advertise(me.ctx, routingDiscovery, string(election.rendezvousNonce))
 	logger.Info("successfully announced!")
 
 	fmt.Println("finding other voters...")
-	peerChan, err := routingDiscovery.FindPeers(me.ctx, string(rendezvousNonce))
+	peerChan, err := routingDiscovery.FindPeers(me.ctx, string(election.rendezvousNonce))
 	if err != nil {
 		panic(err)
 	}


@@ 229,7 232,7 @@ func findPeers(closeElection <-chan int) {
 	logger.Info("done finding peers")
 }
 
-func (voter *Voter) fetchNumber(cmd string, args ...string) *big.Int {
+func (voter *Voter) fetchNumber(election *Election, me *Me, 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)


@@ 279,7 282,7 @@ retry:
 	return new(big.Int).SetBytes(retBytes)
 }
 
-func startVoting() {
+func startVoting(election *Election, me *Me) {
 	var err error
 	me.inputMu.Lock()
 	me.input, err = RandomBigInt(128, false)


@@ 289,7 292,7 @@ func startVoting() {
 	}
 	logger.Infof("our input: %s", me.input)
 
-	ballot := vote(candidates)
+	ballot := vote(election.candidates)
 	logger.Infof("our ballot: %v", ballot)
 
 	// no +1 since we want degree k-1 where k is total number of voters


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


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


@@ 331,7 334,7 @@ func startVoting() {
 	}
 	wg.Wait()
 
-	mat := constructPolyMatrix()
+	mat := constructPolyMatrix(election, me)
 	mat.RREF()
 
 	constant := mat[0][len(mat[0])-1]


@@ 342,16 345,16 @@ func startVoting() {
 	result := constant.Num().Bytes()
 
 	// number of bytes we need to insert at the front since they're zero
-	diff := (len(candidates)*len(candidates)) - len(result)
+	diff := (len(election.candidates)*len(election.candidates)) - len(result)
 	result = append(make([]byte, diff), result...)
 
-	printResults(result, candidates)
+	printResults(result, election.candidates)
 
 	// temporary
 	select {}
 }
 
-func constructPolyMatrix() Matrix {
+func constructPolyMatrix(election *Election, me *Me) Matrix {
 	mat := make(Matrix, len(election.remoteVoters) + 1) // row for everyone (including ourselves)
 
 	i := 0