fix last commit

extmod/modTrezorCrypto/modTrezorCrypto-base58.h

@@ -24,8 +24,11 @@ STATIC mp_obj_t mod_TrezorCrypto_Base58_encode(mp_obj_t self, mp_obj_t data) {
     mp_buffer_info_t databuf;
     mp_get_buffer_raise(data, &databuf, MP_BUFFER_READ);
     vstr_t vstr;
-    vstr_init(&vstr, databuf.len * 8000 / 5857 + 1); // 256 = 2^8 ; 58 > 2^5.857
-    b58enc(vstr.buf, &vstr.len, databuf.buf, databuf.len);
+    vstr_init_len(&vstr, databuf.len * 8000 / 5857 + 1); // 256 = 2^8 ; 58 > 2^5.857
+    bool r = b58enc(vstr.buf, &vstr.len, databuf.buf, databuf.len);
+    if (!r) {
+        nlr_raise(mp_obj_new_exception_msg(&mp_type_ValueError, "Invalid input"));
+    }
     vstr.len--; // b58enc returns length including the trailing zero
     return mp_obj_new_str_from_vstr(&mp_type_str, &vstr);
 }
@@ -36,8 +39,11 @@ STATIC mp_obj_t mod_TrezorCrypto_Base58_decode(mp_obj_t self, mp_obj_t string) {
     mp_buffer_info_t stringbuf;
     mp_get_buffer_raise(string, &stringbuf, MP_BUFFER_READ);
     vstr_t vstr;
-    vstr_init(&vstr, stringbuf.len * 5858 / 8000 + 1); // 256 = 2^8 ; 58 < 2^5.858
-    b58tobin(vstr.buf, &vstr.len, stringbuf.buf);
+    vstr_init_len(&vstr, stringbuf.len * 5858 / 8000 + 1); // 256 = 2^8 ; 58 < 2^5.858
+    bool r = b58tobin(vstr.buf, &vstr.len, stringbuf.buf);
+    if (!r) {
+        nlr_raise(mp_obj_new_exception_msg(&mp_type_ValueError, "Invalid input"));
+    }
     return mp_obj_new_str_from_vstr(&mp_type_bytes, &vstr);
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_2(mod_TrezorCrypto_Base58_decode_obj, mod_TrezorCrypto_Base58_decode);

src/tests/test_crypto_base58.py

@@ -64,11 +64,11 @@ class TestCryptoBase58(unittest.TestCase):
 
     def test_decode_check(self):
         for a, b in self.vectors:
-            self.assertEqual(trezor.crypto.base58.decode_check(b), trezor.utils.unhexlify('a'))
+            self.assertEqual(trezor.crypto.base58.decode_check(b), trezor.utils.unhexlify(a))
 
     def test_encode_check(self):
         for a, b in self.vectors:
-            self.assertEqual(trezor.crypto.base58.encode_check(trezor.utils.unhexlify('a')), b)
+            self.assertEqual(trezor.crypto.base58.encode_check(trezor.utils.unhexlify(a)), b)
 
 if __name__ == '__main__':
     unittest.main()

src/trezor/crypto/base58.py

@@ -10,12 +10,13 @@ def decode(string):
     return _base58.decode(string)
 
 def encode_check(data, hashlen=4):
-    h = sha256.hash(data)
+    h = sha256.hash(sha256.hash(data))
     return encode(data + h[:hashlen])
 
 def decode_check(string, hashlen=4):
     data = decode(string)
-    d, h = data[:-hashlen], data[-hashlen:]
-    if sha256.hash(d) != h:
+    d, h1 = data[:-hashlen], data[-hashlen:]
+    h2 = sha256.hash(sha256.hash(d))[:4]
+    if h1 != h2:
         raise RuntimeError('Checksum error')
     return d