From cccea04e29ea42842da1cb576b1ebf8070ad16af Mon Sep 17 00:00:00 2001 From: Marcos Gaeta Date: Tue, 27 Aug 2024 11:46:50 -0700 Subject: [PATCH 1/5] add new annotation --- pkg/connector/sso_group.go | 60 +++++++++++++++++++++----------------- 1 file changed, 34 insertions(+), 26 deletions(-) diff --git a/pkg/connector/sso_group.go b/pkg/connector/sso_group.go index a1f2603e..593dee35 100644 --- a/pkg/connector/sso_group.go +++ b/pkg/connector/sso_group.go @@ -211,7 +211,7 @@ func (g *ssoGroupResourceType) createOrGetMembership( userID string, ) ( *GroupMembershipOutput, - *annotations.Annotations, + annotations.Annotations, error, ) { logger := ctxzap.Extract(ctx).With( @@ -222,6 +222,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 +235,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 +244,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 +260,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 +300,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 { From 3c9732b8f53c523f25ff7ccea98627adcfe157f8 Mon Sep 17 00:00:00 2001 From: Marcos Gaeta Date: Tue, 27 Aug 2024 17:57:51 -0700 Subject: [PATCH 2/5] move resource types to a file --- pkg/connector/connector.go | 45 ------------------------------- pkg/connector/resource_types.go | 48 +++++++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 45 deletions(-) create mode 100644 pkg/connector/resource_types.go diff --git a/pkg/connector/connector.go b/pkg/connector/connector.go index 016f57a2..04921255 100644 --- a/pkg/connector/connector.go +++ b/pkg/connector/connector.go @@ -29,51 +29,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 diff --git a/pkg/connector/resource_types.go b/pkg/connector/resource_types.go new file mode 100644 index 00000000..9e435deb --- /dev/null +++ b/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"), + } +) From f7c4b2381fc05a01a8deb34b018bf7e636de1b14 Mon Sep 17 00:00:00 2001 From: Marcos Gaeta Date: Tue, 27 Aug 2024 17:59:25 -0700 Subject: [PATCH 3/5] Moke identity store client an interface --- pkg/connector/account.go | 5 +++-- pkg/connector/client/identityStore.go | 31 +++++++++++++++++++++++++++ pkg/connector/client/sso.go | 6 ++++++ pkg/connector/connector.go | 3 ++- pkg/connector/sso_group.go | 5 +++-- pkg/connector/sso_user.go | 5 +++-- 6 files changed, 48 insertions(+), 7 deletions(-) create mode 100644 pkg/connector/client/identityStore.go create mode 100644 pkg/connector/client/sso.go diff --git a/pkg/connector/account.go b/pkg/connector/account.go index 7656116f..3ccadd1f 100644 --- a/pkg/connector/account.go +++ b/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, diff --git a/pkg/connector/client/identityStore.go b/pkg/connector/client/identityStore.go new file mode 100644 index 00000000..da1af47a --- /dev/null +++ b/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) +} diff --git a/pkg/connector/client/sso.go b/pkg/connector/client/sso.go new file mode 100644 index 00000000..ccca282b --- /dev/null +++ b/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 { +} diff --git a/pkg/connector/connector.go b/pkg/connector/connector.go index 04921255..e1b0c0eb 100644 --- a/pkg/connector/connector.go +++ b/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" @@ -75,7 +76,7 @@ type AWS struct { orgClient *awsOrgs.Client ssoAdminClient *awsSsoAdmin.Client ssoSCIMClient *awsIdentityCenterSCIMClient - identityStoreClient *awsIdentityStore.Client + identityStoreClient client.IdentityStoreClient identityInstance *awsSsoAdminTypes.InstanceMetadata } diff --git a/pkg/connector/sso_group.go b/pkg/connector/sso_group.go index 593dee35..efcf7922 100644 --- a/pkg/connector/sso_group.go +++ b/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{ diff --git a/pkg/connector/sso_user.go b/pkg/connector/sso_user.go index d86f42c6..41cd5622 100644 --- a/pkg/connector/sso_user.go +++ b/pkg/connector/sso_user.go @@ -10,6 +10,7 @@ import ( awsIdentityStoreTypes "github.com/aws/aws-sdk-go-v2/service/identitystore/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" v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" "github.com/conductorone/baton-sdk/pkg/annotations" "github.com/conductorone/baton-sdk/pkg/pagination" @@ -19,7 +20,7 @@ import ( type ssoUserResourceType struct { resourceType *v2.ResourceType ssoClient *awsSsoAdmin.Client - identityStoreClient *awsIdentityStore.Client + identityStoreClient client.IdentityStoreClient identityInstance *awsSsoAdminTypes.InstanceMetadata scimClient *awsIdentityCenterSCIMClient region string @@ -114,7 +115,7 @@ func (o *ssoUserResourceType) Grants(_ context.Context, _ *v2.Resource, _ *pagin func ssoUserBuilder( region string, ssoClient *awsSsoAdmin.Client, - identityStoreClient *awsIdentityStore.Client, + identityStoreClient client.IdentityStoreClient, identityInstance *awsSsoAdminTypes.InstanceMetadata, scimClient *awsIdentityCenterSCIMClient, ) *ssoUserResourceType { From adba34a5c1766325df27b6fb45ad6ae0c88eee76 Mon Sep 17 00:00:00 2001 From: Marcos Gaeta Date: Tue, 27 Aug 2024 18:17:11 -0700 Subject: [PATCH 4/5] add unit tests --- pkg/connector/sso_group_test.go | 217 ++++++++++++++++++++++++++++++++ 1 file changed, 217 insertions(+) create mode 100644 pkg/connector/sso_group_test.go diff --git a/pkg/connector/sso_group_test.go b/pkg/connector/sso_group_test.go new file mode 100644 index 00000000..5bdfa47f --- /dev/null +++ b/pkg/connector/sso_group_test.go @@ -0,0 +1,217 @@ +package connector + +import ( + "context" + "slices" + "strconv" + "testing" + + "github.com/aws/aws-sdk-go-v2/aws" + awsIdentityStore "github.com/aws/aws-sdk-go-v2/service/identitystore" + awsIdentityStoreTypes "github.com/aws/aws-sdk-go-v2/service/identitystore/types" + 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" + v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" + "github.com/conductorone/baton-sdk/pkg/pagination" + "github.com/conductorone/baton-sdk/pkg/test" + "github.com/conductorone/baton-sdk/pkg/types/entitlement" + "github.com/stretchr/testify/require" +) + +const ( + mockSSOGroupID = "arn:aws:identitystore:us-east-1::d-90679d1878/group/9458d408-40b1-709f-4f45-92be754928e5" + mockSSOUserID = "arn:aws:identitystore:us-east-1::d-90679d1878/user/54982488-f0d1-70c1-1dd5-6db47f7add45" + mockMembershipID = "1" +) + +type mockedIdentityStoreClient struct { + awsIdentityStore.Client +} + +var ( + store = map[string][]string{} + hasPermission = true +) + +func (c *mockedIdentityStoreClient) ListGroups( + ctx context.Context, + params *awsIdentityStore.ListGroupsInput, + optFns ...func(*awsIdentityStore.Options), +) ( + *awsIdentityStore.ListGroupsOutput, + error, +) { + groups := make([]awsIdentityStoreTypes.Group, 0) + for groupId := range store { + groups = append(groups, awsIdentityStoreTypes.Group{ + DisplayName: aws.String(groupId), + GroupId: aws.String(groupId), + ExternalIds: []awsIdentityStoreTypes.ExternalId{ + { + Id: aws.String("external id"), + }, + }, + }) + } + return &awsIdentityStore.ListGroupsOutput{Groups: groups}, nil +} + +func (c *mockedIdentityStoreClient) ListGroupMemberships( + ctx context.Context, + params *awsIdentityStore.ListGroupMembershipsInput, + optFns ...func(*awsIdentityStore.Options), +) ( + *awsIdentityStore.ListGroupMembershipsOutput, + error, +) { + var startIndex = 0 + var nextToken = aws.String("") + token := params.NextToken + if token != nil && *token != "" { + parsed, err := strconv.Atoi(*token) + if err != nil { + return nil, err + } + startIndex = parsed + } + + memberships := make([]awsIdentityStoreTypes.GroupMembership, 0) + found, _ := store[*params.GroupId] + for i, id := range found { + if i == startIndex { + memberships = append(memberships, awsIdentityStoreTypes.GroupMembership{ + MembershipId: aws.String(id), + }) + nextToken = aws.String(strconv.Itoa(i + 1)) + break + } + } + + output := awsIdentityStore.ListGroupMembershipsOutput{ + GroupMemberships: memberships, + NextToken: nextToken, + } + return &output, nil +} + +func (c *mockedIdentityStoreClient) CreateGroupMembership( + ctx context.Context, + params *awsIdentityStore.CreateGroupMembershipInput, + optFns ...func(*awsIdentityStore.Options), +) (*awsIdentityStore.CreateGroupMembershipOutput, error) { + groupId := params.GroupId + userId := params.MemberId.(*awsIdentityStoreTypes.MemberIdMemberUserId).Value + found, ok := store[*groupId] + if !ok { + found = []string{} + } + + if slices.Contains(found, userId) { + return nil, &awsIdentityStoreTypes.ConflictException{} + } + + store[*groupId] = append(found, userId) + return &awsIdentityStore.CreateGroupMembershipOutput{ + MembershipId: aws.String(userId), + ResultMetadata: middleware.Metadata{}, + }, nil +} + +func (c *mockedIdentityStoreClient) GetGroupMembershipId( + ctx context.Context, + params *awsIdentityStore.GetGroupMembershipIdInput, + optFns ...func(*awsIdentityStore.Options), +) (*awsIdentityStore.GetGroupMembershipIdOutput, error) { + if hasPermission { + return &awsIdentityStore.GetGroupMembershipIdOutput{ + MembershipId: aws.String(mockMembershipID), + ResultMetadata: middleware.Metadata{}, + }, nil + } + return nil, &awsIdentityStoreTypes.AccessDeniedException{} +} + +func resetStore() { + store = map[string][]string{ + mockSSOGroupID: { + mockMembershipID, + }, + } +} + +func TestSSOGroups(t *testing.T) { + ctx := context.Background() + ssoClient := &awsSsoAdmin.Client{} + identityInstance := &awsSsoAdminTypes.InstanceMetadata{ + IdentityStoreId: aws.String(mockMembershipID), + } + c := ssoGroupBuilder( + "", + ssoClient, + &mockedIdentityStoreClient{}, + identityInstance, + ) + + group := &v2.Resource{ + Id: &v2.ResourceId{ + ResourceType: resourceTypeSSOGroup.Id, + Resource: mockSSOGroupID, + }, + } + user := &v2.Resource{ + Id: &v2.ResourceId{ + ResourceType: resourceTypeSSOUser.Id, + Resource: mockSSOUserID, + }, + } + + entitlement := v2.Entitlement{ + Id: entitlement.NewEntitlementID(group, groupMemberEntitlement), + Resource: group, + } + t.Run("should paginate when listing grants", func(t *testing.T) { + resetStore() + store[mockSSOGroupID] = []string{"1", "2", "3"} + resources := make([]*v2.Grant, 0) + pToken := pagination.Token{ + Token: "", + Size: 1, + } + for { + nextResources, nextToken, listAnnotations, err := c.Grants(ctx, nil, &pToken) + resources = append(resources, nextResources...) + + require.Nil(t, err) + test.AssertNoRatelimitAnnotations(t, listAnnotations) + if nextToken == "" { + break + } + + pToken.Token = nextToken + } + + require.NotNil(t, resources) + require.Len(t, resources, 3) + require.NotEmpty(t, resources[0].Id) + }) + + t.Run("should fallback with GetGroupMembershipId if grant already exists", func(t *testing.T) { + resetStore() + // Create the same grant again. + grantsAgain, grantAnnotations, err := c.Grant(ctx, user, &entitlement) + require.Nil(t, err) + test.AssertNoRatelimitAnnotations(t, grantAnnotations) + require.Len(t, grantsAgain, 1) + }) + + t.Run("should fallback to empty grant list when the client does not have access to GetGroupMembershipId", func(t *testing.T) { + resetStore() + hasPermission = false + // Create the same grant again. + grantsAgain, grantAnnotations, err := c.Grant(ctx, user, &entitlement) + require.Nil(t, err) + test.AssertNoRatelimitAnnotations(t, grantAnnotations) + require.Len(t, grantsAgain, 0) + }) +} From b6c07094be4f69e2831cd911b22bfa785825a899 Mon Sep 17 00:00:00 2001 From: Marcos Gaeta Date: Wed, 28 Aug 2024 10:49:32 -0700 Subject: [PATCH 5/5] fix tests --- pkg/connector/sso_group_test.go | 156 ++++---------------------------- test/mock.go | 150 ++++++++++++++++++++++++++++++ 2 files changed, 166 insertions(+), 140 deletions(-) create mode 100644 test/mock.go diff --git a/pkg/connector/sso_group_test.go b/pkg/connector/sso_group_test.go index 5bdfa47f..2eeb1ffd 100644 --- a/pkg/connector/sso_group_test.go +++ b/pkg/connector/sso_group_test.go @@ -2,167 +2,42 @@ package connector import ( "context" - "slices" - "strconv" "testing" "github.com/aws/aws-sdk-go-v2/aws" - awsIdentityStore "github.com/aws/aws-sdk-go-v2/service/identitystore" - awsIdentityStoreTypes "github.com/aws/aws-sdk-go-v2/service/identitystore/types" 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/test" v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" "github.com/conductorone/baton-sdk/pkg/pagination" - "github.com/conductorone/baton-sdk/pkg/test" + testSdk "github.com/conductorone/baton-sdk/pkg/test" "github.com/conductorone/baton-sdk/pkg/types/entitlement" "github.com/stretchr/testify/require" ) -const ( - mockSSOGroupID = "arn:aws:identitystore:us-east-1::d-90679d1878/group/9458d408-40b1-709f-4f45-92be754928e5" - mockSSOUserID = "arn:aws:identitystore:us-east-1::d-90679d1878/user/54982488-f0d1-70c1-1dd5-6db47f7add45" - mockMembershipID = "1" -) - -type mockedIdentityStoreClient struct { - awsIdentityStore.Client -} - -var ( - store = map[string][]string{} - hasPermission = true -) - -func (c *mockedIdentityStoreClient) ListGroups( - ctx context.Context, - params *awsIdentityStore.ListGroupsInput, - optFns ...func(*awsIdentityStore.Options), -) ( - *awsIdentityStore.ListGroupsOutput, - error, -) { - groups := make([]awsIdentityStoreTypes.Group, 0) - for groupId := range store { - groups = append(groups, awsIdentityStoreTypes.Group{ - DisplayName: aws.String(groupId), - GroupId: aws.String(groupId), - ExternalIds: []awsIdentityStoreTypes.ExternalId{ - { - Id: aws.String("external id"), - }, - }, - }) - } - return &awsIdentityStore.ListGroupsOutput{Groups: groups}, nil -} - -func (c *mockedIdentityStoreClient) ListGroupMemberships( - ctx context.Context, - params *awsIdentityStore.ListGroupMembershipsInput, - optFns ...func(*awsIdentityStore.Options), -) ( - *awsIdentityStore.ListGroupMembershipsOutput, - error, -) { - var startIndex = 0 - var nextToken = aws.String("") - token := params.NextToken - if token != nil && *token != "" { - parsed, err := strconv.Atoi(*token) - if err != nil { - return nil, err - } - startIndex = parsed - } - - memberships := make([]awsIdentityStoreTypes.GroupMembership, 0) - found, _ := store[*params.GroupId] - for i, id := range found { - if i == startIndex { - memberships = append(memberships, awsIdentityStoreTypes.GroupMembership{ - MembershipId: aws.String(id), - }) - nextToken = aws.String(strconv.Itoa(i + 1)) - break - } - } - - output := awsIdentityStore.ListGroupMembershipsOutput{ - GroupMemberships: memberships, - NextToken: nextToken, - } - return &output, nil -} - -func (c *mockedIdentityStoreClient) CreateGroupMembership( - ctx context.Context, - params *awsIdentityStore.CreateGroupMembershipInput, - optFns ...func(*awsIdentityStore.Options), -) (*awsIdentityStore.CreateGroupMembershipOutput, error) { - groupId := params.GroupId - userId := params.MemberId.(*awsIdentityStoreTypes.MemberIdMemberUserId).Value - found, ok := store[*groupId] - if !ok { - found = []string{} - } - - if slices.Contains(found, userId) { - return nil, &awsIdentityStoreTypes.ConflictException{} - } - - store[*groupId] = append(found, userId) - return &awsIdentityStore.CreateGroupMembershipOutput{ - MembershipId: aws.String(userId), - ResultMetadata: middleware.Metadata{}, - }, nil -} - -func (c *mockedIdentityStoreClient) GetGroupMembershipId( - ctx context.Context, - params *awsIdentityStore.GetGroupMembershipIdInput, - optFns ...func(*awsIdentityStore.Options), -) (*awsIdentityStore.GetGroupMembershipIdOutput, error) { - if hasPermission { - return &awsIdentityStore.GetGroupMembershipIdOutput{ - MembershipId: aws.String(mockMembershipID), - ResultMetadata: middleware.Metadata{}, - }, nil - } - return nil, &awsIdentityStoreTypes.AccessDeniedException{} -} - -func resetStore() { - store = map[string][]string{ - mockSSOGroupID: { - mockMembershipID, - }, - } -} - func TestSSOGroups(t *testing.T) { ctx := context.Background() ssoClient := &awsSsoAdmin.Client{} identityInstance := &awsSsoAdminTypes.InstanceMetadata{ - IdentityStoreId: aws.String(mockMembershipID), + IdentityStoreId: aws.String(test.MockMembershipID), } c := ssoGroupBuilder( "", ssoClient, - &mockedIdentityStoreClient{}, + &test.MockedIdentityStoreClient{}, identityInstance, ) group := &v2.Resource{ Id: &v2.ResourceId{ ResourceType: resourceTypeSSOGroup.Id, - Resource: mockSSOGroupID, + Resource: test.MockSSOGroupIDARN, }, } user := &v2.Resource{ Id: &v2.ResourceId{ ResourceType: resourceTypeSSOUser.Id, - Resource: mockSSOUserID, + Resource: test.MockSSOUserID, }, } @@ -170,20 +45,21 @@ func TestSSOGroups(t *testing.T) { Id: entitlement.NewEntitlementID(group, groupMemberEntitlement), Resource: group, } + t.Run("should paginate when listing grants", func(t *testing.T) { - resetStore() - store[mockSSOGroupID] = []string{"1", "2", "3"} + test.ResetMock() + test.SetStore(test.MockSSOGroupID, []string{"1", "2", "3"}) resources := make([]*v2.Grant, 0) pToken := pagination.Token{ Token: "", Size: 1, } for { - nextResources, nextToken, listAnnotations, err := c.Grants(ctx, nil, &pToken) + nextResources, nextToken, listAnnotations, err := c.Grants(ctx, group, &pToken) resources = append(resources, nextResources...) require.Nil(t, err) - test.AssertNoRatelimitAnnotations(t, listAnnotations) + testSdk.AssertNoRatelimitAnnotations(t, listAnnotations) if nextToken == "" { break } @@ -197,21 +73,21 @@ func TestSSOGroups(t *testing.T) { }) t.Run("should fallback with GetGroupMembershipId if grant already exists", func(t *testing.T) { - resetStore() + test.ResetMock() // Create the same grant again. grantsAgain, grantAnnotations, err := c.Grant(ctx, user, &entitlement) require.Nil(t, err) - test.AssertNoRatelimitAnnotations(t, grantAnnotations) + testSdk.AssertNoRatelimitAnnotations(t, grantAnnotations) require.Len(t, grantsAgain, 1) }) t.Run("should fallback to empty grant list when the client does not have access to GetGroupMembershipId", func(t *testing.T) { - resetStore() - hasPermission = false + test.ResetMock() + test.SetPermission(false) // Create the same grant again. grantsAgain, grantAnnotations, err := c.Grant(ctx, user, &entitlement) require.Nil(t, err) - test.AssertNoRatelimitAnnotations(t, grantAnnotations) + testSdk.AssertNoRatelimitAnnotations(t, grantAnnotations) require.Len(t, grantsAgain, 0) }) } diff --git a/test/mock.go b/test/mock.go new file mode 100644 index 00000000..8d9fb2c1 --- /dev/null +++ b/test/mock.go @@ -0,0 +1,150 @@ +package test + +import ( + "context" + "slices" + "strconv" + + "github.com/aws/aws-sdk-go-v2/aws" + awsIdentityStore "github.com/aws/aws-sdk-go-v2/service/identitystore" + awsIdentityStoreTypes "github.com/aws/aws-sdk-go-v2/service/identitystore/types" + "github.com/aws/smithy-go/middleware" +) + +const ( + MockSSOGroupID = "9458d408-40b1-709f-4f45-92be754928e5" + MockSSOGroupIDARN = "arn:aws:identitystore:us-east-1::d-90679d1878/group/9458d408-40b1-709f-4f45-92be754928e5" + MockSSOUserID = "arn:aws:identitystore:us-east-1::d-90679d1878/user/54982488-f0d1-70c1-1dd5-6db47f7add45" + MockMembershipID = "54982488-f0d1-70c1-1dd5-6db47f7add45" +) + +var ( + store = map[string][]string{} + hasPermission = true +) + +func ResetMock() { + store = map[string][]string{ + MockSSOGroupID: { + MockMembershipID, + }, + } + hasPermission = true +} + +func SetStore(key string, value []string) { + store[key] = value +} + +func SetPermission(state bool) { + hasPermission = state +} + +type MockedIdentityStoreClient struct { + awsIdentityStore.Client +} + +func (c *MockedIdentityStoreClient) ListGroups( + ctx context.Context, + params *awsIdentityStore.ListGroupsInput, + optFns ...func(*awsIdentityStore.Options), +) ( + *awsIdentityStore.ListGroupsOutput, + error, +) { + groups := make([]awsIdentityStoreTypes.Group, 0) + for groupId := range store { + groups = append(groups, awsIdentityStoreTypes.Group{ + DisplayName: aws.String(groupId), + GroupId: aws.String(groupId), + ExternalIds: []awsIdentityStoreTypes.ExternalId{ + { + Id: aws.String("external id"), + }, + }, + }) + } + return &awsIdentityStore.ListGroupsOutput{Groups: groups}, nil +} + +func (c *MockedIdentityStoreClient) ListGroupMemberships( + ctx context.Context, + params *awsIdentityStore.ListGroupMembershipsInput, + optFns ...func(*awsIdentityStore.Options), +) ( + *awsIdentityStore.ListGroupMembershipsOutput, + error, +) { + var startIndex = 0 + var nextToken = aws.String("") + token := params.NextToken + if token != nil && *token != "" { + parsed, err := strconv.Atoi(*token) + if err != nil { + return nil, err + } + startIndex = parsed + } + + memberships := make([]awsIdentityStoreTypes.GroupMembership, 0) + found := store[*params.GroupId] + for i, id := range found { + if i == startIndex { + memberships = append( + memberships, + awsIdentityStoreTypes.GroupMembership{ + MembershipId: aws.String(id), + MemberId: &awsIdentityStoreTypes.MemberIdMemberUserId{ + Value: id, + }, + }, + ) + nextToken = aws.String(strconv.Itoa(i + 1)) + break + } + } + + output := awsIdentityStore.ListGroupMembershipsOutput{ + GroupMemberships: memberships, + NextToken: nextToken, + } + return &output, nil +} + +func (c *MockedIdentityStoreClient) CreateGroupMembership( + ctx context.Context, + params *awsIdentityStore.CreateGroupMembershipInput, + optFns ...func(*awsIdentityStore.Options), +) (*awsIdentityStore.CreateGroupMembershipOutput, error) { + groupId := params.GroupId + userId := params.MemberId.(*awsIdentityStoreTypes.MemberIdMemberUserId).Value + found, ok := store[*groupId] + if !ok { + found = []string{} + } + + if slices.Contains(found, userId) { + return nil, &awsIdentityStoreTypes.ConflictException{} + } + + store[*groupId] = append(found, userId) + return &awsIdentityStore.CreateGroupMembershipOutput{ + MembershipId: aws.String(userId), + ResultMetadata: middleware.Metadata{}, + }, nil +} + +func (c *MockedIdentityStoreClient) GetGroupMembershipId( + ctx context.Context, + params *awsIdentityStore.GetGroupMembershipIdInput, + optFns ...func(*awsIdentityStore.Options), +) (*awsIdentityStore.GetGroupMembershipIdOutput, error) { + if hasPermission { + return &awsIdentityStore.GetGroupMembershipIdOutput{ + MembershipId: aws.String(MockMembershipID), + ResultMetadata: middleware.Metadata{}, + }, nil + } + + return nil, &awsIdentityStoreTypes.AccessDeniedException{} +}