[develop] check dcv sessions using uid to avoid username truncation

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
         )
 

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.