codexLB-up 1.25.0b10__py3-none-any.whl

app/core/upstream_proxy/cache.py

@@ -0,0 +1,130 @@
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+
+from app.core.cache.invalidation import NAMESPACE_UPSTREAM_ROUTE, get_cache_invalidation_poller
+from app.core.config.settings import get_settings
+from app.core.upstream_proxy.resolver import UpstreamProxyRouteError
+from app.core.upstream_proxy.types import ResolvedUpstreamRoute
+
+
+@dataclass(frozen=True, slots=True)
+class CachedRouteOutcome:
+    """Terminal resolver outcome for one account, reproduced verbatim on a hit.
+
+    Exactly one of the resolver's three outcomes is held: a resolved route, a
+    permitted direct-egress ``None`` (``route is None`` with no error), or a
+    fail-closed error reason. ``unwrap`` re-raises errors with the original
+    reason and pool so a cache hit can never change the degradation path the
+    resolver chose.
+    """
+
+    route: ResolvedUpstreamRoute | None
+    error_reason: str | None
+    error_pool_id: str | None
+    expires_at: float
+
+    def unwrap(self, account_id: str | None) -> ResolvedUpstreamRoute | None:
+        if self.error_reason is not None:
+            raise UpstreamProxyRouteError(self.error_reason, account_id=account_id, pool_id=self.error_pool_id)
+        return self.route
+
+
+class UpstreamRouteCache:
+    """Per-account cache of upstream proxy route resolution outcomes.
+
+    Admin mutations of resolver inputs (account bindings, pool members, the
+    upstream-proxy dashboard settings) invalidate durably through the
+    ``upstream_route`` / ``settings`` cache-invalidation namespaces; the TTL is
+    only a backstop for out-of-band database edits. A TTL of 0 disables the
+    cache entirely (the test-suite default).
+    """
+
+    def __init__(self) -> None:
+        self._entries: dict[str, CachedRouteOutcome] = {}
+        self._generation = 0
+
+    @property
+    def generation(self) -> int:
+        return self._generation
+
+    @staticmethod
+    def _ttl_seconds() -> float:
+        return get_settings().upstream_route_cache_ttl_seconds
+
+    def get(self, account_id: str) -> CachedRouteOutcome | None:
+        if self._ttl_seconds() <= 0:
+            return None
+        entry = self._entries.get(account_id)
+        if entry is None:
+            return None
+        if time.monotonic() >= entry.expires_at:
+            self._entries.pop(account_id, None)
+            return None
+        return entry
+
+    def store_route(self, account_id: str, route: ResolvedUpstreamRoute | None, *, generation: int) -> None:
+        self._store(account_id, route=route, error_reason=None, error_pool_id=None, generation=generation)
+
+    def store_error(self, account_id: str, error: UpstreamProxyRouteError, *, generation: int) -> None:
+        self._store(
+            account_id,
+            route=None,
+            error_reason=error.reason,
+            error_pool_id=error.pool_id,
+            generation=generation,
+        )
+
+    def _store(
+        self,
+        account_id: str,
+        *,
+        route: ResolvedUpstreamRoute | None,
+        error_reason: str | None,
+        error_pool_id: str | None,
+        generation: int,
+    ) -> None:
+        ttl = self._ttl_seconds()
+        if ttl <= 0:
+            return
+        if generation != self._generation:
+            # An invalidation landed while this outcome was being resolved; the
+            # outcome may predate the mutation, so drop it.
+            return
+        self._entries[account_id] = CachedRouteOutcome(
+            route=route,
+            error_reason=error_reason,
+            error_pool_id=error_pool_id,
+            expires_at=time.monotonic() + ttl,
+        )
+
+    def clear(self) -> None:
+        self._generation += 1
+        self._entries.clear()
+
+    async def invalidate(self, *, propagate: bool = True) -> None:
+        """Clear locally and, unless ``propagate`` is False, durably bump the
+        cross-replica ``upstream_route`` namespace before returning.
+
+        Route mutations are security-bearing (a stale entry could keep an
+        account on direct egress after a binding lands), so mutation paths
+        await the bump rather than coalescing it. Poller callbacks use
+        ``clear`` directly, so a remote bump never re-bumps.
+        """
+        self.clear()
+        if propagate:
+            poller = get_cache_invalidation_poller()
+            if poller is not None and not await poller.bump(NAMESPACE_UPSTREAM_ROUTE):
+                # bump() never raises; on failure the coalesced pending-set
+                # retries on every poll cycle until the write lands, so peers
+                # still converge once the database recovers. The TTL backstop
+                # bounds staleness in the interim.
+                poller.request_bump(NAMESPACE_UPSTREAM_ROUTE)
+
+
+_upstream_route_cache = UpstreamRouteCache()
+
+
+def get_upstream_route_cache() -> UpstreamRouteCache:
+    return _upstream_route_cache

app/core/upstream_proxy/resolver.py

@@ -0,0 +1,189 @@
+from __future__ import annotations
+
+import logging
+
+from cryptography.fernet import InvalidToken
+from sqlalchemy import select
+from sqlalchemy.exc import OperationalError
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import selectinload
+
+from app.core.crypto import TokenEncryptor
+from app.core.upstream_proxy.types import ResolvedProxyEndpoint, ResolvedUpstreamRoute, sends_plaintext_credentials
+from app.db.models import AccountProxyBinding, DashboardSettings, ProxyEndpoint, ProxyPool, ProxyPoolMember
+
+_ACCOUNT_BOUND_MODE = "account_bound"
+_DEFAULT_POOL_MODE = "default_pool"
+_SUPPORTED_SCHEMES = frozenset({"http", "https", "socks5", "socks5h"})
+
+logger = logging.getLogger(__name__)
+
+# Endpoint ids already warned about in this process; the resolver runs on every
+# routed request, so the plaintext-credential warning is emitted once per
+# endpoint rather than per request.
+_PLAINTEXT_CREDENTIAL_WARNINGS_EMITTED: set[str] = set()
+
+
+class UpstreamProxyRouteError(RuntimeError):
+    def __init__(self, reason: str, *, account_id: str | None = None, pool_id: str | None = None) -> None:
+        super().__init__(reason)
+        self.reason = reason
+        self.account_id = account_id
+        self.pool_id = pool_id
+
+
+async def resolve_upstream_route(
+    session: AsyncSession,
+    *,
+    account_id: str | None,
+    operation: str,
+    scope: str = "account",
+    strict: bool | None = None,
+    encryptor: TokenEncryptor | None = None,
+) -> ResolvedUpstreamRoute | None:
+    del operation
+    try:
+        binding = await _lookup_account_binding(session, account_id) if account_id else None
+    except OperationalError as exc:
+        if strict is True or not _is_missing_upstream_proxy_schema(exc):
+            raise
+        return None
+    if binding is not None:
+        return await _resolve_pool(
+            session,
+            pool_id=binding.pool_id,
+            mode=_ACCOUNT_BOUND_MODE,
+            account_id=account_id,
+            encryptor=encryptor,
+        )
+
+    try:
+        settings = await _lookup_dashboard_settings(session)
+    except OperationalError as exc:
+        if strict is True or not _is_missing_upstream_proxy_schema(exc):
+            raise
+        return None
+    strict_enabled = settings.upstream_proxy_routing_enabled if settings is not None else False
+    if strict is not None:
+        strict_enabled = strict
+    if not strict_enabled:
+        return None
+    pool_id = settings.upstream_proxy_default_pool_id if settings is not None else None
+    if not pool_id:
+        raise UpstreamProxyRouteError("default_pool_unconfigured", account_id=account_id)
+    if scope not in {"account", "service", "bootstrap"}:
+        raise UpstreamProxyRouteError("unsupported_route_scope", account_id=account_id, pool_id=pool_id)
+    return await _resolve_pool(
+        session,
+        pool_id=pool_id,
+        mode=_DEFAULT_POOL_MODE,
+        account_id=account_id,
+        encryptor=encryptor,
+    )
+
+
+def _is_missing_upstream_proxy_schema(exc: OperationalError) -> bool:
+    message = str(exc.orig).lower() if getattr(exc, "orig", None) is not None else str(exc).lower()
+    # Upstream routing is intentionally optional during rollout and migration.
+    # PostgreSQL reports missing tables/columns with a different wording than
+    # SQLite, so we treat both families as tolerated "schema-not-ready" states.
+    return (
+        "no such table: account_proxy_bindings" in message
+        or "no such column: dashboard_settings.upstream_proxy" in message
+        or "column dashboard_settings.upstream_proxy" in message
+        and "does not exist" in message
+        or 'relation "account_proxy_bindings" does not exist' in message
+        or 'relation "dashboard_settings" does not exist' in message
+        or ('column "dashboard_settings"' in message and "does not exist" in message)
+        or ('column "account_proxy_bindings"' in message and "does not exist" in message)
+    )
+
+
+async def _lookup_account_binding(session: AsyncSession, account_id: str | None) -> AccountProxyBinding | None:
+    if not account_id:
+        return None
+    result = await session.execute(
+        select(AccountProxyBinding)
+        .where(AccountProxyBinding.account_id == account_id, AccountProxyBinding.is_active.is_(True))
+        .limit(1)
+    )
+    return result.scalar_one_or_none()
+
+
+async def _lookup_dashboard_settings(session: AsyncSession) -> DashboardSettings | None:
+    result = await session.execute(select(DashboardSettings).where(DashboardSettings.id == 1).limit(1))
+    return result.scalar_one_or_none()
+
+
+async def _resolve_pool(
+    session: AsyncSession,
+    *,
+    pool_id: str,
+    mode: str,
+    account_id: str | None,
+    encryptor: TokenEncryptor | None,
+) -> ResolvedUpstreamRoute:
+    pool = await session.get(ProxyPool, pool_id)
+    if pool is None or not pool.is_active:
+        raise UpstreamProxyRouteError("pool_unavailable", account_id=account_id, pool_id=pool_id)
+    result = await session.execute(
+        select(ProxyPoolMember)
+        .options(selectinload(ProxyPoolMember.endpoint))
+        .where(ProxyPoolMember.pool_id == pool_id, ProxyPoolMember.is_active.is_(True))
+        .order_by(ProxyPoolMember.sort_order.asc(), ProxyPoolMember.id.asc())
+    )
+    endpoints = [
+        _resolve_endpoint(member.endpoint, encryptor=encryptor)
+        for member in result.scalars().all()
+        if member.endpoint is not None and member.endpoint.is_active
+    ]
+    if not endpoints:
+        raise UpstreamProxyRouteError("pool_has_no_active_endpoints", account_id=account_id, pool_id=pool_id)
+    return ResolvedUpstreamRoute(mode=mode, pool_id=pool_id, endpoint=endpoints[0], fallbacks=tuple(endpoints[1:]))
+
+
+def _warn_plaintext_credentials_once(endpoint: ProxyEndpoint, scheme: str) -> None:
+    if endpoint.id in _PLAINTEXT_CREDENTIAL_WARNINGS_EMITTED:
+        return
+    _PLAINTEXT_CREDENTIAL_WARNINGS_EMITTED.add(endpoint.id)
+    logger.warning(
+        "Upstream proxy endpoint %s (%s://%s:%s) sends its credentials to the proxy in plaintext; "
+        "prefer an https:// proxy or an IP allowlist without credentials",
+        endpoint.id,
+        scheme,
+        endpoint.host,
+        endpoint.port,
+    )
+
+
+def _resolve_endpoint(endpoint: ProxyEndpoint, *, encryptor: TokenEncryptor | None) -> ResolvedProxyEndpoint:
+    scheme = endpoint.scheme.lower().strip()
+    if scheme not in _SUPPORTED_SCHEMES:
+        raise UpstreamProxyRouteError("unsupported_proxy_scheme")
+    if sends_plaintext_credentials(
+        scheme, has_credentials=endpoint.username is not None or endpoint.password_encrypted is not None
+    ):
+        _warn_plaintext_credentials_once(endpoint, scheme)
+    if endpoint.username is not None and ":" in endpoint.username:
+        # RFC 7617 Basic credentials cannot encode a colon in the user-id.
+        raise UpstreamProxyRouteError("invalid_proxy_username")
+    password: str | None = None
+    if endpoint.password_encrypted is not None:
+        try:
+            password = (encryptor or TokenEncryptor()).decrypt(endpoint.password_encrypted)
+        except InvalidToken as exc:
+            raise UpstreamProxyRouteError("proxy_credentials_undecryptable") from exc
+    return ResolvedProxyEndpoint(
+        id=endpoint.id,
+        scheme=scheme,
+        host=endpoint.host,
+        port=endpoint.port,
+        username=endpoint.username,
+        password=password,
+    )
+
+
+def resolve_proxy_endpoint(
+    endpoint: ProxyEndpoint, *, encryptor: TokenEncryptor | None = None
+) -> ResolvedProxyEndpoint:
+    return _resolve_endpoint(endpoint, encryptor=encryptor)

app/core/upstream_proxy/types.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import base64
+from dataclasses import dataclass
+from typing import Any
+from urllib.parse import quote
+
+_PLAINTEXT_SCHEMES = frozenset({"http", "socks5", "socks5h"})
+
+
+def sends_plaintext_credentials(scheme: str, *, has_credentials: bool) -> bool:
+    """Whether proxy credentials cross the LB-to-proxy hop unencrypted.
+
+    ``http`` (CONNECT ``Proxy-Authorization``) and ``socks5``/``socks5h``
+    (username/password sub-negotiation) both send the credential in the clear
+    to the proxy server; only an ``https`` proxy encrypts that hop. Such
+    endpoints are allowed, but the dashboard surfaces a warning and the
+    resolver logs one so operators can move to ``https`` or an IP allowlist.
+    """
+
+    return has_credentials and scheme.lower().strip() in _PLAINTEXT_SCHEMES
+
+
+def _encode_basic_proxy_auth(username: str, password: str) -> str:
+    # RFC 7617 Basic token encoded with latin1, byte-identical to the header
+    # value aiohttp derives from URL userinfo (``BasicAuth`` default encoding);
+    # the CONNECT request differs only in header order. Local so the declared
+    # ``aiohttp>=3.13`` floor holds (``aiohttp.encode_basic_auth`` is 3.14+).
+    if ":" in username:
+        raise ValueError("proxy usernames cannot contain ':'")
+    token = base64.b64encode(f"{username}:{password}".encode("latin1")).decode("ascii")
+    return f"Basic {token}"
+
+
+@dataclass(frozen=True, slots=True)
+class ResolvedProxyEndpoint:
+    id: str
+    scheme: str
+    host: str
+    port: int
+    username: str | None = None
+    password: str | None = None
+
+    @property
+    def plaintext_credentials(self) -> bool:
+        return sends_plaintext_credentials(
+            self.scheme, has_credentials=self.username is not None or self.password is not None
+        )
+
+    @property
+    def proxy_url(self) -> str:
+        scheme = "socks5h" if self.scheme == "socks5" else self.scheme
+        auth = ""
+        if self.username:
+            auth = f"{quote(self.username, safe='')}:{quote(self.password or '', safe='')}@"
+        return f"{scheme}://{auth}{self.host}:{self.port}"
+
+    @property
+    def proxy_url_without_credentials(self) -> str:
+        scheme = "socks5h" if self.scheme == "socks5" else self.scheme
+        return f"{scheme}://{self.host}:{self.port}"
+
+    def aiohttp_proxy_kwargs(self) -> dict[str, Any]:
+        """Return ``proxy``/``proxy_headers`` kwargs for aiohttp request and ws_connect.
+
+        Credentials travel in a ``Proxy-Authorization`` header instead of URL
+        userinfo so aiohttp's ``ConnectionKey``/``Connection`` reprs and
+        ``ClientHttpProxyError.__str__`` never carry the password. ``latin1``
+        matches aiohttp's own userinfo encoding (``BasicAuth`` default), so the
+        CONNECT ``Proxy-Authorization`` value stays byte-identical (aiohttp
+        emits it right after ``Host`` instead of last, which no proxy
+        observes). aiohttp forwards ``proxy_headers``
+        only on the CONNECT tunnel request, i.e. for TLS (``https``/``wss``)
+        targets; callers must not use these kwargs for plaintext targets.
+        """
+        kwargs: dict[str, Any] = {"proxy": self.proxy_url_without_credentials}
+        if self.username:
+            kwargs["proxy_headers"] = {
+                "Proxy-Authorization": _encode_basic_proxy_auth(self.username, self.password or ""),
+            }
+        return kwargs
+
+
+@dataclass(frozen=True, slots=True)
+class ResolvedUpstreamRoute:
+    mode: str
+    pool_id: str
+    endpoint: ResolvedProxyEndpoint
+    fallbacks: tuple[ResolvedProxyEndpoint, ...] = ()
+
+    @property
+    def endpoint_id(self) -> str:
+        return self.endpoint.id
+
+    @property
+    def proxy_url(self) -> str:
+        return self.endpoint.proxy_url
+
+    def with_endpoint(
+        self,
+        endpoint: ResolvedProxyEndpoint,
+        fallbacks: tuple[ResolvedProxyEndpoint, ...],
+    ) -> "ResolvedUpstreamRoute":
+        return ResolvedUpstreamRoute(self.mode, self.pool_id, endpoint, fallbacks)