crypto: Add SigToPub. (#48)

* Add SigToPub function in crypto

2 files changed, 35 insertions, 5 deletions

diff --git a/crypto/eth/eth.go b/crypto/eth/eth.go
index cb20d7d..edf50d3 100644
--- a/crypto/eth/eth.go
+++ b/crypto/eth/eth.go
@@ -79,17 +79,23 @@ func (prv *PrivateKey) PublicKey() crypto.PublicKey {
 // solution is to hash any input before calculating the signature.
 //
 // The produced signature is in the [R || S || V] format where V is 0 or 1.
-func (prv *PrivateKey) Sign(hash common.Hash) (sig crypto.Signature, err error) {
+func (prv *PrivateKey) Sign(hash common.Hash) (
+	sig crypto.Signature, err error) {
 	s, err := ethcrypto.Sign(hash[:], &prv.privateKey)
-	// Magic!
-	sig = crypto.Signature(s[:64])
+	sig = crypto.Signature(s)
 	return
 }
 
 // VerifySignature checks that the given public key created signature over hash.
-// The public key should be in compressed (33 bytes) or uncompressed (65 bytes) format.
+// The public key should be in compressed (33 bytes) or uncompressed (65 bytes)
+// format.
 // The signature should have the 64 byte [R || S] format.
-func (pub PublicKey) VerifySignature(hash common.Hash, signature crypto.Signature) bool {
+func (pub PublicKey) VerifySignature(
+	hash common.Hash, signature crypto.Signature) bool {
+	if len(signature) == 65 {
+		// The last byte is for ecrecover.
+		signature = signature[:64]
+	}
 	return ethcrypto.VerifySignature(pub.publicKey, hash[:], signature)
 }
 
@@ -97,3 +103,13 @@ func (pub PublicKey) VerifySignature(hash common.Hash, signature crypto.Signatur
 func (pub PublicKey) Compress() []byte {
 	return pub.publicKey
 }
+
+// SigToPub returns the PublicKey that created the given signature.
+func SigToPub(
+	hash common.Hash, signature crypto.Signature) (PublicKey, error) {
+	key, err := ethcrypto.SigToPub(hash[:], signature[:])
+	if err != nil {
+		return PublicKey{}, err
+	}
+	return PublicKey{publicKey: ethcrypto.CompressPubkey(key)}, nil
+}
diff --git a/crypto/eth/eth_test.go b/crypto/eth/eth_test.go
index b09297e..f6f48cb 100644
--- a/crypto/eth/eth_test.go
+++ b/crypto/eth/eth_test.go
@@ -21,6 +21,7 @@ import (
 	"testing"
 
 	"github.com/dexon-foundation/dexon-consensus-core/common"
+	"github.com/dexon-foundation/dexon-consensus-core/crypto"
 	"github.com/stretchr/testify/suite"
 )
 
@@ -74,6 +75,19 @@ func (s *ETHCryptoTestSuite) TestSignature() {
 	s.True(decompressPub1.VerifySignature(hash1, sig11))
 }
 
+func (s *ETHCryptoTestSuite) TestSigToPub() {
+	prv, err := NewPrivateKey()
+	s.Require().Nil(err)
+	data := "DEXON is infinitely scalable and low-latency."
+	hash := crypto.Keccak256Hash([]byte(data))
+	sigmsg, err := prv.Sign(hash)
+	s.Require().Nil(err)
+
+	pubkey, err := SigToPub(hash, sigmsg)
+	s.Require().Nil(err)
+	s.Equal(pubkey, prv.PublicKey())
+}
+
 func TestCrypto(t *testing.T) {
 	suite.Run(t, new(ETHCryptoTestSuite))
 }