[v14] MemoryDB IAM auth support

lib/configurators/aws/aws.go

@@ -252,6 +252,8 @@ var (
 			"memorydb:UpdateUser",
 		},
 		requireSecretsManager: true,
+		authBoundary:          []string{"memorydb:Connect"},
+		requireIAMEdit:        true,
 	}
 	// awsKeyspacesActions contains IAM actions for static AWS Keyspaces databases.
 	awsKeyspacesActions = databaseActions{

lib/configurators/aws/aws_test.go

@@ -432,6 +432,9 @@ func TestAWSIAMDocuments(t *testing.T) {
 					},
 					Resources: []string{"arn:aws:secretsmanager:*:123456789012:secret:teleport/*"},
 				},
+				{Effect: awslib.EffectAllow, Resources: []string{roleTarget.String()}, Actions: []string{
+					"iam:GetRolePolicy", "iam:PutRolePolicy", "iam:DeleteRolePolicy",
+				}},
 			},
 			boundaryStatements: []*awslib.Statement{
 				{Effect: awslib.EffectAllow, Resources: []string{"*"}, Actions: []string{
@@ -440,6 +443,7 @@ func TestAWSIAMDocuments(t *testing.T) {
 					"memorydb:DescribeSubnetGroups",
 					"memorydb:DescribeUsers",
 					"memorydb:UpdateUser",
+					"memorydb:Connect",
 				}},
 				{
 					Effect: awslib.EffectAllow,
@@ -451,6 +455,9 @@ func TestAWSIAMDocuments(t *testing.T) {
 					},
 					Resources: []string{"arn:aws:secretsmanager:*:123456789012:secret:teleport/*"},
 				},
+				{Effect: awslib.EffectAllow, Resources: []string{roleTarget.String()}, Actions: []string{
+					"iam:GetRolePolicy", "iam:PutRolePolicy", "iam:DeleteRolePolicy",
+				}},
 			},
 		},
 		"MemoryDB static database": {
@@ -506,6 +513,9 @@ func TestAWSIAMDocuments(t *testing.T) {
 						"arn:aws:kms:*:123456789012:key/my-kms-id",
 					},
 				},
+				{Effect: awslib.EffectAllow, Resources: []string{roleTarget.String()}, Actions: []string{
+					"iam:GetRolePolicy", "iam:PutRolePolicy", "iam:DeleteRolePolicy",
+				}},
 			},
 			boundaryStatements: []*awslib.Statement{
 				{Effect: awslib.EffectAllow, Resources: []string{"*"}, Actions: []string{
@@ -514,6 +524,7 @@ func TestAWSIAMDocuments(t *testing.T) {
 					"memorydb:DescribeSubnetGroups",
 					"memorydb:DescribeUsers",
 					"memorydb:UpdateUser",
+					"memorydb:Connect",
 				}},
 				{
 					Effect: "Allow",
@@ -535,6 +546,9 @@ func TestAWSIAMDocuments(t *testing.T) {
 						"arn:aws:kms:*:123456789012:key/my-kms-id",
 					},
 				},
+				{Effect: awslib.EffectAllow, Resources: []string{roleTarget.String()}, Actions: []string{
+					"iam:GetRolePolicy", "iam:PutRolePolicy", "iam:DeleteRolePolicy",
+				}},
 			},
 		},
 		"AutoDiscovery EC2": {
@@ -1339,12 +1353,17 @@ func TestAWSIAMDocuments(t *testing.T) {
 						},
 						Resources: []string{"arn:aws:secretsmanager:*:123456789012:secret:teleport/*"},
 					},
+					{
+						Effect:    awslib.EffectAllow,
+						Resources: []string{roleTarget.String()},
+						Actions:   []string{"iam:GetRolePolicy", "iam:PutRolePolicy", "iam:DeleteRolePolicy"},
+					},
 				},
 				boundaryStatements: []*awslib.Statement{
 					{
 						Effect:    awslib.EffectAllow,
 						Resources: []string{"*"},
-						Actions:   []string{"memorydb:DescribeClusters", "memorydb:DescribeUsers", "memorydb:UpdateUser"},
+						Actions:   []string{"memorydb:DescribeClusters", "memorydb:DescribeUsers", "memorydb:UpdateUser", "memorydb:Connect"},
 					},
 					{
 						Effect: awslib.EffectAllow,
@@ -1356,6 +1375,11 @@ func TestAWSIAMDocuments(t *testing.T) {
 						},
 						Resources: []string{"arn:aws:secretsmanager:*:123456789012:secret:teleport/*"},
 					},
+					{
+						Effect:    awslib.EffectAllow,
+						Resources: []string{roleTarget.String()},
+						Actions:   []string{"iam:GetRolePolicy", "iam:PutRolePolicy", "iam:DeleteRolePolicy"},
+					},
 				},
 			},
 			"OpenSearch": {

lib/srv/db/access_test.go

@@ -2242,6 +2242,8 @@ type agentParams struct {
 	GCPSQL *mocks.GCPSQLAdminClientMock
 	// ElastiCache defines the AWS ElastiCache mock to use for ElastiCache API calls.
 	ElastiCache *mocks.ElastiCacheMock
+	// MemoryDB defines the AWS MemoryDB mock to use for MemoryDB API calls.
+	MemoryDB *mocks.MemoryDBMock
 	// OnHeartbeat defines a heartbeat function that generates heartbeat events.
 	OnHeartbeat func(error)
 	// CADownloader defines the CA downloader.
@@ -2275,6 +2277,9 @@ func (p *agentParams) setDefaults(c *testContext) {
 	if p.ElastiCache == nil {
 		p.ElastiCache = &mocks.ElastiCacheMock{}
 	}
+	if p.MemoryDB == nil {
+		p.MemoryDB = &mocks.MemoryDBMock{}
+	}
 	if p.CADownloader == nil {
 		p.CADownloader = &fakeDownloader{
 			cert: []byte(fixtures.TLSCACertPEM),
@@ -2288,7 +2293,7 @@ func (p *agentParams) setDefaults(c *testContext) {
 			Redshift:           &mocks.RedshiftMock{},
 			RedshiftServerless: &mocks.RedshiftServerlessMock{},
 			ElastiCache:        p.ElastiCache,
-			MemoryDB:           &mocks.MemoryDBMock{},
+			MemoryDB:           p.MemoryDB,
 			SecretsManager:     secrets.NewMockSecretsManagerClient(secrets.MockSecretsManagerClientConfig{}),
 			IAM:                &mocks.IAMMock{},
 			GCPSQL:             p.GCPSQL,
@@ -3078,6 +3083,43 @@ func withElastiCacheRedis(name string, token, engineVersion string) withDatabase
 	}
 }
 
+func withMemoryDBRedis(name string, token, engineVersion string) withDatabaseOption {
+	return func(t testing.TB, ctx context.Context, testCtx *testContext) types.Database {
+		redisServer, err := redis.NewTestServer(t, common.TestServerConfig{
+			Name:       name,
+			AuthClient: testCtx.authClient,
+		}, redis.TestServerPassword(token))
+		require.NoError(t, err)
+
+		database, err := types.NewDatabaseV3(types.Metadata{
+			Name: name,
+			Labels: map[string]string{
+				"engine-version": engineVersion,
+			},
+		}, types.DatabaseSpecV3{
+			Protocol:      defaults.ProtocolRedis,
+			URI:           fmt.Sprintf("rediss://%s", net.JoinHostPort("localhost", redisServer.Port())),
+			DynamicLabels: dynamicLabels,
+			AWS: types.AWS{
+				Region: "us-west-1",
+				MemoryDB: types.MemoryDB{
+					ClusterName: "example-cluster",
+				},
+			},
+			// Set CA cert to pass cert validation.
+			TLS: types.DatabaseTLS{
+				CACert: string(testCtx.databaseCA.GetActiveKeys().TLS[0].Cert),
+			},
+		})
+		require.NoError(t, err)
+		testCtx.redis[name] = testRedis{
+			db:       redisServer,
+			resource: database,
+		}
+		return database
+	}
+}
+
 func withAzureRedis(name string, token string) withDatabaseOption {
 	return func(t testing.TB, ctx context.Context, testCtx *testContext) types.Database {
 		redisServer, err := redis.NewTestServer(t, common.TestServerConfig{

lib/srv/db/auth_test.go

@@ -25,6 +25,7 @@ import (
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/elasticache"
+	"github.com/aws/aws-sdk-go/service/memorydb"
 	"github.com/gravitational/trace"
 	"github.com/jonboulle/clockwork"
 	"github.com/sirupsen/logrus"
@@ -64,6 +65,8 @@ func TestAuthTokens(t *testing.T) {
 		withAzureRedis("redis-azure-incorrect-token", "qwe123"),
 		withElastiCacheRedis("redis-elasticache-correct-token", elastiCacheRedisToken, "7.0.0"),
 		withElastiCacheRedis("redis-elasticache-incorrect-token", "qwe123", "7.0.0"),
+		withMemoryDBRedis("redis-memorydb-correct-token", memorydbToken, "7.0"),
+		withMemoryDBRedis("redis-memorydb-incorrect-token", "qwe123", "7.0"),
 	}
 	databases := make([]types.Database, 0, len(withDBs))
 	for _, withDB := range withDBs {
@@ -75,9 +78,16 @@ func TestAuthTokens(t *testing.T) {
 		Authentication: &elasticache.Authentication{Type: aws.String("iam")},
 	}
 	ecMock.AddMockUser(elastiCacheIAMUser, nil)
+	memorydbMock := &mocks.MemoryDBMock{}
+	memorydbIAMUser := &memorydb.User{
+		Name:           aws.String("default"),
+		Authentication: &memorydb.Authentication{Type: aws.String("iam")},
+	}
+	memorydbMock.AddMockUser(memorydbIAMUser, nil)
 	testCtx.server = testCtx.setupDatabaseServer(ctx, t, agentParams{
 		Databases:   databases,
 		ElastiCache: ecMock,
+		MemoryDB:    memorydbMock,
 	})
 	go testCtx.startHandlingConnections()
 
@@ -169,6 +179,18 @@ func TestAuthTokens(t *testing.T) {
 			// Make sure we print a user-friendly IAM auth error.
 			err: "Make sure that IAM auth is enabled",
 		},
+		{
+			desc:     "correct MemoryDB auth token",
+			service:  "redis-memorydb-correct-token",
+			protocol: defaults.ProtocolRedis,
+		},
+		{
+			desc:     "incorrect MemoryDB auth token",
+			service:  "redis-memorydb-incorrect-token",
+			protocol: defaults.ProtocolRedis,
+			// Make sure we print a user-friendly IAM auth error.
+			err: "Make sure that IAM auth is enabled",
+		},
 	}
 
 	for _, test := range tests {
@@ -245,6 +267,8 @@ const (
 	azureRedisToken = "azure-redis-token"
 	// elastiCacheRedisToken is a mock ElastiCache Redis token.
 	elastiCacheRedisToken = "elasticache-redis-token"
+	// memorydbToken is a mock MemoryDB auth token.
+	memorydbToken = "memorydb-token"
 	// atlasAuthUser is a mock Mongo Atlas IAM auth user.
 	atlasAuthUser = "arn:aws:iam::111111111111:role/alice"
 	// atlasAuthToken is a mock Mongo Atlas IAM auth token.
@@ -273,6 +297,10 @@ func (a *testAuth) GetElastiCacheRedisToken(ctx context.Context, sessionCtx *com
 	return elastiCacheRedisToken, nil
 }
 
+func (a *testAuth) GetMemoryDBToken(ctx context.Context, sessionCtx *common.Session) (string, error) {
+	return memorydbToken, nil
+}
+
 // GetCloudSQLAuthToken generates Cloud SQL auth token.
 func (a *testAuth) GetCloudSQLAuthToken(ctx context.Context, sessionCtx *common.Session) (string, error) {
 	a.Infof("Generating Cloud SQL auth token for %v.", sessionCtx)

lib/srv/db/cloud/iam.go

@@ -198,6 +198,14 @@ func (c *IAM) isSetupRequiredForDatabase(database types.Database) bool {
 			return true
 		}
 		return ok
+	case types.DatabaseTypeMemoryDB:
+		ok, err := iam.CheckMemoryDBSupportsIAMAuth(database)
+		if err != nil {
+			c.log.WithError(err).Debugf("Assuming database %s supports IAM auth.",
+				database.GetName())
+			return true
+		}
+		return ok
 	default:
 		return false
 	}

lib/srv/db/cloud/iam_test.go

@@ -129,6 +129,22 @@ func TestAWSIAM(t *testing.T) {
 	})
 	require.NoError(t, err)
 
+	memorydb, err := types.NewDatabaseV3(types.Metadata{
+		Name: "aws-memorydb",
+	}, types.DatabaseSpecV3{
+		Protocol: "redis",
+		URI:      "clustercfg.my-memorydb.xxxxxx.memorydb.us-east-1.amazonaws.com:6379",
+		AWS: types.AWS{
+			AccountID: "123456789012",
+			MemoryDB: types.MemoryDB{
+				ClusterName:  "my-memorydb",
+				TLSEnabled:   true,
+				EndpointType: "cluster",
+			},
+		},
+	})
+	require.NoError(t, err)
+
 	// Make configurator.
 	taskChan := make(chan struct{})
 	waitForTaskProcessed := func(t *testing.T) {
@@ -207,6 +223,13 @@ func TestAWSIAM(t *testing.T) {
 				return true // it always is for ElastiCache.
 			},
 		},
+		"MemoryDB": {
+			database:           memorydb,
+			wantPolicyContains: memorydb.GetAWS().MemoryDB.ClusterName,
+			getIAMAuthEnabled: func() bool {
+				return true // it always is for MemoryDB.
+			},
+		},
 	}
 
 	for testName, tt := range tests {