Add AWS KMS support for account keys

CHANGELOG.md

@@ -7,6 +7,24 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Added
+
+AWS KMS support for account keys. Accounts can now be configured with the `aws-kms` key type, which signs transactions using an AWS KMS asymmetric key referenced by its ARN. It works like the existing `google-kms` key type, only the `resourceID` format differs:
+
+```json
+"accounts": {
+  "my-account": {
+    "address": "0x0000000000000000",
+    "key": {
+      "type": "aws-kms",
+      "resourceID": "arn:aws:kms:us-west-2:111122223333:key/1234abcd-12ab-34cd-56ef-1234567890ab"
+    }
+  }
+}
+```
+
+The region is read from the ARN, credentials are resolved from the default AWS configuration, and the key must be an ECDSA P-256 or secp256k1 key. AWS KMS signs those keys with `ECDSA_SHA_256` only, so the hash algorithm is fixed to `SHA2_256` and setting a different one in the configuration is an error.
+
 ## 1.0.0
 
 ### Changed

accounts/account_test.go

@@ -22,7 +22,10 @@ import (
 	"testing"
 
 	"github.com/onflow/flow-go-sdk"
+	"github.com/onflow/flow-go-sdk/crypto"
 	"github.com/stretchr/testify/assert"
+
+	"github.com/onflow/flowkit/v2/config"
 )
 
 func Test_Accounts(t *testing.T) {
@@ -128,3 +131,35 @@ func Test_Accounts(t *testing.T) {
 	})
 
 }
+
+func Test_AWSKMS_AccountRoundTrip(t *testing.T) {
+	conf := config.Config{
+		Accounts: config.Accounts{
+			{
+				Name:    "aws-kms-account",
+				Address: flow.HexToAddress("0xf8d6e0586b0a20c7"),
+				Key: config.AccountKey{
+					Type:       config.KeyTypeAWSKMS,
+					SigAlgo:    config.DefaultSigAlgo,
+					HashAlgo:   config.DefaultHashAlgo, // must be replaced by SHA2_256
+					ResourceID: "arn:aws:kms:us-west-2:111122223333:key/1234abcd-12ab-34cd-56ef-1234567890ab",
+				},
+			},
+		},
+	}
+
+	accounts, err := FromConfig(&conf)
+	assert.NoError(t, err)
+
+	account, err := accounts.ByName("aws-kms-account")
+	assert.NoError(t, err)
+	assert.Equal(t, config.KeyTypeAWSKMS, account.Key.Type())
+	assert.Equal(t, crypto.SHA2_256, account.Key.HashAlgo())
+
+	roundTripped := ToConfig(accounts)
+	assert.Equal(t, crypto.SHA2_256, roundTripped[0].Key.HashAlgo)
+	assert.Equal(t,
+		"arn:aws:kms:us-west-2:111122223333:key/1234abcd-12ab-34cd-56ef-1234567890ab",
+		roundTripped[0].Key.ResourceID,
+	)
+}

accounts/keys.go

@@ -27,9 +27,12 @@ import (
 	"regexp"
 	"strings"
 
+	"github.com/aws/aws-sdk-go-v2/aws"
+	awsconfig "github.com/aws/aws-sdk-go-v2/config"
 	gethAccounts "github.com/ethereum/go-ethereum/accounts"
 	"github.com/lmars/go-slip10"
 	"github.com/onflow/flow-go-sdk/crypto"
+	"github.com/onflow/flow-go-sdk/crypto/awskms"
 	"github.com/onflow/flow-go-sdk/crypto/cloudkms"
 	"github.com/tyler-smith/go-bip39"
 
@@ -69,6 +72,8 @@ var _ Key = &HexKey{}
 
 var _ Key = &KMSKey{}
 
+var _ Key = &AWSKMSKey{}
+
 var _ Key = &BIP44Key{}
 
 var _ Key = &EnvKey{}
@@ -84,6 +89,8 @@ func keyFromConfig(accountKeyConf config.AccountKey) (Key, error) {
 		return bip44KeyFromConfig(accountKeyConf)
 	case config.KeyTypeGoogleKMS:
 		return kmsKeyFromConfig(accountKeyConf)
+	case config.KeyTypeAWSKMS:
+		return awsKMSKeyFromConfig(accountKeyConf)
 	case config.KeyTypeFile:
 		return fileKeyFromConfig(accountKeyConf)
 	}
@@ -228,6 +235,106 @@ func kmsKeyFromConfig(key config.AccountKey) (Key, error) {
 	}, nil
 }
 
+// AWSKMSKey implements AWS KMS system for signing.
+type AWSKMSKey struct {
+	*baseKey
+	kmsKey awskms.Key
+}
+
+// ToConfig convert account key to configuration.
+func (a *AWSKMSKey) ToConfig() config.AccountKey {
+	return config.AccountKey{
+		Type:       a.keyType,
+		Index:      a.index,
+		SigAlgo:    a.sigAlgo,
+		HashAlgo:   a.hashAlgo,
+		ResourceID: a.kmsKey.ARN(),
+	}
+}
+
+func (a *AWSKMSKey) Signer(ctx context.Context) (crypto.Signer, error) {
+	awsCfg, err := awsKMSConfig(ctx, a.kmsKey)
+	if err != nil {
+		return nil, err
+	}
+
+	accountKMSSigner, err := awskms.NewClient(awsCfg).SignerForKey(
+		ctx,
+		a.kmsKey,
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	return accountKMSSigner, nil
+}
+
+func (a *AWSKMSKey) Validate() error {
+	awsCfg, err := awsKMSConfig(context.Background(), a.kmsKey)
+	if err != nil {
+		return err
+	}
+
+	if _, err = awsCfg.Credentials.Retrieve(context.Background()); err != nil {
+		return fmt.Errorf(
+			"could not load AWS credentials for KMS key %s, configure them with `aws configure sso` or `aws sso login`: %w",
+			a.kmsKey.ARN(),
+			err,
+		)
+	}
+
+	return nil
+}
+
+func (a *AWSKMSKey) PrivateKey() (*crypto.PrivateKey, error) {
+	return nil, fmt.Errorf("private key not accessible")
+}
+
+// awsKMSKeyARNPattern matches the canonical AWS KMS key ARN, for example
+// arn:aws:kms:us-west-2:111122223333:key/1234abcd-12ab-34cd-56ef-1234567890ab
+//
+// Key aliases and other KMS resource types are not supported because they cannot be
+// reconstructed from a parsed key, and the signing key must be addressed by its ARN.
+//
+// Only the standard `aws` partition is accepted. awskms.Key.ARN() rebuilds the ARN as
+// arn:aws:kms:<region>:<account>:key/<id> and the SDK sends that as the KeyId, so a
+// aws-cn or aws-us-gov key would be addressed with the wrong partition and rejected by
+// the service. Supporting those partitions needs a change in the SDK.
+var awsKMSKeyARNPattern = regexp.MustCompile(`^arn:aws:kms:[a-z0-9-]+:[0-9]{12}:key/[A-Za-z0-9-]+$`)
+
+func awsKMSKeyFromConfig(key config.AccountKey) (Key, error) {
+	if !awsKMSKeyARNPattern.MatchString(key.ResourceID) {
+		return nil, fmt.Errorf(
+			"invalid AWS KMS resource ID %q, expected the ARN format arn:aws:kms:<region>:<account-id>:key/<key-id> (the aws-cn and aws-us-gov partitions are not supported)",
+			key.ResourceID,
+		)
+	}
+
+	accountKMSKey, err := awskms.KeyFromResourceARN(key.ResourceID)
+	if err != nil {
+		return nil, err
+	}
+
+	return &AWSKMSKey{
+		baseKey: &baseKey{
+			keyType: config.KeyTypeAWSKMS,
+			index:   key.Index,
+			sigAlgo: key.SigAlgo,
+			// AWS KMS signs ECDSA keys with ECDSA_SHA_256 only, so the configured hash
+			// algorithm is overridden rather than trusted. The JSON layer rejects a
+			// conflicting value; this covers programmatic construction of the config.
+			hashAlgo: crypto.SHA2_256,
+		},
+		kmsKey: accountKMSKey,
+	}, nil
+}
+
+// awsKMSConfig loads the default AWS configuration, defaulting the region to the one
+// the KMS key lives in so a region does not have to be configured separately.
+func awsKMSConfig(ctx context.Context, key awskms.Key) (aws.Config, error) {
+	return awsconfig.LoadDefaultConfig(ctx, awsconfig.WithRegion(key.Region))
+}
+
 // HexKey implements account key in hex representation.
 type HexKey struct {
 	*baseKey

accounts/keys_test.go

@@ -47,6 +47,66 @@ func Test_KMS_Keys(t *testing.T) {
 	assert.Equal(t, confKey, kmsKey.ToConfig())
 }
 
+func Test_AWSKMS_Keys(t *testing.T) {
+	confKey := config.AccountKey{
+		Type:       config.KeyTypeAWSKMS,
+		Index:      0,
+		SigAlgo:    config.DefaultSigAlgo,
+		HashAlgo:   crypto.SHA2_256,
+		ResourceID: "arn:aws:kms:us-west-2:111122223333:key/1234abcd-12ab-34cd-56ef-1234567890ab",
+	}
+
+	kmsKey, err := keyFromConfig(confKey)
+	assert.NoError(t, err)
+
+	_, err = kmsKey.PrivateKey()
+	assert.EqualError(t, err, "private key not accessible")
+	assert.Equal(t, confKey, kmsKey.ToConfig())
+}
+
+// AWS KMS signs ECDSA keys with ECDSA_SHA_256 only, so the configured hash algorithm must
+// never reach the signer.
+func Test_AWSKMS_KeysHashAlgorithm(t *testing.T) {
+	confKey := config.AccountKey{
+		Type:       config.KeyTypeAWSKMS,
+		SigAlgo:    config.DefaultSigAlgo,
+		HashAlgo:   crypto.SHA3_256,
+		ResourceID: "arn:aws:kms:us-west-2:111122223333:key/1234abcd-12ab-34cd-56ef-1234567890ab",
+	}
+
+	kmsKey, err := keyFromConfig(confKey)
+	assert.NoError(t, err)
+
+	assert.Equal(t, crypto.SHA2_256, kmsKey.HashAlgo())
+	assert.Equal(t, crypto.SHA2_256, kmsKey.ToConfig().HashAlgo)
+}
+
+func Test_AWSKMS_KeysInvalidResourceID(t *testing.T) {
+	tests := map[string]string{
+		"empty":            "",
+		"google resource":  "projects/my-project/locations/global/keyRings/flow/cryptoKeys/my-account/cryptoKeyVersions/1",
+		"key without id":   "arn:aws:kms:us-west-2:111122223333:key",
+		"alias":            "arn:aws:kms:us-west-2:111122223333:alias/my-key",
+		"short account id": "arn:aws:kms:us-west-2:1111222233333:key/1234abcd",
+		"missing region":   "arn:aws:kms::111122223333:key/1234abcd",
+		"missing account":  "arn:aws:kms:us-west-2::key/1234abcd",
+		"govcloud":         "arn:aws-us-gov:kms:us-gov-west-1:111122223333:key/1234abcd",
+		"china":            "arn:aws-cn:kms:cn-north-1:111122223333:key/1234abcd",
+	}
+
+	for name, resourceID := range tests {
+		t.Run(name, func(t *testing.T) {
+			_, err := keyFromConfig(config.AccountKey{
+				Type:       config.KeyTypeAWSKMS,
+				SigAlgo:    config.DefaultSigAlgo,
+				HashAlgo:   config.DefaultHashAlgo,
+				ResourceID: resourceID,
+			})
+			assert.Error(t, err)
+		})
+	}
+}
+
 func Test_File_key(t *testing.T) {
 	confKey := config.AccountKey{
 		Type:     config.KeyTypeFile,

config/config.go

@@ -46,6 +46,7 @@ type KeyType string
 const (
 	KeyTypeHex       KeyType = "hex"
 	KeyTypeGoogleKMS KeyType = "google-kms"
+	KeyTypeAWSKMS    KeyType = "aws-kms"
 	KeyTypeBip44     KeyType = "bip44"
 	KeyTypeFile      KeyType = "file"
 )

config/json/account.go

@@ -154,11 +154,20 @@ func transformAdvancedToConfig(accountName string, a advancedAccount) (*config.A
 		return nil, fmt.Errorf("invalid hash algorithm for account %s", accountName)
 	}
 
-	validTypes := []config.KeyType{config.KeyTypeHex, config.KeyTypeFile, config.KeyTypeBip44, config.KeyTypeGoogleKMS}
+	validTypes := []config.KeyType{config.KeyTypeHex, config.KeyTypeFile, config.KeyTypeBip44, config.KeyTypeGoogleKMS, config.KeyTypeAWSKMS}
 	if !slices.Contains(validTypes, a.Key.Type) {
 		return nil, fmt.Errorf("invalid key type for account %s", accountName)
 	}
 
+	// AWS KMS signs ECDSA keys with ECDSA_SHA_256 only, so the hash algorithm is not a
+	// choice for that key type.
+	if a.Key.Type == config.KeyTypeAWSKMS {
+		if a.Key.HashAlgo != "" && hashAlgo != crypto.SHA2_256 {
+			return nil, fmt.Errorf("invalid hash algorithm for account %s: aws-kms keys only support SHA2_256", accountName)
+		}
+		hashAlgo = crypto.SHA2_256
+	}
+
 	// check that only one is provided because the values are mutually exclusive
 	set := false
 	for _, v := range []string{a.Key.ResourceID, a.Key.PrivateKey, a.Key.Location} {
@@ -217,7 +226,7 @@ func transformAdvancedToConfig(accountName string, a advancedAccount) (*config.A
 			key.DerivationPath = "m/44'/539'/0'/0/0"
 		}
 
-	case config.KeyTypeGoogleKMS:
+	case config.KeyTypeGoogleKMS, config.KeyTypeAWSKMS:
 		if a.Key.ResourceID == "" {
 			return nil, fmt.Errorf("missing resource ID value for key on account %s", accountName)
 		}
@@ -323,7 +332,7 @@ func transformAdvancedKeyToJSON(key config.AccountKey) advanceKey {
 	case config.KeyTypeBip44:
 		advancedKey.Mnemonic = key.Mnemonic
 		advancedKey.DerivationPath = key.DerivationPath
-	case config.KeyTypeGoogleKMS:
+	case config.KeyTypeGoogleKMS, config.KeyTypeAWSKMS:
 		advancedKey.ResourceID = key.ResourceID
 	case config.KeyTypeFile:
 		advancedKey.Location = filepath.ToSlash(key.Location)
@@ -357,7 +366,7 @@ type advanceKey struct {
 	// bip44 key type
 	Mnemonic       string `json:"mnemonic,omitempty"`
 	DerivationPath string `json:"derivationPath,omitempty"`
-	// kms key type
+	// kms key types (google-kms, aws-kms)
 	ResourceID string `json:"resourceID,omitempty"`
 	// key location
 	Location string `json:"location,omitempty"`