From a0cfdb69853b8245a2deae1a12721b123f39da89 Mon Sep 17 00:00:00 2001
From: David Florness <david@florness.com>
Date: Sun, 17 Jan 2021 22:06:51 -0500
Subject: [PATCH] A few fixes

---
 election/msg.go   |  8 +++++---
 election/voter.go |  3 ++-
 ui/tui.go         | 29 +++++++++++++++--------------
 3 files changed, 22 insertions(+), 18 deletions(-)

diff --git a/election/msg.go b/election/msg.go
index 47c4807..5ffd711 100644
--- a/election/msg.go
+++ b/election/msg.go
@@ -219,11 +219,12 @@ func OnSumMessage(client *mautrix.Client, evt *event.Event, elections *Elections
 		log.Warnf("ignoring %s's sum since we couldn't cast message content to SumMessageContent", evt.Sender)
 		return
 	}
-	sum, ok := new(big.Int).SetString(content.Sum, 64)
-	if !ok {
-		log.Warnf("ignoring %s's sum since we couldn't base64 decode the sum", evt.Sender)
+	bytes, err := base64.StdEncoding.DecodeString(content.Sum)
+	if err != nil {
+		log.Warnf("ignoring %s's sum since we couldn't decode their sum", evt.Sender)
 		return
 	}
+	sum := new(big.Int).SetBytes(bytes)
 	el := getElection(client, evt.RoomID, content.CreateEventId, elections)
 	if el == nil {
 		log.Warnf("ignoring %s's sum since the election does not exist", evt.Sender)
@@ -231,6 +232,7 @@ func OnSumMessage(client *mautrix.Client, evt *event.Event, elections *Elections
 	}
 	el.Lock()
 	defer el.Unlock()
+	// TODO: ensure voter exists
 	voter := el.Voters[evt.Sender]
 	voter.Sum = sum
 }
diff --git a/election/voter.go b/election/voter.go
index 9d2a5fe..8806a54 100644
--- a/election/voter.go
+++ b/election/voter.go
@@ -121,6 +121,7 @@ func (locelVoter *LocalVoter) SendSum(client *mautrix.Client, el *Election) erro
 				time.Sleep(time.Millisecond * 100)
 			}
 			sum.Add(sum, voter.Output)
+			wg.Done()
 		}(voter)
 	}
 	el.RUnlock()
@@ -141,7 +142,7 @@ func GetSums(client *mautrix.Client, el *Election) {
 			for voter.Sum == nil {
 				time.Sleep(time.Millisecond * 100)
 			}
-
+			wg.Done()
 		}(voter)
 	}
 	el.RUnlock()
diff --git a/ui/tui.go b/ui/tui.go
index 1e838f4..3ad20b1 100644
--- a/ui/tui.go
+++ b/ui/tui.go
@@ -138,6 +138,7 @@ func RoomTUI(client *mautrix.Client, roomID id.RoomID, elections *election.Elect
 		if err != nil {
 			panic(err)
 		}
+		ballot = ElectionTUI(client, el, localVoter)
 	})
 	if err := app.SetRoot(list, true).SetFocus(list).Run(); err != nil {
 		panic(err)
@@ -224,27 +225,27 @@ func CreateElectionTUI() (title string, candidates []election.Candidate) {
 func ElectionTUI(client *mautrix.Client, el *election.Election, localVoter *election.LocalVoter) (ballot []byte) {
 	votersTextView := tview.NewTextView()
 	frame := tview.NewFrame(votersTextView)
+	app := tview.NewApplication()
 	el.RLock()
 	if el.Creator == localVoter.UserID {
 		frame.AddText("Press enter to start the election", false, tview.AlignCenter, tcell.ColorWhite)
+		app.SetInputCapture(func(event *tcell.EventKey) *tcell.EventKey {
+			if event.Key() == tcell.KeyEnter {
+				app.QueueUpdateDraw(func() {
+					frame.Clear()
+					frame.AddText("Starting election...", false, tview.AlignCenter, tcell.ColorWhite)
+				})
+				err := election.StartElection(client, el)
+				if err != nil {
+					panic(err)
+				}
+			}
+			return event
+		})
 	} else {
 		frame.AddText("Waiting for election to start...", false, tview.AlignCenter, tcell.ColorWhite)
 	}
 	el.RUnlock()
-	app := tview.NewApplication()
-	app.SetInputCapture(func(event *tcell.EventKey) *tcell.EventKey {
-		if event.Key() == tcell.KeyEnter {
-			app.QueueUpdateDraw(func() {
-				frame.Clear()
-				frame.AddText("Starting election...", false, tview.AlignCenter, tcell.ColorWhite)
-			})
-			err := election.StartElection(client, el)
-			if err != nil {
-				panic(err)
-			}
-		}
-		return event
-	})
 	update := func() {
 		el.RLock()
 		voters := make([]string, 0, len(el.Voters))
-- 
2.38.4