Coverage for daklib/rpc_peer.py: 100%

1# SPDX-License-Identifier: GPL-2.0-or-later 

2# © 2026, Ansgar 🙀 <ansgar@debian.org> 

3 

4""" 

5Peer address handling for the DAK RPC server. 

6""" 

7 

8import contextvars 

9import ipaddress 

10import logging 

11from collections.abc import Callable 

12from dataclasses import dataclass 

13from typing import Any, override 

14 

15import grpc 

16 

17logger = logging.getLogger(__name__) 

18 

19current_peer: contextvars.ContextVar["PeerAddress | None"] = contextvars.ContextVar( 

20 "current_peer", default=None 

21) 

22 

23 

24@dataclass(frozen=True) 

25class PeerAddress: 

26 ip: ipaddress.IPv4Address | ipaddress.IPv6Address 

27 port: int = 0 

28 

29 @override 

30 def __str__(self) -> str: 

31 if isinstance(self.ip, ipaddress.IPv6Address): 

32 return f"[{self.ip}]:{self.port}" 

33 return f"{self.ip}:{self.port}" 

34 

35 @classmethod 

36 def from_forwarded_header(cls, value: str) -> "PeerAddress | None": 

37 """Parse the client address from a header like X-Forwarded-For. 

38 

39 Only the last entry of a comma-separated list is used: it is the 

40 one appended by the trusted reverse proxy; earlier entries are 

41 client-controlled. Entries must be bare IP addresses; the port is 

42 not available (0). 

43 """ 

44 addr = value.rsplit(",", 1)[-1].strip() 

45 try: 

46 return cls(ip=ipaddress.ip_address(addr)) 

47 except ValueError: 

48 return None 

49 

50 @classmethod 

51 def from_grpc_peer(cls, peer: str) -> "PeerAddress | None": 

52 """Parse a gRPC peer string like `ipv4:1.2.3.4:5678` or 

53 `ipv6:[::1]:5678`. 

54 

55 Returns None for peers without an IP address (unix sockets). 

56 """ 

57 if peer.startswith("ipv4:"): 

58 host, sep, port_str = peer.removeprefix("ipv4:").rpartition(":") 

59 elif peer.startswith("ipv6:"): 

60 host, sep, port_str = peer.removeprefix("ipv6:").rpartition(":") 

61 if not (host.startswith("[") and host.endswith("]")): 

62 return None 

63 host = host[1:-1] 

64 else: 

65 return None 

66 if not sep: 

67 return None 

68 try: 

69 port = int(port_str) 

70 if not 0 <= port <= 65535: 

71 return None 

72 return cls(ip=ipaddress.ip_address(host), port=port) 

73 except ValueError: 

74 return None 

75 

76 

77class PeerAddressInterceptor(grpc.ServerInterceptor): 

78 """gRPC interceptor that records the client's peer address. 

79 

80 The address is taken from `peer_header` if configured, falling back to 

81 the connection's transport address. The header must only be configured 

82 when all connections arrive via a trusted reverse proxy that sets it 

83 (e.g. via a unix socket only the proxy can reach); otherwise clients 

84 can spoof their address. 

85 """ 

86 

87 def __init__(self, peer_header: str | None = None) -> None: 

88 # gRPC metadata keys are always lowercase 

89 self._peer_header = peer_header.lower() if peer_header else None 

90 

91 @override 

92 def intercept_service( 

93 self, 

94 continuation: "Callable[[grpc.HandlerCallDetails], grpc.RpcMethodHandler[Any, Any] | None]", 

95 handler_call_details: grpc.HandlerCallDetails, 

96 ) -> "grpc.RpcMethodHandler[Any, Any] | None": 

97 peer: PeerAddress | None = None 

98 if self._peer_header is not None: 

99 metadata = dict(handler_call_details.invocation_metadata) 

100 value = metadata.get(self._peer_header) 

101 if isinstance(value, str): 

102 peer = PeerAddress.from_forwarded_header(value) 

103 if peer is None: 

104 logger.warning( 

105 "invalid peer address in %s header: %r", 

106 self._peer_header, 

107 value, 

108 ) 

109 

110 current_peer.set(peer) 

111 

112 handler = continuation(handler_call_details) 

113 

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

115 return handler 

116 

117 # No trusted header: resolve the transport address, which is only 

118 # reachable via the ServicerContext inside the handler. 

119 original_fn = handler.unary_unary 

120 

121 def peer_wrapper(request: Any, context: grpc.ServicerContext) -> Any: 

122 current_peer.set(PeerAddress.from_grpc_peer(context.peer())) 

123 return original_fn(request, context) 

124 

125 return grpc.unary_unary_rpc_method_handler( 

126 peer_wrapper, 

127 request_deserializer=handler.request_deserializer, 

128 response_serializer=handler.response_serializer, 

129 )