Add grant/revoke to CI tests. (#31)

* Add grant/revoke to CI tests.

* Extract request id correctly.

* Add request id annotation to all grant/revoke calls.

* Check that grants were granted/revoked correctly.

.github/workflows/ci.yaml

@@ -52,13 +52,42 @@ jobs:
           BATON_GLOBAL_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
           BATON_GLOBAL_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
         run: ./baton-aws
-      # - name: Revoke grants
-      #   env:
-      #     BATON_GLOBAL_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
-      #     BATON_GLOBAL_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
-      #   run: ./baton-aws --revoke-grant 'group:arn:aws:iam::425848093043:group/ci-test-group:member:iam_user:arn:aws:iam::425848093043:user/ci-test-user'
-      # - name: Grant entitlements
-      #   env:
-      #     BATON_GLOBAL_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
-      #     BATON_GLOBAL_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
-      #   run: ./baton-aws --grant-entitlement 'group:arn:aws:iam::425848093043:group/ci-test-group:member' --grant-principal 'arn:aws:iam::425848093043:user/ci-test-user' --grant-principal-type 'iam_user'
+      - name: Checkout baton
+        uses: actions/checkout@v3
+        with:
+          repository: ConductorOne/baton
+          ref: main
+          path: ./baton
+      - name: Build baton
+        working-directory: ./baton
+        run: go build ./cmd/baton
+      - name: Grant entitlements
+        env:
+          BATON_GLOBAL_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          BATON_GLOBAL_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+        run: ./baton-aws --grant-entitlement 'group:arn:aws:iam::425848093043:group/ci-test-group:member' --grant-principal 'arn:aws:iam::425848093043:user/ci-test-user' --grant-principal-type 'iam_user'
+      - name: Check for grant before revoking
+        env:
+          BATON_GLOBAL_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          BATON_GLOBAL_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+        run: ./baton-aws && ./baton/baton grants --entitlement='group:arn:aws:iam::425848093043:group/ci-test-group:member' --output-format=json | jq --exit-status '.grants[].principal.id.resource == "arn:aws:iam::425848093043:user/ci-test-user"'
+      - name: Revoke grants
+        env:
+          BATON_GLOBAL_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          BATON_GLOBAL_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+        run: ./baton-aws --revoke-grant 'group:arn:aws:iam::425848093043:group/ci-test-group:member:iam_user:arn:aws:iam::425848093043:user/ci-test-user'
+      - name: Check grant was revoked
+        env:
+          BATON_GLOBAL_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          BATON_GLOBAL_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+        run: ./baton-aws && ./baton/baton grants --entitlement='group:arn:aws:iam::425848093043:group/ci-test-group:member' --output-format=json | jq --exit-status 'if .grants then .grants[]?.principal.id.resource != "arn:aws:iam::425848093043:user/ci-test-user" else . end'
+      - name: Grant entitlements
+        env:
+          BATON_GLOBAL_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          BATON_GLOBAL_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+        run: ./baton-aws --grant-entitlement 'group:arn:aws:iam::425848093043:group/ci-test-group:member' --grant-principal 'arn:aws:iam::425848093043:user/ci-test-user' --grant-principal-type 'iam_user'
+      - name: Check grant was re-granted
+        env:
+          BATON_GLOBAL_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          BATON_GLOBAL_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+        run: ./baton-aws && ./baton/baton grants --entitlement='group:arn:aws:iam::425848093043:group/ci-test-group:member' --output-format=json | jq --exit-status '.grants[].principal.id.resource == "arn:aws:iam::425848093043:user/ci-test-user"'

go.mod

@@ -11,6 +11,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/organizations v1.22.3
 	github.com/aws/aws-sdk-go-v2/service/ssoadmin v1.22.1
 	github.com/aws/aws-sdk-go-v2/service/sts v1.25.4
+	github.com/aws/smithy-go v1.17.0
 	github.com/conductorone/baton-sdk v0.1.13
 	github.com/grpc-ecosystem/go-grpc-middleware v1.4.0
 	github.com/spf13/cobra v1.8.0
@@ -33,7 +34,6 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/s3 v1.44.0 // indirect
 	github.com/aws/aws-sdk-go-v2/service/sso v1.17.3 // indirect
 	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.20.1 // indirect
-	github.com/aws/smithy-go v1.17.0 // indirect
 	github.com/benbjohnson/clock v1.3.5 // indirect
 	github.com/doug-martin/goqu/v9 v9.19.0 // indirect
 	github.com/dustin/go-humanize v1.0.1 // indirect

go.sum

@@ -97,8 +97,6 @@ github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDk
 github.com/cncf/udpa/go v0.0.0-20191209042840-269d4d468f6f/go.mod h1:M8M6+tZqaGXZJjfX53e64911xZQV5JYwmTeXPW+k8Sc=
 github.com/cncf/udpa/go v0.0.0-20200629203442-efcf912fb354/go.mod h1:WmhPx2Nbnhtbo57+VJT5O0JRkEi1Wbu0z5j0R8u5Hbk=
 github.com/cncf/udpa/go v0.0.0-20201120205902-5459f2c99403/go.mod h1:WmhPx2Nbnhtbo57+VJT5O0JRkEi1Wbu0z5j0R8u5Hbk=
-github.com/conductorone/baton-sdk v0.1.13-0.20231129230901-6a4866e4e0ec h1:kslx2qju4psh9Em6TUhZaJXxpc9aLqzaL9ScwqgmnVA=
-github.com/conductorone/baton-sdk v0.1.13-0.20231129230901-6a4866e4e0ec/go.mod h1:mxMxyna5UNmh9T+uQL4UApk0Ers0unkL1SGGnwjoCoc=
 github.com/conductorone/baton-sdk v0.1.13 h1:FO+HzH32TSH+CragU5R/dG+07nEescHatbc+D5Sol8Y=
 github.com/conductorone/baton-sdk v0.1.13/go.mod h1:mxMxyna5UNmh9T+uQL4UApk0Ers0unkL1SGGnwjoCoc=
 github.com/cpuguy83/go-md2man/v2 v2.0.3/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=

pkg/connector/account.go

@@ -290,6 +290,11 @@ func (o *accountResourceType) Grant(ctx context.Context, principal *v2.Resource,
 		return nil, err
 	}
 
+	annos := annotations.New()
+	if reqId := extractRequestID(&createOut.ResultMetadata); reqId != nil {
+		annos.Append(reqId)
+	}
+
 	l := ctxzap.Extract(ctx).With(
 		zap.String("request_id", awsSdk.ToString(createOut.AccountAssignmentCreationStatus.RequestId)),
 		zap.String("principal_id", awsSdk.ToString(createOut.AccountAssignmentCreationStatus.PrincipalId)),
@@ -325,7 +330,7 @@ func (o *accountResourceType) Grant(ctx context.Context, principal *v2.Resource,
 		}
 	}
 
-	return nil, nil
+	return annos, nil
 }
 
 // checkCreateAccountAssignmentStatus checks the status of the account assignment creation request. It returns true if the request is complete, false if it is still in progress.
@@ -425,6 +430,11 @@ func (o *accountResourceType) Revoke(ctx context.Context, grant *v2.Grant) (anno
 		return nil, err
 	}
 
+	annos := annotations.New()
+	if reqId := extractRequestID(&deleteOut.ResultMetadata); reqId != nil {
+		annos.Append(reqId)
+	}
+
 	l := ctxzap.Extract(ctx).With(
 		zap.String("request_id", awsSdk.ToString(deleteOut.AccountAssignmentDeletionStatus.RequestId)),
 		zap.String("principal_id", awsSdk.ToString(deleteOut.AccountAssignmentDeletionStatus.PrincipalId)),
@@ -460,7 +470,7 @@ func (o *accountResourceType) Revoke(ctx context.Context, grant *v2.Grant) (anno
 		}
 	}
 
-	return nil, nil
+	return annos, nil
 }
 
 func (o *accountResourceType) getPermissionSet(ctx context.Context, permissionSetId string) (*awsSsoAdminTypes.PermissionSet, error) {

pkg/connector/helpers.go

@@ -5,6 +5,7 @@ import (
 	"strings"
 
 	"github.com/aws/aws-sdk-go-v2/aws/arn"
+	aws_middleware "github.com/aws/aws-sdk-go-v2/aws/middleware"
 	"github.com/aws/smithy-go/middleware"
 	v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2"
 	"github.com/conductorone/baton-sdk/pkg/annotations"
@@ -128,12 +129,12 @@ func extractRequestID(md *middleware.Metadata) proto.Message {
 		return nil
 	}
 
-	reqId, ok := md.Get("RequestId").(string)
-	if !ok {
-		return nil
+	reqId, hasReqId := aws_middleware.GetRequestIDMetadata(*md)
+	if hasReqId {
+		return &v2.RequestId{
+			RequestId: reqId,
+		}
 	}
 
-	return &v2.RequestId{
-		RequestId: reqId,
-	}
+	return nil
 }

pkg/connector/iam_group.go

@@ -181,12 +181,12 @@ func (o *iamGroupResourceType) Grant(ctx context.Context, principal *v2.Resource
 		),
 	)
 
-	rv := annotations.New()
+	annos := annotations.New()
 	if reqId := extractRequestID(&resp.ResultMetadata); reqId != nil {
-		rv.Append(reqId)
+		annos.Append(reqId)
 	}
 
-	return []*v2.Grant{grant}, nil, nil
+	return []*v2.Grant{grant}, annos, nil
 }
 
 func (o *iamGroupResourceType) Revoke(ctx context.Context, grant *v2.Grant) (annotations.Annotations, error) {
@@ -212,10 +212,10 @@ func (o *iamGroupResourceType) Revoke(ctx context.Context, grant *v2.Grant) (ann
 		return nil, fmt.Errorf("baton-aws: error removing iam user from iam group: %w", err)
 	}
 
-	rv := annotations.New()
+	annos := annotations.New()
 	if reqId := extractRequestID(&resp.ResultMetadata); reqId != nil {
-		rv.Append(reqId)
+		annos.Append(reqId)
 	}
 
-	return rv, nil
+	return annos, nil
 }

pkg/connector/sso_group.go

@@ -225,7 +225,13 @@ func (g *ssoGroupResourceType) Grant(ctx context.Context, principal *v2.Resource
 		l.Error("aws-connector: Failed to create grant", zap.Error(err), zap.String("membership_id", awsSdk.ToString(membership.MembershipId)))
 		return nil, nil, err
 	}
-	return []*v2.Grant{grant}, nil, nil
+
+	annos := annotations.New()
+	if reqId := extractRequestID(&membership.ResultMetadata); reqId != nil {
+		annos.Append(reqId)
+	}
+
+	return []*v2.Grant{grant}, annos, nil
 }
 func (g *ssoGroupResourceType) Revoke(ctx context.Context, grant *v2.Grant) (annotations.Annotations, error) {
 	if grant.Principal.Id.ResourceType != resourceTypeSSOUser.Id {
@@ -237,15 +243,21 @@ func (g *ssoGroupResourceType) Revoke(ctx context.Context, grant *v2.Grant) (ann
 		zap.String("identity_store_id", awsSdk.ToString(g.identityInstance.IdentityStoreId)),
 	)
 
-	if _, err := g.identityStoreClient.DeleteGroupMembership(ctx, &awsIdentityStore.DeleteGroupMembershipInput{
+	resp, err := g.identityStoreClient.DeleteGroupMembership(ctx, &awsIdentityStore.DeleteGroupMembershipInput{
 		IdentityStoreId: g.identityInstance.IdentityStoreId,
 		MembershipId:    awsSdk.String(grant.Id),
-	}); err != nil {
+	})
+	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)
 	}
 
-	return nil, nil
+	annos := annotations.New()
+	if reqId := extractRequestID(&resp.ResultMetadata); reqId != nil {
+		annos.Append(reqId)
+	}
+
+	return annos, nil
 }
 
 func ssoGroupProfile(ctx context.Context, group awsIdentityStoreTypes.Group) map[string]interface{} {