| # Copyright (C) 2019 Garmin Ltd. |
| # |
| # SPDX-License-Identifier: GPL-2.0-only |
| # |
| |
| import json |
| import logging |
| import socket |
| import os |
| from . import chunkify, DEFAULT_MAX_CHUNK |
| |
| |
| logger = logging.getLogger('hashserv.client') |
| |
| |
| class HashConnectionError(Exception): |
| pass |
| |
| |
| class Client(object): |
| MODE_NORMAL = 0 |
| MODE_GET_STREAM = 1 |
| |
| def __init__(self): |
| self._socket = None |
| self.reader = None |
| self.writer = None |
| self.mode = self.MODE_NORMAL |
| self.max_chunk = DEFAULT_MAX_CHUNK |
| |
| def connect_tcp(self, address, port): |
| def connect_sock(): |
| s = socket.create_connection((address, port)) |
| |
| s.setsockopt(socket.SOL_TCP, socket.TCP_NODELAY, 1) |
| s.setsockopt(socket.SOL_TCP, socket.TCP_QUICKACK, 1) |
| s.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) |
| return s |
| |
| self._connect_sock = connect_sock |
| |
| def connect_unix(self, path): |
| def connect_sock(): |
| s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) |
| # AF_UNIX has path length issues so chdir here to workaround |
| cwd = os.getcwd() |
| try: |
| os.chdir(os.path.dirname(path)) |
| s.connect(os.path.basename(path)) |
| finally: |
| os.chdir(cwd) |
| return s |
| |
| self._connect_sock = connect_sock |
| |
| def connect(self): |
| if self._socket is None: |
| self._socket = self._connect_sock() |
| |
| self.reader = self._socket.makefile('r', encoding='utf-8') |
| self.writer = self._socket.makefile('w', encoding='utf-8') |
| |
| self.writer.write('OEHASHEQUIV 1.1\n\n') |
| self.writer.flush() |
| |
| # Restore mode if the socket is being re-created |
| cur_mode = self.mode |
| self.mode = self.MODE_NORMAL |
| self._set_mode(cur_mode) |
| |
| return self._socket |
| |
| def close(self): |
| if self._socket is not None: |
| self._socket.close() |
| self._socket = None |
| self.reader = None |
| self.writer = None |
| |
| def _send_wrapper(self, proc): |
| count = 0 |
| while True: |
| try: |
| self.connect() |
| return proc() |
| except (OSError, HashConnectionError, json.JSONDecodeError, UnicodeDecodeError) as e: |
| logger.warning('Error talking to server: %s' % e) |
| if count >= 3: |
| if not isinstance(e, HashConnectionError): |
| raise HashConnectionError(str(e)) |
| raise e |
| self.close() |
| count += 1 |
| |
| def send_message(self, msg): |
| def get_line(): |
| line = self.reader.readline() |
| if not line: |
| raise HashConnectionError('Connection closed') |
| |
| if not line.endswith('\n'): |
| raise HashConnectionError('Bad message %r' % message) |
| |
| return line |
| |
| def proc(): |
| for c in chunkify(json.dumps(msg), self.max_chunk): |
| self.writer.write(c) |
| self.writer.flush() |
| |
| l = get_line() |
| |
| m = json.loads(l) |
| if 'chunk-stream' in m: |
| lines = [] |
| while True: |
| l = get_line().rstrip('\n') |
| if not l: |
| break |
| lines.append(l) |
| |
| m = json.loads(''.join(lines)) |
| |
| return m |
| |
| return self._send_wrapper(proc) |
| |
| def send_stream(self, msg): |
| def proc(): |
| self.writer.write("%s\n" % msg) |
| self.writer.flush() |
| l = self.reader.readline() |
| if not l: |
| raise HashConnectionError('Connection closed') |
| return l.rstrip() |
| |
| return self._send_wrapper(proc) |
| |
| def _set_mode(self, new_mode): |
| if new_mode == self.MODE_NORMAL and self.mode == self.MODE_GET_STREAM: |
| r = self.send_stream('END') |
| if r != 'ok': |
| raise HashConnectionError('Bad response from server %r' % r) |
| elif new_mode == self.MODE_GET_STREAM and self.mode == self.MODE_NORMAL: |
| r = self.send_message({'get-stream': None}) |
| if r != 'ok': |
| raise HashConnectionError('Bad response from server %r' % r) |
| elif new_mode != self.mode: |
| raise Exception('Undefined mode transition %r -> %r' % (self.mode, new_mode)) |
| |
| self.mode = new_mode |
| |
| def get_unihash(self, method, taskhash): |
| self._set_mode(self.MODE_GET_STREAM) |
| r = self.send_stream('%s %s' % (method, taskhash)) |
| if not r: |
| return None |
| return r |
| |
| def report_unihash(self, taskhash, method, outhash, unihash, extra={}): |
| self._set_mode(self.MODE_NORMAL) |
| m = extra.copy() |
| m['taskhash'] = taskhash |
| m['method'] = method |
| m['outhash'] = outhash |
| m['unihash'] = unihash |
| return self.send_message({'report': m}) |
| |
| def report_unihash_equiv(self, taskhash, method, unihash, extra={}): |
| self._set_mode(self.MODE_NORMAL) |
| m = extra.copy() |
| m['taskhash'] = taskhash |
| m['method'] = method |
| m['unihash'] = unihash |
| return self.send_message({'report-equiv': m}) |
| |
| def get_taskhash(self, method, taskhash, all_properties=False): |
| self._set_mode(self.MODE_NORMAL) |
| return self.send_message({'get': { |
| 'taskhash': taskhash, |
| 'method': method, |
| 'all': all_properties |
| }}) |
| |
| def get_stats(self): |
| self._set_mode(self.MODE_NORMAL) |
| return self.send_message({'get-stats': None}) |
| |
| def reset_stats(self): |
| self._set_mode(self.MODE_NORMAL) |
| return self.send_message({'reset-stats': None}) |