From d7cefb522b7886df5c4e694cc1bef57442b7863f Mon Sep 17 00:00:00 2001
From: Szymon Szlachtowicz <38212223+Szymx95@users.noreply.github.com>
Date: Mon, 6 Sep 2021 11:34:42 +0200
Subject: [PATCH] Check tokens when creating polls and votes (#53)

---
 packages/core/src/classes/WakuPolling.ts      | 45 ++++++++++++--
 .../src/components/Poll.tsx                   | 20 ++++++-
 .../src/components/PollCreation.tsx           | 58 +++++++++++++------
 3 files changed, 96 insertions(+), 27 deletions(-)

diff --git a/packages/core/src/classes/WakuPolling.ts b/packages/core/src/classes/WakuPolling.ts
index ca1c3e4..6a4a539 100644
--- a/packages/core/src/classes/WakuPolling.ts
+++ b/packages/core/src/classes/WakuPolling.ts
@@ -16,6 +16,13 @@ const ABI = [
   'function aggregate(tuple(address target, bytes callData)[] calls) view returns (uint256 blockNumber, bytes[] returnData)',
 ]
 
+export enum MESSEGAGE_SENDING_RESULT {
+  ok = 0,
+  notEnoughToken = 1,
+  errorCreatingMessage = 2,
+  pollNotFound = 3,
+}
+
 export class WakuPolling extends WakuVoting {
   protected multicall: string
 
@@ -78,8 +85,19 @@ export class WakuPolling extends WakuVoting {
     minToken?: BigNumber,
     endTime?: number
   ) {
-    const pollInit = await PollInitMsg.create(signer, question, answers, pollType, this.chainId, minToken, endTime)
-    await this.sendWakuMessage(this.wakuMessages['pollInit'], pollInit)
+    const address = await signer.getAddress()
+    await this.updateBalances(address)
+    if (this.addressesBalances[address] && this.addressesBalances[address]?.gt(minToken ?? BigNumber.from(0))) {
+      const pollInit = await PollInitMsg.create(signer, question, answers, pollType, this.chainId, minToken, endTime)
+      if (pollInit) {
+        await this.sendWakuMessage(this.wakuMessages['pollInit'], pollInit)
+        return MESSEGAGE_SENDING_RESULT.ok
+      } else {
+        return MESSEGAGE_SENDING_RESULT.errorCreatingMessage
+      }
+    } else {
+      return MESSEGAGE_SENDING_RESULT.notEnoughToken
+    }
   }
 
   public async sendTimedPollVote(
@@ -88,18 +106,35 @@ export class WakuPolling extends WakuVoting {
     selectedAnswer: number,
     tokenAmount?: BigNumber
   ) {
-    const pollVote = await TimedPollVoteMsg.create(signer, pollId, selectedAnswer, this.chainId, tokenAmount)
-    await this.sendWakuMessage(this.wakuMessages['pollVote'], pollVote)
+    const address = await signer.getAddress()
+    const poll = this.wakuMessages['pollInit'].arr.find((poll: PollInitMsg): poll is PollInitMsg => poll.id === pollId)
+    if (poll) {
+      await this.updateBalances(address)
+      if (this.addressesBalances[address] && this.addressesBalances[address]?.gt(poll.minToken ?? BigNumber.from(0))) {
+        const pollVote = await TimedPollVoteMsg.create(signer, pollId, selectedAnswer, this.chainId, tokenAmount)
+        if (pollVote) {
+          await this.sendWakuMessage(this.wakuMessages['pollVote'], pollVote)
+        } else {
+          return MESSEGAGE_SENDING_RESULT.errorCreatingMessage
+        }
+      } else {
+        return MESSEGAGE_SENDING_RESULT.notEnoughToken
+      }
+    } else {
+      return MESSEGAGE_SENDING_RESULT.pollNotFound
+    }
   }
 
   protected addressesBalances: { [address: string]: BigNumber } = {}
   protected lastBlockBalances = 0
 
-  protected async updateBalances() {
+  protected async updateBalances(newAddress?: string) {
     const addresses: string[] = [
       ...this.wakuMessages['pollInit'].arr.map((msg) => msg.owner),
       ...this.wakuMessages['pollVote'].arr.map((msg) => msg.voter),
     ]
+
+    if (newAddress) addresses.push(newAddress)
     const addressesToUpdate: { [addr: string]: boolean } = {}
 
     const addAddressToUpdate = (addr: string) => {
diff --git a/packages/polling-components/src/components/Poll.tsx b/packages/polling-components/src/components/Poll.tsx
index 5f345be..cba786f 100644
--- a/packages/polling-components/src/components/Poll.tsx
+++ b/packages/polling-components/src/components/Poll.tsx
@@ -8,7 +8,7 @@ import styled from 'styled-components'
 import { RadioGroup, SmallButton } from '@status-waku-voting/react-components'
 import { PollResults } from './PollResults'
 import { useEthers } from '@usedapp/core'
-
+import { Modal } from '@status-waku-voting/react-components'
 type PollProps = {
   poll: DetailedTimedPoll
   wakuPolling: WakuPolling | undefined
@@ -21,6 +21,7 @@ export function Poll({ poll, wakuPolling, signer }: PollProps) {
   const [tokenAmount, setTokenAmount] = useState(0)
   const [address, setAddress] = useState('')
   const [userInVoters, setUserInVoters] = useState(-1)
+  const [showNotEnoughTokens, setShowNotEnoughTokens] = useState(false)
 
   useEffect(() => {
     if (signer) {
@@ -63,14 +64,17 @@ export function Poll({ poll, wakuPolling, signer }: PollProps) {
       {userInVoters < 0 && (
         <SmallButton
           disabled={!signer || !account}
-          onClick={() => {
+          onClick={async () => {
             if (wakuPolling && signer) {
-              wakuPolling.sendTimedPollVote(
+              const result = await wakuPolling.sendTimedPollVote(
                 signer,
                 poll.poll.id,
                 selectedAnswer ?? 0,
                 poll.poll.pollType === PollType.WEIGHTED ? BigNumber.from(tokenAmount) : undefined
               )
+              if (result === 1) {
+                setShowNotEnoughTokens(true)
+              }
             }
           }}
         >
@@ -78,10 +82,20 @@ export function Poll({ poll, wakuPolling, signer }: PollProps) {
           Vote
         </SmallButton>
       )}
+      {showNotEnoughTokens && (
+        <Modal heading={''} setShowModal={setShowNotEnoughTokens}>
+          <ModalTextWrapper>You don't have enough tokens to vote</ModalTextWrapper>
+        </Modal>
+      )}
     </PollWrapper>
   )
 }
 
+const ModalTextWrapper = styled.div`
+  text-align: center;
+  margin: 32px 0;
+`
+
 const VotingWrapper = styled.div`
   margin-top: 40px;
   margin-left: 32px;
diff --git a/packages/polling-components/src/components/PollCreation.tsx b/packages/polling-components/src/components/PollCreation.tsx
index 6a3e467..d4e98e3 100644
--- a/packages/polling-components/src/components/PollCreation.tsx
+++ b/packages/polling-components/src/components/PollCreation.tsx
@@ -1,4 +1,4 @@
-import React, { useEffect, useState } from 'react'
+import React, { ReactNode, useEffect, useState } from 'react'
 import { Wallet } from 'ethers'
 import { JsonRpcSigner } from '@ethersproject/providers'
 import styled from 'styled-components'
@@ -13,6 +13,27 @@ function getLocaleIsoTime(dateTime: Date) {
   return newDate.toISOString().slice(0, -8)
 }
 
+type ConfirmScreenProps = {
+  setShowModal: (val: boolean) => void
+  children: ReactNode
+}
+
+function ConfirmScreen({ children, setShowModal }: ConfirmScreenProps) {
+  return (
+    <Confirmation>
+      <ConfirmationText>{children}</ConfirmationText>
+      <SmallButton
+        onClick={(e) => {
+          e.preventDefault()
+          setShowModal(false)
+        }}
+      >
+        Close
+      </SmallButton>
+    </Confirmation>
+  )
+}
+
 type PollCreationProps = {
   signer: JsonRpcSigner | Wallet
   wakuPolling: WakuPolling | undefined
@@ -23,13 +44,14 @@ export function PollCreation({ signer, wakuPolling, setShowPollCreation }: PollC
   const [answers, setAnswers] = useState<string[]>(['', ''])
   const [question, setQuestion] = useState('')
   const [showCreateConfirmation, setShowCreateConfirmation] = useState(false)
+  const [showNotEnoughTokens, setShowNotEnoughTokens] = useState(false)
   const [selectedType, setSelectedType] = useState(PollType.NON_WEIGHTED)
   const [endTimePicker, setEndTimePicker] = useState(new Date(new Date().getTime() + 10000000))
 
   return (
     <Modal heading="Create a poll" setShowModal={setShowPollCreation}>
       <NewPollBox>
-        {!showCreateConfirmation && (
+        {!showCreateConfirmation && !showNotEnoughTokens && (
           <PollForm>
             <Input
               label={'Question or title of the poll'}
@@ -65,7 +87,7 @@ export function PollCreation({ signer, wakuPolling, setShowPollCreation }: PollC
             <SmallButton
               onClick={async (e) => {
                 e.preventDefault()
-                await wakuPolling?.createTimedPoll(
+                const result = await wakuPolling?.createTimedPoll(
                   signer,
                   question,
                   answers,
@@ -73,7 +95,12 @@ export function PollCreation({ signer, wakuPolling, setShowPollCreation }: PollC
                   undefined,
                   endTimePicker.getTime()
                 )
-                setShowCreateConfirmation(true)
+                if (result === 0) {
+                  setShowCreateConfirmation(true)
+                }
+                if (result === 1) {
+                  setShowNotEnoughTokens(true)
+                }
               }}
             >
               Create a poll
@@ -82,21 +109,14 @@ export function PollCreation({ signer, wakuPolling, setShowPollCreation }: PollC
         )}
 
         {showCreateConfirmation && (
-          <Confirmation>
-            <ConfirmationText>
-              Your poll has been created!
-              <br />
-              It will appear at the top of the poll list.
-            </ConfirmationText>
-            <SmallButton
-              onClick={(e) => {
-                e.preventDefault()
-                setShowPollCreation(false)
-              }}
-            >
-              Close
-            </SmallButton>
-          </Confirmation>
+          <ConfirmScreen setShowModal={setShowPollCreation}>
+            Your poll has been created!
+            <br />
+            It will appear at the top of the poll list.
+          </ConfirmScreen>
+        )}
+        {showNotEnoughTokens && (
+          <ConfirmScreen setShowModal={setShowPollCreation}>You don't have enough tokens to create poll</ConfirmScreen>
         )}
       </NewPollBox>{' '}
     </Modal>