core/sd_salt: decorate sd_salt functions to ensure proper finalization and power-off

core/src/apps/common/sd_salt.py

@@ -1,5 +1,5 @@
 import storage.sd_salt
-from storage.sd_salt import SD_CARD_HOT_SWAPPABLE, SdSaltMismatch
+from storage.sd_salt import SD_CARD_HOT_SWAPPABLE
 from trezor import io, ui, wire
 from trezor.ui.text import Text
 
@@ -29,7 +29,7 @@ async def _wrong_card_dialog(ctx: wire.GenericContext) -> bool:
     return await confirm(ctx, text, confirm=btn_confirm, cancel=btn_cancel)
 
 
-async def _insert_card_dialog(ctx: wire.GenericContext) -> None:
+async def insert_card_dialog(ctx: wire.GenericContext) -> bool:
     text = Text("SD card protection", ui.ICON_WRONG)
     text.bold("SD card required.")
     text.br_half()
@@ -42,20 +42,20 @@ async def _insert_card_dialog(ctx: wire.GenericContext) -> None:
         btn_confirm = None
         btn_cancel = "Close"
 
-    if not await confirm(ctx, text, confirm=btn_confirm, cancel=btn_cancel):
-        raise SdProtectCancelled
+    return await confirm(ctx, text, confirm=btn_confirm, cancel=btn_cancel)
 
 
-async def sd_write_failed_dialog(ctx: wire.GenericContext) -> bool:
+async def sd_problem_dialog(ctx: wire.GenericContext) -> bool:
     text = Text("SD card protection", ui.ICON_WRONG, ui.RED)
-    text.normal("Failed to write data to", "the SD card.")
+    text.normal("There was a problem", "accessing the SD card.")
     return await confirm(ctx, text, confirm="Retry", cancel="Abort")
 
 
 async def ensure_sd_card(ctx: wire.GenericContext) -> None:
     sd = io.SDCard()
-    while not sd.power(True):
-        await _insert_card_dialog(ctx)
+    while not sd.present():
+        if not await insert_card_dialog(ctx):
+            raise SdProtectCancelled
 
 
 async def request_sd_salt(
@@ -65,11 +65,14 @@ async def request_sd_salt(
         ensure_sd_card(ctx)
         try:
             return storage.sd_salt.load_sd_salt()
-        except SdSaltMismatch as e:
+        except storage.sd_salt.WrongSdCard:
             if not await _wrong_card_dialog(ctx):
-                raise SdProtectCancelled from e
+                raise SdProtectCancelled
         except OSError:
-            # This happens when there is both old and new salt file, and we can't move
-            # new salt over the old salt. If the user clicks Retry, we will try again.
-            if not await sd_write_failed_dialog(ctx):
+            # Either the SD card did not power on, or the filesystem could not be
+            # mounted (card is not formatted?), or there is a staged salt file and
+            # we could not commit it.
+            # In either case, there is no good way to recover. If the user clicks Retry,
+            # we will try again.
+            if not await sd_problem_dialog(ctx):
                 raise

core/src/apps/management/sd_protect.py

@@ -14,7 +14,7 @@ from apps.common.request_pin import (
     request_pin_and_sd_salt,
     show_pin_invalid,
 )
-from apps.common.sd_salt import ensure_sd_card, sd_write_failed_dialog
+from apps.common.sd_salt import ensure_sd_card, sd_problem_dialog
 
 if False:
     from typing import Awaitable, Tuple
@@ -32,10 +32,11 @@ async def _set_salt(
     ctx: wire.Context, salt: bytes, salt_tag: bytes, stage: bool = False
 ) -> None:
     while True:
+        ensure_sd_card(ctx)
         try:
             return storage.sd_salt.set_sd_salt(salt, salt_tag, stage)
         except OSError:
-            if not await sd_write_failed_dialog(ctx):
+            if not await sd_problem_dialog(ctx):
                 raise
 
 
@@ -60,7 +61,7 @@ async def sd_protect_enable(ctx: wire.Context, msg: SdProtect) -> Success:
     # Confirm that user wants to proceed with the operation.
     await require_confirm_sd_protect(ctx, msg)
 
-    # Make sure SD card is available.
+    # Make sure SD card is present.
     await ensure_sd_card(ctx)
 
     # Get the current PIN.
@@ -95,7 +96,7 @@ async def sd_protect_disable(ctx: wire.Context, msg: SdProtect) -> Success:
     if not storage.sd_salt.is_enabled():
         raise wire.ProcessError("SD card protection not enabled")
 
-    # Note that the SD card doesn't need to be accessible in order to disable SD
+    # Note that the SD card doesn't need to be present in order to disable SD
     # protection. The cleanup will not happen in such case, but that does not matter.
 
     # Confirm that user wants to proceed with the operation.
@@ -131,7 +132,7 @@ async def sd_protect_refresh(ctx: wire.Context, msg: SdProtect) -> Success:
     # Confirm that user wants to proceed with the operation.
     await require_confirm_sd_protect(ctx, msg)
 
-    # Make sure SD card is available.
+    # Make sure SD card is present.
     await ensure_sd_card(ctx)
 
     # Get the current PIN and salt from the SD card.

core/src/storage/sd_salt.py

@@ -7,14 +7,16 @@ from trezor.crypto.hashlib import sha256
 from trezor.utils import consteq
 
 if False:
-    from typing import Optional
+    from typing import Optional, TypeVar, Callable
+
+    T = TypeVar("T", bound=Callable)
 
 SD_CARD_HOT_SWAPPABLE = False
 SD_SALT_LEN_BYTES = const(32)
 SD_SALT_AUTH_TAG_LEN_BYTES = const(16)
 
 
-class SdSaltMismatch(Exception):
+class WrongSdCard(Exception):
     pass
 
 
@@ -35,6 +37,48 @@ def _get_salt_path(new: bool = False) -> str:
     return "{}/salt{}".format(_get_device_dir(), ".new" if new else "")
 
 
+_ensure_filesystem_nesting_counter = 0
+
+
+def ensure_filesystem(func: T) -> T:
+    """Ensure the decorated function has access to SD card filesystem.
+
+    Usage:
+    >>> @ensure_filesystem
+    >>> def do_something(arg):
+    >>>     fs = io.FatFS()
+    >>>     # the decorator guarantees that `fs` is mounted
+    >>>     fs.unlink("/dir/" + arg)
+    """
+    # XXX
+    # A slightly better design would be to make the decorated function take the `fs`
+    # as an argument, but that is currently untypeable with mypy.
+    # (see https://github.com/python/mypy/issues/3157)
+    def wrapped_func(*args, **kwargs):  # type: ignore
+        global _ensure_filesystem_nesting_counter
+
+        sd = io.SDCard()
+        if _ensure_filesystem_nesting_counter == 0:
+            if not sd.power(True):
+                raise OSError
+
+        try:
+            _ensure_filesystem_nesting_counter += 1
+            fs = io.FatFS()
+            fs.mount()
+            # XXX do we need to differentiate failure types?
+            # If yes, can the problem be derived from the type of OSError raised?
+            return func(*args, **kwargs)
+        finally:
+            _ensure_filesystem_nesting_counter -= 1
+            assert _ensure_filesystem_nesting_counter >= 0
+            if _ensure_filesystem_nesting_counter == 0:
+                fs.unmount()
+                sd.power(False)
+
+    return wrapped_func  # type: ignore
+
+
 def _load_salt(fs: io.FatFS, auth_key: bytes, path: str) -> Optional[bytearray]:
     # Load the salt file if it exists.
     try:
@@ -54,107 +98,70 @@ def _load_salt(fs: io.FatFS, auth_key: bytes, path: str) -> Optional[bytearray]:
     return salt
 
 
+@ensure_filesystem
 def load_sd_salt() -> Optional[bytearray]:
     salt_auth_key = storage.device.get_sd_salt_auth_key()
     if salt_auth_key is None:
         return None
 
-    sd = io.SDCard()
-    if not sd.power(True):
-        raise OSError
-
     salt_path = _get_salt_path()
     new_salt_path = _get_salt_path(new=True)
 
-    try:
-        fs = io.FatFS()
-        try:
-            fs.mount()
-        except OSError as e:
-            # SD card is probably not formatted. For purposes of loading SD salt, this
-            # is identical to having the wrong card in.
-            raise SdSaltMismatch from e
+    fs = io.FatFS()
 
-        salt = _load_salt(fs, salt_auth_key, salt_path)
-        if salt is not None:
-            return salt
-
-        # Check if there is a new salt.
-        salt = _load_salt(fs, salt_auth_key, new_salt_path)
-        if salt is None:
-            # No valid salt file on this SD card.
-            raise SdSaltMismatch
-
-        # Normal salt file does not exist, but new salt file exists. That means that
-        # 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 can fail with a write error, which falls through as an OSError.
-        # This should be handled in calling code, by allowing the user to retry.
-        fs.rename(new_salt_path, salt_path)
+    salt = _load_salt(fs, salt_auth_key, salt_path)
+    if salt is not None:
         return salt
-    finally:
-        fs.unmount()
-        sd.power(False)
+
+    # Check if there is a new salt.
+    salt = _load_salt(fs, salt_auth_key, new_salt_path)
+    if salt is None:
+        # No valid salt file on this SD card.
+        raise WrongSdCard
+
+    # Normal salt file does not exist, but new salt file exists. That means that
+    # 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 can fail with a write error, which falls through as an OSError.
+    # This should be handled in calling code, by allowing the user to retry.
+    fs.rename(new_salt_path, salt_path)
+    return salt
 
 
+@ensure_filesystem
 def set_sd_salt(salt: bytes, salt_tag: bytes, stage: bool = False) -> None:
     salt_path = _get_salt_path(stage)
-    sd = io.SDCard()
-    if not sd.power(True):
-        raise OSError
-
-    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)
-    finally:
-        fs.unmount()
-        sd.power(False)
+    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)
 
 
+@ensure_filesystem
 def commit_sd_salt() -> None:
     salt_path = _get_salt_path(new=False)
     new_salt_path = _get_salt_path(new=True)
 
-    sd = io.SDCard()
     fs = io.FatFS()
-    if not sd.power(True):
-        raise OSError
-
     try:
-        fs.mount()
-        # TODO Possibly overwrite salt file with random data.
-        try:
-            fs.unlink(salt_path)
-        except OSError:
-            pass
-        fs.rename(new_salt_path, salt_path)
-    finally:
-        fs.unmount()
-        sd.power(False)
+        fs.unlink(salt_path)
+    except OSError:
+        pass
+    fs.rename(new_salt_path, salt_path)
 
 
+@ensure_filesystem
 def remove_sd_salt() -> None:
     salt_path = _get_salt_path()
 
-    sd = io.SDCard()
     fs = io.FatFS()
-    if not sd.power(True):
-        raise OSError
-
-    try:
-        fs.mount()
-        # TODO Possibly overwrite salt file with random data.
-        fs.unlink(salt_path)
-    finally:
-        fs.unmount()
-        sd.power(False)
+    # TODO Possibly overwrite salt file with random data.
+    fs.unlink(salt_path)