Add region into IAM clients as well as EC2 clients (#18)

plugin/plugin.go

@@ -70,6 +70,7 @@ func (p *AwsPlugin) OnCreateCatalog(ctx context.Context, req *pb.OnCreateCatalog
 	state, err := newAwsCatalogPersistedState(append([]awsCatalogPersistedStateOption{
 		withAccessKeyId(catalogSecrets.AccessKeyId),
 		withSecretAccessKey(catalogSecrets.SecretAccessKey),
+		withRegion(catalogAttributes.Region),
 	}, p.testStateOpts...)...)
 	if err != nil {
 		return nil, status.Errorf(codes.InvalidArgument, "error setting up persisted state: %s", err)
@@ -136,7 +137,7 @@ func (p *AwsPlugin) OnUpdateCatalog(ctx context.Context, req *pb.OnUpdateCatalog
 	// (implicitly through awsCatalogPersistedStateFromProto) is that
 	// the state will exist and be populated. Personally I think this
 	// is fine and important, but this may change in the future.
-	state, err := awsCatalogPersistedStateFromProto(req.GetPersisted(), p.testStateOpts...)
+	state, err := awsCatalogPersistedStateFromProto(req.GetPersisted(), append(p.testStateOpts, withRegion(catalogAttributes.Region))...)
 	if err != nil {
 		return nil, status.Errorf(codes.InvalidArgument, "error loading persisted state: %s", err)
 	}
@@ -184,13 +185,31 @@ func (p *AwsPlugin) OnUpdateCatalog(ctx context.Context, req *pb.OnUpdateCatalog
 }
 
 func (p *AwsPlugin) OnDeleteCatalog(ctx context.Context, req *pb.OnDeleteCatalogRequest) (*pb.OnDeleteCatalogResponse, error) {
+	catalog := req.GetCatalog()
+	if catalog == nil {
+		return nil, status.Error(codes.InvalidArgument, "new catalog is nil")
+	}
+
+	attrs := catalog.GetAttributes()
+	if attrs == nil {
+		return nil, status.Error(codes.InvalidArgument, "new catalog missing attributes")
+	}
+
+	catalogAttributes, err := getCatalogAttributes(attrs)
+	if err != nil {
+		return nil, err
+	}
+
+	if err := validateRegion(catalogAttributes.Region); err != nil {
+		return nil, err
+	}
 	// Get the persisted data.
 	// NOTE: We return on error here, blocking the delete. This may or
 	// may not be an overzealous approach to maintaining database/state
 	// integrity. May need to be changed at later time if there are
 	// scenarios where we might be deleting things and any secret state
 	// may be corrupt/and or legitimately missing.
-	state, err := awsCatalogPersistedStateFromProto(req.GetPersisted(), p.testStateOpts...)
+	state, err := awsCatalogPersistedStateFromProto(req.GetPersisted(), append(p.testStateOpts, withRegion(catalogAttributes.Region))...)
 	if err != nil {
 		return nil, status.Errorf(codes.InvalidArgument, "error loading persisted state: %s", err)
 	}
@@ -254,7 +273,7 @@ func (p *AwsPlugin) OnCreateSet(ctx context.Context, req *pb.OnCreateSetRequest)
 		return nil, err
 	}
 
-	state, err := awsCatalogPersistedStateFromProto(req.GetPersisted(), p.testStateOpts...)
+	state, err := awsCatalogPersistedStateFromProto(req.GetPersisted(), append(p.testStateOpts, withRegion(catalogAttributes.Region))...)
 	if err != nil {
 		return nil, status.Errorf(codes.InvalidArgument, "error loading persisted state: %s", err)
 	}
@@ -315,7 +334,7 @@ func (p *AwsPlugin) OnUpdateSet(ctx context.Context, req *pb.OnUpdateSetRequest)
 		return nil, err
 	}
 
-	state, err := awsCatalogPersistedStateFromProto(req.GetPersisted(), p.testStateOpts...)
+	state, err := awsCatalogPersistedStateFromProto(req.GetPersisted(), append(p.testStateOpts, withRegion(catalogAttributes.Region))...)
 	if err != nil {
 		return nil, status.Errorf(codes.InvalidArgument, "error loading persisted state: %s", err)
 	}
@@ -384,7 +403,7 @@ func (p *AwsPlugin) ListHosts(ctx context.Context, req *pb.ListHostsRequest) (*p
 		return nil, err
 	}
 
-	state, err := awsCatalogPersistedStateFromProto(req.GetPersisted(), p.testStateOpts...)
+	state, err := awsCatalogPersistedStateFromProto(req.GetPersisted(), append(p.testStateOpts, withRegion(catalogAttributes.Region))...)
 	if err != nil {
 		return nil, status.Errorf(codes.InvalidArgument, "error loading persisted state: %s", err)
 	}
@@ -427,7 +446,7 @@ func (p *AwsPlugin) ListHosts(ctx context.Context, req *pb.ListHostsRequest) (*p
 		}
 	}
 
-	ec2Client, err := state.EC2Client(WithRegion(catalogAttributes.Region))
+	ec2Client, err := state.EC2Client()
 	if err != nil {
 		return nil, status.Errorf(codes.InvalidArgument, "error getting EC2 client: %s", err)
 	}

plugin/plugin_test.go

@@ -651,7 +651,7 @@ func TestPluginOnCreateSetErr(t *testing.T) {
 					awsutil.MockOptionErr(errors.New(testOptionErr)),
 				}),
 			},
-			expectedErrContains: fmt.Sprintf("error getting EC2 client: error getting AWS session: error reading options in NewCredentialsConfig: %s", testOptionErr),
+			expectedErrContains: fmt.Sprintf("error getting EC2 client: error getting AWS session when fetching EC2 client: error reading options in NewCredentialsConfig: %s", testOptionErr),
 			expectedErrCode:     codes.InvalidArgument,
 		},
 		{
@@ -950,7 +950,7 @@ func TestPluginOnUpdateSetErr(t *testing.T) {
 					awsutil.MockOptionErr(errors.New(testOptionErr)),
 				}),
 			},
-			expectedErrContains: fmt.Sprintf("error getting EC2 client: error getting AWS session: error reading options in NewCredentialsConfig: %s", testOptionErr),
+			expectedErrContains: fmt.Sprintf("error getting EC2 client: error getting AWS session when fetching EC2 client: error reading options in NewCredentialsConfig: %s", testOptionErr),
 			expectedErrCode:     codes.InvalidArgument,
 		},
 		{
@@ -1248,7 +1248,7 @@ func TestPluginListHostsErr(t *testing.T) {
 					awsutil.MockOptionErr(errors.New(testOptionErr)),
 				}),
 			},
-			expectedErrContains: fmt.Sprintf("error getting EC2 client: error getting AWS session: error reading options in NewCredentialsConfig: %s", testOptionErr),
+			expectedErrContains: fmt.Sprintf("error getting EC2 client: error getting AWS session when fetching EC2 client: error reading options in NewCredentialsConfig: %s", testOptionErr),
 			expectedErrCode:     codes.InvalidArgument,
 		},
 		{

plugin/state.go

@@ -73,6 +73,17 @@ func withCredsLastRotatedTime(t time.Time) awsCatalogPersistedStateOption {
 	}
 }
 
+func withRegion(with string) awsCatalogPersistedStateOption {
+	return func(s *awsCatalogPersistedState) error {
+		if s.region != "" {
+			return errors.New("region already set")
+		}
+
+		s.region = with
+		return nil
+	}
+}
+
 type awsCatalogPersistedState struct {
 	AccessKeyId          string
 	SecretAccessKey      string
@@ -83,6 +94,10 @@ type awsCatalogPersistedState struct {
 
 	// testEC2APIFunc provides a way to provide a factory for a mock EC2 client
 	testEC2APIFunc ec2APIFunc
+
+	// Region is not a part of persisted state but if set is passed to clients
+	// created with this state
+	region string
 }
 
 func newAwsCatalogPersistedState(opts ...awsCatalogPersistedStateOption) (*awsCatalogPersistedState, error) {
@@ -146,6 +161,7 @@ func (s *awsCatalogPersistedState) ValidateCreds() error {
 	c, err := awsutil.NewCredentialsConfig(append([]awsutil.Option{
 		awsutil.WithAccessKey(s.AccessKeyId),
 		awsutil.WithSecretKey(s.SecretAccessKey),
+		awsutil.WithRegion(s.region),
 	}, s.testOpts...)...)
 	if err != nil {
 		return fmt.Errorf("error loading credentials: %w", err)
@@ -162,6 +178,7 @@ func (s *awsCatalogPersistedState) RotateCreds() error {
 	c, err := awsutil.NewCredentialsConfig(append([]awsutil.Option{
 		awsutil.WithAccessKey(s.AccessKeyId),
 		awsutil.WithSecretKey(s.SecretAccessKey),
+		awsutil.WithRegion(s.region),
 	}, s.testOpts...)...)
 	if err != nil {
 		return fmt.Errorf("error loading credentials: %w", err)
@@ -219,6 +236,7 @@ func (s *awsCatalogPersistedState) DeleteCreds() error {
 	c, err := awsutil.NewCredentialsConfig(append([]awsutil.Option{
 		awsutil.WithAccessKey(s.AccessKeyId),
 		awsutil.WithSecretKey(s.SecretAccessKey),
+		awsutil.WithRegion(s.region),
 	}, s.testOpts...)...)
 	if err != nil {
 		return fmt.Errorf("error loading credentials: %w", err)
@@ -253,6 +271,7 @@ func (s *awsCatalogPersistedState) GetSession() (*session.Session, error) {
 	c, err := awsutil.NewCredentialsConfig(append([]awsutil.Option{
 		awsutil.WithAccessKey(s.AccessKeyId),
 		awsutil.WithSecretKey(s.SecretAccessKey),
+		awsutil.WithRegion(s.region),
 	}, s.testOpts...)...)
 	if err != nil {
 		return nil, err
@@ -280,18 +299,23 @@ func (s *awsCatalogPersistedState) EC2Client(opt ...Option) (ec2iface.EC2API, er
 		return nil, fmt.Errorf("error parsing options when fetching EC2 client: %w", err)
 	}
 
+	region := opts.withRegion
+	if region == "" {
+		region = s.region
+	}
+
 	sess, err := s.GetSession()
 	if err != nil {
 		return nil, fmt.Errorf("error getting AWS session when fetching EC2 client: %w", err)
 	}
 
 	if s.testEC2APIFunc != nil {
-		return s.testEC2APIFunc(sess, aws.NewConfig().WithRegion(opts.withRegion))
+		return s.testEC2APIFunc(sess, aws.NewConfig().WithRegion(region))
 	}
 
 	cfg := aws.NewConfig()
-	if opts.withRegion != "" {
+	if region != "" {
-		cfg = cfg.WithRegion(opts.withRegion)
+		cfg = cfg.WithRegion(region)
 	}
 
 	cfg.HTTPClient = customClient

plugin/state_test.go

@@ -14,10 +14,12 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-const testOptionErr = "test option error"
+const (
-const testGetCallerIdentityErr = "test error for GetCallerIdentity"
+	testOptionErr            = "test option error"
-const testGetUserErr = "test error for GetUser"
+	testGetCallerIdentityErr = "test error for GetCallerIdentity"
-const testDeleteAccessKeyErr = "test error for DeleteAccessKey"
+	testGetUserErr           = "test error for GetUser"
+	testDeleteAccessKeyErr   = "test error for DeleteAccessKey"
+)
 
 func TestNewAwsCatalogPersistedState(t *testing.T) {
 	staticTime := time.Now()
@@ -64,6 +66,15 @@ func TestNewAwsCatalogPersistedState(t *testing.T) {
 				CredsLastRotatedTime: staticTime,
 			},
 		},
+		{
+			name: "region key",
+			opts: []awsCatalogPersistedStateOption{
+				withRegion("foobar"),
+			},
+			expected: &awsCatalogPersistedState{
+				region: "foobar",
+			},
+		},
 		{
 			name: "double set access key id",
 			opts: []awsCatalogPersistedStateOption{
@@ -73,7 +84,7 @@ func TestNewAwsCatalogPersistedState(t *testing.T) {
 			expectedErr: "access key id already set",
 		},
 		{
-			name: "secret access key",
+			name: "double set secret access key",
 			opts: []awsCatalogPersistedStateOption{
 				withSecretAccessKey("bazqux"),
 				withSecretAccessKey("threefour"),
@@ -81,13 +92,21 @@ func TestNewAwsCatalogPersistedState(t *testing.T) {
 			expectedErr: "secret access key already set",
 		},
 		{
-			name: "rotation time",
+			name: "double set rotation time",
 			opts: []awsCatalogPersistedStateOption{
 				withCredsLastRotatedTime(staticTime),
 				withCredsLastRotatedTime(time.Now()),
 			},
 			expectedErr: "last rotation time already set",
 		},
+		{
+			name: "double set region",
+			opts: []awsCatalogPersistedStateOption{
+				withRegion("foobar"),
+				withRegion("barfoo"),
+			},
+			expectedErr: "region already set",
+		},
 	}
 
 	for _, tc := range cases {
@@ -219,11 +238,13 @@ func TestAwsCatalogPersistedStateFromProto(t *testing.T) {
 			},
 			opts: []awsCatalogPersistedStateOption{
 				withAccessKeyId("ignored"),
+				withRegion("someregion"),
 			},
 			expected: &awsCatalogPersistedState{
 				AccessKeyId:          "foobar",
 				SecretAccessKey:      "bazqux",
 				CredsLastRotatedTime: time.Time{},
+				region:               "someregion",
 			},
 		},
 	}
@@ -641,8 +662,8 @@ func TestAwsCatalogPersistedStateEC2ClientErr(t *testing.T) {
 	)
 	require.NoError(err)
 
-	_, err = state.EC2Client("someregion")
+	_, err = state.EC2Client(WithRegion("someregion"))
-	require.EqualError(err, fmt.Sprintf("error getting AWS session: error reading options in NewCredentialsConfig: %s", testOptionErr))
+	require.EqualError(err, fmt.Sprintf("error getting AWS session when fetching EC2 client: error reading options in NewCredentialsConfig: %s", testOptionErr))
 }
 
 func TestAwsCatalogPersistedStateEC2ClientGood(t *testing.T) {
@@ -652,7 +673,7 @@ func TestAwsCatalogPersistedStateEC2ClientGood(t *testing.T) {
 	)
 	require.NoError(err)
 
-	clientRaw, err := state.EC2Client("someregion")
+	clientRaw, err := state.EC2Client(WithRegion("someregion"))
 	require.NoError(err)
 	require.NotNil(clientRaw)
 	client, ok := clientRaw.(*testMockEC2)