[v16] kms: allow multi-region kms keys to be enabled (#47672) (#48062)

* kms: allow multi-region kms keys to be enabled * add unit tests * fix logging key lint --------- Co-authored-by: Nic Klaassen <[email protected]>
gravitational · Oct 30, 2024 · 9c66230 · 9c66230
1 parent bf31ed6
commit 9c66230
Show file tree

Hide file tree

Showing 6 changed files with 140 additions and 47 deletions.
diff --git a/lib/auth/keystore/aws_kms.go b/lib/auth/keystore/aws_kms.go
@@ -63,12 +63,13 @@ type CloudClientProvider interface {
 }
 
 type awsKMSKeystore struct {
-	kms         kmsiface.KMSAPI
-	clusterName types.ClusterName
-	awsAccount  string
-	awsRegion   string
-	clock       clockwork.Clock
-	logger      *slog.Logger
+	kms                kmsiface.KMSAPI
+	clusterName        types.ClusterName
+	awsAccount         string
+	awsRegion          string
+	multiRegionEnabled bool
+	clock              clockwork.Clock
+	logger             *slog.Logger
 }
 
 func newAWSKMSKeystore(ctx context.Context, cfg *servicecfg.AWSKMSConfig, opts *Options) (*awsKMSKeystore, error) {
@@ -93,12 +94,13 @@ func newAWSKMSKeystore(ctx context.Context, cfg *servicecfg.AWSKMSConfig, opts *
 		clock = clockwork.NewRealClock()
 	}
 	return &awsKMSKeystore{
-		clusterName: opts.ClusterName,
-		awsAccount:  cfg.AWSAccount,
-		awsRegion:   cfg.AWSRegion,
-		kms:         kmsClient,
-		clock:       clock,
-		logger:      opts.Logger,
+		clusterName:        opts.ClusterName,
+		awsAccount:         cfg.AWSAccount,
+		awsRegion:          cfg.AWSRegion,
+		multiRegionEnabled: cfg.MultiRegion.Enabled,
+		kms:                kmsClient,
+		clock:              clock,
+		logger:             opts.Logger,
 	}, nil
 }
 
@@ -125,6 +127,7 @@ func (a *awsKMSKeystore) generateRSA(ctx context.Context, _ ...rsaKeyOption) ([]
 				TagValue: aws.String(a.clusterName.GetClusterName()),
 			},
 		},
+		MultiRegion: aws.Bool(a.multiRegionEnabled),
 	})
 	if err != nil {
 		return nil, nil, trace.Wrap(err)

diff --git a/lib/auth/keystore/aws_kms_test.go b/lib/auth/keystore/aws_kms_test.go
@@ -214,6 +214,61 @@ func TestAWSKMS_RetryWhilePending(t *testing.T) {
 	require.Error(t, err)
 }
 
+// TestMultiRegionKeys asserts that a keystore created with multi-region enabled
+// correctly passes this argument to the AWS client. This gives very little real
+// coverage since the AWS KMS service here is faked, but at least we know the
+// keystore passed the bool to the client correctly. TestBackends and
+// TestManager are both able to run with a real AWS KMS client and you can
+// confirm the keys are really multi-region there.
+func TestMultiRegionKeys(t *testing.T) {
+	ctx := context.Background()
+	clock := clockwork.NewFakeClock()
+
+	const pageSize int = 4
+	fakeKMS := newFakeAWSKMSService(t, clock, "123456789012", "us-west-2", pageSize)
+	clusterName, err := services.NewClusterNameWithRandomID(types.ClusterNameSpecV2{ClusterName: "test-cluster"})
+	require.NoError(t, err)
+	opts := &Options{
+		ClusterName: clusterName,
+		HostUUID:    "uuid",
+		CloudClients: &cloud.TestCloudClients{
+			KMS: fakeKMS,
+			STS: &fakeAWSSTSClient{
+				account: "123456789012",
+			},
+		},
+		clockworkOverride: clock,
+	}
+
+	for _, multiRegion := range []bool{false, true} {
+		t.Run(fmt.Sprint(multiRegion), func(t *testing.T) {
+			cfg := servicecfg.KeystoreConfig{
+				AWSKMS: servicecfg.AWSKMSConfig{
+					AWSAccount: "123456789012",
+					AWSRegion:  "us-west-2",
+					MultiRegion: struct{ Enabled bool }{
+						Enabled: multiRegion,
+					},
+				},
+			}
+			keyStore, err := NewManager(ctx, &cfg, opts)
+			require.NoError(t, err)
+
+			sshKeyPair, err := keyStore.NewSSHKeyPair(ctx)
+			require.NoError(t, err)
+
+			keyID, err := parseAWSKMSKeyID(sshKeyPair.PrivateKey)
+			require.NoError(t, err)
+
+			if multiRegion {
+				assert.Contains(t, keyID.arn, "mrk-")
+			} else {
+				assert.NotContains(t, keyID.arn, "mrk-")
+			}
+		})
+	}
+}
+
 type fakeAWSKMSService struct {
 	kmsiface.KMSAPI
 
@@ -243,6 +298,10 @@ type fakeAWSKMSKey struct {
 
 func (f *fakeAWSKMSService) CreateKey(input *kms.CreateKeyInput) (*kms.CreateKeyOutput, error) {
 	id := uuid.NewString()
+	if aws.BoolValue(input.MultiRegion) {
+		// AWS does this https://docs.aws.amazon.com/kms/latest/developerguide/concepts.html#key-id-key-ARN
+		id = "mrk-" + id
+	}
 	a := arn.ARN{
 		Partition: "aws",
 		Service:   "kms",

diff --git a/lib/auth/keystore/keystore_test.go b/lib/auth/keystore/keystore_test.go
@@ -486,6 +486,7 @@ func newTestPack(ctx context.Context, t *testing.T) *testPack {
 		})
 	}
 
+	// Test with real GCP account if environment is enabled.
 	if config, ok := gcpKMSTestConfig(t); ok {
 		backend, err := newGCPKMSKeyStore(ctx, &config.GCPKMS, opts)
 		require.NoError(t, err)
@@ -499,6 +500,7 @@ func newTestPack(ctx context.Context, t *testing.T) *testPack {
 			}.marshal(),
 		})
 	}
+	// Always test with fake GCP client.
 	fakeGCPKMSConfig := servicecfg.KeystoreConfig{
 		GCPKMS: servicecfg.GCPKMSConfig{
 			ProtectionLevel: "HSM",
@@ -517,54 +519,72 @@ func newTestPack(ctx context.Context, t *testing.T) *testPack {
 		}.marshal(),
 	})
 
-	if config, ok := awsKMSTestConfig(t); ok {
-		backend, err := newAWSKMSKeystore(ctx, &config.AWSKMS, opts)
+	// Test AWS with and without multi-region keys
+	for _, multiRegion := range []bool{false, true} {
+		// Test with real AWS account if environment is enabled.
+		if config, ok := awsKMSTestConfig(t); ok {
+			config.AWSKMS.MultiRegion.Enabled = multiRegion
+
+			backend, err := newAWSKMSKeystore(ctx, &config.AWSKMS, opts)
+			require.NoError(t, err)
+			name := "aws_kms"
+			if multiRegion {
+				name += "_multi_region"
+			}
+			backends = append(backends, &backendDesc{
+				name:            name,
+				config:          config,
+				backend:         backend,
+				expectedKeyType: types.PrivateKeyType_AWS_KMS,
+				unusedRawKey: awsKMSKeyID{
+					arn: arn.ARN{
+						Partition: "aws",
+						Service:   "kms",
+						Region:    config.AWSKMS.AWSRegion,
+						AccountID: config.AWSKMS.AWSAccount,
+						Resource:  "unused",
+					}.String(),
+					account: config.AWSKMS.AWSAccount,
+					region:  config.AWSKMS.AWSRegion,
+				}.marshal(),
+			})
+		}
+
+		// Always test with fake AWS client.
+		fakeAWSKMSConfig := servicecfg.KeystoreConfig{
+			AWSKMS: servicecfg.AWSKMSConfig{
+				AWSAccount: "123456789012",
+				AWSRegion:  "us-west-2",
+				MultiRegion: struct{ Enabled bool }{
+					Enabled: multiRegion,
+				},
+			},
+		}
+		fakeAWSKMSBackend, err := newAWSKMSKeystore(ctx, &fakeAWSKMSConfig.AWSKMS, opts)
 		require.NoError(t, err)
+		name := "fake_aws_kms"
+		if multiRegion {
+			name += "_multi_region"
+		}
 		backends = append(backends, &backendDesc{
-			name:            "aws_kms",
-			config:          config,
-			backend:         backend,
+			name:            name,
+			config:          fakeAWSKMSConfig,
+			backend:         fakeAWSKMSBackend,
 			expectedKeyType: types.PrivateKeyType_AWS_KMS,
 			unusedRawKey: awsKMSKeyID{
 				arn: arn.ARN{
 					Partition: "aws",
 					Service:   "kms",
-					Region:    config.AWSKMS.AWSRegion,
-					AccountID: config.AWSKMS.AWSAccount,
+					Region:    "us-west-2",
+					AccountID: "123456789012",
 					Resource:  "unused",
 				}.String(),
-				account: config.AWSKMS.AWSAccount,
-				region:  config.AWSKMS.AWSRegion,
+				account: "123456789012",
+				region:  "us-west-2",
 			}.marshal(),
 		})
 	}
 
-	fakeAWSKMSConfig := servicecfg.KeystoreConfig{
-		AWSKMS: servicecfg.AWSKMSConfig{
-			AWSAccount: "123456789012",
-			AWSRegion:  "us-west-2",
-		},
-	}
-	fakeAWSKMSBackend, err := newAWSKMSKeystore(ctx, &fakeAWSKMSConfig.AWSKMS, opts)
-	require.NoError(t, err)
-	backends = append(backends, &backendDesc{
-		name:            "fake_aws_kms",
-		config:          fakeAWSKMSConfig,
-		backend:         fakeAWSKMSBackend,
-		expectedKeyType: types.PrivateKeyType_AWS_KMS,
-		unusedRawKey: awsKMSKeyID{
-			arn: arn.ARN{
-				Partition: "aws",
-				Service:   "kms",
-				Region:    "us-west-2",
-				AccountID: "123456789012",
-				Resource:  "unused",
-			}.String(),
-			account: "123456789012",
-			region:  "us-west-2",
-		}.marshal(),
-	})
-
 	return &testPack{
 		opts:     opts,
 		backends: backends,

diff --git a/lib/config/configuration.go b/lib/config/configuration.go
@@ -1170,6 +1170,7 @@ func applyAWSKMSConfig(kmsConfig *AWSKMS, cfg *servicecfg.Config) error {
 		return trace.BadParameter("must set region in ca_key_params.aws_kms")
 	}
 	cfg.Auth.KeyStore.AWSKMS.AWSRegion = kmsConfig.Region
+	cfg.Auth.KeyStore.AWSKMS.MultiRegion = kmsConfig.MultiRegion
 	return nil
 }
 

diff --git a/lib/config/fileconf.go b/lib/config/fileconf.go
@@ -917,6 +917,11 @@ type AWSKMS struct {
 	Account string `yaml:"account"`
 	// Region is the AWS region to use.
 	Region string `yaml:"region"`
+	// MultiRegion contains configuration for multi-region AWS KMS.
+	MultiRegion struct {
+		// Enabled configures new keys to be multi-region.
+		Enabled bool
+	} `yaml:"multi_region,omitempty"`
 }
 
 // TrustedCluster struct holds configuration values under "trusted_clusters" key

diff --git a/lib/service/servicecfg/auth.go b/lib/service/servicecfg/auth.go
@@ -278,6 +278,11 @@ type AWSKMSConfig struct {
 	AWSAccount string
 	// AWSRegion is the AWS region where the keys will reside.
 	AWSRegion string
+	// MultiRegion contains configuration for multi-region AWS KMS.
+	MultiRegion struct {
+		// Enabled configures new keys to be multi-region.
+		Enabled bool
+	}
 }
 
 // CheckAndSetDefaults checks that required parameters of the config are