From 3817ba1d6d0b24783296b684dedf74aed6cb0ec0 Mon Sep 17 00:00:00 2001 From: Justin Date: Fri, 16 Oct 2026 20:24:42 -0400 Subject: [PATCH] fix(apps): give the account api key only to workers that call resources Deploy and dev copied the client's account API key into every endpoint's template env, so any endpoint's workers (and anyone who can read its template) held full account credentials, including endpoints that never call another resource. runpod-flash injected the key only for endpoints whose functions make remote calls. Analyze each resource's module source: the resource body (not its decorators) and the module-level helpers it reaches. The key is added when that code references another resource, calls .remote or .spawn, uses anything imported from runpod, or uses a project-local import the analysis cannot follow. Unreadable source and nested definitions also get the key, so a worker that needs it never goes without. --- runpod/apps/deploy.py | 13 ++- runpod/apps/dev.py | 12 +- runpod/apps/remote_calls.py | 160 ++++++++++++++++++++++++++ tests/test_apps/test_remote_calls.py | 163 +++++++++++++++++++++++++++ 4 files changed, 343 insertions(+), 5 deletions(-) create mode 100644 runpod/apps/remote_calls.py create mode 100644 tests/test_apps/test_remote_calls.py diff --git a/runpod/apps/deploy.py b/runpod/apps/deploy.py index a0ef47ca..b968f972 100644 --- a/runpod/apps/deploy.py +++ b/runpod/apps/deploy.py @@ -308,10 +308,17 @@ def _deployed_endpoint_input( } # cross-resource calls from inside a worker go through the sentinel, - # which authenticates with the account api key; without it any - # .remote() from worker code would fail + # which authenticates with the account api key. only resources that + # may make such a call get it (the analysis fails open) + from .remote_calls import makes_remote_calls + api_key = runpod.api_key - if api_key and "RUNPOD_API_KEY" not in template_env: + handle = app.resources.get(spec.name) + if ( + api_key + and "RUNPOD_API_KEY" not in template_env + and (handle is None or makes_remote_calls(app, handle)) + ): template_env["RUNPOD_API_KEY"] = api_key # nested tasks use the same runtime and sdk channels as the endpoint diff --git a/runpod/apps/dev.py b/runpod/apps/dev.py index c4e95ad1..43a49e9a 100644 --- a/runpod/apps/dev.py +++ b/runpod/apps/dev.py @@ -107,6 +107,13 @@ def _cpu_locations(instance_ids: List[str]) -> str: GENERATION_ENV = "RUNPOD_DEV_GENERATION" +def _needs_api_key(app: App, spec: ResourceSpec) -> bool: + from .remote_calls import makes_remote_calls + + handle = app.resources.get(spec.name) + return handle is None or makes_remote_calls(app, handle) + + def _endpoint_input(app: App, spec: ResourceSpec, generation: int = 1) -> Dict: """build a live endpoint payload with an atomically bound container config.""" @@ -133,10 +140,11 @@ def _endpoint_input(app: App, spec: ResourceSpec, generation: int = 1) -> Dict: {"key": GENERATION_ENV, "value": str(generation)}, # nested .remote() calls from inside a dev worker need # credentials and the dev-session marker to resolve - # sibling dev endpoints by name + # sibling dev endpoints by name. only resources that may + # make one get the key (the analysis fails open) *( [{"key": "RUNPOD_API_KEY", "value": key}] - if (key := _client_api_key()) + if _needs_api_key(app, spec) and (key := _client_api_key()) else [] ), {"key": "RUNPOD_DEV_APP", "value": app.name}, diff --git a/runpod/apps/remote_calls.py b/runpod/apps/remote_calls.py new file mode 100644 index 00000000..c8143414 --- /dev/null +++ b/runpod/apps/remote_calls.py @@ -0,0 +1,160 @@ +"""decide which resources call other resources from inside a worker. + +workers authenticate nested .remote() calls with the account api key, so +deploy and dev give the key to every resource that may make one. the +analysis reads the resource's module source and fails open: anything it +cannot follow counts as a remote call, so no worker that needs the key +goes without it. +""" + +import ast +import inspect +import logging +from pathlib import Path +from typing import Dict, Optional, Set + +log = logging.getLogger(__name__) + +# handle methods that start a call on another resource +_REMOTE_ATTRS = frozenset({"remote", "spawn"}) + + +def resource_target(handle): + """the user function or class behind a resource handle.""" + return ( + getattr(handle, "_fn", None) + or getattr(handle, "_cls", None) + or getattr(handle, "__wrapped__", None) + ) + + +class _BodyNames(ast.NodeVisitor): + """names and attributes used in a definition's body, not its decorators.""" + + def __init__(self): + self.names: Set[str] = set() + self.attrs: Set[str] = set() + + def visit_FunctionDef(self, node): + self._visit_def(node) + + visit_AsyncFunctionDef = visit_FunctionDef + + def visit_ClassDef(self, node): + for base in node.bases: + self.visit(base) + for stmt in node.body: + self.visit(stmt) + + def visit_Lambda(self, node): + self.visit(node.body) + + def _visit_def(self, node): + for default in node.args.defaults + node.args.kw_defaults: + if default is not None: + self.visit(default) + for stmt in node.body: + self.visit(stmt) + + def visit_Name(self, node): + self.names.add(node.id) + + def visit_Attribute(self, node): + self.attrs.add(node.attr) + self.generic_visit(node) + + +def _body_names(node) -> _BodyNames: + visitor = _BodyNames() + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + visitor.visit(node) + else: + visitor.generic_visit(node) + return visitor + + +def _module_tree(target) -> Optional[ast.Module]: + try: + path = inspect.getsourcefile(target) + return ast.parse(Path(path).read_text(encoding="utf-8")) if path else None + except (OSError, TypeError, SyntaxError, ValueError): + return None + + +def makes_remote_calls(app, handle) -> bool: + """true unless the resource provably cannot reach another resource.""" + target = resource_target(handle) + name = getattr(target, "__name__", None) + if target is None or getattr(target, "__qualname__", name) != name: + # nested or anonymous definitions are not analyzed + return True + tree = _module_tree(target) + if tree is None: + log.debug("source for %s unavailable; giving it the api key", name) + return True + + defs: Dict[str, ast.AST] = {} + assigns: Dict[str, ast.AST] = {} + runpod_names: Set[str] = set() + local_imports: Set[str] = set() + module_file = inspect.getsourcefile(target) or "" + project_dir = Path(module_file).resolve().parent if module_file else None + + def _is_local(module_name: str) -> bool: + if not module_name or project_dir is None: + return True + top = module_name.split(".")[0] + return (project_dir / f"{top}.py").exists() or (project_dir / top).is_dir() + + for node in tree.body: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + defs[node.name] = node + elif isinstance(node, (ast.Assign, ast.AnnAssign)): + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + for t in targets: + for n in ast.walk(t): + if isinstance(n, ast.Name): + assigns[n.id] = node + elif isinstance(node, ast.Import): + for alias in node.names: + bound = alias.asname or alias.name.split(".")[0] + if alias.name.split(".")[0] == "runpod": + runpod_names.add(bound) + elif _is_local(alias.name): + local_imports.add(bound) + elif isinstance(node, ast.ImportFrom): + module_name = node.module or "" + for alias in node.names: + bound = alias.asname or alias.name + if module_name.split(".")[0] == "runpod": + runpod_names.add(bound) + elif node.level or _is_local(module_name): + local_imports.add(bound) + + if name not in defs: + return True + + other_resources = { + resource_target(h).__name__ + for h in app.resources.values() + if h is not handle and resource_target(h) is not None + } + + seen: Set[str] = set() + pending = [name] + while pending: + current = pending.pop() + if current in seen: + continue + seen.add(current) + used = _body_names(defs.get(current) or assigns[current]) + if used.attrs & _REMOTE_ATTRS: + return True + for ref in used.names: + if ref in other_resources or ref in runpod_names or ref in local_imports: + return True + if ref in seen: + continue + if ref in defs or ref in assigns: + pending.append(ref) + return False diff --git a/tests/test_apps/test_remote_calls.py b/tests/test_apps/test_remote_calls.py new file mode 100644 index 00000000..1d2b1745 --- /dev/null +++ b/tests/test_apps/test_remote_calls.py @@ -0,0 +1,163 @@ +"""tests for the worker api key scoping analysis.""" + +import importlib.util +import sys +import textwrap + +import pytest + +from runpod.apps.app import _clear_registry +from runpod.apps.remote_calls import makes_remote_calls + + +@pytest.fixture(autouse=True) +def clean_registry(): + _clear_registry() + yield + _clear_registry() + + +@pytest.fixture +def load(tmp_path, monkeypatch): + monkeypatch.syspath_prepend(str(tmp_path)) + + def _load(source, name="remote_calls_app"): + path = tmp_path / f"{name}.py" + path.write_text(textwrap.dedent(source)) + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, name, module) + spec.loader.exec_module(module) + return module + + return _load + + +SOURCE = ''' +import json + +import runpod + +app = runpod.App("calls") + + +def fmt(value): + return json.dumps({"value": value}) + + +def fan_out(items): + return [leaf.remote(x=i) for i in items] + + +@app.queue(name="leaf", cpu="cpu3c-1-2") +def leaf(x: int) -> str: + return fmt(x) + + +@app.queue(name="caller", cpu="cpu3c-1-2") +def caller(n: int) -> list: + return [leaf(x=i) for i in range(n)] + + +@app.queue(name="indirect", cpu="cpu3c-1-2") +def indirect(n: int) -> list: + return fan_out(range(n)) + + +@app.queue(name="client", cpu="cpu3c-1-2") +def client(n: int): + return runpod.Queue(app="other", name="x").remote(n=n) + + +@app.api(name="service", cpu="cpu3c-1-2") +class Service: + @runpod.init + def setup(self): + self.ready = True + + @runpod.post("/echo") + def echo(self, body: dict) -> dict: + return body + + +@runpod.local_entrypoint +def main(): + print(caller.remote(n=2)) +''' + + +class TestMakesRemoteCalls: + def test_classification(self, load): + mod = load(SOURCE) + app = mod.app + result = { + name: makes_remote_calls(app, handle) + for name, handle in app.resources.items() + } + assert result == { + "leaf": False, # helper only, no other resource + "caller": True, # references another resource + "indirect": True, # through a module-level helper + "client": True, # runpod client inside the body + "service": False, # @runpod.* decorators are not calls + } + + def test_local_project_import_counts(self, tmp_path, load): + (tmp_path / "helpers_mod.py").write_text("def run():\n return 1\n") + mod = load( + """ + import runpod + from helpers_mod import run + + app = runpod.App("local") + + @app.queue(name="uses", cpu="cpu3c-1-2") + def uses() -> int: + return run() + """, + name="remote_calls_local", + ) + assert makes_remote_calls(mod.app, mod.app.resources["uses"]) is True + + def test_unanalyzable_targets_fail_open(self): + import runpod + + app = runpod.App("nested") + + @app.queue(name="inner", cpu="cpu3c-1-2") + def inner() -> int: + return 1 + + # defined inside a function: qualname differs, not analyzed + assert makes_remote_calls(app, app.resources["inner"]) is True + + +class TestPayloadsScopeTheKey: + def test_deploy_gives_key_only_to_callers(self, load, monkeypatch): + import runpod + from runpod.apps.deploy import _deployed_endpoint_input + + monkeypatch.setattr(runpod, "api_key", "acct-key") + mod = load(SOURCE, name="remote_calls_deploy") + + def env_keys(name): + payload = _deployed_endpoint_input( + mod.app, mod.app.resources[name].spec, "env-1", "b1", "3.12" + ) + return {e["key"] for e in payload["template"]["env"]} + + assert "RUNPOD_API_KEY" not in env_keys("leaf") + assert "RUNPOD_API_KEY" in env_keys("caller") + + def test_dev_gives_key_only_to_callers(self, load, monkeypatch): + from runpod.apps import dev + + monkeypatch.setattr(dev, "_client_api_key", lambda: "acct-key") + mod = load(SOURCE, name="remote_calls_dev") + + def env_keys(name): + payload = dev._endpoint_input(mod.app, mod.app.resources[name].spec) + return {e["key"] for e in payload["template"]["env"]} + + assert "RUNPOD_API_KEY" not in env_keys("leaf") + assert "RUNPOD_API_KEY" in env_keys("indirect")