address feedback

Signed-off-by: Ryan S <267728323+ironcommit@users.noreply.github.com>

Author: ironcommit

packages/nemo_platform_ext/src/nemo_platform_ext/auth/workload_exchange.py

@@ -16,12 +16,12 @@
 from urllib.parse import urlparse
 
 import httpx
-from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR
-from nmp.common.auth import (
+from nemo_platform_plugin.client.constants import (
     DOCKER_OPAQUE_WORKLOAD_PROOF_TOKEN_TYPE as _DOCKER_OPAQUE_WORKLOAD_PROOF_TOKEN_TYPE,
 )
-from nmp.common.auth import (
+from nemo_platform_plugin.client.constants import (
     JWT_WORKLOAD_SUBJECT_TOKEN_TYPE,
+    WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR,
     subject_token_type_for_exchange,
 )
 

packages/nemo_platform_ext/tests/auth/test_workload_exchange.py

@@ -3,12 +3,15 @@
 
 """Tests for RFC 8693 workload identity token exchange."""
 
+import ast
 import json
 import time
 from base64 import urlsafe_b64encode
+from pathlib import Path
 from unittest.mock import MagicMock, patch
 
 import httpx
+import nemo_platform_ext.auth.workload_exchange as workload_exchange_module
 import pytest
 from nemo_platform_ext.auth.workload_exchange import (
     ACCESS_TOKEN_TYPE,
@@ -24,6 +27,19 @@
 from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR
 
 
+def test_workload_exchange_module_has_no_nmp_common_dependency():
+    assert workload_exchange_module.__file__ is not None
+    tree = ast.parse(Path(workload_exchange_module.__file__).read_text(encoding="utf-8"))
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Import):
+            imported = [alias.name for alias in node.names]
+        elif isinstance(node, ast.ImportFrom):
+            imported = [node.module or ""]
+        else:
+            continue
+        assert not any(name == "nmp.common" or name.startswith("nmp.common.") for name in imported)
+
+
 def _make_jwt(claims: dict) -> str:
     header = {"alg": "RS256", "typ": "JWT"}
     h = urlsafe_b64encode(json.dumps(header).encode()).rstrip(b"=").decode()

packages/nemo_platform_plugin/src/nemo_platform_plugin/client/constants.py

@@ -6,7 +6,17 @@
 import os
 
 WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR = "NMP_WORKLOAD_IDENTITY_TOKEN_FILE"
+JWT_WORKLOAD_SUBJECT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:jwt"
+DOCKER_OPAQUE_WORKLOAD_PROOF_TOKEN_TYPE = "urn:nvidia:nemo:params:oauth:token-type:docker-opaque-workload-proof"
+OPAQUE_DOCKER_PROOF_PREFIX = "nmp_obo_v1"
 
 
 def is_workload_identity_token_file_set() -> bool:
     return bool(os.environ.get(WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR))
+
+
+def subject_token_type_for_exchange(subject_token: str) -> str:
+    """Return the RFC 8693 subject_token_type for a workload identity subject token."""
+    if subject_token.startswith(f"{OPAQUE_DOCKER_PROOF_PREFIX}."):
+        return DOCKER_OPAQUE_WORKLOAD_PROOF_TOKEN_TYPE
+    return JWT_WORKLOAD_SUBJECT_TOKEN_TYPE

packages/nemo_platform_plugin/src/nemo_platform_plugin/client/oidc.py

@@ -30,15 +30,15 @@
 from urllib.parse import urlparse
 
 import httpx
-from nemo_platform_plugin.client.constants import WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR
-from nemo_platform_plugin.client.tls import client_verify_from_env
-from nmp.common.auth import (
+from nemo_platform_plugin.client.constants import (
     DOCKER_OPAQUE_WORKLOAD_PROOF_TOKEN_TYPE as _DOCKER_OPAQUE_WORKLOAD_PROOF_TOKEN_TYPE,
 )
-from nmp.common.auth import (
+from nemo_platform_plugin.client.constants import (
     JWT_WORKLOAD_SUBJECT_TOKEN_TYPE,
+    WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR,
     subject_token_type_for_exchange,
 )
+from nemo_platform_plugin.client.tls import client_verify_from_env
 
 logger = logging.getLogger(__name__)
 

packages/nemo_platform_plugin/tests/test_client_auth.py

@@ -5,11 +5,14 @@
 
 from __future__ import annotations
 
+import ast
 import asyncio
 import time
+from pathlib import Path
 from unittest.mock import patch
 
 import httpx
+import nemo_platform_plugin.client.oidc as oidc_module
 import pytest
 import respx
 import yaml
@@ -45,6 +48,24 @@
     get_nemo_client,
 )
 
+
+def _assert_no_nmp_common_import(module_file: str | None) -> None:
+    assert module_file is not None
+    tree = ast.parse(Path(module_file).read_text(encoding="utf-8"))
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Import):
+            imported = [alias.name for alias in node.names]
+        elif isinstance(node, ast.ImportFrom):
+            imported = [node.module or ""]
+        else:
+            continue
+        assert not any(name == "nmp.common" or name.startswith("nmp.common.") for name in imported)
+
+
+def test_oidc_module_has_no_nmp_common_dependency():
+    _assert_no_nmp_common_import(oidc_module.__file__)
+
+
 # ---------------------------------------------------------------------------
 # StaticToken
 # ---------------------------------------------------------------------------

sdk/python/nemo-platform/src/nemo_platform/auth/workload_exchange.py

@@ -9,7 +9,7 @@
 import jwt
 from fastapi import APIRouter, Depends, HTTPException, Request, Response, status
 from nmp.common.auth.bearer import MalformedBearerTokenError, parse_bearer_authorization_header
-from nmp.common.auth.jwt import TokenClaims
+from nmp.common.auth.jwt import ActorClaims, TokenClaims
 from nmp.common.auth.token_resolver import ResolvedBearerToken, ResolvedTokenKind, resolve_bearer_token
 from nmp.common.config import AuthConfig, get_auth_config
 from nmp.core.auth.api.v2.workload_token_exchange import (
@@ -80,6 +80,25 @@ def _scopes_from_claims(claims: dict[str, object]) -> list[str]:
     return []
 
 
+def _actor_from_claims(claims: dict[str, object]) -> ActorClaims | None:
+    actor_claims = claims.get("act")
+    if not isinstance(actor_claims, dict):
+        return None
+
+    actor_subject = actor_claims.get("sub")
+    if not isinstance(actor_subject, str):
+        return None
+
+    actor_subject = actor_subject.strip()
+    if not actor_subject:
+        return None
+
+    return ActorClaims(
+        subject=actor_subject,
+        groups=_groups_from_claim(actor_claims.get("groups", [])),
+    )
+
+
 def _stamp_principal_headers(response: Response, resolved: ResolvedBearerToken) -> None:
     for header_name, header_value in resolved.principal_headers().items():
         response.headers[header_name] = header_value
@@ -138,6 +157,7 @@ async def _validate_workload_access_token(
             groups=_groups_from_claim(claims.get("groups", [])),
             scopes=_scopes_from_claims(claims),
             raw_claims=claims,
+            actor=_actor_from_claims(claims),
         )
     except jwt.PyJWTError:
         return None

sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/auth/test_workload_exchange.py

@@ -197,6 +197,59 @@ def test_authenticate_workload_access_token_returns_principal_headers(tmp_path):
     assert response.headers["X-NMP-Scopes"] == "openid email groups"
 
 
+def test_authenticate_delegated_workload_access_token_returns_obo_principal_headers(tmp_path):
+    private_key_file = tmp_path / "private.pem"
+    private_key_file.write_bytes(_private_key_pem())
+    config = AuthConfig(
+        enabled=True,
+        token_signing=TokenSigningConfig(
+            issuer="http://testserver/apis/auth",
+            key_id="test-workload",
+            private_key_file=str(private_key_file),
+        ),
+        oidc=OIDCConfig(
+            workload_token_exchange_enabled=True,
+            workload_audience="nemo-platform",
+        ),
+    )
+    signing_key = WorkloadTokenExchangeService().workload_signing_key(config)
+    now = datetime.now(tz=UTC)
+    token = jwt.encode(
+        {
+            "iss": "http://testserver/apis/auth",
+            "sub": "submitter@example.com",
+            "email": "submitter@example.com",
+            "aud": "nemo-platform",
+            "iat": now,
+            "nbf": now,
+            "exp": now + timedelta(minutes=5),
+            "scope": "openid email groups",
+            "groups": "workspace-editors",
+            "act": {
+                "sub": "system:serviceaccount:nemo:job",
+                "groups": ["system:serviceaccounts", "nemo-jobs"],
+            },
+        },
+        signing_key.private_key,
+        algorithm="RS256",
+        headers={"kid": signing_key.kid},
+    )
+    with _test_client(config) as client:
+        response = client.get(
+            "/authenticate",
+            headers={"Authorization": f"Bearer {token}"},
+        )
+
+    assert response.status_code == 200
+    assert response.json()["token_kind"] == "workload_access_token"
+    assert response.headers["X-NMP-Principal-Id"] == "system:serviceaccount:nemo:job"
+    assert response.headers["X-NMP-Principal-Groups"] == "system:serviceaccounts,nemo-jobs"
+    assert response.headers["X-NMP-Principal-On-Behalf-Of"] == "submitter@example.com"
+    assert response.headers["X-NMP-Principal-On-Behalf-Of-Email"] == "submitter@example.com"
+    assert response.headers["X-NMP-Principal-On-Behalf-Of-Groups"] == "workspace-editors"
+    assert response.headers["X-NMP-Scopes"] == "openid email groups"
+
+
 def test_authenticate_workload_subject_token_uses_resolver_callback(tmp_path):
     config = AuthConfig(
         enabled=True,

services/core/auth/src/nmp/core/auth/api/v2/authenticate.py

@@ -561,6 +561,51 @@ def _revoke_workload_delegation_after_failed_start(self, delegation_name: str |
                 extra={"delegation_name": delegation_name},
             )
 
+    def _workload_delegation_name_for_step(self, step: PlatformJobStepWithContext) -> str:
+        return docker_delegation_name(
+            workload_workspace=step.workspace,
+            job_id=step.job,
+            attempt_id=step.attempt_id,
+            step_id=step.id,
+        )
+
+    @staticmethod
+    def _workload_delegation_name_from_container(container: Container) -> str | None:
+        labels = getattr(container, "labels", None) or {}
+        if DOCKER_WORKLOAD_IDENTITY_TOKEN_FILE_LABEL not in labels:
+            return None
+
+        workload_workspace = labels.get(JOB_WORKSPACE_ID_LABEL)
+        job_id = labels.get(JOB_ID_LABEL)
+        attempt_id = labels.get(JOB_ATTEMPT_ID_LABEL)
+        step_id = labels.get(JOB_STEP_ID_LABEL)
+        if not isinstance(workload_workspace, str) or not workload_workspace:
+            return None
+        if not isinstance(job_id, str) or not job_id:
+            return None
+        if not isinstance(attempt_id, str) or not attempt_id:
+            return None
+        if not isinstance(step_id, str) or not step_id:
+            return None
+
+        return docker_delegation_name(
+            workload_workspace=workload_workspace,
+            job_id=job_id,
+            attempt_id=attempt_id,
+            step_id=step_id,
+        )
+
+    def _revoke_workload_delegation_after_terminal(self, delegation_name: str | None) -> None:
+        if not delegation_name:
+            return
+        try:
+            self._revoke_workload_delegation(delegation_name)
+        except Exception:
+            logger.exception(
+                "Failed to revoke Docker workload delegation after terminal status",
+                extra={"delegation_name": delegation_name},
+            )
+
     def _write_workload_identity_subject_token(self, volume_name: str, token: str) -> None:
         storage_config = self._execution_profile_config.storage
         permissions_image = (
@@ -1917,6 +1962,8 @@ def create_step_update(self, step: PlatformJobStepWithContext, container: Contai
         error_details = {}
         if status == PlatformJobStatus.ERROR:
             error_details["message"] = status_details.get("message", "Job encountered an error")
+        if status in PlatformJobStatus.terminals() and self._should_enable_workload_identity_for_step(step):
+            self._revoke_workload_delegation_after_terminal(self._workload_delegation_name_for_step(step))
 
         logger.debug(
             "Docker container status mapped to platform status",
@@ -2166,8 +2213,10 @@ def cleanup_single_container(self, container: Container) -> None:
         job = self.get_label_from_container(container, JOB_ID_LABEL)
         task = self.get_label_from_container(container, JOB_TASK_ID_LABEL)
         exit_code = container.attrs.get("State", {}).get("ExitCode", 0)
+        delegation_name = self._workload_delegation_name_from_container(container)
 
         self._stop_workload_identity_refresher(container.name)
+        self._revoke_workload_delegation_after_terminal(delegation_name)
         self.cleanup_container(container)
         logger.debug(
             "Cleaned up container",