tests.py 63 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545
  1. #! /usr/bin/env python3
  2. #
  3. # Copyright (C) 2018-2019 Garmin Ltd.
  4. #
  5. # SPDX-License-Identifier: GPL-2.0-only
  6. #
  7. from . import create_server, create_client
  8. from .server import DEFAULT_ANON_PERMS, ALL_PERMISSIONS
  9. from bb.asyncrpc import InvokeError
  10. import hashlib
  11. import logging
  12. import multiprocessing
  13. import os
  14. import sys
  15. import tempfile
  16. import threading
  17. import unittest
  18. import socket
  19. import time
  20. import signal
  21. import subprocess
  22. import json
  23. import re
  24. from pathlib import Path
  25. THIS_DIR = Path(__file__).parent
  26. BIN_DIR = THIS_DIR.parent.parent / "bin"
  27. def server_prefunc(server, idx):
  28. logging.basicConfig(level=logging.DEBUG, filename='bbhashserv-%d.log' % idx, filemode='w',
  29. format='%(levelname)s %(filename)s:%(lineno)d %(message)s')
  30. server.logger.debug("Running server %d" % idx)
  31. sys.stdout = open('bbhashserv-stdout-%d.log' % idx, 'w')
  32. sys.stderr = sys.stdout
  33. class HashEquivalenceTestSetup(object):
  34. METHOD = 'TestMethod'
  35. server_index = 0
  36. client_index = 0
  37. def start_server(self, dbpath=None, upstream=None, read_only=False, prefunc=server_prefunc, anon_perms=DEFAULT_ANON_PERMS, admin_username=None, admin_password=None):
  38. self.server_index += 1
  39. if dbpath is None:
  40. dbpath = self.make_dbpath()
  41. def cleanup_server(server):
  42. if server.process.exitcode is not None:
  43. return
  44. server.process.terminate()
  45. server.process.join()
  46. server = create_server(self.get_server_addr(self.server_index),
  47. dbpath,
  48. upstream=upstream,
  49. read_only=read_only,
  50. anon_perms=anon_perms,
  51. admin_username=admin_username,
  52. admin_password=admin_password)
  53. server.dbpath = dbpath
  54. server.serve_as_process(prefunc=prefunc, args=(self.server_index,))
  55. self.addCleanup(cleanup_server, server)
  56. return server
  57. def make_dbpath(self):
  58. return os.path.join(self.temp_dir.name, "db%d.sqlite" % self.server_index)
  59. def start_client(self, server_address, username=None, password=None):
  60. def cleanup_client(client):
  61. client.close()
  62. client = create_client(server_address, username=username, password=password)
  63. self.addCleanup(cleanup_client, client)
  64. return client
  65. def start_test_server(self):
  66. self.server = self.start_server()
  67. return self.server.address
  68. def start_auth_server(self):
  69. auth_server = self.start_server(self.server.dbpath, anon_perms=[], admin_username="admin", admin_password="password")
  70. self.auth_server_address = auth_server.address
  71. self.admin_client = self.start_client(auth_server.address, username="admin", password="password")
  72. return self.admin_client
  73. def auth_client(self, user):
  74. return self.start_client(self.auth_server_address, user["username"], user["token"])
  75. def setUp(self):
  76. self.temp_dir = tempfile.TemporaryDirectory(prefix='bb-hashserv')
  77. self.addCleanup(self.temp_dir.cleanup)
  78. self.server_address = self.start_test_server()
  79. self.client = self.start_client(self.server_address)
  80. def assertClientGetHash(self, client, taskhash, unihash):
  81. result = client.get_unihash(self.METHOD, taskhash)
  82. self.assertEqual(result, unihash)
  83. def assertUserPerms(self, user, permissions):
  84. with self.auth_client(user) as client:
  85. info = client.get_user()
  86. self.assertEqual(info, {
  87. "username": user["username"],
  88. "permissions": permissions,
  89. })
  90. def assertUserCanAuth(self, user):
  91. with self.start_client(self.auth_server_address) as client:
  92. client.auth(user["username"], user["token"])
  93. def assertUserCannotAuth(self, user):
  94. with self.start_client(self.auth_server_address) as client, self.assertRaises(InvokeError):
  95. client.auth(user["username"], user["token"])
  96. def create_test_hash(self, client):
  97. # Simple test that hashes can be created
  98. taskhash = '35788efcb8dfb0a02659d81cf2bfd695fb30faf9'
  99. outhash = '2765d4a5884be49b28601445c2760c5f21e7e5c0ee2b7e3fce98fd7e5970796f'
  100. unihash = 'f46d3fbb439bd9b921095da657a4de906510d2cd'
  101. self.assertClientGetHash(client, taskhash, None)
  102. result = client.report_unihash(taskhash, self.METHOD, outhash, unihash)
  103. self.assertEqual(result['unihash'], unihash, 'Server returned bad unihash')
  104. return taskhash, outhash, unihash
  105. def run_hashclient(self, args, **kwargs):
  106. try:
  107. p = subprocess.run(
  108. [BIN_DIR / "bitbake-hashclient"] + args,
  109. stdout=subprocess.PIPE,
  110. stderr=subprocess.STDOUT,
  111. encoding="utf-8",
  112. **kwargs
  113. )
  114. except subprocess.CalledProcessError as e:
  115. print(e.output)
  116. raise e
  117. print(p.stdout)
  118. return p
  119. class HashEquivalenceCommonTests(object):
  120. def auth_perms(self, *permissions):
  121. self.client_index += 1
  122. user = self.create_user(f"user-{self.client_index}", permissions)
  123. return self.auth_client(user)
  124. def create_user(self, username, permissions, *, client=None):
  125. def remove_user(username):
  126. try:
  127. self.admin_client.delete_user(username)
  128. except bb.asyncrpc.InvokeError:
  129. pass
  130. if client is None:
  131. client = self.admin_client
  132. user = client.new_user(username, permissions)
  133. self.addCleanup(remove_user, username)
  134. return user
  135. def test_create_hash(self):
  136. return self.create_test_hash(self.client)
  137. def test_create_equivalent(self):
  138. # Tests that a second reported task with the same outhash will be
  139. # assigned the same unihash
  140. taskhash = '53b8dce672cb6d0c73170be43f540460bfc347b4'
  141. outhash = '5a9cb1649625f0bf41fc7791b635cd9c2d7118c7f021ba87dcd03f72b67ce7a8'
  142. unihash = 'f37918cc02eb5a520b1aff86faacbc0a38124646'
  143. result = self.client.report_unihash(taskhash, self.METHOD, outhash, unihash)
  144. self.assertEqual(result['unihash'], unihash, 'Server returned bad unihash')
  145. # Report a different task with the same outhash. The returned unihash
  146. # should match the first task
  147. taskhash2 = '3bf6f1e89d26205aec90da04854fbdbf73afe6b4'
  148. unihash2 = 'af36b199320e611fbb16f1f277d3ee1d619ca58b'
  149. result = self.client.report_unihash(taskhash2, self.METHOD, outhash, unihash2)
  150. self.assertEqual(result['unihash'], unihash, 'Server returned bad unihash')
  151. def test_duplicate_taskhash(self):
  152. # Tests that duplicate reports of the same taskhash with different
  153. # outhash & unihash always return the unihash from the first reported
  154. # taskhash
  155. taskhash = '8aa96fcffb5831b3c2c0cb75f0431e3f8b20554a'
  156. outhash = 'afe240a439959ce86f5e322f8c208e1fedefea9e813f2140c81af866cc9edf7e'
  157. unihash = '218e57509998197d570e2c98512d0105985dffc9'
  158. self.client.report_unihash(taskhash, self.METHOD, outhash, unihash)
  159. self.assertClientGetHash(self.client, taskhash, unihash)
  160. outhash2 = '0904a7fe3dc712d9fd8a74a616ddca2a825a8ee97adf0bd3fc86082c7639914d'
  161. unihash2 = 'ae9a7d252735f0dafcdb10e2e02561ca3a47314c'
  162. self.client.report_unihash(taskhash, self.METHOD, outhash2, unihash2)
  163. self.assertClientGetHash(self.client, taskhash, unihash)
  164. outhash3 = '77623a549b5b1a31e3732dfa8fe61d7ce5d44b3370f253c5360e136b852967b4'
  165. unihash3 = '9217a7d6398518e5dc002ed58f2cbbbc78696603'
  166. self.client.report_unihash(taskhash, self.METHOD, outhash3, unihash3)
  167. self.assertClientGetHash(self.client, taskhash, unihash)
  168. def test_remove_taskhash(self):
  169. taskhash, outhash, unihash = self.create_test_hash(self.client)
  170. result = self.client.remove({"taskhash": taskhash})
  171. self.assertGreater(result["count"], 0)
  172. self.assertClientGetHash(self.client, taskhash, None)
  173. result_outhash = self.client.get_outhash(self.METHOD, outhash, taskhash)
  174. self.assertIsNone(result_outhash)
  175. def test_remove_unihash(self):
  176. taskhash, outhash, unihash = self.create_test_hash(self.client)
  177. result = self.client.remove({"unihash": unihash})
  178. self.assertGreater(result["count"], 0)
  179. self.assertClientGetHash(self.client, taskhash, None)
  180. def test_remove_outhash(self):
  181. taskhash, outhash, unihash = self.create_test_hash(self.client)
  182. result = self.client.remove({"outhash": outhash})
  183. self.assertGreater(result["count"], 0)
  184. result_outhash = self.client.get_outhash(self.METHOD, outhash, taskhash)
  185. self.assertIsNone(result_outhash)
  186. def test_remove_method(self):
  187. taskhash, outhash, unihash = self.create_test_hash(self.client)
  188. result = self.client.remove({"method": self.METHOD})
  189. self.assertGreater(result["count"], 0)
  190. self.assertClientGetHash(self.client, taskhash, None)
  191. result_outhash = self.client.get_outhash(self.METHOD, outhash, taskhash)
  192. self.assertIsNone(result_outhash)
  193. def test_clean_unused(self):
  194. taskhash, outhash, unihash = self.create_test_hash(self.client)
  195. # Clean the database, which should not remove anything because all hashes an in-use
  196. result = self.client.clean_unused(0)
  197. self.assertEqual(result["count"], 0)
  198. self.assertClientGetHash(self.client, taskhash, unihash)
  199. # Remove the unihash. The row in the outhash table should still be present
  200. self.client.remove({"unihash": unihash})
  201. result_outhash = self.client.get_outhash(self.METHOD, outhash, taskhash, False)
  202. self.assertIsNotNone(result_outhash)
  203. # Now clean with no minimum age which will remove the outhash
  204. result = self.client.clean_unused(0)
  205. self.assertEqual(result["count"], 1)
  206. result_outhash = self.client.get_outhash(self.METHOD, outhash, taskhash, False)
  207. self.assertIsNone(result_outhash)
  208. def test_huge_message(self):
  209. # Simple test that hashes can be created
  210. taskhash = 'c665584ee6817aa99edfc77a44dd853828279370'
  211. outhash = '3c979c3db45c569f51ab7626a4651074be3a9d11a84b1db076f5b14f7d39db44'
  212. unihash = '90e9bc1d1f094c51824adca7f8ea79a048d68824'
  213. self.assertClientGetHash(self.client, taskhash, None)
  214. siginfo = "0" * (self.client.max_chunk * 4)
  215. result = self.client.report_unihash(taskhash, self.METHOD, outhash, unihash, {
  216. 'outhash_siginfo': siginfo
  217. })
  218. self.assertEqual(result['unihash'], unihash, 'Server returned bad unihash')
  219. result_unihash = self.client.get_taskhash(self.METHOD, taskhash, True)
  220. self.assertEqual(result_unihash['taskhash'], taskhash)
  221. self.assertEqual(result_unihash['unihash'], unihash)
  222. self.assertEqual(result_unihash['method'], self.METHOD)
  223. result_outhash = self.client.get_outhash(self.METHOD, outhash, taskhash)
  224. self.assertEqual(result_outhash['taskhash'], taskhash)
  225. self.assertEqual(result_outhash['method'], self.METHOD)
  226. self.assertEqual(result_outhash['unihash'], unihash)
  227. self.assertEqual(result_outhash['outhash'], outhash)
  228. self.assertEqual(result_outhash['outhash_siginfo'], siginfo)
  229. def test_stress(self):
  230. def query_server(failures):
  231. client = Client(self.server_address)
  232. try:
  233. for i in range(1000):
  234. taskhash = hashlib.sha256()
  235. taskhash.update(str(i).encode('utf-8'))
  236. taskhash = taskhash.hexdigest()
  237. result = client.get_unihash(self.METHOD, taskhash)
  238. if result != taskhash:
  239. failures.append("taskhash mismatch: %s != %s" % (result, taskhash))
  240. finally:
  241. client.close()
  242. # Report hashes
  243. for i in range(1000):
  244. taskhash = hashlib.sha256()
  245. taskhash.update(str(i).encode('utf-8'))
  246. taskhash = taskhash.hexdigest()
  247. self.client.report_unihash(taskhash, self.METHOD, taskhash, taskhash)
  248. failures = []
  249. threads = [threading.Thread(target=query_server, args=(failures,)) for t in range(100)]
  250. for t in threads:
  251. t.start()
  252. for t in threads:
  253. t.join()
  254. self.assertFalse(failures)
  255. def test_upstream_server(self):
  256. # Tests upstream server support. This is done by creating two servers
  257. # that share a database file. The downstream server has it upstream
  258. # set to the test server, whereas the side server doesn't. This allows
  259. # verification that the hash requests are being proxied to the upstream
  260. # server by verifying that they appear on the downstream client, but not
  261. # the side client. It also verifies that the results are pulled into
  262. # the downstream database by checking that the downstream and side servers
  263. # match after the downstream is done waiting for all backfill tasks
  264. down_server = self.start_server(upstream=self.server_address)
  265. down_client = self.start_client(down_server.address)
  266. side_server = self.start_server(dbpath=down_server.dbpath)
  267. side_client = self.start_client(side_server.address)
  268. def check_hash(taskhash, unihash, old_sidehash):
  269. nonlocal down_client
  270. nonlocal side_client
  271. # check upstream server
  272. self.assertClientGetHash(self.client, taskhash, unihash)
  273. # Hash should *not* be present on the side server
  274. self.assertClientGetHash(side_client, taskhash, old_sidehash)
  275. # Hash should be present on the downstream server, since it
  276. # will defer to the upstream server. This will trigger
  277. # the backfill in the downstream server
  278. self.assertClientGetHash(down_client, taskhash, unihash)
  279. # After waiting for the downstream client to finish backfilling the
  280. # task from the upstream server, it should appear in the side server
  281. # since the database is populated
  282. down_client.backfill_wait()
  283. self.assertClientGetHash(side_client, taskhash, unihash)
  284. # Basic report
  285. taskhash = '8aa96fcffb5831b3c2c0cb75f0431e3f8b20554a'
  286. outhash = 'afe240a439959ce86f5e322f8c208e1fedefea9e813f2140c81af866cc9edf7e'
  287. unihash = '218e57509998197d570e2c98512d0105985dffc9'
  288. self.client.report_unihash(taskhash, self.METHOD, outhash, unihash)
  289. check_hash(taskhash, unihash, None)
  290. # Duplicated taskhash with multiple output hashes and unihashes.
  291. # All servers should agree with the originally reported hash
  292. outhash2 = '0904a7fe3dc712d9fd8a74a616ddca2a825a8ee97adf0bd3fc86082c7639914d'
  293. unihash2 = 'ae9a7d252735f0dafcdb10e2e02561ca3a47314c'
  294. self.client.report_unihash(taskhash, self.METHOD, outhash2, unihash2)
  295. check_hash(taskhash, unihash, unihash)
  296. # Report an equivalent task. The sideload will originally report
  297. # no unihash until backfilled
  298. taskhash3 = "044c2ec8aaf480685a00ff6ff49e6162e6ad34e1"
  299. unihash3 = "def64766090d28f627e816454ed46894bb3aab36"
  300. self.client.report_unihash(taskhash3, self.METHOD, outhash, unihash3)
  301. check_hash(taskhash3, unihash, None)
  302. # Test that reporting a unihash in the downstream client isn't
  303. # propagating to the upstream server
  304. taskhash4 = "e3da00593d6a7fb435c7e2114976c59c5fd6d561"
  305. outhash4 = "1cf8713e645f491eb9c959d20b5cae1c47133a292626dda9b10709857cbe688a"
  306. unihash4 = "3b5d3d83f07f259e9086fcb422c855286e18a57d"
  307. down_client.report_unihash(taskhash4, self.METHOD, outhash4, unihash4)
  308. down_client.backfill_wait()
  309. self.assertClientGetHash(down_client, taskhash4, unihash4)
  310. self.assertClientGetHash(side_client, taskhash4, unihash4)
  311. self.assertClientGetHash(self.client, taskhash4, None)
  312. # Test that reporting a unihash in the downstream is able to find a
  313. # match which was previously reported to the upstream server
  314. taskhash5 = '35788efcb8dfb0a02659d81cf2bfd695fb30faf9'
  315. outhash5 = '2765d4a5884be49b28601445c2760c5f21e7e5c0ee2b7e3fce98fd7e5970796f'
  316. unihash5 = 'f46d3fbb439bd9b921095da657a4de906510d2cd'
  317. result = self.client.report_unihash(taskhash5, self.METHOD, outhash5, unihash5)
  318. taskhash6 = '35788efcb8dfb0a02659d81cf2bfd695fb30fafa'
  319. unihash6 = 'f46d3fbb439bd9b921095da657a4de906510d2ce'
  320. result = down_client.report_unihash(taskhash6, self.METHOD, outhash5, unihash6)
  321. self.assertEqual(result['unihash'], unihash5, 'Server failed to copy unihash from upstream')
  322. # Tests read through from server with
  323. taskhash7 = '9d81d76242cc7cfaf7bf74b94b9cd2e29324ed74'
  324. outhash7 = '8470d56547eea6236d7c81a644ce74670ca0bbda998e13c629ef6bb3f0d60b69'
  325. unihash7 = '05d2a63c81e32f0a36542ca677e8ad852365c538'
  326. self.client.report_unihash(taskhash7, self.METHOD, outhash7, unihash7)
  327. result = down_client.get_taskhash(self.METHOD, taskhash7, True)
  328. self.assertEqual(result['unihash'], unihash7, 'Server failed to copy unihash from upstream')
  329. self.assertEqual(result['outhash'], outhash7, 'Server failed to copy unihash from upstream')
  330. self.assertEqual(result['taskhash'], taskhash7, 'Server failed to copy unihash from upstream')
  331. self.assertEqual(result['method'], self.METHOD)
  332. taskhash8 = '86978a4c8c71b9b487330b0152aade10c1ee58aa'
  333. outhash8 = 'ca8c128e9d9e4a28ef24d0508aa20b5cf880604eacd8f65c0e366f7e0cc5fbcf'
  334. unihash8 = 'd8bcf25369d40590ad7d08c84d538982f2023e01'
  335. self.client.report_unihash(taskhash8, self.METHOD, outhash8, unihash8)
  336. result = down_client.get_outhash(self.METHOD, outhash8, taskhash8)
  337. self.assertEqual(result['unihash'], unihash8, 'Server failed to copy unihash from upstream')
  338. self.assertEqual(result['outhash'], outhash8, 'Server failed to copy unihash from upstream')
  339. self.assertEqual(result['taskhash'], taskhash8, 'Server failed to copy unihash from upstream')
  340. self.assertEqual(result['method'], self.METHOD)
  341. taskhash9 = 'ae6339531895ddf5b67e663e6a374ad8ec71d81c'
  342. outhash9 = 'afc78172c81880ae10a1fec994b5b4ee33d196a001a1b66212a15ebe573e00b5'
  343. unihash9 = '6662e699d6e3d894b24408ff9a4031ef9b038ee8'
  344. self.client.report_unihash(taskhash9, self.METHOD, outhash9, unihash9)
  345. result = down_client.get_taskhash(self.METHOD, taskhash9, False)
  346. self.assertEqual(result['unihash'], unihash9, 'Server failed to copy unihash from upstream')
  347. self.assertEqual(result['taskhash'], taskhash9, 'Server failed to copy unihash from upstream')
  348. self.assertEqual(result['method'], self.METHOD)
  349. def test_unihash_exsits(self):
  350. taskhash, outhash, unihash = self.create_test_hash(self.client)
  351. self.assertTrue(self.client.unihash_exists(unihash))
  352. self.assertFalse(self.client.unihash_exists('6662e699d6e3d894b24408ff9a4031ef9b038ee8'))
  353. def test_ro_server(self):
  354. rw_server = self.start_server()
  355. rw_client = self.start_client(rw_server.address)
  356. ro_server = self.start_server(dbpath=rw_server.dbpath, read_only=True)
  357. ro_client = self.start_client(ro_server.address)
  358. # Report a hash via the read-write server
  359. taskhash = '35788efcb8dfb0a02659d81cf2bfd695fb30faf9'
  360. outhash = '2765d4a5884be49b28601445c2760c5f21e7e5c0ee2b7e3fce98fd7e5970796f'
  361. unihash = 'f46d3fbb439bd9b921095da657a4de906510d2cd'
  362. result = rw_client.report_unihash(taskhash, self.METHOD, outhash, unihash)
  363. self.assertEqual(result['unihash'], unihash, 'Server returned bad unihash')
  364. # Check the hash via the read-only server
  365. self.assertClientGetHash(ro_client, taskhash, unihash)
  366. # Ensure that reporting via the read-only server fails
  367. taskhash2 = 'c665584ee6817aa99edfc77a44dd853828279370'
  368. outhash2 = '3c979c3db45c569f51ab7626a4651074be3a9d11a84b1db076f5b14f7d39db44'
  369. unihash2 = '90e9bc1d1f094c51824adca7f8ea79a048d68824'
  370. result = ro_client.report_unihash(taskhash2, self.METHOD, outhash2, unihash2)
  371. self.assertEqual(result['unihash'], unihash2)
  372. # Ensure that the database was not modified
  373. self.assertClientGetHash(rw_client, taskhash2, None)
  374. def test_slow_server_start(self):
  375. # Ensures that the server will exit correctly even if it gets a SIGTERM
  376. # before entering the main loop
  377. event = multiprocessing.Event()
  378. def prefunc(server, idx):
  379. nonlocal event
  380. server_prefunc(server, idx)
  381. event.wait()
  382. def do_nothing(signum, frame):
  383. pass
  384. old_signal = signal.signal(signal.SIGTERM, do_nothing)
  385. self.addCleanup(signal.signal, signal.SIGTERM, old_signal)
  386. server = self.start_server(prefunc=prefunc)
  387. server.process.terminate()
  388. time.sleep(30)
  389. event.set()
  390. server.process.join(300)
  391. self.assertIsNotNone(server.process.exitcode, "Server did not exit in a timely manner!")
  392. def test_diverging_report_race(self):
  393. # Tests that a reported task will correctly pick up an updated unihash
  394. # This is a baseline report added to the database to ensure that there
  395. # is something to match against as equivalent
  396. outhash1 = 'afd11c366050bcd75ad763e898e4430e2a60659b26f83fbb22201a60672019fa'
  397. taskhash1 = '3bde230c743fc45ab61a065d7a1815fbfa01c4740e4c895af2eb8dc0f684a4ab'
  398. unihash1 = '3bde230c743fc45ab61a065d7a1815fbfa01c4740e4c895af2eb8dc0f684a4ab'
  399. result = self.client.report_unihash(taskhash1, self.METHOD, outhash1, unihash1)
  400. # Add a report that is equivalent to Task 1. It should ignore the
  401. # provided unihash and report the unihash from task 1
  402. taskhash2 = '6259ae8263bd94d454c086f501c37e64c4e83cae806902ca95b4ab513546b273'
  403. unihash2 = taskhash2
  404. result = self.client.report_unihash(taskhash2, self.METHOD, outhash1, unihash2)
  405. self.assertEqual(result['unihash'], unihash1)
  406. # Add another report for Task 2, but with a different outhash (e.g. the
  407. # task is non-deterministic). It should still be marked with the Task 1
  408. # unihash because it has the Task 2 taskhash, which is equivalent to
  409. # Task 1
  410. outhash3 = 'd2187ee3a8966db10b34fe0e863482288d9a6185cb8ef58a6c1c6ace87a2f24c'
  411. result = self.client.report_unihash(taskhash2, self.METHOD, outhash3, unihash2)
  412. self.assertEqual(result['unihash'], unihash1)
  413. def test_diverging_report_reverse_race(self):
  414. # Same idea as the previous test, but Tasks 2 and 3 are reported in
  415. # reverse order the opposite order
  416. outhash1 = 'afd11c366050bcd75ad763e898e4430e2a60659b26f83fbb22201a60672019fa'
  417. taskhash1 = '3bde230c743fc45ab61a065d7a1815fbfa01c4740e4c895af2eb8dc0f684a4ab'
  418. unihash1 = '3bde230c743fc45ab61a065d7a1815fbfa01c4740e4c895af2eb8dc0f684a4ab'
  419. result = self.client.report_unihash(taskhash1, self.METHOD, outhash1, unihash1)
  420. taskhash2 = '6259ae8263bd94d454c086f501c37e64c4e83cae806902ca95b4ab513546b273'
  421. unihash2 = taskhash2
  422. # Report Task 3 first. Since there is nothing else in the database it
  423. # will use the client provided unihash
  424. outhash3 = 'd2187ee3a8966db10b34fe0e863482288d9a6185cb8ef58a6c1c6ace87a2f24c'
  425. result = self.client.report_unihash(taskhash2, self.METHOD, outhash3, unihash2)
  426. self.assertEqual(result['unihash'], unihash2)
  427. # Report Task 2. This is equivalent to Task 1 but there is already a mapping for
  428. # taskhash2 so it will report unihash2
  429. result = self.client.report_unihash(taskhash2, self.METHOD, outhash1, unihash2)
  430. self.assertEqual(result['unihash'], unihash2)
  431. # The originally reported unihash for Task 3 should be unchanged even if it
  432. # shares a taskhash with Task 2
  433. self.assertClientGetHash(self.client, taskhash2, unihash2)
  434. def test_get_unihash_batch(self):
  435. TEST_INPUT = (
  436. # taskhash outhash unihash
  437. ('8aa96fcffb5831b3c2c0cb75f0431e3f8b20554a', 'afe240a439959ce86f5e322f8c208e1fedefea9e813f2140c81af866cc9edf7e','218e57509998197d570e2c98512d0105985dffc9'),
  438. # Duplicated taskhash with multiple output hashes and unihashes.
  439. ('8aa96fcffb5831b3c2c0cb75f0431e3f8b20554a', '0904a7fe3dc712d9fd8a74a616ddca2a825a8ee97adf0bd3fc86082c7639914d', 'ae9a7d252735f0dafcdb10e2e02561ca3a47314c'),
  440. # Equivalent hash
  441. ("044c2ec8aaf480685a00ff6ff49e6162e6ad34e1", '0904a7fe3dc712d9fd8a74a616ddca2a825a8ee97adf0bd3fc86082c7639914d', "def64766090d28f627e816454ed46894bb3aab36"),
  442. ("e3da00593d6a7fb435c7e2114976c59c5fd6d561", "1cf8713e645f491eb9c959d20b5cae1c47133a292626dda9b10709857cbe688a", "3b5d3d83f07f259e9086fcb422c855286e18a57d"),
  443. ('35788efcb8dfb0a02659d81cf2bfd695fb30faf9', '2765d4a5884be49b28601445c2760c5f21e7e5c0ee2b7e3fce98fd7e5970796f', 'f46d3fbb439bd9b921095da657a4de906510d2cd'),
  444. ('35788efcb8dfb0a02659d81cf2bfd695fb30fafa', '2765d4a5884be49b28601445c2760c5f21e7e5c0ee2b7e3fce98fd7e5970796f', 'f46d3fbb439bd9b921095da657a4de906510d2ce'),
  445. ('9d81d76242cc7cfaf7bf74b94b9cd2e29324ed74', '8470d56547eea6236d7c81a644ce74670ca0bbda998e13c629ef6bb3f0d60b69', '05d2a63c81e32f0a36542ca677e8ad852365c538'),
  446. )
  447. EXTRA_QUERIES = (
  448. "6b6be7a84ab179b4240c4302518dc3f6",
  449. )
  450. for taskhash, outhash, unihash in TEST_INPUT:
  451. self.client.report_unihash(taskhash, self.METHOD, outhash, unihash)
  452. result = self.client.get_unihash_batch(
  453. [(self.METHOD, data[0]) for data in TEST_INPUT] +
  454. [(self.METHOD, e) for e in EXTRA_QUERIES]
  455. )
  456. self.assertListEqual(result, [
  457. "218e57509998197d570e2c98512d0105985dffc9",
  458. "218e57509998197d570e2c98512d0105985dffc9",
  459. "218e57509998197d570e2c98512d0105985dffc9",
  460. "3b5d3d83f07f259e9086fcb422c855286e18a57d",
  461. "f46d3fbb439bd9b921095da657a4de906510d2cd",
  462. "f46d3fbb439bd9b921095da657a4de906510d2cd",
  463. "05d2a63c81e32f0a36542ca677e8ad852365c538",
  464. None,
  465. ])
  466. def test_unihash_exists_batch(self):
  467. TEST_INPUT = (
  468. # taskhash outhash unihash
  469. ('8aa96fcffb5831b3c2c0cb75f0431e3f8b20554a', 'afe240a439959ce86f5e322f8c208e1fedefea9e813f2140c81af866cc9edf7e','218e57509998197d570e2c98512d0105985dffc9'),
  470. # Duplicated taskhash with multiple output hashes and unihashes.
  471. ('8aa96fcffb5831b3c2c0cb75f0431e3f8b20554a', '0904a7fe3dc712d9fd8a74a616ddca2a825a8ee97adf0bd3fc86082c7639914d', 'ae9a7d252735f0dafcdb10e2e02561ca3a47314c'),
  472. # Equivalent hash
  473. ("044c2ec8aaf480685a00ff6ff49e6162e6ad34e1", '0904a7fe3dc712d9fd8a74a616ddca2a825a8ee97adf0bd3fc86082c7639914d', "def64766090d28f627e816454ed46894bb3aab36"),
  474. ("e3da00593d6a7fb435c7e2114976c59c5fd6d561", "1cf8713e645f491eb9c959d20b5cae1c47133a292626dda9b10709857cbe688a", "3b5d3d83f07f259e9086fcb422c855286e18a57d"),
  475. ('35788efcb8dfb0a02659d81cf2bfd695fb30faf9', '2765d4a5884be49b28601445c2760c5f21e7e5c0ee2b7e3fce98fd7e5970796f', 'f46d3fbb439bd9b921095da657a4de906510d2cd'),
  476. ('35788efcb8dfb0a02659d81cf2bfd695fb30fafa', '2765d4a5884be49b28601445c2760c5f21e7e5c0ee2b7e3fce98fd7e5970796f', 'f46d3fbb439bd9b921095da657a4de906510d2ce'),
  477. ('9d81d76242cc7cfaf7bf74b94b9cd2e29324ed74', '8470d56547eea6236d7c81a644ce74670ca0bbda998e13c629ef6bb3f0d60b69', '05d2a63c81e32f0a36542ca677e8ad852365c538'),
  478. )
  479. EXTRA_QUERIES = (
  480. "6b6be7a84ab179b4240c4302518dc3f6",
  481. )
  482. result_unihashes = set()
  483. for taskhash, outhash, unihash in TEST_INPUT:
  484. result = self.client.report_unihash(taskhash, self.METHOD, outhash, unihash)
  485. result_unihashes.add(result["unihash"])
  486. query = []
  487. expected = []
  488. for _, _, unihash in TEST_INPUT:
  489. query.append(unihash)
  490. expected.append(unihash in result_unihashes)
  491. for unihash in EXTRA_QUERIES:
  492. query.append(unihash)
  493. expected.append(False)
  494. result = self.client.unihash_exists_batch(query)
  495. self.assertListEqual(result, expected)
  496. def test_auth_read_perms(self):
  497. admin_client = self.start_auth_server()
  498. # Create hashes with non-authenticated server
  499. taskhash, outhash, unihash = self.create_test_hash(self.client)
  500. # Validate hash can be retrieved using authenticated client
  501. with self.auth_perms("@read") as client:
  502. self.assertClientGetHash(client, taskhash, unihash)
  503. with self.auth_perms() as client, self.assertRaises(InvokeError):
  504. self.assertClientGetHash(client, taskhash, unihash)
  505. def test_auth_report_perms(self):
  506. admin_client = self.start_auth_server()
  507. # Without read permission, the user is completely denied
  508. with self.auth_perms() as client, self.assertRaises(InvokeError):
  509. self.create_test_hash(client)
  510. # Read permission allows the call to succeed, but it doesn't record
  511. # anythin in the database
  512. with self.auth_perms("@read") as client:
  513. taskhash, outhash, unihash = self.create_test_hash(client)
  514. self.assertClientGetHash(client, taskhash, None)
  515. # Report permission alone is insufficient
  516. with self.auth_perms("@report") as client, self.assertRaises(InvokeError):
  517. self.create_test_hash(client)
  518. # Read and report permission actually modify the database
  519. with self.auth_perms("@read", "@report") as client:
  520. taskhash, outhash, unihash = self.create_test_hash(client)
  521. self.assertClientGetHash(client, taskhash, unihash)
  522. def test_auth_no_token_refresh_from_anon_user(self):
  523. self.start_auth_server()
  524. with self.start_client(self.auth_server_address) as client, self.assertRaises(InvokeError):
  525. client.refresh_token()
  526. def test_auth_self_token_refresh(self):
  527. admin_client = self.start_auth_server()
  528. # Create a new user with no permissions
  529. user = self.create_user("test-user", [])
  530. with self.auth_client(user) as client:
  531. new_user = client.refresh_token()
  532. self.assertEqual(user["username"], new_user["username"])
  533. self.assertNotEqual(user["token"], new_user["token"])
  534. self.assertUserCanAuth(new_user)
  535. self.assertUserCannotAuth(user)
  536. # Explicitly specifying with your own username is fine also
  537. with self.auth_client(new_user) as client:
  538. new_user2 = client.refresh_token(user["username"])
  539. self.assertEqual(user["username"], new_user2["username"])
  540. self.assertNotEqual(user["token"], new_user2["token"])
  541. self.assertUserCanAuth(new_user2)
  542. self.assertUserCannotAuth(new_user)
  543. self.assertUserCannotAuth(user)
  544. def test_auth_token_refresh(self):
  545. admin_client = self.start_auth_server()
  546. user = self.create_user("test-user", [])
  547. with self.auth_perms() as client, self.assertRaises(InvokeError):
  548. client.refresh_token(user["username"])
  549. with self.auth_perms("@user-admin") as client:
  550. new_user = client.refresh_token(user["username"])
  551. self.assertEqual(user["username"], new_user["username"])
  552. self.assertNotEqual(user["token"], new_user["token"])
  553. self.assertUserCanAuth(new_user)
  554. self.assertUserCannotAuth(user)
  555. def test_auth_self_get_user(self):
  556. admin_client = self.start_auth_server()
  557. user = self.create_user("test-user", [])
  558. user_info = user.copy()
  559. del user_info["token"]
  560. with self.auth_client(user) as client:
  561. info = client.get_user()
  562. self.assertEqual(info, user_info)
  563. # Explicitly asking for your own username is fine also
  564. info = client.get_user(user["username"])
  565. self.assertEqual(info, user_info)
  566. def test_auth_get_user(self):
  567. admin_client = self.start_auth_server()
  568. user = self.create_user("test-user", [])
  569. user_info = user.copy()
  570. del user_info["token"]
  571. with self.auth_perms() as client, self.assertRaises(InvokeError):
  572. client.get_user(user["username"])
  573. with self.auth_perms("@user-admin") as client:
  574. info = client.get_user(user["username"])
  575. self.assertEqual(info, user_info)
  576. info = client.get_user("nonexist-user")
  577. self.assertIsNone(info)
  578. def test_auth_reconnect(self):
  579. admin_client = self.start_auth_server()
  580. user = self.create_user("test-user", [])
  581. user_info = user.copy()
  582. del user_info["token"]
  583. with self.auth_client(user) as client:
  584. info = client.get_user()
  585. self.assertEqual(info, user_info)
  586. client.disconnect()
  587. info = client.get_user()
  588. self.assertEqual(info, user_info)
  589. def test_auth_delete_user(self):
  590. admin_client = self.start_auth_server()
  591. user = self.create_user("test-user", [])
  592. # self service
  593. with self.auth_client(user) as client:
  594. client.delete_user(user["username"])
  595. self.assertIsNone(admin_client.get_user(user["username"]))
  596. user = self.create_user("test-user", [])
  597. with self.auth_perms() as client, self.assertRaises(InvokeError):
  598. client.delete_user(user["username"])
  599. with self.auth_perms("@user-admin") as client:
  600. client.delete_user(user["username"])
  601. # User doesn't exist, so even though the permission is correct, it's an
  602. # error
  603. with self.auth_perms("@user-admin") as client, self.assertRaises(InvokeError):
  604. client.delete_user(user["username"])
  605. def test_auth_set_user_perms(self):
  606. admin_client = self.start_auth_server()
  607. user = self.create_user("test-user", [])
  608. self.assertUserPerms(user, [])
  609. # No self service to change permissions
  610. with self.auth_client(user) as client, self.assertRaises(InvokeError):
  611. client.set_user_perms(user["username"], ["@all"])
  612. self.assertUserPerms(user, [])
  613. with self.auth_perms() as client, self.assertRaises(InvokeError):
  614. client.set_user_perms(user["username"], ["@all"])
  615. self.assertUserPerms(user, [])
  616. with self.auth_perms("@user-admin") as client:
  617. client.set_user_perms(user["username"], ["@all"])
  618. self.assertUserPerms(user, sorted(list(ALL_PERMISSIONS)))
  619. # Bad permissions
  620. with self.auth_perms("@user-admin") as client, self.assertRaises(InvokeError):
  621. client.set_user_perms(user["username"], ["@this-is-not-a-permission"])
  622. self.assertUserPerms(user, sorted(list(ALL_PERMISSIONS)))
  623. def test_auth_get_all_users(self):
  624. admin_client = self.start_auth_server()
  625. user = self.create_user("test-user", [])
  626. with self.auth_client(user) as client, self.assertRaises(InvokeError):
  627. client.get_all_users()
  628. # Give the test user the correct permission
  629. admin_client.set_user_perms(user["username"], ["@user-admin"])
  630. with self.auth_client(user) as client:
  631. all_users = client.get_all_users()
  632. # Convert to a dictionary for easier comparison
  633. all_users = {u["username"]: u for u in all_users}
  634. self.assertEqual(all_users,
  635. {
  636. "admin": {
  637. "username": "admin",
  638. "permissions": sorted(list(ALL_PERMISSIONS)),
  639. },
  640. "test-user": {
  641. "username": "test-user",
  642. "permissions": ["@user-admin"],
  643. }
  644. }
  645. )
  646. def test_auth_new_user(self):
  647. self.start_auth_server()
  648. permissions = ["@read", "@report", "@db-admin", "@user-admin"]
  649. permissions.sort()
  650. with self.auth_perms() as client, self.assertRaises(InvokeError):
  651. self.create_user("test-user", permissions, client=client)
  652. with self.auth_perms("@user-admin") as client:
  653. user = self.create_user("test-user", permissions, client=client)
  654. self.assertIn("token", user)
  655. self.assertEqual(user["username"], "test-user")
  656. self.assertEqual(user["permissions"], permissions)
  657. def test_auth_become_user(self):
  658. admin_client = self.start_auth_server()
  659. user = self.create_user("test-user", ["@read", "@report"])
  660. user_info = user.copy()
  661. del user_info["token"]
  662. with self.auth_perms() as client, self.assertRaises(InvokeError):
  663. client.become_user(user["username"])
  664. with self.auth_perms("@user-admin") as client:
  665. become = client.become_user(user["username"])
  666. self.assertEqual(become, user_info)
  667. info = client.get_user()
  668. self.assertEqual(info, user_info)
  669. # Verify become user is preserved across disconnect
  670. client.disconnect()
  671. info = client.get_user()
  672. self.assertEqual(info, user_info)
  673. # test-user doesn't have become_user permissions, so this should
  674. # not work
  675. with self.assertRaises(InvokeError):
  676. client.become_user(user["username"])
  677. # No self-service of become
  678. with self.auth_client(user) as client, self.assertRaises(InvokeError):
  679. client.become_user(user["username"])
  680. # Give test user permissions to become
  681. admin_client.set_user_perms(user["username"], ["@user-admin"])
  682. # It's possible to become yourself (effectively a noop)
  683. with self.auth_perms("@user-admin") as client:
  684. become = client.become_user(client.username)
  685. def test_auth_gc(self):
  686. admin_client = self.start_auth_server()
  687. with self.auth_perms() as client, self.assertRaises(InvokeError):
  688. client.gc_mark("ABC", {"unihash": "123"})
  689. with self.auth_perms() as client, self.assertRaises(InvokeError):
  690. client.gc_status()
  691. with self.auth_perms() as client, self.assertRaises(InvokeError):
  692. client.gc_sweep("ABC")
  693. with self.auth_perms("@db-admin") as client:
  694. client.gc_mark("ABC", {"unihash": "123"})
  695. with self.auth_perms("@db-admin") as client:
  696. client.gc_status()
  697. with self.auth_perms("@db-admin") as client:
  698. client.gc_sweep("ABC")
  699. def test_get_db_usage(self):
  700. usage = self.client.get_db_usage()
  701. self.assertTrue(isinstance(usage, dict))
  702. for name in usage.keys():
  703. self.assertTrue(isinstance(usage[name], dict))
  704. self.assertIn("rows", usage[name])
  705. self.assertTrue(isinstance(usage[name]["rows"], int))
  706. def test_get_db_query_columns(self):
  707. columns = self.client.get_db_query_columns()
  708. self.assertTrue(isinstance(columns, list))
  709. self.assertTrue(len(columns) > 0)
  710. for col in columns:
  711. self.client.remove({col: ""})
  712. def test_auth_is_owner(self):
  713. admin_client = self.start_auth_server()
  714. user = self.create_user("test-user", ["@read", "@report"])
  715. with self.auth_client(user) as client:
  716. taskhash, outhash, unihash = self.create_test_hash(client)
  717. data = client.get_taskhash(self.METHOD, taskhash, True)
  718. self.assertEqual(data["owner"], user["username"])
  719. def test_gc(self):
  720. taskhash = '53b8dce672cb6d0c73170be43f540460bfc347b4'
  721. outhash = '5a9cb1649625f0bf41fc7791b635cd9c2d7118c7f021ba87dcd03f72b67ce7a8'
  722. unihash = 'f37918cc02eb5a520b1aff86faacbc0a38124646'
  723. result = self.client.report_unihash(taskhash, self.METHOD, outhash, unihash)
  724. self.assertEqual(result['unihash'], unihash, 'Server returned bad unihash')
  725. taskhash2 = '3bf6f1e89d26205aec90da04854fbdbf73afe6b4'
  726. outhash2 = '77623a549b5b1a31e3732dfa8fe61d7ce5d44b3370f253c5360e136b852967b4'
  727. unihash2 = 'af36b199320e611fbb16f1f277d3ee1d619ca58b'
  728. result = self.client.report_unihash(taskhash2, self.METHOD, outhash2, unihash2)
  729. self.assertClientGetHash(self.client, taskhash2, unihash2)
  730. # Mark the first unihash to be kept
  731. ret = self.client.gc_mark("ABC", {"unihash": unihash, "method": self.METHOD})
  732. self.assertEqual(ret, {"count": 1})
  733. ret = self.client.gc_status()
  734. self.assertEqual(ret, {"mark": "ABC", "keep": 1, "remove": 1})
  735. # Second hash is still there; mark doesn't delete hashes
  736. self.assertClientGetHash(self.client, taskhash2, unihash2)
  737. ret = self.client.gc_sweep("ABC")
  738. self.assertEqual(ret, {"count": 1})
  739. # Hash is gone. Taskhash is returned for second hash
  740. self.assertClientGetHash(self.client, taskhash2, None)
  741. # First hash is still present
  742. self.assertClientGetHash(self.client, taskhash, unihash)
  743. def test_gc_stream(self):
  744. taskhash = '53b8dce672cb6d0c73170be43f540460bfc347b4'
  745. outhash = '5a9cb1649625f0bf41fc7791b635cd9c2d7118c7f021ba87dcd03f72b67ce7a8'
  746. unihash = 'f37918cc02eb5a520b1aff86faacbc0a38124646'
  747. result = self.client.report_unihash(taskhash, self.METHOD, outhash, unihash)
  748. self.assertEqual(result['unihash'], unihash, 'Server returned bad unihash')
  749. taskhash2 = '3bf6f1e89d26205aec90da04854fbdbf73afe6b4'
  750. outhash2 = '77623a549b5b1a31e3732dfa8fe61d7ce5d44b3370f253c5360e136b852967b4'
  751. unihash2 = 'af36b199320e611fbb16f1f277d3ee1d619ca58b'
  752. result = self.client.report_unihash(taskhash2, self.METHOD, outhash2, unihash2)
  753. self.assertClientGetHash(self.client, taskhash2, unihash2)
  754. taskhash3 = 'a1117c1f5a7c9ab2f5a39cc6fe5e6152169d09c0'
  755. outhash3 = '7289c414905303700a1117c1f5a7c9ab2f5a39cc6fe5e6152169d09c04f9a53c'
  756. unihash3 = '905303700a1117c1f5a7c9ab2f5a39cc6fe5e615'
  757. result = self.client.report_unihash(taskhash3, self.METHOD, outhash3, unihash3)
  758. self.assertClientGetHash(self.client, taskhash3, unihash3)
  759. # Mark the first unihash to be kept
  760. ret = self.client.gc_mark_stream("ABC", (f"unihash {h}" for h in [unihash, unihash2]))
  761. self.assertEqual(ret, {"count": 2})
  762. ret = self.client.gc_status()
  763. self.assertEqual(ret, {"mark": "ABC", "keep": 2, "remove": 1})
  764. # Third hash is still there; mark doesn't delete hashes
  765. self.assertClientGetHash(self.client, taskhash3, unihash3)
  766. ret = self.client.gc_sweep("ABC")
  767. self.assertEqual(ret, {"count": 1})
  768. # Hash is gone. Taskhash is returned for second hash
  769. self.assertClientGetHash(self.client, taskhash3, None)
  770. # First hash is still present
  771. self.assertClientGetHash(self.client, taskhash, unihash)
  772. # Second hash is still present
  773. self.assertClientGetHash(self.client, taskhash2, unihash2)
  774. def test_gc_switch_mark(self):
  775. taskhash = '53b8dce672cb6d0c73170be43f540460bfc347b4'
  776. outhash = '5a9cb1649625f0bf41fc7791b635cd9c2d7118c7f021ba87dcd03f72b67ce7a8'
  777. unihash = 'f37918cc02eb5a520b1aff86faacbc0a38124646'
  778. result = self.client.report_unihash(taskhash, self.METHOD, outhash, unihash)
  779. self.assertEqual(result['unihash'], unihash, 'Server returned bad unihash')
  780. taskhash2 = '3bf6f1e89d26205aec90da04854fbdbf73afe6b4'
  781. outhash2 = '77623a549b5b1a31e3732dfa8fe61d7ce5d44b3370f253c5360e136b852967b4'
  782. unihash2 = 'af36b199320e611fbb16f1f277d3ee1d619ca58b'
  783. result = self.client.report_unihash(taskhash2, self.METHOD, outhash2, unihash2)
  784. self.assertClientGetHash(self.client, taskhash2, unihash2)
  785. # Mark the first unihash to be kept
  786. ret = self.client.gc_mark("ABC", {"unihash": unihash, "method": self.METHOD})
  787. self.assertEqual(ret, {"count": 1})
  788. ret = self.client.gc_status()
  789. self.assertEqual(ret, {"mark": "ABC", "keep": 1, "remove": 1})
  790. # Second hash is still there; mark doesn't delete hashes
  791. self.assertClientGetHash(self.client, taskhash2, unihash2)
  792. # Switch to a different mark and mark the second hash. This will start
  793. # a new collection cycle
  794. ret = self.client.gc_mark("DEF", {"unihash": unihash2, "method": self.METHOD})
  795. self.assertEqual(ret, {"count": 1})
  796. ret = self.client.gc_status()
  797. self.assertEqual(ret, {"mark": "DEF", "keep": 1, "remove": 1})
  798. # Both hashes are still present
  799. self.assertClientGetHash(self.client, taskhash2, unihash2)
  800. self.assertClientGetHash(self.client, taskhash, unihash)
  801. # Sweep with the new mark
  802. ret = self.client.gc_sweep("DEF")
  803. self.assertEqual(ret, {"count": 1})
  804. # First hash is gone, second is kept
  805. self.assertClientGetHash(self.client, taskhash2, unihash2)
  806. self.assertClientGetHash(self.client, taskhash, None)
  807. def test_gc_switch_sweep_mark(self):
  808. taskhash = '53b8dce672cb6d0c73170be43f540460bfc347b4'
  809. outhash = '5a9cb1649625f0bf41fc7791b635cd9c2d7118c7f021ba87dcd03f72b67ce7a8'
  810. unihash = 'f37918cc02eb5a520b1aff86faacbc0a38124646'
  811. result = self.client.report_unihash(taskhash, self.METHOD, outhash, unihash)
  812. self.assertEqual(result['unihash'], unihash, 'Server returned bad unihash')
  813. taskhash2 = '3bf6f1e89d26205aec90da04854fbdbf73afe6b4'
  814. outhash2 = '77623a549b5b1a31e3732dfa8fe61d7ce5d44b3370f253c5360e136b852967b4'
  815. unihash2 = 'af36b199320e611fbb16f1f277d3ee1d619ca58b'
  816. result = self.client.report_unihash(taskhash2, self.METHOD, outhash2, unihash2)
  817. self.assertClientGetHash(self.client, taskhash2, unihash2)
  818. # Mark the first unihash to be kept
  819. ret = self.client.gc_mark("ABC", {"unihash": unihash, "method": self.METHOD})
  820. self.assertEqual(ret, {"count": 1})
  821. ret = self.client.gc_status()
  822. self.assertEqual(ret, {"mark": "ABC", "keep": 1, "remove": 1})
  823. # Sweeping with a different mark raises an error
  824. with self.assertRaises(InvokeError):
  825. self.client.gc_sweep("DEF")
  826. # Both hashes are present
  827. self.assertClientGetHash(self.client, taskhash2, unihash2)
  828. self.assertClientGetHash(self.client, taskhash, unihash)
  829. def test_gc_new_hashes(self):
  830. taskhash = '53b8dce672cb6d0c73170be43f540460bfc347b4'
  831. outhash = '5a9cb1649625f0bf41fc7791b635cd9c2d7118c7f021ba87dcd03f72b67ce7a8'
  832. unihash = 'f37918cc02eb5a520b1aff86faacbc0a38124646'
  833. result = self.client.report_unihash(taskhash, self.METHOD, outhash, unihash)
  834. self.assertEqual(result['unihash'], unihash, 'Server returned bad unihash')
  835. # Start a new garbage collection
  836. ret = self.client.gc_mark("ABC", {"unihash": unihash, "method": self.METHOD})
  837. self.assertEqual(ret, {"count": 1})
  838. ret = self.client.gc_status()
  839. self.assertEqual(ret, {"mark": "ABC", "keep": 1, "remove": 0})
  840. # Add second hash. It should inherit the mark from the current garbage
  841. # collection operation
  842. taskhash2 = '3bf6f1e89d26205aec90da04854fbdbf73afe6b4'
  843. outhash2 = '77623a549b5b1a31e3732dfa8fe61d7ce5d44b3370f253c5360e136b852967b4'
  844. unihash2 = 'af36b199320e611fbb16f1f277d3ee1d619ca58b'
  845. result = self.client.report_unihash(taskhash2, self.METHOD, outhash2, unihash2)
  846. self.assertClientGetHash(self.client, taskhash2, unihash2)
  847. # Sweep should remove nothing
  848. ret = self.client.gc_sweep("ABC")
  849. self.assertEqual(ret, {"count": 0})
  850. # Both hashes are present
  851. self.assertClientGetHash(self.client, taskhash2, unihash2)
  852. self.assertClientGetHash(self.client, taskhash, unihash)
  853. class TestHashEquivalenceClient(HashEquivalenceTestSetup, unittest.TestCase):
  854. def get_server_addr(self, server_idx):
  855. return "unix://" + os.path.join(self.temp_dir.name, 'sock%d' % server_idx)
  856. def test_get(self):
  857. taskhash, outhash, unihash = self.create_test_hash(self.client)
  858. p = self.run_hashclient(["--address", self.server_address, "get", self.METHOD, taskhash])
  859. data = json.loads(p.stdout)
  860. self.assertEqual(data["unihash"], unihash)
  861. self.assertEqual(data["outhash"], outhash)
  862. self.assertEqual(data["taskhash"], taskhash)
  863. self.assertEqual(data["method"], self.METHOD)
  864. def test_get_outhash(self):
  865. taskhash, outhash, unihash = self.create_test_hash(self.client)
  866. p = self.run_hashclient(["--address", self.server_address, "get-outhash", self.METHOD, outhash, taskhash])
  867. data = json.loads(p.stdout)
  868. self.assertEqual(data["unihash"], unihash)
  869. self.assertEqual(data["outhash"], outhash)
  870. self.assertEqual(data["taskhash"], taskhash)
  871. self.assertEqual(data["method"], self.METHOD)
  872. def test_stats(self):
  873. p = self.run_hashclient(["--address", self.server_address, "stats"], check=True)
  874. json.loads(p.stdout)
  875. def test_stress(self):
  876. self.run_hashclient(["--address", self.server_address, "stress"], check=True)
  877. def test_unihash_exsits(self):
  878. taskhash, outhash, unihash = self.create_test_hash(self.client)
  879. p = self.run_hashclient([
  880. "--address", self.server_address,
  881. "unihash-exists", unihash,
  882. ], check=True)
  883. self.assertEqual(p.stdout.strip(), "true")
  884. p = self.run_hashclient([
  885. "--address", self.server_address,
  886. "unihash-exists", '6662e699d6e3d894b24408ff9a4031ef9b038ee8',
  887. ], check=True)
  888. self.assertEqual(p.stdout.strip(), "false")
  889. def test_unihash_exsits_quiet(self):
  890. taskhash, outhash, unihash = self.create_test_hash(self.client)
  891. p = self.run_hashclient([
  892. "--address", self.server_address,
  893. "unihash-exists", unihash,
  894. "--quiet",
  895. ])
  896. self.assertEqual(p.returncode, 0)
  897. self.assertEqual(p.stdout.strip(), "")
  898. p = self.run_hashclient([
  899. "--address", self.server_address,
  900. "unihash-exists", '6662e699d6e3d894b24408ff9a4031ef9b038ee8',
  901. "--quiet",
  902. ])
  903. self.assertEqual(p.returncode, 1)
  904. self.assertEqual(p.stdout.strip(), "")
  905. def test_remove_taskhash(self):
  906. taskhash, outhash, unihash = self.create_test_hash(self.client)
  907. self.run_hashclient([
  908. "--address", self.server_address,
  909. "remove",
  910. "--where", "taskhash", taskhash,
  911. ], check=True)
  912. self.assertClientGetHash(self.client, taskhash, None)
  913. result_outhash = self.client.get_outhash(self.METHOD, outhash, taskhash)
  914. self.assertIsNone(result_outhash)
  915. def test_remove_unihash(self):
  916. taskhash, outhash, unihash = self.create_test_hash(self.client)
  917. self.run_hashclient([
  918. "--address", self.server_address,
  919. "remove",
  920. "--where", "unihash", unihash,
  921. ], check=True)
  922. self.assertClientGetHash(self.client, taskhash, None)
  923. def test_remove_outhash(self):
  924. taskhash, outhash, unihash = self.create_test_hash(self.client)
  925. self.run_hashclient([
  926. "--address", self.server_address,
  927. "remove",
  928. "--where", "outhash", outhash,
  929. ], check=True)
  930. result_outhash = self.client.get_outhash(self.METHOD, outhash, taskhash)
  931. self.assertIsNone(result_outhash)
  932. def test_remove_method(self):
  933. taskhash, outhash, unihash = self.create_test_hash(self.client)
  934. self.run_hashclient([
  935. "--address", self.server_address,
  936. "remove",
  937. "--where", "method", self.METHOD,
  938. ], check=True)
  939. self.assertClientGetHash(self.client, taskhash, None)
  940. result_outhash = self.client.get_outhash(self.METHOD, outhash, taskhash)
  941. self.assertIsNone(result_outhash)
  942. def test_clean_unused(self):
  943. taskhash, outhash, unihash = self.create_test_hash(self.client)
  944. # Clean the database, which should not remove anything because all hashes an in-use
  945. self.run_hashclient([
  946. "--address", self.server_address,
  947. "clean-unused", "0",
  948. ], check=True)
  949. self.assertClientGetHash(self.client, taskhash, unihash)
  950. # Remove the unihash. The row in the outhash table should still be present
  951. self.run_hashclient([
  952. "--address", self.server_address,
  953. "remove",
  954. "--where", "unihash", unihash,
  955. ], check=True)
  956. result_outhash = self.client.get_outhash(self.METHOD, outhash, taskhash, False)
  957. self.assertIsNotNone(result_outhash)
  958. # Now clean with no minimum age which will remove the outhash
  959. self.run_hashclient([
  960. "--address", self.server_address,
  961. "clean-unused", "0",
  962. ], check=True)
  963. result_outhash = self.client.get_outhash(self.METHOD, outhash, taskhash, False)
  964. self.assertIsNone(result_outhash)
  965. def test_refresh_token(self):
  966. admin_client = self.start_auth_server()
  967. user = admin_client.new_user("test-user", ["@read", "@report"])
  968. p = self.run_hashclient([
  969. "--address", self.auth_server_address,
  970. "--login", user["username"],
  971. "--password", user["token"],
  972. "refresh-token"
  973. ], check=True)
  974. new_token = None
  975. for l in p.stdout.splitlines():
  976. l = l.rstrip()
  977. m = re.match(r'Token: +(.*)$', l)
  978. if m is not None:
  979. new_token = m.group(1)
  980. self.assertTrue(new_token)
  981. print("New token is %r" % new_token)
  982. self.run_hashclient([
  983. "--address", self.auth_server_address,
  984. "--login", user["username"],
  985. "--password", new_token,
  986. "get-user"
  987. ], check=True)
  988. def test_set_user_perms(self):
  989. admin_client = self.start_auth_server()
  990. user = admin_client.new_user("test-user", ["@read"])
  991. self.run_hashclient([
  992. "--address", self.auth_server_address,
  993. "--login", admin_client.username,
  994. "--password", admin_client.password,
  995. "set-user-perms",
  996. "-u", user["username"],
  997. "@read", "@report",
  998. ], check=True)
  999. new_user = admin_client.get_user(user["username"])
  1000. self.assertEqual(set(new_user["permissions"]), {"@read", "@report"})
  1001. def test_get_user(self):
  1002. admin_client = self.start_auth_server()
  1003. user = admin_client.new_user("test-user", ["@read"])
  1004. p = self.run_hashclient([
  1005. "--address", self.auth_server_address,
  1006. "--login", admin_client.username,
  1007. "--password", admin_client.password,
  1008. "get-user",
  1009. "-u", user["username"],
  1010. ], check=True)
  1011. self.assertIn("Username:", p.stdout)
  1012. self.assertIn("Permissions:", p.stdout)
  1013. p = self.run_hashclient([
  1014. "--address", self.auth_server_address,
  1015. "--login", user["username"],
  1016. "--password", user["token"],
  1017. "get-user",
  1018. ], check=True)
  1019. self.assertIn("Username:", p.stdout)
  1020. self.assertIn("Permissions:", p.stdout)
  1021. def test_get_all_users(self):
  1022. admin_client = self.start_auth_server()
  1023. admin_client.new_user("test-user1", ["@read"])
  1024. admin_client.new_user("test-user2", ["@read"])
  1025. p = self.run_hashclient([
  1026. "--address", self.auth_server_address,
  1027. "--login", admin_client.username,
  1028. "--password", admin_client.password,
  1029. "get-all-users",
  1030. ], check=True)
  1031. self.assertIn("admin", p.stdout)
  1032. self.assertIn("test-user1", p.stdout)
  1033. self.assertIn("test-user2", p.stdout)
  1034. def test_new_user(self):
  1035. admin_client = self.start_auth_server()
  1036. p = self.run_hashclient([
  1037. "--address", self.auth_server_address,
  1038. "--login", admin_client.username,
  1039. "--password", admin_client.password,
  1040. "new-user",
  1041. "-u", "test-user",
  1042. "@read", "@report",
  1043. ], check=True)
  1044. new_token = None
  1045. for l in p.stdout.splitlines():
  1046. l = l.rstrip()
  1047. m = re.match(r'Token: +(.*)$', l)
  1048. if m is not None:
  1049. new_token = m.group(1)
  1050. self.assertTrue(new_token)
  1051. user = {
  1052. "username": "test-user",
  1053. "token": new_token,
  1054. }
  1055. self.assertUserPerms(user, ["@read", "@report"])
  1056. def test_delete_user(self):
  1057. admin_client = self.start_auth_server()
  1058. user = admin_client.new_user("test-user", ["@read"])
  1059. p = self.run_hashclient([
  1060. "--address", self.auth_server_address,
  1061. "--login", admin_client.username,
  1062. "--password", admin_client.password,
  1063. "delete-user",
  1064. "-u", user["username"],
  1065. ], check=True)
  1066. self.assertIsNone(admin_client.get_user(user["username"]))
  1067. def test_get_db_usage(self):
  1068. p = self.run_hashclient([
  1069. "--address", self.server_address,
  1070. "get-db-usage",
  1071. ], check=True)
  1072. def test_get_db_query_columns(self):
  1073. p = self.run_hashclient([
  1074. "--address", self.server_address,
  1075. "get-db-query-columns",
  1076. ], check=True)
  1077. def test_gc(self):
  1078. taskhash = '53b8dce672cb6d0c73170be43f540460bfc347b4'
  1079. outhash = '5a9cb1649625f0bf41fc7791b635cd9c2d7118c7f021ba87dcd03f72b67ce7a8'
  1080. unihash = 'f37918cc02eb5a520b1aff86faacbc0a38124646'
  1081. result = self.client.report_unihash(taskhash, self.METHOD, outhash, unihash)
  1082. self.assertEqual(result['unihash'], unihash, 'Server returned bad unihash')
  1083. taskhash2 = '3bf6f1e89d26205aec90da04854fbdbf73afe6b4'
  1084. outhash2 = '77623a549b5b1a31e3732dfa8fe61d7ce5d44b3370f253c5360e136b852967b4'
  1085. unihash2 = 'af36b199320e611fbb16f1f277d3ee1d619ca58b'
  1086. result = self.client.report_unihash(taskhash2, self.METHOD, outhash2, unihash2)
  1087. self.assertClientGetHash(self.client, taskhash2, unihash2)
  1088. # Mark the first unihash to be kept
  1089. self.run_hashclient([
  1090. "--address", self.server_address,
  1091. "gc-mark", "ABC",
  1092. "--where", "unihash", unihash,
  1093. "--where", "method", self.METHOD
  1094. ], check=True)
  1095. # Second hash is still there; mark doesn't delete hashes
  1096. self.assertClientGetHash(self.client, taskhash2, unihash2)
  1097. self.run_hashclient([
  1098. "--address", self.server_address,
  1099. "gc-sweep", "ABC",
  1100. ], check=True)
  1101. # Hash is gone. Taskhash is returned for second hash
  1102. self.assertClientGetHash(self.client, taskhash2, None)
  1103. # First hash is still present
  1104. self.assertClientGetHash(self.client, taskhash, unihash)
  1105. class TestHashEquivalenceUnixServer(HashEquivalenceTestSetup, HashEquivalenceCommonTests, unittest.TestCase):
  1106. def get_server_addr(self, server_idx):
  1107. return "unix://" + os.path.join(self.temp_dir.name, 'sock%d' % server_idx)
  1108. class TestHashEquivalenceUnixServerLongPath(HashEquivalenceTestSetup, unittest.TestCase):
  1109. DEEP_DIRECTORY = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa/bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb/ccccccccccccccccccccccccccccccccccccccccccc"
  1110. def get_server_addr(self, server_idx):
  1111. os.makedirs(os.path.join(self.temp_dir.name, self.DEEP_DIRECTORY), exist_ok=True)
  1112. return "unix://" + os.path.join(self.temp_dir.name, self.DEEP_DIRECTORY, 'sock%d' % server_idx)
  1113. def test_long_sock_path(self):
  1114. # Simple test that hashes can be created
  1115. taskhash = '35788efcb8dfb0a02659d81cf2bfd695fb30faf9'
  1116. outhash = '2765d4a5884be49b28601445c2760c5f21e7e5c0ee2b7e3fce98fd7e5970796f'
  1117. unihash = 'f46d3fbb439bd9b921095da657a4de906510d2cd'
  1118. self.assertClientGetHash(self.client, taskhash, None)
  1119. result = self.client.report_unihash(taskhash, self.METHOD, outhash, unihash)
  1120. self.assertEqual(result['unihash'], unihash, 'Server returned bad unihash')
  1121. class TestHashEquivalenceTCPServer(HashEquivalenceTestSetup, HashEquivalenceCommonTests, unittest.TestCase):
  1122. def get_server_addr(self, server_idx):
  1123. # Some hosts cause asyncio module to misbehave, when IPv6 is not enabled.
  1124. # If IPv6 is enabled, it should be safe to use localhost directly, in general
  1125. # case it is more reliable to resolve the IP address explicitly.
  1126. return socket.gethostbyname("localhost") + ":0"
  1127. class TestHashEquivalenceWebsocketServer(HashEquivalenceTestSetup, HashEquivalenceCommonTests, unittest.TestCase):
  1128. def setUp(self):
  1129. try:
  1130. import websockets
  1131. except ImportError as e:
  1132. self.skipTest(str(e))
  1133. super().setUp()
  1134. def get_server_addr(self, server_idx):
  1135. # Some hosts cause asyncio module to misbehave, when IPv6 is not enabled.
  1136. # If IPv6 is enabled, it should be safe to use localhost directly, in general
  1137. # case it is more reliable to resolve the IP address explicitly.
  1138. host = socket.gethostbyname("localhost")
  1139. return "ws://%s:0" % host
  1140. class TestHashEquivalenceWebsocketsSQLAlchemyServer(TestHashEquivalenceWebsocketServer):
  1141. def setUp(self):
  1142. try:
  1143. import sqlalchemy
  1144. import aiosqlite
  1145. except ImportError as e:
  1146. self.skipTest(str(e))
  1147. super().setUp()
  1148. def make_dbpath(self):
  1149. return "sqlite+aiosqlite:///%s" % os.path.join(self.temp_dir.name, "db%d.sqlite" % self.server_index)
  1150. class TestHashEquivalenceExternalServer(HashEquivalenceTestSetup, HashEquivalenceCommonTests, unittest.TestCase):
  1151. def get_env(self, name):
  1152. v = os.environ.get(name)
  1153. if not v:
  1154. self.skipTest(f'{name} not defined to test an external server')
  1155. return v
  1156. def start_test_server(self):
  1157. return self.get_env('BB_TEST_HASHSERV')
  1158. def start_server(self, *args, **kwargs):
  1159. self.skipTest('Cannot start local server when testing external servers')
  1160. def start_auth_server(self):
  1161. self.auth_server_address = self.server_address
  1162. self.admin_client = self.start_client(
  1163. self.server_address,
  1164. username=self.get_env('BB_TEST_HASHSERV_USERNAME'),
  1165. password=self.get_env('BB_TEST_HASHSERV_PASSWORD'),
  1166. )
  1167. return self.admin_client
  1168. def setUp(self):
  1169. super().setUp()
  1170. if "BB_TEST_HASHSERV_USERNAME" in os.environ:
  1171. self.client = self.start_client(
  1172. self.server_address,
  1173. username=os.environ["BB_TEST_HASHSERV_USERNAME"],
  1174. password=os.environ["BB_TEST_HASHSERV_PASSWORD"],
  1175. )
  1176. self.client.remove({"method": self.METHOD})
  1177. def tearDown(self):
  1178. self.client.remove({"method": self.METHOD})
  1179. super().tearDown()
  1180. def test_auth_get_all_users(self):
  1181. self.skipTest("Cannot test all users with external server")