Fix s256 public key encoding (#112)

Author: anhthii

pkg/encoding/ecdsa.go

@@ -2,9 +2,53 @@ package encoding
 
 import (
 	"crypto/ecdsa"
+	"errors"
+	"math/big"
+
+	"github.com/btcsuite/btcd/btcec/v2"
 )
 
 func EncodeS256PubKey(pubKey *ecdsa.PublicKey) ([]byte, error) {
-	publicKeyBytes := append(pubKey.X.Bytes(), pubKey.Y.Bytes()...)
-	return publicKeyBytes, nil
+	if pubKey == nil {
+		return nil, errors.New("public key is nil")
+	}
+
+	params := pubKey.Curve.Params()
+	expected := btcec.S256().Params()
+	if params.P.Cmp(expected.P) != 0 || params.N.Cmp(expected.N) != 0 {
+		return nil, errors.New("unsupported curve, expected secp256k1")
+	}
+
+	const coordSize = 32
+	xBytes := pubKey.X.Bytes()
+	yBytes := pubKey.Y.Bytes()
+
+	if len(xBytes) > coordSize || len(yBytes) > coordSize {
+		return nil, errors.New("coordinate length exceeds 32 bytes")
+	}
+
+	encoded := make([]byte, coordSize*2)
+	copy(encoded[coordSize-len(xBytes):coordSize], xBytes)
+	copy(encoded[coordSize*2-len(yBytes):], yBytes)
+
+	return encoded, nil
+}
+
+func DecodeECDSAPubKey(encodedKey []byte) (*ecdsa.PublicKey, error) {
+	if len(encodedKey) == 65 && encodedKey[0] == 0x04 {
+		encodedKey = encodedKey[1:] // Strip uncompressed prefix
+	}
+	if len(encodedKey) != 64 {
+		return nil, errors.New("invalid encoded key length, expected 64 bytes")
+	}
+
+	x := new(big.Int).SetBytes(encodedKey[:32])
+	y := new(big.Int).SetBytes(encodedKey[32:])
+
+	curve := btcec.S256()
+	if !curve.IsOnCurve(x, y) {
+		return nil, errors.New("invalid public key: point not on secp256k1 curve")
+	}
+
+	return &ecdsa.PublicKey{Curve: curve, X: x, Y: y}, nil
 }

pkg/encoding/encoding_test.go

@@ -7,52 +7,42 @@ import (
 	"math/big"
 	"testing"
 
+	"github.com/btcsuite/btcd/btcec/v2"
 	"github.com/decred/dcrd/dcrec/edwards/v2"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
 func TestEncodeS256PubKey(t *testing.T) {
-	// Generate a test ECDSA key pair
-	privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
-	require.NoError(t, err)
-
-	pubKey := &privateKey.PublicKey
+	pubKey := &ecdsa.PublicKey{
+		Curve: btcec.S256(),
+		X:     big.NewInt(1),
+		Y:     big.NewInt(2),
+	}
 
-	// Test encoding
 	encoded, err := EncodeS256PubKey(pubKey)
 	require.NoError(t, err)
-	assert.NotEmpty(t, encoded)
-
-	// The encoded key should contain both X and Y coordinates appended together
-	xBytes := pubKey.X.Bytes()
-	yBytes := pubKey.Y.Bytes()
-	expectedLength := len(xBytes) + len(yBytes)
-	assert.Equal(t, expectedLength, len(encoded))
+	require.Len(t, encoded, 64)
 
-	// Verify the encoded data contains the coordinates
-	assert.Equal(t, xBytes, encoded[:len(xBytes)])
-	assert.Equal(t, yBytes, encoded[len(xBytes):])
+	expectedX := pubKey.X.FillBytes(make([]byte, 32))
+	expectedY := pubKey.Y.FillBytes(make([]byte, 32))
+	expected := append(expectedX, expectedY...)
+	assert.Equal(t, expected, encoded)
 }
 
-func TestEncodeS256PubKey_SpecificValues(t *testing.T) {
-	// Create a public key with specific values
-	x := big.NewInt(12345)
-	y := big.NewInt(67890)
-	pubKey := &ecdsa.PublicKey{
-		Curve: elliptic.P256(),
-		X:     x,
-		Y:     y,
-	}
+func TestEncodeS256PubKey_WithValidKey(t *testing.T) {
+	privateKey, err := ecdsa.GenerateKey(btcec.S256(), rand.Reader)
+	require.NoError(t, err)
+
+	pubKey := &privateKey.PublicKey
 
 	encoded, err := EncodeS256PubKey(pubKey)
 	require.NoError(t, err)
+	require.Len(t, encoded, 64)
 
-	// Verify the encoding - should be X bytes followed by Y bytes
-	xBytes := x.Bytes()
-	yBytes := y.Bytes()
-	expected := append(xBytes, yBytes...)
-
+	expectedX := pubKey.X.FillBytes(make([]byte, 32))
+	expectedY := pubKey.Y.FillBytes(make([]byte, 32))
+	expected := append(expectedX, expectedY...)
 	assert.Equal(t, expected, encoded)
 }
 
@@ -135,34 +125,25 @@ func TestEncodeDecodeEDDSA_RoundTrip(t *testing.T) {
 }
 
 func TestEncodeS256PubKey_NilPublicKey(t *testing.T) {
-	// Test with nil public key - this should panic or return an error
-	// depending on the implementation
-	defer func() {
-		if r := recover(); r != nil {
-			// Expected panic due to nil pointer
-			t.Log("Correctly panicked on nil public key")
-		}
-	}()
-
 	_, err := EncodeS256PubKey(nil)
-	if err == nil {
-		t.Error("Expected error or panic with nil public key")
-	}
+	require.EqualError(t, err, "public key is nil")
 }
 
-func TestEncodeS256PubKey_ZeroCoordinates(t *testing.T) {
-	// Test with zero coordinates
-	x := big.NewInt(0)
-	y := big.NewInt(0)
+func TestEncodeS256PubKey_UnsupportedCurve(t *testing.T) {
+	privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
+	require.NoError(t, err)
+
+	_, err = EncodeS256PubKey(&privateKey.PublicKey)
+	require.ErrorContains(t, err, "unsupported curve")
+}
+
+func TestEncodeS256PubKey_CoordinateTooLarge(t *testing.T) {
 	pubKey := &ecdsa.PublicKey{
-		Curve: elliptic.P256(),
-		X:     x,
-		Y:     y,
+		Curve: btcec.S256(),
+		X:     new(big.Int).Lsh(big.NewInt(1), 256),
+		Y:     big.NewInt(0),
 	}
 
-	encoded, err := EncodeS256PubKey(pubKey)
-	require.NoError(t, err)
-
-	// Should still work, though the result will be a very short byte array
-	assert.NotNil(t, encoded)
+	_, err := EncodeS256PubKey(pubKey)
+	require.EqualError(t, err, "coordinate length exceeds 32 bytes")
 }