fix: harden private Storage control-plane authentication

src/tower/_storage.py

@@ -5,11 +5,20 @@
 import time
 from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
-from http import HTTPStatus
-from typing import Any, Optional
+from typing import TYPE_CHECKING
+
+import httpx
+
+if TYPE_CHECKING:
+    from pyiceberg.catalog import Catalog
 
-from ._client import _env_client
 from ._context import TowerContext
+from .exceptions import (
+    StorageConnectionError,
+    StorageInvalidCredentialError,
+    StorageMissingAuthenticationError,
+)
+from .tower_api_client import AuthenticatedClient
 from .tower_api_client.api.default import describe_catalog as describe_catalog_api
 from .tower_api_client.api.default import (
     describe_default_catalog as describe_default_catalog_api,
@@ -33,13 +42,137 @@
 # only cache failed catalog type describe requests for this long
 # retry only after this period
 CATALOG_TYPE_FAILURE_CACHE_TTL_SECONDS = 30.0
+DEFAULT_STORAGE_TIMEOUT_SECONDS = 30.0
 DEFAULT_CATALOG_PROVISION_RETRY_DELAYS = (0.25, 0.5, 1.0, 2.0)
 DEFAULT_CATALOG_NAME = "default"
 DEFAULT_ENVIRONMENT_NAME = "default"
 TOWER_CATALOG_TYPE = "tower-catalog"
+INVALID_CREDENTIAL_SENTINEL = "<redacted>"
 logger = logging.getLogger("tower.storage")
 
 
+def _auth_from_context(context: TowerContext) -> tuple[str, str, str]:
+    if context.jwt is not None:
+        token = context.jwt
+        auth_header_name = "Authorization"
+        prefix = "Bearer"
+        source = "TOWER_JWT"
+    elif context.api_key is not None:
+        token = context.api_key
+        auth_header_name = "X-API-Key"
+        prefix = ""
+        source = "TOWER_API_KEY"
+    else:
+        raise StorageMissingAuthenticationError(
+            "No Tower authentication found. Set TOWER_API_KEY or TOWER_JWT."
+        )
+
+    if token.strip() == INVALID_CREDENTIAL_SENTINEL:
+        raise StorageInvalidCredentialError(
+            f"{source} contains the {INVALID_CREDENTIAL_SENTINEL!r} placeholder, "
+            "not a usable Tower credential."
+        )
+    if not token.strip():
+        raise StorageMissingAuthenticationError(
+            "No Tower authentication found. Set TOWER_API_KEY or TOWER_JWT."
+        )
+
+    return token, auth_header_name, prefix
+
+
+def _new_tower_control_plane_client(
+    *,
+    base_url: str,
+    token: str,
+    auth_header_name: str,
+    prefix: str,
+    timeout: float,
+) -> AuthenticatedClient:
+    return AuthenticatedClient(
+        verify_ssl=True,
+        base_url=base_url,
+        token=token,
+        auth_header_name=auth_header_name,
+        prefix=prefix,
+        timeout=httpx.Timeout(timeout),
+        raise_on_unexpected_status=True,
+    )
+
+
+def _auth_hash(token: str, auth_header_name: str, prefix: str) -> str:
+    presented_auth = f"{auth_header_name}\0{prefix}\0{token}"
+    return hashlib.sha256(presented_auth.encode("utf-8")).hexdigest()
+
+
+class _StorageResolver:
+    """Private Tower configuration and authentication for catalog resolution."""
+
+    def __init__(
+        self,
+        *,
+        environment: str | None = None,
+    ) -> None:
+        context = TowerContext.build()
+
+        if environment is not None and not isinstance(environment, str):
+            raise TypeError("environment must be a string or None")
+
+        if environment is not None and not environment.strip():
+            raise ValueError("environment must not be blank")
+
+        self._target_environment = (
+            environment or context.environment or DEFAULT_ENVIRONMENT_NAME
+        )
+        self._base_url = _api_base_url(context.tower_url)
+        self._token, self._auth_header_name, self._auth_prefix = _auth_from_context(
+            context
+        )
+        self._auth_hash = _auth_hash(
+            self._token,
+            self._auth_header_name,
+            self._auth_prefix,
+        )
+
+    def _new_client(self) -> AuthenticatedClient:
+        return _new_tower_control_plane_client(
+            base_url=self._base_url,
+            token=self._token,
+            auth_header_name=self._auth_header_name,
+            prefix=self._auth_prefix,
+            timeout=DEFAULT_STORAGE_TIMEOUT_SECONDS,
+        )
+
+    def _request_catalog_credentials(
+        self,
+        name: str,
+        mode: str,
+    ) -> ErrorModel | VendCatalogCredentialsResponse | None:
+        body = VendCatalogCredentialsBody(mode=_vend_mode(mode))
+
+        try:
+            with self._new_client() as client:
+                return vend_catalog_credentials_api.sync(
+                    name=name,
+                    client=client,
+                    environment=self._target_environment,
+                    body=body,
+                )
+        except httpx.RequestError as error:
+            raise StorageConnectionError(
+                f"Could not connect to Tower at {self._base_url}."
+            ) from error
+
+
+def _api_base_url(tower_url: str) -> str:
+    try:
+        url = httpx.URL(tower_url)
+    except (TypeError, httpx.InvalidURL) as error:
+        raise ValueError(f"Invalid Tower URL: {tower_url}") from error
+    if not url.is_absolute_url or url.scheme not in ("http", "https"):
+        raise ValueError(f"Invalid Tower URL: {tower_url}")
+    return str(url.copy_with(path="/v1", query=None, fragment=None))
+
+
 @dataclass
 class _CachedCredentials:
     credentials: CatalogCredentials
@@ -56,14 +189,14 @@ class _CachedCatalogType:
 
 
 _credential_cache: dict[tuple[str, str, str, str, str], _CachedCredentials] = {}
-_catalog_type_cache: dict[tuple[str, str, str], _CachedCatalogType] = {}
+_catalog_type_cache: dict[tuple[str, str, str, str], _CachedCatalogType] = {}
 
 
 def get_tower_catalog(
     name: str = DEFAULT_CATALOG_NAME,
-    environment: Optional[str] = None,
+    environment: str | None = None,
     mode: str = "read",
-) -> Any:
+) -> Catalog:
     """
     Load a PyIceberg REST catalog using short-lived credentials vended by Tower.
     """
@@ -73,26 +206,25 @@ def get_tower_catalog(
 
 def get_tower_catalog_credentials(
     name: str = DEFAULT_CATALOG_NAME,
-    environment: Optional[str] = None,
+    environment: str | None = None,
     mode: str = "read",
 ) -> CatalogCredentials:
-    ctx = TowerContext.build()
-    environment = environment or ctx.environment or DEFAULT_ENVIRONMENT_NAME
+    storage_resolver = _StorageResolver(environment=environment)
     mode = _normalize_mode(mode)
-    cache_key = _cache_key(ctx, name, environment, mode)
+    cache_key = _cache_key(storage_resolver, name, mode)
 
     now = datetime.now(timezone.utc)
     _prune_credential_cache(now)
     cached = _credential_cache.get(cache_key)
     if cached is not None and cached.is_usable(now):
         return cached.credentials
 
-    credentials = _vend_with_default_catalog_fallback(ctx, name, environment, mode)
+    credentials = _vend_with_default_catalog_fallback(storage_resolver, name, mode)
     _credential_cache[cache_key] = _CachedCredentials(credentials)
     return credentials
 
 
-def load_vended_catalog(name: str, credentials: CatalogCredentials) -> Any:
+def load_vended_catalog(name: str, credentials: CatalogCredentials) -> Catalog:
     from pyiceberg.catalog import load_catalog
 
     return load_catalog(
@@ -105,20 +237,23 @@ def load_vended_catalog(name: str, credentials: CatalogCredentials) -> Any:
 
 
 def _vend_with_default_catalog_fallback(
-    ctx: TowerContext, name: str, environment: str, mode: str
+    storage_resolver: _StorageResolver,
+    name: str,
+    mode: str,
 ) -> CatalogCredentials:
-    result = _vend_catalog_credentials(ctx, name, environment, mode)
+    environment = storage_resolver._target_environment
+    result = storage_resolver._request_catalog_credentials(name, mode)
     if not _is_not_found(result):
         return _unwrap_vend_result(result, name, environment)
 
     if name == DEFAULT_CATALOG_NAME and environment == DEFAULT_ENVIRONMENT_NAME:
-        _ensure_legacy_default_catalog(ctx)
+        _ensure_legacy_default_catalog(storage_resolver)
         for delay in DEFAULT_CATALOG_PROVISION_RETRY_DELAYS:
             time.sleep(delay)
-            result = _vend_catalog_credentials(ctx, name, environment, mode)
+            result = storage_resolver._request_catalog_credentials(name, mode)
             if not _is_not_found(result):
                 return _unwrap_vend_result(result, name, environment)
-            _ensure_legacy_default_catalog(ctx)
+            _ensure_legacy_default_catalog(storage_resolver)
 
         return _unwrap_vend_result(result, name, environment)
 
@@ -127,26 +262,16 @@ def _vend_with_default_catalog_fallback(
     )
 
 
-def _vend_catalog_credentials(
-    ctx: TowerContext, name: str, environment: str, mode: str
-) -> ErrorModel | VendCatalogCredentialsResponse | None:
-    _ensure_tower_auth(ctx)
-    body = VendCatalogCredentialsBody(mode=_vend_mode(mode))
-    return vend_catalog_credentials_api.sync(
-        name=name,
-        client=_env_client(ctx),
-        environment=environment,
-        body=body,
-    )
-
-
 def _describe_tower_catalog_type(
     ctx: TowerContext, name: str, environment: str
 ) -> str | None:
-    if not (ctx.api_key or ctx.jwt):
+    if ctx.jwt is None and ctx.api_key is None:
         return None
 
-    cache_key = (ctx.tower_url, name, environment)
+    token, auth_header_name, prefix = _auth_from_context(ctx)
+    base_url = _api_base_url(ctx.tower_url)
+    auth_hash = _auth_hash(token, auth_header_name, prefix)
+    cache_key = (base_url, auth_hash, name, environment)
     cached = _catalog_type_cache.get(cache_key)
     if cached is not None:
         if cached.retry_at is None:
@@ -157,12 +282,21 @@ def _describe_tower_catalog_type(
 
         _catalog_type_cache.pop(cache_key, None)
 
+    tower_client = _new_tower_control_plane_client(
+        base_url=base_url,
+        token=token,
+        auth_header_name=auth_header_name,
+        prefix=prefix,
+        timeout=CATALOG_TYPE_DESCRIBE_TIMEOUT_SECONDS,
+    )
+
     try:
-        result = describe_catalog_api.sync(
-            name=name,
-            client=_env_client(ctx, timeout=CATALOG_TYPE_DESCRIBE_TIMEOUT_SECONDS),
-            environment=environment,
-        )
+        with tower_client:
+            result = describe_catalog_api.sync(
+                name=name,
+                client=tower_client,
+                environment=environment,
+            )
     except Exception:
         logger.debug(
             "Failed to describe Tower catalog %r in environment %r; "
@@ -199,11 +333,10 @@ def _failed_catalog_type_cache_entry() -> _CachedCatalogType:
     )
 
 
-def _ensure_legacy_default_catalog(ctx: TowerContext) -> None:
+def _ensure_legacy_default_catalog(storage_resolver: _StorageResolver) -> None:
     try:
-        response = describe_default_catalog_api.sync_detailed(client=_env_client(ctx))
-        if response.status_code not in (HTTPStatus.OK, HTTPStatus.ACCEPTED):
-            return
+        with storage_resolver._new_client() as client:
+            describe_default_catalog_api.sync(client=client)
     except Exception:
         # The following vend retry will surface the actionable backend/auth error.
         return
@@ -230,19 +363,18 @@ def _unwrap_vend_result(
     )
 
 
-def _ensure_tower_auth(ctx: TowerContext) -> None:
-    if ctx.api_key or ctx.jwt:
-        return
-
-    raise RuntimeError("No Tower authentication found. Set TOWER_API_KEY or TOWER_JWT.")
-
-
 def _cache_key(
-    ctx: TowerContext, name: str, environment: str, mode: str
+    storage_resolver: _StorageResolver,
+    name: str,
+    mode: str,
 ) -> tuple[str, str, str, str, str]:
-    token = ctx.api_key or ctx.jwt or ""
-    principal_hash = hashlib.sha256(token.encode("utf-8")).hexdigest()
-    return (ctx.tower_url, principal_hash, name, environment, mode)
+    return (
+        storage_resolver._base_url,
+        storage_resolver._auth_hash,
+        name,
+        storage_resolver._target_environment,
+        mode,
+    )
 
 
 def _prune_credential_cache(now: datetime) -> None:

src/tower/exceptions.py

@@ -43,3 +43,23 @@ def __init__(self):
             "replace [a, b] with a & b. You can also pass a PyIceberg "
             "BooleanExpression or a SQL-like filter string."
         )
+
+
+class StorageError(RuntimeError):
+    """Base error for Tower Storage control-plane operations."""
+
+
+class StorageAuthenticationError(StorageError):
+    """Base error for Storage authentication failures."""
+
+
+class StorageMissingAuthenticationError(StorageAuthenticationError):
+    """No supported Tower API key or JWT was available."""
+
+
+class StorageInvalidCredentialError(StorageAuthenticationError):
+    """A configured credential is a known placeholder rather than a secret."""
+
+
+class StorageConnectionError(StorageError):
+    """Tower's control-plane API could not be reached."""