73 lines
2.4 KiB
Python
73 lines
2.4 KiB
Python
#!/usr/bin/env python3
|
|
"""Bounded-memory stream helpers for the Hermes SCM broker."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import tempfile
|
|
import time
|
|
import urllib.request
|
|
from typing import BinaryIO
|
|
|
|
from gitea_api_policy import PolicyError
|
|
|
|
|
|
class RejectRedirect(urllib.request.HTTPRedirectHandler):
|
|
"""Reject every upstream redirect before credentials can be forwarded."""
|
|
|
|
def redirect_request(self, req, fp, code, msg, headers, newurl):
|
|
raise PolicyError("SCM upstream redirects are not allowed")
|
|
|
|
|
|
def response_status(response: object) -> int | None:
|
|
"""Read an HTTP status from either urllib response interface."""
|
|
value = getattr(response, "status", None)
|
|
if value is None and hasattr(response, "getcode"):
|
|
value = response.getcode() # type: ignore[attr-defined]
|
|
return value
|
|
|
|
|
|
def spool_response(
|
|
response,
|
|
maximum: int,
|
|
forbidden: tuple[bytes, ...],
|
|
*,
|
|
memory_limit: int,
|
|
chunk_size: int,
|
|
deadline_seconds: float,
|
|
) -> tuple[BinaryIO, int]:
|
|
"""Spool an upstream response to disk while scanning chunk boundaries."""
|
|
spool = tempfile.SpooledTemporaryFile( # noqa: SIM115 - caller owns returned spool
|
|
max_size=memory_limit, dir="/tmp"
|
|
)
|
|
carry = b""
|
|
total = 0
|
|
deadline = time.monotonic() + deadline_seconds
|
|
try:
|
|
while True:
|
|
remaining = deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
raise PolicyError("SCM upstream response deadline exceeded")
|
|
raw = getattr(getattr(response, "fp", None), "raw", None)
|
|
sock = getattr(raw, "_sock", raw)
|
|
if hasattr(sock, "settimeout"):
|
|
sock.settimeout(remaining)
|
|
chunk = response.read(chunk_size)
|
|
if not chunk:
|
|
break
|
|
total += len(chunk)
|
|
if total > maximum:
|
|
raise PolicyError("SCM upstream response exceeds the safe size limit")
|
|
candidate = carry + chunk
|
|
if any(value in candidate for value in forbidden):
|
|
raise PolicyError(
|
|
"Git upstream response contains runtime credential material"
|
|
)
|
|
width = max(len(value) for value in forbidden) - 1
|
|
carry = candidate[-width:] if width else b""
|
|
spool.write(chunk)
|
|
spool.seek(0)
|
|
return spool, total
|
|
except Exception:
|
|
spool.close()
|
|
raise
|