fix: Return GrantAlreadyExists annotation

pkg/connector/account.go

@@ -15,6 +15,7 @@ import (
 	"github.com/aws/aws-sdk-go-v2/service/organizations/types"
 	awsSsoAdmin "github.com/aws/aws-sdk-go-v2/service/ssoadmin"
 	awsSsoAdminTypes "github.com/aws/aws-sdk-go-v2/service/ssoadmin/types"
+	"github.com/conductorone/baton-aws/pkg/connector/client"
 	"github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap"
 	"go.uber.org/zap"
 
@@ -36,7 +37,7 @@ type accountResourceType struct {
 	ssoAdminClient   *awsSsoAdmin.Client
 	roleArn          string
 	identityInstance *awsSsoAdminTypes.InstanceMetadata
-	identityClient   *awsIdentityStore.Client
+	identityClient   client.IdentityStoreClient
 	region           string
 
 	_permissionSetsCacheMtx    sync.Mutex
@@ -496,7 +497,7 @@ func accountBuilder(
 	ssoAdminClient *awsSsoAdmin.Client,
 	identityInstance *awsSsoAdminTypes.InstanceMetadata,
 	region string,
-	identityClient *awsIdentityStore.Client,
+	identityClient client.IdentityStoreClient,
 ) *accountResourceType {
 	return &accountResourceType{
 		resourceType:     resourceTypeAccount,

pkg/connector/client/identityStore.go

@@ -0,0 +1,31 @@
+package client
+
+import (
+	"context"
+
+	"github.com/aws/aws-sdk-go-v2/service/identitystore"
+)
+
+// IdentityStoreClient is a wrapper interface around `identitystore.Client` so
+// that we can hook in with mocks for unit tests.
+type IdentityStoreClient interface {
+	identitystore.ListGroupMembershipsAPIClient
+	identitystore.ListGroupMembershipsForMemberAPIClient
+	identitystore.ListGroupsAPIClient
+	identitystore.ListUsersAPIClient
+	CreateGroupMembership(
+		ctx context.Context,
+		params *identitystore.CreateGroupMembershipInput,
+		optFns ...func(*identitystore.Options),
+	) (*identitystore.CreateGroupMembershipOutput, error)
+	DeleteGroupMembership(
+		ctx context.Context,
+		params *identitystore.DeleteGroupMembershipInput,
+		optFns ...func(*identitystore.Options),
+	) (*identitystore.DeleteGroupMembershipOutput, error)
+	GetGroupMembershipId(
+		ctx context.Context,
+		params *identitystore.GetGroupMembershipIdInput,
+		optFns ...func(*identitystore.Options),
+	) (*identitystore.GetGroupMembershipIdOutput, error)
+}

pkg/connector/client/sso.go

@@ -0,0 +1,6 @@
+package client
+
+// SSOClient is a wrapper interface around `identitystore.Client` so
+// that we can hook in with mocks for unit tests.
+type SSOClient interface {
+}

pkg/connector/connector.go

@@ -19,6 +19,7 @@ import (
 	awsSsoAdmin "github.com/aws/aws-sdk-go-v2/service/ssoadmin"
 	awsSsoAdminTypes "github.com/aws/aws-sdk-go-v2/service/ssoadmin/types"
 	"github.com/aws/aws-sdk-go-v2/service/sts"
+	"github.com/conductorone/baton-aws/pkg/connector/client"
 	"github.com/conductorone/baton-sdk/pkg/annotations"
 	"github.com/conductorone/baton-sdk/pkg/connectorbuilder"
 	"github.com/conductorone/baton-sdk/pkg/uhttp"
@@ -29,51 +30,6 @@ import (
 	v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2"
 )
 
-var (
-	resourceTypeRole = &v2.ResourceType{
-		Id:          "role",
-		DisplayName: "IAM Role",
-		Traits:      []v2.ResourceType_Trait{v2.ResourceType_TRAIT_ROLE},
-		Annotations: v1AnnotationsForResourceType("role"),
-	}
-	resourceTypeIAMGroup = &v2.ResourceType{
-		Id:          "group",
-		DisplayName: "Group",
-		Traits:      []v2.ResourceType_Trait{v2.ResourceType_TRAIT_GROUP},
-		Annotations: v1AnnotationsForResourceType("group"),
-	}
-	resourceTypeSSOGroup = &v2.ResourceType{
-		Id:          "sso_group",
-		DisplayName: "SSO Group",
-		Traits: []v2.ResourceType_Trait{
-			v2.ResourceType_TRAIT_GROUP,
-		},
-		Annotations: v1AnnotationsForResourceType("sso_group"),
-	}
-	resourceTypeAccount = &v2.ResourceType{
-		Id:          "account", // this is "application" in c1
-		DisplayName: "Account",
-		Traits:      []v2.ResourceType_Trait{v2.ResourceType_TRAIT_APP},
-		Annotations: v1AnnotationsForResourceType("account"),
-	}
-	resourceTypeSSOUser = &v2.ResourceType{
-		Id:          "sso_user",
-		DisplayName: "SSO User",
-		Traits: []v2.ResourceType_Trait{
-			v2.ResourceType_TRAIT_USER,
-		},
-		Annotations: v1AnnotationsForResourceType("sso_user"),
-	}
-	resourceTypeIAMUser = &v2.ResourceType{
-		Id:          "iam_user",
-		DisplayName: "IAM User",
-		Traits: []v2.ResourceType_Trait{
-			v2.ResourceType_TRAIT_USER,
-		},
-		Annotations: v1AnnotationsForResourceType("iam_user"),
-	}
-)
-
 type Config struct {
 	UseAssumeRole           bool
 	GlobalBindingExternalID string
@@ -120,7 +76,7 @@ type AWS struct {
 	orgClient           *awsOrgs.Client
 	ssoAdminClient      *awsSsoAdmin.Client
 	ssoSCIMClient       *awsIdentityCenterSCIMClient
-	identityStoreClient *awsIdentityStore.Client
+	identityStoreClient client.IdentityStoreClient
 	identityInstance    *awsSsoAdminTypes.InstanceMetadata
 }
 

pkg/connector/resource_types.go

@@ -0,0 +1,48 @@
+package connector
+
+import v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2"
+
+var (
+	resourceTypeRole = &v2.ResourceType{
+		Id:          "role",
+		DisplayName: "IAM Role",
+		Traits:      []v2.ResourceType_Trait{v2.ResourceType_TRAIT_ROLE},
+		Annotations: v1AnnotationsForResourceType("role"),
+	}
+	resourceTypeIAMGroup = &v2.ResourceType{
+		Id:          "group",
+		DisplayName: "Group",
+		Traits:      []v2.ResourceType_Trait{v2.ResourceType_TRAIT_GROUP},
+		Annotations: v1AnnotationsForResourceType("group"),
+	}
+	resourceTypeSSOGroup = &v2.ResourceType{
+		Id:          "sso_group",
+		DisplayName: "SSO Group",
+		Traits: []v2.ResourceType_Trait{
+			v2.ResourceType_TRAIT_GROUP,
+		},
+		Annotations: v1AnnotationsForResourceType("sso_group"),
+	}
+	resourceTypeAccount = &v2.ResourceType{
+		Id:          "account", // this is "application" in c1
+		DisplayName: "Account",
+		Traits:      []v2.ResourceType_Trait{v2.ResourceType_TRAIT_APP},
+		Annotations: v1AnnotationsForResourceType("account"),
+	}
+	resourceTypeSSOUser = &v2.ResourceType{
+		Id:          "sso_user",
+		DisplayName: "SSO User",
+		Traits: []v2.ResourceType_Trait{
+			v2.ResourceType_TRAIT_USER,
+		},
+		Annotations: v1AnnotationsForResourceType("sso_user"),
+	}
+	resourceTypeIAMUser = &v2.ResourceType{
+		Id:          "iam_user",
+		DisplayName: "IAM User",
+		Traits: []v2.ResourceType_Trait{
+			v2.ResourceType_TRAIT_USER,
+		},
+		Annotations: v1AnnotationsForResourceType("iam_user"),
+	}
+)

pkg/connector/sso_group.go

@@ -11,6 +11,7 @@ import (
 	awsSsoAdmin "github.com/aws/aws-sdk-go-v2/service/ssoadmin"
 	awsSsoAdminTypes "github.com/aws/aws-sdk-go-v2/service/ssoadmin/types"
 	"github.com/aws/smithy-go/middleware"
+	"github.com/conductorone/baton-aws/pkg/connector/client"
 	v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2"
 	"github.com/conductorone/baton-sdk/pkg/annotations"
 	"github.com/conductorone/baton-sdk/pkg/pagination"
@@ -24,7 +25,7 @@ import (
 type ssoGroupResourceType struct {
 	resourceType        *v2.ResourceType
 	ssoClient           *awsSsoAdmin.Client
-	identityStoreClient *awsIdentityStore.Client
+	identityStoreClient client.IdentityStoreClient
 	identityInstance    *awsSsoAdminTypes.InstanceMetadata
 	region              string
 }
@@ -184,7 +185,7 @@ func (o *ssoGroupResourceType) Grants(ctx context.Context, resource *v2.Resource
 func ssoGroupBuilder(
 	region string,
 	ssoClient *awsSsoAdmin.Client,
-	identityStoreClient *awsIdentityStore.Client,
+	identityStoreClient client.IdentityStoreClient,
 	identityInstance *awsSsoAdminTypes.InstanceMetadata,
 ) *ssoGroupResourceType {
 	return &ssoGroupResourceType{
@@ -211,7 +212,7 @@ func (g *ssoGroupResourceType) createOrGetMembership(
 	userID string,
 ) (
 	*GroupMembershipOutput,
-	*annotations.Annotations,
+	annotations.Annotations,
 	error,
 ) {
 	logger := ctxzap.Extract(ctx).With(
@@ -222,6 +223,7 @@ func (g *ssoGroupResourceType) createOrGetMembership(
 			awsSdk.ToString(g.identityInstance.IdentityStoreId),
 		),
 	)
+	outputAnnotations := annotations.New()
 	groupIdString := awsSdk.String(groupID)
 	memberId := awsIdentityStoreTypes.MemberIdMemberUserId{Value: userID}
 	createInput := &awsIdentityStore.CreateGroupMembershipInput{
@@ -234,7 +236,7 @@ func (g *ssoGroupResourceType) createOrGetMembership(
 		return &GroupMembershipOutput{
 			MembershipId:   createdMembership.MembershipId,
 			ResultMetadata: createdMembership.ResultMetadata,
-		}, nil, nil
+		}, outputAnnotations, nil
 	}
 
 	// Forward along the error if it is an unknown type.
@@ -243,6 +245,8 @@ func (g *ssoGroupResourceType) createOrGetMembership(
 		return nil, nil, err
 	}
 
+	outputAnnotations.Append(&v2.GrantAlreadyExists{})
+
 	logger.Info("ConflictException when creating group, falling back to GET")
 
 	getInput := awsIdentityStore.GetGroupMembershipIdInput{
@@ -257,16 +261,15 @@ func (g *ssoGroupResourceType) createOrGetMembership(
 		var accessDeniedException *awsIdentityStoreTypes.AccessDeniedException
 		if errors.As(err, &accessDeniedException) {
 			logger.Info("Not authorized to perform `GetGroupMembershipId`, falling back to empty membership")
-			// TODO(marcos): Create an annotation that marks this grant as "already exists".
-			return nil, nil, nil
+			return nil, outputAnnotations, nil
 		}
 
-		return nil, nil, err
+		return nil, outputAnnotations, err
 	}
 
 	return &GroupMembershipOutput{
 		MembershipId: foundMembership.MembershipId,
-	}, nil, nil
+	}, outputAnnotations, nil
 }
 
 func (g *ssoGroupResourceType) Grant(
@@ -298,35 +301,41 @@ func (g *ssoGroupResourceType) Grant(
 		zap.String("identity_store_id", awsSdk.ToString(g.identityInstance.IdentityStoreId)),
 	)
 
-	// TODO(marcos): If we get a nil membership and an annotation, return that annotation.
-	membership, _, err := g.createOrGetMembership(ctx, groupID, userID)
+	annos := annotations.New()
+	outputGrants := make([]*v2.Grant, 0)
+
+	membership, annotationsFromGet, err := g.createOrGetMembership(ctx, groupID, userID)
 	if err != nil {
 		l.Error("aws-connector: Failed to create group membership", zap.Error(err))
 		return nil, nil, fmt.Errorf("baton-aws: error adding sso user to sso group: %w", err)
 	}
 
-	annos := annotations.New()
-	if membership == nil {
-		return []*v2.Grant{}, annos, nil
-	}
+	annos.Merge(annotationsFromGet...)
 
-	grant, err := createUserSSOGroupMembershipGrant(
-		g.region,
-		awsSdk.ToString(g.identityInstance.IdentityStoreId),
-		userID,
-		membership.MembershipId,
-		entitlement.Resource,
-	)
-	if err != nil {
-		l.Error("aws-connector: Failed to create grant", zap.Error(err), zap.String("membership_id", awsSdk.ToString(membership.MembershipId)))
-		return nil, nil, err
-	}
+	if membership != nil {
+		grant, err := createUserSSOGroupMembershipGrant(
+			g.region,
+			awsSdk.ToString(g.identityInstance.IdentityStoreId),
+			userID,
+			membership.MembershipId,
+			entitlement.Resource,
+		)
+		if err != nil {
+			l.Error(
+				"aws-connector: Failed to create grant",
+				zap.Error(err),
+				zap.String("membership_id", awsSdk.ToString(membership.MembershipId)),
+			)
+			return nil, nil, err
+		}
 
-	if reqId := extractRequestID(&membership.ResultMetadata); reqId != nil {
-		annos.Append(reqId)
+		if reqId := extractRequestID(&membership.ResultMetadata); reqId != nil {
+			annos.Append(reqId)
+		}
+		outputGrants = append(outputGrants, grant)
 	}
 
-	return []*v2.Grant{grant}, annos, nil
+	return outputGrants, annos, nil
 }
 func (g *ssoGroupResourceType) Revoke(ctx context.Context, grant *v2.Grant) (annotations.Annotations, error) {
 	if grant.Principal.Id.ResourceType != resourceTypeSSOUser.Id {