Coverage for kilter/service/runner.py: 95.62%

1# Copyright 2022-2023, 2025 Dominik Sekotill <dom.sekotill@kodo.org.uk> 

2# 

3# This Source Code Form is subject to the terms of the Mozilla Public 

4# License, v. 2.0. If a copy of the MPL was not distributed with this 

5# file, You can obtain one at http://mozilla.org/MPL/2.0/. 

6 

7""" 

8Coordinate receiving and sending raw messages with a filter and Session object 

9 

10The primary class in this module (`Runner`) is intended to be used with an 

11`anyio.abc.Listener`, which can be obtained, for instance, from 

12`anyio.create_tcp_listener()`. 

13""" 

14 

15from __future__ import annotations (empty)

16 

17import enum (empty)

18import logging (empty)

19from collections import defaultdict (empty)

20from collections.abc import Iterable (empty)

21from typing import Final (empty)

22from typing import TypeAlias (empty)

23from warnings import warn (empty)

24 

25import anyio.abc (empty)

26from async_generator import aclosing (empty)

27 

28from kilter.protocol.buffer import SimpleBuffer (empty)

29from kilter.protocol.core import EventMessage (empty)

30from kilter.protocol.core import FilterMessage (empty)

31from kilter.protocol.core import FilterProtocol (empty)

32from kilter.protocol.core import ResponseMessage (empty)

33from kilter.protocol.messages import * (empty)

34 

35from .options import get_flags (empty)

36from .options import get_macros (empty)

37from .session import Aborted (empty)

38from .session import Filter (empty)

39from .session import FilterResponse (empty)

40from .session import Session (empty)

41from .util import Broadcast (empty)

42from .util import qualname (empty)

43 

44__all__ = [ (empty)

45 "Runner", 

46 "NegotiationError", 

47] 

48 

49FinalResponse: TypeAlias = FilterResponse | TemporaryFailure (empty)

50 

51kiB: Final = 2**10 (empty)

52MiB: Final = 2**20 (empty)

53 

54_logger = logging.getLogger(__package__) (empty)

55 

56 

57class NegotiationError(Exception): (empty)

58 """ 

59 An error raised when MTAs are not compatible with the filter 

60 """ 

61 

62 

63class State(enum.Enum): (empty)

64 

65 CONNECTED = enum.auto() (empty)

66 SESSION = enum.auto() (empty)

67 SESSION_ABORTED = enum.auto() (empty)

68 MESSAGE = enum.auto() (empty)

69 MESSAGE_ABORTED = enum.auto() (empty)

70 

71 

72class _Broadcast(Broadcast[EventMessage]): (empty)

73 

74 def __init__(self) -> None: (empty)

75 super().__init__() (empty)

76 self.task_status = list[anyio.abc.TaskStatus[None]]() (empty)

77 

78 async def shutdown_hook(self) -> None: (empty)

79 await self.pre_receive_hook() (empty)

80 

81 async def pre_receive_hook(self) -> None: (empty)

82 while self.task_status: (empty)

83 self.task_status.pop().started() (empty)

84 

85 

86class Sender: (empty)

87 """ 

88 Concrete implementation of `kilter.service.session.Sender` 

89 """ 

90 

91 def __init__(self, client: anyio.abc.ByteSendStream, proto: FilterProtocol) -> None: (empty)

92 self.client = client (empty)

93 self.proto = proto (empty)

94 

95 async def send(self, message: FilterMessage) -> None: (empty)

96 """ 

97 Encode and send a message to the client stream 

98 """ 

99 buffer = SimpleBuffer(1*kiB) (empty)

100 self.proto.write_to(buffer, message) (empty)

101 await self.client.send(buffer[:]) (empty)

102 if __debug__: (empty)

103 _logger.debug(f"sent: {message}") (empty)

104 

105 

106class Runner: (empty)

107 """ 

108 A filter runner that coordinates passing data between a stream and multiple filters 

109 

110 Instances can be used as handlers that can be passed to `anyio.abc.Listener.serve()` or 

111 used with any `anyio.abc.ByteStream`. 

112 """ 

113 

114 def __init__(self, *filters: Filter): (empty)

115 if len(filters) == 0: # pragma: no-cover (empty)

116 raise TypeError("Runner requires at least one filter to run") 

117 self.filters = set(filters) (empty)

118 if len(filters) != len(self.filters): (empty)

119 warn("Repeated filters will only be run once", stacklevel=2) (empty)

120 self.use_skip = True (empty)

121 

122 async def __call__(self, client: anyio.abc.ByteStream) -> None: (empty)

123 """ 

124 Return an awaitable that starts and coordinates filters 

125 """ 

126 buff = SimpleBuffer(1*MiB) (empty)

127 proto = FilterProtocol(abort_on_unknown=True) (empty)

128 sender = Sender(client, proto) (empty)

129 session = Session(sender, _Broadcast()) (empty)

130 runner = SessionRunner(session) (empty)

131 state = State.CONNECTED (empty)

132 

133 async with ( (empty)

134 aclosing(client), 

135 anyio.create_task_group() as tasks, 

136 ): 

137 while 1: (empty)

138 try: (empty)

139 buff[:] = await client.receive(buff.available) (empty)

140 except ( (empty)

141 anyio.EndOfStream, 

142 anyio.ClosedResourceError, 

143 anyio.BrokenResourceError, 

144 ): 

145 return (empty)

146 for message in proto.read_from(buff): (empty)

147 if __debug__: (empty)

148 _logger.debug(f"received: {message}") (empty)

149 

150 # If previous message was Abort, restart filters for any non-Abort/Close 

151 # message 

152 if state in (State.SESSION_ABORTED, State.MESSAGE_ABORTED): (empty)

153 if not isinstance(message, Abort|Close): (empty)

154 await runner.start(self.filters, tasks) (empty)

155 state = ( (empty)

156 State.CONNECTED if state == State.SESSION_ABORTED else 

157 State.SESSION 

158 ) 

159 

160 match message: (empty)

161 case Negotiate(): (empty)

162 await sender.send(await self._negotiate(message)) (empty)

163 continue (empty)

164 case Connect(): (empty)

165 _logger.info(f"Client connected from {message.hostname}") (empty)

166 await session.deliver(message) (empty)

167 await runner.start(self.filters, tasks) (empty)

168 if proto.needs_response(message): 168 ↛ 170line 168 didn't jump to line 170 because the condition on line 168 was always true(empty)

169 await sender.send(await runner.check_response() or Continue()) (empty)

170 continue (empty)

171 case Helo(): (empty)

172 state = State.SESSION (empty)

173 case EnvelopeFrom(): (empty)

174 state = State.MESSAGE (empty)

175 case Abort() if state in (State.SESSION, State.MESSAGE): (empty)

176 state = ( (empty)

177 State.SESSION_ABORTED if state == State.SESSION else 

178 State.MESSAGE_ABORTED 

179 ) 

180 case Abort(): (empty)

181 _logger.warning("Unexpected Abort received") (empty)

182 state = State.CONNECTED (empty)

183 case Close(): (empty)

184 tasks.cancel_scope.cancel() (empty)

185 return (empty)

186 

187 skip_or_cont = await session.deliver(message) (empty)

188 if not proto.needs_response(message): (empty)

189 continue (empty)

190 if (resp := await runner.check_response()): (empty)

191 await sender.send(resp) (empty)

192 elif self.use_skip: (empty)

193 await sender.send(skip_or_cont()) (empty)

194 else: 

195 await sender.send(Continue()) (empty)

196 

197 async def _negotiate(self, message: Negotiate) -> Negotiate: (empty)

198 _logger.info("Negotiating with MTA") (empty)

199 

200 optmask = ProtocolFlags.NONE (empty)

201 options = \ (empty)

202 ProtocolFlags.SKIP | \ 

203 ProtocolFlags.NO_HELO | \ 

204 ProtocolFlags.NO_SENDER | ProtocolFlags.NO_RECIPIENT | \ 

205 ProtocolFlags.NO_DATA | ProtocolFlags.NO_BODY | \ 

206 ProtocolFlags.NO_HEADERS | ProtocolFlags.NO_END_OF_HEADERS | \ 

207 ProtocolFlags.NR_CONNECT | ProtocolFlags.NR_HELO | \ 

208 ProtocolFlags.NR_SENDER | ProtocolFlags.NR_RECIPIENT | \ 

209 ProtocolFlags.NR_DATA | ProtocolFlags.NR_BODY | \ 

210 ProtocolFlags.NR_HEADER | ProtocolFlags.NR_END_OF_HEADERS 

211 actions = ActionFlags.NONE (empty)

212 macros = defaultdict(set) (empty)

213 

214 options &= message.protocol_flags # Remove unoffered initial flags, they are not required (empty)

215 

216 for filtr in self.filters: (empty)

217 flags = get_flags(filtr) (empty)

218 optmask |= flags.unset_options (empty)

219 options |= flags.set_options (empty)

220 actions |= flags.set_actions (empty)

221 

222 for stage, names in get_macros(filtr).items(): 222 ↛ 223line 222 didn't jump to line 223 because the loop on line 222 never started(empty)

223 macros[stage].update(names) 

224 

225 options &= ~optmask (empty)

226 

227 if (missing_actions := actions & ~message.action_flags): (empty)

228 raise NegotiationError(f"MTA does not accept {missing_actions}") (empty)

229 

230 if (missing_options := options & ~message.protocol_flags): (empty)

231 raise NegotiationError(f"MTA does not offer {missing_options}") (empty)

232 

233 self.use_skip = ProtocolFlags.SKIP in options (empty)

234 

235 return Negotiate(6, actions, options, dict(macros)) (empty)

236 

237 

238class SessionRunner: (empty)

239 

240 def __init__(self, session: Session): (empty)

241 self.session = session (empty)

242 self.filters = dict[Filter, FinalResponse|None]() (empty)

243 

244 async def start(self, filters: Iterable[Filter], task_group: anyio.abc.TaskGroup) -> None: (empty)

245 """ 

246 Run all the given filters in a task group 

247 

248 The session MUST have been primed by the delivery of a Connect message beforehand or 

249 filters will be unable to access the connection details. 

250 """ 

251 _logger.debug("Starting filters") (empty)

252 for flter in filters: (empty)

253 await task_group.start(self.run_filter, flter) (empty)

254 

255 async def run_filter( (empty)

256 self, 

257 flter: Filter, 

258 task_status: anyio.abc.TaskStatus[None], 

259 ) -> None: 

260 """ 

261 Run a filter as a subtask in a task group 

262 

263 A `Future` for returning the filter's response is added to the 

264 `SessionRunner.filter` dict. 

265 """ 

266 if flter in self.filters: 266 ↛ 267line 266 didn't jump to line 267 because the condition on line 266 was never true(empty)

267 raise RuntimeError 

268 self.filters[flter] = None (empty)

269 

270 async with self.session: (empty)

271 assert isinstance(self.session.broadcast, _Broadcast) (empty)

272 status_notifiers = self.session.broadcast.task_status (empty)

273 status_notifiers.append(task_status) (empty)

274 

275 try: (empty)

276 resp: FinalResponse = await flter(self.session) (empty)

277 except Aborted: (empty)

278 _logger.debug(f"Aborted filter {qualname(flter)}") (empty)

279 del self.filters[flter] (empty)

280 return (empty)

281 except Exception: 

282 _logger.exception(f"Error in filter {qualname(flter)}") 

283 resp = TemporaryFailure() 

284 if not isinstance(resp, FinalResponse): (empty)

285 warn(f"expected a valid response from {qualname(flter)}, got {resp}") # type: ignore # Don't fully trust users… (empty)

286 resp = TemporaryFailure() (empty)

287 self.filters[flter] = resp (empty)

288 if task_status in status_notifiers: (empty)

289 status_notifiers.remove(task_status) (empty)

290 task_status.started() (empty)

291 

292 async def check_response(self) -> ResponseMessage|None: (empty)

293 assert self.filters, "no filters when checking for a response" (empty)

294 response: ResponseMessage|None = None (empty)

295 complete = list[Filter]() (empty)

296 for flter, result in self.filters.items(): (empty)

297 # If a filter has not finished or no response is expected, continue without 

298 # removing from filter container; remove failed filters and filters that have 

299 # accepted; return a response for rejections; 

300 match result: (empty)

301 case None: (empty)

302 continue (empty)

303 case Accept(): (empty)

304 _logger.info("Accept from %s, waiting for remaining", qualname(flter)) (empty)

305 case TemporaryFailure() as response: (empty)

306 _logger.warning("Filter failed: %s", flter) (empty)

307 case Reject()|Discard()|ReplyCode() as response: 307 ↛ 311line 307 didn't jump to line 311 because the pattern on line 307 always matched(empty)

308 _logger.info("Returning response %s from %s", type(response).__name__, qualname(flter)) (empty)

309 complete[:] = self.filters (empty)

310 break (empty)

311 case msg: 

312 raise AssertionError(f"unexpected filter result: {msg}") 

313 complete.append(flter) (empty)

314 for flter in complete: (empty)

315 del self.filters[flter] (empty)

316 return response if response else None if self.filters else Accept() (empty)