From d5b4553bb36176dfbbeaf69d11a0ba78240a07ce Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Wed, 2 Dec 2020 15:52:31 +0300
Subject: [PATCH] keys: allow to create keys on arbitrary curve

---
 pkg/crypto/keys/private_key.go      | 17 ++++++++++++++---
 pkg/crypto/keys/private_key_test.go | 17 +++++++++++++++++
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/pkg/crypto/keys/private_key.go b/pkg/crypto/keys/private_key.go
index c790980f2..111f3283e 100644
--- a/pkg/crypto/keys/private_key.go
+++ b/pkg/crypto/keys/private_key.go
@@ -10,6 +10,7 @@ import (
 	"fmt"
 	"math/big"
 
+	"github.com/btcsuite/btcd/btcec"
 	"github.com/nspcc-dev/neo-go/pkg/util"
 	"github.com/nspcc-dev/rfc6979"
 )
@@ -20,16 +21,26 @@ type PrivateKey struct {
 	ecdsa.PrivateKey
 }
 
-// NewPrivateKey creates a new random Secp256k1 private key.
+// NewPrivateKey creates a new random Secp256r1 private key.
 func NewPrivateKey() (*PrivateKey, error) {
-	priv, x, y, err := elliptic.GenerateKey(elliptic.P256(), rand.Reader)
+	return newPrivateKeyOnCurve(elliptic.P256())
+}
+
+// NewSecp256k1PrivateKey creates a new random Secp256k1 private key.
+func NewSecp256k1PrivateKey() (*PrivateKey, error) {
+	return newPrivateKeyOnCurve(btcec.S256())
+}
+
+// newPrivateKeyOnCurve creates a new random private key using curve c.
+func newPrivateKeyOnCurve(c elliptic.Curve) (*PrivateKey, error) {
+	priv, x, y, err := elliptic.GenerateKey(c, rand.Reader)
 	if err != nil {
 		return nil, err
 	}
 	return &PrivateKey{
 		ecdsa.PrivateKey{
 			PublicKey: ecdsa.PublicKey{
-				Curve: elliptic.P256(),
+				Curve: c,
 				X:     x,
 				Y:     y,
 			},
diff --git a/pkg/crypto/keys/private_key_test.go b/pkg/crypto/keys/private_key_test.go
index 9089aa302..b8b806080 100644
--- a/pkg/crypto/keys/private_key_test.go
+++ b/pkg/crypto/keys/private_key_test.go
@@ -6,7 +6,9 @@ import (
 	"testing"
 
 	"github.com/nspcc-dev/neo-go/internal/keytestcases"
+	"github.com/nspcc-dev/neo-go/pkg/crypto/hash"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestPrivateKey(t *testing.T) {
@@ -28,6 +30,21 @@ func TestPrivateKey(t *testing.T) {
 	}
 }
 
+func TestNewPrivateKeyOnCurve(t *testing.T) {
+	msg := []byte{1, 2, 3}
+	h := hash.Sha256(msg).BytesBE()
+	t.Run("Secp256r1", func(t *testing.T) {
+		p, err := NewPrivateKey()
+		require.NoError(t, err)
+		p.PublicKey().Verify(p.Sign(msg), h)
+	})
+	t.Run("Secp256k1", func(t *testing.T) {
+		p, err := NewSecp256k1PrivateKey()
+		require.NoError(t, err)
+		p.PublicKey().Verify(p.Sign(msg), h)
+	})
+}
+
 func TestPrivateKeyFromWIF(t *testing.T) {
 	for _, testCase := range keytestcases.Arr {
 		key, err := NewPrivateKeyFromWIF(testCase.Wif)