Merge pull request #612 from trezor/andrewkozlik/sd-protect-ui

SD-protect error handling

core/src/apps/common/sd_salt.py

@@ -3,12 +3,12 @@ from micropython import const
 from trezor import io, ui, wire
 from trezor.crypto import hmac
 from trezor.crypto.hashlib import sha256
-from trezor.ui.confirm import Confirm
+from trezor.ui.confirm import CONFIRMED, Confirm
 from trezor.ui.text import Text
 from trezor.utils import consteq
 
 from apps.common import storage
-from apps.common.confirm import require_confirm
+from apps.common.confirm import confirm
 
 if False:
     from typing import Optional
@@ -18,6 +18,7 @@ class SdProtectCancelled(Exception):
     pass
 
 
+SD_CARD_HOT_SWAPPABLE = False
 SD_SALT_LEN_BYTES = const(32)
 SD_SALT_AUTH_TAG_LEN_BYTES = const(16)
 SD_SALT_AUTH_KEY_LEN_BYTES = const(16)
@@ -27,22 +28,53 @@ async def _wrong_card_dialog(ctx: Optional[wire.Context]) -> None:
     text = Text("SD card protection", ui.ICON_WRONG)
     text.bold("Wrong SD card.")
     text.br_half()
-    text.normal("Please unplug the", "device and insert a", "different card.")
+    if SD_CARD_HOT_SWAPPABLE:
-    if ctx is None:
+        text.normal("Please insert the", "correct SD card for", "this device.")
-        await Confirm(text, confirm=None, cancel="Close")
+        btn_confirm = "Retry"
+        btn_cancel = "Abort"
     else:
-        await require_confirm(ctx, text, confirm=None, cancel="Close")
+        text.normal("Please unplug the", "device and insert the", "correct SD card.")
+        btn_confirm = None
+        btn_cancel = "Close"
+
+    if ctx is None:
+        if await Confirm(text, confirm=btn_confirm, cancel=btn_cancel) is not CONFIRMED:
+            raise SdProtectCancelled
+    else:
+        if not await confirm(ctx, text, confirm=btn_confirm, cancel=btn_cancel):
+            raise wire.ProcessError("Wrong SD card.")
 
 
 async def _insert_card_dialog(ctx: Optional[wire.Context]) -> None:
-    text = Text("SD card protection")
+    text = Text("SD card protection", ui.ICON_WRONG)
     text.bold("SD card required.")
     text.br_half()
-    text.normal("Please unplug the", "device and insert your", "SD card.")
+    if SD_CARD_HOT_SWAPPABLE:
-    if ctx is None:
+        text.normal("Please insert your", "SD card.")
-        await Confirm(text, confirm=None, cancel="Close")
+        btn_confirm = "Retry"
+        btn_cancel = "Abort"
     else:
-        await require_confirm(ctx, text, confirm=None, cancel="Close")
+        text.normal("Please unplug the", "device and insert your", "SD card.")
+        btn_confirm = None
+        btn_cancel = "Close"
+
+    if ctx is None:
+        if await Confirm(text, confirm=btn_confirm, cancel=btn_cancel) is not CONFIRMED:
+            raise SdProtectCancelled
+    else:
+        if not await confirm(ctx, text, confirm=btn_confirm, cancel=btn_cancel):
+            raise wire.ProcessError("SD card required.")
+
+
+async def _write_failed_dialog(ctx: Optional[wire.Context]) -> None:
+    text = Text("SD card protection", ui.ICON_WRONG, ui.RED)
+    text.normal("Failed to write data to", "the SD card.")
+    if ctx is None:
+        if await Confirm(text, confirm="Retry", cancel="Abort") is not CONFIRMED:
+            raise OSError
+    else:
+        if not await confirm(ctx, text, confirm="Retry", cancel="Abort"):
+            raise wire.ProcessError("Failed to write to SD card.")
 
 
 def _get_device_dir() -> str:
@@ -62,63 +94,64 @@ async def request_sd_salt(
     salt_path = _get_salt_path()
     new_salt_path = _get_salt_path(True)
 
-    sd = io.SDCard()
+    while True:
-    fs = io.FatFS()
+        sd = io.SDCard()
-    if not sd.power(True):
+        fs = io.FatFS()
-        await _insert_card_dialog(ctx)
+        while not sd.power(True):
-        raise SdProtectCancelled
+            await _insert_card_dialog(ctx)
 
-    try:
-        fs.mount()
-
-        # Load salt if it exists.
-        salt = None  # type: Optional[bytearray]
         try:
-            with fs.open(salt_path, "r") as f:
+            fs.mount()
-                salt = bytearray(SD_SALT_LEN_BYTES)
-                salt_tag = bytearray(SD_SALT_AUTH_TAG_LEN_BYTES)
-                f.read(salt)
-                f.read(salt_tag)
-        except OSError:
-            salt = None
 
-        if salt is not None and consteq(
+            # Load salt if it exists.
-            hmac.new(salt_auth_key, salt, sha256).digest()[:SD_SALT_AUTH_TAG_LEN_BYTES],
+            salt = None  # type: Optional[bytearray]
-            salt_tag,
-        ):
-            return salt
-
-        # Load salt.new if it exists.
-        new_salt = None  # type: Optional[bytearray]
-        try:
-            with fs.open(new_salt_path, "r") as f:
-                new_salt = bytearray(SD_SALT_LEN_BYTES)
-                new_salt_tag = bytearray(SD_SALT_AUTH_TAG_LEN_BYTES)
-                f.read(new_salt)
-                f.read(new_salt_tag)
-        except OSError:
-            new_salt = None
-
-        if new_salt is not None and consteq(
-            hmac.new(salt_auth_key, new_salt, sha256).digest()[
-                :SD_SALT_AUTH_TAG_LEN_BYTES
-            ],
-            new_salt_tag,
-        ):
-            # SD salt regeneration was interrupted earlier. Bring into consistent state.
-            # TODO Possibly overwrite salt file with random data.
             try:
-                fs.unlink(salt_path)
+                with fs.open(salt_path, "r") as f:
+                    salt = bytearray(SD_SALT_LEN_BYTES)
+                    salt_tag = bytearray(SD_SALT_AUTH_TAG_LEN_BYTES)
+                    f.read(salt)
+                    f.read(salt_tag)
             except OSError:
-                pass
+                salt = None
-            fs.rename(new_salt_path, salt_path)
-            return new_salt
-    finally:
-        fs.unmount()
-        sd.power(False)
 
-    await _wrong_card_dialog(ctx)
+            if salt is not None and consteq(
-    raise SdProtectCancelled
+                hmac.new(salt_auth_key, salt, sha256).digest()[
+                    :SD_SALT_AUTH_TAG_LEN_BYTES
+                ],
+                salt_tag,
+            ):
+                return salt
+
+            # Load salt.new if it exists.
+            new_salt = None  # type: Optional[bytearray]
+            try:
+                with fs.open(new_salt_path, "r") as f:
+                    new_salt = bytearray(SD_SALT_LEN_BYTES)
+                    new_salt_tag = bytearray(SD_SALT_AUTH_TAG_LEN_BYTES)
+                    f.read(new_salt)
+                    f.read(new_salt_tag)
+            except OSError:
+                new_salt = None
+
+            if new_salt is not None and consteq(
+                hmac.new(salt_auth_key, new_salt, sha256).digest()[
+                    :SD_SALT_AUTH_TAG_LEN_BYTES
+                ],
+                new_salt_tag,
+            ):
+                # SD salt regeneration was interrupted earlier. Bring into consistent state.
+                # TODO Possibly overwrite salt file with random data.
+                try:
+                    fs.unlink(salt_path)
+                except OSError:
+                    pass
+                fs.rename(new_salt_path, salt_path)
+                return new_salt
+        finally:
+            fs.unmount()
+            sd.power(False)
+
+        await _wrong_card_dialog(ctx)
 
 
 async def set_sd_salt(
@@ -126,23 +159,27 @@ async def set_sd_salt(
 ) -> None:
     salt_path = _get_salt_path(new)
 
-    sd = io.SDCard()
+    while True:
-    if not sd.power(True):
+        sd = io.SDCard()
-        await _insert_card_dialog(ctx)
+        while not sd.power(True):
-        raise SdProtectCancelled
+            await _insert_card_dialog(ctx)
 
-    fs = io.FatFS()
+        try:
+            fs = io.FatFS()
+            fs.mount()
+            fs.mkdir("/trezor", True)
+            fs.mkdir(_get_device_dir(), True)
+            with fs.open(salt_path, "w") as f:
+                f.write(salt)
+                f.write(salt_tag)
+            break
+        except Exception:
+            fs.unmount()
+            sd.power(False)
+            await _write_failed_dialog(ctx)
 
-    try:
+    fs.unmount()
-        fs.mount()
+    sd.power(False)
-        fs.mkdir("/trezor", True)
-        fs.mkdir(_get_device_dir(), True)
-        with fs.open(salt_path, "w") as f:
-            f.write(salt)
-            f.write(salt_tag)
-    finally:
-        fs.unmount()
-        sd.power(False)
 
 
 async def stage_sd_salt(
@@ -158,8 +195,7 @@ async def commit_sd_salt(ctx: Optional[wire.Context]) -> None:
     sd = io.SDCard()
     fs = io.FatFS()
     if not sd.power(True):
-        await _insert_card_dialog(ctx)
+        raise IOError
-        raise SdProtectCancelled
 
     try:
         fs.mount()
@@ -180,8 +216,7 @@ async def remove_sd_salt(ctx: Optional[wire.Context]) -> None:
     sd = io.SDCard()
     fs = io.FatFS()
     if not sd.power(True):
-        await _insert_card_dialog(ctx)
+        raise IOError
-        raise SdProtectCancelled
 
     try:
         fs.mount()

core/src/apps/management/sd_protect.py

@@ -62,10 +62,7 @@ async def sd_protect_enable(ctx: wire.Context, msg: SdProtect) -> Success:
     salt_tag = hmac.new(salt_auth_key, salt, sha256).digest()[
         :SD_SALT_AUTH_TAG_LEN_BYTES
     ]
-    try:
+    await set_sd_salt(ctx, salt, salt_tag)
-        await set_sd_salt(ctx, salt, salt_tag)
-    except Exception:
-        raise wire.ProcessError("Failed to write to SD card")
 
     if not config.change_pin(pin, pin, None, salt):
         # Wrong PIN. Clean up the prepared salt file.
@@ -131,10 +128,7 @@ async def sd_protect_refresh(ctx: wire.Context, msg: SdProtect) -> Success:
     new_salt_tag = hmac.new(new_salt_auth_key, new_salt, sha256).digest()[
         :SD_SALT_AUTH_TAG_LEN_BYTES
     ]
-    try:
+    await stage_sd_salt(ctx, new_salt, new_salt_tag)
-        await stage_sd_salt(ctx, new_salt, new_salt_tag)
-    except Exception:
-        raise wire.ProcessError("Failed to write to SD card")
 
     if not config.change_pin(pin_to_int(pin), pin_to_int(pin), old_salt, new_salt):
         await show_pin_invalid(ctx)