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
9import hashlib
10import logging
11import multiprocessing
12import sys
13import tempfile
14import threading
15import unittest
16
17
18class TestHashEquivalenceServer(object):
19    METHOD = 'TestMethod'
20
21    def _run_server(self):
22        # logging.basicConfig(level=logging.DEBUG, filename='bbhashserv.log', filemode='w',
23        #                     format='%(levelname)s %(filename)s:%(lineno)d %(message)s')
24        self.server.serve_forever()
25
26    def setUp(self):
27        if sys.version_info < (3, 5, 0):
28            self.skipTest('Python 3.5 or later required')
29
30        self.temp_dir = tempfile.TemporaryDirectory(prefix='bb-hashserv')
31        self.dbfile = os.path.join(self.temp_dir.name, 'db.sqlite')
32
33        self.server = create_server(self.get_server_addr(), self.dbfile)
34        self.server_thread = multiprocessing.Process(target=self._run_server)
35        self.server_thread.start()
36        self.client = create_client(self.server.address)
37
38    def tearDown(self):
39        # Shutdown server
40        s = getattr(self, 'server', None)
41        if s is not None:
42            self.server_thread.terminate()
43            self.server_thread.join()
44        self.client.close()
45        self.temp_dir.cleanup()
46
47    def test_create_hash(self):
48        # Simple test that hashes can be created
49        taskhash = '35788efcb8dfb0a02659d81cf2bfd695fb30faf9'
50        outhash = '2765d4a5884be49b28601445c2760c5f21e7e5c0ee2b7e3fce98fd7e5970796f'
51        unihash = 'f46d3fbb439bd9b921095da657a4de906510d2cd'
52
53        result = self.client.get_unihash(self.METHOD, taskhash)
54        self.assertIsNone(result, msg='Found unexpected task, %r' % result)
55
56        result = self.client.report_unihash(taskhash, self.METHOD, outhash, unihash)
57        self.assertEqual(result['unihash'], unihash, 'Server returned bad unihash')
58
59    def test_create_equivalent(self):
60        # Tests that a second reported task with the same outhash will be
61        # assigned the same unihash
62        taskhash = '53b8dce672cb6d0c73170be43f540460bfc347b4'
63        outhash = '5a9cb1649625f0bf41fc7791b635cd9c2d7118c7f021ba87dcd03f72b67ce7a8'
64        unihash = 'f37918cc02eb5a520b1aff86faacbc0a38124646'
65
66        result = self.client.report_unihash(taskhash, self.METHOD, outhash, unihash)
67        self.assertEqual(result['unihash'], unihash, 'Server returned bad unihash')
68
69        # Report a different task with the same outhash. The returned unihash
70        # should match the first task
71        taskhash2 = '3bf6f1e89d26205aec90da04854fbdbf73afe6b4'
72        unihash2 = 'af36b199320e611fbb16f1f277d3ee1d619ca58b'
73        result = self.client.report_unihash(taskhash2, self.METHOD, outhash, unihash2)
74        self.assertEqual(result['unihash'], unihash, 'Server returned bad unihash')
75
76    def test_duplicate_taskhash(self):
77        # Tests that duplicate reports of the same taskhash with different
78        # outhash & unihash always return the unihash from the first reported
79        # taskhash
80        taskhash = '8aa96fcffb5831b3c2c0cb75f0431e3f8b20554a'
81        outhash = 'afe240a439959ce86f5e322f8c208e1fedefea9e813f2140c81af866cc9edf7e'
82        unihash = '218e57509998197d570e2c98512d0105985dffc9'
83        self.client.report_unihash(taskhash, self.METHOD, outhash, unihash)
84
85        result = self.client.get_unihash(self.METHOD, taskhash)
86        self.assertEqual(result, unihash)
87
88        outhash2 = '0904a7fe3dc712d9fd8a74a616ddca2a825a8ee97adf0bd3fc86082c7639914d'
89        unihash2 = 'ae9a7d252735f0dafcdb10e2e02561ca3a47314c'
90        self.client.report_unihash(taskhash, self.METHOD, outhash2, unihash2)
91
92        result = self.client.get_unihash(self.METHOD, taskhash)
93        self.assertEqual(result, unihash)
94
95        outhash3 = '77623a549b5b1a31e3732dfa8fe61d7ce5d44b3370f253c5360e136b852967b4'
96        unihash3 = '9217a7d6398518e5dc002ed58f2cbbbc78696603'
97        self.client.report_unihash(taskhash, self.METHOD, outhash3, unihash3)
98
99        result = self.client.get_unihash(self.METHOD, taskhash)
100        self.assertEqual(result, unihash)
101
102    def test_stress(self):
103        def query_server(failures):
104            client = Client(self.server.address)
105            try:
106                for i in range(1000):
107                    taskhash = hashlib.sha256()
108                    taskhash.update(str(i).encode('utf-8'))
109                    taskhash = taskhash.hexdigest()
110                    result = client.get_unihash(self.METHOD, taskhash)
111                    if result != taskhash:
112                        failures.append("taskhash mismatch: %s != %s" % (result, taskhash))
113            finally:
114                client.close()
115
116        # Report hashes
117        for i in range(1000):
118            taskhash = hashlib.sha256()
119            taskhash.update(str(i).encode('utf-8'))
120            taskhash = taskhash.hexdigest()
121            self.client.report_unihash(taskhash, self.METHOD, taskhash, taskhash)
122
123        failures = []
124        threads = [threading.Thread(target=query_server, args=(failures,)) for t in range(100)]
125
126        for t in threads:
127            t.start()
128
129        for t in threads:
130            t.join()
131
132        self.assertFalse(failures)
133
134
135class TestHashEquivalenceUnixServer(TestHashEquivalenceServer, unittest.TestCase):
136    def get_server_addr(self):
137        return "unix://" + os.path.join(self.temp_dir.name, 'sock')
138
139
140class TestHashEquivalenceTCPServer(TestHashEquivalenceServer, unittest.TestCase):
141    def get_server_addr(self):
142        return "localhost:0"
143