From 8b104e456beb22f2993ab27bcd2ccabadba95183 Mon Sep 17 00:00:00 2001
From: Thane Thomson <connect@thanethomson.com>
Date: Wed, 13 Apr 2022 18:49:47 -0400
Subject: [PATCH] privval: update signing logic to cater for vote extensions

Signed-off-by: Thane Thomson <connect@thanethomson.com>
---
 privval/file.go | 40 ++++++++++++++++++++++++++--------------
 1 file changed, 26 insertions(+), 14 deletions(-)

diff --git a/privval/file.go b/privval/file.go
index a0835cfeb..1e31b8c78 100644
--- a/privval/file.go
+++ b/privval/file.go
@@ -117,6 +117,14 @@ type FilePVLastSignState struct {
 	filePath string
 }
 
+func (lss *FilePVLastSignState) reset() {
+	lss.Height = 0
+	lss.Round = 0
+	lss.Step = 0
+	lss.Signature = nil
+	lss.SignBytes = nil
+}
+
 // checkHRS checks the given height, round, step (HRS) against that of the
 // FilePVLastSignState. It returns an error if the arguments constitute a regression,
 // or if they match but the SignBytes are empty.
@@ -328,12 +336,7 @@ func (pv *FilePV) Save() error {
 // Reset resets all fields in the FilePV.
 // NOTE: Unsafe!
 func (pv *FilePV) Reset() error {
-	var sig []byte
-	pv.LastSignState.Height = 0
-	pv.LastSignState.Round = 0
-	pv.LastSignState.Step = 0
-	pv.LastSignState.Signature = sig
-	pv.LastSignState.SignBytes = nil
+	pv.LastSignState.reset()
 	return pv.Save()
 }
 
@@ -370,6 +373,11 @@ func (pv *FilePV) signVote(chainID string, vote *tmproto.Vote) error {
 
 	signBytes := types.VoteSignBytes(chainID, vote)
 	extSignBytes := types.VoteExtensionSignBytes(chainID, vote)
+	// We always sign the vote extension. See below for details.
+	extSig, err := pv.Key.PrivKey.Sign(extSignBytes)
+	if err != nil {
+		return err
+	}
 
 	// We might crash before writing to the wal,
 	// causing us to try to re-sign for the same HRS.
@@ -380,6 +388,8 @@ func (pv *FilePV) signVote(chainID string, vote *tmproto.Vote) error {
 		if bytes.Equal(signBytes, lss.SignBytes) {
 			vote.Signature = lss.Signature
 		} else {
+			// Compares the canonicalized votes (i.e. without vote extensions
+			// or vote extension signatures).
 			timestamp, ok, err := checkVotesOnlyDifferByTimestamp(lss.SignBytes, signBytes)
 			if err != nil {
 				return err
@@ -391,6 +401,12 @@ func (pv *FilePV) signVote(chainID string, vote *tmproto.Vote) error {
 			vote.Timestamp = timestamp
 			vote.Signature = lss.Signature
 		}
+
+		// Vote extensions are non-deterministic, so it's possible that an
+		// application may have created a different extension. We therefore
+		// always re-sign the vote extension.
+		vote.ExtensionSignature = extSig
+
 		return nil
 	}
 
@@ -403,12 +419,6 @@ func (pv *FilePV) signVote(chainID string, vote *tmproto.Vote) error {
 		return err
 	}
 	vote.Signature = sig
-
-	// Sign the vote extension, regardless of whether it's present or not
-	extSig, err := pv.Key.PrivKey.Sign(extSignBytes)
-	if err != nil {
-		return err
-	}
 	vote.ExtensionSignature = extSig
 
 	return nil
@@ -462,8 +472,10 @@ func (pv *FilePV) saveSigned(height int64, round int32, step int8, signBytes []b
 
 //-----------------------------------------------------------------------------------------
 
-// returns the timestamp from the lastSignBytes.
-// returns true if the only difference in the votes is their timestamp.
+// Returns the timestamp from the lastSignBytes.
+// Returns true if the only difference in the votes is their timestamp.
+// Performs these checks on the canonical votes (excluding the vote extension
+// and vote extension signatures).
 func checkVotesOnlyDifferByTimestamp(lastSignBytes, newSignBytes []byte) (time.Time, bool, error) {
 	var lastVote, newVote tmproto.CanonicalVote
 	if err := protoio.UnmarshalDelimited(lastSignBytes, &lastVote); err != nil {