mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-15 04:19:08 +00:00
feat(proxy): accept inline PEM for the billing-metrics mTLS credentials
LITELLM_BILLING_METRICS_CLIENT_CERT, _CLIENT_KEY and _CA_CERT took a filesystem path. ECS injects Secrets Manager values as environment content and cannot mount them as files, so a licensed deployment there could not turn metering on. Each variable now takes either a path or the PEM itself. Inline PEM, detected by the "-----BEGIN" prefix, is written once when the recorder is built into a 0700 temp dir as a 0600 file, and the config points at that path. The OTLP exporter still only ever sees paths. A write failure disables metering through the existing failure-as-None path rather than raising, and path-valued variables are passed through untouched, so nothing changes for deployments that mount files. The mixed case works too: mount the CA, inject the client credentials
This commit is contained in:
parent
4f7f706a63
commit
1b91dd4c0b
2 changed files with 159 additions and 4 deletions
|
|
@ -14,6 +14,7 @@ payload; the secret license key is never sent as an attribute or header.
|
|||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Optional, Union
|
||||
|
||||
|
|
@ -40,6 +41,15 @@ DEFAULT_EXPORT_INTERVAL_MS = 60_000
|
|||
SHUTDOWN_FLUSH_TIMEOUT_MS = 5_000
|
||||
_METRICS_PATH = "/v1/metrics"
|
||||
|
||||
# The cert env vars take a path or the PEM itself. Secret stores that inject
|
||||
# values as env content cannot mount them as files, so inline PEM is written out.
|
||||
_PEM_PREFIX = "-----BEGIN"
|
||||
_PEM_DIR_PREFIX = "litellm-billing-mtls-"
|
||||
_PEM_FILE_MODE = 0o600
|
||||
_CLIENT_CERT_FILENAME = "client.crt"
|
||||
_CLIENT_KEY_FILENAME = "client.key"
|
||||
_CA_CERT_FILENAME = "ca.crt"
|
||||
|
||||
METRIC_NAME = "litellm.enterprise.billable_requests"
|
||||
METER_NAME = "litellm.enterprise.billing"
|
||||
|
||||
|
|
@ -136,6 +146,55 @@ def _export_interval_ms() -> int:
|
|||
return DEFAULT_EXPORT_INTERVAL_MS
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _CredentialPaths:
|
||||
client_cert_path: str
|
||||
client_key_path: str
|
||||
ca_cert_path: Optional[str]
|
||||
|
||||
|
||||
def _is_pem_content(value: str) -> bool:
|
||||
return value.lstrip().startswith(_PEM_PREFIX)
|
||||
|
||||
|
||||
def _write_pem(directory: str, filename: str, pem: str) -> str:
|
||||
path = os.path.join(directory, filename)
|
||||
with open(path, "w", encoding="utf-8") as handle:
|
||||
handle.write(pem if pem.endswith("\n") else f"{pem}\n")
|
||||
os.chmod(path, _PEM_FILE_MODE)
|
||||
return path
|
||||
|
||||
|
||||
def _resolve_credential_paths(*, client_cert: str, client_key: str, ca_cert: Optional[str]) -> _CredentialPaths:
|
||||
"""
|
||||
Accept either a filesystem path or inline PEM content for each credential.
|
||||
|
||||
Secret stores that inject values as environment content rather than mounted
|
||||
files (ECS tasks reading AWS Secrets Manager, Cloud Run reading Secret
|
||||
Manager) can only deliver the certificate as a string. The OTLP exporter
|
||||
takes paths, so inline PEM is written to a private directory once, when the
|
||||
recorder is built. Raises OSError if that write fails; the caller disables
|
||||
metering rather than propagating.
|
||||
"""
|
||||
inline = tuple(value for value in (client_cert, client_key, ca_cert) if value and _is_pem_content(value))
|
||||
if not inline:
|
||||
return _CredentialPaths(client_cert, client_key, ca_cert)
|
||||
|
||||
# mkdtemp is 0o700, so the 0o600 key file it holds is unreachable by other users.
|
||||
directory = tempfile.mkdtemp(prefix=_PEM_DIR_PREFIX)
|
||||
return _CredentialPaths(
|
||||
client_cert_path=(
|
||||
_write_pem(directory, _CLIENT_CERT_FILENAME, client_cert) if _is_pem_content(client_cert) else client_cert
|
||||
),
|
||||
client_key_path=(
|
||||
_write_pem(directory, _CLIENT_KEY_FILENAME, client_key) if _is_pem_content(client_key) else client_key
|
||||
),
|
||||
ca_cert_path=(
|
||||
_write_pem(directory, _CA_CERT_FILENAME, ca_cert) if ca_cert and _is_pem_content(ca_cert) else ca_cert
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def load_billing_metrics_config(
|
||||
*, license_data: Optional["EnterpriseLicenseData"], litellm_version: str
|
||||
) -> Optional[BillingMetricsConfig]:
|
||||
|
|
@ -162,7 +221,17 @@ def load_billing_metrics_config(
|
|||
)
|
||||
return None
|
||||
|
||||
required_paths = [client_cert, client_key] + ([ca_cert] if ca_cert else [])
|
||||
try:
|
||||
paths = _resolve_credential_paths(client_cert=client_cert, client_key=client_key, ca_cert=ca_cert)
|
||||
except OSError as exc:
|
||||
verbose_proxy_logger.warning(
|
||||
"Enterprise billing metrics disabled: could not write inline certificate content to disk: %s", exc
|
||||
)
|
||||
return None
|
||||
|
||||
required_paths = [paths.client_cert_path, paths.client_key_path] + (
|
||||
[paths.ca_cert_path] if paths.ca_cert_path else []
|
||||
)
|
||||
unreadable = [path for path in required_paths if not os.path.isfile(path)]
|
||||
if unreadable:
|
||||
verbose_proxy_logger.warning(
|
||||
|
|
@ -173,9 +242,9 @@ def load_billing_metrics_config(
|
|||
|
||||
return BillingMetricsConfig(
|
||||
endpoint=endpoint,
|
||||
client_cert_path=client_cert,
|
||||
client_key_path=client_key,
|
||||
ca_cert_path=ca_cert,
|
||||
client_cert_path=paths.client_cert_path,
|
||||
client_key_path=paths.client_key_path,
|
||||
ca_cert_path=paths.ca_cert_path,
|
||||
export_interval_ms=_export_interval_ms(),
|
||||
litellm_version=litellm_version,
|
||||
license_id=(license_data or {}).get("user_id"),
|
||||
|
|
|
|||
|
|
@ -8,7 +8,9 @@ collectors), the metric attribute mapping, and that recording produces the
|
|||
expected OTLP counter via an in-memory reader.
|
||||
"""
|
||||
|
||||
import os
|
||||
import socket
|
||||
import stat
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
|
|
@ -163,6 +165,90 @@ def test_load_config_with_empty_string_env_is_disabled(monkeypatch, tmp_path):
|
|||
assert bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") is None
|
||||
|
||||
|
||||
_CLIENT_CERT_PEM = "-----BEGIN CERTIFICATE-----\nclient-cert-body\n-----END CERTIFICATE-----"
|
||||
_CLIENT_KEY_PEM = "-----BEGIN PRIVATE KEY-----\nclient-key-body\n-----END PRIVATE KEY-----"
|
||||
_CA_CERT_PEM = "-----BEGIN CERTIFICATE-----\nca-body\n-----END CERTIFICATE-----"
|
||||
|
||||
|
||||
def test_load_config_materializes_inline_pem_content(monkeypatch):
|
||||
"""
|
||||
ECS and Cloud Run inject secrets as env content, not as mounted files, so the
|
||||
cert env vars must accept PEM directly. The exporter takes paths, so the PEM
|
||||
is written to disk and the config points at those files.
|
||||
"""
|
||||
monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4317")
|
||||
monkeypatch.setenv(bm.CLIENT_CERT_ENV, _CLIENT_CERT_PEM)
|
||||
monkeypatch.setenv(bm.CLIENT_KEY_ENV, _CLIENT_KEY_PEM)
|
||||
monkeypatch.setenv(bm.CA_CERT_ENV, _CA_CERT_PEM)
|
||||
|
||||
config = bm.load_billing_metrics_config(license_data=None, litellm_version="1.0")
|
||||
|
||||
assert config is not None
|
||||
assert config.ca_cert_path is not None
|
||||
written = {
|
||||
config.client_cert_path: _CLIENT_CERT_PEM,
|
||||
config.client_key_path: _CLIENT_KEY_PEM,
|
||||
config.ca_cert_path: _CA_CERT_PEM,
|
||||
}
|
||||
for path, pem in written.items():
|
||||
assert path != pem, "config must carry a file path, not the PEM itself"
|
||||
assert os.path.isfile(path)
|
||||
assert Path(path).read_text(encoding="utf-8") == f"{pem}\n"
|
||||
|
||||
# The private key must not be world- or group-readable.
|
||||
assert stat.S_IMODE(os.stat(config.client_key_path).st_mode) == 0o600
|
||||
|
||||
|
||||
def test_load_config_accepts_a_mix_of_pem_content_and_file_paths(monkeypatch, tmp_path):
|
||||
"""A deployment may mount the CA but inject the client credentials inline."""
|
||||
paths = _write_certs(tmp_path)
|
||||
monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4317")
|
||||
monkeypatch.setenv(bm.CLIENT_CERT_ENV, _CLIENT_CERT_PEM)
|
||||
monkeypatch.setenv(bm.CLIENT_KEY_ENV, _CLIENT_KEY_PEM)
|
||||
monkeypatch.setenv(bm.CA_CERT_ENV, paths[bm.CA_CERT_ENV])
|
||||
|
||||
config = bm.load_billing_metrics_config(license_data=None, litellm_version="1.0")
|
||||
|
||||
assert config is not None
|
||||
assert config.ca_cert_path == paths[bm.CA_CERT_ENV]
|
||||
assert Path(config.client_cert_path).read_text(encoding="utf-8") == f"{_CLIENT_CERT_PEM}\n"
|
||||
|
||||
|
||||
def test_load_config_leaves_file_paths_untouched(monkeypatch, tmp_path):
|
||||
"""Path-valued env vars keep working; nothing is copied or rewritten."""
|
||||
paths = _set_full_env(monkeypatch, tmp_path)
|
||||
|
||||
config = bm.load_billing_metrics_config(license_data=None, litellm_version="1.0")
|
||||
|
||||
assert config is not None
|
||||
assert config.client_cert_path == paths[bm.CLIENT_CERT_ENV]
|
||||
assert config.client_key_path == paths[bm.CLIENT_KEY_ENV]
|
||||
assert config.ca_cert_path == paths[bm.CA_CERT_ENV]
|
||||
|
||||
|
||||
def test_load_config_with_inline_pem_disabled_when_unwritable(monkeypatch):
|
||||
"""A failure to materialize the PEM disables metering instead of raising."""
|
||||
monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4317")
|
||||
monkeypatch.setenv(bm.CLIENT_CERT_ENV, _CLIENT_CERT_PEM)
|
||||
monkeypatch.setenv(bm.CLIENT_KEY_ENV, _CLIENT_KEY_PEM)
|
||||
|
||||
def _explode(prefix=None):
|
||||
raise OSError("read-only filesystem")
|
||||
|
||||
monkeypatch.setattr(bm.tempfile, "mkdtemp", _explode)
|
||||
|
||||
assert bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") is None
|
||||
|
||||
|
||||
def test_load_config_with_empty_pem_env_is_disabled(monkeypatch):
|
||||
"""Empty stays empty: an unset secret must not be mistaken for inline PEM."""
|
||||
monkeypatch.setenv(bm.ENDPOINT_ENV, "https://collector.example:4317")
|
||||
monkeypatch.setenv(bm.CLIENT_CERT_ENV, "")
|
||||
monkeypatch.setenv(bm.CLIENT_KEY_ENV, "")
|
||||
|
||||
assert bm.load_billing_metrics_config(license_data=None, litellm_version="1.0") is None
|
||||
|
||||
|
||||
def test_export_interval_default_and_override(monkeypatch):
|
||||
assert bm._export_interval_ms() == bm.DEFAULT_EXPORT_INTERVAL_MS
|
||||
monkeypatch.setenv(bm.EXPORT_INTERVAL_ENV, "5000")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue