diff --git a/embed/extmod/modtrezorcrypto/modtrezorcrypto-blake2b.h b/embed/extmod/modtrezorcrypto/modtrezorcrypto-blake2b.h index 27e1fb23c0..a686508a32 100644 --- a/embed/extmod/modtrezorcrypto/modtrezorcrypto-blake2b.h +++ b/embed/extmod/modtrezorcrypto/modtrezorcrypto-blake2b.h @@ -33,21 +33,29 @@ typedef struct _mp_obj_Blake2b_t { STATIC mp_obj_t mod_trezorcrypto_Blake2b_update(mp_obj_t self, mp_obj_t data); -/// def __init__(self, data: bytes = None, key: bytes = None) -> None: +/// def __init__(self, data: bytes = None, outlen: int = Blake2b.digest_size, key: bytes = None) -> None: /// ''' /// Creates a hash context object. /// ''' STATIC mp_obj_t mod_trezorcrypto_Blake2b_make_new(const mp_obj_type_t *type, size_t n_args, size_t n_kw, const mp_obj_t *args) { - mp_arg_check_num(n_args, n_kw, 0, 2, false); + mp_arg_check_num(n_args, n_kw, 0, 3, false); mp_obj_Blake2b_t *o = m_new_obj(mp_obj_Blake2b_t); o->base.type = type; + int res = 0; // constructor called with key argument set - if (n_args == 2) { + if (n_args == 3) { + size_t outlen = mp_obj_get_int(args[1]); mp_buffer_info_t key; - mp_get_buffer_raise(args[1], &key, MP_BUFFER_READ); - blake2b_InitKey(&(o->ctx), BLAKE2B_DIGEST_LENGTH, key.buf, key.len); + mp_get_buffer_raise(args[2], &key, MP_BUFFER_READ); + res = blake2b_InitKey(&(o->ctx), outlen, key.buf, key.len); + } else if (n_args == 2) { + size_t outlen = mp_obj_get_int(args[1]); + res = blake2b_Init(&(o->ctx), outlen); } else { - blake2b_Init(&(o->ctx), BLAKE2B_DIGEST_LENGTH); + res = blake2b_Init(&(o->ctx), BLAKE2B_DIGEST_LENGTH); + } + if (res < 0) { + mp_raise_ValueError("Invalid Blake2b parameters"); } // constructor called with data argument set if (n_args >= 1) { @@ -80,9 +88,9 @@ STATIC mp_obj_t mod_trezorcrypto_Blake2b_digest(mp_obj_t self) { uint8_t out[BLAKE2B_DIGEST_LENGTH]; BLAKE2B_CTX ctx; memcpy(&ctx, &(o->ctx), sizeof(BLAKE2B_CTX)); - blake2b_Final(&ctx, out, BLAKE2B_DIGEST_LENGTH); + blake2b_Final(&ctx, out, ctx.outlen); memset(&ctx, 0, sizeof(BLAKE2B_CTX)); - return mp_obj_new_bytes(out, sizeof(out)); + return mp_obj_new_bytes(out, o->ctx.outlen); } STATIC MP_DEFINE_CONST_FUN_OBJ_1(mod_trezorcrypto_Blake2b_digest_obj, mod_trezorcrypto_Blake2b_digest); diff --git a/embed/extmod/modtrezorcrypto/modtrezorcrypto-blake2s.h b/embed/extmod/modtrezorcrypto/modtrezorcrypto-blake2s.h index 927c31a6a2..79ce0f326d 100644 --- a/embed/extmod/modtrezorcrypto/modtrezorcrypto-blake2s.h +++ b/embed/extmod/modtrezorcrypto/modtrezorcrypto-blake2s.h @@ -33,21 +33,29 @@ typedef struct _mp_obj_Blake2s_t { STATIC mp_obj_t mod_trezorcrypto_Blake2s_update(mp_obj_t self, mp_obj_t data); -/// def __init__(self, data: bytes = None, key: bytes = None) -> None: +/// def __init__(self, data: bytes = None, outlen: int = Blake2s.digest_size, key: bytes = None) -> None: /// ''' /// Creates a hash context object. /// ''' STATIC mp_obj_t mod_trezorcrypto_Blake2s_make_new(const mp_obj_type_t *type, size_t n_args, size_t n_kw, const mp_obj_t *args) { - mp_arg_check_num(n_args, n_kw, 0, 2, false); + mp_arg_check_num(n_args, n_kw, 0, 3, false); mp_obj_Blake2s_t *o = m_new_obj(mp_obj_Blake2s_t); o->base.type = type; + int res = 0; // constructor called with key argument set - if (n_args == 2) { + if (n_args == 3) { + size_t outlen = mp_obj_get_int(args[1]); mp_buffer_info_t key; - mp_get_buffer_raise(args[1], &key, MP_BUFFER_READ); - blake2s_InitKey(&(o->ctx), BLAKE2S_DIGEST_LENGTH, key.buf, key.len); + mp_get_buffer_raise(args[2], &key, MP_BUFFER_READ); + res = blake2s_InitKey(&(o->ctx), outlen, key.buf, key.len); + } else if (n_args == 2) { + size_t outlen = mp_obj_get_int(args[1]); + res = blake2s_Init(&(o->ctx), outlen); } else { - blake2s_Init(&(o->ctx), BLAKE2S_DIGEST_LENGTH); + res = blake2s_Init(&(o->ctx), BLAKE2S_DIGEST_LENGTH); + } + if (res < 0) { + mp_raise_ValueError("Invalid Blake2s parameters"); } // constructor called with data argument set if (n_args >= 1) { @@ -80,9 +88,9 @@ STATIC mp_obj_t mod_trezorcrypto_Blake2s_digest(mp_obj_t self) { uint8_t out[BLAKE2S_DIGEST_LENGTH]; BLAKE2S_CTX ctx; memcpy(&ctx, &(o->ctx), sizeof(BLAKE2S_CTX)); - blake2s_Final(&ctx, out, BLAKE2S_DIGEST_LENGTH); + blake2s_Final(&ctx, out, ctx.outlen); memset(&ctx, 0, sizeof(BLAKE2S_CTX)); - return mp_obj_new_bytes(out, sizeof(out)); + return mp_obj_new_bytes(out, o->ctx.outlen); } STATIC MP_DEFINE_CONST_FUN_OBJ_1(mod_trezorcrypto_Blake2s_digest_obj, mod_trezorcrypto_Blake2s_digest); diff --git a/tests/test_trezor.crypto.hashlib.blake2b.py b/tests/test_trezor.crypto.hashlib.blake2b.py index 1f66a631b0..20255e7dc1 100644 --- a/tests/test_trezor.crypto.hashlib.blake2b.py +++ b/tests/test_trezor.crypto.hashlib.blake2b.py @@ -18,11 +18,11 @@ class TestCryptoBlake2b(unittest.TestCase): def test_digest(self): key = unhexlify('000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f202122232425262728292a2b2c2d2e2f303132333435363738393a3b3c3d3e3f') for d, h in self.vectors: - self.assertEqual(hashlib.blake2b(unhexlify(d), key).digest(), unhexlify(h)) + self.assertEqual(hashlib.blake2b(unhexlify(d), hashlib.blake2b.digest_size, key).digest(), unhexlify(h)) def test_update(self): key = unhexlify('000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f202122232425262728292a2b2c2d2e2f303132333435363738393a3b3c3d3e3f') - x = hashlib.blake2b(b'', key) + x = hashlib.blake2b(b'', hashlib.blake2b.digest_size, key) x.update(bytes(range(10))) self.assertEqual(x.digest(), unhexlify('4fe181f54ad63a2983feaaf77d1e7235c2beb17fa328b6d9505bda327df19fc37f02c4b6f0368ce23147313a8e5738b5fa2a95b29de1c7f8264eb77b69f585cd')) x.update(bytes(range(10, 30))) diff --git a/tests/test_trezor.crypto.hashlib.blake2s.py b/tests/test_trezor.crypto.hashlib.blake2s.py index 0d4d229a46..6f389b6f3c 100644 --- a/tests/test_trezor.crypto.hashlib.blake2s.py +++ b/tests/test_trezor.crypto.hashlib.blake2s.py @@ -19,11 +19,11 @@ class TestCryptoBlake2s(unittest.TestCase): def test_digest(self): key = unhexlify('000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f') for d, h in self.vectors: - self.assertEqual(hashlib.blake2s(unhexlify(d), key).digest(), unhexlify(h)) + self.assertEqual(hashlib.blake2s(unhexlify(d), hashlib.blake2s.digest_size, key).digest(), unhexlify(h)) def test_update(self): key = unhexlify('000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f') - x = hashlib.blake2s(b'', key) + x = hashlib.blake2s(b'', hashlib.blake2s.digest_size, key) x.update(bytes(range(10))) self.assertEqual(x.digest(), unhexlify('f5c4b2ba1a00781b13aba0425242c69cb1552f3f71a9a3bb22b4a6b4277b46dd')) x.update(bytes(range(10, 30)))