1
0
mirror of https://github.com/trezor/trezor-firmware.git synced 2024-11-26 17:38:39 +00:00

trezorlib: reentrant session handling

This fixes the breakage introduced by transport reshuffles.
It's still not great and I'd love to see context manager based sessions.
But it's good enough for now.
This commit is contained in:
matejcik 2018-11-08 18:08:02 +01:00
parent daf97afb37
commit 85b85c67b3
6 changed files with 25 additions and 17 deletions

View File

@ -84,15 +84,23 @@ class BaseClient(object):
LOG.info("creating client instance for device: {}".format(transport.get_path()))
self.transport = transport
self.ui = ui
self.session_counter = 0
super(BaseClient, self).__init__() # *args, **kwargs)
def open(self):
if self.session_counter == 0:
self.transport.begin_session()
self.session_counter += 1
def close(self):
pass
if self.session_counter == 1:
self.transport.end_session()
self.session_counter -= 1
def cancel(self):
self._raw_write(proto.Cancel())
@tools.session
def call_raw(self, msg):
__tracebackhide__ = True # for pytest # pylint: disable=W0612
self._raw_write(msg)
@ -174,6 +182,7 @@ class ProtocolMixin(object):
def set_tx_api(self, tx_api):
warnings.warn("set_tx_api is deprecated, use new arguments to sign_tx")
@tools.session
def init_device(self):
resp = self.call(proto.Initialize(state=self.state))
if not isinstance(resp, proto.Features):

View File

@ -25,6 +25,8 @@ from .tools import expect
class DebugLink:
def __init__(self, transport):
self.transport = transport
def open(self):
self.transport.begin_session()
def close(self):
@ -186,10 +188,13 @@ class TrezorClientDebugLink(TrezorClient):
self.set_passphrase("")
super().__init__(transport, ui=self.ui)
def open(self):
super().open()
self.debug.open()
def close(self):
self.debug.close()
super().close()
if self.debug:
self.debug.close()
def set_filter(self, message_type, callback):
self.filters[message_type] = callback

View File

@ -41,10 +41,9 @@ class TrezorTest:
# self.client.set_buttonwait(3)
device.wipe(self.client)
self.client.transport.begin_session()
self.client.open()
def teardown_method(self, method):
self.client.transport.end_session()
self.client.close()
def _setup_mnemonic(self, mnemonic=None, pin="", passphrase=False):

View File

@ -55,15 +55,9 @@ def client():
wirelink = get_device()
client = TrezorClientDebugLink(wirelink)
wipe_device(client)
client.transport.begin_session()
client.open()
yield client
client.transport.end_session()
# XXX debuglink session must also be closed
# client.close accomplishes that for now; going forward, there should
# also be proper session handling for debuglink
client.close()

View File

@ -224,11 +224,11 @@ def session(f):
@functools.wraps(f)
def wrapped_f(client, *args, **kwargs):
__tracebackhide__ = True # for pytest # pylint: disable=W0612
client.transport.begin_session()
client.open()
try:
return f(client, *args, **kwargs)
finally:
client.transport.begin_session()
client.close()
return wrapped_f

View File

@ -85,15 +85,16 @@ class Protocol:
self.handle = handle
self.session_counter = 0
# XXX we might be able to remove this now that TrezorClient does session handling
def begin_session(self) -> None:
if self.session_counter == 0:
self.handle.open()
self.session_counter += 1
def end_session(self) -> None:
self.session_counter = max(self.session_counter - 1, 0)
if self.session_counter == 0:
if self.session_counter == 1:
self.handle.close()
self.session_counter -= 1
def read(self) -> protobuf.MessageType:
raise NotImplementedError