1
0

test_media_retention.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321
  1. # Copyright 2022 The Matrix.org Foundation C.I.C.
  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 io
  15. from typing import Iterable, Optional, Tuple
  16. from twisted.test.proto_helpers import MemoryReactor
  17. from synapse.rest import admin
  18. from synapse.rest.client import login, register, room
  19. from synapse.server import HomeServer
  20. from synapse.types import UserID
  21. from synapse.util import Clock
  22. from tests import unittest
  23. from tests.unittest import override_config
  24. from tests.utils import MockClock
  25. class MediaRetentionTestCase(unittest.HomeserverTestCase):
  26. ONE_DAY_IN_MS = 24 * 60 * 60 * 1000
  27. THIRTY_DAYS_IN_MS = 30 * ONE_DAY_IN_MS
  28. servlets = [
  29. room.register_servlets,
  30. login.register_servlets,
  31. register.register_servlets,
  32. admin.register_servlets_for_client_rest_resource,
  33. ]
  34. def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer:
  35. # We need to be able to test advancing time in the homeserver, so we
  36. # replace the test homeserver's default clock with a MockClock, which
  37. # supports advancing time.
  38. return self.setup_test_homeserver(clock=MockClock())
  39. def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
  40. self.remote_server_name = "remote.homeserver"
  41. self.store = hs.get_datastores().main
  42. # Create a user to upload media with
  43. test_user_id = self.register_user("alice", "password")
  44. # Inject media (recently accessed, old access, never accessed, old access
  45. # quarantined media) into both the local store and the remote cache, plus
  46. # one additional local media that is marked as protected from quarantine.
  47. media_repository = hs.get_media_repository()
  48. test_media_content = b"example string"
  49. def _create_media_and_set_attributes(
  50. last_accessed_ms: Optional[int],
  51. is_quarantined: Optional[bool] = False,
  52. is_protected: Optional[bool] = False,
  53. ) -> str:
  54. # "Upload" some media to the local media store
  55. mxc_uri = self.get_success(
  56. media_repository.create_content(
  57. media_type="text/plain",
  58. upload_name=None,
  59. content=io.BytesIO(test_media_content),
  60. content_length=len(test_media_content),
  61. auth_user=UserID.from_string(test_user_id),
  62. )
  63. )
  64. media_id = mxc_uri.split("/")[-1]
  65. # Set the last recently accessed time for this media
  66. if last_accessed_ms is not None:
  67. self.get_success(
  68. self.store.update_cached_last_access_time(
  69. local_media=(media_id,),
  70. remote_media=(),
  71. time_ms=last_accessed_ms,
  72. )
  73. )
  74. if is_quarantined:
  75. # Mark this media as quarantined
  76. self.get_success(
  77. self.store.quarantine_media_by_id(
  78. server_name=self.hs.config.server.server_name,
  79. media_id=media_id,
  80. quarantined_by="@theadmin:test",
  81. )
  82. )
  83. if is_protected:
  84. # Mark this media as protected from quarantine
  85. self.get_success(
  86. self.store.mark_local_media_as_safe(
  87. media_id=media_id,
  88. safe=True,
  89. )
  90. )
  91. return media_id
  92. def _cache_remote_media_and_set_attributes(
  93. media_id: str,
  94. last_accessed_ms: Optional[int],
  95. is_quarantined: Optional[bool] = False,
  96. ) -> str:
  97. # Pretend to cache some remote media
  98. self.get_success(
  99. self.store.store_cached_remote_media(
  100. origin=self.remote_server_name,
  101. media_id=media_id,
  102. media_type="text/plain",
  103. media_length=1,
  104. time_now_ms=clock.time_msec(),
  105. upload_name="testfile.txt",
  106. filesystem_id="abcdefg12345",
  107. )
  108. )
  109. # Set the last recently accessed time for this media
  110. if last_accessed_ms is not None:
  111. self.get_success(
  112. hs.get_datastores().main.update_cached_last_access_time(
  113. local_media=(),
  114. remote_media=((self.remote_server_name, media_id),),
  115. time_ms=last_accessed_ms,
  116. )
  117. )
  118. if is_quarantined:
  119. # Mark this media as quarantined
  120. self.get_success(
  121. self.store.quarantine_media_by_id(
  122. server_name=self.remote_server_name,
  123. media_id=media_id,
  124. quarantined_by="@theadmin:test",
  125. )
  126. )
  127. return media_id
  128. # Start with the local media store
  129. self.local_recently_accessed_media = _create_media_and_set_attributes(
  130. last_accessed_ms=self.THIRTY_DAYS_IN_MS,
  131. )
  132. self.local_not_recently_accessed_media = _create_media_and_set_attributes(
  133. last_accessed_ms=self.ONE_DAY_IN_MS,
  134. )
  135. self.local_not_recently_accessed_quarantined_media = (
  136. _create_media_and_set_attributes(
  137. last_accessed_ms=self.ONE_DAY_IN_MS,
  138. is_quarantined=True,
  139. )
  140. )
  141. self.local_not_recently_accessed_protected_media = (
  142. _create_media_and_set_attributes(
  143. last_accessed_ms=self.ONE_DAY_IN_MS,
  144. is_protected=True,
  145. )
  146. )
  147. self.local_never_accessed_media = _create_media_and_set_attributes(
  148. last_accessed_ms=None,
  149. )
  150. # And now the remote media store
  151. self.remote_recently_accessed_media = _cache_remote_media_and_set_attributes(
  152. media_id="a",
  153. last_accessed_ms=self.THIRTY_DAYS_IN_MS,
  154. )
  155. self.remote_not_recently_accessed_media = (
  156. _cache_remote_media_and_set_attributes(
  157. media_id="b",
  158. last_accessed_ms=self.ONE_DAY_IN_MS,
  159. )
  160. )
  161. self.remote_not_recently_accessed_quarantined_media = (
  162. _cache_remote_media_and_set_attributes(
  163. media_id="c",
  164. last_accessed_ms=self.ONE_DAY_IN_MS,
  165. is_quarantined=True,
  166. )
  167. )
  168. # Remote media will always have a "last accessed" attribute, as it would not
  169. # be fetched from the remote homeserver unless instigated by a user.
  170. @override_config(
  171. {
  172. "media_retention": {
  173. # Enable retention for local media
  174. "local_media_lifetime": "30d"
  175. # Cached remote media should not be purged
  176. }
  177. }
  178. )
  179. def test_local_media_retention(self) -> None:
  180. """
  181. Tests that local media that have not been accessed recently is purged, while
  182. cached remote media is unaffected.
  183. """
  184. # Advance 31 days (in seconds)
  185. self.reactor.advance(31 * 24 * 60 * 60)
  186. # Check that media has been correctly purged.
  187. # Local media accessed <30 days ago should still exist.
  188. # Remote media should be unaffected.
  189. self._assert_if_mxc_uris_purged(
  190. purged=[
  191. (
  192. self.hs.config.server.server_name,
  193. self.local_not_recently_accessed_media,
  194. ),
  195. (self.hs.config.server.server_name, self.local_never_accessed_media),
  196. ],
  197. not_purged=[
  198. (self.hs.config.server.server_name, self.local_recently_accessed_media),
  199. (
  200. self.hs.config.server.server_name,
  201. self.local_not_recently_accessed_quarantined_media,
  202. ),
  203. (
  204. self.hs.config.server.server_name,
  205. self.local_not_recently_accessed_protected_media,
  206. ),
  207. (self.remote_server_name, self.remote_recently_accessed_media),
  208. (self.remote_server_name, self.remote_not_recently_accessed_media),
  209. (
  210. self.remote_server_name,
  211. self.remote_not_recently_accessed_quarantined_media,
  212. ),
  213. ],
  214. )
  215. @override_config(
  216. {
  217. "media_retention": {
  218. # Enable retention for cached remote media
  219. "remote_media_lifetime": "30d"
  220. # Local media should not be purged
  221. }
  222. }
  223. )
  224. def test_remote_media_cache_retention(self) -> None:
  225. """
  226. Tests that entries from the remote media cache that have not been accessed
  227. recently is purged, while local media is unaffected.
  228. """
  229. # Advance 31 days (in seconds)
  230. self.reactor.advance(31 * 24 * 60 * 60)
  231. # Check that media has been correctly purged.
  232. # Local media should be unaffected.
  233. # Remote media accessed <30 days ago should still exist.
  234. self._assert_if_mxc_uris_purged(
  235. purged=[
  236. (self.remote_server_name, self.remote_not_recently_accessed_media),
  237. ],
  238. not_purged=[
  239. (self.remote_server_name, self.remote_recently_accessed_media),
  240. (self.hs.config.server.server_name, self.local_recently_accessed_media),
  241. (
  242. self.hs.config.server.server_name,
  243. self.local_not_recently_accessed_media,
  244. ),
  245. (
  246. self.hs.config.server.server_name,
  247. self.local_not_recently_accessed_quarantined_media,
  248. ),
  249. (
  250. self.hs.config.server.server_name,
  251. self.local_not_recently_accessed_protected_media,
  252. ),
  253. (
  254. self.remote_server_name,
  255. self.remote_not_recently_accessed_quarantined_media,
  256. ),
  257. (self.hs.config.server.server_name, self.local_never_accessed_media),
  258. ],
  259. )
  260. def _assert_if_mxc_uris_purged(
  261. self, purged: Iterable[Tuple[str, str]], not_purged: Iterable[Tuple[str, str]]
  262. ) -> None:
  263. def _assert_mxc_uri_purge_state(
  264. server_name: str, media_id: str, expect_purged: bool
  265. ) -> None:
  266. """Given an MXC URI, assert whether it has been purged or not."""
  267. if server_name == self.hs.config.server.server_name:
  268. found_media_dict = self.get_success(
  269. self.store.get_local_media(media_id)
  270. )
  271. else:
  272. found_media_dict = self.get_success(
  273. self.store.get_cached_remote_media(server_name, media_id)
  274. )
  275. mxc_uri = f"mxc://{server_name}/{media_id}"
  276. if expect_purged:
  277. self.assertIsNone(
  278. found_media_dict, msg=f"{mxc_uri} unexpectedly not purged"
  279. )
  280. else:
  281. self.assertIsNotNone(
  282. found_media_dict,
  283. msg=f"{mxc_uri} unexpectedly purged",
  284. )
  285. # Assert that the given MXC URIs have either been correctly purged or not.
  286. for server_name, media_id in purged:
  287. _assert_mxc_uri_purge_state(server_name, media_id, expect_purged=True)
  288. for server_name, media_id in not_purged:
  289. _assert_mxc_uri_purge_state(server_name, media_id, expect_purged=False)