From f4397e089513e16db5f363458c3c61b351acb5b9 Mon Sep 17 00:00:00 2001
From: Ross Williams <ross.williams@s12solutions.com>
Date: Fri, 20 Sep 2019 17:50:57 +0100
Subject: [PATCH] feat(amplify-category-auth): allow more than one groupClaim

---
 .../OperationsArgument.test.ts.snap           | 130 ++++++++----------
 .../graphql-auth-transformer/src/constants.ts |   3 +-
 .../graphql-auth-transformer/src/resources.ts | 111 +++++++--------
 .../ModelAuthTransformer.e2e.test.ts          |   2 +-
 4 files changed, 109 insertions(+), 137 deletions(-)

diff --git a/packages/graphql-auth-transformer/src/__tests__/__snapshots__/OperationsArgument.test.ts.snap b/packages/graphql-auth-transformer/src/__tests__/__snapshots__/OperationsArgument.test.ts.snap
index 47366f314ef..2f9158006f7 100644
--- a/packages/graphql-auth-transformer/src/__tests__/__snapshots__/OperationsArgument.test.ts.snap
+++ b/packages/graphql-auth-transformer/src/__tests__/__snapshots__/OperationsArgument.test.ts.snap
@@ -13,16 +13,15 @@ exports[`Test "create", "update", "delete" auth operations 3`] = `
 ## [Start] Check authMode and execute owner/group checks **
 #if( $authMode == \\"userPools\\" )
   ## [Start] Static Group Authorization Checks **
-  ## Authorization rule: { allow: groups, groups: \\"[\\"Admin\\",\\"Dev\\"]\\" } **
+  #set($isStaticGroupAuthorized = $util.defaultIfNull(
+            $isStaticGroupAuthorized, false))
+  ## Authorization rule: { allow: groups, groups: [\\"Admin\\",\\"Dev\\"], groupClaim: \\"cognito:groups\\" } **
   #set( $userGroups = $util.defaultIfNull($ctx.identity.claims.get(\\"cognito:groups\\"), []) )
   #set( $allowedGroups = [\\"Admin\\", \\"Dev\\"] )
-  #set($isStaticGroupAuthorized = $util.defaultIfNull(
-                $isStaticGroupAuthorized, false))
   #foreach( $userGroup in $userGroups )
-    #foreach( $allowedGroup in $allowedGroups )
-      #if( $allowedGroup == $userGroup )
-        #set( $isStaticGroupAuthorized = true )
-      #end
+    #if( $allowedGroups.contains($userGroup) )
+      #set( $isStaticGroupAuthorized = true )
+      #break
     #end
   #end
   ## [End] Static Group Authorization Checks **
@@ -72,16 +71,15 @@ exports[`Test "create", "update", "delete" auth operations 4`] = `
 ## [Start] Check authMode and execute owner/group checks **
 #if( $authMode == \\"userPools\\" )
   ## [Start] Static Group Authorization Checks **
-  ## Authorization rule: { allow: groups, groups: \\"[\\"Admin\\",\\"Dev\\"]\\" } **
+  #set($isStaticGroupAuthorized = $util.defaultIfNull(
+            $isStaticGroupAuthorized, false))
+  ## Authorization rule: { allow: groups, groups: [\\"Admin\\",\\"Dev\\"], groupClaim: \\"cognito:groups\\" } **
   #set( $userGroups = $util.defaultIfNull($ctx.identity.claims.get(\\"cognito:groups\\"), []) )
   #set( $allowedGroups = [\\"Admin\\", \\"Dev\\"] )
-  #set($isStaticGroupAuthorized = $util.defaultIfNull(
-                $isStaticGroupAuthorized, false))
   #foreach( $userGroup in $userGroups )
-    #foreach( $allowedGroup in $allowedGroups )
-      #if( $allowedGroup == $userGroup )
-        #set( $isStaticGroupAuthorized = true )
-      #end
+    #if( $allowedGroups.contains($userGroup) )
+      #set( $isStaticGroupAuthorized = true )
+      #break
     #end
   #end
   ## [End] Static Group Authorization Checks **
@@ -285,16 +283,15 @@ exports[`Test "create", "update", "delete" auth operations 5`] = `
 ## [Start] Check authMode and execute owner/group checks **
 #if( $authMode == \\"userPools\\" )
   ## [Start] Static Group Authorization Checks **
-  ## Authorization rule: { allow: groups, groups: \\"[\\"Admin\\",\\"Dev\\"]\\" } **
+  #set($isStaticGroupAuthorized = $util.defaultIfNull(
+            $isStaticGroupAuthorized, false))
+  ## Authorization rule: { allow: groups, groups: [\\"Admin\\",\\"Dev\\"], groupClaim: \\"cognito:groups\\" } **
   #set( $userGroups = $util.defaultIfNull($ctx.identity.claims.get(\\"cognito:groups\\"), []) )
   #set( $allowedGroups = [\\"Admin\\", \\"Dev\\"] )
-  #set($isStaticGroupAuthorized = $util.defaultIfNull(
-                $isStaticGroupAuthorized, false))
   #foreach( $userGroup in $userGroups )
-    #foreach( $allowedGroup in $allowedGroups )
-      #if( $allowedGroup == $userGroup )
-        #set( $isStaticGroupAuthorized = true )
-      #end
+    #if( $allowedGroups.contains($userGroup) )
+      #set( $isStaticGroupAuthorized = true )
+      #break
     #end
   #end
   ## [End] Static Group Authorization Checks **
@@ -433,16 +430,15 @@ exports[`Test that checks subscription resolvers are generated with auth logic 1
 ## [Start] Check authMode and execute owner/group checks **
 #if( $authMode == \\"userPools\\" )
   ## [Start] Static Group Authorization Checks **
-  ## Authorization rule: { allow: groups, groups: \\"[\\"Admin\\"]\\" } **
+  #set($isStaticGroupAuthorized = $util.defaultIfNull(
+            $isStaticGroupAuthorized, false))
+  ## Authorization rule: { allow: groups, groups: [\\"Admin\\"], groupClaim: \\"cognito:groups\\" } **
   #set( $userGroups = $util.defaultIfNull($ctx.identity.claims.get(\\"cognito:groups\\"), []) )
   #set( $allowedGroups = [\\"Admin\\"] )
-  #set($isStaticGroupAuthorized = $util.defaultIfNull(
-                $isStaticGroupAuthorized, false))
   #foreach( $userGroup in $userGroups )
-    #foreach( $allowedGroup in $allowedGroups )
-      #if( $allowedGroup == $userGroup )
-        #set( $isStaticGroupAuthorized = true )
-      #end
+    #if( $allowedGroups.contains($userGroup) )
+      #set( $isStaticGroupAuthorized = true )
+      #break
     #end
   #end
   ## [End] Static Group Authorization Checks **
@@ -489,16 +485,15 @@ exports[`Test that checks subscription resolvers are generated with auth logic 2
 ## [Start] Check authMode and execute owner/group checks **
 #if( $authMode == \\"userPools\\" )
   ## [Start] Static Group Authorization Checks **
-  ## Authorization rule: { allow: groups, groups: \\"[\\"Admin\\"]\\" } **
+  #set($isStaticGroupAuthorized = $util.defaultIfNull(
+            $isStaticGroupAuthorized, false))
+  ## Authorization rule: { allow: groups, groups: [\\"Admin\\"], groupClaim: \\"cognito:groups\\" } **
   #set( $userGroups = $util.defaultIfNull($ctx.identity.claims.get(\\"cognito:groups\\"), []) )
   #set( $allowedGroups = [\\"Admin\\"] )
-  #set($isStaticGroupAuthorized = $util.defaultIfNull(
-                $isStaticGroupAuthorized, false))
   #foreach( $userGroup in $userGroups )
-    #foreach( $allowedGroup in $allowedGroups )
-      #if( $allowedGroup == $userGroup )
-        #set( $isStaticGroupAuthorized = true )
-      #end
+    #if( $allowedGroups.contains($userGroup) )
+      #set( $isStaticGroupAuthorized = true )
+      #break
     #end
   #end
   ## [End] Static Group Authorization Checks **
@@ -545,16 +540,15 @@ exports[`Test that checks subscription resolvers are generated with auth logic 3
 ## [Start] Check authMode and execute owner/group checks **
 #if( $authMode == \\"userPools\\" )
   ## [Start] Static Group Authorization Checks **
-  ## Authorization rule: { allow: groups, groups: \\"[\\"Admin\\"]\\" } **
+  #set($isStaticGroupAuthorized = $util.defaultIfNull(
+            $isStaticGroupAuthorized, false))
+  ## Authorization rule: { allow: groups, groups: [\\"Admin\\"], groupClaim: \\"cognito:groups\\" } **
   #set( $userGroups = $util.defaultIfNull($ctx.identity.claims.get(\\"cognito:groups\\"), []) )
   #set( $allowedGroups = [\\"Admin\\"] )
-  #set($isStaticGroupAuthorized = $util.defaultIfNull(
-                $isStaticGroupAuthorized, false))
   #foreach( $userGroup in $userGroups )
-    #foreach( $allowedGroup in $allowedGroups )
-      #if( $allowedGroup == $userGroup )
-        #set( $isStaticGroupAuthorized = true )
-      #end
+    #if( $allowedGroups.contains($userGroup) )
+      #set( $isStaticGroupAuthorized = true )
+      #break
     #end
   #end
   ## [End] Static Group Authorization Checks **
@@ -605,16 +599,15 @@ exports[`Test that operation overwrites queries in auth operations 3`] = `
 ## [Start] Check authMode and execute owner/group checks **
 #if( $authMode == \\"userPools\\" )
   ## [Start] Static Group Authorization Checks **
-  ## Authorization rule: { allow: groups, groups: \\"[\\"Admin\\",\\"Dev\\"]\\" } **
+  #set($isStaticGroupAuthorized = $util.defaultIfNull(
+            $isStaticGroupAuthorized, false))
+  ## Authorization rule: { allow: groups, groups: [\\"Admin\\",\\"Dev\\"], groupClaim: \\"cognito:groups\\" } **
   #set( $userGroups = $util.defaultIfNull($ctx.identity.claims.get(\\"cognito:groups\\"), []) )
   #set( $allowedGroups = [\\"Admin\\", \\"Dev\\"] )
-  #set($isStaticGroupAuthorized = $util.defaultIfNull(
-                $isStaticGroupAuthorized, false))
   #foreach( $userGroup in $userGroups )
-    #foreach( $allowedGroup in $allowedGroups )
-      #if( $allowedGroup == $userGroup )
-        #set( $isStaticGroupAuthorized = true )
-      #end
+    #if( $allowedGroups.contains($userGroup) )
+      #set( $isStaticGroupAuthorized = true )
+      #break
     #end
   #end
   ## [End] Static Group Authorization Checks **
@@ -664,16 +657,15 @@ exports[`Test that operation overwrites queries in auth operations 4`] = `
 ## [Start] Check authMode and execute owner/group checks **
 #if( $authMode == \\"userPools\\" )
   ## [Start] Static Group Authorization Checks **
-  ## Authorization rule: { allow: groups, groups: \\"[\\"Admin\\",\\"Dev\\"]\\" } **
+  #set($isStaticGroupAuthorized = $util.defaultIfNull(
+            $isStaticGroupAuthorized, false))
+  ## Authorization rule: { allow: groups, groups: [\\"Admin\\",\\"Dev\\"], groupClaim: \\"cognito:groups\\" } **
   #set( $userGroups = $util.defaultIfNull($ctx.identity.claims.get(\\"cognito:groups\\"), []) )
   #set( $allowedGroups = [\\"Admin\\", \\"Dev\\"] )
-  #set($isStaticGroupAuthorized = $util.defaultIfNull(
-                $isStaticGroupAuthorized, false))
   #foreach( $userGroup in $userGroups )
-    #foreach( $allowedGroup in $allowedGroups )
-      #if( $allowedGroup == $userGroup )
-        #set( $isStaticGroupAuthorized = true )
-      #end
+    #if( $allowedGroups.contains($userGroup) )
+      #set( $isStaticGroupAuthorized = true )
+      #break
     #end
   #end
   ## [End] Static Group Authorization Checks **
@@ -877,16 +869,15 @@ exports[`Test that operation overwrites queries in auth operations 5`] = `
 ## [Start] Check authMode and execute owner/group checks **
 #if( $authMode == \\"userPools\\" )
   ## [Start] Static Group Authorization Checks **
-  ## Authorization rule: { allow: groups, groups: \\"[\\"Admin\\",\\"Dev\\"]\\" } **
+  #set($isStaticGroupAuthorized = $util.defaultIfNull(
+            $isStaticGroupAuthorized, false))
+  ## Authorization rule: { allow: groups, groups: [\\"Admin\\",\\"Dev\\"], groupClaim: \\"cognito:groups\\" } **
   #set( $userGroups = $util.defaultIfNull($ctx.identity.claims.get(\\"cognito:groups\\"), []) )
   #set( $allowedGroups = [\\"Admin\\", \\"Dev\\"] )
-  #set($isStaticGroupAuthorized = $util.defaultIfNull(
-                $isStaticGroupAuthorized, false))
   #foreach( $userGroup in $userGroups )
-    #foreach( $allowedGroup in $allowedGroups )
-      #if( $allowedGroup == $userGroup )
-        #set( $isStaticGroupAuthorized = true )
-      #end
+    #if( $allowedGroups.contains($userGroup) )
+      #set( $isStaticGroupAuthorized = true )
+      #break
     #end
   #end
   ## [End] Static Group Authorization Checks **
@@ -1019,16 +1010,15 @@ exports[`Test that subscriptions are only generated if the respective mutation o
 ## [Start] Check authMode and execute owner/group checks **
 #if( $authMode == \\"userPools\\" )
   ## [Start] Static Group Authorization Checks **
-  ## Authorization rule: { allow: groups, groups: \\"[\\"Moderator\\"]\\" } **
+  #set($isStaticGroupAuthorized = $util.defaultIfNull(
+            $isStaticGroupAuthorized, false))
+  ## Authorization rule: { allow: groups, groups: [\\"Moderator\\"], groupClaim: \\"cognito:groups\\" } **
   #set( $userGroups = $util.defaultIfNull($ctx.identity.claims.get(\\"cognito:groups\\"), []) )
   #set( $allowedGroups = [\\"Moderator\\"] )
-  #set($isStaticGroupAuthorized = $util.defaultIfNull(
-                $isStaticGroupAuthorized, false))
   #foreach( $userGroup in $userGroups )
-    #foreach( $allowedGroup in $allowedGroups )
-      #if( $allowedGroup == $userGroup )
-        #set( $isStaticGroupAuthorized = true )
-      #end
+    #if( $allowedGroups.contains($userGroup) )
+      #set( $isStaticGroupAuthorized = true )
+      #break
     #end
   #end
   ## [End] Static Group Authorization Checks **
diff --git a/packages/graphql-auth-transformer/src/constants.ts b/packages/graphql-auth-transformer/src/constants.ts
index f49eb93e146..1024b55c015 100644
--- a/packages/graphql-auth-transformer/src/constants.ts
+++ b/packages/graphql-auth-transformer/src/constants.ts
@@ -3,6 +3,7 @@ export const DEFAULT_OWNER_FIELD = "owner"
 export const DEFAULT_IDENTITY_FIELD = "username"
 export const GROUPS_AUTH_STRATEGY = "groups"
 export const DEFAULT_GROUPS_FIELD = "groups"
+export const DEFAULT_GROUP_CLAIM = "cognito:groups"
 export const ON_CREATE_FIELD = "onCreate"
 export const ON_UPDATE_FIELD = "onUpdate"
-export const ON_DELETE_FIELD = "onDelete"
\ No newline at end of file
+export const ON_DELETE_FIELD = "onDelete"
diff --git a/packages/graphql-auth-transformer/src/resources.ts b/packages/graphql-auth-transformer/src/resources.ts
index 1b65620b909..e150c0dfb7d 100644
--- a/packages/graphql-auth-transformer/src/resources.ts
+++ b/packages/graphql-auth-transformer/src/resources.ts
@@ -17,7 +17,8 @@ import { FieldDefinitionNode } from 'graphql';
 import {
     DEFAULT_OWNER_FIELD,
     DEFAULT_IDENTITY_FIELD,
-    DEFAULT_GROUPS_FIELD
+    DEFAULT_GROUPS_FIELD,
+    DEFAULT_GROUP_CLAIM
 } from './constants'
 
 function replaceIfUsername(identityClaim: string): string {
@@ -229,40 +230,36 @@ export class ResourceFactory {
         if (!rules || rules.length === 0) {
             return comment(`No Static Group Authorization Rules`)
         }
-        const allowedGroups: string[] = []
-        let customClaim: string;
+        const variableToSet = ResourceConstants.SNIPPETS.IsStaticGroupAuthorizedVariable;
+        let groupAuthorizationExpressions = []
         for (const rule of rules) {
             const groups = rule.groups;
-            for (const group of groups) {
-                if (group) {
-                    allowedGroups.push(group);
-                }
-            }
-            if (rule.groupClaim) {
-                if (customClaim) {
-                    throw new InvalidDirectiveError(`@auth directive currently only supports one source for groupClaim.
-  - Identified '${customClaim}' and '${rule.groupClaim}'`)
-                }
-                customClaim = rule.groupClaim;
+            const groupClaimAttribute = rule.groupClaim || DEFAULT_GROUP_CLAIM
+           
+            if (groups) {
+                groupAuthorizationExpressions = groupAuthorizationExpressions.concat(
+                    comment(`Authorization rule: { allow: groups, groups: ${JSON.stringify(groups)}, groupClaim: "${groupClaimAttribute}" }`),
+                    this.setUserGroups(rule.groupClaim),
+                    set(ref('allowedGroups'), list(groups.map(s => str(s)))),
+                    forEach(ref('userGroup'), ref('userGroups'), [
+                        iff(
+                            raw(`$allowedGroups.contains($userGroup)`),
+                            compoundExpression([
+                                set(ref(variableToSet), raw('true')),
+                                raw('#break')
+                            ])
+                        )
+                    ])
+                );
             }
         }
         const staticGroupAuthorizedVariable = this.getStaticAuthorizationVariable(field);
 
+        // tslint:disable-next-line
         return block('Static Group Authorization Checks', [
-            comment(`Authorization rule: { allow: groups, groups: "${JSON.stringify(allowedGroups)}" }`),
-            this.setUserGroups(customClaim),
-            set(ref('allowedGroups'), list(allowedGroups.map(s => str(s)))),
-            // tslint:disable-next-line
-            raw(`#set($${staticGroupAuthorizedVariable} = $util.defaultIfNull(
-                $${staticGroupAuthorizedVariable}, false))`),
-            forEach(ref('userGroup'), ref('userGroups'), [
-                forEach(ref('allowedGroup'), ref('allowedGroups'), [
-                    iff(
-                        raw('$allowedGroup == $userGroup'),
-                        set(ref(staticGroupAuthorizedVariable), raw('true'))
-                    )
-                ])
-            ])
+        raw(`#set($${staticGroupAuthorizedVariable} = $util.defaultIfNull(
+            $${staticGroupAuthorizedVariable}, false))`),
+            ...groupAuthorizationExpressions
         ])
     }
 
@@ -307,7 +304,7 @@ export class ResourceFactory {
         let groupAuthorizationExpression: Expression = this.dynamicAuthorizationExpressionForCreate(
             rules, variableToCheck, variableToSet,
             rule => `Authorization rule on field "${fieldToCheck}": { allow: ${rule.allow}, \
-groupsField: "${rule.groupsField || DEFAULT_GROUPS_FIELD}" }`
+groupsField: "${rule.groupsField || DEFAULT_GROUPS_FIELD}", groupClaim: "${rule.groupClaim || DEFAULT_GROUP_CLAIM}" }`
         )
         return block(`Dynamic group authorization rules for field "${fieldToCheck}"`, [
             groupAuthorizationExpression
@@ -321,20 +318,15 @@ groupsField: "${rule.groupsField || DEFAULT_GROUPS_FIELD}" }`
         formatComment?: (rule: AuthRule) => string,
     ) {
         let groupAuthorizationExpressions = []
-        let customClaim: string;
         for (const rule of rules) {
-            if (rule.groupClaim) {
-                if (customClaim) {
-                    throw new InvalidDirectiveError('@auth directive currently only supports one source for groupClaim!')
-                }
-                customClaim = rule.groupClaim;
-            }
             // for loop do check of rules here
             const groupsAttribute = rule.groupsField || DEFAULT_GROUPS_FIELD
+            const groupClaimAttribute = rule.groupClaim || DEFAULT_GROUP_CLAIM
             groupAuthorizationExpressions = groupAuthorizationExpressions.concat(
                 formatComment ?
                     comment(formatComment(rule)) :
-                    comment(`Authorization rule: { allow: ${rule.allow}, groupsField: "${groupsAttribute}" }`),
+                    comment(`Authorization rule: { allow: ${rule.allow}, groupsField: "${groupsAttribute}", groupClaim: "${groupClaimAttribute}"`),
+                this.setUserGroups(rule.groupClaim),
                 set(
                     ref(variableToSet),
                     raw(`$util.defaultIfNull($${variableToSet}, false)`)
@@ -358,11 +350,7 @@ groupsField: "${rule.groupsField || DEFAULT_GROUPS_FIELD}" }`
             )
         }
 
-        // adds group claim
-        return compoundExpression([
-            this.setUserGroups(customClaim),
-            ...groupAuthorizationExpressions,
-        ])
+        return compoundExpression(groupAuthorizationExpressions)
     }
 
     /**
@@ -586,20 +574,15 @@ identityClaim: "${rule.identityField || rule.identityClaim || DEFAULT_IDENTITY_F
 
         let groupAuthorizationExpressions = []
         let ruleNumber = 0
-        let customClaim: string;
         for (const rule of rules) {
-            if (rule.groupClaim) {
-                if (customClaim) {
-                    throw new InvalidDirectiveError('@auth directive currently only supports one source for groupClaim!')
-                }
-                customClaim = rule.groupClaim;
-            }
             const groupsAttribute = rule.groupsField || DEFAULT_GROUPS_FIELD
             const groupsAttributeName = fieldBeingProtected ? `${fieldBeingProtected}_groupsAttribute${ruleNumber}` : `groupsAttribute${ruleNumber}`
             const groupName = fieldBeingProtected ? `${fieldBeingProtected}_group${ruleNumber}` : `group${ruleNumber}`
+            const groupClaimAttribute = rule.groupClaim || DEFAULT_GROUP_CLAIM
             groupAuthorizationExpressions = groupAuthorizationExpressions.concat(
-                comment(`Authorization rule${fieldMention}: { allow: ${rule.allow}, groupsField: "${groupsAttribute}" }`),
+                comment(`Authorization rule${fieldMention}: { allow: ${rule.allow}, groupsField: "${groupsAttribute}", groupClaim: "${groupClaimAttribute}"}`),
                 // Add the new auth expression and values
+                this.setUserGroups(rule.groupClaim),
                 forEach(ref('userGroup'), ref('userGroups'), [
                     raw(`$util.qr($groupAuthExpressions.add("contains(#${groupsAttributeName}, :${groupName}$foreach.count)"))`),
                     raw(`$util.qr($groupAuthExpressionValues.put(":${groupName}$foreach.count", { "S": $userGroup }))`),
@@ -610,7 +593,6 @@ identityClaim: "${rule.identityField || rule.identityClaim || DEFAULT_IDENTITY_F
         }
         // check for groupclaim here
         return block('Dynamic group authorization checks', [
-            this.setUserGroups(customClaim),
             set(ref('groupAuthExpressions'), list([])),
             set(ref('groupAuthExpressionValues'), obj({})),
             set(ref('groupAuthExpressionNames'), obj({})),
@@ -693,18 +675,13 @@ identityClaim: "${rule.identityField || rule.identityClaim || DEFAULT_IDENTITY_F
             return comment(`No Dynamic Group Authorization Rules`)
         }
         let groupAuthorizationExpressions = [];
-        let customClaim: string;
         for (const rule of rules) {
-            if (rule.groupClaim) {
-                if (customClaim) {
-                    throw new InvalidDirectiveError('@auth directive currently only supports one source for groupClaim!')
-                }
-                customClaim = rule.groupClaim;
-            }
             const groupsAttribute = rule.groupsField || DEFAULT_GROUPS_FIELD
+            const groupClaimAttribute = rule.groupClaim || DEFAULT_GROUP_CLAIM
             groupAuthorizationExpressions = groupAuthorizationExpressions.concat(
-                comment(`Authorization rule: { allow: ${rule.allow}, groupsField: "${groupsAttribute}" }`),
+                comment(`Authorization rule: { allow: ${rule.allow}, groupsField: "${groupsAttribute}", groupClaim: "${groupClaimAttribute}" }`),
                 set(ref('allowedGroups'), ref(`util.defaultIfNull($${variableToCheck}.${groupsAttribute}, [])`)),
+                this.setUserGroups(rule.groupClaim),
                 forEach(ref('userGroup'), ref('userGroups'), [
                     iff(
                         raw('$util.isList($allowedGroups)'),
@@ -723,7 +700,6 @@ identityClaim: "${rule.identityField || rule.identityClaim || DEFAULT_IDENTITY_F
         }
         // check for group claim here
         return block('Dynamic Group Authorization Checks', [
-            this.setUserGroups(customClaim),
             set(ref(variableToSet), defaultValue),
             ...groupAuthorizationExpressions,
         ])
@@ -899,15 +875,20 @@ identityClaim: "${rule.identityField || rule.identityClaim || DEFAULT_IDENTITY_F
 
     public setUserGroups(customGroup?: string): Expression {
         if (customGroup) {
-            return block( `Using groupClaim: ${customGroup} as source for userGroup`, [
-                set(ref('userGroup'), raw(`$util.defaultIfNull($ctx.identity.claims.get("${customGroup}"), [])`)),
+            return compoundExpression([
+                set(ref('userGroups'), raw(`$util.defaultIfNull($ctx.identity.claims.get("${customGroup}"), [])`)),
                 iff(
-                    raw('$util.isString($userGroup)'),
-                    set(ref('userGroup'), raw('[$userGroup]')),
+                    raw('$util.isString($userGroups)'),
+                    ifElse(raw('$util.isList($util.parseJson($userGroups))'),
+                        set(ref('userGroups'), raw('$util.parseJson($userGroups)')),
+                        set(ref('userGroups'), raw('[$userGroups]'))
+                    )
+                        
+                
                 ),
             ]);
         }
-        return set(ref('userGroups'), raw('$util.defaultIfNull($ctx.identity.claims.get("cognito:groups"), [])'));
+        return set(ref('userGroups'), raw(`$util.defaultIfNull($ctx.identity.claims.get("${DEFAULT_GROUP_CLAIM}"), [])`));
     }
 
     public generateSubscriptionResolver(fieldName: string, subscriptionTypeName: string = 'Subscription') {
diff --git a/packages/graphql-transformers-e2e-tests/src/__tests__/ModelAuthTransformer.e2e.test.ts b/packages/graphql-transformers-e2e-tests/src/__tests__/ModelAuthTransformer.e2e.test.ts
index 2f9ce76b4bf..123aaebdb60 100644
--- a/packages/graphql-transformers-e2e-tests/src/__tests__/ModelAuthTransformer.e2e.test.ts
+++ b/packages/graphql-transformers-e2e-tests/src/__tests__/ModelAuthTransformer.e2e.test.ts
@@ -113,7 +113,7 @@ beforeAll(async () => {
     }
     type AdminNote @model @auth(
         rules: [
-            {allow: groups, groups: ["Admin"]}
+            {allow: groups, groups: ["Admin"], groupClaim: "cognito:groups"}
         ]
     ) {
         id: ID!