Coverage for gco / services / tls_proxy.py: 100.00%
169 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""Hot-reloading TLS termination proxy for ALB-facing GCO API 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. Certificate files are treated as a pluggable projection:
6today cert-manager supplies a Secret volume; a future Kubernetes PodCertificate
7volume can replace it without changing the proxy or Service topology.
8"""
10from __future__ import annotations
12import asyncio
13import contextlib
14import hashlib
15import logging
16import math
17import os
18import signal
19import ssl
20from dataclasses import dataclass
21from pathlib import Path
22from typing import Any
24logger = logging.getLogger(__name__)
26TLS_CERT_FILE_ENV = "GCO_TLS_CERT_FILE"
27TLS_KEY_FILE_ENV = "GCO_TLS_KEY_FILE"
28DEFAULT_CERT_FILE = "/var/run/gco/tls/tls.crt"
29DEFAULT_KEY_FILE = "/var/run/gco/tls/tls.key"
30_BUFFER_BYTES = 64 * 1024
33@dataclass(frozen=True)
34class ProxyConfig:
35 """Validated listener, upstream, rotation, and shutdown settings."""
37 host: str
38 port: int
39 upstream_host: str
40 upstream_port: int
41 cert_file: Path
42 key_file: Path
43 poll_seconds: float
44 graceful_shutdown_seconds: float
47def _positive_port(name: str, default: int) -> int:
48 value = int(os.getenv(name, str(default)))
49 if not 1 <= value <= 65535:
50 raise RuntimeError(f"{name} must be between 1 and 65535")
51 return value
54def _non_negative_number(name: str, default: float) -> float:
55 value = float(os.getenv(name, str(default)))
56 if not math.isfinite(value) or value < 0:
57 raise RuntimeError(f"{name} must be a finite non-negative number")
58 return value
61def load_proxy_config() -> ProxyConfig:
62 """Resolve process configuration, failing closed on an incomplete keypair."""
63 cert_value = os.getenv(TLS_CERT_FILE_ENV, DEFAULT_CERT_FILE).strip()
64 key_value = os.getenv(TLS_KEY_FILE_ENV, DEFAULT_KEY_FILE).strip()
65 if not cert_value or not key_value:
66 raise RuntimeError(f"{TLS_CERT_FILE_ENV} and {TLS_KEY_FILE_ENV} must not be empty")
67 cert_file = Path(cert_value)
68 key_file = Path(key_value)
69 return ProxyConfig(
70 host=os.getenv("TLS_PROXY_HOST", "0.0.0.0"), # nosec B104 — pod listener
71 port=_positive_port("TLS_PROXY_PORT", 8443),
72 upstream_host=os.getenv("TLS_PROXY_UPSTREAM_HOST", "127.0.0.1"),
73 upstream_port=_positive_port("TLS_PROXY_UPSTREAM_PORT", 9000),
74 cert_file=cert_file,
75 key_file=key_file,
76 poll_seconds=_non_negative_number("TLS_PROXY_POLL_SECONDS", 5.0),
77 graceful_shutdown_seconds=_non_negative_number("GRACEFUL_SHUTDOWN_TIMEOUT_SECONDS", 30.0),
78 )
81def _keypair_digest(config: ProxyConfig) -> str:
82 """Return a digest of readable certificate material without logging it."""
83 digest = hashlib.sha256()
84 for variable, path in (
85 (TLS_CERT_FILE_ENV, config.cert_file),
86 (TLS_KEY_FILE_ENV, config.key_file),
87 ):
88 if not path.is_file() or not os.access(path, os.R_OK):
89 raise RuntimeError(f"{variable} does not reference a readable file: {path}")
90 digest.update(path.read_bytes())
91 return digest.hexdigest()
94def _ssl_context(config: ProxyConfig) -> tuple[ssl.SSLContext, str]:
95 """Build a TLS 1.2+ server context and return its keypair digest."""
96 digest = _keypair_digest(config)
97 context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
98 context.minimum_version = ssl.TLSVersion.TLSv1_2
99 context.load_cert_chain(config.cert_file, config.key_file)
100 return context, digest
103async def _close_writer(writer: asyncio.StreamWriter) -> None:
104 writer.close()
105 with contextlib.suppress(ConnectionError, OSError):
106 await writer.wait_closed()
109async def _pump(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
110 while data := await reader.read(_BUFFER_BYTES):
111 writer.write(data)
112 await writer.drain()
115class TlsProxy:
116 """TLS-only TCP proxy with certificate reload and graceful stream drain."""
118 def __init__(self, config: ProxyConfig) -> None:
119 self.config = config
120 self._server: asyncio.Server | None = None
121 self._keypair_digest = ""
122 self._stop = asyncio.Event()
123 self._connections: set[asyncio.Task[Any]] = set()
124 self._retired_acceptors: set[asyncio.Task[Any]] = set()
126 async def _handle_connection(
127 self,
128 client_reader: asyncio.StreamReader,
129 client_writer: asyncio.StreamWriter,
130 ) -> None:
131 task = asyncio.current_task()
132 if task is not None:
133 self._connections.add(task)
134 upstream_writer: asyncio.StreamWriter | None = None
135 try:
136 upstream_reader, upstream_writer = await asyncio.open_connection(
137 self.config.upstream_host,
138 self.config.upstream_port,
139 )
140 pumps = {
141 asyncio.create_task(_pump(client_reader, upstream_writer)),
142 asyncio.create_task(_pump(upstream_reader, client_writer)),
143 }
144 _done, pending = await asyncio.wait(
145 pumps,
146 return_when=asyncio.FIRST_COMPLETED,
147 )
148 for pending_task in pending:
149 pending_task.cancel()
150 await asyncio.gather(*pumps, return_exceptions=True)
151 except (ConnectionError, OSError) as exc:
152 logger.debug("TLS proxy connection ended before upstream was ready: %s", exc)
153 finally:
154 if upstream_writer is not None:
155 await _close_writer(upstream_writer)
156 await _close_writer(client_writer)
157 if task is not None:
158 self._connections.discard(task)
160 async def start(self) -> None:
161 """Start the cert-backed ALB listener."""
162 context, self._keypair_digest = _ssl_context(self.config)
163 self._server = await asyncio.start_server(
164 self._handle_connection,
165 self.config.host,
166 self.config.port,
167 ssl=context,
168 )
169 logger.info(
170 "TLS proxy listening on https://%s:%d; upstream=http://%s:%d",
171 self.config.host,
172 self.config.port,
173 self.config.upstream_host,
174 self.config.upstream_port,
175 )
177 async def _reload_certificate(self, context: ssl.SSLContext, digest: str) -> None:
178 old_server = self._server
179 if old_server is not None:
180 # ``Server.wait_closed`` waits for accepted clients on current
181 # Python releases. Closing the acceptor releases its listening
182 # socket synchronously; retire it in the background so a long-lived
183 # stream cannot block the replacement listener from binding.
184 old_server.close()
185 retired = asyncio.create_task(old_server.wait_closed())
186 self._retired_acceptors.add(retired)
187 retired.add_done_callback(self._retired_acceptors.discard)
188 try:
189 self._server = await asyncio.start_server(
190 self._handle_connection,
191 self.config.host,
192 self.config.port,
193 ssl=context,
194 )
195 except Exception:
196 self._server = None
197 self._stop.set()
198 logger.critical(
199 "TLS listener rebind failed after certificate rotation; exiting for restart",
200 exc_info=True,
201 )
202 raise
203 self._keypair_digest = digest
204 logger.info("Reloaded the TLS listener after workload certificate rotation")
206 async def watch_certificates(self) -> None:
207 """Reload atomically projected certificate changes without dropping streams."""
208 while not self._stop.is_set():
209 try:
210 await asyncio.wait_for(self._stop.wait(), timeout=self.config.poll_seconds)
211 except TimeoutError:
212 # The polling interval elapsed normally; inspect the projected
213 # keypair and reload only when its content digest changed.
214 try:
215 context, digest = _ssl_context(self.config)
216 except OSError, RuntimeError, ssl.SSLError:
217 logger.exception("Rejected an unreadable or invalid rotated TLS keypair")
218 continue
219 if digest != self._keypair_digest:
220 await self._reload_certificate(context, digest)
221 else:
222 break
224 async def shutdown(self) -> None:
225 """Stop accepting connections and drain established streams."""
226 self._stop.set()
227 current_server = self._server
228 if current_server is not None:
229 current_server.close()
231 active = set(self._connections)
232 if active:
233 try:
234 await asyncio.wait_for(
235 asyncio.gather(*active, return_exceptions=True),
236 timeout=self.config.graceful_shutdown_seconds,
237 )
238 except TimeoutError:
239 logger.warning(
240 "Cancelling %d TLS proxy connection(s) after the %.1fs drain budget",
241 len(active),
242 self.config.graceful_shutdown_seconds,
243 )
244 for task in active:
245 task.cancel()
246 await asyncio.gather(*active, return_exceptions=True)
248 acceptor_waiters: list[asyncio.Future[Any]] = list(self._retired_acceptors)
249 if current_server is not None:
250 acceptor_waiters.append(asyncio.ensure_future(current_server.wait_closed()))
251 if acceptor_waiters:
252 await asyncio.gather(*acceptor_waiters, return_exceptions=True)
255async def run_proxy(config: ProxyConfig | None = None) -> None:
256 """Run until SIGTERM/SIGINT, then drain accepted proxy connections."""
257 proxy = TlsProxy(config or load_proxy_config())
258 loop = asyncio.get_running_loop()
259 for signum in (signal.SIGTERM, signal.SIGINT):
260 with contextlib.suppress(NotImplementedError):
261 loop.add_signal_handler(signum, proxy._stop.set)
263 await proxy.start()
264 watcher = asyncio.create_task(proxy.watch_certificates())
265 stop_waiter = asyncio.create_task(proxy._stop.wait())
266 done, _pending = await asyncio.wait(
267 {watcher, stop_waiter},
268 return_when=asyncio.FIRST_COMPLETED,
269 )
270 watcher_error: BaseException | None = None
271 if watcher in done and not watcher.cancelled():
272 watcher_error = watcher.exception()
273 proxy._stop.set()
274 stop_waiter.cancel()
275 watcher.cancel()
276 await asyncio.gather(stop_waiter, watcher, return_exceptions=True)
277 await proxy.shutdown()
278 if watcher_error is not None:
279 raise watcher_error
282def main() -> None:
283 logging.basicConfig(
284 level=os.getenv("LOG_LEVEL", "INFO").upper(),
285 format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
286 )
287 asyncio.run(run_proxy())
290if __name__ == "__main__":
291 main()