fix: Fail jobs configured to run as worker agent when Admin (#230)

Signed-off-by: Sakshi Sakshi <[email protected]>
Signed-off-by: Brian Axelson <[email protected]>

src/deadline_worker_agent/scheduler/scheduler.py

@@ -708,6 +708,22 @@ def _create_new_sessions(
             # Requires some updates to the code below
             try:
                 job_details = job_entities.job_details()
+
+                # For Windows the WA runs as Administrator so fail jobs that were configured to runAs - WORKER_AGENT_USER as that would provide Admin privileges to the job
+                if (
+                    os.name == "nt"
+                    and job_details.job_run_as_user
+                    and job_details.job_run_as_user.is_worker_agent_user
+                ):
+                    err_msg = (
+                        "Job cannot run as WORKER_AGENT_USER as it has administrator privileges."
+                    )
+                    self._fail_all_actions(session_spec, err_msg)
+                    logger.warning("[%s] %s", new_session_id, err_msg)
+                    # Force an immediate UpdateWorkerSchedule request
+                    self._wakeup.set()
+                    continue
+
             except (ValueError, RuntimeError) as error:
                 # Can't even start a session right now if we don't
                 # get valid job_details, so let's fail the actions

src/deadline_worker_agent/sessions/job_entities/job_details.py

@@ -93,7 +93,8 @@ def job_run_as_user_api_model_to_worker_agent(
     expected by the Worker Agent.
     """
     if "runAs" in job_run_as_user_data and job_run_as_user_data["runAs"] == "WORKER_AGENT_USER":
-        return None
+        job_run_as_user = JobRunAsUser(is_worker_agent_user=True)
+        return job_run_as_user
 
     if os.name == "posix":
         user = ""
@@ -154,6 +155,7 @@ class JobRunAsUser:
     posix: PosixSessionUser | None = None
     windows: WindowsSessionUser | None = None
     windows_settings: JobRunAsWindowsUser | None = None
+    is_worker_agent_user: bool = False
 
     def __eq__(self, other: Any) -> bool:
         if other is None:

test/unit/scheduler/test_scheduler.py

@@ -24,6 +24,7 @@
     UPDATE_WORKER_SCHEDULE_MAX_MESSAGE_CHARS,
 )
 from deadline_worker_agent.scheduler.session_action_status import SessionActionStatus
+from deadline_worker_agent.sessions.job_entities.job_details import JobDetails, JobRunAsUser
 from deadline_worker_agent.startup.config import JobsRunAsUserOverride
 from deadline_worker_agent.errors import ServiceShutdown
 import deadline_worker_agent.scheduler.scheduler as scheduler_mod
@@ -34,6 +35,7 @@
     DeadlineRequestInterrupted,
 )
 from deadline_worker_agent.file_system_operations import FileSystemPermissionEnum
+from openjd.model import SpecificationRevision
 
 
 @pytest.fixture
@@ -790,6 +792,79 @@ def test_job_details_error(
             assert action_update.start_time == datetime_now_mock.return_value
             assert action_update.end_time == datetime_now_mock.return_value
 
+    @pytest.mark.skipif(os.name != "nt", reason="Windows-only test.")
+    def test_job_details_run_as_worker_agent_user_windows(
+        self,
+        scheduler: WorkerScheduler,
+    ) -> None:
+        """Tests that when a session encounters a runAs: WORKER_AGENT_USER for Windows os,
+        the first assigned action is marked as FAILED, the rest are marked as NEVER_ATTEPTED,
+        and the scheduler's wakeup event is set so that it makes an
+        immediate follow-up UpdateWorkerSchedule request to signal the failure.
+        """
+        # GIVEN
+        queue_id = "queue-abcdef0123456789abcdef0123456789"
+        session_id = "session-abcdef0123456789abcdef0123456789"
+        assigned_sessions: dict[str, AssignedSession] = {
+            session_id: AssignedSession(
+                queueId=queue_id,
+                jobId="job-abcdef0123456789abcdef0123456789",
+                logConfiguration=LogConfiguration(
+                    logDriver="awslogs",
+                    options={},
+                    parameters={"interval": "15"},
+                ),
+                sessionActions=[
+                    EnvironmentAction(
+                        actionType="ENV_ENTER",
+                        environmentId="env-1",
+                        sessionActionId="action-1",
+                    ),
+                    TaskRunAction(
+                        actionType="TASK_RUN",
+                        parameters={},
+                        sessionActionId="action-2",
+                        stepId="step-1",
+                        taskId="task-1",
+                    ),
+                ],
+            ),
+        }
+        expected_err_msg = "Job cannot run as WORKER_AGENT_USER as it has administrator privileges."
+
+        job_entity_mock = MagicMock()
+        job_entity_mock.job_details.return_value = JobDetails(
+            log_group_name="/aws/deadline/queue-0000",
+            schema_version=SpecificationRevision.v2023_09,
+            job_run_as_user=JobRunAsUser(is_worker_agent_user=True),
+        )
+
+        with (
+            patch.object(scheduler_mod, "datetime") as datetime_mock,
+            patch.object(scheduler_mod, "JobEntities") as job_entities_mock,
+        ):
+            job_entities_mock.return_value = job_entity_mock
+            datetime_now_mock: MagicMock = datetime_mock.now
+
+            # WHEN
+            scheduler._create_new_sessions(assigned_sessions=assigned_sessions)
+
+        # THEN
+        for action_num in (1, 2):
+            action_id = f"action-{action_num}"
+            assert (
+                action_update := scheduler._action_updates_map.get(action_id, None)
+            ), f"no action update for {action_id}"
+            assert action_update.id == action_id
+            assert action_update.completed_status == (
+                "FAILED" if action_num == 1 else "NEVER_ATTEMPTED"
+            )
+            assert action_update.status is not None
+            assert action_update.status.state == ActionState.FAILED
+            assert action_update.status.fail_message == expected_err_msg
+            assert action_update.start_time == datetime_now_mock.return_value
+            assert action_update.end_time == datetime_now_mock.return_value
+
 
 class TestQueueAwsCredentialsManagement:
     """Tests that validate that we are constructing and destroying credentials objects

test/unit/sessions/job_entities/test_job_details.py

@@ -49,6 +49,15 @@ def job_details_no_user() -> JobDetails:
     )
 
 
+@pytest.fixture
+def job_details_only_run_as_worker_agent_user() -> JobDetails:
+    return JobDetails(
+        log_group_name="/aws/deadline/queue-0000",
+        schema_version=SpecificationRevision.v2023_09,
+        job_run_as_user=JobRunAsUser(is_worker_agent_user=True),
+    )
+
+
 @pytest.mark.parametrize(
     "data",
     [
@@ -413,7 +422,7 @@ def test_input_validation_success(data: dict[str, Any]) -> None:
                     "runAs": "WORKER_AGENT_USER",
                 },
             },
-            "job_details_no_user",
+            "job_details_only_run_as_worker_agent_user",
             id="required with runAs WORKER_AGENT_USER",
         ),
     ],