diff --git a/extmod/moduhashlib.c b/extmod/moduhashlib.c index 195a24334..b170de9e5 100644 --- a/extmod/moduhashlib.c +++ b/extmod/moduhashlib.c @@ -60,9 +60,16 @@ typedef struct _mp_obj_hash_t { mp_obj_base_t base; - char state[0]; + bool final; // if set, update and digest raise an exception + uintptr_t state[0]; // must be aligned to a machine word } mp_obj_hash_t; +static void uhashlib_ensure_not_final(mp_obj_hash_t *self) { + if (self->final) { + mp_raise_ValueError(MP_ERROR_TEXT("hash is final")); + } +} + #if MICROPY_PY_UHASHLIB_SHA256 STATIC mp_obj_t uhashlib_sha256_update(mp_obj_t self_in, mp_obj_t arg); @@ -78,6 +85,7 @@ STATIC mp_obj_t uhashlib_sha256_make_new(const mp_obj_type_t *type, size_t n_arg mp_arg_check_num(n_args, n_kw, 0, 1, false); mp_obj_hash_t *o = m_new_obj_var(mp_obj_hash_t, char, sizeof(mbedtls_sha256_context)); o->base.type = type; + o->final = false; mbedtls_sha256_init((mbedtls_sha256_context *)&o->state); mbedtls_sha256_starts_ret((mbedtls_sha256_context *)&o->state, 0); if (n_args == 1) { @@ -88,6 +96,7 @@ STATIC mp_obj_t uhashlib_sha256_make_new(const mp_obj_type_t *type, size_t n_arg STATIC mp_obj_t uhashlib_sha256_update(mp_obj_t self_in, mp_obj_t arg) { mp_obj_hash_t *self = MP_OBJ_TO_PTR(self_in); + uhashlib_ensure_not_final(self); mp_buffer_info_t bufinfo; mp_get_buffer_raise(arg, &bufinfo, MP_BUFFER_READ); mbedtls_sha256_update_ret((mbedtls_sha256_context *)&self->state, bufinfo.buf, bufinfo.len); @@ -96,6 +105,8 @@ STATIC mp_obj_t uhashlib_sha256_update(mp_obj_t self_in, mp_obj_t arg) { STATIC mp_obj_t uhashlib_sha256_digest(mp_obj_t self_in) { mp_obj_hash_t *self = MP_OBJ_TO_PTR(self_in); + uhashlib_ensure_not_final(self); + self->final = true; vstr_t vstr; vstr_init_len(&vstr, 32); mbedtls_sha256_finish_ret((mbedtls_sha256_context *)&self->state, (unsigned char *)vstr.buf); @@ -110,6 +121,7 @@ STATIC mp_obj_t uhashlib_sha256_make_new(const mp_obj_type_t *type, size_t n_arg mp_arg_check_num(n_args, n_kw, 0, 1, false); mp_obj_hash_t *o = m_new_obj_var(mp_obj_hash_t, char, sizeof(CRYAL_SHA256_CTX)); o->base.type = type; + o->final = false; sha256_init((CRYAL_SHA256_CTX *)o->state); if (n_args == 1) { uhashlib_sha256_update(MP_OBJ_FROM_PTR(o), args[0]); @@ -119,6 +131,7 @@ STATIC mp_obj_t uhashlib_sha256_make_new(const mp_obj_type_t *type, size_t n_arg STATIC mp_obj_t uhashlib_sha256_update(mp_obj_t self_in, mp_obj_t arg) { mp_obj_hash_t *self = MP_OBJ_TO_PTR(self_in); + uhashlib_ensure_not_final(self); mp_buffer_info_t bufinfo; mp_get_buffer_raise(arg, &bufinfo, MP_BUFFER_READ); sha256_update((CRYAL_SHA256_CTX *)self->state, bufinfo.buf, bufinfo.len); @@ -127,6 +140,8 @@ STATIC mp_obj_t uhashlib_sha256_update(mp_obj_t self_in, mp_obj_t arg) { STATIC mp_obj_t uhashlib_sha256_digest(mp_obj_t self_in) { mp_obj_hash_t *self = MP_OBJ_TO_PTR(self_in); + uhashlib_ensure_not_final(self); + self->final = true; vstr_t vstr; vstr_init_len(&vstr, SHA256_BLOCK_SIZE); sha256_final((CRYAL_SHA256_CTX *)self->state, (byte *)vstr.buf); @@ -160,6 +175,7 @@ STATIC mp_obj_t uhashlib_sha1_make_new(const mp_obj_type_t *type, size_t n_args, mp_arg_check_num(n_args, n_kw, 0, 1, false); mp_obj_hash_t *o = m_new_obj_var(mp_obj_hash_t, char, sizeof(SHA1_CTX)); o->base.type = type; + o->final = false; SHA1_Init((SHA1_CTX *)o->state); if (n_args == 1) { uhashlib_sha1_update(MP_OBJ_FROM_PTR(o), args[0]); @@ -169,6 +185,7 @@ STATIC mp_obj_t uhashlib_sha1_make_new(const mp_obj_type_t *type, size_t n_args, STATIC mp_obj_t uhashlib_sha1_update(mp_obj_t self_in, mp_obj_t arg) { mp_obj_hash_t *self = MP_OBJ_TO_PTR(self_in); + uhashlib_ensure_not_final(self); mp_buffer_info_t bufinfo; mp_get_buffer_raise(arg, &bufinfo, MP_BUFFER_READ); SHA1_Update((SHA1_CTX *)self->state, bufinfo.buf, bufinfo.len); @@ -177,6 +194,8 @@ STATIC mp_obj_t uhashlib_sha1_update(mp_obj_t self_in, mp_obj_t arg) { STATIC mp_obj_t uhashlib_sha1_digest(mp_obj_t self_in) { mp_obj_hash_t *self = MP_OBJ_TO_PTR(self_in); + uhashlib_ensure_not_final(self); + self->final = true; vstr_t vstr; vstr_init_len(&vstr, SHA1_SIZE); SHA1_Final((byte *)vstr.buf, (SHA1_CTX *)self->state); @@ -196,6 +215,7 @@ STATIC mp_obj_t uhashlib_sha1_make_new(const mp_obj_type_t *type, size_t n_args, mp_arg_check_num(n_args, n_kw, 0, 1, false); mp_obj_hash_t *o = m_new_obj_var(mp_obj_hash_t, char, sizeof(mbedtls_sha1_context)); o->base.type = type; + o->final = false; mbedtls_sha1_init((mbedtls_sha1_context *)o->state); mbedtls_sha1_starts_ret((mbedtls_sha1_context *)o->state); if (n_args == 1) { @@ -206,6 +226,7 @@ STATIC mp_obj_t uhashlib_sha1_make_new(const mp_obj_type_t *type, size_t n_args, STATIC mp_obj_t uhashlib_sha1_update(mp_obj_t self_in, mp_obj_t arg) { mp_obj_hash_t *self = MP_OBJ_TO_PTR(self_in); + uhashlib_ensure_not_final(self); mp_buffer_info_t bufinfo; mp_get_buffer_raise(arg, &bufinfo, MP_BUFFER_READ); mbedtls_sha1_update_ret((mbedtls_sha1_context *)self->state, bufinfo.buf, bufinfo.len); @@ -214,6 +235,8 @@ STATIC mp_obj_t uhashlib_sha1_update(mp_obj_t self_in, mp_obj_t arg) { STATIC mp_obj_t uhashlib_sha1_digest(mp_obj_t self_in) { mp_obj_hash_t *self = MP_OBJ_TO_PTR(self_in); + uhashlib_ensure_not_final(self); + self->final = true; vstr_t vstr; vstr_init_len(&vstr, 20); mbedtls_sha1_finish_ret((mbedtls_sha1_context *)self->state, (byte *)vstr.buf); @@ -247,6 +270,7 @@ STATIC mp_obj_t uhashlib_md5_make_new(const mp_obj_type_t *type, size_t n_args, mp_arg_check_num(n_args, n_kw, 0, 1, false); mp_obj_hash_t *o = m_new_obj_var(mp_obj_hash_t, char, sizeof(MD5_CTX)); o->base.type = type; + o->final = false; MD5_Init((MD5_CTX *)o->state); if (n_args == 1) { uhashlib_md5_update(MP_OBJ_FROM_PTR(o), args[0]); @@ -256,6 +280,7 @@ STATIC mp_obj_t uhashlib_md5_make_new(const mp_obj_type_t *type, size_t n_args, STATIC mp_obj_t uhashlib_md5_update(mp_obj_t self_in, mp_obj_t arg) { mp_obj_hash_t *self = MP_OBJ_TO_PTR(self_in); + uhashlib_ensure_not_final(self); mp_buffer_info_t bufinfo; mp_get_buffer_raise(arg, &bufinfo, MP_BUFFER_READ); MD5_Update((MD5_CTX *)self->state, bufinfo.buf, bufinfo.len); @@ -264,6 +289,8 @@ STATIC mp_obj_t uhashlib_md5_update(mp_obj_t self_in, mp_obj_t arg) { STATIC mp_obj_t uhashlib_md5_digest(mp_obj_t self_in) { mp_obj_hash_t *self = MP_OBJ_TO_PTR(self_in); + uhashlib_ensure_not_final(self); + self->final = true; vstr_t vstr; vstr_init_len(&vstr, MD5_SIZE); MD5_Final((byte *)vstr.buf, (MD5_CTX *)self->state); @@ -283,6 +310,7 @@ STATIC mp_obj_t uhashlib_md5_make_new(const mp_obj_type_t *type, size_t n_args, mp_arg_check_num(n_args, n_kw, 0, 1, false); mp_obj_hash_t *o = m_new_obj_var(mp_obj_hash_t, char, sizeof(mbedtls_md5_context)); o->base.type = type; + o->final = false; mbedtls_md5_init((mbedtls_md5_context *)o->state); mbedtls_md5_starts_ret((mbedtls_md5_context *)o->state); if (n_args == 1) { @@ -293,6 +321,7 @@ STATIC mp_obj_t uhashlib_md5_make_new(const mp_obj_type_t *type, size_t n_args, STATIC mp_obj_t uhashlib_md5_update(mp_obj_t self_in, mp_obj_t arg) { mp_obj_hash_t *self = MP_OBJ_TO_PTR(self_in); + uhashlib_ensure_not_final(self); mp_buffer_info_t bufinfo; mp_get_buffer_raise(arg, &bufinfo, MP_BUFFER_READ); mbedtls_md5_update_ret((mbedtls_md5_context *)self->state, bufinfo.buf, bufinfo.len); @@ -301,6 +330,8 @@ STATIC mp_obj_t uhashlib_md5_update(mp_obj_t self_in, mp_obj_t arg) { STATIC mp_obj_t uhashlib_md5_digest(mp_obj_t self_in) { mp_obj_hash_t *self = MP_OBJ_TO_PTR(self_in); + uhashlib_ensure_not_final(self); + self->final = true; vstr_t vstr; vstr_init_len(&vstr, 16); mbedtls_md5_finish_ret((mbedtls_md5_context *)self->state, (byte *)vstr.buf); diff --git a/tests/extmod/uhashlib_final.py b/tests/extmod/uhashlib_final.py new file mode 100644 index 000000000..f562cc178 --- /dev/null +++ b/tests/extmod/uhashlib_final.py @@ -0,0 +1,35 @@ +try: + import uhashlib +except ImportError: + print("SKIP") + raise SystemExit + + +for algo_name in ("md5", "sha1", "sha256"): + algo = getattr(uhashlib, algo_name, None) + if not algo: + continue + + # Running .digest() several times in row is not supported. + h = algo(b"123") + h.digest() + try: + h.digest() + print("fail") + except ValueError: + # Expected path, don't print anything so test output is the + # same even if the algorithm is not implemented on the port. + pass + + # Partial digests are not supported. + h = algo(b"123") + h.digest() + try: + h.update(b"456") + print("fail") + except ValueError: + # Expected path, don't print anything so test output is the + # same even if the algorithm is not implemented on the port. + pass + +print("done") diff --git a/tests/extmod/uhashlib_final.py.exp b/tests/extmod/uhashlib_final.py.exp new file mode 100644 index 000000000..19f86f493 --- /dev/null +++ b/tests/extmod/uhashlib_final.py.exp @@ -0,0 +1 @@ +done