[ARCH-332] KMS keystore support (#1767)

* Wip

* Simplify

* Simpler

* Fix pubkey

* Testing

* test

* Lint

* Expore ASN utilities

* Linter

* Old cli thing

* Simpler cli

* Lint

* More lint

* Lint and simpler interface comp

* Comments

Author: connorwstein

keystore/cli/cli.go

@@ -12,10 +12,11 @@ import (
 	"time"
 
 	"github.com/jmoiron/sqlx"
-	_ "github.com/lib/pq"
+	_ "github.com/lib/pq" // Register postgres driver
 	"github.com/spf13/cobra"
 
 	ks "github.com/smartcontractkit/chainlink-common/keystore"
+	"github.com/smartcontractkit/chainlink-common/keystore/kms"
 	"github.com/smartcontractkit/chainlink-common/keystore/pgstore"
 )
 
@@ -27,12 +28,15 @@ func NewRootCmd() *cobra.Command {
 	cmd := &cobra.Command{
 		Use: "keys",
 		Long: ` 
-CLI for managing keystore keys. Must specify KEYSTORE_FILE_PATH or KEYSTORE_DB_URL 
-and KEYSTORE_PASSWORD in order to load the keystore. 
+CLI for managing keystore keys. 
 
-KEYSTORE_FILE_PATH: is the path to the keystore file, can be empty for a new keystore.
-File must already exist. Example to create a new keystore file: touch ./keystore.json
+If KEYSTORE_KMS_PROFILE is set, will load the keystore from KMS.
+KEYSTORE_KMS_PROFILE: is the AWS profile to use for KMS (region will be taken from the profile).
 
+Otherwise, will load the keystore from a file or database.
+KEYSTORE_PASSWORD: password used to encrypt the key material before storage, must be provided.
+KEYSTORE_FILE_PATH: is the path to the keystore file, can be empty for a new keystore.
+File must already exist. Example to create a new keystore file: touch ./keystore.json. Either KEYSTORE_FILE_PATH or KEYSTORE_DB_URL must be set.
 KEYSTORE_DB_URL: is the postgres connection URL. Only use this if your keystore is stored in a pg database.
 Requires a pg database with a 'encrypted_keystore' table with the following schema:
 CREATE TABLE IF NOT EXISTS encrypted_keystore (
@@ -42,17 +46,42 @@ CREATE TABLE IF NOT EXISTS encrypted_keystore (
     updated_at timestamptz NOT NULL DEFAULT NOW(),
     encrypted_data BYTEA NOT NULL
 );
-
-KEYSTORE_PASSWORD is the password used to encrypt the key material before storage.
 `,
 		Short:        "CLI for managing keystore keys",
 		SilenceUsage: true,
 	}
-	cmd.PersistentFlags().String("file-path", "", "Overrides KEYSTORE_FILE_PATH environment variable")
-	cmd.PersistentFlags().String("db-url", "", "Overrides KEYSTORE_DB_URL environment variable")
-	cmd.PersistentFlags().String("password", "", "Overrides KEYSTORE_PASSWORD environment variable. Not recommended as will leave shell traces.")
 
-	cmd.AddCommand(NewListCmd(), NewGetCmd(), NewCreateCmd(), NewDeleteCmd(), NewExportCmd(), NewImportCmd(), NewSetMetadataCmd(), NewSignCmd(), NewVerifyCmd(), NewEncryptCmd(), NewDecryptCmd())
+	// Check if KMS profile is set - if so, hide commands that don't work with KMS
+	isKMSMode := os.Getenv("KEYSTORE_KMS_PROFILE") != ""
+
+	// Commands that work with both regular keystore and KMS
+	listCmd := NewListCmd()
+	getCmd := NewGetCmd()
+	signCmd := NewSignCmd()
+	verifyCmd := NewVerifyCmd()
+
+	// Commands that only work with regular keystore (not KMS)
+	createCmd := NewCreateCmd()
+	deleteCmd := NewDeleteCmd()
+	exportCmd := NewExportCmd()
+	importCmd := NewImportCmd()
+	setMetadataCmd := NewSetMetadataCmd()
+	// Note these could potentially be supported with KMS, but not yet implemented.
+	encryptCmd := NewEncryptCmd()
+	decryptCmd := NewDecryptCmd()
+
+	// Hide admin/encryption commands when using KMS (keys are managed externally)
+	if isKMSMode {
+		createCmd.Hidden = true
+		deleteCmd.Hidden = true
+		exportCmd.Hidden = true
+		importCmd.Hidden = true
+		setMetadataCmd.Hidden = true
+		encryptCmd.Hidden = true
+		decryptCmd.Hidden = true
+	}
+
+	cmd.AddCommand(listCmd, getCmd, createCmd, deleteCmd, exportCmd, importCmd, setMetadataCmd, signCmd, verifyCmd, encryptCmd, decryptCmd)
 	return cmd
 }
 
@@ -62,7 +91,7 @@ func NewListCmd() *cobra.Command {
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx, cancel := context.WithTimeout(cmd.Context(), KeystoreLoadTimeout)
 			defer cancel()
-			k, err := loadKeystore(ctx, cmd)
+			k, err := loadKeystoreSignerReader(ctx, cmd)
 			if err != nil {
 				return err
 			}
@@ -85,7 +114,13 @@ func NewGetCmd() *cobra.Command {
 	cmd := cobra.Command{
 		Use: "get", Short: "Get keys",
 		RunE: func(cmd *cobra.Command, args []string) error {
-			return runKeystoreCommand[ks.GetKeysRequest, ks.GetKeysResponse](cmd, args, func(ctx context.Context, k ks.Keystore, req ks.GetKeysRequest) (ks.GetKeysResponse, error) {
+			return runKeystoreCommand[interface {
+				ks.Reader
+				ks.Signer
+			}, ks.GetKeysRequest, ks.GetKeysResponse](cmd, args, loadKeystoreSignerReader, func(ctx context.Context, k interface {
+				ks.Reader
+				ks.Signer
+			}, req ks.GetKeysRequest) (ks.GetKeysResponse, error) {
 				return k.GetKeys(ctx, req)
 			})
 		},
@@ -99,7 +134,7 @@ func NewCreateCmd() *cobra.Command {
 	cmd := cobra.Command{
 		Use: "create", Short: "Create a key",
 		RunE: func(cmd *cobra.Command, args []string) error {
-			return runKeystoreCommand[ks.CreateKeysRequest, ks.CreateKeysResponse](cmd, args, func(ctx context.Context, k ks.Keystore, req ks.CreateKeysRequest) (ks.CreateKeysResponse, error) {
+			return runKeystoreCommand[ks.Keystore, ks.CreateKeysRequest, ks.CreateKeysResponse](cmd, args, loadKeystore, func(ctx context.Context, k ks.Keystore, req ks.CreateKeysRequest) (ks.CreateKeysResponse, error) {
 				return k.CreateKeys(ctx, req)
 			})
 		},
@@ -128,7 +163,7 @@ func NewDeleteCmd() *cobra.Command {
 			}
 			if !confirmYes {
 				// Prompt for confirmation on stdin
-				_, err = cmd.OutOrStderr().Write([]byte(fmt.Sprintf("This will permanently delete keys: %v. Type 'yes' to confirm: ", strings.Join(req.KeyNames, ", "))))
+				_, err = fmt.Fprintf(cmd.OutOrStderr(), "This will permanently delete keys: %v. Type 'yes' to confirm: ", strings.Join(req.KeyNames, ", "))
 				if err != nil {
 					return err
 				}
@@ -159,7 +194,7 @@ func NewExportCmd() *cobra.Command {
 	cmd := cobra.Command{
 		Use: "export", Short: "Export a key to an encrypted JSON file",
 		RunE: func(cmd *cobra.Command, args []string) error {
-			return runKeystoreCommand[ks.ExportKeysRequest, ks.ExportKeysResponse](cmd, args, func(ctx context.Context, k ks.Keystore, req ks.ExportKeysRequest) (ks.ExportKeysResponse, error) {
+			return runKeystoreCommand[ks.Keystore, ks.ExportKeysRequest, ks.ExportKeysResponse](cmd, args, loadKeystore, func(ctx context.Context, k ks.Keystore, req ks.ExportKeysRequest) (ks.ExportKeysResponse, error) {
 				return k.ExportKeys(ctx, req)
 			})
 		},
@@ -173,7 +208,7 @@ func NewImportCmd() *cobra.Command {
 	cmd := cobra.Command{
 		Use: "import", Short: "Import an encrypted key JSON file",
 		RunE: func(cmd *cobra.Command, args []string) error {
-			return runKeystoreCommand[ks.ImportKeysRequest, ks.ImportKeysResponse](cmd, args, func(ctx context.Context, k ks.Keystore, req ks.ImportKeysRequest) (ks.ImportKeysResponse, error) {
+			return runKeystoreCommand[ks.Keystore, ks.ImportKeysRequest, ks.ImportKeysResponse](cmd, args, loadKeystore, func(ctx context.Context, k ks.Keystore, req ks.ImportKeysRequest) (ks.ImportKeysResponse, error) {
 				return k.ImportKeys(ctx, req)
 			})
 		},
@@ -187,7 +222,7 @@ func NewSetMetadataCmd() *cobra.Command {
 	cmd := cobra.Command{
 		Use: "set-metadata", Short: "Set metadata for keys",
 		RunE: func(cmd *cobra.Command, args []string) error {
-			return runKeystoreCommand[ks.SetMetadataRequest, ks.SetMetadataResponse](cmd, args, func(ctx context.Context, k ks.Keystore, req ks.SetMetadataRequest) (ks.SetMetadataResponse, error) {
+			return runKeystoreCommand[ks.Keystore, ks.SetMetadataRequest, ks.SetMetadataResponse](cmd, args, loadKeystore, func(ctx context.Context, k ks.Keystore, req ks.SetMetadataRequest) (ks.SetMetadataResponse, error) {
 				return k.SetMetadata(ctx, req)
 			})
 		},
@@ -197,8 +232,13 @@ func NewSetMetadataCmd() *cobra.Command {
 	return &cmd
 }
 
-func runKeystoreCommand[Req any, Resp any](cmd *cobra.Command, args []string, fn func(ctx context.Context, k ks.Keystore,
-	req Req) (Resp, error)) error {
+// runKeystoreCommandGeneric is a generic helper that runs a keystore command with a custom loader function.
+func runKeystoreCommand[K any, Req any, Resp any](
+	cmd *cobra.Command,
+	args []string,
+	loader func(ctx context.Context, cmd *cobra.Command) (K, error),
+	fn func(ctx context.Context, k K, req Req) (Resp, error),
+) error {
 	jsonBytes, err := readJSONInput(cmd)
 	if err != nil {
 		return err
@@ -210,7 +250,7 @@ func runKeystoreCommand[Req any, Resp any](cmd *cobra.Command, args []string, fn
 	}
 	ctx, cancel := context.WithTimeout(cmd.Context(), KeystoreLoadTimeout)
 	defer cancel()
-	k, err := loadKeystore(ctx, cmd)
+	k, err := loader(ctx, cmd)
 	if err != nil {
 		return err
 	}
@@ -233,7 +273,13 @@ func NewSignCmd() *cobra.Command {
 	cmd := cobra.Command{
 		Use: "sign", Short: "Sign data with a key",
 		RunE: func(cmd *cobra.Command, args []string) error {
-			return runKeystoreCommand[ks.SignRequest, ks.SignResponse](cmd, args, func(ctx context.Context, k ks.Keystore, req ks.SignRequest) (ks.SignResponse, error) {
+			return runKeystoreCommand[interface {
+				ks.Reader
+				ks.Signer
+			}, ks.SignRequest, ks.SignResponse](cmd, args, loadKeystoreSignerReader, func(ctx context.Context, k interface {
+				ks.Reader
+				ks.Signer
+			}, req ks.SignRequest) (ks.SignResponse, error) {
 				return k.Sign(ctx, req)
 			})
 		},
@@ -272,7 +318,13 @@ func NewVerifyCmd() *cobra.Command {
 	cmd := cobra.Command{
 		Use: "verify", Short: "Verify a signature",
 		RunE: func(cmd *cobra.Command, args []string) error {
-			return runKeystoreCommand[ks.VerifyRequest, ks.VerifyResponse](cmd, args, func(ctx context.Context, k ks.Keystore, req ks.VerifyRequest) (ks.VerifyResponse, error) {
+			return runKeystoreCommand[interface {
+				ks.Reader
+				ks.Signer
+			}, ks.VerifyRequest, ks.VerifyResponse](cmd, args, loadKeystoreSignerReader, func(ctx context.Context, k interface {
+				ks.Reader
+				ks.Signer
+			}, req ks.VerifyRequest) (ks.VerifyResponse, error) {
 				return k.Verify(ctx, req)
 			})
 		},
@@ -286,7 +338,7 @@ func NewEncryptCmd() *cobra.Command {
 	cmd := cobra.Command{
 		Use: "encrypt", Short: "Encrypt data to a remote public key",
 		RunE: func(cmd *cobra.Command, args []string) error {
-			return runKeystoreCommand[ks.EncryptRequest, ks.EncryptResponse](cmd, args, func(ctx context.Context, k ks.Keystore, req ks.EncryptRequest) (ks.EncryptResponse, error) {
+			return runKeystoreCommand[ks.Keystore, ks.EncryptRequest, ks.EncryptResponse](cmd, args, loadKeystore, func(ctx context.Context, k ks.Keystore, req ks.EncryptRequest) (ks.EncryptResponse, error) {
 				return k.Encrypt(ctx, req)
 			})
 		},
@@ -325,7 +377,7 @@ func NewDecryptCmd() *cobra.Command {
 	cmd := cobra.Command{
 		Use: "decrypt", Short: "Decrypt data with a key",
 		RunE: func(cmd *cobra.Command, args []string) error {
-			return runKeystoreCommand[ks.DecryptRequest, ks.DecryptResponse](cmd, args, func(ctx context.Context, k ks.Keystore, req ks.DecryptRequest) (ks.DecryptResponse, error) {
+			return runKeystoreCommand[ks.Keystore, ks.DecryptRequest, ks.DecryptResponse](cmd, args, loadKeystore, func(ctx context.Context, k ks.Keystore, req ks.DecryptRequest) (ks.DecryptResponse, error) {
 				return k.Decrypt(ctx, req)
 			})
 		},
@@ -335,38 +387,27 @@ func NewDecryptCmd() *cobra.Command {
 	return &cmd
 }
 
-func loadKeystore(ctx context.Context, cmd *cobra.Command) (ks.Keystore, error) {
-	// Use parent command which has the persistent flags.
-	// This works whether keystore CLI is standalone or embedded as a subcommand.
-	parent := cmd.Parent()
-	filePath, err := parent.Flags().GetString("file-path")
-	if err != nil {
-		return nil, err
-	}
-	dbURL, err := parent.Flags().GetString("db-url")
-	if err != nil {
-		return nil, err
-	}
-	password, err := parent.Flags().GetString("password")
-	if err != nil {
-		return nil, err
-	}
-	// Env var fallbacks (flags override)
-	if filePath == "" {
-		if v := os.Getenv("KEYSTORE_FILE_PATH"); v != "" {
-			filePath = v
-		}
-	}
-	if dbURL == "" {
-		if v := os.Getenv("KEYSTORE_DB_URL"); v != "" {
-			dbURL = v
-		}
-	}
-	if password == "" {
-		if v := os.Getenv("KEYSTORE_PASSWORD"); v != "" {
-			password = v
+func loadKeystoreSignerReader(ctx context.Context, cmd *cobra.Command) (interface {
+	ks.Reader
+	ks.Signer
+}, error) {
+	// Check if KMS mode is enabled
+	kmsProfile := os.Getenv("KEYSTORE_KMS_PROFILE")
+	if kmsProfile != "" {
+		client, err := kms.NewClient(kmsProfile)
+		if err != nil {
+			return nil, fmt.Errorf("create KMS client: %w", err)
 		}
+		return kms.NewKeystore(client)
 	}
+	return loadKeystore(ctx, cmd)
+}
+
+func loadKeystore(ctx context.Context, cmd *cobra.Command) (ks.Keystore, error) {
+	// Read from environment variables only
+	filePath := os.Getenv("KEYSTORE_FILE_PATH")
+	dbURL := os.Getenv("KEYSTORE_DB_URL")
+	password := os.Getenv("KEYSTORE_PASSWORD")
 	if password == "" {
 		return nil, errors.New("keystore password is required")
 	}

keystore/cli/cli_test.go

@@ -22,6 +22,8 @@ func setupKeystore(t *testing.T) func(t *testing.T) {
 	require.NoError(t, err)
 	t.Setenv("KEYSTORE_FILE_PATH", keystoreFile)
 	t.Setenv("KEYSTORE_PASSWORD", "testpassword")
+	// Set to empty string to test regular keystore mode.
+	t.Setenv("KEYSTORE_KMS_PROFILE", "")
 	return func(t *testing.T) {
 		f.Close()
 		os.RemoveAll(tempDir)

keystore/go.mod

@@ -3,6 +3,7 @@ module github.com/smartcontractkit/chainlink-common/keystore
 go 1.25.3
 
 require (
+	github.com/aws/aws-sdk-go v1.55.5
 	github.com/ethereum/go-ethereum v1.16.2
 	github.com/jmoiron/sqlx v1.4.0
 	github.com/lib/pq v1.10.9
@@ -58,6 +59,7 @@ require (
 	github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
 	github.com/jackc/pgtype v1.14.4 // indirect
 	github.com/jackc/pgx/v4 v4.18.3 // indirect
+	github.com/jmespath/go-jmespath v0.4.0 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
 	github.com/klauspost/compress v1.18.0 // indirect
 	github.com/klauspost/cpuid/v2 v2.2.10 // indirect

keystore/go.sum

@@ -14,6 +14,8 @@ github.com/apache/arrow-go/v18 v18.3.1 h1:oYZT8FqONiK74JhlH3WKVv+2NKYoyZ7C2ioD4D
 github.com/apache/arrow-go/v18 v18.3.1/go.mod h1:12QBya5JZT6PnBihi5NJTzbACrDGXYkrgjujz3MRQXU=
 github.com/apache/thrift v0.21.0 h1:tdPmh/ptjE1IJnhbhrcl2++TauVjy242rkV/UzJChnE=
 github.com/apache/thrift v0.21.0/go.mod h1:W1H8aR/QRtYNvrPeFXBtobyRkd0/YVhTc6i07XIAgDw=
+github.com/aws/aws-sdk-go v1.55.5 h1:KKUZBfBoyqy5d3swXyiC7Q76ic40rYcbqH7qjh59kzU=
+github.com/aws/aws-sdk-go v1.55.5/go.mod h1:eRwEWoyTWFMVYVQzKMNHWP5/RV4xIUGMQfXQHfHkpNU=
 github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk=
 github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
@@ -171,6 +173,10 @@ github.com/jackc/puddle v0.0.0-20190413234325-e4ced69a3a2b/go.mod h1:m4B5Dj62Y0f
 github.com/jackc/puddle v0.0.0-20190608224051-11cab39313c9/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk=
 github.com/jackc/puddle v1.1.3/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk=
 github.com/jackc/puddle v1.3.0/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk=
+github.com/jmespath/go-jmespath v0.4.0 h1:BEgLn5cpjn8UN1mAw4NjwDrS35OdebyEtFe+9YPoQUg=
+github.com/jmespath/go-jmespath v0.4.0/go.mod h1:T8mJZnbsbmF+m6zOOFylbeCJqk5+pHWvzYPziyZiYoo=
+github.com/jmespath/go-jmespath/internal/testify v1.5.1 h1:shLQSRRSCCPj3f2gpwzGwWFoC7ycTf1rcQZHOlsJ6N8=
+github.com/jmespath/go-jmespath/internal/testify v1.5.1/go.mod h1:L3OGu8Wl2/fWfCI6z80xFu9LTZmf1ZRjMHUOPmWr69U=
 github.com/jmoiron/sqlx v1.4.0 h1:1PLqN7S1UYp5t4SrVVnt4nUVNemrDAtxlulVe+Qgm3o=
 github.com/jmoiron/sqlx v1.4.0/go.mod h1:ZrZ7UsYB/weZdl2Bxg6jCRO9c3YHl8r3ahlKmRT4JLY=
 github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
@@ -505,6 +511,7 @@ gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EV
 gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI=
 gopkg.in/inconshreveable/log15.v2 v2.0.0-20180818164646-67afb5ed74ec/go.mod h1:aPpfJ7XW+gOuirDoZ8gHhLh3kZ1B08FtV2bbmy7Jv3s=
 gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
 gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
 gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=