Pin S3 and SMB connector peers
This commit is contained in:
@@ -0,0 +1,331 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
import socket
|
||||
import threading
|
||||
from functools import lru_cache
|
||||
from importlib import import_module
|
||||
from typing import Any
|
||||
|
||||
from govoplan_core.security.outbound_http import (
|
||||
OutboundHttpError,
|
||||
create_outbound_connection,
|
||||
)
|
||||
|
||||
|
||||
class SdkPeerPinningError(RuntimeError):
|
||||
"""Raised when an optional SDK cannot be bound to the pinned transport."""
|
||||
|
||||
|
||||
_BOTOCORE_CLIENT_CREATION_LOCK = threading.Lock()
|
||||
|
||||
|
||||
def create_pinned_s3_client(**kwargs: Any) -> Any:
|
||||
"""Construct an S3 client whose first and every later socket is pinned.
|
||||
|
||||
Botocore does not expose the HTTP-session class through boto3's public
|
||||
client API. Its endpoint creator does expose that seam, so this function
|
||||
replaces the creator only for the bounded client-construction operation.
|
||||
The lock avoids an unsafe interleaving with another GovOPlaN client
|
||||
construction. Any unrelated botocore client created during the short
|
||||
replacement window also receives the stricter transport.
|
||||
"""
|
||||
|
||||
try:
|
||||
boto3 = import_module("boto3")
|
||||
botocore_args = import_module("botocore.args")
|
||||
except ImportError as exc:
|
||||
raise SdkPeerPinningError(
|
||||
"S3 connector browsing requires the optional boto3 dependency"
|
||||
) from exc
|
||||
|
||||
pinned_session_cls = _botocore_transport_types()[0]
|
||||
original_creator = getattr(botocore_args, "EndpointCreator", None)
|
||||
if original_creator is None or not hasattr(original_creator, "create_endpoint"):
|
||||
raise SdkPeerPinningError(
|
||||
"The installed botocore release does not expose the endpoint transport seam required for peer pinning"
|
||||
)
|
||||
|
||||
class PinnedEndpointCreator(original_creator): # type: ignore[misc, valid-type]
|
||||
def create_endpoint(self, *args: Any, **endpoint_kwargs: Any) -> Any:
|
||||
endpoint_kwargs["http_session_cls"] = pinned_session_cls
|
||||
return super().create_endpoint(*args, **endpoint_kwargs)
|
||||
|
||||
with _BOTOCORE_CLIENT_CREATION_LOCK:
|
||||
current_creator = getattr(botocore_args, "EndpointCreator", None)
|
||||
if current_creator is not original_creator:
|
||||
raise SdkPeerPinningError(
|
||||
"Botocore endpoint construction changed concurrently; refusing to create an unproven S3 client"
|
||||
)
|
||||
botocore_args.EndpointCreator = PinnedEndpointCreator
|
||||
try:
|
||||
session_type = getattr(getattr(boto3, "session", None), "Session", None)
|
||||
if session_type is None:
|
||||
raise SdkPeerPinningError(
|
||||
"The installed boto3 release does not expose its isolated session constructor"
|
||||
)
|
||||
client = session_type().client("s3", **kwargs)
|
||||
finally:
|
||||
if getattr(botocore_args, "EndpointCreator", None) is PinnedEndpointCreator:
|
||||
botocore_args.EndpointCreator = original_creator
|
||||
|
||||
http_session = getattr(getattr(client, "_endpoint", None), "http_session", None)
|
||||
if http_session is None or not isinstance(http_session, pinned_session_cls):
|
||||
close = getattr(client, "close", None)
|
||||
if callable(close):
|
||||
close()
|
||||
raise SdkPeerPinningError(
|
||||
"Botocore did not install the required pinned HTTP transport; the S3 client was discarded"
|
||||
)
|
||||
return client
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _botocore_transport_types() -> tuple[type[Any], type[Any], type[Any]]:
|
||||
try:
|
||||
awsrequest = import_module("botocore.awsrequest")
|
||||
httpsession = import_module("botocore.httpsession")
|
||||
urllib3_exceptions = import_module("urllib3.exceptions")
|
||||
except ImportError as exc:
|
||||
raise SdkPeerPinningError(
|
||||
"S3 connector browsing requires a compatible botocore HTTP transport"
|
||||
) from exc
|
||||
|
||||
required = (
|
||||
"AWSHTTPConnection",
|
||||
"AWSHTTPSConnection",
|
||||
"AWSHTTPConnectionPool",
|
||||
"AWSHTTPSConnectionPool",
|
||||
)
|
||||
if any(not hasattr(awsrequest, name) for name in required) or not hasattr(
|
||||
httpsession, "URLLib3Session"
|
||||
):
|
||||
raise SdkPeerPinningError(
|
||||
"The installed botocore release is missing the connection classes required for peer pinning"
|
||||
)
|
||||
|
||||
def pinned_new_connection(connection: Any) -> socket.socket:
|
||||
hostname = str(getattr(connection, "_dns_host", "") or "").strip()
|
||||
port = int(getattr(connection, "port", 0) or 0)
|
||||
if not hostname or not port:
|
||||
raise urllib3_exceptions.NewConnectionError(
|
||||
connection, "Pinned S3 connection is missing its target authority"
|
||||
)
|
||||
try:
|
||||
return create_outbound_connection(
|
||||
hostname,
|
||||
port,
|
||||
timeout=getattr(connection, "timeout", None),
|
||||
source_address=getattr(connection, "source_address", None),
|
||||
socket_options=getattr(connection, "socket_options", None),
|
||||
label="S3 connector peer",
|
||||
)
|
||||
except socket.timeout as exc:
|
||||
raise urllib3_exceptions.ConnectTimeoutError(
|
||||
connection,
|
||||
f"Connection to {hostname} timed out while selecting an approved peer",
|
||||
) from exc
|
||||
except (OSError, OutboundHttpError, ValueError) as exc:
|
||||
raise urllib3_exceptions.NewConnectionError(
|
||||
connection,
|
||||
f"S3 connector peer was rejected: {exc}",
|
||||
) from exc
|
||||
|
||||
class PinnedAWSHTTPConnection(awsrequest.AWSHTTPConnection):
|
||||
_new_conn = pinned_new_connection
|
||||
|
||||
class PinnedAWSHTTPSConnection(awsrequest.AWSHTTPSConnection):
|
||||
_new_conn = pinned_new_connection
|
||||
|
||||
class PinnedAWSHTTPConnectionPool(awsrequest.AWSHTTPConnectionPool):
|
||||
ConnectionCls = PinnedAWSHTTPConnection
|
||||
|
||||
class PinnedAWSHTTPSConnectionPool(awsrequest.AWSHTTPSConnectionPool):
|
||||
ConnectionCls = PinnedAWSHTTPSConnection
|
||||
|
||||
class PinnedURLLib3Session(httpsession.URLLib3Session):
|
||||
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
||||
# A proxy would select the final target outside this process. Until
|
||||
# proxy peer delegation is modeled, connector traffic is direct.
|
||||
kwargs["proxies"] = {}
|
||||
super().__init__(*args, **kwargs)
|
||||
self._pool_classes_by_scheme = {
|
||||
"http": PinnedAWSHTTPConnectionPool,
|
||||
"https": PinnedAWSHTTPSConnectionPool,
|
||||
}
|
||||
manager = getattr(self, "_manager", None)
|
||||
if manager is None or not hasattr(manager, "pool_classes_by_scheme"):
|
||||
raise SdkPeerPinningError(
|
||||
"The installed botocore pool manager cannot enforce pinned connection classes"
|
||||
)
|
||||
manager.pool_classes_by_scheme = self._pool_classes_by_scheme
|
||||
|
||||
return PinnedURLLib3Session, PinnedAWSHTTPConnection, PinnedAWSHTTPSConnection
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def pinned_smb_connection_cache() -> dict[str, Any]:
|
||||
"""Return the Files-owned cache; unproven process-global sessions are never reused."""
|
||||
|
||||
return {}
|
||||
|
||||
|
||||
def install_pinned_smb_transport(smbclient: Any) -> Any:
|
||||
"""Install a process-wide fail-closed smbclient session factory.
|
||||
|
||||
smbclient routes initial connections, reconnects, and DFS targets through
|
||||
``smbclient._pool.register_session``. Replacing that single factory with a
|
||||
behavior-compatible implementation ensures every target uses a socket
|
||||
selected by Core's deployment-wide outbound policy.
|
||||
"""
|
||||
|
||||
try:
|
||||
pool = import_module("smbclient._pool")
|
||||
connection_module = import_module("smbprotocol.connection")
|
||||
session_module = import_module("smbprotocol.session")
|
||||
transport_module = import_module("smbprotocol.transport")
|
||||
except ImportError as exc:
|
||||
raise SdkPeerPinningError(
|
||||
"SMB connector browsing requires the optional smbprotocol dependency"
|
||||
) from exc
|
||||
|
||||
current = getattr(pool, "register_session", None)
|
||||
if getattr(current, "__govoplan_peer_pinned__", False):
|
||||
expected_tcp = getattr(current, "__govoplan_pinned_tcp__", None)
|
||||
if expected_tcp is None or getattr(connection_module, "Tcp", None) is not expected_tcp:
|
||||
raise SdkPeerPinningError(
|
||||
"The installed SMB transport changed after peer pinning; refusing to reuse the session factory"
|
||||
)
|
||||
return smbclient
|
||||
required_parameters = {
|
||||
"server",
|
||||
"username",
|
||||
"password",
|
||||
"port",
|
||||
"encrypt",
|
||||
"connection_timeout",
|
||||
"connection_cache",
|
||||
"auth_protocol",
|
||||
"require_signing",
|
||||
}
|
||||
if current is None or not required_parameters.issubset(
|
||||
inspect.signature(current).parameters
|
||||
):
|
||||
raise SdkPeerPinningError(
|
||||
"The installed smbprotocol release does not expose the session seam required for peer pinning"
|
||||
)
|
||||
|
||||
class PinnedTcp(transport_module.Tcp):
|
||||
def connect(self) -> None:
|
||||
with self._sock_lock:
|
||||
if self.connected:
|
||||
return
|
||||
try:
|
||||
self._sock = create_outbound_connection(
|
||||
self.server,
|
||||
int(self.port),
|
||||
timeout=self.timeout,
|
||||
label="SMB connector peer",
|
||||
)
|
||||
except (OSError, OutboundHttpError, ValueError) as exc:
|
||||
raise ValueError(
|
||||
f"SMB connector peer '{self.server}:{self.port}' was rejected: {exc}"
|
||||
) from exc
|
||||
self._sock.settimeout(None)
|
||||
self.connected = True
|
||||
|
||||
# Connection.connect() instantiates the module-level Tcp symbol on every
|
||||
# reconnect. Replacing it process-wide is deliberate: an SMB connection
|
||||
# created by another code path must become stricter, never bypass Files'
|
||||
# peer boundary.
|
||||
if not hasattr(connection_module, "Tcp"):
|
||||
raise SdkPeerPinningError(
|
||||
"The installed smbprotocol release cannot install the pinned TCP transport"
|
||||
)
|
||||
connection_module.Tcp = PinnedTcp
|
||||
|
||||
def pinned_register_session(
|
||||
server: str,
|
||||
username: str | None = None,
|
||||
password: str | None = None,
|
||||
port: int = 445,
|
||||
encrypt: bool | None = None,
|
||||
connection_timeout: float = 60,
|
||||
connection_cache: dict[str, Any] | None = None,
|
||||
auth_protocol: str = "negotiate",
|
||||
require_signing: bool = True,
|
||||
) -> Any:
|
||||
cache = pinned_smb_connection_cache() if connection_cache is None else connection_cache
|
||||
connection_key = f"{server.lower()}:{port}"
|
||||
connection = cache.get(connection_key)
|
||||
transport = getattr(connection, "transport", None)
|
||||
if connection is not None and not isinstance(transport, PinnedTcp):
|
||||
disconnect = getattr(connection, "disconnect", None)
|
||||
if callable(disconnect):
|
||||
try:
|
||||
disconnect(close=True)
|
||||
except Exception:
|
||||
pass
|
||||
cache.pop(connection_key, None)
|
||||
connection = None
|
||||
if connection is None or not getattr(connection.transport, "connected", False):
|
||||
connection = connection_module.Connection(
|
||||
pool.ClientConfig().client_guid,
|
||||
server,
|
||||
port,
|
||||
require_signing=require_signing,
|
||||
)
|
||||
connection.transport = PinnedTcp(server, port)
|
||||
connection.connect(timeout=connection_timeout)
|
||||
if not isinstance(connection.transport, PinnedTcp):
|
||||
disconnect = getattr(connection, "disconnect", None)
|
||||
if callable(disconnect):
|
||||
try:
|
||||
disconnect(close=True)
|
||||
except Exception:
|
||||
pass
|
||||
raise SdkPeerPinningError(
|
||||
"smbprotocol replaced the required pinned TCP transport during connection setup"
|
||||
)
|
||||
cache[connection_key] = connection
|
||||
|
||||
session = next(
|
||||
(
|
||||
item
|
||||
for item in connection.session_table.values()
|
||||
if username is None or item.username == username
|
||||
),
|
||||
None,
|
||||
)
|
||||
if session is None:
|
||||
session = session_module.Session(
|
||||
connection,
|
||||
username=username,
|
||||
password=password,
|
||||
require_encryption=(encrypt is True),
|
||||
auth_protocol=auth_protocol,
|
||||
)
|
||||
session.connect()
|
||||
elif encrypt is not None:
|
||||
if session.encrypt_data and not encrypt:
|
||||
raise ValueError(
|
||||
"Cannot disable encryption on an already negotiated session."
|
||||
)
|
||||
if not session.encrypt_data and encrypt:
|
||||
session.encrypt = True
|
||||
return session
|
||||
|
||||
pinned_register_session.__govoplan_peer_pinned__ = True # type: ignore[attr-defined]
|
||||
pinned_register_session.__govoplan_pinned_tcp__ = PinnedTcp # type: ignore[attr-defined]
|
||||
pool.register_session = pinned_register_session
|
||||
if hasattr(smbclient, "register_session"):
|
||||
smbclient.register_session = pinned_register_session
|
||||
return smbclient
|
||||
|
||||
|
||||
__all__ = [
|
||||
"SdkPeerPinningError",
|
||||
"create_pinned_s3_client",
|
||||
"install_pinned_smb_transport",
|
||||
"pinned_smb_connection_cache",
|
||||
]
|
||||
Reference in New Issue
Block a user