Coverage for lambda / proxy-shared / proxy_utils.py: 100.00%
217 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +0000
1"""Shared, fail-closed utilities for API Gateway backend proxy Lambdas."""
3import hashlib
4import hmac
5import json
6import logging
7import os
8import re
9import secrets
10import threading
11import time
12from collections.abc import Mapping
13from typing import Any
14from urllib.parse import quote, urlencode, urlsplit, urlunsplit
16import boto3
17import urllib3
18from backend_tls import get_backend_http_pool
20# <pyflowchart-code-diagram> BEGIN - auto-inserted, do not edit
21# Generated at (UTC): 2026-08-30T12:00:00Z
22# Generated from Git commit: 8cb8e766461f5b7f35205e12db980988598cf6d5
23# Flowchart(s) generated from this file:
24# * ``build_signed_headers`` -> ``diagrams/code_diagrams/lambda/proxy-shared/proxy_utils.build_signed_headers.html``
25# (PNG: ``diagrams/code_diagrams/lambda/proxy-shared/proxy_utils.build_signed_headers.png``)
26# * ``forward_request`` -> ``diagrams/code_diagrams/lambda/proxy-shared/proxy_utils.forward_request.html``
27# (PNG: ``diagrams/code_diagrams/lambda/proxy-shared/proxy_utils.forward_request.png``)
28# Regenerate with ``SOURCE_DATE_EPOCH=<unix-seconds> GCO_DIAGRAM_SOURCE_COMMIT=<40-char-sha> python diagrams/generate.py --code-only``.
29# <pyflowchart-code-diagram> END
32logger = logging.getLogger(__name__)
34_HOP_BY_HOP_HEADERS = frozenset(
35 {
36 "connection",
37 "keep-alive",
38 "proxy-authenticate",
39 "proxy-authorization",
40 "te",
41 "trailer",
42 "transfer-encoding",
43 "upgrade",
44 }
45)
46_ALLOWED_REQUEST_HEADERS = frozenset(
47 {
48 "accept",
49 "accept-encoding",
50 "cache-control",
51 "content-encoding",
52 "content-type",
53 "idempotency-key",
54 "if-match",
55 "if-none-match",
56 "prefer",
57 "range",
58 "user-agent",
59 "x-request-id",
60 }
61)
62#: The SigV4 principal API Gateway authenticated, forwarded so the cluster API
63#: can attribute work to a caller (the Slurm routes record it in each job's
64#: comment). Only this proxy sets it, from the request context: a client
65#: sending its own ``x-gco-caller`` loses it in ``sanitize_request_headers``
66#: (not an allowed inbound header), and the service trusts the header only
67#: inside a request whose HMAC envelope verifies.
68CALLER_IDENTITY_HEADER = "x-gco-caller"
69_CALLER_ARN_PATTERN = re.compile(
70 r"^arn:aws[a-zA-Z-]*:(?:iam|sts)::\d{12}:[A-Za-z0-9._/@+=,:-]{1,256}$"
71)
72_INTERNAL_SIGNATURE_HEADERS = frozenset(
73 {
74 "x-gco-signature-version",
75 "x-gco-signature",
76 "x-gco-timestamp",
77 "x-gco-nonce",
78 "x-gco-content-sha256",
79 CALLER_IDENTITY_HEADER,
80 }
81)
82_RETRYABLE_STATUS_CODES = frozenset({429, 502, 503, 504})
83_RETRYABLE_METHODS = frozenset({"GET", "HEAD", "OPTIONS"})
84# The one upstream answer a mutating request may be replayed after: the
85# Application Load Balancer's own 503. The load balancer serves it when it has
86# no target to send the request to (no registered or healthy target, or the
87# target optimizer refused the request), which is a decision it makes before
88# selecting a target, so the request provably never reached the backend and a
89# replay cannot duplicate a write. The page below is the fixed one the load
90# balancer serves, captured verbatim from an ALB with an empty target group:
91# ``HTTP/1.1 503 Service Temporarily Unavailable``, ``Server: awselb/2.0``,
92# ``Content-Type: text/html``, ``Content-Length: 162``. A target's 503 (the
93# services answer JSON with their own ``Server`` header) and nginx's page (an
94# extra ``<hr><center>nginx</center>`` line) do not match. 502 and 504 are
95# never in this category: the load balancer produces them after it has handed
96# the request to a target.
97_LOAD_BALANCER_503_PAGE = re.compile(
98 rb"<html>\r?\n"
99 rb"<head><title>503 Service Temporarily Unavailable</title></head>\r?\n"
100 rb"<body>\r?\n"
101 rb"<center><h1>503 Service Temporarily Unavailable</h1></center>\r?\n"
102 rb"</body>\r?\n"
103 rb"</html>(?:\r?\n)?"
104)
105_LOAD_BALANCER_503_MAX_BODY_BYTES = 512
106_LOAD_BALANCER_SERVER_HEADER = "awselb/2.0"
109def _response_header(response: Any, name: str) -> str | None:
110 """Return one response header by case-insensitive name, or None."""
111 headers = getattr(response, "headers", None)
112 if not isinstance(headers, Mapping):
113 return None
114 wanted = name.lower()
115 for key, value in headers.items():
116 if str(key).lower() == wanted:
117 return str(value)
118 return None
121def _is_load_balancer_503(response: Any) -> bool:
122 """True only for the load balancer's own 503 page (see the constants above).
124 Fails closed: any 503 that does not carry the load balancer's fixed page,
125 or that names a different server, is treated as a target's answer. The
126 ``Server`` header is a gate rather than the signal because operators can
127 disable it on the listener (``routing.http.response.server.enabled``).
128 """
129 if response.status != 503:
130 return False
131 server = _response_header(response, "server")
132 if server is not None and server.strip().lower() != _LOAD_BALANCER_SERVER_HEADER:
133 return False
134 data = getattr(response, "data", None)
135 if not isinstance(data, bytes | bytearray) or len(data) > _LOAD_BALANCER_503_MAX_BODY_BYTES:
136 return False
137 return _LOAD_BALANCER_503_PAGE.fullmatch(bytes(data)) is not None
140def _bounded_env_float(name: str, default: float, minimum: float, maximum: float) -> float:
141 try:
142 value = float(os.getenv(name, str(default)))
143 except ValueError:
144 return default
145 return value if minimum <= value <= maximum else default
148def _bounded_env_int(name: str, default: int, minimum: int, maximum: int) -> int:
149 try:
150 value = int(os.getenv(name, str(default)))
151 except ValueError:
152 return default
153 return value if minimum <= value <= maximum else default
156_secret_lock = threading.Lock()
157_cached_secret: str | None = None
158_cache_timestamp = 0.0
159_last_successful_refresh = 0.0
160_last_refresh_attempt = 0.0
161_CACHE_TTL_SECONDS = _bounded_env_float("SECRET_CACHE_TTL_SECONDS", 300.0, 1.0, 3600.0)
162_CACHE_MAX_STALE_SECONDS = max(
163 _CACHE_TTL_SECONDS,
164 _bounded_env_float("SECRET_CACHE_MAX_STALE_SECONDS", 900.0, 1.0, 7200.0),
165)
166_CACHE_RETRY_SECONDS = _bounded_env_float("SECRET_CACHE_RETRY_SECONDS", 5.0, 0.1, 60.0)
169def _secret_region(secret_arn: str) -> str | None:
170 """Return the owning region for a Secrets Manager ARN, if present."""
171 parts = secret_arn.split(":", 5)
172 if len(parts) == 6 and parts[0] == "arn" and parts[2] == "secretsmanager":
173 return parts[3] or None
174 return None
177# The regional VPC proxy runs outside the API Gateway region where the shared
178# HMAC key lives. Secrets Manager clients do not route cross-region ARNs to the
179# owning endpoint automatically, so bind the client to the ARN's region.
180_secrets_client = boto3.client(
181 "secretsmanager",
182 region_name=_secret_region(os.getenv("SECRET_ARN", "")),
183)
186def get_secret_token() -> str:
187 """Return a cached signing key, with a strictly bounded stale grace period."""
188 global _cached_secret, _cache_timestamp, _last_successful_refresh, _last_refresh_attempt
190 now = time.monotonic()
191 age = now - _last_successful_refresh
192 if _cached_secret is not None and age < _CACHE_TTL_SECONDS:
193 return _cached_secret
194 if (
195 _cached_secret is not None
196 and age <= _CACHE_MAX_STALE_SECONDS
197 and now - _last_refresh_attempt < _CACHE_RETRY_SECONDS
198 ):
199 return _cached_secret
201 with _secret_lock:
202 now = time.monotonic()
203 age = now - _last_successful_refresh
204 if _cached_secret is not None and age < _CACHE_TTL_SECONDS:
205 return _cached_secret
206 if (
207 _cached_secret is not None
208 and age <= _CACHE_MAX_STALE_SECONDS
209 and now - _last_refresh_attempt < _CACHE_RETRY_SECONDS
210 ):
211 return _cached_secret
213 _last_refresh_attempt = now
214 try:
215 response = _secrets_client.get_secret_value(SecretId=os.environ["SECRET_ARN"])
216 secret_data = json.loads(response["SecretString"])
217 token = secret_data.get("token")
218 if not isinstance(token, str) or not token:
219 raise ValueError("secret token is missing")
220 except Exception as error:
221 if _cached_secret is not None and age <= _CACHE_MAX_STALE_SECONDS:
222 logger.warning("Secrets Manager refresh failed; using bounded stale signing key")
223 return _cached_secret
224 raise RuntimeError("Authentication signing key is unavailable") from error
226 _cached_secret = token
227 _last_successful_refresh = now
228 _cache_timestamp = now # Backward-compatible observability/test alias.
229 return token
232def sanitize_request_headers(headers: dict[str, Any]) -> dict[str, str]:
233 """Apply a case-insensitive end-to-end allowlist at the IAM trust boundary."""
234 sanitized: dict[str, str] = {}
235 for name, value in headers.items():
236 normalized = str(name).strip().lower()
237 if normalized not in _ALLOWED_REQUEST_HEADERS or value is None:
238 continue
239 sanitized[normalized] = str(value)
240 return sanitized
243def _request_target(target_url: str) -> str:
244 parsed = urlsplit(target_url)
245 return parsed.path + (f"?{parsed.query}" if parsed.query else "")
248def build_signed_headers(
249 signing_key: str,
250 http_method: str,
251 target_url: str,
252 body: str | None,
253) -> dict[str, str]:
254 """Build short-lived request authentication without transmitting the signing key."""
255 timestamp = str(int(time.time()))
256 nonce = secrets.token_hex(16)
257 content_hash = hashlib.sha256((body or "").encode("utf-8")).hexdigest()
258 canonical = "\n".join(
259 ["v1", timestamp, nonce, http_method.upper(), _request_target(target_url), content_hash]
260 )
261 signature = hmac.new(
262 signing_key.encode("utf-8"),
263 canonical.encode("utf-8"),
264 hashlib.sha256,
265 ).hexdigest()
266 return {
267 "x-gco-signature-version": "v1",
268 "x-gco-signature": signature,
269 "x-gco-timestamp": timestamp,
270 "x-gco-nonce": nonce,
271 "x-gco-content-sha256": content_hash,
272 }
275def caller_identity_headers(event: Mapping[str, Any]) -> dict[str, str]:
276 """The ``x-gco-caller`` header for the principal API Gateway authenticated.
278 Reads ``requestContext.identity.userArn`` (the IAM user or assumed-role
279 session ARN API Gateway resolved from the SigV4 signature). Anything that
280 is not a well-formed IAM/STS ARN yields no header, so the service falls
281 back to attributing the work to the API itself rather than to a guess.
282 """
283 context = event.get("requestContext")
284 identity = context.get("identity") if isinstance(context, Mapping) else None
285 arn = identity.get("userArn") if isinstance(identity, Mapping) else None
286 if isinstance(arn, str) and _CALLER_ARN_PATTERN.fullmatch(arn):
287 return {CALLER_IDENTITY_HEADER: arn}
288 return {}
291def _outbound_headers(headers: dict[str, Any]) -> dict[str, str]:
292 """Re-validate caller headers while retaining only generated auth fields."""
293 allowed = _ALLOWED_REQUEST_HEADERS | _INTERNAL_SIGNATURE_HEADERS
294 return {
295 str(name).lower(): str(value)
296 for name, value in headers.items()
297 if str(name).lower() in allowed and value is not None
298 }
301_MAX_RETRIES = _bounded_env_int("PROXY_MAX_RETRIES", 3, 1, 5)
302_RETRY_BACKOFF_BASE = _bounded_env_float("PROXY_RETRY_BACKOFF_BASE", 0.3, 0.0, 5.0)
303_http: urllib3.PoolManager | None = None
306def _tls_failure_response() -> dict[str, Any]:
307 """Return a bounded error without exposing certificate or trust details."""
308 return {
309 "statusCode": 502,
310 "headers": {"Content-Type": "application/json"},
311 "body": json.dumps({"error": "Backend TLS verification failed"}),
312 }
315def forward_request(
316 target_url: str,
317 http_method: str,
318 headers: dict[str, str],
319 body: str | None,
320 timeout: float = 29.0,
321) -> dict[str, Any]:
322 """Forward over authenticated TLS within one deadline.
324 Read-only methods replay retryable statuses and transport failures.
325 Mutating methods are replayed only after the load balancer's own 503,
326 which it serves before selecting a target (see ``_is_load_balancer_503``);
327 every other answer or failure may have reached a target and is relayed
328 after exactly one attempt. Plaintext, non-443, and credential-bearing
329 targets are rejected before any network request.
330 """
331 parsed_target = urlsplit(target_url)
332 try:
333 target_port = parsed_target.port
334 except ValueError as exc:
335 raise ValueError("Backend proxy target has an invalid port") from exc
336 if (
337 parsed_target.scheme.lower() != "https"
338 or not parsed_target.hostname
339 or parsed_target.username is not None
340 or parsed_target.password is not None
341 or target_port not in {None, 443}
342 ):
343 raise ValueError("Backend proxy targets must use HTTPS on port 443")
345 # Anchor the deadline before acquiring transport: the caller computed the
346 # budget from the Lambda's remaining time, so a cold-start trust-bundle
347 # refresh must consume this budget. Anchoring after it extended the wall
348 # clock past the Lambda timeout, killing the function mid-flight instead
349 # of returning its bounded 504 when the backend black-holed.
350 deadline = time.monotonic() + max(timeout, 0.0)
352 try:
353 transport = _http or get_backend_http_pool()
354 except RuntimeError:
355 logger.exception("Backend TLS trust is unavailable")
356 return {
357 "statusCode": 503,
358 "headers": {"Content-Type": "application/json"},
359 "body": json.dumps({"error": "Backend trust is temporarily unavailable"}),
360 }
362 method = http_method.upper()
363 encoded_body = body.encode("utf-8") if body else None
364 # Every method shares one attempt budget; what differs is which upstream
365 # answers may spend it. A read-only request replays any retryable status
366 # or transport failure. A mutating request replays only the load
367 # balancer's own 503, the one answer proven not to have reached a target;
368 # anything else it relays, and after a transport failure it stops, because
369 # the write may already have been applied.
370 read_only = method in _RETRYABLE_METHODS
371 max_attempts = _MAX_RETRIES
372 last_exception: Exception | None = None
373 last_response: urllib3.BaseHTTPResponse | None = None
374 attempts_made = 0
376 for attempt in range(max_attempts):
377 remaining = deadline - time.monotonic()
378 if remaining <= 0:
379 break
380 attempts_made = attempt + 1
381 try:
382 response = transport.request(
383 method,
384 target_url,
385 headers=_outbound_headers(headers),
386 body=encoded_body,
387 timeout=urllib3.Timeout(total=remaining),
388 )
389 if response.status not in _RETRYABLE_STATUS_CODES:
390 return _build_success_response(response)
391 last_response = response
392 if not read_only and not _is_load_balancer_503(response):
393 logger.warning(
394 "Upstream status %d for %s may have reached a target; relaying it unreplayed",
395 response.status,
396 method,
397 )
398 return _build_success_response(response)
399 logger.warning(
400 "Retryable upstream status %d on attempt %d/%d for %s%s",
401 response.status,
402 attempt + 1,
403 max_attempts,
404 method,
405 "" if read_only else " (load balancer answered before selecting a target)",
406 )
407 if attempt == max_attempts - 1:
408 return _build_success_response(response)
409 response.release_conn()
410 except urllib3.exceptions.SSLError:
411 logger.exception("Backend TLS verification failed")
412 return _tls_failure_response()
413 except (urllib3.exceptions.MaxRetryError, urllib3.exceptions.TimeoutError) as error:
414 if isinstance(getattr(error, "reason", None), urllib3.exceptions.SSLError):
415 logger.exception("Backend TLS verification failed")
416 return _tls_failure_response()
417 last_exception = error
418 if not read_only:
419 # The write may have been applied: report this failure
420 # instead of replaying, and instead of relaying an earlier
421 # load balancer 503 that would claim nothing reached the
422 # backend.
423 logger.warning(
424 "Upstream %s failed on attempt %d; not replaying a mutating request",
425 method,
426 attempt + 1,
427 )
428 last_response = None
429 break
430 logger.warning(
431 "Upstream %s failed on attempt %d/%d",
432 method,
433 attempt + 1,
434 max_attempts,
435 )
436 except Exception:
437 logger.exception("Unexpected proxy forwarding failure")
438 return {
439 "statusCode": 500,
440 "headers": {"Content-Type": "application/json"},
441 "body": json.dumps({"error": "Internal server error"}),
442 }
444 if attempt < max_attempts - 1:
445 remaining = deadline - time.monotonic()
446 backoff = _RETRY_BACKOFF_BASE * (2**attempt)
447 if remaining <= backoff:
448 break
449 time.sleep(backoff)
451 if last_response is not None:
452 return _build_success_response(last_response)
453 if last_exception is not None:
454 status_code = 503 if isinstance(last_exception, urllib3.exceptions.MaxRetryError) else 504
455 return {
456 "statusCode": status_code,
457 "headers": {"Content-Type": "application/json"},
458 "body": json.dumps(
459 {
460 "error": "Service unavailable" if status_code == 503 else "Gateway timeout",
461 "message": f"Upstream failed after {attempts_made} attempt(s)",
462 }
463 ),
464 }
465 return {
466 "statusCode": 504,
467 "headers": {"Content-Type": "application/json"},
468 "body": json.dumps({"error": "Gateway timeout"}),
469 }
472def _build_success_response(response: urllib3.BaseHTTPResponse) -> dict[str, Any]:
473 """Build an API Gateway response without hop-by-hop framing headers."""
474 response_headers = {
475 str(name): str(value)
476 for name, value in response.headers.items()
477 if str(name).lower() not in _HOP_BY_HOP_HEADERS
478 }
479 return {
480 "statusCode": response.status,
481 "headers": response_headers,
482 "body": response.data.decode("utf-8"),
483 }
486def build_target_url(
487 endpoint: str,
488 path: str,
489 query_params: dict[str, str | list[str]] | None,
490) -> str:
491 """Build one HTTPS/443 upstream URL without losing repeated query keys."""
492 base_url = endpoint if "://" in endpoint else f"https://{endpoint}"
493 parsed_endpoint = urlsplit(base_url)
494 try:
495 endpoint_port = parsed_endpoint.port
496 except ValueError as exc:
497 raise ValueError(f"Invalid proxy endpoint: {endpoint!r}") from exc
498 if (
499 parsed_endpoint.scheme.lower() != "https"
500 or not parsed_endpoint.hostname
501 or parsed_endpoint.username is not None
502 or parsed_endpoint.password is not None
503 or endpoint_port not in {None, 443}
504 or parsed_endpoint.query
505 or parsed_endpoint.fragment
506 ):
507 raise ValueError(f"Proxy endpoint must use HTTPS on port 443: {endpoint!r}")
509 request_path = path if path.startswith("/") else f"/{path}"
510 request_path = re.sub(r"%(?![0-9A-Fa-f]{2})", "%25", request_path)
511 encoded_path = quote(request_path, safe="/:@-._~!$&'()*+,;=%")
512 endpoint_path = parsed_endpoint.path.rstrip("/")
513 return urlunsplit(
514 (
515 parsed_endpoint.scheme,
516 parsed_endpoint.netloc,
517 f"{endpoint_path}{encoded_path}",
518 urlencode(query_params or {}, doseq=True),
519 "",
520 )
521 )