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