Source code for daklib.rpc_peer

# SPDX-License-Identifier: GPL-2.0-or-later
# © 2026, Ansgar 🙀 <ansgar@debian.org>

"""
Peer address handling for the DAK RPC server.
"""

import contextvars
import ipaddress
import logging
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any, override

import grpc

logger = logging.getLogger(__name__)

current_peer: contextvars.ContextVar["PeerAddress | None"] = contextvars.ContextVar(
    "current_peer", default=None
)




@dataclass(frozen=True)
class PeerAddress:
    ip: ipaddress.IPv4Address | ipaddress.IPv6Address
    port: int = 0

    @override
    def __str__(self) -> str:
        if isinstance(self.ip, ipaddress.IPv6Address):
            return f"[{self.ip}]:{self.port}"
        return f"{self.ip}:{self.port}"



    @classmethod
    def from_forwarded_header(cls, value: str) -> "PeerAddress | None":
        """Parse the client address from a header like X-Forwarded-For.

        Only the last entry of a comma-separated list is used: it is the
        one appended by the trusted reverse proxy; earlier entries are
        client-controlled. Entries must be bare IP addresses; the port is
        not available (0).
        """
        addr = value.rsplit(",", 1)[-1].strip()
        try:
            return cls(ip=ipaddress.ip_address(addr))
        except ValueError:
            return None




    @classmethod
    def from_grpc_peer(cls, peer: str) -> "PeerAddress | None":
        """Parse a gRPC peer string like `ipv4:1.2.3.4:5678` or
        `ipv6:[::1]:5678`.

        Returns None for peers without an IP address (unix sockets).
        """
        if peer.startswith("ipv4:"):
            host, sep, port_str = peer.removeprefix("ipv4:").rpartition(":")
        elif peer.startswith("ipv6:"):
            host, sep, port_str = peer.removeprefix("ipv6:").rpartition(":")
            if not (host.startswith("[") and host.endswith("]")):
                return None
            host = host[1:-1]
        else:
            return None
        if not sep:
            return None
        try:
            port = int(port_str)
            if not 0 <= port <= 65535:
                return None
            return cls(ip=ipaddress.ip_address(host), port=port)
        except ValueError:
            return None






class PeerAddressInterceptor(grpc.ServerInterceptor):
    """gRPC interceptor that records the client's peer address.

    The address is taken from `peer_header` if configured, falling back to
    the connection's transport address. The header must only be configured
    when all connections arrive via a trusted reverse proxy that sets it
    (e.g. via a unix socket only the proxy can reach); otherwise clients
    can spoof their address.
    """

    def __init__(self, peer_header: str | None = None) -> None:
        # gRPC metadata keys are always lowercase
        self._peer_header = peer_header.lower() if peer_header else None



    @override
    def intercept_service(
        self,
        continuation: "Callable[[grpc.HandlerCallDetails], grpc.RpcMethodHandler[Any, Any] | None]",
        handler_call_details: grpc.HandlerCallDetails,
    ) -> "grpc.RpcMethodHandler[Any, Any] | None":
        peer: PeerAddress | None = None
        if self._peer_header is not None:
            metadata = dict(handler_call_details.invocation_metadata)
            value = metadata.get(self._peer_header)
            if isinstance(value, str):
                peer = PeerAddress.from_forwarded_header(value)
                if peer is None:
                    logger.warning(
                        "invalid peer address in %s header: %r",
                        self._peer_header,
                        value,
                    )

        current_peer.set(peer)

        handler = continuation(handler_call_details)

        if peer is not None or handler is None or handler.unary_unary is None:
            return handler

        # No trusted header: resolve the transport address, which is only
        # reachable via the ServicerContext inside the handler.
        original_fn = handler.unary_unary

        def peer_wrapper(request: Any, context: grpc.ServicerContext) -> Any:
            current_peer.set(PeerAddress.from_grpc_peer(context.peer()))
            return original_fn(request, context)

        return grpc.unary_unary_rpc_method_handler(
            peer_wrapper,
            request_deserializer=handler.request_deserializer,
            response_serializer=handler.response_serializer,
        )