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 016f57a2..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" @@ -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 } 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"), + } +) diff --git a/pkg/connector/sso_group.go b/pkg/connector/sso_group.go index a1f2603e..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{ @@ -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 { diff --git a/pkg/connector/sso_group_test.go b/pkg/connector/sso_group_test.go new file mode 100644 index 00000000..2eeb1ffd --- /dev/null +++ b/pkg/connector/sso_group_test.go @@ -0,0 +1,93 @@ +package connector + +import ( + "context" + "testing" + + "github.com/aws/aws-sdk-go-v2/aws" + 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/test" + v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" + "github.com/conductorone/baton-sdk/pkg/pagination" + testSdk "github.com/conductorone/baton-sdk/pkg/test" + "github.com/conductorone/baton-sdk/pkg/types/entitlement" + "github.com/stretchr/testify/require" +) + +func TestSSOGroups(t *testing.T) { + ctx := context.Background() + ssoClient := &awsSsoAdmin.Client{} + identityInstance := &awsSsoAdminTypes.InstanceMetadata{ + IdentityStoreId: aws.String(test.MockMembershipID), + } + c := ssoGroupBuilder( + "", + ssoClient, + &test.MockedIdentityStoreClient{}, + identityInstance, + ) + + group := &v2.Resource{ + Id: &v2.ResourceId{ + ResourceType: resourceTypeSSOGroup.Id, + Resource: test.MockSSOGroupIDARN, + }, + } + user := &v2.Resource{ + Id: &v2.ResourceId{ + ResourceType: resourceTypeSSOUser.Id, + Resource: test.MockSSOUserID, + }, + } + + entitlement := v2.Entitlement{ + Id: entitlement.NewEntitlementID(group, groupMemberEntitlement), + Resource: group, + } + + t.Run("should paginate when listing grants", func(t *testing.T) { + 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, group, &pToken) + resources = append(resources, nextResources...) + + require.Nil(t, err) + testSdk.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) { + test.ResetMock() + // Create the same grant again. + grantsAgain, grantAnnotations, err := c.Grant(ctx, user, &entitlement) + require.Nil(t, err) + 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) { + test.ResetMock() + test.SetPermission(false) + // Create the same grant again. + grantsAgain, grantAnnotations, err := c.Grant(ctx, user, &entitlement) + require.Nil(t, err) + testSdk.AssertNoRatelimitAnnotations(t, grantAnnotations) + require.Len(t, grantsAgain, 0) + }) +} 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 { 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{} +}