From 4c2e4bcb6553fa30c9f4daa395788d3dd2e2aac5 Mon Sep 17 00:00:00 2001
From: Andrew Kozlik <andrew.kozlik@gmail.com>
Date: Thu, 11 Mar 2021 14:00:51 +0100
Subject: [PATCH] fix(storage): Check for overflow in counter increment.

---
 storage/storage.c                   |  8 ++++++++
 storage/tests/c/storage.py          | 14 +++++++-------
 storage/tests/python/src/consts.py  |  3 +++
 storage/tests/python/src/storage.py |  6 ++++++
 storage/tests/tests/test_set_get.py | 16 ++++++++++++++++
 5 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/storage/storage.c b/storage/storage.c
index 01bcd789a..10431f25e 100644
--- a/storage/storage.c
+++ b/storage/storage.c
@@ -1363,9 +1363,17 @@ secbool storage_next_counter(const uint16_t key, uint32_t *count) {
   }
 
   *count = val_stored[0] + 1 + 32 * (i - 1);
+  if (*count < val_stored[0]) {
+    // Value overflow.
+    return secfalse;
+  }
 
   if (i < len_words) {
     *count += hamming_weight(~val_stored[i]);
+    if (*count < val_stored[0]) {
+      // Value overflow.
+      return secfalse;
+    }
     return norcow_update_word(key, sizeof(uint32_t) * i, val_stored[i] >> 1);
   } else {
     return storage_set_counter(key, *count);
diff --git a/storage/tests/c/storage.py b/storage/tests/c/storage.py
index ee652ff99..663887a02 100644
--- a/storage/tests/c/storage.py
+++ b/storage/tests/c/storage.py
@@ -70,17 +70,17 @@ class Storage:
         if sectrue != self.lib.storage_set(c.c_uint16(key), val, c.c_uint16(len(val))):
             raise RuntimeError("Failed to set value in storage.")
 
-    def set_counter(self, key: int, count: int) -> bool:
-        return sectrue == self.lib.storage_set_counter(
+    def set_counter(self, key: int, count: int) -> None:
+        if count > 0xFFFF_FFFF or sectrue != self.lib.storage_set_counter(
             c.c_uint16(key), c.c_uint32(count)
-        )
+        ):
+            raise RuntimeError("Failed to set value in storage.")
 
     def next_counter(self, key: int) -> int:
         count = c.c_uint32()
-        if sectrue == self.lib.storage_next_counter(c.c_uint16(key), c.byref(count)):
-            return count.value
-        else:
-            return None
+        if sectrue != self.lib.storage_next_counter(c.c_uint16(key), c.byref(count)):
+            raise RuntimeError("Failed to set value in storage.")
+        return count.value
 
     def delete(self, key: int) -> bool:
         return sectrue == self.lib.storage_delete(c.c_uint16(key))
diff --git a/storage/tests/python/src/consts.py b/storage/tests/python/src/consts.py
index 349794dee..a02d6a72f 100644
--- a/storage/tests/python/src/consts.py
+++ b/storage/tests/python/src/consts.py
@@ -1,3 +1,6 @@
+# ----- General ----- #
+UINT32_MAX = 0xFFFF_FFFF
+
 # ----- PIN and encryption related ----- #
 
 # App ID where PIN log is stored.
diff --git a/storage/tests/python/src/storage.py b/storage/tests/python/src/storage.py
index 538fa2be9..bc6e940e5 100644
--- a/storage/tests/python/src/storage.py
+++ b/storage/tests/python/src/storage.py
@@ -149,6 +149,10 @@ class Storage:
         app = key >> 8
         if not consts.is_app_public(app):
             raise RuntimeError("Counter can be set only for public items")
+
+        if val > consts.UINT32_MAX:
+            raise RuntimeError("Failed to set value in storage.")
+
         counter = val.to_bytes(4, sys.byteorder) + bytearray(
             b"\xFF" * consts.COUNTER_TAIL_SIZE
         )
@@ -167,6 +171,8 @@ class Storage:
         tail = helpers.to_int_by_words(current[4:])
         tail_count = "{0:064b}".format(tail).count("0")
         increased_count = base + tail_count + 1
+        if increased_count > consts.UINT32_MAX:
+            raise RuntimeError("Failed to set value in storage.")
 
         if tail_count == consts.COUNTER_MAX_TAIL:
             self.set_counter(key, increased_count)
diff --git a/storage/tests/tests/test_set_get.py b/storage/tests/tests/test_set_get.py
index d46c8b12c..5fce93364 100644
--- a/storage/tests/tests/test_set_get.py
+++ b/storage/tests/tests/test_set_get.py
@@ -1,5 +1,7 @@
 import pytest
 
+from python.src import consts
+
 from . import common
 
 # Strings for testing ChaCha20 encryption.
@@ -183,3 +185,17 @@ def test_counter():
         for s in (sc, sp):
             assert i == s.next_counter(0xC001)
     assert common.memory_equals(sc, sp)
+
+    for s in (sc, sp):
+        with pytest.raises(RuntimeError):
+            s.set_counter(0xC001, consts.UINT32_MAX + 1)
+
+        start = consts.UINT32_MAX - 100
+        s.set_counter(0xC001, start)
+        for i in range(start, consts.UINT32_MAX):
+            assert i + 1 == s.next_counter(0xC001)
+
+        with pytest.raises(RuntimeError):
+            s.next_counter(0xC001)
+
+    assert common.memory_equals(sc, sp)