core: use filesystem wrapper instead of the ensure_filesystem decorator

core/src/storage/sd_salt.py

@@ -1,12 +1,13 @@
 from micropython import const
 
 import storage.device
-from trezor import io
 from trezor.crypto import hmac
 from trezor.crypto.hashlib import sha256
+from trezor.sdcard import get_filesystem
 from trezor.utils import consteq
 
 if False:
+    from trezor import io
     from typing import Optional, TypeVar, Callable
 
     T = TypeVar("T", bound=Callable)
@@ -37,48 +38,6 @@ 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:
@@ -98,7 +57,6 @@ 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:
@@ -107,60 +65,55 @@ def load_sd_salt() -> Optional[bytearray]:
     salt_path = _get_salt_path()
     new_salt_path = _get_salt_path(new=True)
 
-    fs = io.FatFS()
+    with get_filesystem() as fs:
+        salt = _load_salt(fs, salt_auth_key, salt_path)
+        if salt is not None:
+            return salt
 
-    salt = _load_salt(fs, salt_auth_key, salt_path)
-    if salt is not None:
+        # 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
 
-    # 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)
-    fs = io.FatFS()
-    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)
+    with get_filesystem() as fs:
+        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)
 
-    fs = io.FatFS()
-    try:
-        fs.unlink(salt_path)
-    except OSError:
-        pass
-    fs.rename(new_salt_path, salt_path)
+    with get_filesystem() as fs:
+        try:
+            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()
-
-    fs = io.FatFS()
-    # TODO Possibly overwrite salt file with random data.
-    fs.unlink(salt_path)
+    with get_filesystem() as fs:
+        # TODO Possibly overwrite salt file with random data.
+        fs.unlink(salt_path)