Fix batch signature verification

zetaclient/tss/crypto.go

@@ -6,6 +6,7 @@ import (
 	"crypto/elliptic"
 	"encoding/base64"
 	"encoding/hex"
+	"fmt"
 	"strings"
 
 	"github.com/btcsuite/btcd/btcutil"
@@ -168,6 +169,49 @@ func SignatureToBytes(input keysign.Signature) (sig [65]byte, err error) {
 	return sig, nil
 }
 
+// apparently go-tss returns res.Signatures in a different order than digests,
+// thus we need to ensure the order AND verify the signatures
+func verifySignatures(digests [][]byte, res keysign.Response, pk PubKey) ([][65]byte, error) {
+	switch {
+	case len(digests) == 0:
+		return nil, errors.New("empty digests list")
+	case len(digests) != len(res.Signatures):
+		return nil, errors.New("length mismatch")
+	case len(digests) == 1:
+		// most common case
+		sig, err := VerifySignature(res.Signatures[0], pk, digests[0])
+		if err != nil {
+			return nil, err
+		}
+
+		return [][65]byte{sig}, nil
+	}
+
+	// map bas64(digest) => slice index
+	cache := make(map[string]int, len(digests))
+	for i, digest := range digests {
+		cache[base64EncodeString(digest)] = i
+	}
+
+	signatures := make([][65]byte, len(res.Signatures))
+
+	for _, sigResponse := range res.Signatures {
+		i, ok := cache[sigResponse.Msg]
+		if !ok {
+			return nil, errors.Errorf("missing digest %s", sigResponse.Msg)
+		}
+
+		sig, err := VerifySignature(sigResponse, pk, digests[i])
+		if err != nil {
+			return nil, fmt.Errorf("unable to verify signature: %w (#%d)", err, i)
+		}
+
+		signatures[i] = sig
+	}
+
+	return signatures, nil
+}
+
 // combineDigests combines the digests
 func combineDigests(digestList []string) []byte {
 	digestConcat := strings.Join(digestList, "")

zetaclient/tss/service.go

@@ -211,17 +211,14 @@ func (s *Service) SignBatch(
 		return nil, fmt.Errorf("keysign fail: signature list length mismatch")
 	}
 
-	signatures := make([][65]byte, len(res.Signatures))
-	for i, sigResponse := range res.Signatures {
-		signatures[i], err = VerifySignature(sigResponse, s.PubKey(), digests[i])
-		if err != nil {
-			return nil, fmt.Errorf("unable to verify signature: %w (#%d)", err, i)
-		}
+	sigs, err := verifySignatures(digests, res, s.PubKey())
+	if err != nil {
+		return nil, errors.Wrap(err, "unable to verify signatures")
 	}
 
 	// todo sig save to LRU cache (chain-id + digest). We need LRU per EACH chain
 
-	return signatures, nil
+	return sigs, nil
 }
 
 var (

zetaclient/tss/service_test.go

@@ -81,6 +81,37 @@ func TestService(t *testing.T) {
 			require.NotEmpty(t, sig)
 		})
 	})
+
+	t.Run("SignBatch", func(t *testing.T) {
+		// ARRANGE
+		ts := newTestSuite(t)
+
+		// Given tss service
+		s, err := tss.NewService(ts, ts.PubKeyBech32(), ts.zetacore, ts.logger)
+		require.NoError(t, err)
+
+		// Given several sample messages to sign
+		digests := [][]byte{
+			ts.SampleDigest(),
+			ts.SampleDigest(),
+			ts.SampleDigest(),
+			ts.SampleDigest(),
+			ts.SampleDigest(),
+			ts.SampleDigest(),
+			ts.SampleDigest(),
+		}
+
+		// Given mock response
+		const blockHeight = 123
+		ts.keySignerMock.AddCall(ts.PubKeyBech32(), digests, blockHeight, true, nil)
+
+		// ACT
+		sig, err := s.SignBatch(ts.ctx, digests, blockHeight, 2, 3)
+
+		// ASSERT
+		require.NoError(t, err)
+		require.NotEmpty(t, sig)
+	})
 }
 
 type testSuite struct {
@@ -188,6 +219,9 @@ func (m *keySignerMock) sign(req keysign.Request) keysign.Response {
 		})
 	}
 
+	// might be random... we should tolerate that
+	signatures = lo.Shuffle(signatures)
+
 	return keysign.Response{
 		Signatures: signatures,
 		Status:     tsscommon.Success,