Merge pull request #700 from trezor/tsusanka/upgrade-u2f-check

Get next u2f counter and add tests for it

common/protob/messages-management.proto

@@ -333,5 +333,21 @@ message WordAck {
  * @next Success
  */
 message SetU2FCounter {
-    optional uint32 u2f_counter = 1;    // counter
+    optional uint32 u2f_counter = 1;
+}
+
+/**
+ * Request: Set U2F counter
+ * @start
+ * @next NextU2FCounter
+ */
+message GetNextU2FCounter {
+}
+
+/**
+ * Request: Set U2F counter
+ * @end
+ */
+message NextU2FCounter {
+    optional uint32 u2f_counter = 1;
 }

common/protob/messages.proto

@@ -62,6 +62,8 @@ enum MessageType {
     MessageType_GetFeatures = 55 [(wire_in) = true];
     MessageType_SetU2FCounter = 63 [(wire_in) = true];
     MessageType_SdProtect = 79 [(wire_in) = true];
+    MessageType_GetNextU2FCounter = 80 [(wire_in) = true];
+    MessageType_NextU2FCounter = 81 [(wire_out) = true];
 
     // Bootloader
     MessageType_FirmwareErase = 6 [(wire_in) = true, (wire_bootloader) = true];

core/src/apps/management/__init__.py

@@ -14,4 +14,5 @@ def boot() -> None:
     wire.add(MessageType.ApplyFlags, __name__, "apply_flags")
     wire.add(MessageType.ChangePin, __name__, "change_pin")
     wire.add(MessageType.SetU2FCounter, __name__, "set_u2f_counter")
+    wire.add(MessageType.GetNextU2FCounter, __name__, "get_next_u2f_counter")
     wire.add(MessageType.SdProtect, __name__, "sd_protect")

core/src/apps/management/get_next_u2f_counter.py

@@ -0,0 +1,20 @@
+import storage.device
+from trezor import ui, wire
+from trezor.messages import ButtonRequestType
+from trezor.messages.GetNextU2FCounter import GetNextU2FCounter
+from trezor.messages.NextU2FCounter import NextU2FCounter
+from trezor.ui.text import Text
+
+from apps.common.confirm import require_confirm
+
+
+async def get_next_u2f_counter(
+    ctx: wire.Context, msg: GetNextU2FCounter
+) -> NextU2FCounter:
+    text = Text("Get next U2F counter", ui.ICON_CONFIG)
+    text.normal("Do you really want to")
+    text.bold("increase and retrieve")
+    text.normal("the U2F counter?")
+    await require_confirm(ctx, text, code=ButtonRequestType.ProtectCall)
+
+    return NextU2FCounter(storage.device.next_u2f_counter())

core/src/apps/management/set_u2f_counter.py

@@ -1,13 +1,14 @@
 import storage.device
 from trezor import ui, wire
 from trezor.messages import ButtonRequestType
+from trezor.messages.SetU2FCounter import SetU2FCounter
 from trezor.messages.Success import Success
 from trezor.ui.text import Text
 
 from apps.common.confirm import require_confirm
 
 
-async def set_u2f_counter(ctx, msg):
+async def set_u2f_counter(ctx: wire.Context, msg: SetU2FCounter) -> Success:
     if msg.u2f_counter is None:
         raise wire.ProcessError("No value provided")
 

core/src/trezor/messages/GetNextU2FCounter.py

@@ -0,0 +1,14 @@
+# Automatically generated by pb2py
+# fmt: off
+import protobuf as p
+
+if __debug__:
+    try:
+        from typing import Dict, List  # noqa: F401
+        from typing_extensions import Literal  # noqa: F401
+    except ImportError:
+        pass
+
+
+class GetNextU2FCounter(p.MessageType):
+    MESSAGE_WIRE_TYPE = 80

core/src/trezor/messages/MessageType.py

@@ -37,6 +37,8 @@ WordAck = 47  # type: Literal[47]
 GetFeatures = 55  # type: Literal[55]
 SetU2FCounter = 63  # type: Literal[63]
 SdProtect = 79  # type: Literal[79]
+GetNextU2FCounter = 80  # type: Literal[80]
+NextU2FCounter = 81  # type: Literal[81]
 FirmwareErase = 6  # type: Literal[6]
 FirmwareUpload = 7  # type: Literal[7]
 FirmwareRequest = 8  # type: Literal[8]

core/src/trezor/messages/NextU2FCounter.py

@@ -0,0 +1,26 @@
+# Automatically generated by pb2py
+# fmt: off
+import protobuf as p
+
+if __debug__:
+    try:
+        from typing import Dict, List  # noqa: F401
+        from typing_extensions import Literal  # noqa: F401
+    except ImportError:
+        pass
+
+
+class NextU2FCounter(p.MessageType):
+    MESSAGE_WIRE_TYPE = 81
+
+    def __init__(
+        self,
+        u2f_counter: int = None,
+    ) -> None:
+        self.u2f_counter = u2f_counter
+
+    @classmethod
+    def get_fields(cls) -> Dict:
+        return {
+            1: ('u2f_counter', p.UVarintType, 0),
+        }

legacy/firmware/fsm.h

@@ -65,6 +65,7 @@ void fsm_msgApplyFlags(const ApplyFlags *msg);
 void fsm_msgRecoveryDevice(const RecoveryDevice *msg);
 void fsm_msgWordAck(const WordAck *msg);
 void fsm_msgSetU2FCounter(const SetU2FCounter *msg);
+void fsm_msgGetNextU2FCounter(void);
 
 // coin
 void fsm_msgGetPublicKey(const GetPublicKey *msg);

legacy/firmware/fsm_msg_common.h

@@ -430,3 +430,21 @@ void fsm_msgSetU2FCounter(const SetU2FCounter *msg) {
   fsm_sendSuccess(_("U2F counter set"));
   layoutHome();
 }
+
+void fsm_msgGetNextU2FCounter() {
+  layoutDialogSwipe(&bmp_icon_question, _("Cancel"), _("Confirm"), NULL,
+                    _("Do you want to"), _("increase and retrieve"),
+                    _("the U2F counter?"), NULL, NULL, NULL);
+  if (!protectButton(ButtonRequestType_ButtonRequest_ProtectCall, false)) {
+    fsm_sendFailure(FailureType_Failure_ActionCancelled, NULL);
+    layoutHome();
+    return;
+  }
+  uint32_t counter = config_nextU2FCounter();
+
+  RESP_INIT(NextU2FCounter);
+  resp->has_u2f_counter = true;
+  resp->u2f_counter = counter;
+  msg_write(MessageType_MessageType_NextU2FCounter, resp);
+  layoutHome();
+}

python/src/trezorlib/cli/trezorctl.py

@@ -412,6 +412,12 @@ def set_u2f_counter(connect, counter):
     return device.set_u2f_counter(connect(), counter)
 
 
+@cli.command(help="Get U2F counter.")
+@click.pass_obj
+def get_next_u2f_counter(connect):
+    return device.get_next_u2f_counter(connect())
+
+
 @cli.command(help="Reset device to factory defaults and remove all private data.")
 @click.option(
     "-b",

python/src/trezorlib/device.py

@@ -99,8 +99,12 @@ def sd_protect(client, operation):
 
 @expect(proto.Success, field="message")
 def set_u2f_counter(client, u2f_counter):
-    ret = client.call(proto.SetU2FCounter(u2f_counter=u2f_counter))
-    return ret
+    return client.call(proto.SetU2FCounter(u2f_counter=u2f_counter))
+
+
+@expect(proto.NextU2FCounter, field="u2f_counter")
+def get_next_u2f_counter(client):
+    return client.call(proto.GetNextU2FCounter())
 
 
 @expect(proto.Success, field="message")

python/src/trezorlib/messages/GetNextU2FCounter.py

@@ -0,0 +1,14 @@
+# Automatically generated by pb2py
+# fmt: off
+from .. import protobuf as p
+
+if __debug__:
+    try:
+        from typing import Dict, List  # noqa: F401
+        from typing_extensions import Literal  # noqa: F401
+    except ImportError:
+        pass
+
+
+class GetNextU2FCounter(p.MessageType):
+    MESSAGE_WIRE_TYPE = 80

python/src/trezorlib/messages/MessageType.py

@@ -35,6 +35,8 @@ WordAck = 47  # type: Literal[47]
 GetFeatures = 55  # type: Literal[55]
 SetU2FCounter = 63  # type: Literal[63]
 SdProtect = 79  # type: Literal[79]
+GetNextU2FCounter = 80  # type: Literal[80]
+NextU2FCounter = 81  # type: Literal[81]
 FirmwareErase = 6  # type: Literal[6]
 FirmwareUpload = 7  # type: Literal[7]
 FirmwareRequest = 8  # type: Literal[8]

python/src/trezorlib/messages/NextU2FCounter.py

@@ -0,0 +1,26 @@
+# Automatically generated by pb2py
+# fmt: off
+from .. import protobuf as p
+
+if __debug__:
+    try:
+        from typing import Dict, List  # noqa: F401
+        from typing_extensions import Literal  # noqa: F401
+    except ImportError:
+        pass
+
+
+class NextU2FCounter(p.MessageType):
+    MESSAGE_WIRE_TYPE = 81
+
+    def __init__(
+        self,
+        u2f_counter: int = None,
+    ) -> None:
+        self.u2f_counter = u2f_counter
+
+    @classmethod
+    def get_fields(cls) -> Dict:
+        return {
+            1: ('u2f_counter', p.UVarintType, 0),
+        }

python/src/trezorlib/messages/__init__.py

@@ -101,6 +101,7 @@ from .GetAddress import GetAddress
 from .GetECDHSessionKey import GetECDHSessionKey
 from .GetEntropy import GetEntropy
 from .GetFeatures import GetFeatures
+from .GetNextU2FCounter import GetNextU2FCounter
 from .GetPublicKey import GetPublicKey
 from .HDNodePathType import HDNodePathType
 from .HDNodeType import HDNodeType
@@ -187,6 +188,7 @@ from .NEMSignTx import NEMSignTx
 from .NEMSignedTx import NEMSignedTx
 from .NEMTransactionCommon import NEMTransactionCommon
 from .NEMTransfer import NEMTransfer
+from .NextU2FCounter import NextU2FCounter
 from .PassphraseAck import PassphraseAck
 from .PassphraseRequest import PassphraseRequest
 from .PassphraseStateAck import PassphraseStateAck

tests/device_tests/test_u2f_counter.py

@@ -0,0 +1,27 @@
+# This file is part of the Trezor project.
+#
+# Copyright (C) 2012-2019 SatoshiLabs and contributors
+#
+# This library is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License version 3
+# as published by the Free Software Foundation.
+#
+# This library is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the License along with this library.
+# If not, see <https://www.gnu.org/licenses/lgpl-3.0.html>.
+
+from trezorlib import device
+
+
+def test_u2f_counter(client):
+    assert device.get_next_u2f_counter(client) == 0
+    assert device.get_next_u2f_counter(client) == 1
+    device.set_u2f_counter(client, 111111)
+    assert device.get_next_u2f_counter(client) == 111112
+    assert device.get_next_u2f_counter(client) == 111113
+    device.set_u2f_counter(client, 0)
+    assert device.get_next_u2f_counter(client) == 1

tests/upgrade_tests/test_firmware_upgrades.py

@@ -40,7 +40,7 @@ LANGUAGE = "english"
 STRENGTH = 128
 
 
-def for_all(*args, minimum_version=(1, 0, 0)):
+def for_all(*args, legacy_minimum_version=(1, 0, 0), core_minimum_version=(2, 0, 0)):
     if not args:
         args = ("core", "legacy")
 
@@ -49,6 +49,13 @@ def for_all(*args, minimum_version=(1, 0, 0)):
 
     all_params = []
     for gen in args:
+        if gen == "legacy":
+            minimum_version = legacy_minimum_version
+        elif gen == "core":
+            minimum_version = core_minimum_version
+        else:
+            raise ValueError
+
         if gen not in enabled_gens:
             continue
         try:
@@ -164,7 +171,7 @@ def test_upgrade_reset_skip_backup(gen, from_tag, to_tag):
         assert btc.get_address(emu.client, "Bitcoin", PATH) == address
 
 
-@for_all(minimum_version=(1, 7, 2))
+@for_all(legacy_minimum_version=(1, 7, 2))
 def test_upgrade_reset_no_backup(gen, from_tag, to_tag):
     def asserts(tag, client):
         assert not client.features.pin_protection
@@ -199,7 +206,7 @@ def test_upgrade_reset_no_backup(gen, from_tag, to_tag):
 
 
 # Although Shamir was introduced in 2.1.2 already, the debug instrumentation was not present until 2.1.9.
-@for_all("core", minimum_version=(2, 1, 9))
+@for_all("core", core_minimum_version=(2, 1, 9))
 def test_upgrade_shamir_recovery(gen, from_tag, to_tag):
     with EmulatorWrapper(gen, from_tag) as emu, BackgroundDeviceHandler(
         emu.client
@@ -240,6 +247,24 @@ def test_upgrade_shamir_recovery(gen, from_tag, to_tag):
         device_handler.check_finalize()
 
 
+@for_all(legacy_minimum_version=(1, 8, 4), core_minimum_version=(2, 1, 9))
+def test_upgrade_u2f(gen, from_tag, to_tag):
+    """
+    Check U2F counter stayed the same after an upgrade.
+    """
+    with EmulatorWrapper(gen, from_tag) as emu:
+        success = device.set_u2f_counter(emu.client, 10)
+        assert "U2F counter set" in success
+
+        counter = device.get_next_u2f_counter(emu.client)
+        assert counter == 11
+        storage = emu.storage()
+
+    with EmulatorWrapper(gen, to_tag, storage=storage) as emu:
+        counter = device.get_next_u2f_counter(emu.client)
+        assert counter == 12
+
+
 if __name__ == "__main__":
     if not ALL_TAGS:
         print("No versions found. Remember to run download_emulators.sh")