Merge pull request #1005 from Yelp/u/mpiano/SEC-19555

KaspariK · web-flow · commit 864cd60ec29b · 2025-03-13T14:15:03.000-04:00
auth support for Tron APIs
diff --git a/tests/api/auth_test.py b/tests/api/auth_test.py
@@ -0,0 +1,76 @@
+from unittest.mock import MagicMock
+from unittest.mock import patch
+
+import pytest
+from twisted.web.server import Request
+
+from tron.api.auth import AuthorizationFilter
+from tron.api.auth import AuthorizationOutcome
+
+
+@pytest.fixture
+def mock_auth_filter():
+    with patch("tron.api.auth.requests"):
+        yield AuthorizationFilter("http://localhost:31337/whatever", True)
+
+
+def mock_request(path: str, token: str, method: str):
+    res = MagicMock(spec=Request, path=path.encode(), method=method.encode())
+    res.getHeader.return_value = token
+    return res
+
+
+def test_is_request_authorized(mock_auth_filter):
+    mock_auth_filter.session.post.return_value.json.return_value = {
+        "result": {"allowed": True, "reason": "User allowed"}
+    }
+    assert mock_auth_filter.is_request_authorized(
+        mock_request("/api/jobs/foobar.run.2", "aaa.bbb.ccc", "get")
+    ) == AuthorizationOutcome(True, "User allowed")
+    mock_auth_filter.session.post.assert_called_once_with(
+        url="http://localhost:31337/whatever",
+        json={
+            "input": {
+                "path": "/api/jobs/foobar.run.2",
+                "backend": "tron",
+                "token": "aaa.bbb.ccc",
+                "method": "get",
+                "service": "foobar",
+            }
+        },
+        timeout=2,
+    )
+
+
+def test_is_request_authorized_fail(mock_auth_filter):
+    mock_auth_filter.session.post.side_effect = Exception
+    assert mock_auth_filter.is_request_authorized(
+        mock_request("/allowed", "eee.ddd.fff", "get")
+    ) == AuthorizationOutcome(False, "Auth backend error")
+
+
+def test_is_request_authorized_malformed(mock_auth_filter):
+    mock_auth_filter.session.post.return_value.json.return_value = {"foo": "bar"}
+    assert mock_auth_filter.is_request_authorized(
+        mock_request("/allowed", "eee.ddd.fff", "post")
+    ) == AuthorizationOutcome(False, "Malformed auth response")
+
+
+def test_is_request_authorized_no_enforce(mock_auth_filter):
+    mock_auth_filter.session.post.return_value.json.return_value = {
+        "result": {"allowed": False, "reason": "Missing token"}
+    }
+    with patch.object(mock_auth_filter, "enforce", False):
+        assert mock_auth_filter.is_request_authorized(mock_request("/foobar", "", "post")) == AuthorizationOutcome(
+            True, "Auth dry-run"
+        )
+
+
+def test_is_request_authorized_disabled(mock_auth_filter):
+    mock_auth_filter.session.post.return_value.json.return_value = {
+        "result": {"allowed": False, "reason": "Missing token"}
+    }
+    with patch.object(mock_auth_filter, "endpoint", None):
+        assert mock_auth_filter.is_request_authorized(mock_request("/buzz", "", "post")) == AuthorizationOutcome(
+            True, "Auth not enabled"
+        )
diff --git a/tron/api/auth.py b/tron/api/auth.py
@@ -0,0 +1,106 @@
+import logging
+import os
+from functools import lru_cache
+from typing import NamedTuple
+from typing import Optional
+
+import cachetools.func
+import requests
+from twisted.web.server import Request
+
+
+logger = logging.getLogger(__name__)
+AUTH_CACHE_SIZE = 50000
+AUTH_CACHE_TTL = 30 * 60
+
+
+class AuthorizationOutcome(NamedTuple):
+    authorized: bool
+    reason: str
+
+
+class AuthorizationFilter:
+    """API request authorization via external system"""
+
+    def __init__(self, endpoint: str, enforce: bool):
+        """Constructor
+
+        :param str endpoint: HTTP endpoint of external authorization system
+        :param bool enforce: whether to enforce authorization decisions
+        """
+        self.endpoint = endpoint
+        self.enforce = enforce
+        self.session = requests.Session()
+
+    @classmethod
+    @lru_cache(maxsize=1)
+    def get_from_env(cls) -> "AuthorizationFilter":
+        return cls(
+            endpoint=os.getenv("API_AUTH_ENDPOINT", ""),
+            enforce=bool(os.getenv("API_AUTH_ENFORCE", "")),
+        )
+
+    def is_request_authorized(self, request: Request) -> AuthorizationOutcome:
+        """Check if API request is authorized
+
+        :param Request request: API request object
+        :return: auth outcome
+        """
+        if not self.endpoint:
+            return AuthorizationOutcome(True, "Auth not enabled")
+        token = (request.getHeader("Authorization") or "").strip()
+        token = token.split()[-1] if token else ""  # removes "Bearer" prefix
+        url_path = request.path.decode()
+        service = url_path.split("/")[-1].split(".", 1)[0] if "/jobs/" in url_path else None
+        auth_outcome = self._is_request_authorized_impl(
+            # path and method are byte arrays in twisted
+            path=url_path,
+            token=token,
+            method=request.method.decode(),
+            service=service,
+        )
+        return auth_outcome if self.enforce else AuthorizationOutcome(True, "Auth dry-run")
+
+    @cachetools.func.ttl_cache(maxsize=AUTH_CACHE_SIZE, ttl=AUTH_CACHE_TTL)
+    def _is_request_authorized_impl(
+        self,
+        path: str,
+        token: str,
+        method: str,
+        service: Optional[str],
+    ) -> AuthorizationOutcome:
+        """Check if API request is authorized
+
+        :param str path: API path
+        :param str token: authentication token
+        :param str method: http method
+        :return: auth outcome
+        """
+        try:
+            response = self.session.post(
+                url=self.endpoint,
+                json={
+                    "input": {
+                        "path": path,
+                        "backend": "tron",
+                        "token": token,
+                        "method": method.lower(),
+                        "service": service,
+                    },
+                },
+                timeout=2,
+            ).json()
+        except Exception as e:
+            logger.exception(f"Issue communicating with auth endpoint: {e}")
+            return AuthorizationOutcome(False, "Auth backend error")
+
+        auth_result_allowed = response.get("result", {}).get("allowed")
+        if auth_result_allowed is None:
+            return AuthorizationOutcome(False, "Malformed auth response")
+
+        if not auth_result_allowed:
+            reason = response["result"].get("reason", "Denied")
+            return AuthorizationOutcome(False, reason)
+
+        reason = response["result"].get("reason", "Ok")
+        return AuthorizationOutcome(True, reason)
diff --git a/tron/api/resource.py b/tron/api/resource.py
@@ -26,6 +26,7 @@
 from tron.api import adapter, controller
 from tron.api import requestargs
 from tron.api.async_resource import AsyncResource
+from tron.api.auth import AuthorizationFilter
 from tron.metrics import view_all_metrics
 from tron.metrics import meter
 from tron.utils import maybe_decode
@@ -514,6 +515,18 @@ def render_GET(self, request):
         }
         return respond(request=request, response=response)
 
+    def render(self, request):
+        """Overriding base `render` method to support auth"""
+        auth_outcome = AuthorizationFilter.get_from_env().is_request_authorized(request)
+        if not auth_outcome.authorized:
+            return respond(
+                request=request,
+                response={"reason": auth_outcome.reason},
+                code=http.FORBIDDEN,
+                headers={"X-Auth-Failure-Reason": auth_outcome.reason},
+            )
+        return super().render(request)
+
 
 class RootResource(resource.Resource):
     def __init__(self, mcp, web_path):
diff --git a/tron/commands/client.py b/tron/commands/client.py
@@ -37,9 +37,27 @@ class RequestError(ValueError):
 }
 
 
+def get_sso_auth_token() -> str:
+    """Generate an authentication token for the calling user from the Single Sign On provider, if configured"""
+
+    # These imports are here because:
+    # - okta-auth is an internal library, so can never be imported or type-checked in public builds
+    # - there's an annoying circular import with the cmd_utils module
+    from okta_auth import get_and_cache_jwt_default  # type: ignore
+    from tron.commands.cmd_utils import get_client_config
+
+    client_id = get_client_config().get("auth_sso_oidc_client_id")
+    return get_and_cache_jwt_default(client_id) if client_id else ""  # type: ignore
+
+
 def build_url_request(uri, data, headers=None, method=None):
     headers = headers or default_headers
     enc_data = urllib.parse.urlencode(data).encode() if data else None
+    # Currently implementing auth only for management actions (i.e. POST requests)
+    if os.getenv("TRONCTL_API_AUTH") and (data or method.upper() == "POST"):
+        token = get_sso_auth_token()
+        if token:
+            headers["Authorization"] = f"Bearer {token}"
     return urllib.request.Request(uri, enc_data, headers=headers, method=method)
 
 
diff --git a/yelp_package/extra_requirements_yelp.txt b/yelp_package/extra_requirements_yelp.txt
@@ -1,4 +1,7 @@
 clusterman-metrics==2.2.1  # used by tron for pre-scaling for Spark runs
 logreader==1.2.0        # used by tron logreader
+okta-auth==1.0.2        # used for API auth
+pyjwt==2.9.0            # required by okta-auth
+saml-helper==2.5.3      # required by okta-auth
 simplejson==3.19.2      # required by tron CLI
 yelp-meteorite==2.1.1   # used by task-processing to emit metrics, clusterman-metrics dependency
