Coverage for gco / services / tls_proxy.py: 100.00%
214 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +0000
1"""Hot-reloading TLS termination proxy for GCO pods.
3The application container listens on pod-loopback HTTP. A second container
4using this module exposes only HTTPS to the pod network and forwards decrypted
5bytes over loopback. It fronts the ALB-facing API pods, the cost monitor, the
6inference monitor's metrics endpoint, the OpenCost, Grafana and MLflow chart pods, and
7managed model pods. Certificate files are treated as a pluggable projection:
8today cert-manager supplies a Secret volume; a future Kubernetes PodCertificate
9volume can replace it without changing the proxy or Service topology.
11Rotation never touches the listening socket. The listener is bound once with a
12dispatcher ``SSLContext`` that carries no keypair; its ``sni_callback`` points
13every handshake at the currently active, fully validated keypair context.
14OpenSSL runs that callback for every ClientHello -- with ``server_name=None``
15when the peer sends no SNI, which is how the ALB connects to its targets -- so
16activating a rotated leaf is a single attribute assignment: no accept gap while
17a replacement listener binds, no rebind that could fail, and established
18streams keep the session they already negotiated.
20The process runs on uvloop when it is importable (the service images ship it
21for uvicorn) and on the stdlib selector loop otherwise.
23Startup fails closed by default: a keypair that is missing or invalid when the
24process starts is an error, because a GCO pod mounts its Secret non-optionally
25and the kubelet does not start the container until the Secret exists. Pods GCO
26does not build (the OpenCost, Grafana and MLflow chart pods) can start before the
27post-Helm Certificates have produced their Secrets, so those sidecars mount the
28Secret ``optional`` and set ``TLS_PROXY_KEYPAIR_WAIT_SECONDS``: the proxy then
29polls for the keypair for that long before binding, and only raises the load
30error once the budget is spent.
32The module is deliberately stdlib-only (uvloop is optional) and imports nothing
33from ``gco``: the inference monitor ships this file's source to managed model
34pods, where it runs as a plain script on a stock Python image.
35"""
37from __future__ import annotations
39import asyncio
40import contextlib
41import hashlib
42import logging
43import math
44import os
45import signal
46import ssl
47import warnings
48from collections.abc import Callable
49from dataclasses import dataclass
50from pathlib import Path
51from typing import Any
53logger = logging.getLogger(__name__)
55TLS_CERT_FILE_ENV = "GCO_TLS_CERT_FILE"
56TLS_KEY_FILE_ENV = "GCO_TLS_KEY_FILE"
57KEYPAIR_WAIT_SECONDS_ENV = "TLS_PROXY_KEYPAIR_WAIT_SECONDS"
58DEFAULT_CERT_FILE = "/var/run/gco/tls/tls.crt"
59DEFAULT_KEY_FILE = "/var/run/gco/tls/tls.key"
60_BUFFER_BYTES = 64 * 1024
61# While waiting for a keypair, say so at most this often: the wait can last
62# half an hour on a fresh cluster and one line per poll would bury the log.
63_KEYPAIR_WAIT_LOG_INTERVAL_SECONDS = 60.0
64# Floor for the keypair wait's poll interval, so TLS_PROXY_POLL_SECONDS=0 (a
65# valid rotation setting) cannot turn the wait into a busy loop.
66_MIN_KEYPAIR_WAIT_POLL_SECONDS = 0.1
69@dataclass(frozen=True)
70class ProxyConfig:
71 """Validated listener, upstream, rotation, and shutdown settings."""
73 host: str
74 port: int
75 upstream_host: str
76 upstream_port: int
77 cert_file: Path
78 key_file: Path
79 poll_seconds: float
80 graceful_shutdown_seconds: float
81 # 0 keeps startup fail-closed: an absent or invalid keypair is an error.
82 keypair_wait_seconds: float = 0.0
85def _positive_port(name: str, default: int) -> int:
86 value = int(os.getenv(name, str(default)))
87 if not 1 <= value <= 65535:
88 raise RuntimeError(f"{name} must be between 1 and 65535")
89 return value
92def _non_negative_number(name: str, default: float) -> float:
93 value = float(os.getenv(name, str(default)))
94 if not math.isfinite(value) or value < 0:
95 raise RuntimeError(f"{name} must be a finite non-negative number")
96 return value
99def load_proxy_config() -> ProxyConfig:
100 """Resolve process configuration, failing closed on an incomplete keypair."""
101 cert_value = os.getenv(TLS_CERT_FILE_ENV, DEFAULT_CERT_FILE).strip()
102 key_value = os.getenv(TLS_KEY_FILE_ENV, DEFAULT_KEY_FILE).strip()
103 if not cert_value or not key_value:
104 raise RuntimeError(f"{TLS_CERT_FILE_ENV} and {TLS_KEY_FILE_ENV} must not be empty")
105 cert_file = Path(cert_value)
106 key_file = Path(key_value)
107 return ProxyConfig(
108 host=os.getenv("TLS_PROXY_HOST", "0.0.0.0"), # pod listener
109 port=_positive_port("TLS_PROXY_PORT", 8443),
110 upstream_host=os.getenv("TLS_PROXY_UPSTREAM_HOST", "127.0.0.1"),
111 upstream_port=_positive_port("TLS_PROXY_UPSTREAM_PORT", 9000),
112 cert_file=cert_file,
113 key_file=key_file,
114 poll_seconds=_non_negative_number("TLS_PROXY_POLL_SECONDS", 5.0),
115 graceful_shutdown_seconds=_non_negative_number("GRACEFUL_SHUTDOWN_TIMEOUT_SECONDS", 30.0),
116 keypair_wait_seconds=_non_negative_number(KEYPAIR_WAIT_SECONDS_ENV, 0.0),
117 )
120def _keypair_digest(config: ProxyConfig) -> str:
121 """Return a digest of readable certificate material without logging it."""
122 digest = hashlib.sha256()
123 for variable, path in (
124 (TLS_CERT_FILE_ENV, config.cert_file),
125 (TLS_KEY_FILE_ENV, config.key_file),
126 ):
127 if not path.is_file() or not os.access(path, os.R_OK):
128 raise RuntimeError(f"{variable} does not reference a readable file: {path}")
129 digest.update(path.read_bytes())
130 return digest.hexdigest()
133def _server_context() -> ssl.SSLContext:
134 """TLS 1.2+ server context with no keypair; handshakes fail closed until one is loaded."""
135 context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
136 context.minimum_version = ssl.TLSVersion.TLSv1_2
137 return context
140def _ssl_context(config: ProxyConfig) -> tuple[ssl.SSLContext, str]:
141 """Build a validated keypair context and return it with its keypair digest."""
142 digest = _keypair_digest(config)
143 context = _server_context()
144 context.load_cert_chain(config.cert_file, config.key_file)
145 return context, digest
148async def _close_writer(writer: asyncio.StreamWriter) -> None:
149 writer.close()
150 with contextlib.suppress(ConnectionError, OSError):
151 await writer.wait_closed()
154async def _pump(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
155 while data := await reader.read(_BUFFER_BYTES):
156 writer.write(data)
157 await writer.drain()
160class TlsProxy:
161 """TLS-only TCP proxy with in-place certificate rotation and graceful stream drain."""
163 def __init__(self, config: ProxyConfig) -> None:
164 self.config = config
165 self._server: asyncio.Server | None = None
166 # Replaced by the projected keypair in start(); until then the
167 # keypair-less placeholder makes any handshake fail closed.
168 self._active_context: ssl.SSLContext = _server_context()
169 self._keypair_digest = ""
170 self._stop = asyncio.Event()
171 self._connections: set[asyncio.Task[Any]] = set()
173 def _select_context(
174 self,
175 ssl_object: ssl.SSLObject,
176 server_name: str | None,
177 dispatcher: ssl.SSLContext,
178 ) -> None:
179 """Point one handshake at the keypair active when its ClientHello arrived.
181 The ALB sends no SNI, so ``server_name`` is ``None`` in production; the
182 active leaf serves every peer regardless of the name presented.
183 """
184 del server_name, dispatcher
185 ssl_object.context = self._active_context
187 async def _handle_connection(
188 self,
189 client_reader: asyncio.StreamReader,
190 client_writer: asyncio.StreamWriter,
191 ) -> None:
192 task = asyncio.current_task()
193 if task is not None:
194 self._connections.add(task)
195 upstream_writer: asyncio.StreamWriter | None = None
196 try:
197 upstream_reader, upstream_writer = await asyncio.open_connection(
198 self.config.upstream_host,
199 self.config.upstream_port,
200 )
201 pumps = {
202 asyncio.create_task(_pump(client_reader, upstream_writer)),
203 asyncio.create_task(_pump(upstream_reader, client_writer)),
204 }
205 _done, pending = await asyncio.wait(
206 pumps,
207 return_when=asyncio.FIRST_COMPLETED,
208 )
209 for pending_task in pending:
210 pending_task.cancel()
211 await asyncio.gather(*pumps, return_exceptions=True)
212 except (ConnectionError, OSError) as exc:
213 logger.debug("TLS proxy connection ended before upstream was ready: %s", exc)
214 finally:
215 if upstream_writer is not None:
216 await _close_writer(upstream_writer)
217 await _close_writer(client_writer)
218 if task is not None:
219 self._connections.discard(task)
221 async def _wait_for_stop(self, timeout: float) -> bool:
222 """Return True when shutdown is requested within ``timeout`` seconds."""
223 try:
224 await asyncio.wait_for(self._stop.wait(), timeout=timeout)
225 except TimeoutError:
226 return False
227 return True
229 async def _load_initial_keypair(self) -> tuple[ssl.SSLContext, str] | None:
230 """Load the keypair the listener starts with, waiting for it when configured.
232 With no wait budget this is the fail-closed load: any error propagates.
233 With a budget the load is retried every poll interval until it
234 succeeds; once the budget is spent the last load error is raised, so
235 a Secret that never appears still fails the container visibly. Returns
236 None when shutdown is requested during the wait (nothing was bound, so
237 there is nothing to drain).
238 """
239 budget = self.config.keypair_wait_seconds
240 if budget <= 0:
241 return _ssl_context(self.config)
242 loop = asyncio.get_running_loop()
243 started = loop.time()
244 deadline = started + budget
245 poll = max(self.config.poll_seconds, _MIN_KEYPAIR_WAIT_POLL_SECONDS)
246 next_log = started
247 waited = False
248 while True:
249 try:
250 keypair = _ssl_context(self.config)
251 except (OSError, RuntimeError, ssl.SSLError) as exc:
252 now = loop.time()
253 remaining = deadline - now
254 if remaining <= 0:
255 logger.error("No usable TLS keypair after waiting %.0fs", budget)
256 raise
257 if now >= next_log:
258 logger.warning(
259 "TLS keypair not loadable yet; retrying for up to %.0fs: %s",
260 remaining,
261 exc,
262 )
263 next_log = now + _KEYPAIR_WAIT_LOG_INTERVAL_SECONDS
264 waited = True
265 else:
266 if waited:
267 logger.info("TLS keypair loaded after %.1fs", loop.time() - started)
268 return keypair
269 if await self._wait_for_stop(min(poll, remaining)):
270 logger.info("Shutdown requested while waiting for the TLS keypair")
271 return None
273 async def start(self) -> None:
274 """Bind the listener once; every handshake is routed to the active keypair.
276 When a keypair wait is configured and shutdown arrives before the
277 keypair does, this returns without binding.
278 """
279 keypair = await self._load_initial_keypair()
280 if keypair is None:
281 return
282 self._active_context, self._keypair_digest = keypair
283 dispatcher = _server_context()
284 dispatcher.sni_callback = self._select_context
285 self._server = await asyncio.start_server(
286 self._handle_connection,
287 self.config.host,
288 self.config.port,
289 ssl=dispatcher,
290 )
291 logger.info(
292 "TLS proxy listening on https://%s:%d; upstream=http://%s:%d",
293 self.config.host,
294 self.config.port,
295 self.config.upstream_host,
296 self.config.upstream_port,
297 )
299 def activate_keypair(self, context: ssl.SSLContext, digest: str) -> None:
300 """Serve a validated rotated keypair to new handshakes.
302 The listener and every established stream are untouched: only the
303 context the next ClientHello is routed to changes.
304 """
305 self._active_context = context
306 self._keypair_digest = digest
307 logger.info("Activated the rotated workload certificate for new TLS handshakes")
309 async def watch_certificates(self) -> None:
310 """Activate atomically projected certificate changes without dropping streams."""
311 while not self._stop.is_set():
312 try:
313 await asyncio.wait_for(self._stop.wait(), timeout=self.config.poll_seconds)
314 except TimeoutError:
315 # The polling interval elapsed normally; inspect the projected
316 # keypair and activate it only when its content digest changed.
317 try:
318 context, digest = _ssl_context(self.config)
319 except OSError, RuntimeError, ssl.SSLError:
320 logger.exception("Rejected an unreadable or invalid rotated TLS keypair")
321 continue
322 if digest != self._keypair_digest:
323 self.activate_keypair(context, digest)
324 else:
325 break
327 async def shutdown(self) -> None:
328 """Stop accepting connections and drain established streams."""
329 self._stop.set()
330 server = self._server
331 if server is not None:
332 server.close()
334 active = set(self._connections)
335 if active:
336 try:
337 await asyncio.wait_for(
338 asyncio.gather(*active, return_exceptions=True),
339 timeout=self.config.graceful_shutdown_seconds,
340 )
341 except TimeoutError:
342 logger.warning(
343 "Cancelling %d TLS proxy connection(s) after the %.1fs drain budget",
344 len(active),
345 self.config.graceful_shutdown_seconds,
346 )
347 for task in active:
348 task.cancel()
349 await asyncio.gather(*active, return_exceptions=True)
351 if server is not None:
352 await server.wait_closed()
355async def run_proxy(config: ProxyConfig | None = None) -> None:
356 """Run until SIGTERM/SIGINT, then drain accepted proxy connections."""
357 proxy = TlsProxy(config or load_proxy_config())
358 loop = asyncio.get_running_loop()
359 logger.info("TLS proxy event loop: %s.%s", type(loop).__module__, type(loop).__qualname__)
360 for signum in (signal.SIGTERM, signal.SIGINT):
361 with contextlib.suppress(NotImplementedError), warnings.catch_warnings():
362 # uvloop 0.22 validates the callback with asyncio.iscoroutinefunction,
363 # which Python 3.14 deprecates. The warning is attributed to this
364 # call site and, with the module running as __main__, would print
365 # into every sidecar's log at startup; nothing here can act on it.
366 warnings.filterwarnings(
367 "ignore",
368 message=r"'asyncio\.iscoroutinefunction' is deprecated",
369 category=DeprecationWarning,
370 )
371 loop.add_signal_handler(signum, proxy._stop.set)
373 await proxy.start()
374 watcher = asyncio.create_task(proxy.watch_certificates())
375 stop_waiter = asyncio.create_task(proxy._stop.wait())
376 done, _pending = await asyncio.wait(
377 {watcher, stop_waiter},
378 return_when=asyncio.FIRST_COMPLETED,
379 )
380 watcher_error: BaseException | None = None
381 if watcher in done and not watcher.cancelled():
382 watcher_error = watcher.exception()
383 proxy._stop.set()
384 stop_waiter.cancel()
385 watcher.cancel()
386 await asyncio.gather(stop_waiter, watcher, return_exceptions=True)
387 await proxy.shutdown()
388 if watcher_error is not None:
389 raise watcher_error
392def _loop_factory() -> Callable[[], asyncio.AbstractEventLoop] | None:
393 """Prefer uvloop (shipped in the service images for uvicorn); else the stdlib loop."""
394 try:
395 import uvloop
396 except ImportError:
397 return None
398 return uvloop.new_event_loop
401def main() -> None:
402 logging.basicConfig(
403 level=os.getenv("LOG_LEVEL", "INFO").upper(),
404 format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
405 )
406 asyncio.run(run_proxy(), loop_factory=_loop_factory())
409if __name__ == "__main__":
410 main()