From 1016b0c98a8237ebfb161aefaaa0de5dde43cd94 Mon Sep 17 00:00:00 2001
From: Andrew Kozlik <andrew.kozlik@gmail.com>
Date: Sat, 6 Jul 2024 23:18:42 +0200
Subject: [PATCH] feat(storage): Do not lock freshly initialized storage.

---
 storage/storage.c                   | 3 +--
 storage/tests/python/src/storage.py | 4 +++-
 storage/tests/tests/test_set_get.py | 1 +
 3 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/storage/storage.c b/storage/storage.c
index e4f4632c3..fe36b52e0 100644
--- a/storage/storage.c
+++ b/storage/storage.c
@@ -792,6 +792,7 @@ void storage_init(PIN_UI_WAIT_CALLBACK callback, const uint8_t *salt,
                   const uint16_t salt_len) {
   initialized = secfalse;
   unlocked = secfalse;
+  memzero(cached_keys, sizeof(cached_keys));
   norcow_init(&norcow_active_version);
   initialized = sectrue;
   ui_callback = callback;
@@ -810,9 +811,7 @@ void storage_init(PIN_UI_WAIT_CALLBACK callback, const uint8_t *salt,
   uint16_t len = 0;
   if (secfalse == norcow_get(EDEK_PVC_KEY, &val, &len)) {
     init_wiped_storage();
-    storage_lock();
   }
-  memzero(cached_keys, sizeof(cached_keys));
 }
 
 secbool storage_pin_fails_increase(void) {
diff --git a/storage/tests/python/src/storage.py b/storage/tests/python/src/storage.py
index 5a0c67aae..ef4f534d7 100644
--- a/storage/tests/python/src/storage.py
+++ b/storage/tests/python/src/storage.py
@@ -42,7 +42,7 @@ class Storage:
         self.pin_log.init()
         self._set_wipe_code(consts.WIPE_CODE_EMPTY)
         self._set_pin(consts.PIN_EMPTY)
-        self.unlocked = False
+        self.unlocked = True
 
     def _set_pin(self, pin: str):
         random_salt = prng.random_buffer(consts.PIN_SALT_SIZE)
@@ -220,6 +220,8 @@ class Storage:
 
     def _decrypt(self, key: int) -> bytes:
         data = self.nc.get(key)
+        if data is None:
+            raise RuntimeError("Key not found")
         iv = data[: consts.CHACHA_IV_SIZE]
         # cipher text with MAC
 
diff --git a/storage/tests/tests/test_set_get.py b/storage/tests/tests/test_set_get.py
index 5236ab1b1..25838d616 100644
--- a/storage/tests/tests/test_set_get.py
+++ b/storage/tests/tests/test_set_get.py
@@ -237,6 +237,7 @@ def test_set_similar(nc_class):
 def test_set_locked(nc_class):
     sc, sp = common.init(nc_class)
     for s in (sc, sp):
+        s.lock()
         with pytest.raises(RuntimeError):
             s.set(0x0303, b"test")
         with pytest.raises(RuntimeError):