Refactor database DiscoveryResourceChecker

lib/srv/db/access_test.go

@@ -71,6 +71,7 @@ import (
 	alpncommon "github.com/gravitational/teleport/lib/srv/alpnproxy/common"
 	"github.com/gravitational/teleport/lib/srv/db/cassandra"
 	"github.com/gravitational/teleport/lib/srv/db/clickhouse"
+	"github.com/gravitational/teleport/lib/srv/db/cloud"
 	"github.com/gravitational/teleport/lib/srv/db/common"
 	"github.com/gravitational/teleport/lib/srv/db/dynamodb"
 	"github.com/gravitational/teleport/lib/srv/db/elasticsearch"
@@ -2203,6 +2204,9 @@ type agentParams struct {
 	AWSMatchers []types.AWSMatcher
 	// AzureMatchers is a list of Azure databases matchers.
 	AzureMatchers []types.AzureMatcher
+	// discoveryResourceChecker performs some pre-checks when creating databases
+	// discovered by the discovery service.
+	DiscoveryResourceChecker cloud.DiscoveryResourceChecker
 }
 
 func (p *agentParams) setDefaults(c *testContext) {
@@ -2242,6 +2246,10 @@ func (p *agentParams) setDefaults(c *testContext) {
 			GCPSQL:             p.GCPSQL,
 		}
 	}
+
+	if p.DiscoveryResourceChecker == nil {
+		p.DiscoveryResourceChecker = &fakeDiscoveryResourceChecker{}
+	}
 }
 
 func (c *testContext) setupDatabaseServer(ctx context.Context, t *testing.T, p agentParams) *Server {
@@ -2325,7 +2333,7 @@ func (c *testContext) setupDatabaseServer(ctx context.Context, t *testing.T, p a
 		AWSMatchers:              p.AWSMatchers,
 		AzureMatchers:            p.AzureMatchers,
 		ShutdownPollPeriod:       100 * time.Millisecond,
-		discoveryResourceChecker: &fakeDiscoveryResourceChecker{},
+		discoveryResourceChecker: p.DiscoveryResourceChecker,
 	})
 	require.NoError(t, err)
 
@@ -2989,9 +2997,15 @@ func withAzureRedis(name string, token string) withDatabaseOption {
 	}
 }
 
-type fakeDiscoveryResourceChecker struct{}
+type fakeDiscoveryResourceChecker struct {
+	errorsByName map[string]error
+}
 
-func (f fakeDiscoveryResourceChecker) check(_ context.Context, _ types.Database) {
+func (f *fakeDiscoveryResourceChecker) Check(_ context.Context, database types.Database) error {
+	if len(f.errorsByName) == 0 {
+		return nil
+	}
+	return trace.Wrap(f.errorsByName[database.GetName()])
 }
 
 var dynamicLabels = types.LabelsToV2(map[string]types.CommandLabel{

lib/srv/db/cloud/resource_checker.go

@@ -0,0 +1,103 @@
+/*
+Copyright 2023 Gravitational, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package cloud
+
+import (
+	"context"
+
+	"github.com/gravitational/trace"
+	"github.com/sirupsen/logrus"
+
+	"github.com/gravitational/teleport"
+	"github.com/gravitational/teleport/api/types"
+	"github.com/gravitational/teleport/lib/cloud"
+	"github.com/gravitational/teleport/lib/services"
+)
+
+// DiscoveryResourceChecker defines an interface for checking database
+// resources created by the discovery service.
+type DiscoveryResourceChecker interface {
+	// Check performs required checks on provided database resource before it
+	// gets registered.
+	Check(ctx context.Context, database types.Database) error
+}
+
+// DiscoveryResourceCheckerConfig is the config for DiscoveryResourceChecker.
+type DiscoveryResourceCheckerConfig struct {
+	// ResourceMatchers is a list of database resource matchers.
+	ResourceMatchers []services.ResourceMatcher
+	// Clients is an interface for retrieving cloud clients.
+	Clients cloud.Clients
+	// Context is the database server close context.
+	Context context.Context
+	// Log is used for logging.
+	Log logrus.FieldLogger
+}
+
+// CheckAndSetDefaults validates the config and sets default values.
+func (c *DiscoveryResourceCheckerConfig) CheckAndSetDefaults() error {
+	if c.Clients == nil {
+		cloudClients, err := cloud.NewClients()
+		if err != nil {
+			return trace.Wrap(err)
+		}
+		c.Clients = cloudClients
+	}
+	if c.Context == nil {
+		c.Context = context.Background()
+	}
+	if c.Log == nil {
+		c.Log = logrus.WithField(trace.Component, teleport.ComponentDatabase)
+	}
+	return nil
+}
+
+// NewDiscoveryResourceChecker creates a new DiscoveryResourceChecker.
+func NewDiscoveryResourceChecker(cfg DiscoveryResourceCheckerConfig) (DiscoveryResourceChecker, error) {
+	if err := cfg.CheckAndSetDefaults(); err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	c := &discoveryResourceChecker{}
+
+	// TODO(greedy52) implement url checker.
+	// TODO(greedy52) implement name checker.
+	if checker, err := newCrednentialsChecker(cfg); err != nil {
+		return nil, trace.Wrap(err)
+	} else {
+		c.checkers = append(c.checkers, checker)
+	}
+	return c, nil
+}
+
+// discoveryResourceChecker is a composite checker.
+type discoveryResourceChecker struct {
+	checkers []DiscoveryResourceChecker
+}
+
+// Check calls Check from all its checkers and aggregate the errors.
+func (c *discoveryResourceChecker) Check(ctx context.Context, database types.Database) error {
+	if database.Origin() != types.OriginCloud {
+		return nil
+	}
+
+	errors := make([]error, 0, len(c.checkers))
+	for _, checker := range c.checkers {
+		errors = append(errors, trace.Wrap(checker.Check(ctx, database)))
+	}
+	return trace.NewAggregate(errors...)
+}

lib/srv/db/cloud/resource_checker_credentials.go

@@ -0,0 +1,164 @@
+/*
+Copyright 2023 Gravitational, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package cloud
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/arm"
+	"github.com/gravitational/trace"
+	"github.com/sirupsen/logrus"
+	"golang.org/x/exp/slices"
+
+	"github.com/gravitational/teleport/api/types"
+	"github.com/gravitational/teleport/lib/cloud"
+	"github.com/gravitational/teleport/lib/cloud/aws"
+	"github.com/gravitational/teleport/lib/services"
+	"github.com/gravitational/teleport/lib/utils"
+)
+
+// credentialsChecker performs some quick checks to see whether this database
+// agent can handle the incoming database wrt to the agent's credentials.
+//
+// Note that this checker warns the user with suggestions on how to configure
+// the credentials correctly instead of returning errors.
+type credentialsChecker struct {
+	cloudClients     cloud.Clients
+	resourceMatchers []services.ResourceMatcher
+	log              logrus.FieldLogger
+	cache            *utils.FnCache
+}
+
+func newCrednentialsChecker(cfg DiscoveryResourceCheckerConfig) (*credentialsChecker, error) {
+	cache, err := utils.NewFnCache(utils.FnCacheConfig{
+		TTL:     10 * time.Minute,
+		Context: cfg.Context,
+	})
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	return &credentialsChecker{
+		cloudClients:     cfg.Clients,
+		resourceMatchers: cfg.ResourceMatchers,
+		log:              cfg.Log,
+		cache:            cache,
+	}, nil
+}
+
+// Check performs some quick checks to see whether this database agent can
+// handle the incoming database wrt to the agent's credentials.
+func (c *credentialsChecker) Check(ctx context.Context, database types.Database) error {
+	switch {
+	case database.IsAWSHosted():
+		c.checkAWS(ctx, database)
+	case database.IsAzure():
+		c.checkAzure(ctx, database)
+	default:
+		c.log.Debugf("Database %q has unknown cloud type %q.", database.GetName(), database.GetType())
+	}
+	return nil
+}
+
+func (c *credentialsChecker) checkAWS(ctx context.Context, database types.Database) {
+	meta := database.GetAWS()
+	identity, err := c.getAWSIdentity(ctx, &meta)
+	if err != nil {
+		c.warn(err, database, "Failed to get AWS identity when checking a database created by the discovery service.")
+		return
+	}
+
+	if meta.AccountID != "" && meta.AccountID != identity.GetAccountID() {
+		c.warn(nil, database, fmt.Sprintf("The database agent's identity and discovered database %q have different AWS account IDs (%s vs %s).",
+			database.GetName(),
+			identity.GetAccountID(),
+			meta.AccountID,
+		))
+		return
+	}
+}
+
+// getAWSIdentity returns the identity used to access the given database,
+// that is either the agent's identity or the database's configured assume-role.
+func (c *credentialsChecker) getAWSIdentity(ctx context.Context, meta *types.AWS) (aws.Identity, error) {
+	if meta.AssumeRoleARN != "" {
+		// If the database has an assume role ARN, use that instead of
+		// agent identity. This avoids an unnecessary sts call too.
+		return aws.IdentityFromArn(meta.AssumeRoleARN)
+	}
+
+	identity, err := utils.FnCacheGet(ctx, c.cache, types.CloudAWS, func(ctx context.Context) (aws.Identity, error) {
+		client, err := c.cloudClients.GetAWSSTSClient(ctx, "")
+		if err != nil {
+			return nil, trace.Wrap(err)
+		}
+		return aws.GetIdentityWithClient(ctx, client)
+	})
+	return identity, trace.Wrap(err)
+}
+
+func (c *credentialsChecker) checkAzure(ctx context.Context, database types.Database) {
+	allSubIDs, err := utils.FnCacheGet(ctx, c.cache, types.CloudAzure, func(ctx context.Context) ([]string, error) {
+		client, err := c.cloudClients.GetAzureSubscriptionClient()
+		if err != nil {
+			return nil, trace.Wrap(err)
+		}
+		return client.ListSubscriptionIDs(ctx)
+	})
+	if err != nil {
+		c.warn(err, database, "Failed to get Azure subscription IDs when checking a database created by the discovery service.")
+		return
+	}
+
+	rid, err := arm.ParseResourceID(database.GetAzure().ResourceID)
+	if err != nil {
+		c.log.Warnf("Failed to parse resource ID of database %q: %v.", database.GetName(), err)
+		return
+	}
+
+	if !slices.Contains(allSubIDs, rid.SubscriptionID) {
+		c.warn(nil, database, fmt.Sprintf("The discovered database %q is in a subscription (ID: %s) that the database agent does not have access to.",
+			database.GetName(),
+			rid.SubscriptionID,
+		))
+		return
+	}
+}
+
+func (c *credentialsChecker) warn(err error, database types.Database, msg string) {
+	log := c.log.WithField("database", database)
+	if err != nil {
+		log = log.WithField("error", err.Error())
+	}
+
+	logLevel := logrus.InfoLevel
+	if c.isWildcardMatcher() {
+		logLevel = logrus.WarnLevel
+	}
+	log.Logf(logLevel, "%s You can update \"db_service.resources\" section of this agent's config file to filter out unwanted resources (see https://goteleport.com/docs/database-access/reference/configuration/ for more details). If this database is intended to be handled by this agent, please verify that valid cloud credentials are configured for the agent.", msg)
+}
+
+func (c *credentialsChecker) isWildcardMatcher() bool {
+	if len(c.resourceMatchers) != 1 {
+		return false
+	}
+
+	wildcardLabels := c.resourceMatchers[0].Labels[types.Wildcard]
+	return len(wildcardLabels) == 1 && wildcardLabels[0] == types.Wildcard
+}