248 lines
7.9 KiB
Python
248 lines
7.9 KiB
Python
from __future__ import annotations
|
|
|
|
import socket
|
|
import ssl
|
|
import select
|
|
from collections.abc import Mapping
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass
|
|
from typing import Any, Iterator
|
|
|
|
import httpcore
|
|
import httpx
|
|
|
|
from govoplan_core.security.outbound_http import (
|
|
OutboundHttpError,
|
|
bounded_chunks_bytes,
|
|
create_outbound_connection,
|
|
response_limit,
|
|
validate_outbound_http_url,
|
|
)
|
|
|
|
|
|
class ConnectorHttpError(RuntimeError):
|
|
pass
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class ConnectorHttpResponse:
|
|
status_code: int
|
|
headers: Mapping[str, str]
|
|
content: bytes
|
|
|
|
|
|
_HTTPCORE_TRANSPORT_ERRORS = (
|
|
httpcore.TimeoutException,
|
|
httpcore.NetworkError,
|
|
httpcore.ProtocolError,
|
|
httpcore.ProxyError,
|
|
httpcore.UnsupportedProtocol,
|
|
)
|
|
|
|
|
|
class _SocketNetworkStream(httpcore.NetworkStream):
|
|
"""Public httpcore NetworkStream adapter around an approved socket."""
|
|
|
|
def __init__(self, sock: socket.socket) -> None:
|
|
self._socket = sock
|
|
|
|
def read(self, max_bytes: int, timeout: float | None = None) -> bytes:
|
|
try:
|
|
self._socket.settimeout(timeout)
|
|
return self._socket.recv(max_bytes)
|
|
except socket.timeout as exc:
|
|
raise httpcore.ReadTimeout(str(exc)) from exc
|
|
except OSError as exc:
|
|
raise httpcore.ReadError(str(exc)) from exc
|
|
|
|
def write(self, buffer: bytes, timeout: float | None = None) -> None:
|
|
try:
|
|
self._socket.settimeout(timeout)
|
|
self._socket.sendall(buffer)
|
|
except socket.timeout as exc:
|
|
raise httpcore.WriteTimeout(str(exc)) from exc
|
|
except OSError as exc:
|
|
raise httpcore.WriteError(str(exc)) from exc
|
|
|
|
def close(self) -> None:
|
|
self._socket.close()
|
|
|
|
def start_tls(
|
|
self,
|
|
ssl_context: ssl.SSLContext,
|
|
server_hostname: str | None = None,
|
|
timeout: float | None = None,
|
|
) -> httpcore.NetworkStream:
|
|
try:
|
|
self._socket.settimeout(timeout)
|
|
tls_socket = ssl_context.wrap_socket(self._socket, server_hostname=server_hostname)
|
|
except socket.timeout as exc:
|
|
self.close()
|
|
raise httpcore.ConnectTimeout(str(exc)) from exc
|
|
except OSError as exc:
|
|
self.close()
|
|
raise httpcore.ConnectError(str(exc)) from exc
|
|
return _SocketNetworkStream(tls_socket)
|
|
|
|
def get_extra_info(self, info: str) -> Any:
|
|
if info == "ssl_object" and isinstance(self._socket, ssl.SSLSocket):
|
|
return self._socket
|
|
if info == "client_addr":
|
|
return self._socket.getsockname()
|
|
if info == "server_addr":
|
|
return self._socket.getpeername()
|
|
if info == "socket":
|
|
return self._socket
|
|
if info == "is_readable":
|
|
try:
|
|
return bool(select.select([self._socket], [], [], 0)[0])
|
|
except (OSError, ValueError):
|
|
return True
|
|
return None
|
|
|
|
|
|
class _OutboundPolicyNetworkBackend(httpcore.NetworkBackend):
|
|
def connect_tcp(
|
|
self,
|
|
host: str,
|
|
port: int,
|
|
timeout: float | None = None,
|
|
local_address: str | None = None,
|
|
socket_options: Any = None,
|
|
) -> httpcore.NetworkStream:
|
|
source_address = None if local_address is None else (local_address, 0)
|
|
try:
|
|
sock = create_outbound_connection(
|
|
host,
|
|
port,
|
|
timeout=timeout,
|
|
source_address=source_address,
|
|
socket_options=socket_options,
|
|
label="File connector HTTP endpoint",
|
|
)
|
|
except socket.timeout as exc:
|
|
raise httpcore.ConnectTimeout(str(exc)) from exc
|
|
except (OSError, OutboundHttpError) as exc:
|
|
raise httpcore.ConnectError(str(exc)) from exc
|
|
return _SocketNetworkStream(sock)
|
|
|
|
def connect_unix_socket(self, path: str, timeout: float | None = None, socket_options: Any = None): # type: ignore[no-untyped-def]
|
|
del path, timeout, socket_options
|
|
raise httpcore.ConnectError("Unix sockets are not supported for file connectors")
|
|
|
|
|
|
class _HttpcoreResponseStream(httpx.SyncByteStream):
|
|
def __init__(self, stream: Any, *, request: httpx.Request) -> None:
|
|
self._stream = stream
|
|
self._request = request
|
|
|
|
def __iter__(self): # type: ignore[no-untyped-def]
|
|
try:
|
|
yield from self._stream
|
|
except _HTTPCORE_TRANSPORT_ERRORS as exc:
|
|
raise httpx.TransportError(str(exc), request=self._request) from exc
|
|
|
|
def close(self) -> None:
|
|
if hasattr(self._stream, "close"):
|
|
self._stream.close()
|
|
|
|
|
|
class _OutboundPolicyHTTPTransport(httpx.BaseTransport):
|
|
def __init__(self) -> None:
|
|
self._connection_pool = httpcore.ConnectionPool(
|
|
ssl_context=httpx.create_ssl_context(verify=True, trust_env=False),
|
|
max_connections=100,
|
|
max_keepalive_connections=20,
|
|
keepalive_expiry=5.0,
|
|
network_backend=_OutboundPolicyNetworkBackend(),
|
|
)
|
|
|
|
def handle_request(self, request: httpx.Request) -> httpx.Response:
|
|
core_request = httpcore.Request(
|
|
method=request.method,
|
|
url=httpcore.URL(
|
|
scheme=request.url.raw_scheme,
|
|
host=request.url.raw_host,
|
|
port=request.url.port,
|
|
target=request.url.raw_path,
|
|
),
|
|
headers=request.headers.raw,
|
|
content=request.stream,
|
|
extensions=request.extensions,
|
|
)
|
|
try:
|
|
response = self._connection_pool.handle_request(core_request)
|
|
except _HTTPCORE_TRANSPORT_ERRORS as exc:
|
|
raise httpx.TransportError(str(exc), request=request) from exc
|
|
return httpx.Response(
|
|
status_code=response.status,
|
|
headers=response.headers,
|
|
stream=_HttpcoreResponseStream(response.stream, request=request),
|
|
extensions=response.extensions,
|
|
)
|
|
|
|
def close(self) -> None:
|
|
self._connection_pool.close()
|
|
|
|
|
|
_CONNECTOR_HTTP_CLIENT = httpx.Client(
|
|
transport=_OutboundPolicyHTTPTransport(),
|
|
follow_redirects=False,
|
|
timeout=15.0,
|
|
)
|
|
|
|
|
|
@contextmanager
|
|
def _stream_connector_request(method: str, url: str, **kwargs: Any) -> Iterator[httpx.Response]:
|
|
timeout = kwargs.pop("timeout", 15.0)
|
|
with _CONNECTOR_HTTP_CLIENT.stream(
|
|
method,
|
|
url,
|
|
timeout=timeout,
|
|
**kwargs,
|
|
) as response:
|
|
yield response
|
|
|
|
|
|
def request_connector_bytes(
|
|
method: str,
|
|
url: str,
|
|
*,
|
|
headers: Mapping[str, str] | None = None,
|
|
params: Mapping[str, str] | None = None,
|
|
data: Mapping[str, str] | bytes | str | None = None,
|
|
content: bytes | str | None = None,
|
|
auth: Any = None,
|
|
timeout: float = 15.0,
|
|
kind: str = "structured",
|
|
max_bytes: int | None = None,
|
|
label: str = "File connector",
|
|
) -> ConnectorHttpResponse:
|
|
try:
|
|
validated_url = validate_outbound_http_url(url, label=f"{label} URL")
|
|
effective_limit = response_limit(kind) if max_bytes is None else min(int(max_bytes), response_limit(kind))
|
|
with _stream_connector_request(
|
|
method,
|
|
validated_url,
|
|
headers=dict(headers or {}),
|
|
params=params,
|
|
data=data,
|
|
content=content,
|
|
auth=auth,
|
|
timeout=timeout,
|
|
) as response:
|
|
body = bounded_chunks_bytes(
|
|
response.iter_bytes(),
|
|
headers=response.headers,
|
|
max_bytes=effective_limit,
|
|
kind=kind,
|
|
label=f"{label} response",
|
|
)
|
|
return ConnectorHttpResponse(
|
|
status_code=response.status_code,
|
|
headers=dict(response.headers),
|
|
content=body,
|
|
)
|
|
except (httpx.HTTPError, OutboundHttpError, ValueError) as exc:
|
|
raise ConnectorHttpError(str(exc)) from exc
|