CAS-204: Add NFT/FT restrictions to One Address One Vote strategy

backend/main/server/app.go

@@ -66,7 +66,7 @@ type Strategy interface {
 	TallyVotes(votes []*models.VoteWithBalance, p *models.ProposalResults) (models.ProposalResults, error)
 	GetVotes(votes []*models.VoteWithBalance, proposal *models.Proposal) ([]*models.VoteWithBalance, error)
 	GetVoteWeightForBalance(vote *models.VoteWithBalance, proposal *models.Proposal) (float64, error)
-	InitStrategy(f *shared.FlowAdapter, db *shared.Database, sc *shared.SnapshotClient)
+	InitStrategy(f *shared.FlowAdapter, db *shared.Database, sc *shared.SnapshotClient, name string)
 	FetchBalance(b *models.Balance, p *models.Proposal) (*models.Balance, error)
 	RequiresSnapshot() bool
 }
@@ -75,6 +75,8 @@ var strategyMap = map[string]Strategy{
 	"token-weighted-default":        &strategies.TokenWeightedDefault{},
 	"staked-token-weighted-default": &strategies.StakedTokenWeightedDefault{},
 	"one-address-one-vote":          &strategies.OneAddressOneVote{},
+	"one-address-one-vote-ft":       &strategies.OneAddressOneVote{},
+	"one-address-one-vote-nft":      &strategies.OneAddressOneVote{},
 	"balance-of-nfts":               &strategies.BalanceOfNfts{},
 }
 
@@ -322,7 +324,7 @@ func (a *App) getResultsForProposal(w http.ResponseWriter, r *http.Request) {
 	}
 
 	proposalWithChoices := models.NewProposalResults(proposalId, p.Choices)
-	s.InitStrategy(a.FlowAdapter, a.DB, a.SnapshotClient)
+	s.InitStrategy(a.FlowAdapter, a.DB, a.SnapshotClient, *p.Strategy)
 
 	proposalResults, err := s.TallyVotes(votes, proposalWithChoices)
 	if err != nil {
@@ -379,7 +381,7 @@ func (a *App) getVotesForProposal(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	s.InitStrategy(a.FlowAdapter, a.DB, a.SnapshotClient)
+	s.InitStrategy(a.FlowAdapter, a.DB, a.SnapshotClient, *proposal.Strategy)
 
 	votes, totalRecords, err := models.GetVotesForProposal(a.DB, start, count, order, proposalId, *proposal.Strategy)
 	if err != nil {
@@ -590,7 +592,7 @@ func (a *App) createVoteForProposal(w http.ResponseWriter, r *http.Request) {
 		respondWithError(w, http.StatusBadRequest, err.Error())
 		return
 	}
-	// validate user signature
+	// validate transaction signature
 	// if err := a.FlowAdapter.UserTransactionValidate(v.Addr, v.Message, v.Composite_signatures, v.TransactionId, a.TxOptionsAddresses, p.Choices); err != nil {
 	// 	respondWithError(w, http.StatusBadRequest, err.Error())
 	// 	return
@@ -612,7 +614,7 @@ func (a *App) createVoteForProposal(w http.ResponseWriter, r *http.Request) {
 		emptyBalance.BlockHeight = *p.Block_height
 	}
 
-	s.InitStrategy(a.FlowAdapter, a.DB, a.SnapshotClient)
+	s.InitStrategy(a.FlowAdapter, a.DB, a.SnapshotClient, *p.Strategy)
 
 	balance, err := s.FetchBalance(emptyBalance, &p)
 	if err != nil {
@@ -795,7 +797,7 @@ func (a *App) createProposal(w http.ResponseWriter, r *http.Request) {
 
 	}
 	s := strategyMap[*p.Strategy]
-	s.InitStrategy(a.FlowAdapter, a.DB, a.SnapshotClient)
+	s.InitStrategy(a.FlowAdapter, a.DB, a.SnapshotClient, *p.Strategy)
 
 	var snapshotResponse *shared.SnapshotResponse
 	if s.RequiresSnapshot() {

backend/main/strategies/balance_of_nfts.go

@@ -4,62 +4,20 @@ import (
 	"github.com/DapperCollectives/CAST/backend/main/models"
 	"github.com/DapperCollectives/CAST/backend/main/shared"
 	s "github.com/DapperCollectives/CAST/backend/main/shared"
-	"github.com/rs/zerolog/log"
 )
 
 type BalanceOfNfts struct {
 	s.StrategyStruct
 	SC s.SnapshotClient
 	DB *s.Database
+	name string
 }
 
 func (b *BalanceOfNfts) FetchBalance(
 	balance *models.Balance,
 	p *models.Proposal,
 ) (*models.Balance, error) {
-
-	v := models.Vote{Proposal_id: balance.Proposal_id, Addr: balance.Addr}
-	vb := &models.VoteWithBalance{
-		NFTs: []*models.NFT{},
-		Vote: v,
-	}
-
-	var c models.Community
-	if err := c.GetCommunityByProposalId(b.DB, balance.Proposal_id); err != nil {
-		return nil, err
-	}
-
-	strategy, err := models.MatchStrategyByProposal(*c.Strategies, *p.Strategy)
-	if err != nil {
-		log.Error().Err(err).Msg("Unable to find strategy for contract.")
-		return nil, err
-	}
-
-	nftIds, err := b.FlowAdapter.GetNFTIds(balance.Addr, &strategy.Contract)
-	if err != nil {
-		return nil, err
-	}
-
-	for _, nftId := range nftIds {
-		nft := &models.NFT{
-			ID: nftId,
-		}
-		vb.NFTs = append(vb.NFTs, nft)
-	}
-
-	doesExist, err := models.DoesNFTExist(b.DB, vb)
-	if err != nil {
-		return nil, err
-	}
-
-	//only if the NFT ID is not already in the DB,
-	//do we add the balance
-	if !doesExist && err == nil {
-		err = models.CreateUserNFTRecord(b.DB, vb)
-		balance.NFTCount = len(vb.NFTs)
-	}
-
-	return balance, nil
+	return FetchNFTBalance(b.DB, b.FlowAdapter, balance, p)
 }
 
 func (b *BalanceOfNfts) TallyVotes(
@@ -104,8 +62,10 @@ func (s *BalanceOfNfts) InitStrategy(
 	f *shared.FlowAdapter,
 	db *shared.Database,
 	sc *s.SnapshotClient,
+	name string,
 ) {
 	s.FlowAdapter = f
 	s.DB = db
 	s.SC = *sc
+	s.name = name
 }

backend/main/strategies/one_address_one_vote.go

@@ -14,26 +14,34 @@ type OneAddressOneVote struct {
 	s.StrategyStruct
 	SC s.SnapshotClient
 	DB *s.Database
+	name string
 }
 
 func (s *OneAddressOneVote) FetchBalance(
 	b *models.Balance,
 	p *models.Proposal,
 ) (*models.Balance, error) {
 
-	if err := b.GetBalanceByAddressAndBlockHeight(s.DB); err != nil && err.Error() != pgx.ErrNoRows.Error() {
-		log.Error().Err(err).Msg("Error querying address b at blockheight.")
-		return nil, err
-	}
-
-	if b.ID == "" {
-		if err := b.CreateBalance(s.DB); err != nil {
-			log.Error().Err(err).Msg("Error saving b to database.")
+	if s.name == "one-address-one-vote-nft" {
+		return FetchNFTBalance(s.DB, s.FlowAdapter, b, p)
+	} else if s.name == "one-address-one-vote-ft" {
+		return FetchFTBalance(s.DB, s.SC, b, p)
+	} else {
+		// DEPRECATED: original one-address-one-vote balance implementation before NFT/FT restrictions were introduced
+		if err := b.GetBalanceByAddressAndBlockHeight(s.DB); err != nil && err.Error() != pgx.ErrNoRows.Error() {
+			log.Error().Err(err).Msg("error querying address b at blockheight")
 			return nil, err
 		}
+
+		if b.ID == "" {
+			if err := b.CreateBalance(s.DB); err != nil {
+				log.Error().Err(err).Msg("error saving b to DB")
+				return nil, err
+			}
+		}
+
+		return b, nil
 	}
-
-	return b, nil
 }
 
 func (s *OneAddressOneVote) TallyVotes(
@@ -87,8 +95,10 @@ func (s *OneAddressOneVote) InitStrategy(
 	f *shared.FlowAdapter,
 	db *shared.Database,
 	sc *s.SnapshotClient,
+	name string,
 ) {
 	s.FlowAdapter = f
 	s.DB = db
 	s.SC = *sc
+	s.name = name
 }

backend/main/strategies/staked_token_weighted_default.go

@@ -14,6 +14,7 @@ type StakedTokenWeightedDefault struct {
 	s.StrategyStruct
 	SC s.SnapshotClient
 	DB *s.Database
+	name string
 }
 
 func (s *StakedTokenWeightedDefault) FetchBalance(
@@ -117,8 +118,10 @@ func (s *StakedTokenWeightedDefault) InitStrategy(
 	f *shared.FlowAdapter,
 	db *shared.Database,
 	sc *s.SnapshotClient,
+	name string,
 ) {
 	s.FlowAdapter = f
 	s.DB = db
 	s.SC = *sc
+	s.name = name
 }

backend/main/strategies/token_weighted_default.go

@@ -16,6 +16,7 @@ type TokenWeightedDefault struct {
 	s.StrategyStruct
 	SC s.SnapshotClient
 	DB *s.Database
+	name string
 }
 
 type FTBalanceResponse struct {
@@ -138,8 +139,10 @@ func (s *TokenWeightedDefault) InitStrategy(
 	f *shared.FlowAdapter,
 	db *shared.Database,
 	sc *s.SnapshotClient,
+	name string,
 ) {
 	s.FlowAdapter = f
 	s.DB = db
 	s.SC = *sc
+	s.name = name
 }

backend/main/strategies/utils.go

@@ -0,0 +1,98 @@
+package strategies
+
+import (
+	"github.com/DapperCollectives/CAST/backend/main/models"
+	s "github.com/DapperCollectives/CAST/backend/main/shared"
+	"github.com/jackc/pgx/v4"
+	"github.com/rs/zerolog/log"
+)
+
+func FetchNFTBalance(
+	db *s.Database,
+	fa *s.FlowAdapter,
+	balance *models.Balance,
+	p *models.Proposal,
+) (*models.Balance, error) {
+
+	vb := &models.VoteWithBalance{
+		NFTs: []*models.NFT{},
+	}
+
+	var c models.Community
+	if err := c.GetCommunityByProposalId(db, balance.Proposal_id); err != nil {
+		return nil, err
+	}
+
+	strategy, err := models.MatchStrategyByProposal(*c.Strategies, *p.Strategy)
+	if err != nil {
+		log.Error().Err(err).Msg("Unable to find strategy for contract")
+		return nil, err
+	}
+
+	nftIds, err := fa.GetNFTIds(balance.Addr, &strategy.Contract)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, nftId := range nftIds {
+		nft := &models.NFT{
+			ID: nftId,
+		}
+		vb.NFTs = append(vb.NFTs, nft)
+	}
+
+	doesExist, err := models.DoesNFTExist(db, vb)
+	if err != nil {
+		return nil, err
+	}
+
+	//only if the NFT ID is not already in the DB,
+	//do we add the balance
+	if !doesExist && err == nil {
+		err = models.CreateUserNFTRecord(db, vb)
+		balance.NFTCount = len(vb.NFTs)
+	}
+
+	return balance, nil
+}
+
+func FetchFTBalance(
+	db *s.Database,
+	sc s.SnapshotClient,
+	b *models.Balance,
+	p *models.Proposal,
+) (*models.Balance, error) {
+
+	var c models.Community
+	if err := c.GetCommunityByProposalId(db, b.Proposal_id); err != nil {
+		return nil, err
+	}
+
+	if err := b.GetBalanceByAddressAndBlockHeight(db); err != nil && err.Error() != pgx.ErrNoRows.Error() {
+		log.Error().Err(err).Msg("error fetching balance from DB")
+		return nil, err
+	}
+
+	strategy, err := models.MatchStrategyByProposal(*c.Strategies, *p.Strategy)
+	if err != nil {
+		log.Error().Err(err).Msg("Unable to find strategy for contract")
+		return nil, err
+	}
+
+	if err := sc.GetAddressBalanceAtBlockHeight(
+		b.Addr,
+		b.BlockHeight,
+		b,
+		&strategy.Contract,
+	); err != nil {
+		log.Error().Err(err).Msg("error fetching balance")
+		return nil, err
+	}
+
+	if err := b.CreateBalance(db); err != nil {
+		log.Error().Err(err).Msg("error creating balance in the DB")
+		return nil, err
+	}
+
+	return b, nil
+}