WIP: IAM auth for db

go.mod

@@ -22,9 +22,12 @@ require (
 	github.com/grpc-ecosystem/go-grpc-middleware v1.2.0
 	github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0
 	github.com/grpc-ecosystem/grpc-gateway v1.14.3
+	github.com/jackc/pgconn v1.6.4
+	github.com/jackc/pgx/v4 v4.8.1
 	github.com/jinzhu/gorm v1.9.12
 	github.com/kelseyhightower/envconfig v1.4.0 // indirect
 	github.com/lib/pq v1.3.0
+	github.com/lyft/datacatalog v0.2.1
 	github.com/lyft/flyteidl v0.18.3
 	github.com/lyft/flytepropeller v0.3.7
 	github.com/lyft/flytestdlib v0.3.9
@@ -41,6 +44,7 @@ require (
 	google.golang.org/genproto v0.0.0-20200331122359-1ee6d9798940
 	google.golang.org/grpc v1.28.0
 	gopkg.in/gormigrate.v1 v1.6.0
+	gorm.io/driver/postgres v1.0.0
 	k8s.io/api v0.17.3
 	k8s.io/apimachinery v0.17.3
 	k8s.io/client-go v11.0.1-0.20190918222721-c0e3722d5cf0+incompatible

pkg/repositories/config/database.go

@@ -8,5 +8,9 @@ type DbConfig struct {
 	DbName       string `json:"dbname"`
 	User         string `json:"user"`
 	Password     string `json:"password"`
+	RootCA       string `json:"rootCA"`
 	ExtraOptions string `json:"options"`
+
+	Region string `json:"region"` // AWS specific
+	UseIAM bool   `json:"useIAM"`
 }

pkg/repositories/config/postgres.go

@@ -1,12 +1,24 @@
 package config
 
 import (
+	"context"
+	"crypto/tls"
+	"crypto/x509"
 	"fmt"
 
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/session"
+	"github.com/aws/aws-sdk-go/service/rds/rdsutils"
+	"github.com/jackc/pgconn"
+	"github.com/jackc/pgx/v4"
+	"github.com/lyft/flyteadmin/pkg/errors"
+	"github.com/lyft/flytestdlib/logger"
 	"github.com/lyft/flytestdlib/promutils"
+	"google.golang.org/grpc/codes"
 
 	"github.com/jinzhu/gorm"
 	_ "github.com/jinzhu/gorm/dialects/postgres" // Required to import database driver.
+	// goGormPostgres "gorm.io/driver/postgres"
 )
 
 const Postgres = "postgres"
@@ -71,11 +83,65 @@ func (p *PostgresConfigProvider) IsDebug() bool {
 
 // Opens a connection to the database specified in the config.
 // You must call CloseDbConnection at the end of your session!
-func OpenDbConnection(config DbConnectionConfigProvider) *gorm.DB {
-	db, err := gorm.Open(config.GetType(), config.GetArgs())
+func OpenDbConnection(ctx context.Context, dbConfig DbConfig) (*gorm.DB, error) {
+	connConfig := pgx.ConnConfig{
+		Config: pgconn.Config{
+			Host:     dbConfig.Host,
+			User:     dbConfig.User,
+			Database: dbConfig.DbName,
+		},
+	}
+
+	if len(dbConfig.RootCA) > 0 {
+		logger.Debugf(ctx, "Preparing to append a root CA")
+		ca := x509.NewCertPool()
+		if ok := ca.AppendCertsFromPEM([]byte(dbConfig.RootCA)); !ok {
+			logger.Errorf(ctx, "Failed to append cert from PEM")
+			return nil, errors.NewFlyteAdminErrorf(codes.Internal, "Failed to append cert from PEM")
+		}
+
+		connConfig.TLSConfig = &tls.Config{
+			RootCAs:    ca,
+			ServerName: dbConfig.Host,
+		}
+	}
+
+	// Use IAM auth when connecting to the database.
+	if dbConfig.UseIAM {
+		logger.Debugf(ctx, "Preparing to use IAM")
+		sess, err := session.NewSession(
+			&aws.Config{
+				Region: aws.String(dbConfig.Region),
+			},
+		)
+		if err != nil {
+			logger.Errorf(ctx, "Failed to create new AWS session with err: [%+v]", err)
+			return nil, errors.NewFlyteAdminErrorf(codes.Internal, "Failed to create new AWS session with err: [%+v]", err)
+		}
+
+		endpoint := fmt.Sprintf("%s:%d", dbConfig.Host, dbConfig.Port)
+		token, err := rdsutils.BuildAuthToken(
+			endpoint,
+			dbConfig.Region,
+			dbConfig.User,
+			sess.Config.Credentials,
+		)
+		if err != nil {
+			logger.Errorf(ctx, "Failed to build auth token with err: [%+v]", err)
+			return nil, errors.NewFlyteAdminErrorf(codes.Internal, "Failed to build auth token with err: [%+v]", err)
+		}
+
+		connConfig.Password = token
+	} else {
+		logger.Debugf(ctx, "Not using IAM")
+		connConfig.Password = dbConfig.Password
+	}
+
+	db, err := gorm.Open("postgres", connConfig.ConnString())
 	if err != nil {
-		panic(err)
+		logger.Errorf(ctx, "Failed to open db connection with err: [%+v]", err)
+		return nil, errors.NewFlyteAdminErrorf(codes.Internal, "Failed to open db connection with err: [%+v]", err)
 	}
-	db.LogMode(config.IsDebug())
-	return db
+	db.LogMode(dbConfig.IsDebug)
+	return db, nil
 }

pkg/repositories/factory.go

@@ -1,6 +1,7 @@
 package repositories
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/lyft/flyteadmin/pkg/repositories/config"
@@ -38,7 +39,11 @@ func GetRepository(repoType RepoConfig, dbConfig config.DbConfig, scope promutil
 	switch repoType {
 	case POSTGRES:
 		postgresScope := scope.NewSubScope("postgres")
-		db := config.OpenDbConnection(config.NewPostgresConfigProvider(dbConfig, postgresScope))
+		// TODO(katrogan): pass a context here?
+		db, err := config.OpenDbConnection(context.Background(), dbConfig)
+		if err != nil {
+			panic(err)
+		}
 		return NewPostgresRepo(
 			db,
 			errors.NewPostgresErrorTransformer(postgresScope.NewSubScope("errors")),

pkg/rpc/adminservice/base.go

@@ -71,7 +71,10 @@ func NewAdminServer(kubeConfig, master string) *AdminService {
 		DbName:       dbConfigValues.DbName,
 		User:         dbConfigValues.User,
 		Password:     dbConfigValues.Password,
+		RootCA:       dbConfigValues.RootCA,
 		ExtraOptions: dbConfigValues.ExtraOptions,
+		Region:       dbConfigValues.Region,
+		UseIAM:       dbConfigValues.UseIAM,
 	}
 	db := repositories.GetRepository(
 		repositories.POSTGRES, dbConfig, adminScope.NewSubScope("database"))

pkg/runtime/application_config_provider.go

@@ -29,25 +29,46 @@ type ApplicationConfigurationProvider struct{}
 
 func (p *ApplicationConfigurationProvider) GetDbConfig() interfaces.DbConfig {
 	dbConfigSection := databaseConfig.GetConfig().(*interfaces.DbConfigSection)
-	password := dbConfigSection.Password
-	if len(dbConfigSection.PasswordPath) > 0 {
-		if _, err := os.Stat(dbConfigSection.PasswordPath); os.IsNotExist(err) {
+	password := dbConfigSection.DBCertSection.Password
+	if len(dbConfigSection.DBCertSection.PasswordPath) > 0 {
+		if _, err := os.Stat(dbConfigSection.DBCertSection.PasswordPath); os.IsNotExist(err) {
 			logger.Fatalf(context.Background(),
-				"missing database password at specified path [%s]", dbConfigSection.PasswordPath)
+				"missing database password at specified path [%s]", dbConfigSection.DBCertSection.PasswordPath)
 		}
-		passwordVal, err := ioutil.ReadFile(dbConfigSection.PasswordPath)
+		passwordVal, err := ioutil.ReadFile(dbConfigSection.DBCertSection.PasswordPath)
 		if err != nil {
 			logger.Fatalf(context.Background(), "failed to read database password from path [%s] with err: %v",
-				dbConfigSection.PasswordPath, err)
+				dbConfigSection.DBCertSection.PasswordPath, err)
 		}
 		password = string(passwordVal)
 	}
+	rootCA := dbConfigSection.DBCertSection.RootCA
+	if len(dbConfigSection.DBCertSection.RootCAPath) > 0 {
+		if _, err := os.Stat(dbConfigSection.DBCertSection.RootCAPath); os.IsNotExist(err) {
+			logger.Fatalf(context.Background(),
+				"missing database root CA at specified path [%s]", dbConfigSection.DBCertSection.RootCAPath)
+		}
+		rootCAVal, err := ioutil.ReadFile(dbConfigSection.DBCertSection.RootCAPath)
+		if err != nil {
+			logger.Fatalf(context.Background(), "failed to read database root CA from path [%s] with err: %v",
+				dbConfigSection.DBCertSection.RootCAPath, err)
+		}
+		rootCA = string(rootCAVal)
+	}
+	var region string
+	if len(dbConfigSection.AWSDbConfig.Region) > 0 {
+		region = dbConfigSection.AWSDbConfig.Region
+	}
+
 	return interfaces.DbConfig{
 		Host:         dbConfigSection.Host,
 		Port:         dbConfigSection.Port,
 		DbName:       dbConfigSection.DbName,
 		User:         dbConfigSection.User,
 		Password:     password,
+		RootCA:       rootCA,
+		Region:       region,
+		UseIAM:       dbConfigSection.AWSDbConfig.UseIAM,
 		ExtraOptions: dbConfigSection.ExtraOptions,
 	}
 }

pkg/runtime/interfaces/application_configuration.go

@@ -1,5 +1,20 @@
 package interfaces
 
+type AWSDbConfig struct {
+	Region string `json:"region"`
+	UseIAM bool   `json:"useIAM"`
+}
+
+type DBCertSection struct {
+	// Either Password or PasswordPath must be set.
+	// The Password resolves to the database password.
+	Password     string `json:"password"`
+	PasswordPath string `json:"passwordPath"`
+	// If using a root CA, either RootCA or RootCAPath should be set.
+	RootCA     string `json:"rootCA"`
+	RootCAPath string `json:"rootCAPath"`
+}
+
 // This configuration section is used to for initiating the database connection with the store that holds registered
 // entities (e.g. workflows, tasks, launch plans...)
 // This struct specifically maps to the flyteadmin config yaml structure.
@@ -12,12 +27,10 @@ type DbConfigSection struct {
 	DbName string `json:"dbname"`
 	// The database user who is connecting to the server.
 	User string `json:"username"`
-	// Either Password or PasswordPath must be set.
-	// The Password resolves to the database password.
-	Password     string `json:"password"`
-	PasswordPath string `json:"passwordPath"`
 	// See http://gorm.io/docs/connecting_to_the_database.html for available options passed, in addition to the above.
-	ExtraOptions string `json:"options"`
+	ExtraOptions  string        `json:"options"`
+	DBCertSection DBCertSection `json:"dbCert"`
+	AWSDbConfig   AWSDbConfig   `json:"awsDbConfig"`
 }
 
 // This represents a configuration used for initiating database connections much like DbConfigSection, however the
@@ -29,7 +42,10 @@ type DbConfig struct {
 	DbName       string `json:"dbname"`
 	User         string `json:"username"`
 	Password     string `json:"password"`
+	RootCA       string `json:"rootCA"`
 	ExtraOptions string `json:"options"`
+	Region       string `json:"region"` // AWS specific
+	UseIAM       bool   `json:"useIAM"`
 }
 
 // This configuration is the base configuration to start admin