trezor-firmware/core/tests/test_storage.cache.py

from common import *  # isort:skip

from mock_storage import mock_storage
from storage import cache, cache_codec, cache_common
from trezor.messages import EndSession, Initialize
from trezor.wire import context
from trezor.wire.codec.codec_context import CodecContext

from apps.base import handle_EndSession, handle_Initialize
from apps.common.cache import stored, stored_async

KEY = 0


# Function moved from cache.py, as it was not used there
def is_session_started() -> bool:
    return cache_codec._active_session_idx is not None


class TestStorageCache(unittest.TestCase):

    def setUpClass(self):
        context.CURRENT_CONTEXT = CodecContext(None, bytearray(64))

    def tearDownClass(self):
        context.CURRENT_CONTEXT = None

    def setUp(self):
        cache.clear_all()

    def test_start_session(self):
        session_id_a = cache_codec.start_session()
        self.assertIsNotNone(session_id_a)
        session_id_b = cache_codec.start_session()
        self.assertNotEqual(session_id_a, session_id_b)

        cache.clear_all()
        with self.assertRaises(cache_common.InvalidSessionError):
            context.cache_set(KEY, "something")
        with self.assertRaises(cache_common.InvalidSessionError):
            context.cache_get(KEY)

    def test_end_session(self):
        session_id = cache_codec.start_session()
        self.assertTrue(is_session_started())
        context.cache_set(KEY, b"A")
        cache_codec.end_current_session()
        self.assertFalse(is_session_started())
        self.assertRaises(cache_common.InvalidSessionError, context.cache_get, KEY)

        # ending an ended session should be a no-op
        cache_codec.end_current_session()
        self.assertFalse(is_session_started())

        session_id_a = cache_codec.start_session(session_id)
        # original session no longer exists
        self.assertNotEqual(session_id_a, session_id)
        # original session data no longer exists
        self.assertIsNone(context.cache_get(KEY))

        # create a new session
        session_id_b = cache_codec.start_session()
        # switch back to original session
        session_id = cache_codec.start_session(session_id_a)
        self.assertEqual(session_id, session_id_a)
        # end original session
        cache_codec.end_current_session()
        # switch back to B
        session_id = cache_codec.start_session(session_id_b)
        self.assertEqual(session_id, session_id_b)

    def test_session_queue(self):
        session_id = cache_codec.start_session()
        self.assertEqual(cache_codec.start_session(session_id), session_id)
        context.cache_set(KEY, b"A")
        for i in range(cache_codec._MAX_SESSIONS_COUNT):
            cache_codec.start_session()
        self.assertNotEqual(cache_codec.start_session(session_id), session_id)
        self.assertIsNone(context.cache_get(KEY))

    def test_get_set(self):
        session_id1 = cache_codec.start_session()
        context.cache_set(KEY, b"hello")
        self.assertEqual(context.cache_get(KEY), b"hello")

        session_id2 = cache_codec.start_session()
        context.cache_set(KEY, b"world")
        self.assertEqual(context.cache_get(KEY), b"world")

        cache_codec.start_session(session_id2)
        self.assertEqual(context.cache_get(KEY), b"world")
        cache_codec.start_session(session_id1)
        self.assertEqual(context.cache_get(KEY), b"hello")

        cache.clear_all()
        with self.assertRaises(cache_common.InvalidSessionError):
            context.cache_get(KEY)

    def test_get_set_int(self):
        session_id1 = cache_codec.start_session()
        context.cache_set_int(KEY, 1234)
        self.assertEqual(context.cache_get_int(KEY), 1234)

        session_id2 = cache_codec.start_session()
        context.cache_set_int(KEY, 5678)
        self.assertEqual(context.cache_get_int(KEY), 5678)

        cache_codec.start_session(session_id2)
        self.assertEqual(context.cache_get_int(KEY), 5678)
        cache_codec.start_session(session_id1)
        self.assertEqual(context.cache_get_int(KEY), 1234)

        cache.clear_all()
        with self.assertRaises(cache_common.InvalidSessionError):
            context.cache_get_int(KEY)

    def test_delete(self):
        session_id1 = cache_codec.start_session()
        self.assertIsNone(context.cache_get(KEY))
        context.cache_set(KEY, b"hello")
        self.assertEqual(context.cache_get(KEY), b"hello")
        context.cache_delete(KEY)
        self.assertIsNone(context.cache_get(KEY))

        context.cache_set(KEY, b"hello")
        cache_codec.start_session()
        self.assertIsNone(context.cache_get(KEY))
        context.cache_set(KEY, b"hello")
        self.assertEqual(context.cache_get(KEY), b"hello")
        context.cache_delete(KEY)
        self.assertIsNone(context.cache_get(KEY))

        cache_codec.start_session(session_id1)
        self.assertEqual(context.cache_get(KEY), b"hello")

    def test_decorators(self):
        run_count = 0
        cache_codec.start_session()

        @stored(KEY)
        def func():
            nonlocal run_count
            run_count += 1
            return b"foo"

        # cache is empty
        self.assertIsNone(context.cache_get(KEY))
        self.assertEqual(run_count, 0)
        self.assertEqual(func(), b"foo")
        # function was run
        self.assertEqual(run_count, 1)
        self.assertEqual(context.cache_get(KEY), b"foo")
        # function does not run again but returns cached value
        self.assertEqual(func(), b"foo")
        self.assertEqual(run_count, 1)

        @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_codec.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_codec.start_session()

        self.assertIsNone(context.cache_get(KEY))
        context.cache_set(KEY, b"")
        self.assertEqual(context.cache_get(KEY), b"")

        context.cache_delete(KEY)
        run_count = 0

        @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):
        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_codec.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"
        context.cache_set(KEY, b"hello")
        # check that it is cleared
        features = call_Initialize()
        session_id = features.session_id
        self.assertIsNone(context.cache_get(KEY))
        # store "hello" again
        context.cache_set(KEY, b"hello")
        self.assertEqual(context.cache_get(KEY), b"hello")

        # supplying a different session ID starts a new cache
        call_Initialize(session_id=b"A" * cache_codec.SESSION_ID_LENGTH)
        self.assertIsNone(context.cache_get(KEY))

        # but resuming a session loads the previous one
        call_Initialize(session_id=session_id)
        self.assertEqual(context.cache_get(KEY), b"hello")

    def test_EndSession(self):
        self.assertRaises(cache_common.InvalidSessionError, context.cache_get, KEY)
        cache_codec.start_session()
        self.assertTrue(is_session_started())
        self.assertIsNone(context.cache_get(KEY))
        await_result(handle_EndSession(EndSession()))
        self.assertFalse(is_session_started())
        self.assertRaises(cache_common.InvalidSessionError, context.cache_get, KEY)


if __name__ == "__main__":
    unittest.main()