From bce8d710c5bb23491bae492cf43efcd295bac1d1 Mon Sep 17 00:00:00 2001
From: Geoff Greer <geoff@greer.fm>
Date: Fri, 6 Sep 2024 10:18:21 -0700
Subject: [PATCH] Test double-revoking grants. (#57)

* Test double-revoking grants.

* Try to fetch the group membership by principal and entitlement ids if grant id is unsuccessful.

* Return success and already revoked annotation if grant was already revoked.

* Change this log from an error to info because it's not that bad of a problem.
---
 pkg/connector/sso_group.go | 78 ++++++++++++++++++++++++++++++++------
 test/grant-revoke.sh       |  5 ++-
 2 files changed, 71 insertions(+), 12 deletions(-)

diff --git a/pkg/connector/sso_group.go b/pkg/connector/sso_group.go
index efcf792..a7791cd 100644
--- a/pkg/connector/sso_group.go
+++ b/pkg/connector/sso_group.go
@@ -202,6 +202,23 @@ type GroupMembershipOutput struct {
 	ResultMetadata middleware.Metadata
 }
 
+func (g *ssoGroupResourceType) getGroupMembership(ctx context.Context, groupId string, userId string) (*awsIdentityStore.GetGroupMembershipIdOutput, error) {
+	groupIdString := awsSdk.String(groupId)
+	memberId := awsIdentityStoreTypes.MemberIdMemberUserId{Value: userId}
+
+	getInput := awsIdentityStore.GetGroupMembershipIdInput{
+		GroupId:         groupIdString,
+		IdentityStoreId: g.identityInstance.IdentityStoreId,
+		MemberId:        &memberId,
+	}
+	foundMembership, err := g.identityStoreClient.GetGroupMembershipId(ctx, &getInput)
+	if err != nil {
+		return nil, err
+	}
+
+	return foundMembership, nil
+}
+
 // createOrGetMembership the `CreateGroupMembership()` method errors when a
 // group membership already exists. Instead of passing along the
 // `ConflictError`, attempt to get the group membership with a call to
@@ -249,12 +266,7 @@ func (g *ssoGroupResourceType) createOrGetMembership(
 
 	logger.Info("ConflictException when creating group, falling back to GET")
 
-	getInput := awsIdentityStore.GetGroupMembershipIdInput{
-		GroupId:         groupIdString,
-		IdentityStoreId: g.identityInstance.IdentityStoreId,
-		MemberId:        &memberId,
-	}
-	foundMembership, err := g.identityStoreClient.GetGroupMembershipId(ctx, &getInput)
+	foundMembership, err := g.getGroupMembership(ctx, groupID, userID)
 	if err != nil {
 		// If we lack permission for the `GetGroupMembershipId` operation, fail
 		// more gracefully by returning nil.
@@ -339,7 +351,7 @@ func (g *ssoGroupResourceType) Grant(
 }
 func (g *ssoGroupResourceType) Revoke(ctx context.Context, grant *v2.Grant) (annotations.Annotations, error) {
 	if grant.Principal.Id.ResourceType != resourceTypeSSOUser.Id {
-		return nil, errors.New("baton-aws: only sso users can be removed from sso group")
+		return nil, errors.New("baton-aws: only sso users can be removed from sso groups")
 	}
 
 	l := ctxzap.Extract(ctx).With(
@@ -347,19 +359,63 @@ func (g *ssoGroupResourceType) Revoke(ctx context.Context, grant *v2.Grant) (ann
 		zap.String("identity_store_id", awsSdk.ToString(g.identityInstance.IdentityStoreId)),
 	)
 
+	annos := annotations.New()
+	membershipId := grant.Id
+
 	resp, err := g.identityStoreClient.DeleteGroupMembership(
 		ctx,
 		&awsIdentityStore.DeleteGroupMembershipInput{
 			IdentityStoreId: g.identityInstance.IdentityStoreId,
-			MembershipId:    awsSdk.String(grant.Id),
+			MembershipId:    awsSdk.String(membershipId),
+		},
+	)
+	if err == nil {
+		l.Debug("revoked grant", zap.String("membership_id", membershipId))
+		if reqId := extractRequestID(&resp.ResultMetadata); reqId != nil {
+			annos.Append(reqId)
+		}
+
+		return annos, nil
+	}
+
+	l.Info("aws-connector: Failed to delete group membership. Trying to fetch group membership in case grant ID is incorrect", zap.Error(err))
+	groupId, err := ssoGroupIdFromARN(grant.Entitlement.Resource.Id.Resource)
+	if err != nil {
+		return annos, err
+	}
+
+	userId, err := ssoUserIdFromARN(grant.Principal.Id.Resource)
+	if err != nil {
+		return annos, err
+	}
+
+	foundMembership, getErr := g.getGroupMembership(ctx, groupId, userId)
+	if getErr != nil {
+		var notFoundException *awsIdentityStoreTypes.ResourceNotFoundException
+		if errors.As(getErr, &notFoundException) {
+			l.Debug("group membership already deleted", zap.String("group_id", groupId), zap.String("user_id", userId))
+			annos.Append(&v2.GrantAlreadyRevoked{})
+			return annos, nil
+		}
+
+		l.Error("aws-connector: Failed to get group membership", zap.Error(getErr))
+		return nil, fmt.Errorf("baton-aws: error removing sso user from sso group: %w %w", err, getErr)
+	}
+
+	membershipId = *foundMembership.MembershipId
+	resp, err = g.identityStoreClient.DeleteGroupMembership(
+		ctx,
+		&awsIdentityStore.DeleteGroupMembershipInput{
+			IdentityStoreId: g.identityInstance.IdentityStoreId,
+			MembershipId:    awsSdk.String(membershipId),
 		},
 	)
 	if err != nil {
-		l.Error("aws-connector: Failed to delete group membership", zap.Error(err))
-		return nil, fmt.Errorf("baton-aws: error removing sso user from sso group: %w", err)
+		l.Error("aws-connector: Failed to delete group membership", zap.Error(err), zap.String("membership_id", membershipId))
+		return nil, err
 	}
 
-	annos := annotations.New()
+	l.Debug("revoked grant", zap.String("membership_id", membershipId))
 	if reqId := extractRequestID(&resp.ResultMetadata); reqId != nil {
 		annos.Append(reqId)
 	}
diff --git a/test/grant-revoke.sh b/test/grant-revoke.sh
index 8e676da..ac17994 100755
--- a/test/grant-revoke.sh
+++ b/test/grant-revoke.sh
@@ -31,7 +31,10 @@ $BATON_AWS --grant-entitlement="$BATON_ENTITLEMENT" --grant-principal="$BATON_PR
 # Get grant ID
 BATON_GRANT=$($BATON grants --entitlement="$BATON_ENTITLEMENT" --output-format=json | jq --raw-output --exit-status ".grants[] | select( .principal.id.resource == \"$BATON_PRINCIPAL\" ).grant.id")
 
-# Revoke grants
+# Revoke grant
+$BATON_AWS --revoke-grant="$BATON_GRANT"
+
+# Revoke already-revoked grant
 $BATON_AWS --revoke-grant="$BATON_GRANT"
 
 # Check grant was revoked