Skip to content

Commit

Permalink
fixed tests
Browse files Browse the repository at this point in the history
  • Loading branch information
dan13ram committed Jul 2, 2024
1 parent aef56ce commit e10dc35
Show file tree
Hide file tree
Showing 21 changed files with 943 additions and 708 deletions.
3 changes: 3 additions & 0 deletions Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,9 @@ clean_tmp_data :; if [ -d "/tmp/data" ]; then sudo rm -rf /tmp/data; fi
.PHONY: install
install :; go mod download && go mod verify

.PHONY: format
format :; go fmt ./...

.PHONY: lint
lint :; golangci-lint run

Expand Down
102 changes: 50 additions & 52 deletions common/gcp_kms_signer.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@ import (
"github.com/ethereum/go-ethereum/common"
"github.com/ethereum/go-ethereum/crypto"

gax "github.com/googleapis/gax-go/v2"

"github.com/cosmos/cosmos-sdk/crypto/types"
)

Expand All @@ -25,11 +27,19 @@ type Signer interface {
CosmosSign(data []byte) ([]byte, error)
EthAddress() common.Address
CosmosPublicKey() types.PubKey
Destroy()
}

type GCPKeyManagementClient interface {
Close() error
GetPublicKey(ctx context.Context, req *kmspb.GetPublicKeyRequest, opts ...gax.CallOption) (*kmspb.PublicKey, error)
AsymmetricSign(ctx context.Context, req *kmspb.AsymmetricSignRequest, opts ...gax.CallOption) (*kmspb.AsymmetricSignResponse, error)
GetCryptoKeyVersion(ctx context.Context, req *kmspb.GetCryptoKeyVersionRequest, opts ...gax.CallOption) (*kmspb.CryptoKeyVersion, error)
}

// Struct Definition
type GcpKmsSigner struct {
client *kms.KeyManagementClient
client GCPKeyManagementClient
keyName string
ethAddress common.Address
cosmosPubKey types.PubKey
Expand All @@ -38,15 +48,19 @@ type GcpKmsSigner struct {

var _ Signer = &GcpKmsSigner{}

var NewGCPKeyManagementClient = func(ctx context.Context) (GCPKeyManagementClient, error) {
return kms.NewKeyManagementClient(ctx)
}

// Constructor Function
func NewGcpKmsSigner(keyName string) (Signer, error) {
client, err := kms.NewKeyManagementClient(context.Background())
client, err := NewGCPKeyManagementClient(context.Background())
if err != nil {
return nil, fmt.Errorf("failed to create KMS client: %w", err)
}

// verify key algorithm
keyVersionDetails, err := getKeyVersionDetails(client, keyName)
keyVersionDetails, err := resolveKeyVersionDetails(client, keyName)
if err != nil {
return nil, fmt.Errorf("failed to get key version details: %w", err)
}
Expand All @@ -55,14 +69,17 @@ func NewGcpKmsSigner(keyName string) (Signer, error) {
return nil, fmt.Errorf("key algorithm is not EC_SIGN_P256_SHA256")
}

ethAddress, err := resolveEthAddr(client, keyName)
// resolve public key
pubKeyBytes, err := resolvePubKeyBytes(client, keyName)
if err != nil {
return nil, fmt.Errorf("failed to resolve Ethereum address: %w", err)
return nil, fmt.Errorf("failed to resolve public key: %w", err)
}

secp256k1PubKey, err := resolveSecp256k1PubKey(client, keyName)
ethAddress := getEthAddr(pubKeyBytes)

secp256k1PubKey, err := getSecp256k1PubKey(pubKeyBytes)
if err != nil {
return nil, fmt.Errorf("failed to resolve Secp256k1 public key: %w", err)
return nil, fmt.Errorf("failed to get secp256k1 public key: %w", err)
}

cosmosPubKey := &secp256k1.PubKey{Key: secp256k1PubKey.SerializeCompressed()}
Expand All @@ -76,6 +93,11 @@ func NewGcpKmsSigner(keyName string) (Signer, error) {
}, nil
}

// Destructor Function
func (s *GcpKmsSigner) Destroy() {
s.client.Close()
}

// Method Implementations
func (s *GcpKmsSigner) EthSign(data []byte) ([]byte, error) {
digest := data
Expand Down Expand Up @@ -103,16 +125,17 @@ func (s *GcpKmsSigner) CosmosPublicKey() types.PubKey {
return s.cosmosPubKey
}

// Helper Functions (same as before)
func resolveEthAddr(client *kms.KeyManagementClient, keyName string) (common.Address, error) {
resp, err := client.GetPublicKey(context.Background(), &kmspb.GetPublicKeyRequest{Name: keyName})
func resolvePubKeyBytes(client GCPKeyManagementClient, keyName string) ([]byte, error) {
publicKeyResp, err := client.GetPublicKey(context.Background(), &kmspb.GetPublicKeyRequest{Name: keyName})
if err != nil {
return common.Address{}, fmt.Errorf("public key %q lookup: %w", keyName, err)
return nil, fmt.Errorf("failed to get public key: %w", err)
}

block, _ := pem.Decode([]byte(resp.Pem))
publicKeyPem := publicKeyResp.Pem

block, _ := pem.Decode([]byte(publicKeyPem))
if block == nil {
return common.Address{}, fmt.Errorf("public key %q PEM empty: %.130q", keyName, resp.Pem)
return nil, fmt.Errorf("public key %q PEM empty: %.130q", keyName, publicKeyPem)
}

var info struct {
Expand All @@ -121,26 +144,26 @@ func resolveEthAddr(client *kms.KeyManagementClient, keyName string) (common.Add
}
_, err = asn1.Unmarshal(block.Bytes, &info)
if err != nil {
return common.Address{}, fmt.Errorf("public key %q PEM block %q: %w", keyName, block.Type, err)
return nil, fmt.Errorf("public key %q PEM block %q: %w", keyName, block.Type, err)
}

wantAlg := asn1.ObjectIdentifier{1, 2, 840, 10045, 2, 1}
if gotAlg := info.AlgID.Algorithm; !gotAlg.Equal(wantAlg) {
return common.Address{}, fmt.Errorf("public key %q ASN.1 algorithm %s intead of %s", keyName, gotAlg, wantAlg)
return nil, fmt.Errorf("public key %q ASN.1 algorithm %s instead of %s", keyName, gotAlg, wantAlg)
}

return ethPubKeyAddr(info.Key.Bytes), nil
return info.Key.Bytes, nil
}

// PubKeyAddr returns the Ethereum address for (uncompressed-)key bytes.
func ethPubKeyAddr(bytes []byte) common.Address {
func getEthAddr(bytes []byte) common.Address {
digest := crypto.Keccak256(bytes[1:])
var addr common.Address
copy(addr[:], digest[12:])
return addr
}

func ethSignHash(hash common.Hash, client *kms.KeyManagementClient, keyName string, ethAddress common.Address) ([]byte, error) {
func ethSignHash(hash common.Hash, client GCPKeyManagementClient, keyName string, ethAddress common.Address) ([]byte, error) {
// Resolve a signature
req := kmspb.AsymmetricSignRequest{
Name: keyName,
Expand All @@ -161,6 +184,7 @@ func ethSignHash(hash common.Hash, client *kms.KeyManagementClient, keyName stri
if err != nil {
return nil, fmt.Errorf("asymmetric signature encoding: %w", err)
}

var rLen, sLen int // byte size
if params.R != nil {
rLen = (params.R.BitLen() + 7) / 8
Expand Down Expand Up @@ -190,7 +214,7 @@ func ethSignHash(hash common.Hash, client *kms.KeyManagementClient, keyName stri
continue
}

if ethPubKeyAddr(pubKey.SerializeUncompressed()) == ethAddress {
if getEthAddr(pubKey.SerializeUncompressed()) == ethAddress {
// Sign the transaction
sig[65] = recoveryID // Ethereum 'v' parameter
return sig[1:], nil // Exclude BitCoin header
Expand All @@ -200,7 +224,7 @@ func ethSignHash(hash common.Hash, client *kms.KeyManagementClient, keyName stri
return nil, fmt.Errorf("asymmetric signature address recovery mis: %w", recoverErr)
}

func cosmosSignHash(client *kms.KeyManagementClient, keyName string, hash [32]byte, pubKey *dcrecSecp256k1.PublicKey) ([]byte, error) {
func cosmosSignHash(client GCPKeyManagementClient, keyName string, hash [32]byte, pubKey *dcrecSecp256k1.PublicKey) ([]byte, error) {
// Sign the hash using KMS
req := &kmspb.AsymmetricSignRequest{
Name: keyName,
Expand Down Expand Up @@ -264,53 +288,27 @@ func signatureFromBytes(sigStr []byte) (*btcecdsa.Signature, error) {
return btcecdsa.NewSignature(&r, &s), nil
}

// func resolveCosmosPubKey(client *kms.KeyManagementClient, keyName string) (*secp256k1.PubKey, error) {
// pubkeyObject, err := resolveSecp256k1PubKey(client, keyName)
// func getCosmosPubKey(pubKeyBytes []byte) (*secp256k1.PubKey, error) {
// pubkeyObject, err := getSecp256k1PubKey(pubKeyBytes)
// if err != nil {
// return nil, fmt.Errorf("failed to resolve public key: %w", err)
// return nil, err
// }
//
// pk := pubkeyObject.SerializeCompressed()
//
// return &secp256k1.PubKey{Key: pk}, nil
// }

func resolveSecp256k1PubKey(client *kms.KeyManagementClient, keyName string) (*dcrecSecp256k1.PublicKey, error) {
publicKeyResp, err := client.GetPublicKey(context.Background(), &kmspb.GetPublicKeyRequest{Name: keyName})
if err != nil {
return nil, fmt.Errorf("failed to get public key: %w", err)
}

publicKeyPem := publicKeyResp.Pem

block, _ := pem.Decode([]byte(publicKeyPem))
if block == nil {
return nil, fmt.Errorf("public key %q PEM empty: %.130q", keyName, publicKeyPem)
}

var info struct {
AlgID pkix.AlgorithmIdentifier
Key asn1.BitString
}
_, err = asn1.Unmarshal(block.Bytes, &info)
if err != nil {
return nil, fmt.Errorf("public key %q PEM block %q: %w", keyName, block.Type, err)
}

wantAlg := asn1.ObjectIdentifier{1, 2, 840, 10045, 2, 1}
if gotAlg := info.AlgID.Algorithm; !gotAlg.Equal(wantAlg) {
return nil, fmt.Errorf("public key %q ASN.1 algorithm %s instead of %s", keyName, gotAlg, wantAlg)
}

pubkeyObject, err := dcrecSecp256k1.ParsePubKey(info.Key.Bytes)
func getSecp256k1PubKey(pubKeyBytes []byte) (*dcrecSecp256k1.PublicKey, error) {
pubkeyObject, err := dcrecSecp256k1.ParsePubKey(pubKeyBytes)
if err != nil {
return nil, fmt.Errorf("failed to parse public key: %w", err)
}

return pubkeyObject, nil
}

func getKeyVersionDetails(client *kms.KeyManagementClient, keyName string) (*kmspb.CryptoKeyVersion, error) {
func resolveKeyVersionDetails(client GCPKeyManagementClient, keyName string) (*kmspb.CryptoKeyVersion, error) {
// Request the key version details
req := &kmspb.GetCryptoKeyVersionRequest{
Name: keyName,
Expand Down
154 changes: 154 additions & 0 deletions common/gcp_kms_signer_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,154 @@
package common

import (
"context"
"encoding/asn1"
"math/big"
"testing"

"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"

"cloud.google.com/go/kms/apiv1/kmspb"
"github.com/cosmos/cosmos-sdk/crypto/keys/secp256k1"
dcrecSecp256k1 "github.com/decred/dcrd/dcrec/secp256k1/v4"
"github.com/ethereum/go-ethereum/common"
gax "github.com/googleapis/gax-go/v2"
)

// MockGCPKeyManagementClient is a mock implementation of GCPKeyManagementClient
type MockGCPKeyManagementClient struct {
mock.Mock
}

func (m *MockGCPKeyManagementClient) Close() error {
args := m.Called()
return args.Error(0)
}

func (m *MockGCPKeyManagementClient) GetPublicKey(ctx context.Context, req *kmspb.GetPublicKeyRequest, opts ...gax.CallOption) (*kmspb.PublicKey, error) {
args := m.Called(ctx, req, opts)
return args.Get(0).(*kmspb.PublicKey), args.Error(1)
}

func (m *MockGCPKeyManagementClient) AsymmetricSign(ctx context.Context, req *kmspb.AsymmetricSignRequest, opts ...gax.CallOption) (*kmspb.AsymmetricSignResponse, error) {
args := m.Called(ctx, req, opts)
return args.Get(0).(*kmspb.AsymmetricSignResponse), args.Error(1)
}

func (m *MockGCPKeyManagementClient) GetCryptoKeyVersion(ctx context.Context, req *kmspb.GetCryptoKeyVersionRequest, opts ...gax.CallOption) (*kmspb.CryptoKeyVersion, error) {
args := m.Called(ctx, req, opts)
return args.Get(0).(*kmspb.CryptoKeyVersion), args.Error(1)
}

func mockASN1Signature() []byte {
r, _ := new(big.Int).SetString("0adb0cad4842abd2003c8dcf4f2663b1ba9106b208a8b9f62aaa628121d7e818", 16)
s, _ := new(big.Int).SetString("05a2e4ace4f4dae5437804d8ff3a67151e96d0efad0427cb49844ea80c3a03a5", 16)
return asn1Bytes(r, s)
}

func mockPublicKey() []byte {
keyHex := "0459fe4b6a8682418cb86df571c0a36b06f1b37e6760985009599c31d783c6ccac6cb07ba459e7d4949d12b178f8468802ebf7d932007e46b02fc7a2bf90761fd6"
return common.Hex2Bytes(keyHex)
}

func asn1Bytes(r, s *big.Int) []byte {
signature, _ := asn1.Marshal(struct {
R, S *big.Int
}{R: r, S: s})
return signature
}

// Unit tests for GcpKmsSigner
func TestNewGcpKmsSigner(t *testing.T) {
mockClient := new(MockGCPKeyManagementClient)
NewGCPKeyManagementClient = func(ctx context.Context) (GCPKeyManagementClient, error) {
return mockClient, nil
}

keyName := "test-key"
expectedKeyVersion := &kmspb.CryptoKeyVersion{
Algorithm: kmspb.CryptoKeyVersion_EC_SIGN_SECP256K1_SHA256,
}
mockClient.On("GetCryptoKeyVersion", mock.Anything, &kmspb.GetCryptoKeyVersionRequest{Name: keyName}, mock.Anything).Return(expectedKeyVersion, nil)

expectedPublicKey := &kmspb.PublicKey{
Pem: "-----BEGIN PUBLIC KEY-----\nMFYwEAYHKoZIzj0CAQYFK4EEAAoDQgAEWf5LaoaCQYy4bfVxwKNrBvGzfmdgmFAJ\nWZwx14PGzKxssHukWefUlJ0SsXj4RogC6/fZMgB+RrAvx6K/kHYf1g==\n-----END PUBLIC KEY-----",
}
mockClient.On("GetPublicKey", mock.Anything, &kmspb.GetPublicKeyRequest{Name: keyName}, mock.Anything).Return(expectedPublicKey, nil)

signer, err := NewGcpKmsSigner(keyName)
assert.NoError(t, err)
assert.NotNil(t, signer)

mockClient.AssertExpectations(t)
}

func TestGcpKmsSigner_EthSign(t *testing.T) {
mockClient := new(MockGCPKeyManagementClient)
keyName := "test-key"
ethAddress := common.HexToAddress("0xE14FE42D9Cf3C39c27E1c1Ca3dAa033C43eD4d88")

signer := &GcpKmsSigner{
client: mockClient,
keyName: keyName,
ethAddress: ethAddress,
}

data := []byte("example transaction data")
// expectedHash := crypto.Keccak256(data)
expectedSignature := &kmspb.AsymmetricSignResponse{
Signature: mockASN1Signature(),
}
mockClient.On("AsymmetricSign", mock.Anything, mock.Anything, mock.Anything).Return(expectedSignature, nil)

sig, err := signer.EthSign(data)
assert.NoError(t, err)
assert.NotNil(t, sig)

mockClient.AssertExpectations(t)
}

func TestGcpKmsSigner_CosmosSign(t *testing.T) {
mockClient := new(MockGCPKeyManagementClient)
keyName := "test-key"

pubKeyBytes := mockPublicKey()
secp256k1PubKey, _ := dcrecSecp256k1.ParsePubKey(pubKeyBytes)
cosmosPubKey := &secp256k1.PubKey{Key: secp256k1PubKey.SerializeCompressed()}

signer := &GcpKmsSigner{
client: mockClient,
keyName: keyName,
secp256k1PubKey: secp256k1PubKey,
cosmosPubKey: cosmosPubKey,
}

data := []byte("example transaction data")
// expectedHash := crypto.Keccak256(data)
expectedSignature := &kmspb.AsymmetricSignResponse{
Signature: mockASN1Signature(),
}
mockClient.On("AsymmetricSign", mock.Anything, mock.Anything, mock.Anything).Return(expectedSignature, nil)

sig, err := signer.CosmosSign(data)
assert.NoError(t, err)
assert.NotNil(t, sig)

mockClient.AssertExpectations(t)
}

func TestGcpKmsSigner_Destroy(t *testing.T) {
mockClient := new(MockGCPKeyManagementClient)
keyName := "test-key"

signer := &GcpKmsSigner{
client: mockClient,
keyName: keyName,
}

mockClient.On("Close").Return(nil)

signer.Destroy()
mockClient.AssertExpectations(t)
}
Loading

0 comments on commit e10dc35

Please sign in to comment.