test_v2.py 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929
  1. # Copyright 2018 New Vector Ltd
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import itertools
  15. from typing import (
  16. Collection,
  17. Dict,
  18. Iterable,
  19. List,
  20. Mapping,
  21. Optional,
  22. Set,
  23. Tuple,
  24. TypeVar,
  25. )
  26. import attr
  27. from twisted.internet import defer
  28. from synapse.api.constants import EventTypes, JoinRules, Membership
  29. from synapse.api.room_versions import RoomVersions
  30. from synapse.event_auth import auth_types_for_event
  31. from synapse.events import EventBase, make_event_from_dict
  32. from synapse.state.v2 import (
  33. _get_auth_chain_difference,
  34. lexicographical_topological_sort,
  35. resolve_events_with_store,
  36. )
  37. from synapse.types import EventID, StateMap
  38. from tests import unittest
  39. ALICE = "@alice:example.com"
  40. BOB = "@bob:example.com"
  41. CHARLIE = "@charlie:example.com"
  42. EVELYN = "@evelyn:example.com"
  43. ZARA = "@zara:example.com"
  44. ROOM_ID = "!test:example.com"
  45. MEMBERSHIP_CONTENT_JOIN = {"membership": Membership.JOIN}
  46. MEMBERSHIP_CONTENT_BAN = {"membership": Membership.BAN}
  47. ORIGIN_SERVER_TS = 0
  48. class FakeClock:
  49. def sleep(self, msec: float) -> "defer.Deferred[None]":
  50. return defer.succeed(None)
  51. class FakeEvent:
  52. """A fake event we use as a convenience.
  53. NOTE: Again as a convenience we use "node_ids" rather than event_ids to
  54. refer to events. The event_id has node_id as localpart and example.com
  55. as domain.
  56. """
  57. def __init__(
  58. self,
  59. id: str,
  60. sender: str,
  61. type: str,
  62. state_key: Optional[str],
  63. content: Mapping[str, object],
  64. ):
  65. self.node_id = id
  66. self.event_id = EventID(id, "example.com").to_string()
  67. self.sender = sender
  68. self.type = type
  69. self.state_key = state_key
  70. self.content = content
  71. self.room_id = ROOM_ID
  72. def to_event(self, auth_events: List[str], prev_events: List[str]) -> EventBase:
  73. """Given the auth_events and prev_events, convert to a Frozen Event
  74. Args:
  75. auth_events: list of event_ids
  76. prev_events: list of event_ids
  77. Returns:
  78. FrozenEvent
  79. """
  80. global ORIGIN_SERVER_TS
  81. ts = ORIGIN_SERVER_TS
  82. ORIGIN_SERVER_TS = ORIGIN_SERVER_TS + 1
  83. event_dict = {
  84. "auth_events": [(a, {}) for a in auth_events],
  85. "prev_events": [(p, {}) for p in prev_events],
  86. "event_id": self.event_id,
  87. "sender": self.sender,
  88. "type": self.type,
  89. "content": self.content,
  90. "origin_server_ts": ts,
  91. "room_id": ROOM_ID,
  92. }
  93. if self.state_key is not None:
  94. event_dict["state_key"] = self.state_key
  95. return make_event_from_dict(event_dict)
  96. # All graphs start with this set of events
  97. INITIAL_EVENTS = [
  98. FakeEvent(
  99. id="CREATE",
  100. sender=ALICE,
  101. type=EventTypes.Create,
  102. state_key="",
  103. content={"creator": ALICE},
  104. ),
  105. FakeEvent(
  106. id="IMA",
  107. sender=ALICE,
  108. type=EventTypes.Member,
  109. state_key=ALICE,
  110. content=MEMBERSHIP_CONTENT_JOIN,
  111. ),
  112. FakeEvent(
  113. id="IPOWER",
  114. sender=ALICE,
  115. type=EventTypes.PowerLevels,
  116. state_key="",
  117. content={"users": {ALICE: 100}},
  118. ),
  119. FakeEvent(
  120. id="IJR",
  121. sender=ALICE,
  122. type=EventTypes.JoinRules,
  123. state_key="",
  124. content={"join_rule": JoinRules.PUBLIC},
  125. ),
  126. FakeEvent(
  127. id="IMB",
  128. sender=BOB,
  129. type=EventTypes.Member,
  130. state_key=BOB,
  131. content=MEMBERSHIP_CONTENT_JOIN,
  132. ),
  133. FakeEvent(
  134. id="IMC",
  135. sender=CHARLIE,
  136. type=EventTypes.Member,
  137. state_key=CHARLIE,
  138. content=MEMBERSHIP_CONTENT_JOIN,
  139. ),
  140. FakeEvent(
  141. id="IMZ",
  142. sender=ZARA,
  143. type=EventTypes.Member,
  144. state_key=ZARA,
  145. content=MEMBERSHIP_CONTENT_JOIN,
  146. ),
  147. FakeEvent(
  148. id="START", sender=ZARA, type=EventTypes.Message, state_key=None, content={}
  149. ),
  150. FakeEvent(
  151. id="END", sender=ZARA, type=EventTypes.Message, state_key=None, content={}
  152. ),
  153. ]
  154. INITIAL_EDGES = ["START", "IMZ", "IMC", "IMB", "IJR", "IPOWER", "IMA", "CREATE"]
  155. class StateTestCase(unittest.TestCase):
  156. def test_ban_vs_pl(self) -> None:
  157. events = [
  158. FakeEvent(
  159. id="PA",
  160. sender=ALICE,
  161. type=EventTypes.PowerLevels,
  162. state_key="",
  163. content={"users": {ALICE: 100, BOB: 50}},
  164. ),
  165. FakeEvent(
  166. id="MA",
  167. sender=ALICE,
  168. type=EventTypes.Member,
  169. state_key=ALICE,
  170. content={"membership": Membership.JOIN},
  171. ),
  172. FakeEvent(
  173. id="MB",
  174. sender=ALICE,
  175. type=EventTypes.Member,
  176. state_key=BOB,
  177. content={"membership": Membership.BAN},
  178. ),
  179. FakeEvent(
  180. id="PB",
  181. sender=BOB,
  182. type=EventTypes.PowerLevels,
  183. state_key="",
  184. content={"users": {ALICE: 100, BOB: 50}},
  185. ),
  186. ]
  187. edges = [["END", "MB", "MA", "PA", "START"], ["END", "PB", "PA"]]
  188. expected_state_ids = ["PA", "MA", "MB"]
  189. self.do_check(events, edges, expected_state_ids)
  190. def test_join_rule_evasion(self) -> None:
  191. events = [
  192. FakeEvent(
  193. id="JR",
  194. sender=ALICE,
  195. type=EventTypes.JoinRules,
  196. state_key="",
  197. content={"join_rules": JoinRules.PRIVATE},
  198. ),
  199. FakeEvent(
  200. id="ME",
  201. sender=EVELYN,
  202. type=EventTypes.Member,
  203. state_key=EVELYN,
  204. content={"membership": Membership.JOIN},
  205. ),
  206. ]
  207. edges = [["END", "JR", "START"], ["END", "ME", "START"]]
  208. expected_state_ids = ["JR"]
  209. self.do_check(events, edges, expected_state_ids)
  210. def test_offtopic_pl(self) -> None:
  211. events = [
  212. FakeEvent(
  213. id="PA",
  214. sender=ALICE,
  215. type=EventTypes.PowerLevels,
  216. state_key="",
  217. content={"users": {ALICE: 100, BOB: 50}},
  218. ),
  219. FakeEvent(
  220. id="PB",
  221. sender=BOB,
  222. type=EventTypes.PowerLevels,
  223. state_key="",
  224. content={"users": {ALICE: 100, BOB: 50, CHARLIE: 50}},
  225. ),
  226. FakeEvent(
  227. id="PC",
  228. sender=CHARLIE,
  229. type=EventTypes.PowerLevels,
  230. state_key="",
  231. content={"users": {ALICE: 100, BOB: 50, CHARLIE: 0}},
  232. ),
  233. ]
  234. edges = [["END", "PC", "PB", "PA", "START"], ["END", "PA"]]
  235. expected_state_ids = ["PC"]
  236. self.do_check(events, edges, expected_state_ids)
  237. def test_topic_basic(self) -> None:
  238. events = [
  239. FakeEvent(
  240. id="T1", sender=ALICE, type=EventTypes.Topic, state_key="", content={}
  241. ),
  242. FakeEvent(
  243. id="PA1",
  244. sender=ALICE,
  245. type=EventTypes.PowerLevels,
  246. state_key="",
  247. content={"users": {ALICE: 100, BOB: 50}},
  248. ),
  249. FakeEvent(
  250. id="T2", sender=ALICE, type=EventTypes.Topic, state_key="", content={}
  251. ),
  252. FakeEvent(
  253. id="PA2",
  254. sender=ALICE,
  255. type=EventTypes.PowerLevels,
  256. state_key="",
  257. content={"users": {ALICE: 100, BOB: 0}},
  258. ),
  259. FakeEvent(
  260. id="PB",
  261. sender=BOB,
  262. type=EventTypes.PowerLevels,
  263. state_key="",
  264. content={"users": {ALICE: 100, BOB: 50}},
  265. ),
  266. FakeEvent(
  267. id="T3", sender=BOB, type=EventTypes.Topic, state_key="", content={}
  268. ),
  269. ]
  270. edges = [["END", "PA2", "T2", "PA1", "T1", "START"], ["END", "T3", "PB", "PA1"]]
  271. expected_state_ids = ["PA2", "T2"]
  272. self.do_check(events, edges, expected_state_ids)
  273. def test_topic_reset(self) -> None:
  274. events = [
  275. FakeEvent(
  276. id="T1", sender=ALICE, type=EventTypes.Topic, state_key="", content={}
  277. ),
  278. FakeEvent(
  279. id="PA",
  280. sender=ALICE,
  281. type=EventTypes.PowerLevels,
  282. state_key="",
  283. content={"users": {ALICE: 100, BOB: 50}},
  284. ),
  285. FakeEvent(
  286. id="T2", sender=BOB, type=EventTypes.Topic, state_key="", content={}
  287. ),
  288. FakeEvent(
  289. id="MB",
  290. sender=ALICE,
  291. type=EventTypes.Member,
  292. state_key=BOB,
  293. content={"membership": Membership.BAN},
  294. ),
  295. ]
  296. edges = [["END", "MB", "T2", "PA", "T1", "START"], ["END", "T1"]]
  297. expected_state_ids = ["T1", "MB", "PA"]
  298. self.do_check(events, edges, expected_state_ids)
  299. def test_topic(self) -> None:
  300. events = [
  301. FakeEvent(
  302. id="T1", sender=ALICE, type=EventTypes.Topic, state_key="", content={}
  303. ),
  304. FakeEvent(
  305. id="PA1",
  306. sender=ALICE,
  307. type=EventTypes.PowerLevels,
  308. state_key="",
  309. content={"users": {ALICE: 100, BOB: 50}},
  310. ),
  311. FakeEvent(
  312. id="T2", sender=ALICE, type=EventTypes.Topic, state_key="", content={}
  313. ),
  314. FakeEvent(
  315. id="PA2",
  316. sender=ALICE,
  317. type=EventTypes.PowerLevels,
  318. state_key="",
  319. content={"users": {ALICE: 100, BOB: 0}},
  320. ),
  321. FakeEvent(
  322. id="PB",
  323. sender=BOB,
  324. type=EventTypes.PowerLevels,
  325. state_key="",
  326. content={"users": {ALICE: 100, BOB: 50}},
  327. ),
  328. FakeEvent(
  329. id="T3", sender=BOB, type=EventTypes.Topic, state_key="", content={}
  330. ),
  331. FakeEvent(
  332. id="MZ1",
  333. sender=ZARA,
  334. type=EventTypes.Message,
  335. state_key=None,
  336. content={},
  337. ),
  338. FakeEvent(
  339. id="T4", sender=ALICE, type=EventTypes.Topic, state_key="", content={}
  340. ),
  341. ]
  342. edges = [
  343. ["END", "T4", "MZ1", "PA2", "T2", "PA1", "T1", "START"],
  344. ["END", "MZ1", "T3", "PB", "PA1"],
  345. ]
  346. expected_state_ids = ["T4", "PA2"]
  347. self.do_check(events, edges, expected_state_ids)
  348. def test_mainline_sort(self) -> None:
  349. """Tests that the mainline ordering works correctly."""
  350. events = [
  351. FakeEvent(
  352. id="T1", sender=ALICE, type=EventTypes.Topic, state_key="", content={}
  353. ),
  354. FakeEvent(
  355. id="PA1",
  356. sender=ALICE,
  357. type=EventTypes.PowerLevels,
  358. state_key="",
  359. content={"users": {ALICE: 100, BOB: 50}},
  360. ),
  361. FakeEvent(
  362. id="T2", sender=ALICE, type=EventTypes.Topic, state_key="", content={}
  363. ),
  364. FakeEvent(
  365. id="PA2",
  366. sender=ALICE,
  367. type=EventTypes.PowerLevels,
  368. state_key="",
  369. content={
  370. "users": {ALICE: 100, BOB: 50},
  371. "events": {EventTypes.PowerLevels: 100},
  372. },
  373. ),
  374. FakeEvent(
  375. id="PB",
  376. sender=BOB,
  377. type=EventTypes.PowerLevels,
  378. state_key="",
  379. content={"users": {ALICE: 100, BOB: 50}},
  380. ),
  381. FakeEvent(
  382. id="T3", sender=BOB, type=EventTypes.Topic, state_key="", content={}
  383. ),
  384. FakeEvent(
  385. id="T4", sender=ALICE, type=EventTypes.Topic, state_key="", content={}
  386. ),
  387. ]
  388. edges = [
  389. ["END", "T3", "PA2", "T2", "PA1", "T1", "START"],
  390. ["END", "T4", "PB", "PA1"],
  391. ]
  392. # We expect T3 to be picked as the other topics are pointing at older
  393. # power levels. Note that without mainline ordering we'd pick T4 due to
  394. # it being sent *after* T3.
  395. expected_state_ids = ["T3", "PA2"]
  396. self.do_check(events, edges, expected_state_ids)
  397. def do_check(
  398. self,
  399. events: List[FakeEvent],
  400. edges: List[List[str]],
  401. expected_state_ids: List[str],
  402. ) -> None:
  403. """Take a list of events and edges and calculate the state of the
  404. graph at END, and asserts it matches `expected_state_ids`
  405. Args:
  406. events
  407. edges: A list of chains of event edges, e.g.
  408. `[[A, B, C]]` are edges A->B and B->C.
  409. expected_state_ids: The expected state at END, (excluding
  410. the keys that haven't changed since START).
  411. """
  412. # We want to sort the events into topological order for processing.
  413. graph: Dict[str, Set[str]] = {}
  414. fake_event_map: Dict[str, FakeEvent] = {}
  415. for ev in itertools.chain(INITIAL_EVENTS, events):
  416. graph[ev.node_id] = set()
  417. fake_event_map[ev.node_id] = ev
  418. for a, b in pairwise(INITIAL_EDGES):
  419. graph[a].add(b)
  420. for edge_list in edges:
  421. for a, b in pairwise(edge_list):
  422. graph[a].add(b)
  423. event_map: Dict[str, EventBase] = {}
  424. state_at_event: Dict[str, StateMap[str]] = {}
  425. # We copy the map as the sort consumes the graph
  426. graph_copy = {k: set(v) for k, v in graph.items()}
  427. for node_id in lexicographical_topological_sort(graph_copy, key=lambda e: e):
  428. fake_event = fake_event_map[node_id]
  429. event_id = fake_event.event_id
  430. prev_events = list(graph[node_id])
  431. state_before: StateMap[str]
  432. if len(prev_events) == 0:
  433. state_before = {}
  434. elif len(prev_events) == 1:
  435. state_before = dict(state_at_event[prev_events[0]])
  436. else:
  437. state_d = resolve_events_with_store(
  438. FakeClock(),
  439. ROOM_ID,
  440. RoomVersions.V2,
  441. [state_at_event[n] for n in prev_events],
  442. event_map=event_map,
  443. state_res_store=TestStateResolutionStore(event_map),
  444. )
  445. state_before = self.successResultOf(defer.ensureDeferred(state_d))
  446. state_after = dict(state_before)
  447. if fake_event.state_key is not None:
  448. state_after[(fake_event.type, fake_event.state_key)] = event_id
  449. # This type ignore is a bit sad. Things we have tried:
  450. # 1. Define a `GenericEvent` Protocol satisfied by FakeEvent, EventBase and
  451. # EventBuilder. But this is Hard because the relevant attributes are
  452. # DictProperty[T] descriptors on EventBase but normal Ts on FakeEvent.
  453. # 2. Define a `GenericEvent` Protocol describing `FakeEvent` only, and
  454. # change this function to accept Union[Event, EventBase, EventBuilder].
  455. # This seems reasonable to me, but mypy isn't happy. I think that's
  456. # a mypy bug, see https://github.com/python/mypy/issues/5570
  457. # Instead, resort to a type-ignore.
  458. auth_types = set(auth_types_for_event(RoomVersions.V6, fake_event)) # type: ignore[arg-type]
  459. auth_events = []
  460. for key in auth_types:
  461. if key in state_before:
  462. auth_events.append(state_before[key])
  463. event = fake_event.to_event(auth_events, prev_events)
  464. state_at_event[node_id] = state_after
  465. event_map[event_id] = event
  466. expected_state = {}
  467. for node_id in expected_state_ids:
  468. # expected_state_ids are node IDs rather than event IDs,
  469. # so we have to convert
  470. event_id = EventID(node_id, "example.com").to_string()
  471. event = event_map[event_id]
  472. key = (event.type, event.state_key)
  473. expected_state[key] = event_id
  474. start_state = state_at_event["START"]
  475. end_state = {
  476. key: value
  477. for key, value in state_at_event["END"].items()
  478. if key in expected_state or start_state.get(key) != value
  479. }
  480. self.assertEqual(expected_state, end_state)
  481. class LexicographicalTestCase(unittest.TestCase):
  482. def test_simple(self) -> None:
  483. graph: Dict[str, Set[str]] = {
  484. "l": {"o"},
  485. "m": {"n", "o"},
  486. "n": {"o"},
  487. "o": set(),
  488. "p": {"o"},
  489. }
  490. res = list(lexicographical_topological_sort(graph, key=lambda x: x))
  491. self.assertEqual(["o", "l", "n", "m", "p"], res)
  492. class SimpleParamStateTestCase(unittest.TestCase):
  493. def setUp(self) -> None:
  494. # We build up a simple DAG.
  495. event_map = {}
  496. create_event = FakeEvent(
  497. id="CREATE",
  498. sender=ALICE,
  499. type=EventTypes.Create,
  500. state_key="",
  501. content={"creator": ALICE},
  502. ).to_event([], [])
  503. event_map[create_event.event_id] = create_event
  504. alice_member = FakeEvent(
  505. id="IMA",
  506. sender=ALICE,
  507. type=EventTypes.Member,
  508. state_key=ALICE,
  509. content=MEMBERSHIP_CONTENT_JOIN,
  510. ).to_event([create_event.event_id], [create_event.event_id])
  511. event_map[alice_member.event_id] = alice_member
  512. join_rules = FakeEvent(
  513. id="IJR",
  514. sender=ALICE,
  515. type=EventTypes.JoinRules,
  516. state_key="",
  517. content={"join_rule": JoinRules.PUBLIC},
  518. ).to_event(
  519. auth_events=[create_event.event_id, alice_member.event_id],
  520. prev_events=[alice_member.event_id],
  521. )
  522. event_map[join_rules.event_id] = join_rules
  523. # Bob and Charlie join at the same time, so there is a fork
  524. bob_member = FakeEvent(
  525. id="IMB",
  526. sender=BOB,
  527. type=EventTypes.Member,
  528. state_key=BOB,
  529. content=MEMBERSHIP_CONTENT_JOIN,
  530. ).to_event(
  531. auth_events=[create_event.event_id, join_rules.event_id],
  532. prev_events=[join_rules.event_id],
  533. )
  534. event_map[bob_member.event_id] = bob_member
  535. charlie_member = FakeEvent(
  536. id="IMC",
  537. sender=CHARLIE,
  538. type=EventTypes.Member,
  539. state_key=CHARLIE,
  540. content=MEMBERSHIP_CONTENT_JOIN,
  541. ).to_event(
  542. auth_events=[create_event.event_id, join_rules.event_id],
  543. prev_events=[join_rules.event_id],
  544. )
  545. event_map[charlie_member.event_id] = charlie_member
  546. self.event_map = event_map
  547. self.create_event = create_event
  548. self.alice_member = alice_member
  549. self.join_rules = join_rules
  550. self.bob_member = bob_member
  551. self.charlie_member = charlie_member
  552. self.state_at_bob = {
  553. (e.type, e.state_key): e.event_id
  554. for e in [create_event, alice_member, join_rules, bob_member]
  555. }
  556. self.state_at_charlie = {
  557. (e.type, e.state_key): e.event_id
  558. for e in [create_event, alice_member, join_rules, charlie_member]
  559. }
  560. self.expected_combined_state = {
  561. (e.type, e.state_key): e.event_id
  562. for e in [
  563. create_event,
  564. alice_member,
  565. join_rules,
  566. bob_member,
  567. charlie_member,
  568. ]
  569. }
  570. def test_event_map_none(self) -> None:
  571. # Test that we correctly handle passing `None` as the event_map
  572. state_d = resolve_events_with_store(
  573. FakeClock(),
  574. ROOM_ID,
  575. RoomVersions.V2,
  576. [self.state_at_bob, self.state_at_charlie],
  577. event_map=None,
  578. state_res_store=TestStateResolutionStore(self.event_map),
  579. )
  580. state = self.successResultOf(defer.ensureDeferred(state_d))
  581. self.assert_dict(self.expected_combined_state, state)
  582. class AuthChainDifferenceTestCase(unittest.TestCase):
  583. """We test that `_get_auth_chain_difference` correctly handles unpersisted
  584. events.
  585. """
  586. def test_simple(self) -> None:
  587. # Test getting the auth difference for a simple chain with a single
  588. # unpersisted event:
  589. #
  590. # Unpersisted | Persisted
  591. # |
  592. # C -|-> B -> A
  593. a = FakeEvent(
  594. id="A",
  595. sender=ALICE,
  596. type=EventTypes.Member,
  597. state_key="",
  598. content={},
  599. ).to_event([], [])
  600. b = FakeEvent(
  601. id="B",
  602. sender=ALICE,
  603. type=EventTypes.Member,
  604. state_key="",
  605. content={},
  606. ).to_event([a.event_id], [])
  607. c = FakeEvent(
  608. id="C",
  609. sender=ALICE,
  610. type=EventTypes.Member,
  611. state_key="",
  612. content={},
  613. ).to_event([b.event_id], [])
  614. persisted_events = {a.event_id: a, b.event_id: b}
  615. unpersited_events = {c.event_id: c}
  616. state_sets = [{"a": a.event_id, "b": b.event_id}, {"c": c.event_id}]
  617. store = TestStateResolutionStore(persisted_events)
  618. diff_d = _get_auth_chain_difference(
  619. ROOM_ID, state_sets, unpersited_events, store
  620. )
  621. difference = self.successResultOf(defer.ensureDeferred(diff_d))
  622. self.assertEqual(difference, {c.event_id})
  623. def test_multiple_unpersisted_chain(self) -> None:
  624. # Test getting the auth difference for a simple chain with multiple
  625. # unpersisted events:
  626. #
  627. # Unpersisted | Persisted
  628. # |
  629. # D -> C -|-> B -> A
  630. a = FakeEvent(
  631. id="A",
  632. sender=ALICE,
  633. type=EventTypes.Member,
  634. state_key="",
  635. content={},
  636. ).to_event([], [])
  637. b = FakeEvent(
  638. id="B",
  639. sender=ALICE,
  640. type=EventTypes.Member,
  641. state_key="",
  642. content={},
  643. ).to_event([a.event_id], [])
  644. c = FakeEvent(
  645. id="C",
  646. sender=ALICE,
  647. type=EventTypes.Member,
  648. state_key="",
  649. content={},
  650. ).to_event([b.event_id], [])
  651. d = FakeEvent(
  652. id="D",
  653. sender=ALICE,
  654. type=EventTypes.Member,
  655. state_key="",
  656. content={},
  657. ).to_event([c.event_id], [])
  658. persisted_events = {a.event_id: a, b.event_id: b}
  659. unpersited_events = {c.event_id: c, d.event_id: d}
  660. state_sets = [
  661. {"a": a.event_id, "b": b.event_id},
  662. {"c": c.event_id, "d": d.event_id},
  663. ]
  664. store = TestStateResolutionStore(persisted_events)
  665. diff_d = _get_auth_chain_difference(
  666. ROOM_ID, state_sets, unpersited_events, store
  667. )
  668. difference = self.successResultOf(defer.ensureDeferred(diff_d))
  669. self.assertEqual(difference, {d.event_id, c.event_id})
  670. def test_unpersisted_events_different_sets(self) -> None:
  671. # Test getting the auth difference for with multiple unpersisted events
  672. # in different branches:
  673. #
  674. # Unpersisted | Persisted
  675. # |
  676. # D --> C -|-> B -> A
  677. # E ----^ -|---^
  678. # |
  679. a = FakeEvent(
  680. id="A",
  681. sender=ALICE,
  682. type=EventTypes.Member,
  683. state_key="",
  684. content={},
  685. ).to_event([], [])
  686. b = FakeEvent(
  687. id="B",
  688. sender=ALICE,
  689. type=EventTypes.Member,
  690. state_key="",
  691. content={},
  692. ).to_event([a.event_id], [])
  693. c = FakeEvent(
  694. id="C",
  695. sender=ALICE,
  696. type=EventTypes.Member,
  697. state_key="",
  698. content={},
  699. ).to_event([b.event_id], [])
  700. d = FakeEvent(
  701. id="D",
  702. sender=ALICE,
  703. type=EventTypes.Member,
  704. state_key="",
  705. content={},
  706. ).to_event([c.event_id], [])
  707. e = FakeEvent(
  708. id="E",
  709. sender=ALICE,
  710. type=EventTypes.Member,
  711. state_key="",
  712. content={},
  713. ).to_event([c.event_id, b.event_id], [])
  714. persisted_events = {a.event_id: a, b.event_id: b}
  715. unpersited_events = {c.event_id: c, d.event_id: d, e.event_id: e}
  716. state_sets = [
  717. {"a": a.event_id, "b": b.event_id, "e": e.event_id},
  718. {"c": c.event_id, "d": d.event_id},
  719. ]
  720. store = TestStateResolutionStore(persisted_events)
  721. diff_d = _get_auth_chain_difference(
  722. ROOM_ID, state_sets, unpersited_events, store
  723. )
  724. difference = self.successResultOf(defer.ensureDeferred(diff_d))
  725. self.assertEqual(difference, {d.event_id, e.event_id})
  726. T = TypeVar("T")
  727. def pairwise(iterable: Iterable[T]) -> Iterable[Tuple[T, T]]:
  728. "s -> (s0,s1), (s1,s2), (s2, s3), ..."
  729. a, b = itertools.tee(iterable)
  730. next(b, None)
  731. return zip(a, b)
  732. @attr.s
  733. class TestStateResolutionStore:
  734. event_map: Dict[str, EventBase] = attr.ib()
  735. def get_events(
  736. self, event_ids: Collection[str], allow_rejected: bool = False
  737. ) -> "defer.Deferred[Dict[str, EventBase]]":
  738. """Get events from the database
  739. Args:
  740. event_ids: The event_ids of the events to fetch
  741. allow_rejected: If True return rejected events.
  742. Returns:
  743. Dict from event_id to event.
  744. """
  745. return defer.succeed(
  746. {eid: self.event_map[eid] for eid in event_ids if eid in self.event_map}
  747. )
  748. def _get_auth_chain(self, event_ids: Iterable[str]) -> List[str]:
  749. """Gets the full auth chain for a set of events (including rejected
  750. events).
  751. Includes the given event IDs in the result.
  752. Note that:
  753. 1. All events must be state events.
  754. 2. For v1 rooms this may not have the full auth chain in the
  755. presence of rejected events
  756. Args:
  757. event_ids: The event IDs of the events to fetch the auth
  758. chain for. Must be state events.
  759. Returns:
  760. List of event IDs of the auth chain.
  761. """
  762. # Simple DFS for auth chain
  763. result = set()
  764. stack = list(event_ids)
  765. while stack:
  766. event_id = stack.pop()
  767. if event_id in result:
  768. continue
  769. result.add(event_id)
  770. event = self.event_map[event_id]
  771. for aid in event.auth_event_ids():
  772. stack.append(aid)
  773. return list(result)
  774. def get_auth_chain_difference(
  775. self, room_id: str, auth_sets: List[Set[str]]
  776. ) -> "defer.Deferred[Set[str]]":
  777. chains = [frozenset(self._get_auth_chain(a)) for a in auth_sets]
  778. common = set(chains[0]).intersection(*chains[1:])
  779. return defer.succeed(set(chains[0]).union(*chains[1:]) - common)