Merge pull request #172 from aws-samples/wip/tmscarla/configurable-oidc-roles

feat: make user roles claim in the jwt configurable, using 'cognito:groups' as default

Author: tmscarla

README.md

@@ -116,6 +116,15 @@ npm run dev
 
 Then navigate to [http://localhost:3000](http://localhost:3000)
 
+## Testing
+
+Launch tests of the API backend by running:
+
+```bash
+pytest
+```
+For detailed information on how to invoke `pytest`, see this [resource](https://docs.pytest.org/en/7.1.x/how-to/usage.html). 
+
 ## Security
 
 See [CONTRIBUTING](CONTRIBUTING.md#security-issue-notifications) for more information.

api/PclusterApiHandler.py

@@ -33,6 +33,7 @@
 SECRET_ID = os.getenv("SECRET_ID")
 ENABLE_MFA = os.getenv("ENABLE_MFA")
 SITE_URL = os.getenv("SITE_URL", API_BASE_URL)
+USER_ROLES_CLAIM = os.getenv("USER_ROLES_CLAIM", "cognito:groups")
 
 try:
     if (not USER_POOL_ID or USER_POOL_ID == "") and SECRET_ID:
@@ -137,7 +138,7 @@ def authenticate(group):
         return auth_redirect()
     except jose.exceptions.JWSSignatureError:
         return logout()
-    if not disable_auth() and (group != "guest") and (group not in set(decoded.get("cognito:groups", []))):
+    if not disable_auth() and (group != "guest") and (group not in set(decoded.get(USER_ROLES_CLAIM, []))):
         return auth_redirect()
 
 
@@ -540,15 +541,24 @@ def get_instance_types():
     return {"instance_types": sorted(instance_types, key=lambda x: x["InstanceType"])}
 
 
+def _get_user_roles(decoded):
+    print(os.environ.get("USER_ROLES_CLAIM"))
+    return decoded[USER_ROLES_CLAIM] if USER_ROLES_CLAIM in decoded else ["user"]
+
+
+
 def get_identity():
     if running_local():
-        return {"cognito:groups": ["user", "admin"], "username": "username", "attributes": {"email": "[email protected]"}}
+        return {"user_roles": ["user", "admin"], "username": "username", "attributes": {"email": "[email protected]"}}
 
     access_token = request.cookies.get("accessToken")
     if not access_token:
         return {"message": "No access token."}, 401
     try:
         decoded = jwt_decode(access_token, USER_POOL_ID)
+        decoded["user_roles"] = _get_user_roles(decoded)
+        decoded.pop(USER_ROLES_CLAIM)
+
         username = decoded.get("username")
         if username:
             cognito = boto3.client("cognito-idp")
@@ -559,7 +569,7 @@ def get_identity():
         return {"message": "Signature expired."}, 401
 
     if disable_auth():
-        decoded["cognito:groups"] = ["user", "admin"]
+        decoded["user_roles"] = ["user", "admin"]
 
     return decoded
 

api/tests/__init__.py

@@ -0,0 +1,20 @@
+import pytest
+from unittest import mock
+from api.PclusterApiHandler import _get_user_roles
+
+
+@mock.patch("api.PclusterApiHandler.USER_ROLES_CLAIM", "user_roles")
+def test_user_roles():
+    user_roles = ["user", "admin"]
+
+    _test_decoded_with_user_roles_claim(decoded={"user_roles": user_roles}, user_roles=user_roles)
+    _test_decoded_without_user_roles_claim(decoded={})
+
+
+
+def _test_decoded_with_user_roles_claim(decoded, user_roles):
+    assert _get_user_roles(decoded) == user_roles
+
+
+def _test_decoded_without_user_roles_claim(decoded):
+    assert _get_user_roles(decoded) == ["user"]

api/tests/test_pcluster_api_handler.py

@@ -0,0 +1 @@
+export const USER_ROLES_CLAIM = "user_roles"

frontend/src/auth/constants.ts

@@ -10,7 +10,7 @@
 // limitations under the License.
 import * as React from 'react';
 import { Link, useLocation } from "react-router-dom"
-import { setState, useState, isAdmin} from '../store'
+import { setState, useState, isGuest, isUser, isAdmin } from '../store'
 
 // UI Elements
 import Divider from '@mui/material/Divider';
@@ -24,18 +24,8 @@ import HomeIcon from '@mui/icons-material/Home';
 import GroupIcon from '@mui/icons-material/Group';
 
 export function SideBarIcons(props) {
-  let identity = useState(['identity']);
-  let groups = useState(['identity', 'cognito:groups']) || [];
   const drawerOpen = useState(['app', 'sidebar', 'drawerOpen']);
-
-  const isGuest = () => {
-    return identity && (!groups || ((!groups.includes("admin")) && (!groups.includes("user"))));
-  }
-
-  const isUser = () => {
-    return groups && ((groups.includes("admin")) || (groups.includes("user")));
-  }
-
+  
   const location = useLocation();
   let defaultPage = isGuest() ? "home" : "clusters";
   let section = location && location.pathname && location.pathname.substring(1);
@@ -72,18 +62,7 @@ export function SideBarIcons(props) {
 }
 
 export default function SideBar(props) {
-  let identity = useState(['identity']);
-  let groups = useState(['identity', 'cognito:groups']);
   const drawerOpen = useState(['app', 'sidebar', 'drawerOpen']);
-
-  const isGuest = () => {
-    return identity && (!groups || ((!groups.includes("admin")) && (!groups.includes("user"))));
-  }
-
-  const isUser = () => {
-    return groups && ((groups.includes("admin")) || (groups.includes("user")));
-  }
-
   useNotifier();
 
   const location = useLocation();

frontend/src/components/SideBar.js

@@ -12,6 +12,7 @@ import axios from 'axios'
 import { store, setState, getState, clearState, updateState, clearAllState } from './store'
 import { enqueueSnackbar as enqueueSnackbarAction } from './redux/snackbar_actions';
 import { closeSnackbar as closeSnackbarAction } from './redux/snackbar_actions';
+import { USER_ROLES_CLAIM } from './auth/constants';
 
 // UI Elements
 import Button from '@mui/material/Button';
@@ -768,7 +769,7 @@ function LoadInitialState() {
   clearAllState();
   GetVersion();
   GetIdentity((identity) => {
-    let groups = identity['cognito:groups'];
+    let groups = identity[USER_ROLES_CLAIM];
     if(groups && (groups.includes("admin") || groups.includes("user")))
     {
       ListUsers();

frontend/src/model.js

@@ -30,14 +30,12 @@ import CssBaseline from '@mui/material/CssBaseline';
 // Components
 import Loading from '../components/Loading'
 
+import { isGuest } from '../store';
+
+
 export default function App() {
   const identity = useState(['identity']);
-  const groups = useState(['identity', 'cognito:groups']);
-  
-  const isGuest = () => {
-    return identity && (!groups || ((!groups.includes("admin")) && (!groups.includes("user"))));
-  }
-  
+
   React.useEffect(() => {
     LoadInitialState();
   }, [])

frontend/src/pages/index.tsx

@@ -12,6 +12,7 @@ import { createStore, combineReducers } from '@reduxjs/toolkit'
 import { useSelector } from 'react-redux'
 import notifications from './redux/snackbar_reducer'
 import { getIn, swapIn } from './util'
+import { USER_ROLES_CLAIM } from './auth/constants'
 
 // These are identity reducers that allow the names to be at the top level for combining
 function clusters(state = {}, action){ return state }
@@ -143,12 +144,12 @@ function useState(path) {
 }
 
 function isAdmin() {
-  let groups = getState(['identity', 'cognito:groups']) || [];
+  let groups = getState(['identity', USER_ROLES_CLAIM]) || [];
   return groups && groups.includes("admin");
 }
 
 function isUser() {
-  let groups = getState(['identity', 'cognito:groups']) || [];
+  let groups = getState(['identity', USER_ROLES_CLAIM]) || [];
   return groups && ((groups.includes("admin")) || (groups.includes("user")));
 }
 

frontend/src/store.js

@@ -4,3 +4,4 @@ boto3
 requests
 python-jose
 pyyaml
+pytest