diff --git a/cookbooks/aws-parallelcluster-platform/files/dcv/pcluster_dcv_authenticator.py b/cookbooks/aws-parallelcluster-platform/files/dcv/pcluster_dcv_authenticator.py index 5ca30726f..5014b9f28 100644 --- a/cookbooks/aws-parallelcluster-platform/files/dcv/pcluster_dcv_authenticator.py +++ b/cookbooks/aws-parallelcluster-platform/files/dcv/pcluster_dcv_authenticator.py @@ -361,13 +361,17 @@ def _is_session_valid(user, session_id): # TODO change this method if DCV updates his behaviour. """ logger.info("Verifying NICE DCV session validity..") + + # Query by uid rather than username to avoid truncation by ps command + uid = subprocess.check_output(["/usr/bin/id", "-u", user]).decode("utf-8").strip() # nosec B603 + # Remove the first and the last because they are the heading and empty, respectively # All commands and arguments in this subprocess call are built as literals - processes = subprocess.check_output(["/bin/ps", "aux"]).decode("utf-8").split("\n")[1:-1] # nosec B603 + processes = subprocess.check_output(["/bin/ps", "aunx"]).decode("utf-8").split("\n")[1:-1] # nosec B603 # Check the filter is empty if not next( - filter(lambda process: DCVAuthenticator.check_dcv_process(process, user, session_id), processes), None + filter(lambda process: DCVAuthenticator.check_dcv_process(process, uid, session_id), processes), None ): raise DCVAuthenticator.IncorrectRequestError("The given session does not exists") logger.info("The NICE DCV session is valid.") @@ -377,21 +381,21 @@ def _verify_session_existence(user, session_id): retry(DCVAuthenticator._is_session_valid, func_args=[user, session_id], attempts=20, wait=1) @staticmethod - def check_dcv_process(row, user, session_id): + def check_dcv_process(row, uid, session_id): """Check if there is a dcvagent process running for the given user and for the given session_id.""" # row example: - # centos 63 0.0 0.0 4348844 3108 ?? Ss 23Jul19 2:32.46 /usr/libexec/dcv/dcvagent --mode full \ + # 1000 63 0.0 0.0 4348844 3108 ?? Ss 23Jul19 2:32.46 /usr/libexec/dcv/dcvagent --mode full \ # --session-id mysession - # ubuntu 2949 0.3 0.4 860568 34328 ? Sl 20:10 0:18 /usr/lib/x86_64-linux-gnu/dcv/dcvagent --mode full \ + # 2000 2949 0.3 0.4 860568 34328 ? Sl 20:10 0:18 /usr/lib/x86_64-linux-gnu/dcv/dcvagent --mode full \ # --session-id mysession fields = row.split() command_index = 10 session_name_index = 14 - user_index = 0 + uid_index = 0 return ( fields[command_index].endswith("/dcv/dcvagent") - and fields[user_index] == user + and fields[uid_index] == uid and fields[session_name_index] == session_id ) diff --git a/test/unit/dcv/test_dcv_authenticator.py b/test/unit/dcv/test_dcv_authenticator.py index 1d0d3f6f7..e01c7195f 100644 --- a/test/unit/dcv/test_dcv_authenticator.py +++ b/test/unit/dcv/test_dcv_authenticator.py @@ -163,6 +163,26 @@ def test_get_request_token_parameter(parameters, keys, result): DCVAuthenticator._extract_parameters_values(parameters, keys) +def test_is_session_valid(mocker): + ps_aunx_output = ( + b"USER PID %CPU %MEM VSZ RSS TTY STAT START TIME COMMAND\n" + b"1000 63 0.0 0.0 4 3108 ?? Ss 23Jul19 2:32.46 /dcv/dcvagent --mode full --session-id sid\n" + b"2000 2949 0.3 0.4 8 34328 ? Sl 20:10 0:18 /dcv/dcvagent --mode full --session-id sid\n" + ) + # Mock subprocess.check_output with realistic responses for `/usr/bin/id` and `/bin/ps aunx` + mocker.patch(AUTH_MODULE_MOCK_PATH + "subprocess.check_output", side_effect=[b"1000", ps_aunx_output]) + + # Test that the session is valid + DCVAuthenticator._is_session_valid("myuser", "sid") + + # Mock subprocess.check_output with realistic responses for `/usr/bin/id` and `/bin/ps aunx` + mocker.patch(AUTH_MODULE_MOCK_PATH + "subprocess.check_output", side_effect=[b"1000", ps_aunx_output]) + + # Test that the session is not valid + with pytest.raises(DCVAuthenticator.IncorrectRequestError): + DCVAuthenticator._is_session_valid("myuser", "wrongsession") + + def test_get_request_token(mocker): """Verify the first step of the authentication process, the retrieval of the Request Token.""" # A nosec comment is appended to the following line in order to disable the B105 check.