path: root/lib/hashserv/tests.py

#! /usr/bin/env python3
#
# Copyright (C) 2018-2019 Garmin Ltd.
#
# SPDX-License-Identifier: GPL-2.0-only
#

from . import create_server, create_client
from .server import DEFAULT_ANON_PERMS, ALL_PERMISSIONS
from bb.asyncrpc import InvokeError
import hashlib
import logging
import multiprocessing
import os
import sys
import tempfile
import threading
import unittest
import socket
import time
import signal

def server_prefunc(server, idx):
    logging.basicConfig(level=logging.DEBUG, filename='bbhashserv-%d.log' % idx, filemode='w',
                        format='%(levelname)s %(filename)s:%(lineno)d %(message)s')
    server.logger.debug("Running server %d" % idx)
    sys.stdout = open('bbhashserv-stdout-%d.log' % idx, 'w')
    sys.stderr = sys.stdout

class HashEquivalenceTestSetup(object):
    METHOD = 'TestMethod'

    server_index = 0
    client_index = 0

    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):
        self.server_index += 1
        if dbpath is None:
            dbpath = self.make_dbpath()

        def cleanup_server(server):
            if server.process.exitcode is not None:
                return

            server.process.terminate()
            server.process.join()

        server = create_server(self.get_server_addr(self.server_index),
                               dbpath,
                               upstream=upstream,
                               read_only=read_only,
                               anon_perms=anon_perms,
                               admin_username=admin_username,
                               admin_password=admin_password)
        server.dbpath = dbpath

        server.serve_as_process(prefunc=prefunc, args=(self.server_index,))
        self.addCleanup(cleanup_server, server)

        return server

    def make_dbpath(self):
        return os.path.join(self.temp_dir.name, "db%d.sqlite" % self.server_index)

    def start_client(self, server_address, username=None, password=None):
        def cleanup_client(client):
            client.close()

        client = create_client(server_address, username=username, password=password)
        self.addCleanup(cleanup_client, client)

        return client

    def start_test_server(self):
        self.server = self.start_server()
        return self.server.address

    def start_auth_server(self):
        self.auth_server = self.start_server(self.server.dbpath, anon_perms=[], admin_username="admin", admin_password="password")
        self.admin_client = self.start_client(self.auth_server.address, username="admin", password="password")
        return self.admin_client

    def auth_client(self, user):
        return self.start_client(self.auth_server.address, user["username"], user["token"])

    def auth_perms(self, *permissions):
        self.client_index += 1
        user = self.admin_client.new_user(f"user-{self.client_index}", permissions)
        return self.auth_client(user)

    def setUp(self):
        if sys.version_info < (3, 5, 0):
            self.skipTest('Python 3.5 or later required')

        self.temp_dir = tempfile.TemporaryDirectory(prefix='bb-hashserv')
        self.addCleanup(self.temp_dir.cleanup)

        self.server_address = self.start_test_server()

        self.client = self.start_client(self.server_address)

    def assertClientGetHash(self, client, taskhash, unihash):
        result = client.get_unihash(self.METHOD, taskhash)
        self.assertEqual(result, unihash)


class HashEquivalenceCommonTests(object):
    def create_test_hash(self, client):
        # Simple test that hashes can be created
        taskhash = '35788efcb8dfb0a02659d81cf2bfd695fb30faf9'
        outhash = '2765d4a5884be49b28601445c2760c5f21e7e5c0ee2b7e3fce98fd7e5970796f'
        unihash = 'f46d3fbb439bd9b921095da657a4de906510d2cd'

        self.assertClientGetHash(client, taskhash, None)

        result = client.report_unihash(taskhash, self.METHOD, outhash, unihash)
        self.assertEqual(result['unihash'], unihash, 'Server returned bad unihash')
        return taskhash, outhash, unihash

    def test_create_hash(self):
        return self.create_test_hash(self.client)

    def test_create_equivalent(self):
        # Tests that a second reported task with the same outhash will be
        # assigned the same unihash
        taskhash = '53b8dce672cb6d0c73170be43f540460bfc347b4'
        outhash = '5a9cb1649625f0bf41fc7791b635cd9c2d7118c7f021ba87dcd03f72b67ce7a8'
        unihash = 'f37918cc02eb5a520b1aff86faacbc0a38124646'

        result = self.client.report_unihash(taskhash, self.METHOD, outhash, unihash)
        self.assertEqual(result['unihash'], unihash, 'Server returned bad unihash')

        # Report a different task with the same outhash. The returned unihash
        # should match the first task
        taskhash2 = '3bf6f1e89d26205aec90da04854fbdbf73afe6b4'
        unihash2 = 'af36b199320e611fbb16f1f277d3ee1d619ca58b'
        result = self.client.report_unihash(taskhash2, self.METHOD, outhash, unihash2)
        self.assertEqual(result['unihash'], unihash, 'Server returned bad unihash')

    def test_duplicate_taskhash(self):
        # Tests that duplicate reports of the same taskhash with different
        # outhash & unihash always return the unihash from the first reported
        # taskhash
        taskhash = '8aa96fcffb5831b3c2c0cb75f0431e3f8b20554a'
        outhash = 'afe240a439959ce86f5e322f8c208e1fedefea9e813f2140c81af866cc9edf7e'
        unihash = '218e57509998197d570e2c98512d0105985dffc9'
        self.client.report_unihash(taskhash, self.METHOD, outhash, unihash)

        self.assertClientGetHash(self.client, taskhash, unihash)

        outhash2 = '0904a7fe3dc712d9fd8a74a616ddca2a825a8ee97adf0bd3fc86082c7639914d'
        unihash2 = 'ae9a7d252735f0dafcdb10e2e02561ca3a47314c'
        self.client.report_unihash(taskhash, self.METHOD, outhash2, unihash2)

        self.assertClientGetHash(self.client, taskhash, unihash)

        outhash3 = '77623a549b5b1a31e3732dfa8fe61d7ce5d44b3370f253c5360e136b852967b4'
        unihash3 = '9217a7d6398518e5dc002ed58f2cbbbc78696603'
        self.client.report_unihash(taskhash, self.METHOD, outhash3, unihash3)

        self.assertClientGetHash(self.client, taskhash, unihash)

    def test_remove_taskhash(self):
        taskhash, outhash, unihash = self.test_create_hash()
        result = self.client.remove({"taskhash": taskhash})
        self.assertGreater(result["count"], 0)
        self.assertClientGetHash(self.client, taskhash, None)

        result_outhash = self.client.get_outhash(self.METHOD, outhash, taskhash)
        self.assertIsNone(result_outhash)

    def test_remove_unihash(self):
        taskhash, outhash, unihash = self.test_create_hash()
        result = self.client.remove({"unihash": unihash})
        self.assertGreater(result["count"], 0)
        self.assertClientGetHash(self.client, taskhash, None)

    def test_remove_outhash(self):
        taskhash, outhash, unihash = self.test_create_hash()
        result = self.client.remove({"outhash": outhash})
        self.assertGreater(result["count"], 0)

        result_outhash = self.client.get_outhash(self.METHOD, outhash, taskhash)
        self.assertIsNone(result_outhash)

    def test_remove_method(self):
        taskhash, outhash, unihash = self.test_create_hash()
        result = self.client.remove({"method": self.METHOD})
        self.assertGreater(result["count"], 0)
        self.assertClientGetHash(self.client, taskhash, None)

        result_outhash = self.client.get_outhash(self.METHOD, outhash, taskhash)
        self.assertIsNone(result_outhash)

    def test_clean_unused(self):
        taskhash, outhash, unihash = self.test_create_hash()

        # Clean the database, which should not remove anything because all hashes an in-use
        result = self.client.clean_unused(0)
        self.assertEqual(result["count"], 0)
        self.assertClientGetHash(self.client, taskhash, unihash)

        # Remove the unihash. The row in the outhash table should still be present
        self.client.remove({"unihash": unihash})
        result_outhash = self.client.get_outhash(self.METHOD, outhash, taskhash, False)
        self.assertIsNotNone(result_outhash)

        # Now clean with no minimum age which will remove the outhash
        result = self.client.clean_unused(0)
        self.assertEqual(result["count"], 1)
        result_outhash = self.client.get_outhash(self.METHOD, outhash, taskhash, False)
        self.assertIsNone(result_outhash)

    def test_huge_message(self):
        # Simple test that hashes can be created
        taskhash = 'c665584ee6817aa99edfc77a44dd853828279370'
        outhash = '3c979c3db45c569f51ab7626a4651074be3a9d11a84b1db076f5b14f7d39db44'
        unihash = '90e9bc1d1f094c51824adca7f8ea79a048d68824'

        self.assertClientGetHash(self.client, taskhash, None)

        siginfo = "0" * (self.client.max_chunk * 4)

        result = self.client.report_unihash(taskhash, self.METHOD, outhash, unihash, {
            'outhash_siginfo': siginfo
        })
        self.assertEqual(result['unihash'], unihash, 'Server returned bad unihash')

        result_unihash = self.client.get_taskhash(self.METHOD, taskhash, True)
        self.assertEqual(result_unihash['taskhash'], taskhash)
        self.assertEqual(result_unihash['unihash'], unihash)
        self.assertEqual(result_unihash['method'], self.METHOD)

        result_outhash = self.client.get_outhash(self.METHOD, outhash, taskhash)
        self.assertEqual(result_outhash['taskhash'], taskhash)
        self.assertEqual(result_outhash['method'], self.METHOD)
        self.assertEqual(result_outhash['unihash'], unihash)
        self.assertEqual(result_outhash['outhash'], outhash)
        self.assertEqual(result_outhash['outhash_siginfo'], siginfo)

    def test_stress(self):
        def query_server(failures):
            client = Client(self.server_address)
            try:
                for i in range(1000):
                    taskhash = hashlib.sha256()
                    taskhash.update(str(i).encode('utf-8'))
                    taskhash = taskhash.hexdigest()
                    result = client.get_unihash(self.METHOD, taskhash)
                    if result != taskhash:
                        failures.append("taskhash mismatch: %s != %s" % (result, taskhash))
            finally:
                client.close()

        # Report hashes
        for i in range(1000):
            taskhash = hashlib.sha256()
            taskhash.update(str(i).encode('utf-8'))
            taskhash = taskhash.hexdigest()
            self.client.report_unihash(taskhash, self.METHOD, taskhash, taskhash)

        failures = []
        threads = [threading.Thread(target=query_server, args=(failures,)) for t in range(100)]

        for t in threads:
            t.start()

        for t in threads:
            t.join()

        self.assertFalse(failures)

    def test_upstream_server(self):
        # Tests upstream server support. This is done by creating two servers
        # that share a database file. The downstream server has it upstream
        # set to the test server, whereas the side server doesn't. This allows
        # verification that the hash requests are being proxied to the upstream
        # server by verifying that they appear on the downstream client, but not
        # the side client. It also verifies that the results are pulled into
        # the downstream database by checking that the downstream and side servers
        # match after the downstream is done waiting for all backfill tasks
        down_server = self.start_server(upstream=self.server_address)
        down_client = self.start_client(down_server.address)
        side_server = self.start_server(dbpath=down_server.dbpath)
        side_client = self.start_client(side_server.address)

        def check_hash(taskhash, unihash, old_sidehash):
            nonlocal down_client
            nonlocal side_client

            # check upstream server
            self.assertClientGetHash(self.client, taskhash, unihash)

            # Hash should *not* be present on the side server
            self.assertClientGetHash(side_client, taskhash, old_sidehash)

            # Hash should be present on the downstream server, since it
            # will defer to the upstream server. This will trigger
            # the backfill in the downstream server
            self.assertClientGetHash(down_client, taskhash, unihash)

            # After waiting for the downstream client to finish backfilling the
            # task from the upstream server, it should appear in the side server
            # since the database is populated
            down_client.backfill_wait()
            self.assertClientGetHash(side_client, taskhash, unihash)

        # Basic report
        taskhash = '8aa96fcffb5831b3c2c0cb75f0431e3f8b20554a'
        outhash = 'afe240a439959ce86f5e322f8c208e1fedefea9e813f2140c81af866cc9edf7e'
        unihash = '218e57509998197d570e2c98512d0105985dffc9'
        self.client.report_unihash(taskhash, self.METHOD, outhash, unihash)

        check_hash(taskhash, unihash, None)

        # Duplicated taskhash with multiple output hashes and unihashes.
        # All servers should agree with the originally reported hash
        outhash2 = '0904a7fe3dc712d9fd8a74a616ddca2a825a8ee97adf0bd3fc86082c7639914d'
        unihash2 = 'ae9a7d252735f0dafcdb10e2e02561ca3a47314c'
        self.client.report_unihash(taskhash, self.METHOD, outhash2, unihash2)

        check_hash(taskhash, unihash, unihash)

        # Report an equivalent task. The sideload will originally report
        # no unihash until backfilled
        taskhash3 = "044c2ec8aaf480685a00ff6ff49e6162e6ad34e1"
        unihash3 = "def64766090d28f627e816454ed46894bb3aab36"
        self.client.report_unihash(taskhash3, self.METHOD, outhash, unihash3)

        check_hash(taskhash3, unihash, None)

        # Test that reporting a unihash in the downstream client isn't
        # propagating to the upstream server
        taskhash4 = "e3da00593d6a7fb435c7e2114976c59c5fd6d561"
        outhash4 = "1cf8713e645f491eb9c959d20b5cae1c47133a292626dda9b10709857cbe688a"
        unihash4 = "3b5d3d83f07f259e9086fcb422c855286e18a57d"
        down_client.report_unihash(taskhash4, self.METHOD, outhash4, unihash4)
        down_client.backfill_wait()

        self.assertClientGetHash(down_client, taskhash4, unihash4)
        self.assertClientGetHash(side_client, taskhash4, unihash4)
        self.assertClientGetHash(self.client, taskhash4, None)

        # Test that reporting a unihash in the downstream is able to find a
        # match which was previously reported to the upstream server
        taskhash5 = '35788efcb8dfb0a02659d81cf2bfd695fb30faf9'
        outhash5 = '2765d4a5884be49b28601445c2760c5f21e7e5c0ee2b7e3fce98fd7e5970796f'
        unihash5 = 'f46d3fbb439bd9b921095da657a4de906510d2cd'
        result = self.client.report_unihash(taskhash5, self.METHOD, outhash5, unihash5)

        taskhash6 = '35788efcb8dfb0a02659d81cf2bfd695fb30fafa'
        unihash6 = 'f46d3fbb439bd9b921095da657a4de906510d2ce'
        result = down_client.report_unihash(taskhash6, self.METHOD, outhash5, unihash6)
        self.assertEqual(result['unihash'], unihash5, 'Server failed to copy unihash from upstream')

        # Tests read through from server with
        taskhash7 = '9d81d76242cc7cfaf7bf74b94b9cd2e29324ed74'
        outhash7 = '8470d56547eea6236d7c81a644ce74670ca0bbda998e13c629ef6bb3f0d60b69'
        unihash7 = '05d2a63c81e32f0a36542ca677e8ad852365c538'
        self.client.report_unihash(taskhash7, self.METHOD, outhash7, unihash7)

        result = down_client.get_taskhash(self.METHOD, taskhash7, True)
        self.assertEqual(result['unihash'], unihash7, 'Server failed to copy unihash from upstream')
        self.assertEqual(result['outhash'], outhash7, 'Server failed to copy unihash from upstream')
        self.assertEqual(result['taskhash'], taskhash7, 'Server failed to copy unihash from upstream')
        self.assertEqual(result['method'], self.METHOD)

        taskhash8 = '86978a4c8c71b9b487330b0152aade10c1ee58aa'
        outhash8 = 'ca8c128e9d9e4a28ef24d0508aa20b5cf880604eacd8f65c0e366f7e0cc5fbcf'
        unihash8 = 'd8bcf25369d40590ad7d08c84d538982f2023e01'
        self.client.report_unihash(taskhash8, self.METHOD, outhash8, unihash8)

        result = down_client.get_outhash(self.METHOD, outhash8, taskhash8)
        self.assertEqual(result['unihash'], unihash8, 'Server failed to copy unihash from upstream')
        self.assertEqual(result['outhash'], outhash8, 'Server failed to copy unihash from upstream')
        self.assertEqual(result['taskhash'], taskhash8, 'Server failed to copy unihash from upstream')
        self.assertEqual(result['method'], self.METHOD)

        taskhash9 = 'ae6339531895ddf5b67e663e6a374ad8ec71d81c'
        outhash9 = 'afc78172c81880ae10a1fec994b5b4ee33d196a001a1b66212a15ebe573e00b5'
        unihash9 = '6662e699d6e3d894b24408ff9a4031ef9b038ee8'
        self.client.report_unihash(taskhash9, self.METHOD, outhash9, unihash9)

        result = down_client.get_taskhash(self.METHOD, taskhash9, False)
        self.assertEqual(result['unihash'], unihash9, 'Server failed to copy unihash from upstream')
        self.assertEqual(result['taskhash'], taskhash9, 'Server failed to copy unihash from upstream')
        self.assertEqual(result['method'], self.METHOD)

    def test_ro_server(self):
        rw_server = self.start_server()
        rw_client = self.start_client(rw_server.address)

        ro_server = self.start_server(dbpath=rw_server.dbpath, read_only=True)
        ro_client = self.start_client(ro_server.address)

        # Report a hash via the read-write server
        taskhash = '35788efcb8dfb0a02659d81cf2bfd695fb30faf9'
        outhash = '2765d4a5884be49b28601445c2760c5f21e7e5c0ee2b7e3fce98fd7e5970796f'
        unihash = 'f46d3fbb439bd9b921095da657a4de906510d2cd'

        result = rw_client.report_unihash(taskhash, self.METHOD, outhash, unihash)
        self.assertEqual(result['unihash'], unihash, 'Server returned bad unihash')

        # Check the hash via the read-only server
        self.assertClientGetHash(ro_client, taskhash, unihash)

        # Ensure that reporting via the read-only server fails
        taskhash2 = 'c665584ee6817aa99edfc77a44dd853828279370'
        outhash2 = '3c979c3db45c569f51ab7626a4651074be3a9d11a84b1db076f5b14f7d39db44'
        unihash2 = '90e9bc1d1f094c51824adca7f8ea79a048d68824'

        result = ro_client.report_unihash(taskhash2, self.METHOD, outhash2, unihash2)
        self.assertEqual(result['unihash'], unihash2)

        # Ensure that the database was not modified
        self.assertClientGetHash(rw_client, taskhash2, None)


    def test_slow_server_start(self):
        # Ensures that the server will exit correctly even if it gets a SIGTERM
        # before entering the main loop

        event = multiprocessing.Event()

        def prefunc(server, idx):
            nonlocal event
            server_prefunc(server, idx)
            event.wait()

        def do_nothing(signum, frame):
            pass

        old_signal = signal.signal(signal.SIGTERM, do_nothing)
        self.addCleanup(signal.signal, signal.SIGTERM, old_signal)

        server = self.start_server(prefunc=prefunc)
        server.process.terminate()
        time.sleep(30)
        event.set()
        server.process.join(300)
        self.assertIsNotNone(server.process.exitcode, "Server did not exit in a timely manner!")

    def test_diverging_report_race(self):
        # Tests that a reported task will correctly pick up an updated unihash

        # This is a baseline report added to the database to ensure that there
        # is something to match against as equivalent
        outhash1 = 'afd11c366050bcd75ad763e898e4430e2a60659b26f83fbb22201a60672019fa'
        taskhash1 = '3bde230c743fc45ab61a065d7a1815fbfa01c4740e4c895af2eb8dc0f684a4ab'
        unihash1 = '3bde230c743fc45ab61a065d7a1815fbfa01c4740e4c895af2eb8dc0f684a4ab'
        result = self.client.report_unihash(taskhash1, self.METHOD, outhash1, unihash1)

        # Add a report that is equivalent to Task 1. It should ignore the
        # provided unihash and report the unihash from task 1
        taskhash2 = '6259ae8263bd94d454c086f501c37e64c4e83cae806902ca95b4ab513546b273'
        unihash2 = taskhash2
        result = self.client.report_unihash(taskhash2, self.METHOD, outhash1, unihash2)
        self.assertEqual(result['unihash'], unihash1)

        # Add another report for Task 2, but with a different outhash (e.g. the
        # task is non-deterministic). It should still be marked with the Task 1
        # unihash because it has the Task 2 taskhash, which is equivalent to
        # Task 1
        outhash3 = 'd2187ee3a8966db10b34fe0e863482288d9a6185cb8ef58a6c1c6ace87a2f24c'
        result = self.client.report_unihash(taskhash2, self.METHOD, outhash3, unihash2)
        self.assertEqual(result['unihash'], unihash1)


    def test_diverging_report_reverse_race(self):
        # Same idea as the previous test, but Tasks 2 and 3 are reported in
        # reverse order the opposite order

        outhash1 = 'afd11c366050bcd75ad763e898e4430e2a60659b26f83fbb22201a60672019fa'
        taskhash1 = '3bde230c743fc45ab61a065d7a1815fbfa01c4740e4c895af2eb8dc0f684a4ab'
        unihash1 = '3bde230c743fc45ab61a065d7a1815fbfa01c4740e4c895af2eb8dc0f684a4ab'
        result = self.client.report_unihash(taskhash1, self.METHOD, outhash1, unihash1)

        taskhash2 = '6259ae8263bd94d454c086f501c37e64c4e83cae806902ca95b4ab513546b273'
        unihash2 = taskhash2

        # Report Task 3 first. Since there is nothing else in the database it
        # will use the client provided unihash
        outhash3 = 'd2187ee3a8966db10b34fe0e863482288d9a6185cb8ef58a6c1c6ace87a2f24c'
        result = self.client.report_unihash(taskhash2, self.METHOD, outhash3, unihash2)
        self.assertEqual(result['unihash'], unihash2)

        # Report Task 2. This is equivalent to Task 1 but there is already a mapping for
        # taskhash2 so it will report unihash2
        result = self.client.report_unihash(taskhash2, self.METHOD, outhash1, unihash2)
        self.assertEqual(result['unihash'], unihash2)

        # The originally reported unihash for Task 3 should be unchanged even if it
        # shares a taskhash with Task 2
        self.assertClientGetHash(self.client, taskhash2, unihash2)

    def test_auth_read_perms(self):
        admin_client = self.start_auth_server()

        # Create hashes with non-authenticated server
        taskhash, outhash, unihash = self.test_create_hash()

        # Validate hash can be retrieved using authenticated client
        with self.auth_perms("@read") as client:
            self.assertClientGetHash(client, taskhash, unihash)

        with self.auth_perms() as client, self.assertRaises(InvokeError):
            self.assertClientGetHash(client, taskhash, unihash)

    def test_auth_report_perms(self):
        admin_client = self.start_auth_server()

        # Without read permission, the user is completely denied
        with self.auth_perms() as client, self.assertRaises(InvokeError):
            self.create_test_hash(client)

        # Read permission allows the call to succeed, but it doesn't record
        # anythin in the database
        with self.auth_perms("@read") as client:
            taskhash, outhash, unihash = self.create_test_hash(client)
            self.assertClientGetHash(client, taskhash, None)

        # Report permission alone is insufficient
        with self.auth_perms("@report") as client, self.assertRaises(InvokeError):
            self.create_test_hash(client)

        # Read and report permission actually modify the database
        with self.auth_perms("@read", "@report") as client:
            taskhash, outhash, unihash = self.create_test_hash(client)
            self.assertClientGetHash(client, taskhash, unihash)

    def test_auth_no_token_refresh_from_anon_user(self):
        self.start_auth_server()

        with self.start_client(self.auth_server.address) as client, self.assertRaises(InvokeError):
            client.refresh_token()

    def assertUserCanAuth(self, user):
        with self.start_client(self.auth_server.address) as client:
            client.auth(user["username"], user["token"])

    def assertUserCannotAuth(self, user):
        with self.start_client(self.auth_server.address) as client, self.assertRaises(InvokeError):
            client.auth(user["username"], user["token"])

    def test_auth_self_token_refresh(self):
        admin_client = self.start_auth_server()

        # Create a new user with no permissions
        user = admin_client.new_user("test-user", [])

        with self.auth_client(user) as client:
            new_user = client.refresh_token()

        self.assertEqual(user["username"], new_user["username"])
        self.assertNotEqual(user["token"], new_user["token"])
        self.assertUserCanAuth(new_user)
        self.assertUserCannotAuth(user)

        # Explicitly specifying with your own username is fine also
        with self.auth_client(new_user) as client:
            new_user2 = client.refresh_token(user["username"])

        self.assertEqual(user["username"], new_user2["username"])
        self.assertNotEqual(user["token"], new_user2["token"])
        self.assertUserCanAuth(new_user2)
        self.assertUserCannotAuth(new_user)
        self.assertUserCannotAuth(user)

    def test_auth_token_refresh(self):
        admin_client = self.start_auth_server()

        user = admin_client.new_user("test-user", [])

        with self.auth_perms() as client, self.assertRaises(InvokeError):
            client.refresh_token(user["username"])

        with self.auth_perms("@user-admin") as client:
            new_user = client.refresh_token(user["username"])

        self.assertEqual(user["username"], new_user["username"])
        self.assertNotEqual(user["token"], new_user["token"])
        self.assertUserCanAuth(new_user)
        self.assertUserCannotAuth(user)

    def test_auth_self_get_user(self):
        admin_client = self.start_auth_server()

        user = admin_client.new_user("test-user", [])
        user_info = user.copy()
        del user_info["token"]

        with self.auth_client(user) as client:
            info = client.get_user()
            self.assertEqual(info, user_info)

            # Explicitly asking for your own username is fine also
            info = client.get_user(user["username"])
            self.assertEqual(info, user_info)

    def test_auth_get_user(self):
        admin_client = self.start_auth_server()

        user = admin_client.new_user("test-user", [])
        user_info = user.copy()
        del user_info["token"]

        with self.auth_perms() as client, self.assertRaises(InvokeError):
            client.get_user(user["username"])

        with self.auth_perms("@user-admin") as client:
            info = client.get_user(user["username"])
            self.assertEqual(info, user_info)

            info = client.get_user("nonexist-user")
            self.assertIsNone(info)

    def test_auth_reconnect(self):
        admin_client = self.start_auth_server()

        user = admin_client.new_user("test-user", [])
        user_info = user.copy()
        del user_info["token"]

        with self.auth_client(user) as client:
            info = client.get_user()
            self.assertEqual(info, user_info)

            client.disconnect()

            info = client.get_user()
            self.assertEqual(info, user_info)

    def test_auth_delete_user(self):
        admin_client = self.start_auth_server()

        user = admin_client.new_user("test-user", [])

        # No self service
        with self.auth_client(user) as client, self.assertRaises(InvokeError):
            client.delete_user(user["username"])

        with self.auth_perms() as client, self.assertRaises(InvokeError):
            client.delete_user(user["username"])

        with self.auth_perms("@user-admin") as client:
            client.delete_user(user["username"])

        # User doesn't exist, so even though the permission is correct, it's an
        # error
        with self.auth_perms("@user-admin") as client, self.assertRaises(InvokeError):
            client.delete_user(user["username"])

    def assertUserPerms(self, user, permissions):
        with self.auth_client(user) as client:
            info = client.get_user()
            self.assertEqual(info, {
                "username": user["username"],
                "permissions": permissions,
            })

    def test_auth_set_user_perms(self):
        admin_client = self.start_auth_server()

        user = admin_client.new_user("test-user", [])

        self.assertUserPerms(user, [])

        # No self service to change permissions
        with self.auth_client(user) as client, self.assertRaises(InvokeError):
            client.set_user_perms(user["username"], ["@all"])
        self.assertUserPerms(user, [])

        with self.auth_perms() as client, self.assertRaises(InvokeError):
            client.set_user_perms(user["username"], ["@all"])
        self.assertUserPerms(user, [])

        with self.auth_perms("@user-admin") as client:
            client.set_user_perms(user["username"], ["@all"])
        self.assertUserPerms(user, sorted(list(ALL_PERMISSIONS)))

        # Bad permissions
        with self.auth_perms("@user-admin") as client, self.assertRaises(InvokeError):
            client.set_user_perms(user["username"], ["@this-is-not-a-permission"])
        self.assertUserPerms(user, sorted(list(ALL_PERMISSIONS)))

    def test_auth_get_all_users(self):
        admin_client = self.start_auth_server()

        user = admin_client.new_user("test-user", [])

        with self.auth_client(user) as client, self.assertRaises(InvokeError):
            client.get_all_users()

        # Give the test user the correct permission
        admin_client.set_user_perms(user["username"], ["@user-admin"])

        with self.auth_client(user) as client:
            all_users = client.get_all_users()

        # Convert to a dictionary for easier comparison
        all_users = {u["username"]: u for u in all_users}

        self.assertEqual(all_users,
            {
                "admin": {
                    "username": "admin",
                    "permissions": sorted(list(ALL_PERMISSIONS)),
                },
                "test-user": {
                    "username": "test-user",
                    "permissions": ["@user-admin"],
                }
            }
        )

    def test_auth_new_user(self):
        self.start_auth_server()

        permissions = ["@read", "@report", "@db-admin", "@user-admin"]
        permissions.sort()

        with self.auth_perms() as client, self.assertRaises(InvokeError):
            client.new_user("test-user", permissions)

        with self.auth_perms("@user-admin") as client:
            user = client.new_user("test-user", permissions)
            self.assertIn("token", user)
            self.assertEqual(user["username"], "test-user")
            self.assertEqual(user["permissions"], permissions)

    def test_auth_become_user(self):
        admin_client = self.start_auth_server()

        user = admin_client.new_user("test-user", ["@read", "@report"])
        user_info = user.copy()
        del user_info["token"]

        with self.auth_perms() as client, self.assertRaises(InvokeError):
            client.become_user(user["username"])

        with self.auth_perms("@user-admin") as client:
            become = client.become_user(user["username"])
            self.assertEqual(become, user_info)

            info = client.get_user()
            self.assertEqual(info, user_info)

            # Verify become user is preserved across disconnect
            client.disconnect()

            info = client.get_user()
            self.assertEqual(info, user_info)

            # test-user doesn't have become_user permissions, so this should
            # not work
            with self.assertRaises(InvokeError):
                client.become_user(user["username"])

        # No self-service of become
        with self.auth_client(user) as client, self.assertRaises(InvokeError):
            client.become_user(user["username"])

        # Give test user permissions to become
        admin_client.set_user_perms(user["username"], ["@user-admin"])

        # It's possible to become yourself (effectively a noop)
        with self.auth_perms("@user-admin") as client:
            become = client.become_user(client.username)

    def test_get_db_usage(self):
        usage = self.client.get_db_usage()

        self.assertTrue(isinstance(usage, dict))
        for name in usage.keys():
            self.assertTrue(isinstance(usage[name], dict))
            self.assertIn("rows", usage[name])
            self.assertTrue(isinstance(usage[name]["rows"], int))


class TestHashEquivalenceUnixServer(HashEquivalenceTestSetup, HashEquivalenceCommonTests, unittest.TestCase):
    def get_server_addr(self, server_idx):
        return "unix://" + os.path.join(self.temp_dir.name, 'sock%d' % server_idx)


class TestHashEquivalenceUnixServerLongPath(HashEquivalenceTestSetup, unittest.TestCase):
    DEEP_DIRECTORY = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa/bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb/ccccccccccccccccccccccccccccccccccccccccccc"
    def get_server_addr(self, server_idx):
        os.makedirs(os.path.join(self.temp_dir.name, self.DEEP_DIRECTORY), exist_ok=True)
        return "unix://" + os.path.join(self.temp_dir.name, self.DEEP_DIRECTORY, 'sock%d' % server_idx)


    def test_long_sock_path(self):
        # Simple test that hashes can be created
        taskhash = '35788efcb8dfb0a02659d81cf2bfd695fb30faf9'
        outhash = '2765d4a5884be49b28601445c2760c5f21e7e5c0ee2b7e3fce98fd7e5970796f'
        unihash = 'f46d3fbb439bd9b921095da657a4de906510d2cd'

        self.assertClientGetHash(self.client, taskhash, None)

        result = self.client.report_unihash(taskhash, self.METHOD, outhash, unihash)
        self.assertEqual(result['unihash'], unihash, 'Server returned bad unihash')


class TestHashEquivalenceTCPServer(HashEquivalenceTestSetup, HashEquivalenceCommonTests, unittest.TestCase):
    def get_server_addr(self, server_idx):
        # Some hosts cause asyncio module to misbehave, when IPv6 is not enabled.
        # If IPv6 is enabled, it should be safe to use localhost directly, in general
        # case it is more reliable to resolve the IP address explicitly.
        return socket.gethostbyname("localhost") + ":0"


class TestHashEquivalenceWebsocketServer(HashEquivalenceTestSetup, HashEquivalenceCommonTests, unittest.TestCase):
    def setUp(self):
        try:
            import websockets
        except ImportError as e:
            self.skipTest(str(e))

        super().setUp()

    def get_server_addr(self, server_idx):
        # Some hosts cause asyncio module to misbehave, when IPv6 is not enabled.
        # If IPv6 is enabled, it should be safe to use localhost directly, in general
        # case it is more reliable to resolve the IP address explicitly.
        host = socket.gethostbyname("localhost")
        return "ws://%s:0" % host


class TestHashEquivalenceWebsocketsSQLAlchemyServer(TestHashEquivalenceWebsocketServer):
    def setUp(self):
        try:
            import sqlalchemy
            import aiosqlite
        except ImportError as e:
            self.skipTest(str(e))

        super().setUp()

    def make_dbpath(self):
        return "sqlite+aiosqlite:///%s" % os.path.join(self.temp_dir.name, "db%d.sqlite" % self.server_index)


class TestHashEquivalenceExternalServer(HashEquivalenceTestSetup, HashEquivalenceCommonTests, unittest.TestCase):
    def start_test_server(self):
        if 'BB_TEST_HASHSERV' not in os.environ:
            self.skipTest('BB_TEST_HASHSERV not defined to test an external server')

        return os.environ['BB_TEST_HASHSERV']

    def start_server(self, *args, **kwargs):
        self.skipTest('Cannot start local server when testing external servers')

    def setUp(self):
        super().setUp()
        self.client.remove({"method": self.METHOD})

    def tearDown(self):
        self.client.remove({"method": self.METHOD})
        super().tearDown()