From 7e2e5e18793188a415edca03692c715c9e5604aa Mon Sep 17 00:00:00 2001
From: Roman Khimov <roman@nspcc.ru>
Date: Wed, 10 Jun 2020 19:14:13 +0300
Subject: [PATCH] keys: add support for uncompressed serialization in PublicKey

---
 pkg/crypto/keys/publickey.go      | 52 ++++++++++++++++++++++++-------
 pkg/crypto/keys/publickey_test.go | 12 ++++---
 2 files changed, 48 insertions(+), 16 deletions(-)

diff --git a/pkg/crypto/keys/publickey.go b/pkg/crypto/keys/publickey.go
index d03bdb22b..6c9853649 100644
--- a/pkg/crypto/keys/publickey.go
+++ b/pkg/crypto/keys/publickey.go
@@ -1,7 +1,6 @@
 package keys
 
 import (
-	"bytes"
 	"crypto/ecdsa"
 	"crypto/elliptic"
 	"crypto/x509"
@@ -19,6 +18,9 @@ import (
 	"github.com/pkg/errors"
 )
 
+// coordLen is the number of bytes in serialized X or Y coordinate.
+const coordLen = 32
+
 // PublicKeys is a list of public keys.
 type PublicKeys []*PublicKey
 
@@ -106,23 +108,49 @@ func NewPublicKeyFromBytes(b []byte) (*PublicKey, error) {
 	return pubKey, nil
 }
 
-// Bytes returns the byte array representation of the public key.
-func (p *PublicKey) Bytes() []byte {
+// getBytes serializes X and Y using compressed or uncompressed format.
+func (p *PublicKey) getBytes(compressed bool) []byte {
 	if p.IsInfinity() {
 		return []byte{0x00}
 	}
 
-	var (
-		x       = p.X.Bytes()
-		paddedX = append(bytes.Repeat([]byte{0x00}, 32-len(x)), x...)
-		prefix  = byte(0x03)
-	)
-
-	if p.Y.Bit(0) == 0 {
-		prefix = byte(0x02)
+	var resLen = 1 + coordLen
+	if !compressed {
+		resLen += coordLen
 	}
+	var res = make([]byte, resLen)
+	var prefix byte
 
-	return append([]byte{prefix}, paddedX...)
+	xBytes := p.X.Bytes()
+	copy(res[1+coordLen-len(xBytes):], xBytes)
+	if compressed {
+		if p.Y.Bit(0) == 0 {
+			prefix = 0x02
+		} else {
+			prefix = 0x03
+		}
+	} else {
+		prefix = 0x04
+		yBytes := p.Y.Bytes()
+		copy(res[1+coordLen+coordLen-len(yBytes):], yBytes)
+
+	}
+	res[0] = prefix
+
+	return res
+}
+
+// Bytes returns byte array representation of the public key in compressed
+// form (33 bytes with 0x02 or 0x03 prefix, except infinity which is always 0).
+func (p *PublicKey) Bytes() []byte {
+	return p.getBytes(true)
+}
+
+// UncompressedBytes returns byte array representation of the public key in
+// uncompressed form (65 bytes with 0x04 prefix, except infinity which is
+// always 0).
+func (p *PublicKey) UncompressedBytes() []byte {
+	return p.getBytes(false)
 }
 
 // NewPublicKeyFromASN1 returns a NEO PublicKey from the ASN.1 serialized key.
diff --git a/pkg/crypto/keys/publickey_test.go b/pkg/crypto/keys/publickey_test.go
index 3ec900784..1c15947be 100644
--- a/pkg/crypto/keys/publickey_test.go
+++ b/pkg/crypto/keys/publickey_test.go
@@ -95,10 +95,14 @@ func TestPubkeyToAddress(t *testing.T) {
 
 func TestDecodeBytes(t *testing.T) {
 	pubKey := getPubKey(t)
-	decodedPubKey := &PublicKey{}
-	err := decodedPubKey.DecodeBytes(pubKey.Bytes())
-	require.NoError(t, err)
-	require.Equal(t, pubKey, decodedPubKey)
+	var testBytesFunction = func(t *testing.T, bytesFunction func() []byte) {
+		decodedPubKey := &PublicKey{}
+		err := decodedPubKey.DecodeBytes(bytesFunction())
+		require.NoError(t, err)
+		require.Equal(t, pubKey, decodedPubKey)
+	}
+	t.Run("compressed", func(t *testing.T) { testBytesFunction(t, pubKey.Bytes) })
+	t.Run("uncompressed", func(t *testing.T) { testBytesFunction(t, pubKey.UncompressedBytes) })
 }
 
 func TestSort(t *testing.T) {