Set tspend & treasury policies on VSP

Update the existing RPCs to update treasury/tspend voting preferences
so that they forward updated preferences to VSP.

Author: jholdstock

internal/rpc/jsonrpc/methods.go

@@ -3811,6 +3811,9 @@ func (s *Server) setDisapprovePercent(ctx context.Context, icmd interface{}) (in
 
 // setTreasuryPolicy saves the voting policy for treasury spends by a particular
 // key, and optionally, setting the key policy used by a specific ticket.
+//
+// If a VSP host is configured in the application settings, the voting
+// preferences will also be set with the VSP.
 func (s *Server) setTreasuryPolicy(ctx context.Context, icmd interface{}) (interface{}, error) {
 	cmd := icmd.(*types.SetTreasuryPolicyCmd)
 	w, ok := s.walletLoader.LoadedWallet()
@@ -3854,6 +3857,16 @@ func (s *Server) setTreasuryPolicy(ctx context.Context, icmd interface{}) (inter
 	}
 
 	err = w.SetTreasuryKeyPolicy(ctx, pikey, policy, ticketHash)
+	if err != nil {
+		return nil, err
+	}
+
+	// Update voting preferences on VSPs if required.
+	policyMap := map[string]string{
+		cmd.Key: cmd.Policy,
+	}
+	err = s.updateVSPVoteChoices(ctx, w, ticketHash, nil, nil, policyMap)
+
 	return nil, err
 }
 
@@ -3929,6 +3942,9 @@ func (s *Server) tspendPolicy(ctx context.Context, icmd interface{}) (interface{
 
 // setTSpendPolicy saves the voting policy for a particular tspend transaction
 // hash, and optionally, setting the tspend policy used by a specific ticket.
+//
+// If a VSP host is configured in the application settings, the voting
+// preferences will also be set with the VSP.
 func (s *Server) setTSpendPolicy(ctx context.Context, icmd interface{}) (interface{}, error) {
 	cmd := icmd.(*types.SetTSpendPolicyCmd)
 	w, ok := s.walletLoader.LoadedWallet()
@@ -3975,6 +3991,15 @@ func (s *Server) setTSpendPolicy(ctx context.Context, icmd interface{}) (interfa
 	}
 
 	err = w.SetTSpendPolicy(ctx, hash, policy, ticketHash)
+	if err != nil {
+		return nil, err
+	}
+
+	// Update voting preferences on VSPs if required.
+	policyMap := map[string]string{
+		cmd.Hash: cmd.Policy,
+	}
+	err = s.updateVSPVoteChoices(ctx, w, ticketHash, nil, policyMap, nil)
 	return nil, err
 }
 
@@ -4635,34 +4660,42 @@ func (s *Server) setVoteChoice(ctx context.Context, icmd interface{}) (interface
 		ticketHash = hash
 	}
 
-	choice := wallet.AgendaChoice{
-		AgendaID: cmd.AgendaID,
-		ChoiceID: cmd.ChoiceID,
+	choice := []wallet.AgendaChoice{
+		{
+			AgendaID: cmd.AgendaID,
+			ChoiceID: cmd.ChoiceID,
+		},
 	}
-	_, err := w.SetAgendaChoices(ctx, ticketHash, choice)
+	_, err := w.SetAgendaChoices(ctx, ticketHash, choice...)
 	if err != nil {
 		return nil, err
 	}
 
 	// Update voting preferences on VSPs if required.
+	err = s.updateVSPVoteChoices(ctx, w, ticketHash, choice, nil, nil)
+	return nil, err
+}
+
+func (s *Server) updateVSPVoteChoices(ctx context.Context, w *wallet.Wallet, ticketHash *chainhash.Hash,
+	choices []wallet.AgendaChoice, tspendPolicy map[string]string, treasuryPolicy map[string]string) error {
 	if ticketHash != nil {
 		vspHost, err := w.VSPHostForTicket(ctx, ticketHash)
 		if err != nil {
 			if errors.Is(err, errors.NotExist) {
 				// Ticket is not registered with a VSP, nothing more to do here.
-				return nil, nil
+				return nil
 			}
-			return nil, err
+			return err
 		}
 		vspClient, err := loader.LookupVSP(vspHost)
 		if err != nil {
-			return nil, err
+			return err
 		}
-		err = vspClient.SetVoteChoice(ctx, ticketHash, choice)
-		return nil, err
+		err = vspClient.SetVoteChoice(ctx, ticketHash, choices, tspendPolicy, treasuryPolicy)
+		return err
 	}
 	var firstErr error
-	err = w.ForUnspentUnexpiredTickets(ctx, func(hash *chainhash.Hash) error {
+	err := w.ForUnspentUnexpiredTickets(ctx, func(hash *chainhash.Hash) error {
 		vspHost, err := w.VSPHostForTicket(ctx, hash)
 		if err != nil && firstErr == nil {
 			if errors.Is(err, errors.NotExist) {
@@ -4679,16 +4712,16 @@ func (s *Server) setVoteChoice(ctx context.Context, icmd interface{}) (interface
 		}
 		// Never return errors here, so all tickets are tried.
 		// The first error will be returned to the user.
-		err = vspClient.SetVoteChoice(ctx, hash, choice)
+		err = vspClient.SetVoteChoice(ctx, hash, choices, tspendPolicy, treasuryPolicy)
 		if err != nil && firstErr == nil {
 			firstErr = err
 		}
 		return nil
 	})
 	if err != nil {
-		return nil, err
+		return err
 	}
-	return nil, firstErr
+	return firstErr
 }
 
 // signMessage signs the given message with the private key for the given

internal/rpc/rpcserver/server.go

@@ -4069,7 +4069,10 @@ func (s *walletServer) SetVspdVoteChoices(ctx context.Context, req *pb.SetVspdVo
 			return err
 		}
 		if ticketHost == vspHost {
-			_ = vspClient.SetVoteChoice(ctx, hash, choices...)
+			tSpendChoices := s.wallet.TSpendPolicyForTicket(hash)
+			treasuryChoices := s.wallet.TreasuryKeyPolicyForTicket(hash)
+
+			_ = vspClient.SetVoteChoice(ctx, hash, choices, tSpendChoices, treasuryChoices)
 		}
 		return nil
 	})

internal/vsp/feepayment.go

@@ -611,6 +611,8 @@ type ticketStatus struct {
 	FeeTxStatus     string            `json:"feetxstatus"`
 	FeeTxHash       string            `json:"feetxhash"`
 	VoteChoices     map[string]string `json:"votechoices"`
+	TSpendPolicy    map[string]string `json:"tspendpolicy"`
+	TreasuryPolicy  map[string]string `json:"treasurypolicy"`
 	Request         []byte            `json:"request"`
 }
 
@@ -662,7 +664,8 @@ func (c *Client) status(ctx context.Context, ticketHash *chainhash.Hash) (*ticke
 	return &resp, nil
 }
 
-func (c *Client) setVoteChoices(ctx context.Context, ticketHash *chainhash.Hash, choices []wallet.AgendaChoice) error {
+func (c *Client) setVoteChoices(ctx context.Context, ticketHash *chainhash.Hash,
+	choices []wallet.AgendaChoice, tspendPolicy map[string]string, treasuryPolicy map[string]string) error {
 	w := c.Wallet
 	params := w.ChainParams()
 
@@ -693,13 +696,17 @@ func (c *Client) setVoteChoices(ctx context.Context, ticketHash *chainhash.Hash,
 
 	var resp ticketStatus
 	requestBody, err := json.Marshal(&struct {
-		Timestamp   int64             `json:"timestamp"`
-		TicketHash  string            `json:"tickethash"`
-		VoteChoices map[string]string `json:"votechoices"`
+		Timestamp      int64             `json:"timestamp"`
+		TicketHash     string            `json:"tickethash"`
+		VoteChoices    map[string]string `json:"votechoices"`
+		TSpendPolicy   map[string]string `json:"tspendpolicy"`
+		TreasuryPolicy map[string]string `json:"treasurypolicy"`
 	}{
-		Timestamp:   time.Now().Unix(),
-		TicketHash:  ticketHash.String(),
-		VoteChoices: agendaChoices,
+		Timestamp:      time.Now().Unix(),
+		TicketHash:     ticketHash.String(),
+		VoteChoices:    agendaChoices,
+		TSpendPolicy:   tspendPolicy,
+		TreasuryPolicy: treasuryPolicy,
 	})
 	if err != nil {
 		return err
@@ -882,17 +889,21 @@ func (fp *feePayment) submitPayment() (err error) {
 		Request   []byte `json:"request"`
 	}
 	requestBody, err := json.Marshal(&struct {
-		Timestamp   int64             `json:"timestamp"`
-		TicketHash  string            `json:"tickethash"`
-		FeeTx       json.Marshaler    `json:"feetx"`
-		VotingKey   string            `json:"votingkey"`
-		VoteChoices map[string]string `json:"votechoices"`
+		Timestamp      int64             `json:"timestamp"`
+		TicketHash     string            `json:"tickethash"`
+		FeeTx          json.Marshaler    `json:"feetx"`
+		VotingKey      string            `json:"votingkey"`
+		VoteChoices    map[string]string `json:"votechoices"`
+		TSpendPolicy   map[string]string `json:"tspendpolicy"`
+		TreasuryPolicy map[string]string `json:"treasurypolicy"`
 	}{
-		Timestamp:   time.Now().Unix(),
-		TicketHash:  fp.ticketHash.String(),
-		FeeTx:       txMarshaler(feeTx),
-		VotingKey:   votingKey,
-		VoteChoices: voteChoices,
+		Timestamp:      time.Now().Unix(),
+		TicketHash:     fp.ticketHash.String(),
+		FeeTx:          txMarshaler(feeTx),
+		VotingKey:      votingKey,
+		VoteChoices:    voteChoices,
+		TSpendPolicy:   w.TSpendPolicyForTicket(&fp.ticketHash),
+		TreasuryPolicy: w.TreasuryKeyPolicyForTicket(&fp.ticketHash),
 	})
 	if err != nil {
 		return err

internal/vsp/vsp.go

@@ -292,10 +292,12 @@ func (c *Client) ProcessWithPolicy(ctx context.Context, ticketHash *chainhash.Ha
 	return nil
 }
 
-// SetVoteChoice takes the provided AgendaChoices and ticket hash, checks the
-// status of the ticket from the connected vsp.  The status provides the
-// current vote choice so we can just update from there if need be.
-func (c *Client) SetVoteChoice(ctx context.Context, hash *chainhash.Hash, choices ...wallet.AgendaChoice) error {
+// SetVoteChoice takes the provided consensus, tspend and treasury key voting
+// preferences, and checks if they match the status of the specified ticket from
+// the connected VSP. The status provides the current voting preferences so we
+// can just update from there if need be.
+func (c *Client) SetVoteChoice(ctx context.Context, hash *chainhash.Hash,
+	choices []wallet.AgendaChoice, tspendPolicy map[string]string, treasuryPolicy map[string]string) error {
 
 	// Retrieve current voting preferences from VSP.
 	status, err := c.status(ctx, hash)
@@ -311,6 +313,7 @@ func (c *Client) SetVoteChoice(ctx context.Context, hash *chainhash.Hash, choice
 	// VSP preferences to determine if VSP needs to be updated.
 	update := false
 
+	// Check consensus vote choices.
 	for _, newChoice := range choices {
 		vspChoice, ok := status.VoteChoices[newChoice.AgendaID]
 		if !ok {
@@ -323,13 +326,41 @@ func (c *Client) SetVoteChoice(ctx context.Context, hash *chainhash.Hash, choice
 		}
 	}
 
+	// Apply the above changes to the two checks below.
+
+	// Check tspend policies.
+	for newTSpend, newChoice := range tspendPolicy {
+		vspChoice, ok := status.TSpendPolicy[newTSpend]
+		if !ok {
+			update = true
+			break
+		}
+		if vspChoice != newChoice {
+			update = true
+			break
+		}
+	}
+
+	// Check treasury policies.
+	for newKey, newChoice := range treasuryPolicy {
+		vspChoice, ok := status.TSpendPolicy[newKey]
+		if !ok {
+			update = true
+			break
+		}
+		if vspChoice != newChoice {
+			update = true
+			break
+		}
+	}
+
 	if !update {
 		log.Debugf("VSP already has correct vote choices for ticket %s", hash)
 		return nil
 	}
 
 	log.Debugf("Updating vote choices on VSP for ticket %s", hash)
-	err = c.setVoteChoices(ctx, hash, choices)
+	err = c.setVoteChoices(ctx, hash, choices, tspendPolicy, treasuryPolicy)
 	if err != nil {
 		return err
 	}

wallet/wallet.go

@@ -549,6 +549,30 @@ func (w *Wallet) SetAgendaChoices(ctx context.Context, ticketHash *chainhash.Has
 	return voteBits, nil
 }
 
+// TreasuryKeyPolicyForTicket returns all of the treasury key policies set for a
+// single ticket. It does not consider the global wallet setting.
+func (w *Wallet) TreasuryKeyPolicyForTicket(ticketHash *chainhash.Hash) map[string]string {
+	w.stakeSettingsLock.Lock()
+	defer w.stakeSettingsLock.Unlock()
+
+	policies := make(map[string]string)
+	for key, value := range w.vspTSpendKeyPolicy {
+		if key.Ticket.IsEqual(ticketHash) {
+			var choice string
+			switch value {
+			case stake.TreasuryVoteYes:
+				choice = "yes"
+			case stake.TreasuryVoteNo:
+				choice = "no"
+			default:
+				choice = "abstain"
+			}
+			policies[key.TreasuryKey] = choice
+		}
+	}
+	return policies
+}
+
 // TreasuryKeyPolicy returns a vote policy for provided Pi key. If there is
 // no policy this method returns TreasuryVoteInvalid.
 // A non-nil ticket hash may be used by a VSP to return per-ticket policies.
@@ -566,6 +590,30 @@ func (w *Wallet) TreasuryKeyPolicy(pikey []byte, ticket *chainhash.Hash) stake.T
 	return w.tspendKeyPolicy[string(pikey)]
 }
 
+// TSpendPolicyForTicket returns all of the tspend policies set for a single
+// ticket. It does not consider the global wallet setting.
+func (w *Wallet) TSpendPolicyForTicket(ticketHash *chainhash.Hash) map[string]string {
+	w.stakeSettingsLock.Lock()
+	defer w.stakeSettingsLock.Unlock()
+
+	policies := make(map[string]string)
+	for key, value := range w.vspTSpendPolicy {
+		if key.Ticket.IsEqual(ticketHash) {
+			var choice string
+			switch value {
+			case stake.TreasuryVoteYes:
+				choice = "yes"
+			case stake.TreasuryVoteNo:
+				choice = "no"
+			default:
+				choice = "abstain"
+			}
+			policies[key.TSpend.String()] = choice
+		}
+	}
+	return policies
+}
+
 // TSpendPolicy returns a vote policy for a tspend.  If a policy is set for a
 // particular tspend transaction, that policy is returned.  Otherwise, if the
 // tspend is known, any policy for the treasury key which signs the tspend is