from common import * # isort:skip # noqa: F403 from mock_storage import mock_storage from storage import cache, cache_codec, cache_thp from storage.cache_common import InvalidSessionError from trezor import utils from trezor.messages import Initialize from trezor.messages import EndSession from apps.base import handle_EndSession, handle_Initialize KEY = 0 if utils.USE_THP: _PROTOCOL_CACHE = cache_thp else: _PROTOCOL_CACHE = cache_codec # Function moved from cache.py, as it was not used there def is_session_started() -> bool: return _PROTOCOL_CACHE.get_active_session() is not None class TestStorageCache( unittest.TestCase ): # noqa: F405 # pyright: ignore[reportUndefinedVariable] def setUp(self): cache.clear_all() def test_start_session(self): session_id_a = cache.start_session() self.assertIsNotNone(session_id_a) session_id_b = cache.start_session() self.assertNotEqual(session_id_a, session_id_b) cache.clear_all() with self.assertRaises(InvalidSessionError): cache.set(KEY, "something") with self.assertRaises(InvalidSessionError): cache.get(KEY) def test_end_session(self): session_id = cache.start_session() self.assertTrue(is_session_started()) cache.set(KEY, b"A") cache.end_current_session() self.assertFalse(is_session_started()) self.assertRaises(InvalidSessionError, cache.get, KEY) # ending an ended session should be a no-op cache.end_current_session() self.assertFalse(is_session_started()) session_id_a = cache.start_session(session_id) # original session no longer exists self.assertNotEqual(session_id_a, session_id) # original session data no longer exists self.assertIsNone(cache.get(KEY)) # create a new session session_id_b = cache.start_session() # switch back to original session session_id = cache.start_session(session_id_a) self.assertEqual(session_id, session_id_a) # end original session cache.end_current_session() # switch back to B session_id = cache.start_session(session_id_b) self.assertEqual(session_id, session_id_b) def test_session_queue(self): session_id = cache.start_session() self.assertEqual(cache.start_session(session_id), session_id) cache.set(KEY, b"A") for i in range(_PROTOCOL_CACHE._MAX_SESSIONS_COUNT): cache.start_session() self.assertNotEqual(cache.start_session(session_id), session_id) self.assertIsNone(cache.get(KEY)) def test_get_set(self): session_id1 = cache.start_session() cache.set(KEY, b"hello") self.assertEqual(cache.get(KEY), b"hello") session_id2 = cache.start_session() cache.set(KEY, b"world") self.assertEqual(cache.get(KEY), b"world") cache.start_session(session_id2) self.assertEqual(cache.get(KEY), b"world") cache.start_session(session_id1) self.assertEqual(cache.get(KEY), b"hello") cache.clear_all() with self.assertRaises(InvalidSessionError): cache.get(KEY) def test_get_set_int(self): session_id1 = cache.start_session() cache.set_int(KEY, 1234) self.assertEqual(cache.get_int(KEY), 1234) session_id2 = cache.start_session() cache.set_int(KEY, 5678) self.assertEqual(cache.get_int(KEY), 5678) cache.start_session(session_id2) self.assertEqual(cache.get_int(KEY), 5678) cache.start_session(session_id1) self.assertEqual(cache.get_int(KEY), 1234) cache.clear_all() with self.assertRaises(InvalidSessionError): cache.get_int(KEY) def test_delete(self): session_id1 = cache.start_session() self.assertIsNone(cache.get(KEY)) cache.set(KEY, b"hello") self.assertEqual(cache.get(KEY), b"hello") cache.delete(KEY) self.assertIsNone(cache.get(KEY)) cache.set(KEY, b"hello") cache.start_session() self.assertIsNone(cache.get(KEY)) cache.set(KEY, b"hello") self.assertEqual(cache.get(KEY), b"hello") cache.delete(KEY) self.assertIsNone(cache.get(KEY)) cache.start_session(session_id1) self.assertEqual(cache.get(KEY), b"hello") def test_decorators(self): run_count = 0 cache.start_session() @cache.stored(KEY) def func(): nonlocal run_count run_count += 1 return b"foo" # cache is empty self.assertIsNone(cache.get(KEY)) self.assertEqual(run_count, 0) self.assertEqual(func(), b"foo") # function was run self.assertEqual(run_count, 1) self.assertEqual(cache.get(KEY), b"foo") # function does not run again but returns cached value self.assertEqual(func(), b"foo") self.assertEqual(run_count, 1) @cache.stored_async(KEY) async def async_func(): nonlocal run_count run_count += 1 return b"bar" # cache is still full self.assertEqual(await_result(async_func()), b"foo") self.assertEqual(run_count, 1) cache.start_session() self.assertEqual(await_result(async_func()), b"bar") self.assertEqual(run_count, 2) # awaitable is also run only once self.assertEqual(await_result(async_func()), b"bar") self.assertEqual(run_count, 2) def test_empty_value(self): cache.start_session() self.assertIsNone(cache.get(KEY)) cache.set(KEY, b"") self.assertEqual(cache.get(KEY), b"") cache.delete(KEY) run_count = 0 @cache.stored(KEY) def func(): nonlocal run_count run_count += 1 return b"" self.assertEqual(func(), b"") # function gets called once self.assertEqual(run_count, 1) self.assertEqual(func(), b"") # function is not called for a second time self.assertEqual(run_count, 1) @mock_storage def test_Initialize(self): if utils.USE_THP: # INITIALIZE SHOULD NOT BE IN THP!!! TODO return def call_Initialize(**kwargs): msg = Initialize(**kwargs) return await_result(handle_Initialize(msg)) # calling Initialize without an ID allocates a new one session_id = cache.start_session() features = call_Initialize() self.assertNotEqual(session_id, features.session_id) # calling Initialize with the current ID does not allocate a new one features = call_Initialize(session_id=session_id) self.assertEqual(session_id, features.session_id) # store "hello" cache.set(KEY, b"hello") # check that it is cleared features = call_Initialize() session_id = features.session_id self.assertIsNone(cache.get(KEY)) # store "hello" again cache.set(KEY, b"hello") self.assertEqual(cache.get(KEY), b"hello") # supplying a different session ID starts a new cache call_Initialize(session_id=b"A" * _PROTOCOL_CACHE.SESSION_ID_LENGTH) self.assertIsNone(cache.get(KEY)) # but resuming a session loads the previous one call_Initialize(session_id=session_id) self.assertEqual(cache.get(KEY), b"hello") def test_EndSession(self): self.assertRaises(InvalidSessionError, cache.get, KEY) session_id = cache.start_session() self.assertTrue(is_session_started()) self.assertIsNone(cache.get(KEY)) await_result(handle_EndSession(EndSession())) self.assertFalse(is_session_started()) self.assertRaises(InvalidSessionError, cache.get, KEY) if __name__ == "__main__": unittest.main()