feat: add support for role based auth flag (#2519)

* feat: add support for role based auth flag Currently we are not able to delete already stored secrets from Config BE DB so we are adding support role based auth flag so that users can easily switch between AWS Access Keys and Roles. * fix: warehouse util tests
rudderlabs · Oct 10, 2022 · 810bb5d · 810bb5d
1 parent 1013a93
commit 810bb5d
Show file tree

Hide file tree

Showing 3 changed files with 148 additions and 43 deletions.
diff --git a/utils/awsutils/session.go b/utils/awsutils/session.go
@@ -21,6 +21,7 @@ type SessionConfig struct {
 	AccessKeyID      string         `mapstructure:"accessKeyID"`
 	AccessKey        string         `mapstructure:"accessKey"`
 	SecretAccessKey  string         `mapstructure:"secretAccessKey"`
+	RoleBasedAuth    bool           `mapstructure:"roleBasedAuth"`
 	IAMRoleARN       string         `mapstructure:"iamRoleARN"`
 	ExternalID       string         `mapstructure:"externalID"`
 	Endpoint         *string        `mapstructure:"endpoint"`
@@ -71,7 +72,7 @@ func CreateSession(config *SessionConfig) (*session.Session, error) {
 		awsCredentials *credentials.Credentials
 		err            error
 	)
-	if config.IAMRoleARN != "" {
+	if config.RoleBasedAuth {
 		awsCredentials, err = createCredentailsForRole(config)
 	} else if config.AccessKey != "" && config.AccessKeyID != "" {
 		awsCredentials, err = credentials.NewStaticCredentials(config.AccessKeyID, config.AccessKey, ""), nil
@@ -90,6 +91,11 @@ func CreateSession(config *SessionConfig) (*session.Session, error) {
 	})
 }
 
+func isRoleBasedAuthFieldExist(config map[string]interface{}) bool {
+	_, ok := config["roleBasedAuth"].(bool)
+	return ok
+}
+
 func NewSimpleSessionConfigForDestination(destination *backendconfig.DestinationT, serviceName string) (*SessionConfig, error) {
 	if destination == nil {
 		return nil, errors.New("destination should not be nil")
@@ -98,6 +104,15 @@ func NewSimpleSessionConfigForDestination(destination *backendconfig.Destination
 	if err := mapstructure.Decode(destination.Config, &sessionConfig); err != nil {
 		return nil, fmt.Errorf("unable to populate session config using destinationConfig: %w", err)
 	}
+
+	if sessionConfig.RoleBasedAuth && sessionConfig.IAMRoleARN == "" {
+		return nil, errors.New("incompatible role configuration")
+	}
+
+	if !isRoleBasedAuthFieldExist(destination.Config) {
+		sessionConfig.RoleBasedAuth = sessionConfig.IAMRoleARN != ""
+	}
+
 	// Some AWS destinations are using SecretAccessKey instead of accessKey
 	if sessionConfig.SecretAccessKey != "" {
 		sessionConfig.AccessKey = sessionConfig.SecretAccessKey

diff --git a/utils/awsutils/session_test.go b/utils/awsutils/session_test.go
@@ -65,22 +65,98 @@ func TestNewSessionConfigWithSecretAccessKey(t *testing.T) {
 
 func TestNewSessionConfigWithRole(t *testing.T) {
 	serviceName := "s3"
-	destinationWithRole := backendconfig.DestinationT{
-		Config: map[string]interface{}{
-			"region":     someRegion,
-			"iamRoleARN": someIAMRoleARN,
-		},
-		WorkspaceID: someWorkspaceID,
-	}
-	sessionConfig, err := NewSessionConfigForDestination(&destinationWithRole, httpTimeout, serviceName)
-	assert.Nil(t, err)
-	assert.NotNil(t, sessionConfig)
-	assert.Equal(t, *sessionConfig, SessionConfig{
-		Region:     someRegion,
-		IAMRoleARN: someIAMRoleARN,
-		ExternalID: someWorkspaceID,
-		Timeout:    &httpTimeout,
-		Service:    serviceName,
+	t.Run("Without RoleBasedAuth", func(t *testing.T) {
+		destinationWithRole := backendconfig.DestinationT{
+			Config: map[string]interface{}{
+				"region":     someRegion,
+				"iamRoleARN": someIAMRoleARN,
+			},
+			WorkspaceID: someWorkspaceID,
+		}
+		sessionConfig, err := NewSessionConfigForDestination(&destinationWithRole, httpTimeout, serviceName)
+		assert.Nil(t, err)
+		assert.NotNil(t, sessionConfig)
+		assert.Equal(t, *sessionConfig, SessionConfig{
+			Region:        someRegion,
+			RoleBasedAuth: true,
+			IAMRoleARN:    someIAMRoleARN,
+			ExternalID:    someWorkspaceID,
+			Timeout:       &httpTimeout,
+			Service:       serviceName,
+		})
+	})
+
+	t.Run("With RoleBasedAuth false", func(t *testing.T) {
+		destinationWithRole := backendconfig.DestinationT{
+			Config: map[string]interface{}{
+				"region":        someRegion,
+				"roleBasedAuth": false,
+				"iamRoleARN":    someIAMRoleARN,
+			},
+			WorkspaceID: someWorkspaceID,
+		}
+		sessionConfig, err := NewSessionConfigForDestination(&destinationWithRole, httpTimeout, serviceName)
+		assert.Nil(t, err)
+		assert.NotNil(t, sessionConfig)
+		assert.Equal(t, *sessionConfig, SessionConfig{
+			Region:        someRegion,
+			RoleBasedAuth: false,
+			IAMRoleARN:    someIAMRoleARN,
+			ExternalID:    someWorkspaceID,
+			Timeout:       &httpTimeout,
+			Service:       serviceName,
+		})
+	})
+}
+
+func TestNewSessionConfigWithRoleBasedAuth(t *testing.T) {
+	serviceName := "s3"
+	t.Run("invalid RoleBasedAuth flag", func(t *testing.T) {
+		destinationWithRole := backendconfig.DestinationT{
+			Config: map[string]interface{}{
+				"region":        someRegion,
+				"iamRoleARN":    someIAMRoleARN,
+				"roleBasedAuth": "no", // should be bool
+			},
+			WorkspaceID: someWorkspaceID,
+		}
+		_, err := NewSessionConfigForDestination(&destinationWithRole, httpTimeout, serviceName)
+		assert.NotNil(t, err)
+		assert.ErrorContains(t, err, "'roleBasedAuth' expected type 'bool'")
+	})
+	t.Run("With iamRoleARN", func(t *testing.T) {
+		destinationWithRole := backendconfig.DestinationT{
+			Config: map[string]interface{}{
+				"region":        someRegion,
+				"iamRoleARN":    someIAMRoleARN,
+				"roleBasedAuth": true,
+			},
+			WorkspaceID: someWorkspaceID,
+		}
+		sessionConfig, err := NewSessionConfigForDestination(&destinationWithRole, httpTimeout, serviceName)
+		assert.Nil(t, err)
+		assert.NotNil(t, sessionConfig)
+		assert.Equal(t, *sessionConfig, SessionConfig{
+			Region:        someRegion,
+			RoleBasedAuth: true,
+			IAMRoleARN:    someIAMRoleARN,
+			ExternalID:    someWorkspaceID,
+			Timeout:       &httpTimeout,
+			Service:       serviceName,
+		})
+	})
+
+	t.Run("Without iamRoleARN", func(t *testing.T) {
+		destinationWithRole := backendconfig.DestinationT{
+			Config: map[string]interface{}{
+				"region":        someRegion,
+				"roleBasedAuth": true,
+			},
+			WorkspaceID: someWorkspaceID,
+		}
+		_, err := NewSessionConfigForDestination(&destinationWithRole, httpTimeout, serviceName)
+		assert.NotNil(t, err)
+		assert.EqualError(t, err, "incompatible role configuration")
 	})
 }
 
@@ -115,30 +191,43 @@ func TestNewSessionConfigWithBadDestination(t *testing.T) {
 }
 
 func TestCreateSessionWithRole(t *testing.T) {
-	sessionConfig := SessionConfig{
-		Region:     someRegion,
-		IAMRoleARN: someIAMRoleARN,
-		ExternalID: someWorkspaceID,
-		Timeout:    &httpTimeout,
-	}
-	awsSession, err := CreateSession(&sessionConfig)
-	assert.Nil(t, err)
-	assert.NotNil(t, awsSession)
-	assert.NotNil(t, awsSession.Config.Credentials)
-	assert.Equal(t, sessionConfig.Region, *awsSession.Config.Region)
-	assert.Equal(t, *sessionConfig.Timeout, awsSession.Config.HTTPClient.Timeout)
-}
+	t.Run("With RoleBasedAuth but without ExternalID", func(t *testing.T) {
+		sessionConfig := SessionConfig{
+			Region:        someRegion,
+			RoleBasedAuth: true,
+			IAMRoleARN:    someIAMRoleARN,
+			Timeout:       &httpTimeout,
+		}
+		awsSession, err := CreateSession(&sessionConfig)
+		assert.NotNil(t, err)
+		assert.Nil(t, awsSession)
+		assert.EqualError(t, err, "externalID is required for IAM role")
+	})
 
-func TestCreateSessionWithRoleButWithoutExternalID(t *testing.T) {
-	sessionConfig := SessionConfig{
-		Region:     someRegion,
-		IAMRoleARN: someIAMRoleARN,
-		Timeout:    &httpTimeout,
-	}
-	awsSession, err := CreateSession(&sessionConfig)
-	assert.NotNil(t, err)
-	assert.Nil(t, awsSession)
-	assert.EqualError(t, err, "externalID is required for IAM role")
+	t.Run("With RoleBasedAuth false and without ExternalID", func(t *testing.T) {
+		sessionConfig := SessionConfig{
+			Region:        someRegion,
+			RoleBasedAuth: false,
+			IAMRoleARN:    someIAMRoleARN,
+			Timeout:       &httpTimeout,
+		}
+		awsSession, err := CreateSession(&sessionConfig)
+		assert.Nil(t, err)
+		assert.NotNil(t, awsSession)
+	})
+
+	t.Run("With RoleBasedAuth true auth and ExternalID", func(t *testing.T) {
+		sessionConfig := SessionConfig{
+			Region:        someRegion,
+			RoleBasedAuth: true,
+			ExternalID:    someWorkspaceID,
+			IAMRoleARN:    someIAMRoleARN,
+			Timeout:       &httpTimeout,
+		}
+		awsSession, err := CreateSession(&sessionConfig)
+		assert.Nil(t, err)
+		assert.NotNil(t, awsSession)
+	})
 }
 
 func TestCreateSessionWithAccessKeys(t *testing.T) {

diff --git a/warehouse/utils/utils_test.go b/warehouse/utils/utils_test.go
@@ -1239,9 +1239,10 @@ func TestCreateAWSSessionConfig(t *testing.T) {
 			},
 			service: "redshift",
 			expectedConfig: &awsutils.SessionConfig{
-				IAMRoleARN: someIAMRoleARN,
-				ExternalID: someWorkspaceID,
-				Service:    "redshift",
+				RoleBasedAuth: true,
+				IAMRoleARN:    someIAMRoleARN,
+				ExternalID:    someWorkspaceID,
+				Service:       "redshift",
 			},
 		},
 	}