ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: round 7 Codex hardening (cross-wrapper env scrub + device URL allowlist)

Browse source 
Two P1 fixes from the round 7 reviewer pass:

1. _ScrubbedEnvAsyncCodex no longer leaks secrets across overlapping
   sessions.

   The fallback env-scrub wrapper (used when the installed SDK build
   does not accept AppServerConfig(env=...)) refcounts deleted env
   vars under a shared lock so concurrent fan-out workers do not
   restore Studio's secrets while a peer is still inside SDK startup.
   The round 6 implementation enumerated keys to scrub via
   _codex_sdk_env_override() which only returns keys currently in
   os.environ. If wrapper A had already deleted HF_TOKEN before
   wrapper B entered, B's overrides dict no longer contained HF_TOKEN,
   B never bumped the refcount for it, and A's exit restored
   HF_TOKEN into the process env while B was still running -- so
   B's spawned codex app-server inherited the secret.

   Round 7 fix: under the lock, the union of (a) the current
   overrides dict and (b) every key still refcounted by an earlier
   wrapper is the set of keys this session must scrub. Originals are
   tracked module-level rather than per-instance so the last wrapper
   to release a key always restores the right pre-scrub value
   regardless of who first saw it. New regression test reproduces
   the leak against the round 6 code (asserts refcount == 2 after
   B enters; old code records 1) and locks the fix in.

2. Device-auth verification URL is now host-allowlisted.

   The codex login --device-auth output parser pulled any https URL
   matching /device, /activate, or /verify out of the CLI's stdout
   and emitted it as a device_url event. The frontend rendered that
   URL as an "Open verification page" CTA the user can click. A
   compromised codex shim earlier on PATH could print
   https://evil.example/activate?code=ABCD and Studio would surface
   the phishing link verbatim, even though every other login output
   line goes through a strict safe-vocabulary filter.

   Round 7 fix: device_url events only fire for URLs whose host is
   on a small allowlist (auth.openai.com / chatgpt.com over https).
   Anything else is logged at warn and dropped. Tests cover the
   known-good upstream URLs, several attacker patterns (lookalike
   subdomains, http downgrade, javascript:), and garbage input.
This commit is contained in:
Daniel Han 2026-05-25 14:43:10 +00:00
commit f05517ac79
2 changed files with 235 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

126
studio/backend/core/inference/codex_provider.py
View file

@ -44,12 +44,52 @@ import shutil
import sys
import time
from typing import Any, AsyncGenerator, Optional
from urllib.parse import urlparse
import structlog
logger = structlog.get_logger(__name__)
# Device-auth verification URLs the upstream codex CLI prints during
# `codex login --device-auth`. Only these hosts are forwarded to the
# browser as "Open verification page". A shimmed codex earlier on PATH
# can still print a phishing URL that matches the regex used to fish
# the verification URL out of stdout (`/device`, `/activate`,
# `/verify`), but Studio refuses to surface anything not on this list,
# so the malicious URL never reaches the user.
_ALLOWED_DEVICE_AUTH_HOSTS: frozenset[str] = frozenset({
"auth.openai.com",
"chatgpt.com",
})
def _safe_host(url: str) -> Optional[str]:
"""Return the lower-case host of ``url`` or None if it cannot be parsed."""
try:
return (urlparse(url).hostname or "").lower() or None
except Exception:
return None
def _is_allowed_device_url(url: str) -> bool:
"""Return True iff ``url`` looks like a real codex device-auth URL.
Requires https, a parseable URL, and a host on the allowlist above.
Codex login URLs are always https in the wild; downgrade to http
is a strong signal of a malicious shim, so we drop both at the
same gate.
"""
try:
parsed = urlparse(url)
except Exception:
return False
if parsed.scheme != "https":
return False
host = (parsed.hostname or "").lower()
return host in _ALLOWED_DEVICE_AUTH_HOSTS
# Hard cap on parallel Codex fan-out. Picked to match the upper bound
# in the request validator -- exceeding this risks the local Codex CLI
# rate-limiting itself or starving the loop.
@ -96,6 +136,14 @@ def _codex_sdk_env_override() -> dict[str, str]:
_SCRUBBED_ENV_LOCK = asyncio.Lock()
_SCRUBBED_ENV_REFCOUNT: dict[str, int] = {}
# Saved originals shared across all wrappers. The first wrapper to scrub
# a key records its pre-scrub value here; later wrappers that pick up the
# same key while it is already absent from os.environ inherit the same
# saved value so the very last wrapper to release a key still restores
# the right thing. Kept module-level (not per-instance) because two
# concurrent wrappers must agree on what the original was even though
# only one of them actually saw it in os.environ.
_SCRUBBED_ENV_ORIGINALS: dict[str, str] = {}
class _ScrubbedEnvAsyncCodex:
@ -111,7 +159,7 @@ class _ScrubbedEnvAsyncCodex:
enter/exit critical section, and a per-key refcount tracks how
many concurrent wrappers are currently "holding" the scrub. A
key is only restored when the last wrapper using it exits. This
fixes two issues round 6 caught:
fixes three issues:
1. Two concurrent fan-out workers used to race: wrapper A could
restore `HF_TOKEN` while wrapper B was still inside SDK
@ -122,6 +170,15 @@ class _ScrubbedEnvAsyncCodex:
Python never called `__aexit__`, so the deleted keys leaked
permanently. Construction now happens INSIDE the try/except
in `__aenter__`, and the scrub is rolled back on failure.
3. ``_codex_sdk_env_override()`` only enumerates keys currently
in ``os.environ``, so a wrapper B entering AFTER wrapper A
already scrubbed (say) ``HF_TOKEN`` would never see that key
in its overrides dict, never bump its refcount, and miss the
scrub for HF_TOKEN entirely. When A then exited it would
restore HF_TOKEN while B was still mid-session. Wrapper B
now also picks up every key currently refcounted by an
earlier wrapper (read under the lock) so the refcount
reflects the true set of holders for every scrubbed key.
"""
def __init__(self, async_codex_cls: Any):
@ -131,26 +188,32 @@ class _ScrubbedEnvAsyncCodex:
# __aexit__ knows exactly which counters to decrement (avoids
# racing with concurrent wrappers that scrub a different set).
self._held_keys: list[str] = []
# Snapshot of the original values at the time of the FIRST
# wrapper that scrubbed each key, so restoration uses the
# real pre-scrub value.
self._restored_via_us: dict[str, str] = {}
async def __aenter__(self) -> Any:
import os
overrides = _codex_sdk_env_override()
async with _SCRUBBED_ENV_LOCK:
for key in overrides:
if key not in os.environ and _SCRUBBED_ENV_REFCOUNT.get(key, 0) == 0:
continue
if _SCRUBBED_ENV_REFCOUNT.get(key, 0) == 0:
# Keys this session would scrub if it were entering first.
keys_to_scrub = set(_codex_sdk_env_override())
# Plus every key still held by an earlier wrapper -- without
# this we would miss keys that are already absent from
# os.environ but ARE still scrubbed and refcounted.
keys_to_scrub.update(
key for key, count in _SCRUBBED_ENV_REFCOUNT.items() if count > 0
)
for key in keys_to_scrub:
current = _SCRUBBED_ENV_REFCOUNT.get(key, 0)
if current == 0:
# First wrapper to scrub this key -- save the
# original so the very last wrapper to release it
# can restore the right value.
self._restored_via_us[key] = os.environ[key]
# original. If the key is somehow not in os.environ
# right now (deleted between override-snapshot and
# here, or simply never set), skip it: nothing to
# scrub and nothing to restore.
if key not in os.environ:
continue
_SCRUBBED_ENV_ORIGINALS[key] = os.environ[key]
del os.environ[key]
_SCRUBBED_ENV_REFCOUNT[key] = _SCRUBBED_ENV_REFCOUNT.get(key, 0) + 1
_SCRUBBED_ENV_REFCOUNT[key] = current + 1
self._held_keys.append(key)
try:
self._inner = self._async_codex_cls()
@ -176,15 +239,18 @@ class _ScrubbedEnvAsyncCodex:
current = _SCRUBBED_ENV_REFCOUNT.get(key, 0)
if current <= 0:
continue
_SCRUBBED_ENV_REFCOUNT[key] = current - 1
if current - 1 == 0:
next_count = current - 1
if next_count == 0:
# Last wrapper holding this key -- restore the
# original value if WE were the first to scrub it,
# or pull from any other wrapper's saved snapshot.
if key in self._restored_via_us:
os.environ.setdefault(key, self._restored_via_us[key])
# original snapshot recorded when it was first
# scrubbed, regardless of which wrapper saved it.
_SCRUBBED_ENV_REFCOUNT.pop(key, None)
original = _SCRUBBED_ENV_ORIGINALS.pop(key, None)
if original is not None:
os.environ.setdefault(key, original)
else:
_SCRUBBED_ENV_REFCOUNT[key] = next_count
self._held_keys.clear()
self._restored_via_us.clear()
def _resolve_codex_bin() -> Optional[str]:
@ -1389,8 +1455,22 @@ async def stream_codex_device_login() -> AsyncGenerator[dict[str, Any], None]:
if not url_emitted:
match = url_re.search(line)
if match:
yield {"type": "device_url", "url": match.group(0)}
url_emitted = True
candidate = match.group(0)
if _is_allowed_device_url(candidate):
yield {"type": "device_url", "url": candidate}
url_emitted = True
else:
# A shimmed codex on PATH could print a phishing
# URL that matches the regex but points at an
# attacker host (e.g. https://evil.example/activate
# ?code=ABCD). Drop it rather than surfacing
# "Open verification page" to a real user.
logger.warning(
"codex_provider.login_url_rejected",
host = (
_safe_host(candidate) or "<unparsable>"
),
)
if not code_emitted:
cm = code_re.search(line)
if cm:

132
studio/backend/tests/test_codex_provider.py
View file

@ -1895,3 +1895,135 @@ async def _consume_first(gen):
"""
async for _ in gen:
return
# ── Round 7: _ScrubbedEnvAsyncCodex cross-wrapper concurrency ──────
class TestScrubbedEnvConcurrency:
"""Reproduce the cross-wrapper concurrency hole the round 7 review
surfaced and lock in the fix: when wrapper B enters AFTER wrapper A
has already deleted ``HF_TOKEN`` from ``os.environ``, B must still
increment the refcount for that key so A's exit does not restore
the secret while B is mid-session.
"""
def test_overlapping_wrappers_keep_keys_scrubbed_until_last_release(
self, monkeypatch
):
import os
from core.inference.codex_provider import (
_SCRUBBED_ENV_REFCOUNT,
_ScrubbedEnvAsyncCodex,
)
# Reset module-level state in case prior tests left residue.
_SCRUBBED_ENV_REFCOUNT.clear()
# _SCRUBBED_ENV_ORIGINALS is the round 7 fix's shared snapshot
# store; older codex_provider builds tracked originals per-
# instance under _restored_via_us. Reset whichever store the
# current build exposes so prior tests cannot leak state in.
from core.inference import codex_provider as _cp
_orig = getattr(_cp, "_SCRUBBED_ENV_ORIGINALS", None)
if isinstance(_orig, dict):
_orig.clear()
monkeypatch.setenv("HF_TOKEN", "sekret-hf")
monkeypatch.setenv("GH_TOKEN", "sekret-gh")
# Keys NOT on the safe-list end up in _codex_sdk_env_override().
class _FakeInner:
async def __aenter__(self_inner):
return self_inner
async def __aexit__(self_inner, *a):
return False
def _fake_async_codex():
return _FakeInner()
async def scenario():
wrapper_a = _ScrubbedEnvAsyncCodex(_fake_async_codex)
wrapper_b = _ScrubbedEnvAsyncCodex(_fake_async_codex)
# Wrapper A enters first and scrubs both secrets.
await wrapper_a.__aenter__()
assert "HF_TOKEN" not in os.environ
assert "GH_TOKEN" not in os.environ
# Wrapper B enters while A is still active. Even though
# os.environ no longer contains HF_TOKEN/GH_TOKEN (A already
# deleted them), B must pick them up from the live refcount
# table so A's later exit does not restore them prematurely.
await wrapper_b.__aenter__()
assert _SCRUBBED_ENV_REFCOUNT.get("HF_TOKEN") == 2
assert _SCRUBBED_ENV_REFCOUNT.get("GH_TOKEN") == 2
# A exits first -- B is still active so the keys MUST remain
# absent from os.environ.
await wrapper_a.__aexit__(None, None, None)
assert "HF_TOKEN" not in os.environ, (
"HF_TOKEN leaked back into os.environ while wrapper B "
"is still active"
)
assert "GH_TOKEN" not in os.environ
assert _SCRUBBED_ENV_REFCOUNT.get("HF_TOKEN") == 1
assert _SCRUBBED_ENV_REFCOUNT.get("GH_TOKEN") == 1
# B exits -- now the keys must be restored from the saved
# originals.
await wrapper_b.__aexit__(None, None, None)
assert os.environ.get("HF_TOKEN") == "sekret-hf"
assert os.environ.get("GH_TOKEN") == "sekret-gh"
assert "HF_TOKEN" not in _SCRUBBED_ENV_REFCOUNT
asyncio.run(scenario())
# ── Round 7: device-auth URL allowlisting ───────────────────────────
class TestDeviceUrlAllowlist:
"""Lock in the device-auth URL allowlist: only `auth.openai.com`
and `chatgpt.com` over https are accepted as `device_url` events.
A shimmed codex earlier on PATH could otherwise print
`https://evil.example/activate?code=ABCD` and Studio would render
a phishing CTA.
"""
def test_known_good_urls_allowed(self):
from core.inference.codex_provider import _is_allowed_device_url
assert _is_allowed_device_url(
"https://auth.openai.com/codex/device?user_code=ABCD-EFGH"
)
assert _is_allowed_device_url(
"https://chatgpt.com/activate?user_code=WXYZ-1234"
)
def test_attacker_hosts_rejected(self):
from core.inference.codex_provider import _is_allowed_device_url
for evil in [
"https://evil.example/activate?code=ABCD",
"https://auth-openai-com.evil.example/codex/device",
"https://chatgpt.com.evil.example/activate",
"https://login.openai.com/codex/device",
]:
assert not _is_allowed_device_url(evil), evil
def test_http_downgrade_rejected(self):
from core.inference.codex_provider import _is_allowed_device_url
assert not _is_allowed_device_url(
"http://auth.openai.com/codex/device?user_code=ABCD-EFGH"
)
def test_garbage_url_rejected(self):
from core.inference.codex_provider import _is_allowed_device_url
assert not _is_allowed_device_url("not a url")
assert not _is_allowed_device_url("")
assert not _is_allowed_device_url("javascript:alert(1)")